[NFC] Remove unused forOp argument from setStageCluster (#5288)

<git-pr-chain> [NFC] Remove unused forOp argument from `setStageCluster` #### [PR chain](https://github.com/jlebar/git-pr-chain) 1. 👉 #5288 👈 **YOU ARE HERE** 1. #5289 1. #5290 </git-pr-chain>
triton-lang · Nov 30, 2024 · 912e595 · 912e595
1 parent 27e11ab
commit 912e595
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 13 deletions.
diff --git a/include/triton/Dialect/TritonGPU/Transforms/PipeliningUtility.h b/include/triton/Dialect/TritonGPU/Transforms/PipeliningUtility.h
@@ -35,7 +35,7 @@ void replaceUsesAndPropagateType(OpBuilder &builder, Operation *oldUse,
 // Return the minClusterId and maxClusterId for the given ForOp.
 std::pair<int, int> getMinMaxCluster(scf::ForOp &forOp);
 std::pair<int, int> getStageCluster(Operation *op);
-void setStageCluster(scf::ForOp &forOp, Operation *op, int stage, int cluster);
+void setStageCluster(Operation *op, int stage, int cluster);
 } // namespace triton
 } // namespace mlir
 

diff --git a/lib/Dialect/TritonGPU/Transforms/Pipeliner/MatmulLoopPipeline.cpp b/lib/Dialect/TritonGPU/Transforms/Pipeliner/MatmulLoopPipeline.cpp
@@ -64,11 +64,7 @@ class OpBuilderWithStage : public OpBuilder {
   OpTy createWithStage(Location location, int stage, int cluster,
                        Args &&...args) {
     OpTy op = OpBuilder::create<OpTy>(location, std::forward<Args>(args)...);
-    auto ctx = getContext();
-    op->setAttr(mlir::triton::kLoopStageAttrName,
-                IntegerAttr::get(IntegerType::get(ctx, 32), stage));
-    op->setAttr(mlir::triton::kLoopClusterAttrName,
-                IntegerAttr::get(IntegerType::get(ctx, 32), cluster));
+    tt::setStageCluster(op, stage, cluster);
     return op;
   }
   using OpBuilder::create;
@@ -204,9 +200,8 @@ static int createAsyncCopy(scf::ForOp forOp, tt::LoadOp loadOp, Value alloc,
     // Prefetch load if is not MMAV3 and is used by the dot.
     if (loadToInfo[loadOp].usedByDot) {
       assert(stageForFirstUse >= 1);
-      tt::setStageCluster(forOp, wait, stageForFirstUse - 1, maxClusterId + 1);
-      tt::setStageCluster(forOp, viewLoad, stageForFirstUse - 1,
-                          maxClusterId + 1);
+      tt::setStageCluster(wait, stageForFirstUse - 1, maxClusterId + 1);
+      tt::setStageCluster(viewLoad, stageForFirstUse - 1, maxClusterId + 1);
       retCode = stageForFirstUse - 1;
     }
   }

diff --git a/lib/Dialect/TritonGPU/Transforms/Pipeliner/PipeliningUtility.cpp b/lib/Dialect/TritonGPU/Transforms/Pipeliner/PipeliningUtility.cpp
@@ -188,9 +188,8 @@ std::pair<int, int> mlir::triton::getStageCluster(Operation *op) {
   return std::make_pair(stage, clusterId);
 }
 
-void mlir::triton::setStageCluster(scf::ForOp &forOp, Operation *op, int stage,
-                                   int cluster) {
-  auto ctx = forOp.getContext();
+void mlir::triton::setStageCluster(Operation *op, int stage, int cluster) {
+  auto ctx = op->getContext();
   op->setAttr(mlir::triton::kLoopStageAttrName,
               IntegerAttr::get(IntegerType::get(ctx, 32), stage));
   op->setAttr(mlir::triton::kLoopClusterAttrName,

diff --git a/lib/Dialect/TritonGPU/Transforms/Pipeliner/Schedule.cpp b/lib/Dialect/TritonGPU/Transforms/Pipeliner/Schedule.cpp
@@ -94,7 +94,7 @@ void tt::CoarseSchedule::dump() {
 // Set <stage, cluster> based on CoarseSchedule.
 void tt::CoarseSchedule::serialize(scf::ForOp &forOp) {
   for (auto [op, stage, cluster] : getOpsInOrder(forOp)) {
-    tt::setStageCluster(forOp, op, stage, *cluster);
+    tt::setStageCluster(op, stage, *cluster);
   }
 }