[Mlir-commits] [mlir] [mlir][tosa] Fix indexing in TosaToTensor (PR #140906)

[Shay Kleiman]

https://github.com/shay-kl updated https://github.com/llvm/llvm-project/pull/140906

>From 216a5f2ab90e1dc9c3fc97ce7694fb4fc841056c Mon Sep 17 00:00:00 2001
From: Shay Kleiman <shay.kleiman at mobileye.com>
Date: Wed, 21 May 2025 17:17:14 +0300
Subject: [PATCH 1/2] [mlir][tosa] Fix indexing in TosaToTensor

Changed the indexing used in the extractOp from one that is intended for
0d tensors to one that is intended for 1d tensors.
---
 mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp b/mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp
index 5f23a33049f87..615c7ca1cfd15 100644
--- a/mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp
+++ b/mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp
@@ -362,7 +362,8 @@ class PadConverter : public OpConversionPattern<tosa::PadOp> {
     // Setup the default constantAttr.
 
     Value padConstant = rewriter.createOrFold<tensor::ExtractOp>(
-        loc, padOp.getPadConst(), ValueRange({}));
+        loc, padOp.getPadConst(),
+        ValueRange({rewriter.create<arith::ConstantIndexOp>(loc, 0)}));
 
     if (!padConstant) {
       return rewriter.notifyMatchFailure(

>From 41151d6f46be4f0824ac15f6f2b2a1a8baab2910 Mon Sep 17 00:00:00 2001
From: Shay Kleiman <shay.kleiman at mobileye.com>
Date: Thu, 22 May 2025 10:16:47 +0300
Subject: [PATCH 2/2] Added test which doesn't pass prior to change

---
 .../TosaToTensor/tosa-to-tensor.mlir          | 22 +++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/mlir/test/Conversion/TosaToTensor/tosa-to-tensor.mlir b/mlir/test/Conversion/TosaToTensor/tosa-to-tensor.mlir
index b19cd2115891c..0a276e2a5c3d2 100644
--- a/mlir/test/Conversion/TosaToTensor/tosa-to-tensor.mlir
+++ b/mlir/test/Conversion/TosaToTensor/tosa-to-tensor.mlir
@@ -700,3 +700,25 @@ func.func @concat_non_axis_dyn_mixed(%arg0: tensor<?x1xf32>, %arg1: tensor<?x1xf
   %0 = "tosa.concat"(%arg0, %arg1, %arg2) <{axis = 1 : i32}> : (tensor<?x1xf32>, tensor<?x1xf32>, tensor<?x1xf32>) -> tensor<5x3xf32>
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @pad_variable_pad_const
+// CHECK-SAME: (%[[ARG0_SSA:.*]]: tensor<2x2xi32>, %[[PAD_INPUT_TENSOR_SSA:.*]]: tensor<1xi32>)
+func.func @pad_variable_pad_const(%arg0: tensor<2x2xi32>, %pad_input_tensor: tensor<1xi32>) -> tensor<4x5xi32> {
+  // CHECK-DAG: %[[C0_INDEX:.*]] = arith.constant 0 : index
+  // CHECK-DAG: %[[EXTRACTED_PAD_VAL:.*]] = tensor.extract %[[PAD_INPUT_TENSOR_SSA]][%[[C0_INDEX]]] : tensor<1xi32>
+
+  // CHECK: %[[C_PAD_LOW_0:.*]] = arith.constant 1 : index
+  // CHECK: %[[C_PAD_HIGH_0:.*]] = arith.constant 1 : index
+  // CHECK: %[[C_PAD_LOW_1:.*]] = arith.constant 0 : index
+  // CHECK: %[[C_PAD_HIGH_1:.*]] = arith.constant 3 : index
+
+  // CHECK: %{{.*}} = tensor.pad %[[ARG0_SSA]] low[%[[C_PAD_LOW_0]], %[[C_PAD_LOW_1]]] high[%[[C_PAD_HIGH_0]], %[[C_PAD_HIGH_1]]] {
+  // CHECK:   tensor.yield %[[EXTRACTED_PAD_VAL]] : i32
+  // CHECK: } : tensor<2x2xi32> to tensor<4x5xi32>
+
+  %padding_indices = tosa.const_shape {values = dense<[1, 1, 0, 3]> : tensor<4xindex>} : () -> !tosa.shape<4>
+  %result = "tosa.pad"(%arg0, %padding_indices, %pad_input_tensor) : (tensor<2x2xi32>, !tosa.shape<4>, tensor<1xi32>) -> tensor<4x5xi32>
+  return %result : tensor<4x5xi32>
+}