[compiler-rt] [PGO][HIP] Support hipModuleLoad in offload PGO (PR #211875)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-amdgpu

Author: Yaxun (Sam) Liu (yxsamliu)

<details>
<summary>Changes</summary>

Offload PGO finds profile sections by inspecting the image passed to
hipModuleLoadData. hipModuleLoad only provides a file name, so its device
profile counters were not collected.

After a successful hipModuleLoad, read the code object with the existing
profile file-buffer helper and register it through the same path used by
in-memory module loads.


---
Full diff: https://github.com/llvm/llvm-project/pull/211875.diff


2 Files Affected:

- (modified) compiler-rt/lib/profile/InstrProfilingPlatformROCm.cpp (+37-3) 
- (added) compiler-rt/test/profile/AMDGPU/device-module-load-file.hip (+63) 


``````````diff
diff --git a/compiler-rt/lib/profile/InstrProfilingPlatformROCm.cpp b/compiler-rt/lib/profile/InstrProfilingPlatformROCm.cpp
index 93d3cbce38de8..cae2e118e9f1a 100644
--- a/compiler-rt/lib/profile/InstrProfilingPlatformROCm.cpp
+++ b/compiler-rt/lib/profile/InstrProfilingPlatformROCm.cpp
@@ -9,6 +9,7 @@
 extern "C" {
 #include "InstrProfiling.h"
 #include "InstrProfilingPort.h"
+#include "InstrProfilingUtil.h"
 }
 
 #include "interception/interception.h"
@@ -1411,11 +1412,44 @@ INTERCEPTOR(int, hipGraphLaunch_spt, HipGraphExec GraphExec, HipStream Stream) {
                                Stream);
 }
 
+static void registerDynamicModuleFromFile(int Rc, void **Module,
+                                          const char *Filename) {
+  if (Rc)
+    return;
+
+  FILE *File = fopen(Filename, "rb");
+  if (!File) {
+    PROF_WARN("failed to open dynamic module file %s\n", Filename);
+    return;
+  }
+
+  if (fseek(File, 0, SEEK_END) != 0) {
+    PROF_WARN("failed to seek dynamic module file %s\n", Filename);
+    fclose(File);
+    return;
+  }
+  long FileSize = ftell(File);
+  if (FileSize <= 0 || fseek(File, 0, SEEK_SET) != 0) {
+    PROF_WARN("failed to get size of dynamic module file %s\n", Filename);
+    fclose(File);
+    return;
+  }
+
+  ManagedMemory Image;
+  lprofGetFileContentBuffer(File, (uint64_t)FileSize, &Image);
+  fclose(File);
+  if (Image.Status == MS_INVALID) {
+    PROF_WARN("failed to read dynamic module file %s\n", Filename);
+    return;
+  }
+
+  __llvm_profile_offload_register_dynamic_module(Rc, Module, Image.Addr);
+  lprofReleaseBuffer(&Image, (size_t)FileSize);
+}
+
 INTERCEPTOR(int, hipModuleLoad, void **module, const char *fname) {
   int rc = REAL(hipModuleLoad)(module, fname);
-  /* Pass NULL image: no in-memory ELF is available for filename loads,
-   * so the register hook skips symbol enumeration. */
-  __llvm_profile_offload_register_dynamic_module(rc, module, nullptr);
+  registerDynamicModuleFromFile(rc, module, fname);
   return rc;
 }
 
diff --git a/compiler-rt/test/profile/AMDGPU/device-module-load-file.hip b/compiler-rt/test/profile/AMDGPU/device-module-load-file.hip
new file mode 100644
index 0000000000000..0ac7751446999
--- /dev/null
+++ b/compiler-rt/test/profile/AMDGPU/device-module-load-file.hip
@@ -0,0 +1,63 @@
+// Device PGO supports code objects loaded by filename through hipModuleLoad.
+//
+// REQUIRES: hip, amdgpu
+// UNSUPPORTED: windows
+
+// RUN: rm -rf %t.dir && mkdir -p %t.dir
+// RUN: %clang -x hip --offload-arch=%amdgpu_arch \
+// RUN:   --rocm-path=%hip_lib_path/.. \
+// RUN:   -fprofile-instr-generate -DBUILD_MODULE %s \
+// RUN:   -o %t.dir/module -L%hip_lib_path -lamdhip64
+// RUN: cd %t.dir && llvm-objdump --offloading module
+// RUN: mv %t.dir/module.0.hip-amdgcn-amd-amdhsa--*gfx* %t.dir/module.co
+// RUN: %clang -x c++ -c %s -D__HIP_PLATFORM_AMD__ \
+// RUN:   -I%hip_lib_path/../include -o %t.dir/loader.o
+// RUN: %clang %t.dir/loader.o --hip-link \
+// RUN:   --rocm-path=%hip_lib_path/.. -fprofile-instr-generate \
+// RUN:   -L%hip_lib_path -lamdhip64 -o %t.dir/loader
+// RUN: env LLVM_PROFILE_FILE=%t.dir/host.%%p.profraw \
+// RUN:   LD_LIBRARY_PATH=%hip_lib_path:$LD_LIBRARY_PATH \
+// RUN:   HIP_VISIBLE_DEVICES=0 %t.dir/loader %t.dir/module.co
+// RUN: ls %t.dir/gfx*.profraw
+// RUN: llvm-profdata merge %t.dir/*.profraw -o %t.dir/module.profdata
+// RUN: llvm-profdata show --all-functions %t.dir/module.profdata \
+// RUN:   | FileCheck %s
+
+#include <hip/hip_runtime.h>
+
+#ifdef BUILD_MODULE
+
+extern "C" __global__ void module_kernel() {}
+
+int main() { return 0; }
+
+#else
+
+int main(int Argc, char **Argv) {
+  if (Argc != 2)
+    return 1;
+
+  hipModule_t Module = nullptr;
+  if (hipModuleLoad(&Module, Argv[1]) != hipSuccess)
+    return 2;
+
+  hipFunction_t Function = nullptr;
+  if (hipModuleGetFunction(&Function, Module, "module_kernel") != hipSuccess)
+    return 3;
+
+  if (hipModuleLaunchKernel(Function, 1, 1, 1, 1, 1, 1, 0, nullptr, nullptr,
+                            nullptr) != hipSuccess)
+    return 4;
+  if (hipDeviceSynchronize() != hipSuccess)
+    return 5;
+  if (hipModuleUnload(Module) != hipSuccess)
+    return 6;
+  return 0;
+}
+
+#endif
+
+// CHECK:      module_kernel:
+// CHECK-NEXT:   Hash:
+// CHECK-NEXT:   Counters: 1
+// CHECK-NEXT:   Function count: 1

``````````

</details>


https://github.com/llvm/llvm-project/pull/211875