[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"
LogicalResult matchAndRewrite(scf::ParallelOp op,
                                 PatternRewriter &rewriter) const override;
 
 private:
   bool asyncDispatch;
   int32_t numWorkerThreads;
   int32_t minTaskSize;
 };
 
-struct ParallelComputeFunctionType {
-  FunctionType type;
-  llvm::SmallVector<Value> captures;
+class ParallelComputeFunctionType {
+public:
+  ParallelComputeFunctionType(FunctionType type, SmallVector<Value> captures,
+                              unsigned num_loops)
+      : type_(type), captures_(std::move(captures)), num_loops_(num_loops) {}
+
+  FunctionType type() const { return type_; }
+  const SmallVector<Value> &captures() const { return captures_; }
+
+  // Helper functions to parse the parallel compute function arguments.
+  BlockArgument blockIndex(ArrayRef<BlockArgument> args) const;
+  BlockArgument blockSize(ArrayRef<BlockArgument> args) const;
+  ArrayRef<BlockArgument> tripCounts(ArrayRef<BlockArgument> args) const;
+  ArrayRef<BlockArgument> lowerBounds(ArrayRef<BlockArgument> args) const;
+  ArrayRef<BlockArgument> upperBounds(ArrayRef<BlockArgument> args) const;
+  ArrayRef<BlockArgument> steps(ArrayRef<BlockArgument> args) const;
+  ArrayRef<BlockArgument> captures(ArrayRef<BlockArgument> args) const;
+
+private:
+  FunctionType type_;           // signatrue of the parallel compute function
+  SmallVector<Value> captures_; // values captured by the parallel operation
+  unsigned num_loops_;          // depth of the scf.for loop nest
 };
 
-struct ParallelComputeFunction {
-  FuncOp func;
-  llvm::SmallVector<Value> captures;
+class ParallelComputeFunction {
+public:
+  ParallelComputeFunction(ParallelComputeFunctionType type, FuncOp func)
+      : type_(std::move(type)), func_(func) {}
+
+  FuncOp func() const { return func_; }
+  const SmallVector<Value> &captures() const { return type_.captures(); }
+
+  // Helper functions to replace block arguments with statically known constants
+  // in the function body. We rely un later dead argument elimination to cleanup
+  // unused arguments.
+  void setTripCounts(ArrayRef<IntegerAttr> attrs);
+  void setLowerBounds(ArrayRef<IntegerAttr> attrs);
+  void setUpperBounds(ArrayRef<IntegerAttr> attrs);
+  void setSteps(ArrayRef<IntegerAttr> attrs);
+
+private:
+  ArrayRef<BlockArgument> args() { return func_.getArguments(); }
+
+  ParallelComputeFunctionType type_;
+  FuncOp func_;
 };
 
 } // namespace
 
+BlockArgument
+ParallelComputeFunctionType::blockIndex(ArrayRef<BlockArgument> args) const {
+  return args[0];
+}
+
+BlockArgument
+ParallelComputeFunctionType::blockSize(ArrayRef<BlockArgument> args) const {
+  return args[1];
+}
+
+ArrayRef<BlockArgument>
+ParallelComputeFunctionType::tripCounts(ArrayRef<BlockArgument> args) const {
+  return args.drop_front(2).take_front(num_loops_);
+}
+
+ArrayRef<BlockArgument>
+ParallelComputeFunctionType::lowerBounds(ArrayRef<BlockArgument> args) const {
+  return args.drop_front(2 + 1 * num_loops_).take_front(num_loops_);
+}
+
+ArrayRef<BlockArgument>
+ParallelComputeFunctionType::upperBounds(ArrayRef<BlockArgument> args) const {
+  return args.drop_front(2 + 2 * num_loops_).take_front(num_loops_);
+}
+
+ArrayRef<BlockArgument>
+ParallelComputeFunctionType::steps(ArrayRef<BlockArgument> args) const {
+  return args.drop_front(2 + 3 * num_loops_).take_front(num_loops_);
+}
+
+ArrayRef<BlockArgument>
+ParallelComputeFunctionType::captures(ArrayRef<BlockArgument> args) const {
+  return args.drop_front(2 + 4 * num_loops_);
+}
+
+// Replace all uses of `args` with constants created in the function body from
+// the `attrs` (if attribute is not null).
+static void replaceWithConstants(FuncOp func, ArrayRef<BlockArgument> args,
+                                 ArrayRef<IntegerAttr> attrs) {
+  // Create constants at the beginning of the entry block.
+  OpBuilder b = OpBuilder::atBlockBegin(&func.body().front());
+  assert(args.size() == attrs.size() && "illegal `args` and `attrs` sizes");
+
+  for (auto tuple : llvm::zip(args, attrs)) {
+    Value value = std::get<0>(tuple);
+    if (IntegerAttr attr = std::get<1>(tuple))
+      value.replaceAllUsesWith(b.create<ConstantOp>(func.getLoc(), attr));
+  }
+}
+
+void ParallelComputeFunction::setTripCounts(ArrayRef<IntegerAttr> attrs) {
+  replaceWithConstants(func_, type_.tripCounts(args()), attrs);
+}
+
+void ParallelComputeFunction::setLowerBounds(ArrayRef<IntegerAttr> attrs) {
+  replaceWithConstants(func_, type_.lowerBounds(args()), attrs);
+}
+
+void ParallelComputeFunction::setUpperBounds(ArrayRef<IntegerAttr> attrs) {
+  replaceWithConstants(func_, type_.upperBounds(args()), attrs);
+}
+
+void ParallelComputeFunction::setSteps(ArrayRef<IntegerAttr> attrs) {
+  replaceWithConstants(func_, type_.steps(args()), 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) {
+                                      ArrayRef<BlockArgument> 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};
+  return {rewriter.getFunctionType(inputs, TypeRange()), capturesVector,
+          op.getNumLoops()};
 }
 
 // Create a parallel compute fuction from the parallel operation.
 static ParallelComputeFunction
 createParallelComputeFunction(scf::ParallelOp op, 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;
+  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()};
-  };
+  ArrayRef<BlockArgument> args = block->getArguments();
 
   // Block iteration position defined by the block index and size.
