Move getPredMask func to Triton/IR/Utility.h (#6697)

Signed-off-by: Anatoly Myachev <[email protected]>

Author: anmyachev

include/triton/Dialect/Triton/IR/Utility.h

@@ -173,6 +173,10 @@ template <typename T> auto seq(T start, T end, T step) {
                          [=](T i) { return start + i * step; });
 }
 
+// Combine the current mask with the given predicate.
+Value getPredMask(RewriterBase &rewriter, Type typeLike, Value currentMask,
+                  Value pred);
+
 } // namespace triton
 } // namespace mlir
 

include/triton/Dialect/TritonGPU/Transforms/PipeliningUtility.h

@@ -64,9 +64,6 @@ Value sinkValueRedefinition(RewriterBase &rewriter, Value in, Value out,
 bool loopHasDistGreaterThanOne(scf::ForOp forOp);
 bool isOuterLoop(scf::ForOp forOp);
 
-Value getPredMask(RewriterBase &rewriter, Type typeLike, Value currentMask,
-                  Value pred);
-
 /// Function to mask operations during scheduling.
 Operation *predicateOp(RewriterBase &rewriter, Operation *op, Value pred);
 

lib/Dialect/Triton/IR/CMakeLists.txt

@@ -8,6 +8,7 @@ add_triton_library(TritonIR
   Traits.cpp
   Types.cpp
   OpInterfaces.cpp
+  Utility.cpp
 
   DEPENDS
   TritonTableGen

lib/Dialect/Triton/IR/Utility.cpp

@@ -0,0 +1,19 @@
+#include "triton/Dialect/Triton/IR/Utility.h"
+#include "triton/Dialect/Triton/IR/Dialect.h"
+
+using namespace mlir;
+namespace tt = mlir::triton;
+
+Value tt::getPredMask(RewriterBase &rewriter, Type typeLike, Value currentMask,
+                      Value pred) {
+  Type maskType = tt::getI1SameShape(typeLike);
+  Location loc = pred.getLoc();
+  Value mask = pred;
+  if (isa<RankedTensorType>(maskType)) {
+    mask = rewriter.create<tt::SplatOp>(loc, maskType, pred);
+  }
+  if (currentMask) {
+    mask = rewriter.create<arith::AndIOp>(loc, mask, currentMask);
+  }
+  return mask;
+}

lib/Dialect/Triton/Transforms/LoopInvariantCodeMotion.cpp

@@ -1,7 +1,7 @@
 #include "mlir/Transforms/LoopInvariantCodeMotionUtils.h"
 #include "triton/Dialect/Triton/IR/Dialect.h"
+#include "triton/Dialect/Triton/IR/Utility.h"
 #include "triton/Dialect/Triton/Transforms/Passes.h"
-#include "triton/Dialect/TritonGPU/Transforms/PipeliningUtility.h"
 #include "llvm/Support/Debug.h"
 
 #define GEN_PASS_CLASSES

lib/Dialect/TritonGPU/Transforms/Pipeliner/PipeliningUtility.cpp

@@ -7,6 +7,7 @@
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Support/LLVM.h"
+#include "triton/Dialect/Triton/IR/Utility.h"
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
 #include "triton/Dialect/TritonGPU/Transforms/Passes.h"
 #include "triton/Dialect/TritonGPU/Transforms/Utility.h"
@@ -154,21 +155,6 @@ bool mlir::triton::isOuterLoop(scf::ForOp forOp) {
   });
 }
 
-// Combine the current mask with the given predicate.
-Value mlir::triton::getPredMask(RewriterBase &rewriter, Type typeLike,
-                                Value currentMask, Value pred) {
-  Type maskType = tt::getI1SameShape(typeLike);
-  Location loc = pred.getLoc();
-  Value mask = pred;
-  if (isa<RankedTensorType>(maskType)) {
-    mask = rewriter.create<tt::SplatOp>(loc, maskType, pred);
-  }
-  if (currentMask) {
-    mask = rewriter.create<arith::AndIOp>(loc, mask, currentMask);
-  }
-  return mask;
-}
-
 // Function to mask operations during scheduling.
 Operation *mlir::triton::predicateOp(RewriterBase &rewriter, Operation *op,
                                      Value pred) {