Skip to content

[libc] Handle differing wavefront sizes correctly in the AMDHSA loader #117788

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Nov 27, 2024
Merged

[libc] Handle differing wavefront sizes correctly in the AMDHSA loader #117788

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3c2b2e6
[libc] Handle differing wavefront sizes correctly in the AMDHSA loader
jhuber6 Nov 26, 2024
5b22893
test
jhuber6 Nov 26, 2024
641143c
test eq
jhuber6 Nov 27, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions libc/utils/gpu/loader/amdgpu/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ set(LLVM_LINK_COMPONENTS
Object
Option
Support
FrontendOffloading
)

add_llvm_executable(amdhsa-loader amdhsa-loader.cpp)
Expand Down
46 changes: 27 additions & 19 deletions libc/utils/gpu/loader/amdgpu/amdhsa-loader.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,8 @@
#include "hsa/hsa_ext_amd.h"
#endif

#include "llvm/Frontend/Offloading/Utility.h"

#include <atomic>
#include <cstdio>
#include <cstdlib>
Expand Down Expand Up @@ -163,17 +165,13 @@ hsa_status_t launch_kernel(hsa_agent_t dev_agent, hsa_executable_t executable,
hsa_queue_t *queue, rpc::Server &server,
const LaunchParameters &params,
const char *kernel_name, args_t kernel_args,
bool print_resource_usage) {
uint32_t wavefront_size, bool print_resource_usage) {
// Look up the kernel in the loaded executable.
hsa_executable_symbol_t symbol;
if (hsa_status_t err = hsa_executable_get_symbol_by_name(
executable, kernel_name, &dev_agent, &symbol))
return err;

uint32_t wavefront_size = 0;
if (hsa_status_t err = hsa_agent_get_info(
dev_agent, HSA_AGENT_INFO_WAVEFRONT_SIZE, &wavefront_size))
handle_error(err);
// Retrieve different properties of the kernel symbol used for launch.
uint64_t kernel;
uint32_t args_size;
Expand Down Expand Up @@ -419,6 +417,16 @@ int load(int argc, const char **argv, const char **envp, void *image,
dev_agent, &coarsegrained_pool))
handle_error(err);

// The AMDGPU target can change its wavefront size. There currently isn't a
// good way to look this up through the HSA API so we use the LLVM interface.
uint16_t abi_version;
llvm::StringRef image_ref(reinterpret_cast<char *>(image), size);
llvm::StringMap<llvm::offloading::amdgpu::AMDGPUKernelMetaData> info_map;
if (llvm::Error err = llvm::offloading::amdgpu::getAMDGPUMetaDataFromImage(
llvm::MemoryBufferRef(image_ref, ""), info_map, abi_version)) {
handle_error(llvm::toString(std::move(err)).c_str());
}

// Allocate fine-grained memory on the host to hold the pointer array for the
// copied argv and allow the GPU agent to access it.
auto allocator = [&](uint64_t size) -> void * {
Expand Down Expand Up @@ -448,10 +456,10 @@ int load(int argc, const char **argv, const char **envp, void *image,
hsa_amd_memory_fill(dev_ret, 0, /*count=*/1);

// Allocate finegrained memory for the RPC server and client to share.
uint32_t wavefront_size = 0;
if (hsa_status_t err = hsa_agent_get_info(
dev_agent, HSA_AGENT_INFO_WAVEFRONT_SIZE, &wavefront_size))
handle_error(err);
uint32_t wavefront_size =
llvm::max_element(info_map, [](auto &&x, auto &&y) {
return x.second.WavefrontSize < y.second.WavefrontSize;
})->second.WavefrontSize;

// Set up the RPC server.
void *rpc_buffer;
Expand Down Expand Up @@ -513,7 +521,6 @@ int load(int argc, const char **argv, const char **envp, void *image,
if (HSA_STATUS_SUCCESS ==
hsa_executable_get_symbol_by_name(executable, "__llvm_libc_clock_freq",
&dev_agent, &freq_sym)) {

void *host_clock_freq;
if (hsa_status_t err =
hsa_amd_memory_pool_allocate(finegrained_pool, sizeof(uint64_t),
Expand Down Expand Up @@ -553,16 +560,17 @@ int load(int argc, const char **argv, const char **envp, void *image,

LaunchParameters single_threaded_params = {1, 1, 1, 1, 1, 1};
begin_args_t init_args = {argc, dev_argv, dev_envp};
if (hsa_status_t err = launch_kernel(dev_agent, executable, kernargs_pool,
coarsegrained_pool, queue, server,
single_threaded_params, "_begin.kd",
init_args, print_resource_usage))
if (hsa_status_t err = launch_kernel(
dev_agent, executable, kernargs_pool, coarsegrained_pool, queue,
server, single_threaded_params, "_begin.kd", init_args,
info_map["_begin"].WavefrontSize, print_resource_usage))
handle_error(err);

start_args_t args = {argc, dev_argv, dev_envp, dev_ret};
if (hsa_status_t err = launch_kernel(
dev_agent, executable, kernargs_pool, coarsegrained_pool, queue,
server, params, "_start.kd", args, print_resource_usage))
server, params, "_start.kd", args, info_map["_start"].WavefrontSize,
print_resource_usage))
handle_error(err);

void *host_ret;
Expand All @@ -580,10 +588,10 @@ int load(int argc, const char **argv, const char **envp, void *image,
int ret = *static_cast<int *>(host_ret);

end_args_t fini_args = {ret};
if (hsa_status_t err = launch_kernel(dev_agent, executable, kernargs_pool,
coarsegrained_pool, queue, server,
single_threaded_params, "_end.kd",
fini_args, print_resource_usage))
if (hsa_status_t err = launch_kernel(
dev_agent, executable, kernargs_pool, coarsegrained_pool, queue,
server, single_threaded_params, "_end.kd", fini_args,
info_map["_end"].WavefrontSize, print_resource_usage))
handle_error(err);

if (hsa_status_t err = hsa_amd_memory_pool_free(rpc_buffer))
Expand Down