[Mlir-commits] [mlir] [MLIR] Create GPU utils library & move distribution utils (PR #119264)

[Kunwar Grover]

================
@@ -0,0 +1,57 @@
+//===- VectorDistributionUtils.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.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_GPU_TRANSFORMS_DISTRIBUTIONUTILS_H_
+#define MLIR_DIALECT_GPU_TRANSFORMS_DISTRIBITIONUTILS_H_
+
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/IR/PatternMatch.h"
+
+#include <utility>
+
+namespace mlir {
+namespace gpu {
+/// 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);
+
+/// Helper to create a new WarpExecuteOnLane0Op with different signature.
+WarpExecuteOnLane0Op moveRegionToNewWarpOpAndReplaceReturns(
+    RewriterBase &rewriter, WarpExecuteOnLane0Op warpOp,
+    ValueRange newYieldedValues, TypeRange newReturnTypes);
+
+/// Helper to create a new WarpExecuteOnLane0Op region with extra outputs.
+/// `indices` return the index of each new output.
+WarpExecuteOnLane0Op moveRegionToNewWarpOpAndAppendReturns(
+    RewriterBase &rewriter, WarpExecuteOnLane0Op warpOp,
+    ValueRange newYieldedValues, TypeRange newReturnTypes,
+    llvm::SmallVector<size_t> &indices);
+
+/// Helper to know if an op can be hoisted out of the region.
+bool canBeHoisted(Operation *op, function_ref<bool(Value)> definedOutside);
----------------
Groverkss wrote:

Note that this moves static functions from a transformation to public functions in the GPU namespace. I'm not a big fan of this approach; a function like "canBeHoisted" is weird without context. These functions were fine when they were constrained to a single transformation like VectorDistribution, but now that they span dialects, it should to be done better.

A solution we used downstream for something similar was to create an inherited pattern class specific to distribution patterns: https://github.com/iree-org/iree/blob/main/compiler/src/iree/compiler/Codegen/Common/GPU/GPUVectorDistribution.h#L28

It would work really well here as well. You could have:

```
template <typename T>
struct WarpDistributionPattern : OpRewritePattern<gpu::WarpExecuteOnLane0Op> {
   virtual LogicalResult matchAndRewrite(T op, PatternRewriter &rewriter) const;
}
```

then you don't need to duplicate the `getWarpResult` logic everywhere. You can expose all of these distribution specific functions as part of the pattern as well.

https://github.com/llvm/llvm-project/pull/119264