Support running different allocation analyses

include/triton/Analysis/Allocation.h

@@ -17,6 +17,9 @@ namespace mlir {
 
 namespace triton {
 class AllocationAnalysis;
+namespace intel {
+class AllocationAnalysis;
+}
 
 // To convert a tensor from one layout to another, we need to allocate a
 // temporary buffer (i.e., scratch buffer) in shared memory. The conversion may
@@ -238,6 +241,7 @@ class Allocation {
   size_t sharedMemorySize = 0;
 
   friend class triton::AllocationAnalysis;
+  friend class triton::intel::AllocationAnalysis;
 };
 
 template <>

third_party/intel/include/Analysis/Allocation.h

@@ -0,0 +1,12 @@
+#ifndef TRITON_INTEL_ANALYSIS_ALLOCATION_H
+#define TRITON_INTEL_ANALYSIS_ALLOCATION_H
+
+#include "triton/Analysis/Allocation.h"
+
+namespace mlir {
+template <>
+void Allocation::run<triton::intel::AllocationAnalysis>(
+    FuncAllocMapT &funcAllocMap);
+} // namespace mlir
+
+#endif

third_party/intel/lib/Analysis/CMakeLists.txt

@@ -1,4 +1,5 @@
 add_triton_library(TritonIntelAnalysis
+    Allocation.cpp
     AxisInfo.cpp
     DPAS.cpp
     Liveness.cpp

third_party/intel/lib/TritonIntelGPUToLLVM/AllocateSharedMemory.cpp

@@ -22,7 +22,8 @@ struct AllocateSharedMemory
   void runOnOperation() override {
     ModuleOp mod = getOperation();
     MLIRContext *ctx = &getContext();
-    ModuleAllocation allocation = ModuleAllocation::get(mod);
+    ModuleAllocation allocation =
+        ModuleAllocation::get<triton::intel::AllocationAnalysis>(mod);
 
     mod.walk([&](FunctionOpInterface funcOp) {
       if (allocation.isRoot(funcOp) && allocation.getSharedMemorySize()) {

third_party/intel/lib/TritonIntelGPUToLLVM/TritonGPUToLLVM.cpp

@@ -94,7 +94,8 @@ struct ConvertTritonGPUToLLVM
 
     // Allocate shared memory and set barrier
     if (!pipelineManager.skipSharedMemoryAllocation()) {
-      ModuleAllocation allocation = ModuleAllocation::get(mod);
+      ModuleAllocation allocation =
+          ModuleAllocation::get<triton::intel::AllocationAnalysis>(mod);
       ModuleMembarAnalysis membarPass(&allocation);
       membarPass.run();
     }