[ROCm] Enable kernarg preloading

ROCm · Jan 8, 2025 · 9174f23 · 9174f23
1 parent 92152a2
commit 9174f23
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/third_party/gpus/crosstool/clang/bin/crosstool_wrapper_driver_rocm.tpl b/third_party/gpus/crosstool/clang/bin/crosstool_wrapper_driver_rocm.tpl
@@ -188,6 +188,7 @@ def InvokeHipcc(argv, log=False):
   hipccopts += defines
   hipccopts += std_options
   hipccopts += m_options
+  hipccopts += ' -mllvm=-amdgpu-kernarg-preload-count=16 '
 
   if depfiles:
     # Generate the dependency file

diff --git a/third_party/xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc b/third_party/xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc
@@ -1048,6 +1048,10 @@ std::string GetROCDLDir(const DebugOptions& debug_options) {
 
 void AMDGPUBackendInit(const DebugOptions& debug_options,
                        std::string& rocdl_dir_path) {
+  FeedLLVMWithFlags({
+      "-amdgpu-kernarg-preload-count=16",
+  });
+
   llvm_ir::InitializeLLVMCommandLineOptions(
       debug_options.xla_backend_extra_options());