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

Author: sweiglbosker

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)

mlir/test/Conversion/TensorToSPIRV/tensor-ops-to-spirv.mlir

@@ -79,3 +79,10 @@ func.func @tensor_2d_empty() -> () {
   %x = arith.constant dense<> : tensor<2x0xi32>
   return
 }
+
+// CHECK-LABEL: func @very_large_tensor
+// CHECK-NEXT:    arith.constant dense<1>
+func.func @very_large_tensor() -> () {
+  %x = arith.constant dense<1> : tensor<4294967296xi32>
+  return
+}