[MLIR][Spirv] Don't lower tensors that can't be represented by an ArrayType

[sweiglbosker]

I noticed this because of #159738, though it was only caught by his fuzzer because it wrapped to 0.
Also, is there a reason for the usage of unsigned for sizes in spirv types? I believe most of the builtin types use int64_t for sizes, so it may make sense to do the same for spirv.

[llvmbot]

@llvm/pr-subscribers-mlir-spirv

Author: Stefan Weigl-Bosker (sweiglbosker)

Changes

I noticed this because of #159738, though it was only caught by his fuzzer because it wrapped to 0.
Also, is there a reason for the usage of unsigned for sizes in spirv types? I believe most of the builtin types use int64_t for sizes, so it may make sense to do the same for spirv.

---

Full diff: https://github.com/llvm/llvm-project/pull/171002.diff

1 Files Affected:

• (modified) mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp (+5)

diff --git a/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp b/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
index cb9b7f6ec2fd2..f07307fcd2f9d 100644
--- a/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
+++ b/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
@@ -502,6 +502,11 @@ static Type convertTensorType(const spirv::TargetEnv &targetEnv,
                << type << " illegal: cannot handle zero-element tensors\n");
     return nullptr;
   }
+  if (arrayElemCount > std::numeric_limits<unsigned>::max()) {
+    LLVM_DEBUG(llvm::dbgs()
+               << type << " illegal: cannot fit tensor into target type\n");
+    return nullptr;
+  }
 
   Type arrayElemType = convertScalarType(targetEnv, options, scalarType);
   if (!arrayElemType)

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Stefan Weigl-Bosker (sweiglbosker)

Changes

I noticed this because of #159738, though it was only caught by his fuzzer because it wrapped to 0.
Also, is there a reason for the usage of unsigned for sizes in spirv types? I believe most of the builtin types use int64_t for sizes, so it may make sense to do the same for spirv.

---

Full diff: https://github.com/llvm/llvm-project/pull/171002.diff

1 Files Affected:

• (modified) mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp (+5)

diff --git a/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp b/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
index cb9b7f6ec2fd2..f07307fcd2f9d 100644
--- a/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
+++ b/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
@@ -502,6 +502,11 @@ static Type convertTensorType(const spirv::TargetEnv &targetEnv,
                << type << " illegal: cannot handle zero-element tensors\n");
     return nullptr;
   }
+  if (arrayElemCount > std::numeric_limits<unsigned>::max()) {
+    LLVM_DEBUG(llvm::dbgs()
+               << type << " illegal: cannot fit tensor into target type\n");
+    return nullptr;
+  }
 
   Type arrayElemType = convertScalarType(targetEnv, options, scalarType);
   if (!arrayElemType)

[kuhar]

Thanks for fixing this. Would it be possible to add a lit test?

[sweiglbosker]

Thanks for fixing this. Would it be possible to add a lit test?

Done!

[kuhar]

Suggested change

	if (arrayElemCount > std::numeric_limits<unsigned>::max()) {
	if (arrayElemCount > std::numeric_limits<uint32_t>::max()) {

to make it platform-independent

[sweiglbosker]

I think this is actually the accurate way to do things, as ArrayType.elementCount is an unsigned.