[mlir] AsyncParallelFor: sink constants into the parallel compute function

[ezhulenev]

Summary

With complex recursive structure of async dispatch function LLVM can't always propagate constants to the parallel_compute_fn and it often prevents optimizations like loop unrolling and vectorization. We help LLVM by pushing known constants into the parallel_compute_fn explicitly.

[ezhulenev]

Fix LLVM style violations

[ezhulenev]

Cleanup

[bkramer]

looks good

mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp

 #include "PassDetail.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Async/IR/Async.h"
 #include "mlir/Dialect/Async/Passes.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
+#include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/RegionUtils.h"
 
 using namespace mlir;
 using namespace mlir::async;
 
 #define DEBUG_TYPE "async-parallel-for"

 private:
   bool asyncDispatch;
   int32_t numWorkerThreads;
   int32_t minTaskSize;
 };
 
 struct ParallelComputeFunctionType {
   FunctionType type;
-  llvm::SmallVector<Value> captures;
+  SmallVector<Value> captures;
+};
+
+// Helper struct to parse parallel compute function argument list.
+struct ParallelComputeFunctionArgs {
+  BlockArgument blockIndex();
+  BlockArgument blockSize();
+  ArrayRef<BlockArgument> tripCounts();
+  ArrayRef<BlockArgument> lowerBounds();
+  ArrayRef<BlockArgument> upperBounds();
+  ArrayRef<BlockArgument> steps();
+  ArrayRef<BlockArgument> captures();
+
+  unsigned numLoops;
+  ArrayRef<BlockArgument> args;
+};
+
+struct ParallelComputeFunctionBounds {
+  SmallVector<IntegerAttr> tripCounts;
+  SmallVector<IntegerAttr> lowerBounds;
+  SmallVector<IntegerAttr> upperBounds;
+  SmallVector<IntegerAttr> steps;
 };
 
 struct ParallelComputeFunction {
+  unsigned numLoops;
   FuncOp func;
   llvm::SmallVector<Value> captures;
 };
 
 } // namespace
 
+BlockArgument ParallelComputeFunctionArgs::blockIndex() { return args[0]; }
+BlockArgument ParallelComputeFunctionArgs::blockSize() { return args[1]; }
+
+ArrayRef<BlockArgument> ParallelComputeFunctionArgs::tripCounts() {
+  return args.drop_front(2).take_front(numLoops);
+}
+
+ArrayRef<BlockArgument> ParallelComputeFunctionArgs::lowerBounds() {
+  return args.drop_front(2 + 1 * numLoops).take_front(numLoops);
+}
+
+ArrayRef<BlockArgument> ParallelComputeFunctionArgs::upperBounds() {
+  return args.drop_front(2 + 2 * numLoops).take_front(numLoops);
+}
+
+ArrayRef<BlockArgument> ParallelComputeFunctionArgs::steps() {
+  return args.drop_front(2 + 3 * numLoops).take_front(numLoops);
+}
+
+ArrayRef<BlockArgument> ParallelComputeFunctionArgs::captures() {
+  return args.drop_front(2 + 4 * numLoops);
+}
+
+template <typename ValueRange>
+static SmallVector<IntegerAttr> integerConstants(ValueRange values) {
+  SmallVector<IntegerAttr> attrs(values.size());
+  for (unsigned i = 0; i < values.size(); ++i)
+    matchPattern(values[i], m_Constant(&attrs[i]));
+  return attrs;
+}
+
 // Converts one-dimensional iteration index in the [0, tripCount) interval
 // into multidimensional iteration coordinate.
 static SmallVector<Value> delinearize(ImplicitLocOpBuilder &b, Value index,
                                       ArrayRef<Value> tripCounts) {
   SmallVector<Value> coords(tripCounts.size());
   assert(!tripCounts.empty() && "tripCounts must be not empty");
 
   for (ssize_t i = tripCounts.size() - 1; i >= 0; --i) {
     coords[i] = b.create<arith::RemSIOp>(index, tripCounts[i]);
     index = b.create<arith::DivSIOp>(index, tripCounts[i]);
   }
 
   return coords;
 }
 
 // Returns a function type and implicit captures for a parallel compute
 // function. We'll need a list of implicit captures to setup block and value
 // mapping when we'll clone the body of the parallel operation.
 static ParallelComputeFunctionType
 getParallelComputeFunctionType(scf::ParallelOp op, PatternRewriter &rewriter) {
   // Values implicitly captured by the parallel operation.
   llvm::SetVector<Value> captures;
   getUsedValuesDefinedAbove(op.region(), op.region(), captures);
 
-  llvm::SmallVector<Type> inputs;
+  SmallVector<Type> inputs;
   inputs.reserve(2 + 4 * op.getNumLoops() + captures.size());
 
   Type indexTy = rewriter.getIndexType();
 
   // One-dimensional iteration space defined by the block index and size.
   inputs.push_back(indexTy); // blockIndex
   inputs.push_back(indexTy); // blockSize
 
   // Multi-dimensional parallel iteration space defined by the loop trip counts.
   for (unsigned i = 0; i < op.getNumLoops(); ++i)
     inputs.push_back(indexTy); // loop tripCount
 
-  // Parallel operation lower bound, upper bound and step.
+  // Parallel operation lower bound, upper bound and step. Lower bound, upper
+  // bound and step passed as contiguous arguments:
+  //   call @compute(%lb0, %lb1, ..., %ub0, %ub1, ..., %step0, %step1, ...)
   for (unsigned i = 0; i < op.getNumLoops(); ++i) {
     inputs.push_back(indexTy); // lower bound
     inputs.push_back(indexTy); // upper bound
     inputs.push_back(indexTy); // step
   }
 
   // Types of the implicit captures.
   for (Value capture : captures)
     inputs.push_back(capture.getType());
 
   // Convert captures to vector for later convenience.
   SmallVector<Value> capturesVector(captures.begin(), captures.end());
   return {rewriter.getFunctionType(inputs, TypeRange()), capturesVector};
 }
 
 // Create a parallel compute fuction from the parallel operation.
 static ParallelComputeFunction
