From 4c6f33bdd55febea1f1d5da36ceea05dfa87e5ea Mon Sep 17 00:00:00 2001 From: Sachin Pisal Date: Thu, 1 Aug 2024 14:09:42 -0700 Subject: [PATCH 1/2] migrating changes from CuStateVecCircuitSimulator.cu --- .../custatevec/CuStateVecCircuitSimulator.cpp | 61 +++++++++++++++++-- 1 file changed, 55 insertions(+), 6 deletions(-) diff --git a/runtime/nvqir/custatevec/CuStateVecCircuitSimulator.cpp b/runtime/nvqir/custatevec/CuStateVecCircuitSimulator.cpp index 03044c436b..69be83ec9d 100644 --- a/runtime/nvqir/custatevec/CuStateVecCircuitSimulator.cpp +++ b/runtime/nvqir/custatevec/CuStateVecCircuitSimulator.cpp @@ -190,16 +190,43 @@ class CuStateVecCircuitSimulator } // User state provided... - - // FIXME handle case where pointer is a device pointer + // Check if the pointer is a device pointer + cudaPointerAttributes attributes; + HANDLE_CUDA_ERROR(cudaPointerGetAttributes(&attributes, state)); // First allocation, so just set the user provided data here ScopedTraceWithContext( "CuStateVecCircuitSimulator::addQubitsToState cudaMemcpy", stateDimension * sizeof(CudaDataType)); - HANDLE_CUDA_ERROR(cudaMemcpy(deviceStateVector, state, + + if (attributes.type == cudaMemoryTypeDevice) { + int currentDevice; + HANDLE_CUDA_ERROR(cudaGetDevice(¤tDevice)); + + if (attributes.device != currentDevice) { + // Memory is on a different GPU + // Set the device to the device where the memory is located + HANDLE_CUDA_ERROR(cudaSetDevice(attributes.device)); + + // Perform device to device copy + HANDLE_CUDA_ERROR(cudaMemcpy(deviceStateVector, state, + stateDimension * sizeof(CudaDataType), + cudaMemcpyDeviceToDevice)); + + // Restore the current device + HANDLE_CUDA_ERROR(cudaSetDevice(currentDevice)); + } else { + // Memory is on the same GPU + HANDLE_CUDA_ERROR(cudaMemcpy(deviceStateVector, state, + stateDimension * sizeof(CudaDataType), + cudaMemcpyDeviceToDevice)); + } + } else { + // Else, copy from host to device + HANDLE_CUDA_ERROR(cudaMemcpy(deviceStateVector, state, stateDimension * sizeof(CudaDataType), cudaMemcpyHostToDevice)); + } return; } @@ -221,11 +248,33 @@ class CuStateVecCircuitSimulator n_blocks, threads_per_block, otherState, (1UL << count)); HANDLE_CUDA_ERROR(cudaGetLastError()); } else { - - // FIXME Handle case where data is already on GPU - HANDLE_CUDA_ERROR(cudaMemcpy(otherState, state, + // Check if the pointer is a device pointer + cudaPointerAttributes attributes; + HANDLE_CUDA_ERROR(cudaPointerGetAttributes(&attributes, state)); + + if (attributes.type == cudaMemoryTypeDevice) { + int currentDevice; + HANDLE_CUDA_ERROR(cudaGetDevice(¤tDevice)); + + if (attributes.device != currentDevice) { + // Memory is on a different GPU + // Set the device to the device where the memory is located + HANDLE_CUDA_ERROR(cudaSetDevice(attributes.device)); + + // Perform device to device copy + HANDLE_CUDA_ERROR(cudaMemcpy(otherState, state, + stateDimension * sizeof(CudaDataType), + cudaMemcpyDeviceToDevice)); + + // Restore the current device + HANDLE_CUDA_ERROR(cudaSetDevice(currentDevice)); + } + } else { + // Else, copy from host to device + HANDLE_CUDA_ERROR(cudaMemcpy(otherState, state, (1UL << count) * sizeof(CudaDataType), cudaMemcpyHostToDevice)); + } } { From b5ccda4b07c6234808cd07876336cc1391d4ecb4 Mon Sep 17 00:00:00 2001 From: Sachin Pisal Date: Thu, 1 Aug 2024 15:32:02 -0700 Subject: [PATCH 2/2] formatting --- .../custatevec/CuStateVecCircuitSimulator.cpp | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/runtime/nvqir/custatevec/CuStateVecCircuitSimulator.cpp b/runtime/nvqir/custatevec/CuStateVecCircuitSimulator.cpp index 69be83ec9d..5c2f3a01c3 100644 --- a/runtime/nvqir/custatevec/CuStateVecCircuitSimulator.cpp +++ b/runtime/nvqir/custatevec/CuStateVecCircuitSimulator.cpp @@ -210,22 +210,22 @@ class CuStateVecCircuitSimulator // Perform device to device copy HANDLE_CUDA_ERROR(cudaMemcpy(deviceStateVector, state, - stateDimension * sizeof(CudaDataType), - cudaMemcpyDeviceToDevice)); + stateDimension * sizeof(CudaDataType), + cudaMemcpyDeviceToDevice)); // Restore the current device HANDLE_CUDA_ERROR(cudaSetDevice(currentDevice)); } else { // Memory is on the same GPU HANDLE_CUDA_ERROR(cudaMemcpy(deviceStateVector, state, - stateDimension * sizeof(CudaDataType), - cudaMemcpyDeviceToDevice)); + stateDimension * sizeof(CudaDataType), + cudaMemcpyDeviceToDevice)); } } else { // Else, copy from host to device HANDLE_CUDA_ERROR(cudaMemcpy(deviceStateVector, state, - stateDimension * sizeof(CudaDataType), - cudaMemcpyHostToDevice)); + stateDimension * sizeof(CudaDataType), + cudaMemcpyHostToDevice)); } return; } @@ -263,8 +263,8 @@ class CuStateVecCircuitSimulator // Perform device to device copy HANDLE_CUDA_ERROR(cudaMemcpy(otherState, state, - stateDimension * sizeof(CudaDataType), - cudaMemcpyDeviceToDevice)); + stateDimension * sizeof(CudaDataType), + cudaMemcpyDeviceToDevice)); // Restore the current device HANDLE_CUDA_ERROR(cudaSetDevice(currentDevice)); @@ -272,8 +272,8 @@ class CuStateVecCircuitSimulator } else { // Else, copy from host to device HANDLE_CUDA_ERROR(cudaMemcpy(otherState, state, - (1UL << count) * sizeof(CudaDataType), - cudaMemcpyHostToDevice)); + (1UL << count) * sizeof(CudaDataType), + cudaMemcpyHostToDevice)); } }