[mlir] Add ValueBoundsOpInterfaceImpl for scf.forall

[Max191]

Adds a ValueBoundsOpInterface implementation for scf.forall ops. The implementation supports bounding for both induction variables, results, and block args of the forall op. Induction variables are given upper and lower bounds based on the lower and upper loop bounds, and dimensions of the results and init block arguments are constrained to be equal to the matching dims of the shared_outs operand.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-scf

Author: None (Max191)

Changes

Adds a ValueBoundsOpInterface implementation for scf.forall ops. The implementation supports bounding for both induction variables, results, and block args of the forall op. Induction variables are given upper and lower bounds based on the lower and upper loop bounds, and dimensions of the results and init block arguments are constrained to be equal to the matching dims of the shared_outs operand.

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp (+42)
• (modified) mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir (+36)

diff --git a/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
index 17a1c016ea16d5..fbd236b648cb8a 100644
--- a/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -95,6 +95,47 @@ struct ForOpInterface
   }
 };
 
+struct ForallOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<ForallOpInterface,
+                                                   ForallOp> {
+
+  void populateBoundsForIndexValue(Operation *op, Value value,
+                                   ValueBoundsConstraintSet &cstr) const {
+    auto forallOp = cast<ForallOp>(op);
+
+    // Index values should be induction variables, since the semantics of
+    // tensor::ParallelInsertSliceOp requires forall outputs to be ranked
+    // tensors.
+    auto blockArg = cast<BlockArgument>(value);
+    assert(blockArg.getArgNumber() < forallOp.getInductionVars().size() &&
+           "expected index value to be an induction var");
+    int64_t idx = blockArg.getArgNumber();
+    // TODO: Take into account step size.
+    AffineExpr lb = cstr.getExpr(forallOp.getMixedLowerBound()[idx]);
+    AffineExpr ub = cstr.getExpr(forallOp.getMixedUpperBound()[idx]);
+    cstr.bound(value) >= lb;
+    cstr.bound(value) < ub;
+  }
+
+  void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
+                                       ValueBoundsConstraintSet &cstr) const {
+    auto forallOp = cast<ForallOp>(op);
+
+    // `value` is an iter_arg or an OpResult.
+    int64_t iterArgIdx;
+    if (auto iterArg = llvm::dyn_cast<BlockArgument>(value)) {
+      iterArgIdx = iterArg.getArgNumber() - forallOp.getInductionVars().size();
+    } else {
+      iterArgIdx = llvm::cast<OpResult>(value).getResultNumber();
+    }
+
+    // The forall results and output arguments have the same sizes as the output
+    // operands.
+    Value outputOperand = forallOp.getOutputs()[iterArgIdx];
+    cstr.bound(value)[dim] == cstr.getExpr(outputOperand, dim);
+  }
+};
+
 struct IfOpInterface
     : public ValueBoundsOpInterface::ExternalModel<IfOpInterface, IfOp> {
 
@@ -161,6 +202,7 @@ void mlir::scf::registerValueBoundsOpInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, scf::SCFDialect *dialect) {
     scf::ForOp::attachInterface<scf::ForOpInterface>(*ctx);
+    scf::ForallOp::attachInterface<scf::ForallOpInterface>(*ctx);
     scf::IfOp::attachInterface<scf::IfOpInterface>(*ctx);
   });
 }
diff --git a/mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir b/mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir
index 9ab03da1c9a94f..65e1017e62c1a4 100644
--- a/mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir
+++ b/mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir
@@ -107,6 +107,42 @@ func.func @scf_for_swapping_yield(%t1: tensor<?xf32>, %t2: tensor<?xf32>, %a: in
 
 // -----
 
+// CHECK-LABEL: func @scf_forall(
+//  CHECK-SAME:     %[[a:.*]]: index, %[[b:.*]]: index, %[[c:.*]]: index
+//       CHECK:   "test.some_use"(%[[a]], %[[b]])
+func.func @scf_forall(%a: index, %b: index, %c: index) {
+  scf.forall (%iv) = (%a) to (%b) step (%c) {
+    %0 = "test.reify_bound"(%iv) {type = "LB"} : (index) -> (index)
+    %1 = "test.reify_bound"(%iv) {type = "UB"} : (index) -> (index)
+    "test.some_use"(%0, %1) : (index, index) -> ()
+  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @scf_forall_tensor_result(
+//  CHECK-SAME:     %[[size:.*]]: index, %[[a:.*]]: index, %[[b:.*]]: index, %[[c:.*]]: index
+//       CHECK:   "test.some_use"(%[[size]])
+//       CHECK:   "test.some_use"(%[[size]])
+func.func @scf_forall_tensor_result(%size: index, %a: index, %b: index, %c: index) {
+  %cst = arith.constant 5.0 : f32
+  %empty = tensor.empty(%size) : tensor<?xf32>
+  %0 = scf.forall (%iv) = (%a) to (%b) step (%c) shared_outs(%arg = %empty) -> tensor<?xf32> {
+    %filled = linalg.fill ins(%cst : f32) outs(%arg : tensor<?xf32>) -> tensor<?xf32>
+    %1 = "test.reify_bound"(%arg) {type = "EQ", dim = 0} : (tensor<?xf32>) -> (index)
+    "test.some_use"(%1) : (index) -> ()
+    scf.forall.in_parallel {
+      tensor.parallel_insert_slice %filled into %arg[0][%size][1] : tensor<?xf32> into tensor<?xf32>
+    }
+  }
+  %2 = "test.reify_bound"(%0) {type = "EQ", dim = 0} : (tensor<?xf32>) -> (index)
+  "test.some_use"(%2) : (index) -> ()
+  return
+}
+
+// -----
+
 // CHECK-LABEL: func @scf_if_constant(
 func.func @scf_if_constant(%c : i1) {
   // CHECK: arith.constant 4 : index