Merge commit 'f05cdc4cec480f31e79f89e2a8a26c4b51614ac2'

Author: whitneywhtsang

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
 

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/Triton/Transforms/Utility.h"
 #include "llvm/Support/Debug.h"
 
 #define GEN_PASS_CLASSES

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

@@ -7,7 +7,7 @@
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Support/LLVM.h"
-#include "triton/Dialect/Triton/Transforms/Utility.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"