-createParallelComputeFunction(scf::ParallelOp op, PatternRewriter &rewriter) {
+createParallelComputeFunction(scf::ParallelOp op,
+                              ParallelComputeFunctionBounds bounds,
+                              PatternRewriter &rewriter) {
   OpBuilder::InsertionGuard guard(rewriter);
   ImplicitLocOpBuilder b(op.getLoc(), rewriter);
 
   ModuleOp module = op->getParentOfType<ModuleOp>();
 
-  // Make sure that all constants will be inside the parallel operation body to
-  // reduce the number of parallel compute function arguments.
-  cloneConstantsIntoTheRegion(op.getLoopBody(), rewriter);
-
   ParallelComputeFunctionType computeFuncType =
       getParallelComputeFunctionType(op, rewriter);
 
   FunctionType type = computeFuncType.type;
   FuncOp func = FuncOp::create(op.getLoc(), "parallel_compute_fn", type);
   func.setPrivate();
 
   // Insert function into the module symbol table and assign it unique name.
   SymbolTable symbolTable(module);
   symbolTable.insert(func);
   rewriter.getListener()->notifyOperationInserted(func);
 
   // Create function entry block.
   Block *block = b.createBlock(&func.getBody(), func.begin(), type.getInputs());
   b.setInsertionPointToEnd(block);
 
-  unsigned offset = 0; // argument offset for arguments decoding
-
-  // Returns `numArguments` arguments starting from `offset` and updates offset
-  // by moving forward to the next argument.
-  auto getArguments = [&](unsigned numArguments) -> ArrayRef<Value> {
-    auto args = block->getArguments();
-    auto slice = args.drop_front(offset).take_front(numArguments);
-    offset += numArguments;
-    return {slice.begin(), slice.end()};
-  };
+  ParallelComputeFunctionArgs args = {op.getNumLoops(), func.getArguments()};
 
   // Block iteration position defined by the block index and size.
-  Value blockIndex = block->getArgument(offset++);
-  Value blockSize = block->getArgument(offset++);
+  BlockArgument blockIndex = args.blockIndex();
+  BlockArgument blockSize = args.blockSize();
 
   // Constants used below.
   Value c0 = b.create<arith::ConstantIndexOp>(0);
   Value c1 = b.create<arith::ConstantIndexOp>(1);
 
+  // Materialize known constants as constant operation in the function body.
+  auto values = [&](ArrayRef<BlockArgument> args, ArrayRef<IntegerAttr> attrs) {
+    return llvm::to_vector(
+        llvm::map_range(llvm::zip(args, attrs), [&](auto tuple) -> Value {
+          if (IntegerAttr attr = std::get<1>(tuple))
+            return b.create<ConstantOp>(attr);
+          return std::get<0>(tuple);
+        }));
+  };
+
   // Multi-dimensional parallel iteration space defined by the loop trip counts.
-  ArrayRef<Value> tripCounts = getArguments(op.getNumLoops());
+  auto tripCounts = values(args.tripCounts(), bounds.tripCounts);
+
+  // Parallel operation lower bound and step.
+  auto lowerBounds = values(args.lowerBounds(), bounds.lowerBounds);
+  auto steps = values(args.steps(), bounds.steps);
+
+  // Remaining arguments are implicit captures of the parallel operation.
+  ArrayRef<BlockArgument> captures = args.captures();
 
   // Compute a product of trip counts to get the size of the flattened
   // one-dimensional iteration space.
   Value tripCount = tripCounts[0];
   for (unsigned i = 1; i < tripCounts.size(); ++i)
     tripCount = b.create<arith::MulIOp>(tripCount, tripCounts[i]);
 
-  // Parallel operation lower bound and step.
-  ArrayRef<Value> lowerBound = getArguments(op.getNumLoops());
-  offset += op.getNumLoops(); // skip upper bound arguments
-  ArrayRef<Value> step = getArguments(op.getNumLoops());
-
-  // Remaining arguments are implicit captures of the parallel operation.
-  ArrayRef<Value> captures = getArguments(block->getNumArguments() - offset);
-
   // Find one-dimensional iteration bounds: [blockFirstIndex, blockLastIndex]:
   //   blockFirstIndex = blockIndex * blockSize
   Value blockFirstIndex = b.create<arith::MulIOp>(blockIndex, blockSize);
 
   // The last one-dimensional index in the block defined by the `blockIndex`:
   //   blockLastIndex = min(blockFirstIndex + blockSize, tripCount) - 1
   Value blockEnd0 = b.create<arith::AddIOp>(blockFirstIndex, blockSize);
   Value blockEnd1 = b.create<arith::MinSIOp>(blockEnd0, tripCount);
createParallelComputeFunction(scf::ParallelOp op,
   // work concurrently.
   LoopNestBuilder workLoopBuilder = [&](size_t loopIdx) -> LoopBodyBuilder {
     return [&, loopIdx](OpBuilder &nestedBuilder, Location loc, Value iv,
                         ValueRange args) {
       ImplicitLocOpBuilder nb(loc, nestedBuilder);
 
       // Compute induction variable for `loopIdx`.
       computeBlockInductionVars[loopIdx] = nb.create<arith::AddIOp>(
-          lowerBound[loopIdx], nb.create<arith::MulIOp>(iv, step[loopIdx]));
+          lowerBounds[loopIdx], nb.create<arith::MulIOp>(iv, steps[loopIdx]));
 
       // Check if we are inside first or last iteration of the loop.
       isBlockFirstCoord[loopIdx] = nb.create<arith::CmpIOp>(
           arith::CmpIPredicate::eq, iv, blockFirstCoord[loopIdx]);
       isBlockLastCoord[loopIdx] = nb.create<arith::CmpIOp>(
           arith::CmpIPredicate::eq, iv, blockLastCoord[loopIdx]);
 
       // Check if the previous loop is in its first or last iteration.
return [&, loopIdx](OpBuilder &nestedBuilder, Location loc, Value iv,
         nb.clone(bodyOp, mapping);
     };
   };
 
   b.create<scf::ForOp>(blockFirstCoord[0], blockEndCoord[0], c1, ValueRange(),
                        workLoopBuilder(0));
   b.create<ReturnOp>(ValueRange());
 
-  return {func, std::move(computeFuncType.captures)};
+  return {op.getNumLoops(), func, std::move(computeFuncType.captures)};
 }
 
 // Creates recursive async dispatch function for the given parallel compute
 // function. Dispatch function keeps splitting block range into halves until it
 // reaches a single block, and then excecutes it inline.
 //
 // Function pseudocode (mix of C++ and MLIR):
 //

 AsyncParallelForRewrite::matchAndRewrite(scf::ParallelOp op,
                                          PatternRewriter &rewriter) const {
   // We do not currently support rewrite for parallel op with reductions.
   if (op.getNumReductions() != 0)
     return failure();
 
   ImplicitLocOpBuilder b(op.getLoc(), rewriter);
 
+  // Make sure that all constants will be inside the parallel operation body to
+  // reduce the number of parallel compute function arguments.
+  cloneConstantsIntoTheRegion(op.getLoopBody(), rewriter);
+
   // Compute trip count for each loop induction variable:
   //   tripCount = ceil_div(upperBound - lowerBound, step);
   SmallVector<Value> tripCounts(op.getNumLoops());
   for (size_t i = 0; i < op.getNumLoops(); ++i) {
     auto lb = op.lowerBound()[i];
     auto ub = op.upperBound()[i];
     auto step = op.step()[i];
-    auto range = b.create<arith::SubIOp>(ub, lb);
-    tripCounts[i] = b.create<arith::CeilDivSIOp>(range, step);
+    auto range = b.createOrFold<arith::SubIOp>(ub, lb);
+    tripCounts[i] = b.createOrFold<arith::CeilDivSIOp>(range, step);
   }
 
   // Compute a product of trip counts to get the 1-dimensional iteration space
   // for the scf.parallel operation.
   Value tripCount = tripCounts[0];
   for (size_t i = 1; i < tripCounts.size(); ++i)
     tripCount = b.create<arith::MulIOp>(tripCount, tripCounts[i]);
 
auto dispatch = [&](OpBuilder &nestedBuilder, Location loc) {
     //   blockSize = min(tripCount,
     //                   max(ceil_div(tripCount, maxComputeBlocks),
     //                       ceil_div(minTaskSize, bodySize)))
     Value bs0 = b.create<arith::CeilDivSIOp>(tripCount, maxComputeBlocks);
     Value bs1 = b.create<arith::MaxSIOp>(bs0, minTaskSizeCst);
     Value blockSize = b.create<arith::MinSIOp>(tripCount, bs1);
     Value blockCount = b.create<arith::CeilDivSIOp>(tripCount, blockSize);
 
+    // Collect statically known constants defining the loop nest in the parallel
+    // compute function. LLVM can't always push constants across the non-trivial
+    // async dispatch call graph, by providing these values explicitly we can
+    // choose to build more efficient loop nest, and rely on a better constant
+    // folding, loop unrolling and vectorization.
+    ParallelComputeFunctionBounds staticBounds = {
+        integerConstants(tripCounts),
+        integerConstants(op.lowerBound()),
+        integerConstants(op.upperBound()),
+        integerConstants(op.step()),
+    };
+
     // Create a parallel compute function that takes a block id and computes the
     // parallel operation body for a subset of iteration space.
     ParallelComputeFunction parallelComputeFunction =
-        createParallelComputeFunction(op, rewriter);
+        createParallelComputeFunction(op, staticBounds, rewriter);
 
     // Dispatch parallel compute function using async recursive work splitting,
     // or by submitting compute task sequentially from a caller thread.
     if (asyncDispatch) {
       doAsyncDispatch(b, rewriter, parallelComputeFunction, op, blockSize,
                       blockCount, tripCounts);
     } else {
       doSequentialDispatch(b, rewriter, parallelComputeFunction, op, blockSize,

mlir/test/Dialect/Async/async-parallel-for-compute-fn.mlir

-// RUN: mlir-opt %s                                                            \
+// RUN: mlir-opt %s -split-input-file                                          \
 // RUN:    -async-parallel-for=async-dispatch=true                             \
 // RUN: | FileCheck %s
 
-// RUN: mlir-opt %s                                                            \
+// RUN: mlir-opt %s -split-input-file                                          \
 // RUN:    -async-parallel-for=async-dispatch=false                            \
 // RUN:    -canonicalize -inline -symbol-dce                                   \
 // RUN: | FileCheck %s
 
 // Check that constants defined outside of the `scf.parallel` body will be
 // sunk into the parallel compute function to avoid blowing up the number
 // of parallel compute function arguments.
 

 // CHECK-SAME:   %[[LB:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[UB:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[STEP:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[MEMREF:arg[0-9]+]]: memref<?xf32>
 // CHECK-SAME: ) {
 // CHECK:        %[[CST:.*]] = arith.constant 1.0{{.*}} : f32
 // CHECK:        scf.for
 // CHECK:          memref.store %[[CST]], %[[MEMREF]]
+
+// -----
+
+// Check that constant loop bound sunk into the parallel compute function.
+
+// CHECK-LABEL: func @sink_constant_step(
+func @sink_constant_step(%arg0: memref<?xf32>, %lb: index, %ub: index) {
+  %one = arith.constant 1.0 : f32
+  %st = arith.constant 123 : index
+
+  scf.parallel (%i) = (%lb) to (%ub) step (%st) {
+    memref.store %one, %arg0[%i] : memref<?xf32>
+  }
+
+  return
+}
+
+// CHECK-LABEL: func private @parallel_compute_fn(
+// CHECK-SAME:   %[[BLOCK_INDEX:arg[0-9]+]]: index,
+// CHECK-SAME:   %[[BLOCK_SIZE:arg[0-9]+]]: index,
+// CHECK-SAME:   %[[TRIP_COUNT:arg[0-9]+]]: index,
+// CHECK-SAME:   %[[LB:arg[0-9]+]]: index,
+// CHECK-SAME:   %[[UB:arg[0-9]+]]: index,
+// CHECK-SAME:   %[[STEP:arg[0-9]+]]: index,
+// CHECK-SAME:   %[[MEMREF:arg[0-9]+]]: memref<?xf32>
+// CHECK-SAME: ) {
+// CHECK:        %[[CSTEP:.*]] = arith.constant 123 : index
+// CHECK-NOT:    %[[STEP]]
+// CHECK:        scf.for %[[I:arg[0-9]+]]
+// CHECK:          %[[TMP:.*]] = arith.muli %[[I]], %[[CSTEP]]
+// CHECK:          %[[IDX:.*]] = arith.addi %[[LB]], %[[TMP]]
+// CHECK:          memref.store
 No newline at end of file