address review comments

Author: kurapov-peter

mlir/include/mlir/Dialect/GPU/Utils/DistributionUtils.h

@@ -22,38 +22,39 @@ namespace mlir {
 namespace gpu {
 struct WarpDistributionPattern : OpRewritePattern<WarpExecuteOnLane0Op> {
   using OpRewritePattern<WarpExecuteOnLane0Op>::OpRewritePattern;
+
   virtual LogicalResult
   matchAndRewrite(WarpExecuteOnLane0Op op,
                   PatternRewriter &rewriter) const override = 0;
 
 protected:
   /// Return a value yielded by `warpOp` which statifies the filter lamdba
   /// condition and is not dead.
-  static OpOperand *getWarpResult(WarpExecuteOnLane0Op warpOp,
-                                  const std::function<bool(Operation *)> &fn);
+  OpOperand *getWarpResult(WarpExecuteOnLane0Op warpOp,
+                           const std::function<bool(Operation *)> &fn) const;
 
   /// Helper to create a new WarpExecuteOnLane0Op with different signature.
-  static WarpExecuteOnLane0Op moveRegionToNewWarpOpAndReplaceReturns(
+  WarpExecuteOnLane0Op moveRegionToNewWarpOpAndReplaceReturns(
       RewriterBase &rewriter, WarpExecuteOnLane0Op warpOp,
-      ValueRange newYieldedValues, TypeRange newReturnTypes);
+      ValueRange newYieldedValues, TypeRange newReturnTypes) const;
 
   /// Helper to create a new WarpExecuteOnLane0Op region with extra outputs.
   /// `indices` return the index of each new output.
-  static WarpExecuteOnLane0Op moveRegionToNewWarpOpAndAppendReturns(
+  WarpExecuteOnLane0Op moveRegionToNewWarpOpAndAppendReturns(
       RewriterBase &rewriter, WarpExecuteOnLane0Op warpOp,
       ValueRange newYieldedValues, TypeRange newReturnTypes,
-      llvm::SmallVector<size_t> &indices);
+      llvm::SmallVector<size_t> &indices) const;
 
   /// Delinearize the given `laneId` into multiple dimensions, where each
   /// dimension's size is determined by `originalShape` and `distributedShape`
   /// together. This function expects the total numbers of threads needed for
   /// distribution is equal to `warpSize`. Returns true and updates
   /// `delinearizedIds` if so.
-  static bool delinearizeLaneId(OpBuilder &builder, Location loc,
-                                ArrayRef<int64_t> originalShape,
-                                ArrayRef<int64_t> distributedShape,
-                                int64_t warpSize, Value laneId,
-                                SmallVectorImpl<Value> &delinearizedIds);
+  bool delinearizeLaneId(OpBuilder &builder, Location loc,
+                         ArrayRef<int64_t> originalShape,
+                         ArrayRef<int64_t> distributedShape, int64_t warpSize,
+                         Value laneId,
+                         SmallVectorImpl<Value> &delinearizedIds) const;
 };
 
 } // namespace gpu

mlir/lib/Dialect/GPU/Utils/DistributionUtils.cpp

@@ -23,7 +23,7 @@ using namespace mlir::gpu;
 WarpExecuteOnLane0Op
 WarpDistributionPattern::moveRegionToNewWarpOpAndReplaceReturns(
     RewriterBase &rewriter, WarpExecuteOnLane0Op warpOp,
-    ValueRange newYieldedValues, TypeRange newReturnTypes) {
+    ValueRange newYieldedValues, TypeRange newReturnTypes) const {
   // Create a new op before the existing one, with the extra operands.
   OpBuilder::InsertionGuard g(rewriter);
   rewriter.setInsertionPoint(warpOp);
@@ -51,7 +51,7 @@ WarpExecuteOnLane0Op
 WarpDistributionPattern::moveRegionToNewWarpOpAndAppendReturns(
     RewriterBase &rewriter, WarpExecuteOnLane0Op warpOp,
     ValueRange newYieldedValues, TypeRange newReturnTypes,
-    llvm::SmallVector<size_t> &indices) {
+    llvm::SmallVector<size_t> &indices) const {
   SmallVector<Type> types(warpOp.getResultTypes().begin(),
                           warpOp.getResultTypes().end());
   auto yield = cast<gpu::YieldOp>(
@@ -82,7 +82,8 @@ WarpDistributionPattern::moveRegionToNewWarpOpAndAppendReturns(
 }
 
 OpOperand *WarpDistributionPattern::getWarpResult(
-    WarpExecuteOnLane0Op warpOp, const std::function<bool(Operation *)> &fn) {
+    WarpExecuteOnLane0Op warpOp,
+    const std::function<bool(Operation *)> &fn) const {
   auto yield = cast<gpu::YieldOp>(
       warpOp.getBodyRegion().getBlocks().begin()->getTerminator());
   for (OpOperand &yieldOperand : yield->getOpOperands()) {
@@ -99,7 +100,7 @@ OpOperand *WarpDistributionPattern::getWarpResult(
 bool WarpDistributionPattern::delinearizeLaneId(
     OpBuilder &builder, Location loc, ArrayRef<int64_t> originalShape,
     ArrayRef<int64_t> distributedShape, int64_t warpSize, Value laneId,
-    SmallVectorImpl<Value> &delinearizedIds) {
+    SmallVectorImpl<Value> &delinearizedIds) const {
   // If the original shape and the distributed shape is the same, we don't
   // distribute at all--every thread is handling the whole. For such case, we
   // should not rely on lane IDs later. So just return an empty lane ID vector.

mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp

@@ -536,11 +536,11 @@ struct WarpOpTransferWrite : public WarpDistributionPattern {
   /// Clone `writeOp` assumed to be nested under `warpOp` into a new warp
   /// execute op with the proper return type. The new write op is updated to
   /// write the result of the new warp execute op. The old `writeOp` is deleted.
-  static vector::TransferWriteOp cloneWriteOp(RewriterBase &rewriter,
-                                              WarpExecuteOnLane0Op warpOp,
-                                              vector::TransferWriteOp writeOp,
-                                              VectorType targetType,
-                                              VectorType maybeMaskType) {
+  vector::TransferWriteOp cloneWriteOp(RewriterBase &rewriter,
+                                       WarpExecuteOnLane0Op warpOp,
+                                       vector::TransferWriteOp writeOp,
+                                       VectorType targetType,
+                                       VectorType maybeMaskType) const {
     assert(writeOp->getParentOp() == warpOp &&
            "write must be nested immediately under warp");
     OpBuilder::InsertionGuard g(rewriter);