[mlir] Refactor transform dialect's gpu block func

[guraypp]

Summary

This revision refactors gpu block id generator lambda that is used in the transform dialect. It removes the lambda and instead uses a static function that's name generateGpuBlockIds.

It also simplifies arguments that the function takes.

[nicolasvasilache]

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

1487

doc this please.

[guraypp]

add documentation on the block generator

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 
 /// Maps the top level `scf.foreach_thread` op to GPU Thread Blocks. Mapping is
 /// one-to-one and the induction variables of `scf.foreach_thread` are rewritten
 /// to gpu.block_id according to the thread_dim_apping attribute. Dynamic,
 /// `scf.foreach_thread` trip counts are currently not supported. Dynamic block
 /// dim sizes are currently not supported.
 LogicalResult rewriteTopLevelForeachThreadToGpuBlocks(
     RewriterBase &rewriter, scf::ForeachThreadOp foreachThreadOp,
-    function_ref<void(Operation *, const SmallVector<int64_t> &, IndexType,
+    function_ref<void(RewriterBase &, scf::ForeachThreadOp,
                       SmallVector<Value> &)>
         blockIdGenerator,
     SmallVector<int64_t> &gridDims);
 
 /// Finds the top level scf::ForeachThreadOp of given target.
 FailureOr<scf::ForeachThreadOp> findTopLevelForeachThreadOp(Operation *target);
 
 /// Searches `scf.foreach_thread` ops nested under `target` and maps each such

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

 }
 
 //===----------------------------------------------------------------------===//
 // MapNestedForeachThreadToGpuBlocks
 //===----------------------------------------------------------------------===//
 
 LogicalResult mlir::linalg::rewriteTopLevelForeachThreadToGpuBlocks(
     RewriterBase &rewriter, scf::ForeachThreadOp foreachThreadOp,
-    function_ref<void(Operation *, const SmallVector<int64_t> &, IndexType,
+    function_ref<void(RewriterBase &, scf::ForeachThreadOp,
                       SmallVector<Value> &)>
         blockIdGenerator,
     SmallVector<int64_t> &gridDims) {
   if (foreachThreadOp.getNumResults() > 0)
     return foreachThreadOp->emitError(
         "only bufferized scf.foreach_thread lowers to gpu.block_id");
   if (foreachThreadOp.getNumThreads().size() > 3)
     return foreachThreadOp->emitError(
         "scf.foreach_thread with rank > 3 does not lower to gpu.block_id");
 
   // Step 0. Outline the compute workload region and set up the workload
   // operands.
   auto potentialGridDim = foreachThreadOp.getPermutedNumThreads(rewriter);
   if (failed(potentialGridDim) ||
       llvm::any_of(*potentialGridDim, [](OpFoldResult ofr) {
         return !getConstantIntValue(ofr).has_value();
       }))
     return foreachThreadOp->emitError("unsupported dynamic gridDim");
 
   for (OpFoldResult ofr : *potentialGridDim)
     gridDims.push_back(getConstantIntValue(ofr).value());
 
-  IndexType indexType = rewriter.getIndexType();
   SmallVector<Value> blockOps;
-  blockIdGenerator(foreachThreadOp, gridDims, indexType, blockOps);
+  blockIdGenerator(rewriter, foreachThreadOp, blockOps);
 
   // Step 1. Move the body of foreachThreadOp.
   // Erase the terminator first, it will not be used since we are on buffers.
   rewriter.eraseOp(foreachThreadOp.getTerminator());
   Block *targetBlock = foreachThreadOp->getBlock();
   Block::iterator insertionPoint = Block::iterator(foreachThreadOp);
   Block &sourceBlock = foreachThreadOp.getRegion().front();
   targetBlock->getOperations().splice(insertionPoint,
Value blockSizeZ =
       blockDimZ.has_value() ? createConstant(blockDimZ.value()) : one;
   auto launchOp = rewriter.create<gpu::LaunchOp>(
       loc, gridSizeX, gridSizeY, gridSizeZ, blockSizeX, blockSizeY, blockSizeZ);
   rewriter.setInsertionPointToEnd(&launchOp.body().front());
   rewriter.create<gpu::TerminatorOp>(loc);
   return launchOp;
 }
 
+/// This is an helper that is only used in
+/// rewriteTopLevelForeachThreadToGpuBlocks. It generates GPU dialects block_id
+static void generateGpuBlockIds(RewriterBase &rewriter,
+                                scf::ForeachThreadOp foreachOp,
+                                SmallVector<Value> &blockOps) {
+  Location loc = foreachOp->getLoc();
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPoint(foreachOp);
+  IndexType indexType = rewriter.getIndexType();
+  SmallVector<gpu::Dimension> gpuDims{gpu::Dimension::x, gpu::Dimension::y,
+                                      gpu::Dimension::z};
+  for (int64_t idx : llvm::seq<int64_t>(0, gpuDims.size())) {
+    blockOps.push_back(
+        rewriter.create<gpu::BlockIdOp>(loc, indexType, gpuDims[idx]));
+  }
+}
+
 DiagnosedSilenceableFailure
 transform::MapNestedForeachThreadToGpuBlocks::applyToOne(
     Operation *target, SmallVectorImpl<Operation *> &results,
     transform::TransformState &state) {
   gpu::LaunchOp gpuLaunch = dyn_cast<gpu::LaunchOp>(target);
   SimpleRewriter rewriter(getContext());
 
   if (!getGenerateGpuLaunch() && !gpuLaunch) {
if (getGenerateGpuLaunch()) {
     gpuLaunch = *maybeGpuLaunch;
     rewriter.setInsertionPointToStart(&gpuLaunch.body().front());
     Operation *newForeachThreadOp = rewriter.clone(*topLevelForeachThreadOp);
     rewriter.eraseOp(topLevelForeachThreadOp);
     topLevelForeachThreadOp =
         dyn_cast<scf::ForeachThreadOp>(newForeachThreadOp);
   }
 
-  auto generateBlocks = [&](Operation *op, const SmallVector<int64_t> &gridDims,
-                            IndexType indexType, SmallVector<Value> &blockOps) {
-    Location loc = op->getLoc();
-    OpBuilder::InsertionGuard guard(rewriter);
-    rewriter.setInsertionPoint(op);
-    SmallVector<gpu::Dimension> gpuDims{gpu::Dimension::x, gpu::Dimension::y,
-                                        gpu::Dimension::z};
-    for (int64_t idx : llvm::seq<int64_t>(0, gridDims.size())) {
-      blockOps.push_back(
-          rewriter.create<gpu::BlockIdOp>(loc, indexType, gpuDims[idx]));
-    }
-  };
-
   SmallVector<int64_t> gridDim = extractFromI64ArrayAttr(getGridDim());
   if (failed(mlir::linalg::rewriteTopLevelForeachThreadToGpuBlocks(
-          rewriter, topLevelForeachThreadOp, generateBlocks, gridDim)))
+          rewriter, topLevelForeachThreadOp, generateGpuBlockIds, gridDim)))
     return DiagnosedSilenceableFailure(reportUnknownTransformError(target));
 
   if (failed(alterGpuLaunch(rewriter, gpuLaunch, gridDim[0], gridDim[1],
                             gridDim[2])))
     return DiagnosedSilenceableFailure::definiteFailure();
 
   results.assign({gpuLaunch});
   return DiagnosedSilenceableFailure(success());