tongbuxiugai

GPUBaseLib/GPUTool/GPUTool.cu

@@ -353,18 +353,40 @@ extern "C"  void FreeCUDAHost(void* ptr) {
 	ptr = nullptr;
 }
 
-// GPU参数内存声明
-extern "C" void* mallocCUDADevice(size_t memsize) {
-	void* ptr;
-	cudaMalloc(&ptr, memsize);
-#ifdef __CUDADEBUG__
-	cudaError_t err = cudaGetLastError();
-	if (err != cudaSuccess) {
-		printf("mallocCUDADevice CUDA Error: %s, malloc memory : %d byte\n", cudaGetErrorString(err), memsize);
-		exit(2);
+// 多GPU内存分配函数
+void* mallocCUDADevice(size_t memsize, int device_id = 0)
+{
+	void* ptr = nullptr;
+	cudaError_t err;
+
+	// 1. 检查设备ID有效性
+	int num_devices;
+	cudaGetDeviceCount(&num_devices);
+	if (device_id < 0 || device_id >= num_devices)
+	{
+		printfinfo("Invalid device ID: %d\n", device_id);
+		return nullptr;
 	}
-#endif // __CUDADEBUG__
-	cudaDeviceSynchronize();
+
+	// 2. 切换目标GPU设备
+	err = cudaSetDevice(device_id);
+	if (err != cudaSuccess)
+	{
+		PrintLasterError("cudaSetDevice");
+		return nullptr;
+	}
+
+	// 3. 分配显存
+	err = cudaMalloc(&ptr, memsize);
+	if (err != cudaSuccess)
+	{
+		PrintLasterError("cudaMalloc");
+		return nullptr;
+	}
+
+	// 4. 可选：同步设备（视需求决定是否保留）
+	// cudaDeviceSynchronize();
+
 	return ptr;
 }