[mlir][tosa] Verify zero-dim tensors in the input

[Lewuathe]

Summary

As TOSA does not support the tensor with zero dimensions, we can check the zero value for the static shape input.

Ideally, we should be able to check the tensor shape more broadly, such as using CPred in the TOSA type definition. But based on the discussion here It makes input type verification complicated and hard to maintain and still only applies to the case the input is statically shaped. Therefore, in this change, we have put the zero dimension check in the verification of each op, which would be flexible and maintainable.

See: https://github.com/llvm/llvm-project/issues/63212

[mgehre-amd]

Great, thanks for addressing my bug report!

[TinaAMD]

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

103

Maybe we can name it "hasZeroDimension" (without the "s")? When I first saw the name, I was a bit confused if this checks for rank == 0 (which is permitted).

[jpienaar]

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

104

Is static on purpose or do you want to see if ranked? E.g.,. ?x0x10 would pass verification as not static, but you already know it's invalid

[Lewuathe]

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

104

Good point. I'm going to omit the check for the static type and check the zero as long as the input is ranked.

[Lewuathe]

Support all ranked tensors.

[eric-k256]

I agree with the concept here, just a comment on having a clearer error message and then I think it will be ready.

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

132

Here and below, I think this error is confusing in the same way the function name was confusing. I suggest

"tensor has a dimension with size zero. Each dimension of a tensor must have size >= 1"

as an example.

[Lewuathe]

Improve error message.

[Lewuathe]

@eric-k256 I have improved the error message accordingly. Can I ask you to review it again?

[eric-k256]

Looks good, thanks!

[Peiming]

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

104

Could the ShapedType be unranked? We observed an (*this).hasRank() && "cannot query rank of unranked shaped type" error from upstream.

[Peiming]

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

104

https://reviews.llvm.org/D154912

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

return builder.create<tosa::ConstOp>(loc, type,
                                          llvm::cast<ElementsAttr>(value));
   return nullptr;
 }
 
 //===----------------------------------------------------------------------===//
 // TOSA Operator Verifiers.
 //===----------------------------------------------------------------------===//
 
