[mlir] Add a folder for lbs, ubs, steps of scf.forall.

mlir/include/mlir/Dialect/Arith/Utils/Utils.h

 
 namespace mlir {
 
 /// Matches a ConstantIndexOp.
 detail::op_matcher<arith::ConstantIndexOp> matchConstantIndex();
 
 /// Detects the `values` produced by a ConstantIndexOp and places the new
 /// constant in place of the corresponding sentinel value.
+/// TODO(pifon2a): Remove this function and use foldDynamicIndexList.
 void canonicalizeSubViewPart(SmallVectorImpl<OpFoldResult> &values,
                              function_ref<bool(int64_t)> isDynamic);
 
+/// Returns `success` when any of the elements in `ofrs` was produced by
+/// arith::ConstantIndexOp. In that case the constant attribute replaces the
+/// Value. Returns `failure` when no folding happened.
+LogicalResult foldDynamicIndexList(Builder &b,
+                                   SmallVectorImpl<OpFoldResult> &ofrs);
+
 llvm::SmallBitVector getPositionsOfShapeOne(unsigned rank,
                                             ArrayRef<int64_t> shape);
 
 /// Pattern to rewrite a subview op with constant arguments.
 template <typename OpType, typename ResultTypeFunc, typename CastOpFunc>
 class OpWithOffsetSizesAndStridesConstantArgumentFolder final
     : public OpRewritePattern<OpType> {
 public:

mlir/lib/Dialect/Arith/Utils/Utils.cpp

 
 /// Matches a ConstantIndexOp.
 /// TODO: This should probably just be a general matcher that uses matchConstant
 /// and checks the operation for an index type.
 detail::op_matcher<arith::ConstantIndexOp> mlir::matchConstantIndex() {
   return detail::op_matcher<arith::ConstantIndexOp>();
 }
 
-/// Detects the `values` produced by a ConstantIndexOp and places the new
-/// constant in place of the corresponding sentinel value.
+// Detects the `values` produced by a ConstantIndexOp and places the new
+// constant in place of the corresponding sentinel value.
 void mlir::canonicalizeSubViewPart(
     SmallVectorImpl<OpFoldResult> &values,
     llvm::function_ref<bool(int64_t)> isDynamic) {
   for (OpFoldResult &ofr : values) {
     if (ofr.is<Attribute>())
       continue;
     // Newly static, move from Value to constant.
     if (auto cstOp =
             ofr.dyn_cast<Value>().getDefiningOp<arith::ConstantIndexOp>())
       ofr = OpBuilder(cstOp).getIndexAttr(cstOp.value());
   }
 }
 
+// Returns `success` when any of the elements in `ofrs` was produced by
+// arith::ConstantIndexOp. In that case the constant attribute replaces the
+// Value. Returns `failure` when no folding happened.
+LogicalResult mlir::foldDynamicIndexList(Builder &b,
+                                         SmallVectorImpl<OpFoldResult> &ofrs) {
+  bool valuesChanged = false;
+  for (OpFoldResult &ofr : ofrs) {
+    if (ofr.is<Attribute>())
+      continue;
+    // Newly static, move from Value to constant.
+    if (auto cstOp =
+            ofr.dyn_cast<Value>().getDefiningOp<arith::ConstantIndexOp>()) {
+      ofr = b.getIndexAttr(cstOp.value());
+      valuesChanged = true;
+    }
+  }
+  return success(valuesChanged);
+}
+
 llvm::SmallBitVector mlir::getPositionsOfShapeOne(unsigned rank,
                                                   ArrayRef<int64_t> shape) {
   llvm::SmallBitVector dimsToProject(shape.size());
   for (unsigned pos = 0, e = shape.size(); pos < e && rank > 0; ++pos) {
     if (shape[pos] == 1) {
       dimsToProject.set(pos);
       --rank;
     }

mlir/lib/Dialect/SCF/IR/SCF.cpp

if (!forallOp)
       return failure();
     Value sharedOut =
         forallOp.getTiedOpOperand(dimOp.getSource().cast<OpResult>())->get();
     rewriter.updateRootInPlace(
         dimOp, [&]() { dimOp.getSourceMutable().assign(sharedOut); });
     return success();
   }
 };
+
+class ForallOpControlOperandsFolder : public OpRewritePattern<ForallOp> {
+public:
+  using OpRewritePattern<ForallOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ForallOp op,
+                                PatternRewriter &rewriter) const override {
+    SmallVector<OpFoldResult> mixedLowerBound(op.getMixedLowerBound());
+    SmallVector<OpFoldResult> mixedUpperBound(op.getMixedUpperBound());
+    SmallVector<OpFoldResult> mixedStep(op.getMixedStep());
+    if (failed(foldDynamicIndexList(rewriter, mixedLowerBound)) &&
+        failed(foldDynamicIndexList(rewriter, mixedUpperBound)) &&
+        failed(foldDynamicIndexList(rewriter, mixedStep)))
+      return failure();
+
+    SmallVector<Value> dynamicLowerBound, dynamicUpperBound, dynamicStep;
+    SmallVector<int64_t> staticLowerBound, staticUpperBound, staticStep;
+    dispatchIndexOpFoldResults(mixedLowerBound, dynamicLowerBound,
+                               staticLowerBound);
+    op.getDynamicLowerBoundMutable().assign(dynamicLowerBound);
+    op.setStaticLowerBound(staticLowerBound);
+
+    dispatchIndexOpFoldResults(mixedUpperBound, dynamicUpperBound,
+                               staticUpperBound);
+    op.getDynamicUpperBoundMutable().assign(dynamicUpperBound);
+    op.setStaticUpperBound(staticUpperBound);
+
+    dispatchIndexOpFoldResults(mixedStep, dynamicStep, staticStep);
+    op.getDynamicStepMutable().assign(dynamicStep);
+    op.setStaticStep(staticStep);
+    return success();
+  }
+};
+
 } // namespace
 
 void ForallOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                            MLIRContext *context) {
-  results.add<DimOfForallOp>(context);
+  results.add<DimOfForallOp, ForallOpControlOperandsFolder>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // InParallelOp
 //===----------------------------------------------------------------------===//
 
 // Build a InParallelOp with mixed static and dynamic entries.
 void InParallelOp::build(OpBuilder &b, OperationState &result) {

mlir/test/Dialect/SCF/canonicalize.mlir

%2 = scf.forall (%tidx) in (%num_threads) shared_outs(%o = %arg1) -> (tensor<?x?xf32>) {
     }
   }
 
   // CHECK: %[[dim:.*]] = tensor.dim %[[arg1]], %[[c1]]
   %dim = tensor.dim %2, %c1 : tensor<?x?xf32>
   // CHECK: return %[[dim]]
   return %dim : index
 }
