diff --git a/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp b/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp --- a/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp +++ b/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp @@ -390,8 +390,14 @@ 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 arrayElemSize = diff --git a/mlir/test/Conversion/TensorToSPIRV/tensor-ops-to-spirv.mlir b/mlir/test/Conversion/TensorToSPIRV/tensor-ops-to-spirv.mlir --- a/mlir/test/Conversion/TensorToSPIRV/tensor-ops-to-spirv.mlir +++ b/mlir/test/Conversion/TensorToSPIRV/tensor-ops-to-spirv.mlir @@ -1,4 +1,5 @@ -// 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 @@ -27,3 +28,38 @@ // 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 + 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 +}