[mlir][SCF] Add scf::ForeachThread canonicalization.

[nicolasvasilache]

Summary

This revision adds the necessary plumbing for canonicalizing scf::ForeachThread with the
AffineOpSCFCanonicalizationPattern.
In the process the loopMatcher helper is updated to take OpFoldResult instead of just values.
This allows composing various scenarios without the need for an artificial builder.

[nicolasvasilache]

Rebase

mlir/include/mlir/Dialect/SCF/IR/SCF.h

 /// Returns the loop parent of an induction variable. If the provided value is
 /// not an induction variable, then return nullptr.
 ForOp getForInductionVarOwner(Value val);
 
 /// Returns the parallel loop parent of an induction variable. If the provided
 /// value is not an induction variable, then return nullptr.
 ParallelOp getParallelForInductionVarOwner(Value val);
 
+/// Returns the ForeachThreadOp parent of an thread index variable.
+/// If the provided value is not a thread index variable, then return nullptr.
+ForeachThreadOp getForeachThreadOpThreadIndexOwner(Value val);
+
 /// Return true if ops a and b (or their ancestors) are in mutually exclusive
 /// regions/blocks of an IfOp.
 // TODO: Consider moving this functionality to RegionBranchOpInterface.
 bool insideMutuallyExclusiveBranches(Operation *a, Operation *b);
 
 /// An owning vector of values, handy to return from functions.
 using ValueVector = std::vector<Value>;
 using LoopVector = std::vector<scf::ForOp>;

