address review comments

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

Author: whitneywhtsang

test/TritonIntelGPU/loop-pipeline.mlir

@@ -81,7 +81,7 @@ module attributes {"ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 16 : i32,
     // COM: There are 3 stages in loop pipelining, the first 2 prefetching stages are before the loop and the last one is inside the loop.
     // CHECK: %[[LOAD_MASK:.*]] = arith.cmpi slt, {{.*}}
     // CHECK: %[[LOOP_MASK:.*]] = tt.splat %[[LOAD_MASK]] : i1 -> tensor<64x32xi1, #[[$BLOCK_0]]>
-    // CHECK: triton_intel_gpu.prefetch {{.*}} : tensor<64x32x!tt.ptr<f16>, #[[$BLOCK_0]]>
+    // CHECK: triton_intel_gpu.prefetch {{.*}}, %[[LOOP_MASK]] {{.*}} : tensor<64x32x!tt.ptr<f16>, #[[$BLOCK_0]]>
     // CHECK: triton_intel_gpu.prefetch {{.*}} : tensor<32x256x!tt.ptr<f16>, #[[$BLOCK_1]]>
     // CHECK: triton_intel_gpu.prefetch {{.*}} : tensor<64x32x!tt.ptr<f16>, #[[$BLOCK_0]]>
     // CHECK: triton_intel_gpu.prefetch {{.*}} : tensor<32x256x!tt.ptr<f16>, #[[$BLOCK_1]]>

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

@@ -150,9 +150,10 @@ static void collectOpsToPipeline(scf::ForOp forOp,
   }
 }
 
-/// Combine the current mask with the given predicate.
-static Value getPredMask(RewriterBase &rewriter, Type typeLike,
-                         Value currentMask, Value pred) {
+/// 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));
@@ -175,7 +176,7 @@ static Operation *predicateOp(RewriterBase &rewriter, Operation *op,
       .Case<tt::LoadOp, ttgi::PrefetchOp>([&](auto op) {
         rewriter.setInsertionPoint(op);
         Value mask =
-            getPredMask(rewriter, op.getPtr().getType(), op.getMask(), pred);
+            computeNewMask(rewriter, op.getPtr().getType(), op.getMask(), pred);
         op.getMaskMutable().assign(mask);
         return op;
       });