[mlir] Add ValueBoundsOpInterfaceImpl for scf.forall (#118817)

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.

Signed-off-by: Max Dawkins <[email protected]>
Co-authored-by: Max Dawkins <[email protected]>

Author: Max191

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);
   });
 }

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