mlir/include/mlir/Dialect/SCF/Utils/AffineCanonicalizationUtils.h

 #define MLIR_DIALECT_SCF_UTILS_AFFINECANONICALIZATIONUTILS_H_
 
 #include "mlir/Support/LLVM.h"
 
 namespace mlir {
 class AffineMap;
 struct LogicalResult;
 class Operation;
+class OpFoldResult;
 class RewriterBase;
 class Value;
 class ValueRange;
 
 namespace scf {
 class IfOp;
 
 /// Match "for loop"-like operations: If the first parameter is an iteration
 /// variable, return lower/upper bounds via the second/third parameter and the
 /// step size via the last parameter. The function should return `success` in
 /// that case. If the first parameter is not an iteration variable, return
 /// `failure`.
-using LoopMatcherFn =
-    function_ref<LogicalResult(Value, Value &, Value &, Value &)>;
+using LoopMatcherFn = function_ref<LogicalResult(
+    Value, OpFoldResult &, OpFoldResult &, OpFoldResult &)>;
 
 /// Try to canonicalize an min/max operations in the context of for `loops` with
 /// a known range.
 ///
 /// `map` is the body of the min/max operation and `operands` are the SSA values
 /// that the dimensions and symbols are bound to; dimensions are listed first.
 /// If `isMin`, the operation is a min operation; otherwise, a max operation.
 /// `loopMatcher` is used to retrieve loop bounds and the step size for a given

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

void ForeachThreadOp::ensureTerminator(Region &region, OpBuilder &builder,
   if (terminator.getRegion().empty())
     builder.createBlock(&terminator.getRegion());
 }
 
 PerformConcurrentlyOp ForeachThreadOp::getTerminator() {
   return cast<PerformConcurrentlyOp>(getBody()->getTerminator());
 }
 
+ForeachThreadOp mlir::scf::getForeachThreadOpThreadIndexOwner(Value val) {
+  auto tidxArg = val.dyn_cast<BlockArgument>();
+  if (!tidxArg)
+    return ForeachThreadOp();
+  assert(tidxArg.getOwner() && "unlinked block argument");
+  auto *containingOp = tidxArg.getOwner()->getParentOp();
+  return dyn_cast<ForeachThreadOp>(containingOp);
+}
+
 //===----------------------------------------------------------------------===//
 // ParallelInsertSliceOp
 //===----------------------------------------------------------------------===//
 
 // Build a ParallelInsertSliceOp with mixed static and dynamic entries.
 void ParallelInsertSliceOp::build(OpBuilder &b, OperationState &result,
                                   Value source, Value dest,
                                   ArrayRef<OpFoldResult> offsets,

mlir/lib/Dialect/SCF/Transforms/LoopCanonicalization.cpp

LogicalResult matchAndRewrite(OpTy dimOp,
                                 PatternRewriter &rewriter) const override {
     auto forOp = dimOp.source().template getDefiningOp<scf::ForOp>();
     if (!forOp)
       return failure();
     auto opResult = dimOp.source().template cast<OpResult>();
     unsigned resultNumber = opResult.getResultNumber();
     if (!isShapePreserving(forOp, resultNumber))
       return failure();
-    rewriter.updateRootInPlace(dimOp, [&](){
+    rewriter.updateRootInPlace(dimOp, [&]() {
       dimOp.sourceMutable().assign(forOp.getIterOperands()[resultNumber]);
     });
     return success();
   }
 };
 
 /// Canonicalize AffineMinOp/AffineMaxOp operations in the context of scf.for
 /// and scf.parallel loops with a known range.
 template <typename OpTy, bool IsMin>
 struct AffineOpSCFCanonicalizationPattern : public OpRewritePattern<OpTy> {
   using OpRewritePattern<OpTy>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(OpTy op,
                                 PatternRewriter &rewriter) const override {
-    auto loopMatcher = [](Value iv, Value &lb, Value &ub, Value &step) {
+    auto loopMatcher = [](Value iv, OpFoldResult &lb, OpFoldResult &ub,
+                          OpFoldResult &step) {
       if (scf::ForOp forOp = scf::getForInductionVarOwner(iv)) {
         lb = forOp.getLowerBound();
         ub = forOp.getUpperBound();
         step = forOp.getStep();
         return success();
       }
       if (scf::ParallelOp parOp = scf::getParallelForInductionVarOwner(iv)) {
         for (unsigned idx = 0; idx < parOp.getNumLoops(); ++idx) {
           if (parOp.getInductionVars()[idx] == iv) {
             lb = parOp.getLowerBound()[idx];
             ub = parOp.getUpperBound()[idx];
             step = parOp.getStep()[idx];
             return success();
           }
         }
         return failure();
       }
+      if (scf::ForeachThreadOp foreachThreadOp =
+              scf::getForeachThreadOpThreadIndexOwner(iv)) {
+        for (int64_t idx = 0; idx < foreachThreadOp.getRank(); ++idx) {
+          if (foreachThreadOp.getThreadIndices()[idx] == iv) {
+            lb = OpBuilder(iv.getContext()).getIndexAttr(0);
+            ub = foreachThreadOp.getNumThreads()[idx];
+            step = OpBuilder(iv.getContext()).getIndexAttr(1);
+            return success();
+          }
+        }
+        return failure();
+      }
       return failure();
     };
 
     return scf::canonicalizeMinMaxOpInLoop(rewriter, op, op.getAffineMap(),
                                            op.operands(), IsMin, loopMatcher);
   }
 };
 

mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp

canonicalizeMinMaxOp(RewriterBase &rewriter, Operation *op, AffineMap map,
   mlir::canonicalizeMapAndOperands(&newMap, &newOperands);
   rewriter.setInsertionPoint(op);
   rewriter.replaceOpWithNewOp<AffineApplyOp>(op, newMap, newOperands);
   return success();
 }
 
 static LogicalResult
 addLoopRangeConstraints(FlatAffineValueConstraints &constraints, Value iv,
-                        Value lb, Value ub, Value step,
+                        OpFoldResult lb, OpFoldResult ub, OpFoldResult step,
                         RewriterBase &rewriter) {
   // IntegerPolyhedron does not support semi-affine expressions.
   // Therefore, only constant step values are supported.
   auto stepInt = getConstantIntValue(step);
   if (!stepInt)
     return failure();
 
   unsigned dimIv = constraints.appendDimId(iv);
-  unsigned dimLb = constraints.appendDimId(lb);
-  unsigned dimUb = constraints.appendDimId(ub);
+  auto lbv = lb.dyn_cast<Value>();
+  unsigned dimLb =
+      lbv ? constraints.appendDimId(lbv) : constraints.appendDimId(/*num=*/1);
+  auto ubv = ub.dyn_cast<Value>();
+  unsigned dimUb =
+      ubv ? constraints.appendDimId(ubv) : constraints.appendDimId(/*num=*/1);
 
   // If loop lower/upper bounds are constant: Add EQ constraint.
   Optional<int64_t> lbInt = getConstantIntValue(lb);
   Optional<int64_t> ubInt = getConstantIntValue(ub);
   if (lbInt)
     constraints.addBound(IntegerPolyhedron::EQ, dimLb, *lbInt);
   if (ubInt)
     constraints.addBound(IntegerPolyhedron::EQ, dimUb, *ubInt);
LogicalResult scf::canonicalizeMinMaxOpInLoop(RewriterBase &rewriter,
   for (Value operand : operands) {
     // Skip duplicate ivs.
     if (llvm::is_contained(allIvs, operand))
       continue;
 
     // If `operand` is an iteration variable: Find corresponding loop
     // bounds and step.
     Value iv = operand;
-    Value lb, ub, step;
+    OpFoldResult lb, ub, step;
     if (failed(loopMatcher(operand, lb, ub, step)))
       continue;
     allIvs.insert(iv);
 
     if (failed(
             addLoopRangeConstraints(constraints, iv, lb, ub, step, rewriter)))
       return failure();
   }

mlir/test/Dialect/SCF/foreach-thread-canonicalization.mlir

+// RUN: mlir-opt %s -scf-for-loop-canonicalization -canonicalize | FileCheck %s
+
+func.func @reduce() -> tensor<128xf32> {
+  %c2 = arith.constant 2 : index
+  %cst = arith.constant dense<1.000000e+00> : tensor<1x128x384xf32>
+  %cst_0 = arith.constant -0.000000e+00 : f32
+  %0 = linalg.init_tensor [128, 384] : tensor<128x384xf32>
+  %1 = linalg.fill ins(%cst_0 : f32) outs(%0 : tensor<128x384xf32>) -> tensor<128x384xf32>
+  %2 = linalg.init_tensor [128] : tensor<128xf32>
+  %3 = linalg.fill ins(%cst_0 : f32) outs(%2 : tensor<128xf32>) -> tensor<128xf32>
+  %4 = scf.foreach_thread (%arg0) in (%c2) -> (tensor<128xf32>) {
+    %7 = affine.min affine_map<(d0) -> (d0 * -64 + 128, 64)>(%arg0)
+    %8 = affine.max affine_map<(d0) -> (0, d0)>(%7)
+    %9 = affine.apply affine_map<(d0) -> (d0 * 64)>(%arg0)
+    %10 = affine.min affine_map<(d0, d1) -> (d1 * -64 + 128, d0)>(%8, %arg0)
+
+    // CHECK: tensor.extract_slice %{{.*}}[%{{.*}}, 0] [64, 384] [1, 1] : tensor<128x384xf32> to tensor<64x384xf32>
+    // CHECK: tensor.extract_slice %{{.*}}[%{{.*}}] [64] [1] : tensor<128xf32> to tensor<64xf32>
+    %11 = tensor.extract_slice %1[%9, 0] [%10, 384] [1, 1] : tensor<128x384xf32> to tensor<?x384xf32>
+    %12 = tensor.extract_slice %3[%9] [%10] [1] : tensor<128xf32> to tensor<?xf32>
+
+    // CHECK: linalg.generic {{.*}} ins(%{{.*}} : tensor<64x384xf32>) outs(%{{.*}} : tensor<64xf32>) {
+    %13 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>], iterator_types = ["parallel", "reduction"]} ins(%11 : tensor<?x384xf32>) outs(%12 : tensor<?xf32>) {
+    ^bb0(%arg1: f32, %arg2: f32):
+      %14 = arith.addf %arg1, %arg2 : f32
+      linalg.yield %14 : f32
+    } -> tensor<?xf32>
+
+    // TODO: canonicalize this cast away.
+    // CHECK: %[[dyn_casted:.*]] = tensor.cast %{{.*}} : tensor<64xf32> to tensor<?xf32>
+    // CHECK: scf.foreach_thread.parallel_insert_slice %[[dyn_casted:.*]] into %{{.*}}[%{{.*}}] [64] [1] : tensor<?xf32> into tensor<128xf32>
+    scf.foreach_thread.perform_concurrently {
+      scf.foreach_thread.parallel_insert_slice %13 into %3[%9] [%10] [1] : tensor<?xf32> into tensor<128xf32>
+    }
+  }
+  return %4 : tensor<128xf32>
+}