Fix FlashAttension SLM path failure (#2182)

intel · Sep 10, 2024 · fed5856 · fed5856
1 parent 0a4d503
commit fed5856
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/third_party/intel/lib/TritonIntelGPUToLLVM/AllocateSharedMemory.cpp b/third_party/intel/lib/TritonIntelGPUToLLVM/AllocateSharedMemory.cpp
@@ -49,9 +49,14 @@ struct AllocateSharedMemory
                     IntegerAttr::get(IntegerType::get(ctx, 32), offset));
       });
     });
-    mod->setAttr("triton_gpu.shared",
-                 IntegerAttr::get(IntegerType::get(ctx, 32),
-                                  allocation.getSharedMemorySize()));
+    int32_t originSharedSize = 0;
+    if (IntegerAttr sharedAttr =
+            mod->getAttrOfType<IntegerAttr>("triton_gpu.shared"))
+      originSharedSize = sharedAttr.getInt();
+    mod->setAttr(
+        "triton_gpu.shared",
+        IntegerAttr::get(IntegerType::get(ctx, 32),
+                         originSharedSize + allocation.getSharedMemorySize()));
   }
 };