+static bool hasZeroDimension(ShapedType shapedType) {
+  auto rank = shapedType.getRank();
+
+  for (int i = 0; i < rank; i++) {
+    if (shapedType.isDynamicDim(i))
+      continue;
+    if (shapedType.getDimSize(i) == 0)
+      return true;
+  }
+
+  return false;
+}
+
 template <typename T> static LogicalResult verifyConvOp(T op) {
   // All TOSA conv ops have an input() and weight().
   auto inputType = llvm::dyn_cast<RankedTensorType>(op.getInput().getType());
   auto weightType = llvm::dyn_cast<RankedTensorType>(op.getWeight().getType());
 
   // Must be ranked tensor types
   if (!inputType) {
     op.emitOpError("expect a ranked tensor for input, got ") << op.getInput();
     return failure();
   }
   if (!weightType) {
     op.emitOpError("expect a ranked tensor for weight, got ") << op.getWeight();
     return failure();
   }
 
+  if (hasZeroDimension(inputType))
+    return op.emitOpError() << "tensor has a dimension with size zero. Each "
+                               "dimension of a tensor must have size >= 1";
+
   auto inputEType = inputType.getElementType();
   auto weightEType = weightType.getElementType();
 
   bool inputIsQuant = !llvm::isa<FloatType>(inputEType);
   bool weightIsQuant = !llvm::isa<FloatType>(weightEType);
 
   // Either both must be quantized or both unquantized.
   if (inputIsQuant != weightIsQuant) {
op.emitOpError("quantizationattr is required for quantized type, and not "
                    "allowed for float type");
     return failure();
   }
 
   return success();
 }
 
 LogicalResult tosa::AvgPool2dOp::verify() {
-  auto inputETy = llvm::cast<ShapedType>(getInput().getType()).getElementType();
+  auto inputType = llvm::cast<ShapedType>(getInput().getType());
+  if (hasZeroDimension(inputType))
+    return emitOpError() << "tensor has a dimension with size zero. Each "
+                            "dimension of a tensor must have size >= 1";
+
+  auto inputETy = inputType.getElementType();
   auto resultETy = llvm::cast<ShapedType>(getType()).getElementType();
 
   if (auto quantType =
           llvm::dyn_cast<mlir::quant::UniformQuantizedType>(inputETy))
     inputETy = quantType.getStorageType();
 
   if (auto quantType =
           llvm::dyn_cast<mlir::quant::UniformQuantizedType>(resultETy))
inferredReturnShapes.push_back(
       ShapedTypeComponents(newShapeValue, inputType));
   return success();
 }
 
 mlir::LogicalResult tosa::ReshapeOp::verify() {
   ShapedType inputType = llvm::cast<ShapedType>(getInput1().getType());
   ShapedType outputType = llvm::cast<ShapedType>(getType());
 
+  if (hasZeroDimension(inputType) || hasZeroDimension(outputType))
+    return emitOpError() << "tensor has a dimension with size zero. Each "
+                            "dimension of a tensor must have size >= 1";
+
   if (inputType.hasStaticShape() && outputType.hasStaticShape()) {
     int64_t inputElementsNum = inputType.getNumElements();
     int64_t outputElementsNum = outputType.getNumElements();
     if (inputElementsNum != outputElementsNum) {
       return emitOpError() << "Cannot reshape " << inputElementsNum
                            << " elements into " << outputElementsNum;
     }
   }

mlir/test/Dialect/Tosa/invalid.mlir

 
 // -----
 
 func.func @test_const_attribute_type_mismatch() -> tensor<100x100xf32> {
   // expected-error@+1 {{'tosa.const' op failed to verify that all of {value, output} have same shape}}
   %0 = "tosa.const"() {value = dense<0.000000e+00> : tensor<1x1xf32>} : () -> tensor<100x100xf32>
   return %0 : tensor<100x100xf32>
 }
+
+// -----
+
+func.func @test_reshape_static_zero_dim_input(%arg0 : tensor<13x0x3xf32>) -> () {
+  // expected-error@+1 {{'tosa.reshape' op tensor has a dimension with size zero. Each dimension of a tensor must have size >= 1}}
+  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 13, 21, 3>} : (tensor<13x0x3xf32>) -> tensor<13x0x3xf32>
+  return
+}
+
+// -----
+
+func.func @test_reshape_zero_dim_input(%arg0 : tensor<?x0x3xf32>) -> () {
+  // expected-error@+1 {{'tosa.reshape' op tensor has a dimension with size zero. Each dimension of a tensor must have size >= 1}}
+  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 13, 21, 3>} : (tensor<?x0x3xf32>) -> tensor<13x0x3xf32>
+  return
+}
+
+// -----
+
+func.func @test_conv2d_static_zero_dim_input(%arg0: tensor<1x29x0x4xf32>, %arg1: tensor<16x3x3x4xf32>, %arg2: tensor<16xf32>) -> tensor<1x27x27x16xf32> {
+  // expected-error@+1 {{'tosa.conv2d' op tensor has a dimension with size zero. Each dimension of a tensor must have size >= 1}}
+  %0 = "tosa.conv2d"(%arg0, %arg1, %arg2) {dilation = array<i64: 1, 1>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>}
+           : (tensor<1x29x0x4xf32>, tensor<16x3x3x4xf32>, tensor<16xf32>) -> tensor<1x27x27x16xf32>
+  return %0 : tensor<1x27x27x16xf32>
+}
+
+// -----
+
+func.func @test_conv2d_zero_dim_input(%arg0: tensor<1x?x0x4xf32>, %arg1: tensor<16x3x3x4xf32>, %arg2: tensor<16xf32>) -> tensor<1x27x27x16xf32> {
+  // expected-error@+1 {{'tosa.conv2d' op tensor has a dimension with size zero. Each dimension of a tensor must have size >= 1}}
+  %0 = "tosa.conv2d"(%arg0, %arg1, %arg2) {dilation = array<i64: 1, 1>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>}
+           : (tensor<1x?x0x4xf32>, tensor<16x3x3x4xf32>, tensor<16xf32>) -> tensor<1x27x27x16xf32>
+  return %0 : tensor<1x27x27x16xf32>
+}
+
+
+// -----
+
+func.func @test_avg_pool2d_static_zero_dim_input(%arg0: tensor<1x0x7x9xf32>) -> tensor<1x7x7x9xf32> {
+  // expected-error@+1 {{'tosa.avg_pool2d' op tensor has a dimension with size zero. Each dimension of a tensor must have size >= 1}}
+    %0 = "tosa.avg_pool2d"(%arg0) {acc_type = f32, kernel = array<i64: 2, 2>, pad = array<i64: 0, 1, 0, 1>, stride = array<i64: 1, 1>} 
+      : (tensor<1x0x7x9xf32>) -> tensor<1x7x7x9xf32>
+    return %0 : tensor<1x7x7x9xf32>
+}
+
+// -----
+
+func.func @test_avg_pool2d_zero_dim_input(%arg0: tensor<1x0x?x9xf32>) -> tensor<1x7x7x9xf32> {
+  // expected-error@+1 {{'tosa.avg_pool2d' op tensor has a dimension with size zero. Each dimension of a tensor must have size >= 1}}  
+    %0 = "tosa.avg_pool2d"(%arg0) {acc_type = f32, kernel = array<i64: 2, 2>, pad = array<i64: 0, 1, 0, 1>, stride = array<i64: 1, 1>} 
+      : (tensor<1x0x?x9xf32>) -> tensor<1x7x7x9xf32>
+    return %0 : tensor<1x7x7x9xf32>
+}