[Intel] Refactor numWarps lookups to be from the op

whitneywhtsang · whitneywhtsang · commit 6b1642e117a8 · 2025-02-15T00:04:50.000Z
Signed-off-by: Whitney Tsang &lt;whitney.tsang@intel.com&gt;
diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
@@ -1062,8 +1062,7 @@ class DecomposeScaledBlocked
 
     RankedTensorType oldRetType = dotOp.getType();
     auto retShapePerCTA = getShapePerCTA(oldRetType);
-    auto mod = dotOp->getParentOfType<mlir::ModuleOp>();
-    int numWarps = TritonGPUDialect::getNumWarps(mod);
+    int numWarps = lookupNumWarps(dotOp);
     auto CTALayout = getCTALayout(oldRetType.getEncoding());
 
     auto instrShape = mmaVersionToInstrShape(
diff --git a/third_party/amd/lib/TritonAMDGPUTransforms/AccelerateAMDMatmul.cpp b/third_party/amd/lib/TritonAMDGPUTransforms/AccelerateAMDMatmul.cpp
@@ -562,7 +562,7 @@ class ScaledBlockedToMFMA final : public OpRewritePattern<triton::DotScaledOp> {
     auto moduleOp = dotOp->getParentOfType<ModuleOp>();
 
     ttg::CTALayoutAttr ctaLayout = ttg::getCTALayout(oldRetType.getEncoding());
-    int numWarps = ttg::TritonGPUDialect::getNumWarps(moduleOp);
+    int numWarps = ttg::lookupNumWarps(dotOp);
     int numThreads = ttg::TritonGPUDialect::getThreadsPerWarp(moduleOp);
 
     // Choose a suitable MFMA instruction for this scaled dot op.
diff --git a/third_party/intel/lib/Analysis/DPAS.cpp b/third_party/intel/lib/Analysis/DPAS.cpp
@@ -61,8 +61,7 @@ DPASAnalysis::canUseDPAS(FunctionOpInterface funcOp) const {
   // Verify whether the module has the correct number of threads per warp.
   // Note: if the module doesn't then return 'Result::Maybe' to allow the caller
   // to set warp size.
-  Attribute threadsPerWarpAttr =
-      mod->getDiscardableAttr(TritonGPUDialect::getThreadsPerWarpAttrName());
+  Attribute threadsPerWarpAttr = mod->getDiscardableAttr(AttrNumThreadsPerWarp);
   if (!threadsPerWarpAttr)
     return Result::Maybe;
 
diff --git a/third_party/intel/lib/TritonAnnotateModule/TritonAnnotateModule.cpp b/third_party/intel/lib/TritonAnnotateModule/TritonAnnotateModule.cpp
@@ -48,9 +48,6 @@ struct TritonAnnotateModule
   void setThreadsPerWarp(ModuleOp &mod,
                          const DPASAnalysis &dpasAnalysis) const {
     Builder builder(mod);
-    const std::string &AttrNumThreadsPerWarp =
-        TritonGPUDialect::getThreadsPerWarpAttrName();
-
     mod.walk([&](FunctionOpInterface funcOp) {
       // FIXME: DPAS lowering only implemented for 16 threads per warp, i.e.,
       // DPAS is not used for devices like ATS.
diff --git a/third_party/intel/lib/TritonIntelGPUToLLVM/HistogramOpToLLVM.cpp b/third_party/intel/lib/TritonIntelGPUToLLVM/HistogramOpToLLVM.cpp
@@ -170,7 +170,7 @@ struct HistogramOpConversion
     assert((numThreadsPerWarp == 16 || numThreadsPerWarp == 32 ||
             numThreadsPerWarp == 64) &&
            "Only supports 16, 32 or 64 threads per warp");
-    int numWarps = triton::gpu::TritonGPUDialect::getNumWarps(mod);
+    int numWarps = triton::gpu::lookupNumWarps(op);
     // Pad out the bins so that we have at least one bin per thread within a
     // warp.
     numBins = std::max(numBins, numThreadsPerWarp);
diff --git a/third_party/intel/lib/TritonIntelGPUToLLVM/LoadStoreOpToLLVM.cpp b/third_party/intel/lib/TritonIntelGPUToLLVM/LoadStoreOpToLLVM.cpp
@@ -366,7 +366,7 @@ struct PrefetchOpConversion
       std::swap(tensorShape[0], tensorShape[1]);
     }
 
-    unsigned numWarps = triton::gpu::TritonGPUDialect::getNumWarps(mod);
+    unsigned numWarps = triton::gpu::lookupNumWarps(op);
 
     SmallVector<unsigned, 2> shapePerWarp =
         get2DPrefetchShapePerWarp(tensorType);
diff --git a/third_party/intel/lib/TritonIntelGPUToLLVM/ReduceOpToLLVM.cpp b/third_party/intel/lib/TritonIntelGPUToLLVM/ReduceOpToLLVM.cpp
@@ -289,7 +289,7 @@ struct ReduceOpConversion
 
     auto mod = op.getOperation()->getParentOfType<ModuleOp>();
     unsigned numLanes = triton::gpu::TritonGPUDialect::getThreadsPerWarp(mod);
-    int numWarps = triton::gpu::TritonGPUDialect::getNumWarps(mod);
+    int numWarps = triton::gpu::lookupNumWarps(op.getOperation());
     int numThreads = numLanes * numWarps;
 
     Value threadId = getThreadId(rewriter, loc);
diff --git a/third_party/intel/lib/TritonIntelGPUToLLVM/TritonGPUToLLVM.cpp b/third_party/intel/lib/TritonIntelGPUToLLVM/TritonGPUToLLVM.cpp
@@ -101,7 +101,7 @@ struct ConvertTritonGPUToLLVM
     TritonIntelGPUToLLVMTypeConverter typeConverter(context, option, targetInfo,
                                                     isAdvancedPathEnabled);
     TritonLLVMConversionTarget convTarget(*context);
-    int numWarps = triton::gpu::TritonGPUDialect::getNumWarps(mod);
+    int numWarps = triton::gpu::lookupNumWarps(&*mod.getOps().begin());
     int numCTAs = triton::gpu::TritonGPUDialect::getNumCTAs(mod);
     int threadsPerWarp = triton::gpu::TritonGPUDialect::getThreadsPerWarp(mod);
 
diff --git a/third_party/intel/lib/TritonIntelGPUTransforms/AccelerateMatmul.cpp b/third_party/intel/lib/TritonIntelGPUTransforms/AccelerateMatmul.cpp
@@ -105,14 +105,14 @@ class BlockedToDPAS : public OpRewritePattern<tt::DotOp> {
 
     // Create DPAS encoding for the given number of warps
     ArrayRef<int64_t> retShape = oldRetType.getShape();
-    ModuleOp mod = funcOp->getParentOfType<ModuleOp>();
-    unsigned numWarps = ttg::TritonGPUDialect::getNumWarps(mod);
+    unsigned numWarps = ttg::lookupNumWarps(funcOp);
 
     TensorValue a = dotOp.getA();
     TensorValue b = dotOp.getB();
     auto oldAType = cast<RankedTensorType>(a.getType());
     auto oldBType = cast<RankedTensorType>(b.getType());
 
+    ModuleOp mod = funcOp->getParentOfType<ModuleOp>();
     auto dpasCap = ttgi::DpasEncodingAttr::getDPASCapability(mod);
     Type elemType = oldAType.getElementType();
     unsigned opsPerChan = ttgi::DpasEncodingAttr::getOpsPerChannel(elemType);
@@ -295,7 +295,7 @@ class DecomposeScaledBlocked : public OpRewritePattern<tt::DotScaledOp> {
     assert(opDesc.scale && "Expecting valid operand & scale");
 
     MLIRContext *ctx = opDesc.op.getContext();
-    unsigned numWarps = ttg::TritonGPUDialect::getNumWarps(mod);
+    unsigned numWarps = ttg::lookupNumWarps(&*rewriter.getInsertionPoint());
     unsigned warpSize = ttg::TritonGPUDialect::getThreadsPerWarp(mod);
     unsigned opsPerChannel = dpasEnc.getOpsPerChannel();
     unsigned rank = retType.getRank();
@@ -372,7 +372,7 @@ class DecomposeScaledBlocked : public OpRewritePattern<tt::DotScaledOp> {
         aScale ? b.getType().getElementType() : a.getType().getElementType();
     unsigned opsPerChan =
         ttg::intel::DpasEncodingAttr::getOpsPerChannel(elemType);
-    unsigned numWarps = ttg::TritonGPUDialect::getNumWarps(mod);
+    unsigned numWarps = ttg::lookupNumWarps(scaledDotOp);
     SmallVector<unsigned> warpsPerTile = {numWarps, 1};
 
     ArrayRef<int64_t> retShape = scaledDotOp.getType().getShape();
diff --git a/third_party/intel/lib/TritonIntelGPUTransforms/Coalesce.cpp b/third_party/intel/lib/TritonIntelGPUTransforms/Coalesce.cpp
@@ -372,7 +372,7 @@ struct CoalescePass
       if (!refTensorType || !refTensorType.getEncoding())
         return;
 
-      int numWarps = ttg::TritonGPUDialect::getNumWarps(moduleOp);
+      int numWarps = ttg::lookupNumWarps(curr);
       int threadsPerWarp = ttg::TritonGPUDialect::getThreadsPerWarp(moduleOp);
       setCoalescedEncoding(axisInfoAnalysis, curr, numWarps, threadsPerWarp,
                            layoutMap);
diff --git a/third_party/intel/lib/TritonIntelGPUTransforms/MatchTargetSize.cpp b/third_party/intel/lib/TritonIntelGPUTransforms/MatchTargetSize.cpp
@@ -220,7 +220,7 @@ static void rewriteLoadWithSLM(ModuleOp &m, DenseSet<Value> &dotWithSLMOperands,
   OpBuilder b(load);
   auto type = cast<RankedTensorType>(load.getType());
   unsigned bytes = type.getNumElements() * type.getElementTypeBitWidth() / 8;
-  unsigned numWarps = ttg::TritonGPUDialect::getNumWarps(m);
+  unsigned numWarps = ttg::lookupNumWarps(load);
   unsigned slmSize = numWarps * bytes;
 
   // TODO: use LocalAllocOp for SLM allocation
@@ -849,16 +849,13 @@ static SmallVector<Value> glueForReduction(OpBuilder &builder, Location loc,
 
 static Value allocateSLMForTransposedReduction(tt::ReduceOp op, unsigned step,
                                                OpBuilder &b) {
-  auto m = op->getParentOfType<ModuleOp>();
-
   Value src = op.getSrcs().front();
   auto srcTy = cast<RankedTensorType>(src.getType());
   Location loc = op.getLoc();
 
   // Fixed size for num_warps matrices of sg_size^2 shape.
   int64_t size = static_cast<int64_t>(step) * step *
-                 srcTy.getElementTypeBitWidth() / 8 *
-                 ttg::TritonGPUDialect::getNumWarps(m);
+                 srcTy.getElementTypeBitWidth() / 8 * ttg::lookupNumWarps(op);
   Type allocTy = cast<RankedTensorType>(src.getType()).getElementType();
   Type ptrTy = tt::PointerType::get(allocTy, tt::TritonGEN::kWorkgroup);
   return hackAlloc(b, loc, ptrTy, size);
diff --git a/third_party/intel/lib/TritonIntelGPUTransforms/PrefetchBlock.cpp b/third_party/intel/lib/TritonIntelGPUTransforms/PrefetchBlock.cpp
@@ -339,12 +339,10 @@ void PrefetchBlockPass::injectPrefetchOpsInPreheader(
     scf::ForOp loop, SmallVectorImpl<Value> &prefetchPtrs) const {
   assert(prefetchPtrs.empty() && "Expecting an empty vector");
 
-  ModuleOp mod = loop->getParentOfType<ModuleOp>();
   OpBuilder b(loop);
-
   for (tt::LoadOp load : loopLoads.at(loop)) {
     const LoadInfo &loadInfo = loadToLoadInfo.at(load);
-    const unsigned numWarps = ttg::TritonGPUDialect::getNumWarps(mod);
+    const unsigned numWarps = ttg::lookupNumWarps(loop);
 
     b.setInsertionPoint(loadInfo.getBlockPtr());
     auto ptr = cast<tt::MakeTensorPtrOp>(
diff --git a/third_party/intel/lib/TritonIntelGPUTransforms/Utility.cpp b/third_party/intel/lib/TritonIntelGPUTransforms/Utility.cpp
@@ -102,8 +102,8 @@ bool isExpensiveLoadOrStore(Operation *op) {
   // Loads that use more threads than elements can be presumed to have a high
   // hit-rate that makes them cheap to load.
   if (auto ptrType = getRankedTensorType(base.getType())) {
+    int numWarps = ttg::lookupNumWarps(op);
     auto mod = op->getParentOfType<ModuleOp>();
-    int numWarps = ttg::TritonGPUDialect::getNumWarps(mod);
     int threadsPerWarp = ttg::TritonGPUDialect::getThreadsPerWarp(mod);
     return ptrType.getNumElements() >= numWarps * threadsPerWarp;
   }
diff --git a/third_party/intel/lib/TritonToTritonGPUWarp/TritonToTritonGPUWarpPass.cpp b/third_party/intel/lib/TritonToTritonGPUWarp/TritonToTritonGPUWarpPass.cpp
@@ -438,11 +438,11 @@ class ConvertTritonToTritonGPUWarp
     }
 
     /// adding module attributes
-    mod->setAttr(tt::AttrNumWarpsName,
+    mod->setAttr(ttg::AttrNumWarpsName,
                  IntegerAttr::get(i32Ty, llvm::APInt(32, numWarps.getValue())));
-    mod->setAttr(tt::AttrNumThreadsPerWarp,
+    mod->setAttr(ttg::AttrNumThreadsPerWarp,
                  IntegerAttr::get(i32Ty, llvm::APInt(32, 1)));
-    mod->setAttr(tt::AttrNumCTAsName,
+    mod->setAttr(ttg::AttrNumCTAsName,
                  IntegerAttr::get(i32Ty, llvm::APInt(32, 1)));
   }
 

Original file line number	Diff line number	Diff line change
`@@ -366,7 +366,7 @@ struct PrefetchOpConversion`
`366`	`366`	`std::swap(tensorShape[0], tensorShape[1]);`
`367`	`367`	`}`
`368`	`368`
`369`		`- unsigned numWarps = triton::gpu::TritonGPUDialect::getNumWarps(mod);`
	`369`	`+ unsigned numWarps = triton::gpu::lookupNumWarps(op);`
`370`	`370`
`371`	`371`	`SmallVector<unsigned, 2> shapePerWarp =`
`372`	`372`	`get2DPrefetchShapePerWarp(tensorType);`