[Mlir-commits] [mlir] [mlir][tosa]-Edit the verifier of tosa constShapeOp (PR #126962)

[Amir Bishara]

https://github.com/amirBish created https://github.com/llvm/llvm-project/pull/126962

Add verification for rank 1 for the elements' attribute of the tosa const_shape operation.

>From 0cce021dba823cd39e9f1e96b3071923a60a98a5 Mon Sep 17 00:00:00 2001
From: Amir Bishara <amir.bishara at mobileye.com>
Date: Wed, 12 Feb 2025 21:44:14 +0200
Subject: [PATCH] [mlir][tosa]-Edit verifier for tosa constShapeOp

Add verification for rank 1 for the elements' attribute
of the tosa const_shape operation.
---
 mlir/lib/Dialect/Tosa/IR/TosaOps.cpp | 4 ++++
 mlir/test/Dialect/Tosa/invalid.mlir  | 8 ++++++++
 2 files changed, 12 insertions(+)

diff --git a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
index 4928be38476a9..e782838ad97f9 100644
--- 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
index 913191be86f85..ff874af5c5f07 100644
--- 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 at +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 at +1 {{'tosa.sub' op operands don't have matching ranks}}
   %0 = tosa.sub %arg0, %arg1 : (tensor<1x21x3xf32>, tensor<1x13x21x3xf32>) -> tensor<1x13x21x3xf32>