[SYCL][ESIMD] Refactor lowering of slm allocation into a separate pass (#13221)

Co-authored-by: Nick Sarnie <[email protected]>

Author: fineg74

llvm/include/llvm/SYCLLowerIR/ESIMD/LowerESIMD.h

@@ -66,10 +66,6 @@ class ESIMDOptimizeVecArgCallConvPass
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &);
 };
 
-// Lowers calls __esimd_slm_alloc, __esimd_slm_free and __esimd_slm_init APIs.
-// See more details in the .cpp file.
-size_t lowerSLMReservationCalls(Module &M);
-
 // Lowers calls to __esimd_set_kernel_properties
 class SYCLLowerESIMDKernelPropsPass
     : public PassInfoMixin<SYCLLowerESIMDKernelPropsPass> {
@@ -95,6 +91,14 @@ class ESIMDRemoveOptnoneNoinlinePass
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &);
 };
 
+// Lowers calls __esimd_slm_alloc, __esimd_slm_free and __esimd_slm_init APIs.
+// See more details in the .cpp file.
+class ESIMDLowerSLMReservationCalls
+    : public PassInfoMixin<ESIMDLowerSLMReservationCalls> {
+public:
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &);
+};
+
 } // namespace llvm
 
 #endif // LLVM_SYCLLOWERIR_LOWERESIMD_H

llvm/lib/Passes/PassRegistry.def

