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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-linalg

Author: Alexey Kreshchuk (krsch-mobileye)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/230962.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.cpp (+6-5) 
- (modified) mlir/test/Dialect/Linalg/value-bounds-op-interface-impl.mlir (+48) 


``````````diff
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>
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/230962