address more comments

Author: kurapov-peter

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

@@ -1,4 +1,4 @@
-//===- VectorDistributionUtils.h - Distribution Utilities -------*- C++ -*-===//
+//===- DistributionUtils.h - Distribution Utilities -------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -15,13 +15,10 @@
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/Value.h"
 
-#include <numeric>
-#include <utility>
-
-namespace mlir {
-namespace gpu {
+namespace mlir::gpu {
 struct WarpDistributionPattern : OpRewritePattern<WarpExecuteOnLane0Op> {
-  using OpRewritePattern<WarpExecuteOnLane0Op>::OpRewritePattern;
+  using OpRewritePattern::OpRewritePattern;
+  using Base = WarpDistributionPattern;
 
   virtual LogicalResult
   matchAndRewrite(WarpExecuteOnLane0Op op,
@@ -30,8 +27,9 @@ struct WarpDistributionPattern : OpRewritePattern<WarpExecuteOnLane0Op> {
 protected:
   /// Return a value yielded by `warpOp` which statifies the filter lamdba
   /// condition and is not dead.
-  OpOperand *getWarpResult(WarpExecuteOnLane0Op warpOp,
-                           const std::function<bool(Operation *)> &fn) const;
+  OpOperand *
+  getWarpResult(WarpExecuteOnLane0Op warpOp,
+                const llvm::function_ref<bool(Operation *)> fn) const;
 
   /// Helper to create a new WarpExecuteOnLane0Op with different signature.
   WarpExecuteOnLane0Op moveRegionToNewWarpOpAndReplaceReturns(
@@ -43,7 +41,7 @@ struct WarpDistributionPattern : OpRewritePattern<WarpExecuteOnLane0Op> {
   WarpExecuteOnLane0Op moveRegionToNewWarpOpAndAppendReturns(
       RewriterBase &rewriter, WarpExecuteOnLane0Op warpOp,
       ValueRange newYieldedValues, TypeRange newReturnTypes,
-      llvm::SmallVector<size_t> &indices) const;
+      SmallVector<size_t> &indices) const;
 
   /// Delinearize the given `laneId` into multiple dimensions, where each
   /// dimension's size is determined by `originalShape` and `distributedShape`
@@ -57,7 +55,6 @@ struct WarpDistributionPattern : OpRewritePattern<WarpExecuteOnLane0Op> {
                          SmallVectorImpl<Value> &delinearizedIds) const;
 };
 
-} // namespace gpu
-} // namespace mlir
+} // namespace mlir::gpu
 
 #endif // MLIR_DIALECT_GPU_TRANSFORMS_DISTRIBUTIONUTILS_H_

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

@@ -51,22 +51,22 @@ WarpExecuteOnLane0Op
 WarpDistributionPattern::moveRegionToNewWarpOpAndAppendReturns(
     RewriterBase &rewriter, WarpExecuteOnLane0Op warpOp,
     ValueRange newYieldedValues, TypeRange newReturnTypes,
-    llvm::SmallVector<size_t> &indices) const {
+    SmallVector<size_t> &indices) const {
   SmallVector<Type> types(warpOp.getResultTypes().begin(),
                           warpOp.getResultTypes().end());
   auto yield = cast<gpu::YieldOp>(
       warpOp.getBodyRegion().getBlocks().begin()->getTerminator());
   llvm::SmallSetVector<Value, 32> yieldValues(yield.getOperands().begin(),
                                               yield.getOperands().end());
-  for (auto newRet : llvm::zip(newYieldedValues, newReturnTypes)) {
-    if (yieldValues.insert(std::get<0>(newRet))) {
-      types.push_back(std::get<1>(newRet));
+  for (auto [value, type] : llvm::zip_equal(newYieldedValues, newReturnTypes)) {
+    if (yieldValues.insert(value)) {
+      types.push_back(type);
       indices.push_back(yieldValues.size() - 1);
     } else {
       // If the value already exit the region don't create a new output.
       for (auto [idx, yieldOperand] :
            llvm::enumerate(yieldValues.getArrayRef())) {
-        if (yieldOperand == std::get<0>(newRet)) {
+        if (yieldOperand == value) {
           indices.push_back(idx);
           break;
         }
@@ -83,7 +83,7 @@ WarpDistributionPattern::moveRegionToNewWarpOpAndAppendReturns(
 
 OpOperand *WarpDistributionPattern::getWarpResult(
     WarpExecuteOnLane0Op warpOp,
-    const std::function<bool(Operation *)> &fn) const {
+    const llvm::function_ref<bool(Operation *)> fn) const {
   auto yield = cast<gpu::YieldOp>(
       warpOp.getBodyRegion().getBlocks().begin()->getTerminator());
   for (OpOperand &yieldOperand : yield->getOpOperands()) {
@@ -94,7 +94,7 @@ OpOperand *WarpDistributionPattern::getWarpResult(
         return &yieldOperand;
     }
   }
-  return {};
+  return nullptr;
 }
 
 bool WarpDistributionPattern::delinearizeLaneId(