@@ -157,6 +157,7 @@ MODULE_PASS("sycl-propagate-joint-matrix-usage", SYCLPropagateJointMatrixUsagePa
 MODULE_PASS("sycl-add-opt-level-attribute", SYCLAddOptLevelAttributePass())
 MODULE_PASS("compile-time-properties", CompileTimePropertiesPass())
 MODULE_PASS("cleanup-sycl-metadata", CleanupSYCLMetadataPass())
+MODULE_PASS("lower-slm-reservation-calls", ESIMDLowerSLMReservationCalls())
 #undef MODULE_PASS
 
 #ifndef MODULE_PASS_WITH_PARAMS

llvm/lib/SYCLLowerIR/ESIMD/LowerESIMD.cpp

@@ -2013,7 +2013,7 @@ PreservedAnalyses SYCLLowerESIMDPass::run(Module &M,
   generateKernelMetadata(M);
   // This function needs to run after generateKernelMetadata, as it
   // uses the generated metadata:
-  size_t AmountOfESIMDIntrCalls = lowerSLMReservationCalls(M);
+  size_t AmountOfESIMDIntrCalls = 0;
   SmallPtrSet<Type *, 4> GVTS = collectGenXVolatileTypes(M);
   lowerGlobalStores(M, GVTS);
   lowerGlobalsToVector(M);
@@ -2082,6 +2082,13 @@ size_t SYCLLowerESIMDPass::runOnFunction(Function &F,
       // process ESIMD builtins that go through special handling instead of
       // the translation procedure
 
+      // SLM allocation API will be lowered in LowerESIMDSlmReservationCalls
+      // pass
+      if (Name.starts_with("__esimd_slm_alloc") ||
+          Name.starts_with("__esimd_slm_free")) {
+        continue;
+      }
+
       if (Name.starts_with("__esimd_svm_block_ld") ||
           Name.starts_with("__esimd_slm_block_ld")) {
         translateBlockLoad(*CI, Name.starts_with("__esimd_slm_block_ld"));

llvm/lib/SYCLLowerIR/ESIMD/LowerESIMDSlmReservation.cpp

@@ -67,6 +67,7 @@
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Module.h"
+#include "llvm/Pass.h"
 
 #include <set>
 #include <unordered_map>
@@ -81,11 +82,19 @@ namespace {
 constexpr int DebugLevel = 0;
 #endif
 
+bool isGenXSLMInit(const Function &F) {
+  constexpr char SLM_INIT_PREFIX[] = "llvm.genx.slm.init";
+  return F.getName().starts_with(SLM_INIT_PREFIX);
+}
+
 bool isSlmInitCall(const CallInst *CI) {
   if (!CI)
     return false;
   Function *F = CI->getCalledFunction();
-  return F && esimd::isSlmInit(*F);
+  if (!F)
+    return false;
+  assert(!esimd::isSlmInit(*F) && "Should have been translated already");
+  return isGenXSLMInit(*F);
 }
 
 bool isSlmAllocCall(const CallInst *CI) {
@@ -543,7 +552,8 @@ TraversalResult findMaxSLMUsageAlongAllPaths(const ScopedCallGraph::Node *Cur,
   return Res;
 }
 
-size_t lowerSLMReservationCalls(Module &M) {
+PreservedAnalyses
+ESIMDLowerSLMReservationCalls::run(Module &M, ModuleAnalysisManager &MAM) {
   // Create a detailed "scoped" call graph. Scope start/end is marked with
   // x = __esimd_slm_alloc / __esimd_slm_free(x)
   //
@@ -567,7 +577,7 @@ size_t lowerSLMReservationCalls(Module &M) {
 #endif
   if (SCG.getNumSLMScopes() == 0) {
     // Early bail out if nothing to analyze.
-    return 0;
+    return PreservedAnalyses::none();
   }
   // Use the detailed call graph nodes to calculate maximum possible SLM usage
   // at any "scope start" node, and record this info in the result map.
@@ -669,7 +679,8 @@ size_t lowerSLMReservationCalls(Module &M) {
 #endif
   }
   // Return the number of API calls transformed.
-  return SLMAllocCallCnt;
+  return SLMAllocCallCnt == 0 ? PreservedAnalyses::none()
+                              : PreservedAnalyses::all();
 }
 
 } // namespace llvm

llvm/test/SYCLLowerIR/ESIMD/slm_init_same_bb.ll

@@ -1,7 +1,7 @@
 ; This test confirms we don't assert for a single slm_init call
 ; in a basic block with two predecessors.
 ;
-; RUN: opt < %s -passes=LowerESIMD -S | FileCheck %s
+; RUN: opt < %s -passes=LowerESIMD,lower-slm-reservation-calls -S | FileCheck %s
 
 target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-n8:16:32:64"
 target triple = "spir64-unknown-unknown"

llvm/tools/sycl-post-link/sycl-post-link.cpp

@@ -669,6 +669,7 @@ bool lowerEsimdConstructs(module_split::ModuleDesc &MD) {
     MainFPM.addPass(InstCombinePass{});
     MainFPM.addPass(DCEPass{});
   }
+  MPM.addPass(ESIMDLowerSLMReservationCalls{});
   MPM.addPass(createModuleToFunctionPassAdaptor(std::move(MainFPM)));
   MPM.addPass(GenXSPIRVWriterAdaptor(/*RewriteTypes=*/true,
                                      /*RewriteSingleElementVectorsIn*/ false));

sycl/test-e2e/ESIMD/slm_alloc.cpp

@@ -1,11 +1,16 @@
 //
+// REQUIRES-INTEL-DRIVER: lin: 28454, win: 101.5333
 // RUN: %{build} -o %t.1.out
 // RUN: %{run} %t.1.out
 //
 // Vary the test case by forcing inlining of the functions with slm_allocator:
 // RUN: %{build} -DFORCE_INLINE -o %t.2.out
 // RUN: %{run} %t.2.out
 
+// Check if the test sill passes with O0
+// RUN: %{build} -O0 -o %t.3.out
+// RUN: %{run} %t.3.out
+
 // This is end-to-end test for the slm_allocator API used together with the
 // slm_init. The call graph is:
 //         Test1(kernel) - uses slm_init(SLM_IN_KERNEL)

sycl/test-e2e/ESIMD/slm_alloc_many_kernels_many_funcs.cpp

@@ -1,6 +1,7 @@
 //
 // Windows doesn't yet have full shutdown().
 // UNSUPPORTED: ze_debug && windows
+// REQUIRES-INTEL-DRIVER: lin: 28454, win: 101.5333
 //
 // RUN: %{build} -o %t.1.out
 // RUN: %{run} %t.1.out
@@ -9,6 +10,10 @@
 // RUN: %{build} -DFORCE_INLINE -o %t.2.out
 // RUN: %{run} %t.2.out
 
+// Check if the test sill passes with O0
+// RUN: %{build} -O0 -o %t.3.out
+// RUN: %{run} %t.3.out
+
 // Checks validity of SLM frame offsets in case of complex call graph with two
 // kernels and 2 functions all using SLM, and one of the functions using two
 // slm_allocator objects with nested liveness ranges.

sycl/test-e2e/ESIMD/slm_alloc_many_kernels_one_func.cpp

@@ -1,6 +1,7 @@
 //
 // Windows doesn't yet have full shutdown().
 // UNSUPPORTED: ze_debug && windows
+// REQUIRES-INTEL-DRIVER: lin: 28454, win: 101.5333
 //
 // RUN: %{build} -o %t.1.out
 // RUN: %{run} %t.1.out
@@ -9,6 +10,10 @@
 // RUN: %{build} -DFORCE_INLINE -o %t.2.out
 // RUN: %{run} %t.2.out
 
+// Check if the test sill passes with O0
+// RUN: %{build} -O0 -o %t.3.out
+// RUN: %{run} %t.3.out
+
 // Check that SLM frame offset of a function foo called from two kernels Test1
 // and Test2 is the maximum of the SLM size used in both kernels.