[TransferEngine] Use dynamic loader to support tcp over cuda

[alogfans]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello @alogfans, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the TransferEngine's TCP transport capabilities by introducing dynamic loading for GPU runtime libraries. This change allows the system to detect and efficiently transfer data to and from GPU memory over TCP, even when the TransferEngine is not explicitly compiled with CUDA or MUSA support. The primary goal is to improve flexibility and robustness by decoupling GPU memory handling from strict compile-time configurations, ensuring that GPU-aware transfers can proceed as long as the required runtime libraries are present on the system.

Highlights

• Dynamic GPU Runtime Loading: Introduced a new GpuRuntime class that dynamically loads either CUDA (libcudart.so) or MUSA (libmusa_runtime.so) runtime libraries at runtime, allowing for GPU memory operations without compile-time linking.
• Runtime GPU Memory Detection: The isCudaMemory function has been updated to utilize the GpuRuntime instance. This enables the detection of GPU device pointers even when the TransferEngine is not built with USE_CUDA or USE_MUSA flags, as long as the necessary runtime libraries are available.
• Enhanced TCP Transfer for GPU Memory: The tcp_transport mechanism now supports transferring GPU memory (both reading and writing) by staging data through DRAM using the dynamically loaded GPU runtime functions. This applies when compile-time CUDA/MUSA support is not explicitly enabled.
• Improved Error Handling: Added specific fatal error logging for 'Bad address' errors (ec.value() == 14) that occur during GPU memory transfers. This provides clearer guidance to users, suggesting a rebuild with -DUSE_CUDA=ON if such failures arise without explicit CUDA support.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a GpuRuntime class to dynamically load CUDA/MUSA libraries and handle GPU memory transfers without requiring compile-time linking. This is a good approach to make the binary more portable. The changes correctly adapt writeBody and readBody to use this new dynamic loading mechanism. I have a few suggestions to improve code clarity and maintainability, such as replacing std::cout with structured logging, defining constants for magic numbers, and simplifying some logic.

[gemini-code-assist]

This debug message uses std::cout. For consistency with the rest of the codebase which uses glog, it's better to use LOG(INFO) for logging.

Suggested change

	std::cout << "[GpuRuntime] Loaded GPU runtime: " << lib << "\n";
	LOG(INFO) << "[GpuRuntime] Loaded GPU runtime: " << lib;

[gemini-code-assist]

The magic number 4 corresponds to cudaMemcpyDefault. To improve readability and maintainability, it's better to define this as a named constant in the GpuRuntime class and use it here. A similar change should be applied in readBody as well.

For example, you can add this to GpuRuntime:

static constexpr int kMemcpyDefault = 4; // cudaMemcpyDefault

Suggested change

	gpu.copy(dram_buffer, addr + total_transferred_bytes_, buffer_size,
	4);
	gpu.copy(dram_buffer, addr + total_transferred_bytes_, buffer_size,
	GpuRuntime::kMemcpyDefault);

[gemini-code-assist]

The logic in this #else block can be simplified. The is_cuda_memory variable is already captured by the lambda and holds the result of isCudaMemory(addr), so you can use it in the if condition to avoid a redundant function call. Additionally, the inner if (is_cuda_memory) check before delete[] dram_buffer is redundant.

I'd also recommend using a named constant for the magic number 4 as mentioned in another comment.

Suggested change

	if (isCudaMemory(addr)) {
	auto &gpu = GpuRuntime::instance();
	gpu.copy(addr + total_transferred_bytes_, dram_buffer,
	transferred_bytes, 4);
	if (is_cuda_memory) delete[] dram_buffer;
	}
	if (is_cuda_memory) {
	auto &gpu = GpuRuntime::instance();
	gpu.copy(addr + total_transferred_bytes_, dram_buffer,
	transferred_bytes, 4 /* kMemcpyDefault */);
	delete[] dram_buffer;
	}