[NFC][MatmulLoopPipeline] Use existing utility (#4084)

Signed-off-by: Whitney Tsang <[email protected]>

Author: whitneywhtsang

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

@@ -5,7 +5,11 @@
 
 using namespace mlir;
 
+namespace mlir::triton {
+
 Value getPredMask(RewriterBase &rewriter, Type typeLike, Value currentMask,
                   Value pred);
 
+} // namespace mlir::triton
+
 #endif // TRITON_TRANSFORMS_UTILITY_H

lib/Dialect/Triton/Transforms/Utility.cpp

@@ -5,8 +5,8 @@ using namespace mlir;
 namespace tt = mlir::triton;
 
 // Combine the current mask with the given predicate.
-Value getPredMask(RewriterBase &rewriter, Type typeLike, Value currentMask,
-                  Value pred) {
+Value tt::getPredMask(RewriterBase &rewriter, Type typeLike, Value currentMask,
+                      Value pred) {
   Type maskType = tt::getI1SameShape(typeLike);
   Location loc = pred.getLoc();
   Value mask = pred;

third_party/intel/lib/TritonIntelGPUTransforms/Pipeliner/MatmulLoopPipeline.cpp

@@ -6,6 +6,7 @@
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "triton/Analysis/AxisInfo.h"
 #include "triton/Dialect/Triton/IR/Dialect.h"
+#include "triton/Dialect/Triton/Transforms/Utility.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Debug.h"
@@ -150,21 +151,6 @@ static void collectOpsToPipeline(scf::ForOp forOp,
   }
 }
 
-/// Return a new mask of type of shape \p typeLike, and value combining the
-/// current mask \p currentMask with the given predicate \p pred.
-static Value computeNewMask(RewriterBase &rewriter, Type typeLike,
-                            Value currentMask, Value pred) {
-  Location loc = pred.getLoc();
-  Value mask = pred;
-  Type maskType = tt::getI1SameShape(tt::getPointeeType(typeLike));
-
-  if (isa<RankedTensorType>(maskType))
-    mask = rewriter.create<tt::SplatOp>(loc, maskType, pred);
-
-  return currentMask ? rewriter.create<arith::AndIOp>(loc, mask, currentMask)
-                     : mask;
-}
-
 /// Function to mask operations during scheduling.
 static Operation *predicateOp(RewriterBase &rewriter, Operation *op,
                               Value pred) {
@@ -176,7 +162,8 @@ static Operation *predicateOp(RewriterBase &rewriter, Operation *op,
       .Case<tt::LoadOp, ttgi::PrefetchOp>([&](auto op) {
         rewriter.setInsertionPoint(op);
         Value mask =
-            computeNewMask(rewriter, op.getPtr().getType(), op.getMask(), pred);
+            tt::getPredMask(rewriter, tt::getPointeeType(op.getPtr().getType()),
+                            op.getMask(), pred);
         op.getMaskMutable().assign(mask);
         return op;
       });