[mlir][spirv] Handle zero-element tensors in spirv type conversion

[kuhar]

Summary

Return gracefully instead of crashing. Add missing type conversion
tests.

Fixes: https://github.com/llvm/llvm-project/issues/61044

mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp

static Type convertTensorType(const spirv::TargetEnv &targetEnv,
   std::optional<int64_t> scalarSize = getTypeNumBytes(options, scalarType);
   std::optional<int64_t> tensorSize = getTypeNumBytes(options, type);
   if (!scalarSize || !tensorSize) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: cannot deduce element count\n");
     return nullptr;
   }
 
-  auto arrayElemCount = *tensorSize / *scalarSize;
-  auto arrayElemType = convertScalarType(targetEnv, options, scalarType);
+  int64_t arrayElemCount = *tensorSize / *scalarSize;
+  if (arrayElemCount == 0) {
+    LLVM_DEBUG(llvm::dbgs()
+               << type << " illegal: cannot handle zero-element tensors\n");
+    return nullptr;
+  }
+
+  Type arrayElemType = convertScalarType(targetEnv, options, scalarType);
   if (!arrayElemType)
     return nullptr;
   std::optional<int64_t> arrayElemSize =
       getTypeNumBytes(options, arrayElemType);
   if (!arrayElemSize) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: cannot deduce converted element size\n");
     return nullptr;

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

-// RUN: mlir-opt -split-input-file -convert-tensor-to-spirv -verify-diagnostics %s | FileCheck %s
+// RUN: mlir-opt --split-input-file --convert-tensor-to-spirv \
+// RUN:   --verify-diagnostics %s | FileCheck %s
 
 //===----------------------------------------------------------------------===//
 // tensor.extract
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: func @tensor_extract_constant
 // CHECK-SAME: (%[[A:.+]]: i32, %[[B:.+]]: i32, %[[C:.+]]: i32)
 func.func @tensor_extract_constant(%a : index, %b: index, %c: index) -> i32 {
func.func @tensor_extract_constant(%a : index, %b: index, %c: index) -> i32 {
   // CHECK: %[[MUL2:.+]] = spirv.IMul %[[C1]], %[[C]] : i32
   // CHECK: %[[ADD2:.+]] = spirv.IAdd %[[ADD1]], %[[MUL2]] : i32
   // CHECK: %[[AC:.+]] = spirv.AccessChain %[[VAR]][%[[ADD2]]]
   // CHECK: %[[VAL:.+]] = spirv.Load "Function" %[[AC]] : i32
   %extract = tensor.extract %cst[%a, %b, %c] : tensor<2x2x3xi32>
   // CHECK: spirv.ReturnValue %[[VAL]]
   return %extract : i32
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// Type conversion
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: func @tensor_0d
+// CHECK-NEXT:    spirv.Constant 1 : i32
+func.func @tensor_0d() -> () {
+  %x = arith.constant dense<1> : tensor<i32>
+  return
+}
+
+// CHECK-LABEL: func @tensor_1d
+// CHECK-NEXT:    spirv.Constant dense<[1, 2, 3]> : tensor<3xi32> : !spirv.array<3 x i32>
+func.func @tensor_1d() -> () {
+  %x = arith.constant dense<[1, 2, 3]> : tensor<3xi32>
+  return
+}
+
+// CHECK-LABEL: func @tensor_2d
+// CHECK-NEXT:    spirv.Constant dense<[1, 2, 3, 4, 5, 6]> : tensor<6xi32> : !spirv.array<6 x i32>
+func.func @tensor_2d() -> () {
+  %x = arith.constant dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32>
+  return
+}
+
+// We do not handle zero-element tensors yet. Just make we do not crash on them.
+// CHECK-LABEL: func @tensor_2d_empty
+// CHECK-NEXT:    arith.constant dense<>
+func.func @tensor_2d_empty() -> () {
+  %x = arith.constant dense<> : tensor<2x0xi32>
+  return
+}