-  Value blockIndex = block->getArgument(offset++);
-  Value blockSize = block->getArgument(offset++);
+  BlockArgument blockIndex = computeFuncType.blockIndex(args);
+  BlockArgument blockSize = computeFuncType.blockSize(args);
 
   // Constants used below.
   Value c0 = b.create<arith::ConstantIndexOp>(0);
   Value c1 = b.create<arith::ConstantIndexOp>(1);
 
   // Multi-dimensional parallel iteration space defined by the loop trip counts.
-  ArrayRef<Value> tripCounts = getArguments(op.getNumLoops());
+  ArrayRef<BlockArgument> tripCounts = computeFuncType.tripCounts(args);
 
   // 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());
+  ArrayRef<BlockArgument> lowerBound = computeFuncType.lowerBounds(args);
+  ArrayRef<BlockArgument> step = computeFuncType.steps(args);
 
   // Remaining arguments are implicit captures of the parallel operation.
-  ArrayRef<Value> captures = getArguments(block->getNumArguments() - offset);
+  ArrayRef<BlockArgument> captures = computeFuncType.captures(args);
 
   // 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`:
   Value blockEnd0 = b.create<arith::AddIOp>(blockFirstIndex, blockSize);
   Value blockEnd1 = b.create<arith::MinSIOp>(blockEnd0, tripCount);
return [&, loopIdx](OpBuilder &nestedBuilder, Location loc, Value iv,
                               workLoopBuilder(loopIdx + 1));
         nb.create<scf::YieldOp>(loc);
         return;
       }
 
       // Copy the body of the parallel op into the inner-most loop.
       BlockAndValueMapping mapping;
       mapping.map(op.getInductionVars(), computeBlockInductionVars);
