[mlir][tosa]-Edit the verifier of tosa constShapeOp (llvm#126962)

Add verification for rank 1 for the elements' attribute of the tosa const_shape operation.
sivan-shani · Feb 21, 2025 · b537d3a · b537d3a
1 parent 0f3ec9c
commit b537d3a
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
@@ -2618,6 +2618,10 @@ OpTrait::tosa::verifyTosaShapeOperatorWithSameRanks(Operation *op) {
 //===----------------------------------------------------------------------===//
 
 LogicalResult tosa::ConstShapeOp::verify() {
+  // check one dimensional rank
+  auto valuesRank = getValue().getType().getRank();
+  if (valuesRank != 1)
+    return emitOpError("expect elements in attribute value with rank 1");
   // check that number of elements in value attr equal to rank of result shape
   auto count = getValue().getNumElements();
   auto rank = (cast<tosa::shapeType>(getResult().getType())).getRank();

diff --git a/mlir/test/Dialect/Tosa/invalid.mlir b/mlir/test/Dialect/Tosa/invalid.mlir
@@ -1087,6 +1087,14 @@ func.func @test_const_shape_value() -> !tosa.shape<5> {
 
 // -----
 
+func.func @test_const_shape_value() -> !tosa.shape<4> {
+  // expected-error@+1 {{'tosa.const_shape' op expect elements in attribute value with rank 1}}
+  %cst = tosa.const_shape {value = dense<[[1, 2], [3, 4]]> : tensor<2x2xindex>} : () -> !tosa.shape<4>
+  return %cst : !tosa.shape<4>
+}
+
+// -----
+
 func.func @test_sub_with_unequal_operand_ranks(%arg0: tensor<1x21x3xf32>, %arg1: tensor<1x13x21x3xf32>) -> tensor<1x13x21x3xf32> {
   // expected-error@+1 {{'tosa.sub' op operands don't have matching ranks}}
   %0 = tosa.sub %arg0, %arg1 : (tensor<1x21x3xf32>, tensor<1x13x21x3xf32>) -> tensor<1x13x21x3xf32>