[Mlir-commits] [mlir] [MLIR] Support scalars in ValueBounds of linalg.index (PR #230962)

[Alexey Kreshchuk]

https://github.com/krsch-mobileye updated https://github.com/llvm/llvm-project/pull/230962

>From 35fd1f20369ac9aa4045077381addbffd0cacb90 Mon Sep 17 00:00:00 2001
From: Alexey Kreshchuk <alexey.kreshchuk at mobileye.com>
Date: Sun, 11 Oct 2026 16:47:44 +0300
Subject: [PATCH 1/2] [MLIR] Support scalars in ValueBounds of linalg.index

ValueBoundsOpInterfaceImpl for linalg.index expects every operand to be
ShapedType, but linalg.generic also supports scalars as operands. I've
made it query the rank with linalgOp.getRank
---
 .../Linalg/IR/ValueBoundsOpInterfaceImpl.cpp  | 11 +++--
 .../value-bounds-op-interface-impl.mlir       | 48 +++++++++++++++++++
 2 files changed, 54 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.cpp
index f56ef485069f8..12cad71d403b1 100644
--- a/mlir/lib/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -35,14 +35,15 @@ struct IndexOpInterface
             .getPosition();
     // Find the `flatDimPos`-th operand dimension.
     int64_t flatDimCtr = 0;
-    for (Value operand : linalgOp->getOperands()) {
+    for (OpOperand &operand : linalgOp->getOpOperands()) {
       assert(flatDimPos >= flatDimCtr && "invalid pos");
-      auto shapedType = llvm::cast<ShapedType>(operand.getType());
-      if (flatDimPos < flatDimCtr + shapedType.getRank()) {
-        cstr.bound(value) < cstr.getExpr(operand, flatDimPos - flatDimCtr);
+      // Scalar operands do not contribute dimensions to the flattened shape.
+      int64_t rank = linalgOp.getRank(&operand);
+      if (flatDimPos < flatDimCtr + rank) {
+        cstr.bound(value) < cstr.getExpr(operand.get(), flatDimPos - flatDimCtr);
         break;
       }
-      flatDimCtr += shapedType.getRank();
+      flatDimCtr += rank;
     }
   }
 };
diff --git a/mlir/test/Dialect/Linalg/value-bounds-op-interface-impl.mlir b/mlir/test/Dialect/Linalg/value-bounds-op-interface-impl.mlir
index bcd330443cc44..e68d1332bcc03 100644
--- a/mlir/test/Dialect/Linalg/value-bounds-op-interface-impl.mlir
+++ b/mlir/test/Dialect/Linalg/value-bounds-op-interface-impl.mlir
@@ -62,3 +62,51 @@ func.func @linalg_index(%arg0: memref<?x?xf32>,
   }
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @linalg_index_scalar_input(
+func.func @linalg_index_scalar_input(%scalar: i32, %out: tensor<4xi32>) -> tensor<4xi32> {
+  %result = linalg.generic {
+      indexing_maps = [affine_map<(d0) -> ()>, affine_map<(d0) -> (d0)>],
+      iterator_types = ["parallel"]}
+      ins(%scalar : i32) outs(%out : tensor<4xi32>) {
+  ^bb0(%in: i32, %old: i32):
+    %i = linalg.index 0 : index
+    // CHECK: %[[UB:.*]] = arith.constant 4 : index
+    // CHECK: %[[VALUE:.*]] = arith.index_cast %[[UB]] : index to i32
+    // CHECK: linalg.yield %[[VALUE]] : i32
+    %ub = "test.reify_bound"(%i) {type = "UB"} : (index) -> index
+    %value = arith.index_cast %ub : index to i32
+    linalg.yield %value : i32
+  } -> tensor<4xi32>
+  return %result : tensor<4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func @linalg_index_scalar_input_dynamic(
+// CHECK-SAME: %[[OUT:[^ ,:]+]]: tensor<?x?xi32>
+func.func @linalg_index_scalar_input_dynamic(%input: tensor<?xi32>, %scalar: i32,
+                                           %rank_zero: tensor<i32>,
+                                           %out: tensor<?x?xi32>) -> tensor<?x?xi32> {
+  %result = linalg.generic {
+      indexing_maps = [affine_map<(d0, d1) -> (d0)>,
+                       affine_map<(d0, d1) -> ()>,
+                       affine_map<(d0, d1) -> ()>,
+                       affine_map<(d0, d1) -> (d0, d1)>],
+      iterator_types = ["parallel", "parallel"]}
+      ins(%input, %scalar, %rank_zero : tensor<?xi32>, i32, tensor<i32>)
+      outs(%out : tensor<?x?xi32>) {
+  ^bb0(%in: i32, %s: i32, %r: i32, %old: i32):
+    %i = linalg.index 1 : index
+    // CHECK: %[[C1:.*]] = arith.constant 1 : index
+    // CHECK: %[[UB:.*]] = tensor.dim %[[OUT]], %[[C1]]
+    // CHECK: %[[VALUE:.*]] = arith.index_cast %[[UB]] : index to i32
+    // CHECK: linalg.yield %[[VALUE]] : i32
+    %ub = "test.reify_bound"(%i) {type = "UB"} : (index) -> index
+    %value = arith.index_cast %ub : index to i32
+    linalg.yield %value : i32
+  } -> tensor<?x?xi32>
+  return %result : tensor<?x?xi32>
+}

>From d02b463053ff9aa1baa82287c13924e658dc8444 Mon Sep 17 00:00:00 2001
From: Alexey Kreshchuk <alexey.kreshchuk at mobileye.com>
Date: Sun, 11 Oct 2026 17:01:49 +0300
Subject: [PATCH 2/2] format

---
 mlir/lib/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.cpp
index 12cad71d403b1..dbcc3fa63f03a 100644
--- a/mlir/lib/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -40,7 +40,8 @@ struct IndexOpInterface
       // Scalar operands do not contribute dimensions to the flattened shape.
       int64_t rank = linalgOp.getRank(&operand);
       if (flatDimPos < flatDimCtr + rank) {
-        cstr.bound(value) < cstr.getExpr(operand.get(), flatDimPos - flatDimCtr);
+        cstr.bound(value) <
+            cstr.getExpr(operand.get(), flatDimPos - flatDimCtr);
         break;
       }
       flatDimCtr += rank;