-      mapping.map(computeFuncType.captures, captures);
+      mapping.map(computeFuncType.captures(), captures);
 
       for (auto &bodyOp : op.getLoopBody().getOps())
         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 {computeFuncType, func};
 }
 
 // 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):
 //
 //   func @async_dispatch(%block_start : index, %block_end : index, ...) {
 //
 //     // Keep splitting block range until we reached a range of size 1.
 //     while (%block_end - %block_start > 1) {
 //       %mid_index = block_start + (block_end - block_start) / 2;
 //       async.execute { call @async_dispatch(%mid_index, %block_end); }
 //       %block_end = %mid_index
 //     }
 //
 //     // Call parallel compute function for a single block.
 //     call @parallel_compute_fn(%block_start, %block_size, ...);
 //   }
 //
 static FuncOp createAsyncDispatchFunction(ParallelComputeFunction &computeFunc,
                                           PatternRewriter &rewriter) {
   OpBuilder::InsertionGuard guard(rewriter);
-  Location loc = computeFunc.func.getLoc();
+  Location loc = computeFunc.func().getLoc();
   ImplicitLocOpBuilder b(loc, rewriter);
 
-  ModuleOp module = computeFunc.func->getParentOfType<ModuleOp>();
+  ModuleOp module = computeFunc.func()->getParentOfType<ModuleOp>();
 
   ArrayRef<Type> computeFuncInputTypes =
-      computeFunc.func.type().cast<FunctionType>().getInputs();
+      computeFunc.func().type().cast<FunctionType>().getInputs();
 
   // Compared to the parallel compute function async dispatch function takes
   // additional !async.group argument. Also instead of a single `blockIndex` it
   // takes `blockStart` and `blockEnd` arguments to define the range of
   // dispatched blocks.
   SmallVector<Type> inputTypes;
   inputTypes.push_back(async::GroupType::get(rewriter.getContext()));
   inputTypes.push_back(rewriter.getIndexType()); // add blockStart argument
static FuncOp createAsyncDispatchFunction(ParallelComputeFunction &computeFunc,
   // call the parallel compute function for the first block of the range.
   b.setInsertionPointAfter(whileOp);
 
   // Drop async dispatch specific arguments: async group, block start and end.
   auto forwardedInputs = block->getArguments().drop_front(3);
   SmallVector<Value> computeFuncOperands = {blockStart};
   computeFuncOperands.append(forwardedInputs.begin(), forwardedInputs.end());
 
-  b.create<CallOp>(computeFunc.func.sym_name(),
-                   computeFunc.func.getCallableResults(), computeFuncOperands);
+  b.create<CallOp>(computeFunc.func().sym_name(),
+                   computeFunc.func().getCallableResults(),
+                   computeFuncOperands);
   b.create<ReturnOp>(ValueRange());
 
   return func;
 }
 
 // Launch async dispatch of the parallel compute function.
 static void doAsyncDispatch(ImplicitLocOpBuilder &b, PatternRewriter &rewriter,
                             ParallelComputeFunction &parallelComputeFunction,
static void doAsyncDispatch(ImplicitLocOpBuilder &b, PatternRewriter &rewriter,
 
   // Appends operands shared by async dispatch and parallel compute functions to
   // the given operands vector.
   auto appendBlockComputeOperands = [&](SmallVector<Value> &operands) {
     operands.append(tripCounts);
     operands.append(op.lowerBound().begin(), op.lowerBound().end());
     operands.append(op.upperBound().begin(), op.upperBound().end());
     operands.append(op.step().begin(), op.step().end());
-    operands.append(parallelComputeFunction.captures);
+    operands.append(parallelComputeFunction.captures());
   };
 
   // Check if the block size is one, in this case we can skip the async dispatch
   // completely. If this will be known statically, then canonicalization will
   // erase async group operations.
   Value isSingleBlock =
       b.create<arith::CmpIOp>(arith::CmpIPredicate::eq, blockCount, c1);
 
   auto syncDispatch = [&](OpBuilder &nestedBuilder, Location loc) {
     ImplicitLocOpBuilder nb(loc, nestedBuilder);
 
     // Call parallel compute function for the single block.
     SmallVector<Value> operands = {c0, blockSize};
     appendBlockComputeOperands(operands);
 
-    nb.create<CallOp>(parallelComputeFunction.func.sym_name(),
-                      parallelComputeFunction.func.getCallableResults(),
+    nb.create<CallOp>(parallelComputeFunction.func().sym_name(),
+                      parallelComputeFunction.func().getCallableResults(),
                       operands);
     nb.create<scf::YieldOp>();
   };
 
   auto asyncDispatch = [&](OpBuilder &nestedBuilder, Location loc) {
     // Create an async.group to wait on all async tokens from the concurrent
     // execution of multiple parallel compute function. First block will be
     // executed synchronously in the caller thread.

 // from a simple for loop in the caller thread.
 static void
 doSequentialDispatch(ImplicitLocOpBuilder &b, PatternRewriter &rewriter,
                      ParallelComputeFunction &parallelComputeFunction,
                      scf::ParallelOp op, Value blockSize, Value blockCount,
                      const SmallVector<Value> &tripCounts) {
   MLIRContext *ctx = op->getContext();
 
-  FuncOp compute = parallelComputeFunction.func;
+  FuncOp compute = parallelComputeFunction.func();
 
   Value c0 = b.create<arith::ConstantIndexOp>(0);
   Value c1 = b.create<arith::ConstantIndexOp>(1);
 
   // Create an async.group to wait on all async tokens from the concurrent
   // execution of multiple parallel compute function. First block will be
   // executed synchronously in the caller thread.
   Value groupSize = b.create<arith::SubIOp>(blockCount, c1);
   Value group = b.create<CreateGroupOp>(GroupType::get(ctx), groupSize);
 
   // Call parallel compute function for all blocks.
   using LoopBodyBuilder =
       std::function<void(OpBuilder &, Location, Value, ValueRange)>;
 
   // Returns parallel compute function operands to process the given block.
   auto computeFuncOperands = [&](Value blockIndex) -> SmallVector<Value> {
     SmallVector<Value> computeFuncOperands = {blockIndex, blockSize};
     computeFuncOperands.append(tripCounts);
     computeFuncOperands.append(op.lowerBound().begin(), op.lowerBound().end());
     computeFuncOperands.append(op.upperBound().begin(), op.upperBound().end());
     computeFuncOperands.append(op.step().begin(), op.step().end());
-    computeFuncOperands.append(parallelComputeFunction.captures);
+    computeFuncOperands.append(parallelComputeFunction.captures());
     return computeFuncOperands;
   };
 
   // Induction variable is the index of the block: [0, blockCount).
   LoopBodyBuilder loopBuilder = [&](OpBuilder &loopBuilder, Location loc,
                                     Value iv, ValueRange args) {
     ImplicitLocOpBuilder nb(loc, loopBuilder);
 

 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) {
     Value blockSize = b.create<arith::MinSIOp>(tripCount, bs1);
     Value blockCount = b.create<arith::CeilDivSIOp>(tripCount, blockSize);
 
     // 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);
 
+    // Sink known loop nest constants into the parallel compute function to
+    // simplify the compute loop. LLVM can't always do that later because of the
+    // complex IR structure, especially in the async dispatch case.
+    auto constants = [&](auto values) -> SmallVector<IntegerAttr> {
+      SmallVector<IntegerAttr> attrs(values.size());
+      for (unsigned i = 0; i < values.size(); ++i)
+        matchPattern(values[i], m_Constant(&attrs[i]));
+      return attrs;
+    };
+
+    parallelComputeFunction.setTripCounts(constants(tripCounts));
+    parallelComputeFunction.setLowerBounds(constants(op.lowerBound()));
+    parallelComputeFunction.setUpperBounds(constants(op.upperBound()));
+    parallelComputeFunction.setSteps(constants(op.step()));
+
     // 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,
                            blockCount, tripCounts);

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