[Mlir-commits] [mlir] [mlir][tosa] Add missing verifier for `tosa.pad` (PR #120934)

[Longsheng Mou]

https://github.com/CoTinker created https://github.com/llvm/llvm-project/pull/120934

This PR adds a missing verifier for `tosa.pad`, ensuring that the padding shape matches [rank(input), 2]. Fixes #119840.

>From bab92d0f5d203a97f9ad45fe09be0040ca6ba5c5 Mon Sep 17 00:00:00 2001
From: Longsheng Mou <longshengmou at gmail.com>
Date: Mon, 23 Dec 2024 14:46:50 +0800
Subject: [PATCH 1/2] [mlir][tosa] Add missing verifier for `tosa.pad`

This PR adds a missing verifier for `tosa.pad`, ensuring that the padding shape matches [rank(input), 2].
---
 mlir/lib/Dialect/Tosa/IR/TosaOps.cpp | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
index 631d3c48f2df02..6d9d2badc2970b 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
@@ -831,6 +831,11 @@ LogicalResult tosa::PadOp::verify() {
   if (paddingType.hasRank() && paddingType.getRank() != 2)
     return emitOpError() << "expect 'padding' tensor rank equal to 2.";
 
+  auto paddingShape = paddingType.getShape();
+  if (paddingShape.front() != inputType.getRank() || paddingShape.back() != 2)
+    return emitOpError()
+           << "expect 'padding' tensor shape to match [rank(input), 2]";
+
   return success();
 }
 

>From 4f51fbdd6a69fe6f585e323abc8826ba46368ade Mon Sep 17 00:00:00 2001
From: Longsheng Mou <longshengmou at gmail.com>
Date: Mon, 23 Dec 2024 14:49:05 +0800
Subject: [PATCH 2/2] Add tests

---
 mlir/test/Dialect/Tosa/invalid.mlir | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/mlir/test/Dialect/Tosa/invalid.mlir b/mlir/test/Dialect/Tosa/invalid.mlir
index cca50b25d14d6b..8a8fcbed6a67f3 100644
--- a/mlir/test/Dialect/Tosa/invalid.mlir
+++ b/mlir/test/Dialect/Tosa/invalid.mlir
@@ -119,6 +119,22 @@ func.func @test_pad_invalid_padConst_rank(%arg0: tensor<13x21xf32>, %arg1: tenso
 
 // -----
 
+func.func @test_pad_padding_shape_mismatch(%arg0: tensor<13x21x3xf32>, %arg1: tensor<2x2xi32>) -> tensor<13x21x3xf32> {
+  // expected-error at +1 {{'tosa.pad' op expect 'padding' tensor shape to match [rank(input), 2]}}
+  %0 = tosa.pad %arg0, %arg1 : (tensor<13x21x3xf32>, tensor<2x2xi32>) -> tensor<13x21x3xf32>
+  return %0 : tensor<13x21x3xf32>
+}
+
+// -----
+
+func.func @test_pad_padding_shape_mismatch(%arg0: tensor<13x21x3xf32>, %arg1: tensor<3x1xi32>) -> tensor<13x21x3xf32> {
+  // expected-error at +1 {{'tosa.pad' op expect 'padding' tensor shape to match [rank(input), 2]}}
+  %0 = tosa.pad %arg0, %arg1 : (tensor<13x21x3xf32>, tensor<3x1xi32>) -> tensor<13x21x3xf32>
+  return %0 : tensor<13x21x3xf32>
+}
+
+// -----
+
 func.func @test_transpose_non_const(%arg0: tensor<13x21x3xf32>, %arg1: tensor<3xi32>) -> tensor<3x13x21xf32> {
   // expected-error at +1 {{'tosa.transpose' op perms of transpose is not constant}}
   %0 = tosa.transpose %arg0, %arg1 : (tensor<13x21x3xf32>, tensor<3xi32>) -> tensor<3x13x21xf32>