Fixes #1607 (#1609)

Signed-off-by: Tiotto, Ettore <ettore.tiotto@intel.com>
intel · Jul 11, 2024 · b39b566 · b39b566
1 parent 8dfb4e3
commit b39b566
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/third_party/intel/lib/TritonAnnotateModule/TritonAnnotateModule.cpp b/third_party/intel/lib/TritonAnnotateModule/TritonAnnotateModule.cpp
@@ -44,13 +44,13 @@ struct TritonAnnotateModule
     Builder builder(mod);
     const std::string &AttrNumThreadsPerWarp =
         TritonGPUDialect::getThreadsPerWarpAttrName();
-    unsigned reqThreadsPerWarp =
-        DPASAnalysis::supportedThreadsPerWarp(intel::getDeviceArch(mod));
 
     auto result = mod.walk([&](FunctionOpInterface funcOp) {
       if (dpasAnalysis.canUseDPAS(funcOp) == DPASAnalysis::Result::Maybe) {
         // Set the threads per warp attribute to allow dot operation to be
         // lowered to DPAS instructions.
+        unsigned reqThreadsPerWarp =
+            DPASAnalysis::supportedThreadsPerWarp(intel::getDeviceArch(mod));
         mod->setAttr(AttrNumThreadsPerWarp,
                      builder.getI32IntegerAttr(reqThreadsPerWarp));
         assert(dpasAnalysis.canUseDPAS(funcOp) == DPASAnalysis::Result::True &&