gemini-code-assist[bot] commented on code in PR #19770:
URL: https://github.com/apache/tvm/pull/19770#discussion_r3410129639
##########
src/runtime/extra/disco/cuda_ipc/cuda_ipc_memory.cc:
##########
@@ -45,20 +45,22 @@ using tvm::runtime::memory::Buffer;
std::vector<cudaIpcMemHandle_t>
AllGatherIPCHandles(nccl::CCLThreadLocalContext* ctx,
cudaIpcMemHandle_t
local_handle) {
void *d_src, *d_dst;
- CUDA_CALL(cudaMalloc(&d_src, CUDA_IPC_HANDLE_SIZE));
- CUDA_CALL(cudaMalloc(&d_dst, CUDA_IPC_HANDLE_SIZE *
ctx->worker->num_workers));
- CUDA_CALL(cudaMemcpy(d_src, &local_handle, CUDA_IPC_HANDLE_SIZE,
cudaMemcpyHostToDevice));
+ TVM_FFI_CHECK_CUDA_ERROR(cudaMalloc(&d_src, CUDA_IPC_HANDLE_SIZE));
+ TVM_FFI_CHECK_CUDA_ERROR(cudaMalloc(&d_dst, CUDA_IPC_HANDLE_SIZE *
ctx->worker->num_workers));
+ TVM_FFI_CHECK_CUDA_ERROR(
+ cudaMemcpy(d_src, &local_handle, CUDA_IPC_HANDLE_SIZE,
cudaMemcpyHostToDevice));
NCCL_CALL(ncclAllGather(d_src, d_dst, CUDA_IPC_HANDLE_SIZE, ncclChar,
ctx->global_comm,
/*stream=*/nullptr));
std::vector<char> serial_handles(CUDA_IPC_HANDLE_SIZE *
ctx->worker->num_workers, 0);
- CUDA_CALL(cudaMemcpy(serial_handles.data(), d_dst,
- CUDA_IPC_HANDLE_SIZE * ctx->worker->num_workers,
cudaMemcpyDefault));
+ TVM_FFI_CHECK_CUDA_ERROR(cudaMemcpy(serial_handles.data(), d_dst,
+ CUDA_IPC_HANDLE_SIZE *
ctx->worker->num_workers,
+ cudaMemcpyDefault));
std::vector<cudaIpcMemHandle_t> handles(ctx->worker->num_workers);
for (int i = 0; i < ctx->worker->num_workers; ++i) {
memcpy(handles[i].reserved, &serial_handles[i * CUDA_IPC_HANDLE_SIZE],
CUDA_IPC_HANDLE_SIZE);
}
- CUDA_CALL(cudaFree(d_src));
- CUDA_CALL(cudaFree(d_dst));
+ TVM_FFI_CHECK_CUDA_ERROR(cudaFree(d_src));
+ TVM_FFI_CHECK_CUDA_ERROR(cudaFree(d_dst));
Review Comment:
If `cudaMalloc` for `d_dst`, `cudaMemcpy`, or `ncclAllGather` throws an
exception, the allocated GPU memory `d_src` and `d_dst` will be leaked. We
should use a custom deleter with `std::unique_ptr` to ensure `cudaFree` is
called on exception paths.
```c
void *d_src = nullptr, *d_dst = nullptr;
struct CUDAFreeDeleter {
void operator()(void* ptr) const { if (ptr) cudaFree(ptr); }
};
TVM_FFI_CHECK_CUDA_ERROR(cudaMalloc(&d_src, CUDA_IPC_HANDLE_SIZE));
std::unique_ptr<void, CUDAFreeDeleter> src_guard(d_src);
TVM_FFI_CHECK_CUDA_ERROR(cudaMalloc(&d_dst, CUDA_IPC_HANDLE_SIZE *
ctx->worker->num_workers));
std::unique_ptr<void, CUDAFreeDeleter> dst_guard(d_dst);
TVM_FFI_CHECK_CUDA_ERROR(
cudaMemcpy(d_src, &local_handle, CUDA_IPC_HANDLE_SIZE,
cudaMemcpyHostToDevice));
NCCL_CALL(ncclAllGather(d_src, d_dst, CUDA_IPC_HANDLE_SIZE, ncclChar,
ctx->global_comm,
/*stream=*/nullptr));
std::vector<char> serial_handles(CUDA_IPC_HANDLE_SIZE *
ctx->worker->num_workers, 0);
TVM_FFI_CHECK_CUDA_ERROR(cudaMemcpy(serial_handles.data(), d_dst,
CUDA_IPC_HANDLE_SIZE *
ctx->worker->num_workers,
cudaMemcpyDefault));
std::vector<cudaIpcMemHandle_t> handles(ctx->worker->num_workers);
for (int i = 0; i < ctx->worker->num_workers; ++i) {
memcpy(handles[i].reserved, &serial_handles[i * CUDA_IPC_HANDLE_SIZE],
CUDA_IPC_HANDLE_SIZE);
}
return handles;
```
##########
src/backend/cuda/runtime/cuda_device_api.cc:
##########
@@ -50,55 +51,70 @@ class CUDADeviceAPI final : public DeviceAPI {
break;
}
case kMaxThreadsPerBlock: {
- CUDA_CALL(cudaDeviceGetAttribute(&value,
cudaDevAttrMaxThreadsPerBlock, dev.device_id));
+ TVM_FFI_CHECK_CUDA_ERROR(
+ cudaDeviceGetAttribute(&value, cudaDevAttrMaxThreadsPerBlock,
dev.device_id));
break;
}
case kWarpSize: {
- CUDA_CALL(cudaDeviceGetAttribute(&value, cudaDevAttrWarpSize,
dev.device_id));
+ TVM_FFI_CHECK_CUDA_ERROR(
+ cudaDeviceGetAttribute(&value, cudaDevAttrWarpSize,
dev.device_id));
break;
}
case kMaxSharedMemoryPerBlock: {
- CUDA_CALL(
+ TVM_FFI_CHECK_CUDA_ERROR(
cudaDeviceGetAttribute(&value, cudaDevAttrMaxSharedMemoryPerBlock,
dev.device_id));
break;
}
case kComputeVersion: {
std::ostringstream os;
- CUDA_CALL(cudaDeviceGetAttribute(&value,
cudaDevAttrComputeCapabilityMajor, dev.device_id));
+ TVM_FFI_CHECK_CUDA_ERROR(
+ cudaDeviceGetAttribute(&value, cudaDevAttrComputeCapabilityMajor,
dev.device_id));
os << value << ".";
- CUDA_CALL(cudaDeviceGetAttribute(&value,
cudaDevAttrComputeCapabilityMinor, dev.device_id));
+ TVM_FFI_CHECK_CUDA_ERROR(
+ cudaDeviceGetAttribute(&value, cudaDevAttrComputeCapabilityMinor,
dev.device_id));
os << value;
*rv = os.str();
return;
}
case kDeviceName: {
std::string name(256, 0);
- CUDA_DRIVER_CALL(cuDeviceGetName(&name[0], name.size(),
dev.device_id));
+ CUresult result = cuDeviceGetName(&name[0], name.size(),
dev.device_id);
+ if (result != CUDA_SUCCESS && result != CUDA_ERROR_DEINITIALIZED) {
+ const char* msg;
+ cuGetErrorName(result, &msg);
+ TVM_FFI_THROW(CUDAError) << "cuDeviceGetName failed with error: " <<
msg;
+ }
Review Comment:
The `cuGetErrorName` function can return `CUDA_ERROR_INVALID_VALUE` if the
error code is unrecognized, which sets `msg` to `nullptr`. Streaming a null
pointer to `TVM_FFI_THROW` results in undefined behavior (and potential
crashes). It is safer to check if `msg` is valid or use a fallback string.
```c
CUresult result = cuDeviceGetName(&name[0], name.size(),
dev.device_id);
if (result != CUDA_SUCCESS && result != CUDA_ERROR_DEINITIALIZED) {
const char* msg = nullptr;
cuGetErrorName(result, &msg);
TVM_FFI_THROW(CUDAError) << "cuDeviceGetName failed with error: "
<< (msg ? msg : "unknown error");
}
```
##########
src/runtime/vm/cuda/cuda_graph_builtin.cc:
##########
@@ -116,11 +116,11 @@ class ScopedCUDAStream {
class CUDACaptureStream {
public:
explicit CUDACaptureStream(cudaGraph_t* graph) : output_graph_(graph) {
- CUDA_CALL(cudaGetDevice(&device_id_));
+ TVM_FFI_CHECK_CUDA_ERROR(cudaGetDevice(&device_id_));
TVM_FFI_CHECK_SAFE_CALL(
TVMFFIEnvSetStream(kDLCUDA, device_id_, capture_stream_,
reinterpret_cast<TVMFFIStreamHandle*>(&prev_default_stream_)));
- CUDA_CALL(cudaStreamBeginCapture(capture_stream_,
cudaStreamCaptureModeGlobal));
+ TVM_FFI_CHECK_CUDA_ERROR(cudaStreamBeginCapture(capture_stream_,
cudaStreamCaptureModeGlobal));
Review Comment:
If `cudaStreamBeginCapture` throws an exception, the constructor of
`CUDACaptureStream` fails and the destructor is never called. This leaves the
thread-local stream environment pointing to the destroyed `capture_stream_`
member, causing subsequent operations to crash or exhibit undefined behavior.
We should catch any exception and restore the stream environment before
rethrowing.
```c
TVM_FFI_CHECK_CUDA_ERROR(cudaGetDevice(&device_id_));
TVM_FFI_CHECK_SAFE_CALL(
TVMFFIEnvSetStream(kDLCUDA, device_id_, capture_stream_,
reinterpret_cast<TVMFFIStreamHandle*>(&prev_default_stream_)));
try {
TVM_FFI_CHECK_CUDA_ERROR(cudaStreamBeginCapture(capture_stream_,
cudaStreamCaptureModeGlobal));
} catch (...) {
TVMFFIEnvSetStream(kDLCUDA, device_id_, prev_default_stream_, nullptr);
throw;
}
```
##########
src/runtime/extra/contrib/tensorrt/tensorrt_calibrator.h:
##########
@@ -55,8 +56,9 @@ class TensorRTCalibrator : public
nvinfer1::IInt8EntropyCalibrator2 {
std::vector<float*> data_host(bindings.size(), nullptr);
for (size_t i = 0; i < bindings.size(); ++i) {
data_host[i] = new float[batch_size_ * binding_sizes[i]];
- CUDA_CALL(cudaMemcpy(static_cast<void*>(data_host[i]), bindings[i],
- batch_size_ * binding_sizes[i] * sizeof(float),
cudaMemcpyDeviceToHost));
+ TVM_FFI_CHECK_CUDA_ERROR(cudaMemcpy(static_cast<void*>(data_host[i]),
bindings[i],
+ batch_size_ * binding_sizes[i] *
sizeof(float),
+ cudaMemcpyDeviceToHost));
}
Review Comment:
If `cudaMemcpy` throws an exception, any previously allocated `data_host[i]`
in the loop will be leaked because they are not yet pushed to `data_` and the
destructor only cleans up `data_`. We should use a `try-catch` block to free
them on failure.
```c
std::vector<float*> data_host(bindings.size(), nullptr);
try {
for (size_t i = 0; i < bindings.size(); ++i) {
data_host[i] = new float[batch_size_ * binding_sizes[i]];
TVM_FFI_CHECK_CUDA_ERROR(cudaMemcpy(static_cast<void*>(data_host[i]),
bindings[i],
batch_size_ * binding_sizes[i] *
sizeof(float),
cudaMemcpyDeviceToHost));
}
} catch (...) {
for (float* ptr : data_host) {
delete[] ptr;
}
throw;
}
```
##########
src/backend/cuda/runtime/cuda_module.cc:
##########
@@ -127,7 +127,12 @@ class CUDAModuleNode : public ffi::ModuleObj {
std::lock_guard<std::mutex> lock(mutex_);
// must recheck under the lock scope
if (module_[device_id] == nullptr) {
- CUDA_DRIVER_CALL(cuModuleLoadData(&(module_[device_id]), code_.data()));
+ CUresult result = cuModuleLoadData(&(module_[device_id]), code_.data());
+ if (result != CUDA_SUCCESS && result != CUDA_ERROR_DEINITIALIZED) {
+ const char* msg;
+ cuGetErrorName(result, &msg);
+ TVM_FFI_THROW(CUDAError) << "cuModuleLoadData failed with error: " <<
msg;
+ }
Review Comment:
The `cuGetErrorName` function can return `CUDA_ERROR_INVALID_VALUE` if the
error code is unrecognized, which sets `msg` to `nullptr`. Streaming a null
pointer to `TVM_FFI_THROW` results in undefined behavior (and potential
crashes). It is safer to check if `msg` is valid or use a fallback string.
```suggestion
CUresult result = cuModuleLoadData(&(module_[device_id]),
code_.data());
if (result != CUDA_SUCCESS && result != CUDA_ERROR_DEINITIALIZED) {
const char* msg = nullptr;
cuGetErrorName(result, &msg);
TVM_FFI_THROW(CUDAError) << "cuModuleLoadData failed with error: "
<< (msg ? msg : "unknown error");
}
```
##########
src/runtime/extra/contrib/cublas/cublas_utils.cc:
##########
@@ -51,7 +50,7 @@ CuBlasThreadEntry* CuBlasThreadEntry::ThreadLocal(DLDevice
curr_device) {
CuBlasLtThreadEntry::CuBlasLtThreadEntry() {
CHECK_CUBLAS_ERROR(cublasLtCreate(&handle));
CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceCreate(&matmul_pref_desc));
- CUDA_CALL(cudaMalloc(&workspace_ptr, workspace_size));
+ TVM_FFI_CHECK_CUDA_ERROR(cudaMalloc(&workspace_ptr, workspace_size));
}
Review Comment:
If `cudaMalloc` throws an exception, the constructor fails and the
destructor is not called. This causes `handle` and `matmul_pref_desc` to be
leaked. We should wrap the constructor body in a `try-catch` block to perform
cleanup on failure.
```c
CuBlasLtThreadEntry::CuBlasLtThreadEntry() {
CHECK_CUBLAS_ERROR(cublasLtCreate(&handle));
try {
CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceCreate(&matmul_pref_desc));
TVM_FFI_CHECK_CUDA_ERROR(cudaMalloc(&workspace_ptr, workspace_size));
} catch (...) {
if (matmul_pref_desc) {
cublasLtMatmulPreferenceDestroy(matmul_pref_desc);
}
if (handle) {
cublasLtDestroy(handle);
}
throw;
}
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