+
+// -----
+
+// CHECK-LABEL: func @forall_fold_control_operands
+func.func @forall_fold_control_operands(
+    %arg0 : tensor<?x10xf32>, %arg1: tensor<?x10xf32>) -> tensor<?x10xf32> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %dim0 = tensor.dim %arg0, %c0 : tensor<?x10xf32>
+  %dim1 = tensor.dim %arg0, %c1 : tensor<?x10xf32>
+
+  %result = scf.forall (%i, %j) = (%c0, %c0) to (%dim0, %dim1)
+      step (%c1, %c1) shared_outs(%o = %arg1) -> (tensor<?x10xf32>) {
+    %slice = tensor.extract_slice %arg1[%i, %j] [1, 1] [1, 1]
+      : tensor<?x10xf32> to tensor<1x1xf32>
+
+    scf.forall.in_parallel {
+      tensor.parallel_insert_slice %slice into %o[%i, %j] [1, 1] [1, 1]
+        : tensor<1x1xf32> into tensor<?x10xf32>
+    }
+  }
+
+  return %result : tensor<?x10xf32>
+}
+// CHECK: forall (%{{.*}}, %{{.*}}) in (%{{.*}}, 10)

mlir/test/Dialect/SCF/one-shot-bufferize.mlir

 func.func @scf_foreach_private_var(%t: tensor<10xf32>) -> f32 {
   %c2 = arith.constant 2 : index
   %c5 = arith.constant 5 : index
 
   // A copy is inserted for the uses of %t in the loop.
   // CHECK: %[[t_copy:.*]] = memref.alloc() {{.*}} : memref<10xf32>
   // CHECK: memref.copy %[[t]], %[[t_copy]]
 
-  // CHECK: scf.forall (%{{.*}}) in (%{{.*}}) {
+  // CHECK: scf.forall (%{{.*}}) in (2) {
 
   // Load from the copy and store into the shared output.
   // CHECK:   %[[subview:.*]] = memref.subview %[[t]]
   // CHECK:   memref.load %[[t_copy]]
   // CHECK:   memref.store %{{.*}}, %[[subview]]
   %0 = scf.forall (%tid) in (%c2) shared_outs(%o = %t) -> tensor<10xf32> {
     %offset = arith.muli %c5, %tid : index
     %slice = tensor.extract_slice %o[%offset] [5] [1]