Fix flaky build (#4082)

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

Author: anmyachev

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

@@ -0,0 +1,11 @@
+#ifndef TRITON_TRANSFORMS_UTILITY_H
+#define TRITON_TRANSFORMS_UTILITY_H
+
+#include "triton/Dialect/Triton/IR/Dialect.h"
+
+using namespace mlir;
+
+Value getPredMask(RewriterBase &rewriter, Type typeLike, Value currentMask,
+                  Value pred);
+
+#endif // TRITON_TRANSFORMS_UTILITY_H

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/Transforms/CMakeLists.txt

@@ -8,6 +8,7 @@ add_triton_library(TritonTransforms
   LoopUnroll.cpp
   ReorderBroadcast.cpp
   RewriteTensorPointer.cpp
+  Utility.cpp
 
   DEPENDS
   TritonTransformsIncGen

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/Transforms/Passes.h"
-#include "triton/Dialect/TritonGPU/Transforms/PipeliningUtility.h"
+#include "triton/Dialect/Triton/Transforms/Utility.h"
 #include "llvm/Support/Debug.h"
 
 #define GEN_PASS_CLASSES

lib/Dialect/Triton/Transforms/Utility.cpp

@@ -0,0 +1,20 @@
+#include "triton/Dialect/Triton/Transforms/Utility.h"
+#include "triton/Dialect/Triton/IR/Dialect.h"
+
+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) {
+  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/TritonGPU/Transforms/CMakeLists.txt

@@ -44,6 +44,7 @@ add_triton_library(TritonGPUTransforms
   MLIRTransforms
   MLIRTransformUtils
   TritonAnalysis
+  TritonTransforms
   TritonIR
   TritonGPUIR
   TritonNvidiaGPUIR

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/Transforms/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) {