diff --git a/mops/CMakeLists.txt b/mops/CMakeLists.txt index 809ef03b..6c8cc76d 100644 --- a/mops/CMakeLists.txt +++ b/mops/CMakeLists.txt @@ -33,6 +33,7 @@ check_language(CUDA) if(CMAKE_CUDA_COMPILER) enable_language(CUDA) set(CUDA_USE_STATIC_CUDA_RUNTIME OFF CACHE BOOL "" FORCE) + set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -lineinfo") else() message(STATUS "Could not find a CUDA compiler") endif() @@ -118,6 +119,7 @@ if(CMAKE_CUDA_COMPILER) "src/opsa/opsa.cu" "src/hpe/hpe.cu" "src/sap/sap.cu" + "src/sasaw/sasaw.cu" ) endif() diff --git a/mops/include/mops/cuda_first_occurences.hpp b/mops/include/mops/cuda_first_occurences.hpp new file mode 100644 index 00000000..8cabb2f0 --- /dev/null +++ b/mops/include/mops/cuda_first_occurences.hpp @@ -0,0 +1,18 @@ +#ifndef FIRST_OCCURENCES_HPP +#define FIRST_OCCURENCES_HPP + +#include "mops/tensor.hpp" +#include + +/* + * Computes the indexes at which the sorted input array (receiver_list) change in value, with 0 + * prepended. For example, if the receiver list is [0, 0, 0, 1, 1, 1, 2, 2], then the output would + * be [0, 3, 6]. nelements_input refers to the size of the input receiver_list, nelements_output + * refers, to the number of output elements. The elements of receiver_list **must** be sorted such + * that all references to each index appear contiguously and continuously. + */ +int32_t* calculate_first_occurences_cuda( + const int32_t* receiver_list, int32_t nelements_input, int32_t nelements_output +); + +#endif // FIRST_OCCURENCES_HP \ No newline at end of file diff --git a/mops/include/mops/cuda_utils.cuh b/mops/include/mops/cuda_utils.cuh new file mode 100644 index 00000000..13a83d1e --- /dev/null +++ b/mops/include/mops/cuda_utils.cuh @@ -0,0 +1,45 @@ +#ifndef CUDA_UTILS_CUH +#define CUDA_UTILS_CUH + +#include +#include +#include + +using namespace std; + +#define CUDA_CHECK_ERROR(err) \ + do { \ + cudaError_t err_cuda = err; \ + if (err_cuda != cudaSuccess) { \ + fprintf( \ + stderr, \ + "CUDA error in file '%s' in line %i: %s\n", \ + __FILE__, \ + __LINE__, \ + cudaGetErrorString(err_cuda) \ + ); \ + exit(EXIT_FAILURE); \ + } \ + } while (0) + +__host__ __device__ int32_t find_integer_divisor(int32_t x, int32_t bdim); + +/* + * helper function to allocate correctly sized shared memory buffers. creates a pointer reference to + * a shared memory array with n_elements number of elements. On exit, ptr is shifted right by + * nelements * sizeof(T), and space is incremented by this same amount. + * + */ +template +__host__ __device__ T* shared_array(std::size_t n_elements, void*& ptr, std::size_t* space) noexcept; + +/* + * helper function to allocate correctly sized shared memory buffers identically to the shared_array + * method, but in addition are aligned to a certain byte boundary given by alignment. + */ +template +__host__ __device__ T* align_array( + std::size_t n_elements, void*& ptr, const std::size_t alignment, std::size_t* space +) noexcept; + +#endif // CUDA_UTILS_CUH \ No newline at end of file diff --git a/mops/src/sasaw/cuda.tpp b/mops/src/sasaw/cuda.tpp deleted file mode 100644 index 9322d480..00000000 --- a/mops/src/sasaw/cuda.tpp +++ /dev/null @@ -1,38 +0,0 @@ -#include - -#include "mops/sasaw.hpp" - -template -void mops::cuda::sparse_accumulation_scatter_add_with_weights( - Tensor, - Tensor, - Tensor, - Tensor, - Tensor, - Tensor, - Tensor, - Tensor, - Tensor, - Tensor -) { - throw std::runtime_error("CUDA implementation does not exist yet"); -} - -template -void mops::cuda::sparse_accumulation_scatter_add_with_weights_vjp( - Tensor, - Tensor , - Tensor, - Tensor, - Tensor , - Tensor , - Tensor , - Tensor , - Tensor, - Tensor, - Tensor, - Tensor, - Tensor -) { - throw std::runtime_error("CUDA implementation does not exist yet"); -} diff --git a/mops/src/sasaw/sasaw.cpp b/mops/src/sasaw/sasaw.cpp index 1bc107f5..5d42c6c2 100644 --- a/mops/src/sasaw/sasaw.cpp +++ b/mops/src/sasaw/sasaw.cpp @@ -59,9 +59,7 @@ template void mops::sparse_accumulation_scatter_add_with_weights_vjp( Tensor indices_output_2 ); -#ifdef MOPS_CUDA_ENABLED -#include "cuda.tpp" -#else +#ifndef MOPS_CUDA_ENABLED template void mops::cuda:: sparse_accumulation_scatter_add_with_weights(Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor) { @@ -74,8 +72,6 @@ void mops::cuda:: throw std::runtime_error("MOPS was not compiled with CUDA support"); } -#endif - // explicit instantiations of CUDA templates template void mops::cuda::sparse_accumulation_scatter_add_with_weights( Tensor output, @@ -134,3 +130,5 @@ template void mops::cuda::sparse_accumulation_scatter_add_with_weights_vjp indices_output_1, Tensor indices_output_2 ); + +#endif \ No newline at end of file diff --git a/mops/src/sasaw/sasaw.cu b/mops/src/sasaw/sasaw.cu index 8421bac6..4345037a 100644 --- a/mops/src/sasaw/sasaw.cu +++ b/mops/src/sasaw/sasaw.cu @@ -1 +1,225 @@ -// todo: cuda device code +#include "mops/sasaw.hpp" + +#include "internal/checks.hpp" +#include "internal/cuda_first_occurences.cuh" +#include "internal/cuda_utils.cuh" + +#define WARP_SIZE 32 +#define NWARPS_PER_BLOCK 4 + +using namespace mops; +using namespace mops::cuda; + +template +__global__ __launch_bounds__(WARP_SIZE* NWARPS_PER_BLOCK) void sparse_accumulation_scatter_add_with_weights_kernel( + Tensor output, + Tensor A, + Tensor B, + Tensor C, + Tensor W, + Tensor indices_A, + Tensor indices_W_1, + Tensor indices_W_2, + Tensor indices_output_1, + Tensor indices_output_2, + Tensor first_occurences +) { + + extern __shared__ char buffer[]; + + void* sptr = buffer; + size_t space = 0; + + scalar_t* buffer_out = shared_array(output.shape[1] * output.shape[2], sptr, &space); + scalar_t* buffer_A = shared_array(NWARPS_PER_BLOCK * A.shape[1], sptr, &space); + scalar_t* buffer_B = shared_array(NWARPS_PER_BLOCK * B.shape[1], sptr, &space); + scalar_t* buffer_W = + shared_array(NWARPS_PER_BLOCK * *W.shape[1] * W.shape[2], sptr, &space); + + scalar_t* buffer_C = shared_array(C.shape[0], sptr, &space); + int8_t* buffer_indices_A = shared_array(indices_A.shape[0], sptr, &space); + int8_t* buffer_indices_W_2 = shared_array(indices_W_2.shape[0], sptr, &space); + int8_t* buffer_indices_output_2 = shared_array(indices_output_2.shape[0], sptr, &space); + + int laneID = threadIdx.x % WARP_SIZE; + int warpID = threadIdx.x / WARP_SIZE; + + int32_t sample_start = first_occurences.data[blockIdx.x]; + int32_t sample_end = -1; + int32_t node_index = -1; + + if (sample_start != -1) { + node_index = indices_output_1.data[sample_start]; + sample_end = (blockIdx.x == first_occurences.shape[0] - 1) + ? indices_output_1.shape[0] + : (first_occurences.data[blockIdx.x + 1] == -1 + ? indices_output_1.shape[0] + : first_occurences.data[blockIdx.x + 1]); + } + + int32_t nsamples = sample_end - sample_start; + + if (nsamples == 0) { + return; + } + + for (int tid = threadIdx.x; tid < indices_A.shape[0]; tid += blockDim.x) { + buffer_indices_A[tid] = (int8_t)indices_A.data[tid]; + buffer_indices_W_2[tid] = (int8_t)indices_W_2.data[tid]; + buffer_C[tid] = C.data[tid]; + buffer_indices_output_2[tid] = (int8_t)indices_output_2.data[tid]; + } + + for (int tid = threadIdx.x; tid < output.shape[1] * output.shape[2]; tid += blockDim.x) { + buffer_out[tid] = 0.0; + } + + __syncthreads(); + + for (int sample_idx = 0; sample_idx < nsamples; sample_idx++) { + + int sample = sample_start + sample_idx + warpID; + + if (sample >= sample_end) { + break; + } + + // load in temporary buffers for each sample + + for (int tid = laneID; tid < A.shape[1]; tid += WARP_SIZE) { + buffer_A[warpID * A.shape[1] + tid] = A.data[sample * A.shape[1] + tid]; + } + + for (int tid = laneID; tid < B.shape[1]; tid += WARP_SIZE) { + buffer_B[warpID * B.shape[1] + tid] = B.data[sample * B.shape[1] + tid]; + } + + for (int j = 0; j < W.shape[1]; j++) { + for (int tid = laneID; tid < W.shape[2]; tid += WARP_SIZE) { + buffer_W[warpID * W.shape[1] * W.shape[2] + j * W.shape[2] + tid] = + W.data[indices_W_1.data[sample] * W.shape[1] * W.shape[2] + j * W.shape[2] + tid]; + } + } + + __syncwarp(); + + for (int k = 0; k < C.shape[0]; k++) { + for (int tid = laneID; tid < W.shape[2]; tid += WARP_SIZE) { + scalar_t w = + buffer_W[warpID * W.shape[1] * W.shape[2] + buffer_indices_W_2[k] * W.shape[2] + tid]; + + int32_t a_idx = buffer_indices_A[k]; + int32_t index_output_2 = buffer_indices_output_2[k]; + + atomicAdd( + &buffer_out[index_output_2 * output.shape[2] + tid], + A.data[sample * A.shape[1] + a_idx] * B[sample * B.shape[1] + tid] * + buffer_C[k] * w + ); + } + } + } + + __syncthreads(); + + for (int i = warpID; i < output.shape[1]; i += NWARPS_PER_BLOCK) { + for (int j = laneID; j < output.shape[2]; j += WARP_SIZE) { + output.data[node_index * output.shape[1] * output.shape[2] + i * output.shape[2] + j] = + buffer_out[i * output.shape[2] + j]; + } + } +} + +template +void mops::cuda::sparse_accumulation_scatter_add_with_weights( + Tensor output, + Tensor A, + Tensor B, + Tensor C, + Tensor W, + Tensor indices_A, + Tensor indices_W_1, + Tensor indices_W_2, + Tensor indices_output_1, + Tensor indices_output_2 +) { + // TODO +} + +// explicit instantiations of CUDA templates +template void mops::cuda::sparse_accumulation_scatter_add_with_weights( + Tensor output, + Tensor A, + Tensor B, + Tensor C, + Tensor W, + Tensor indices_A, + Tensor indices_W_1, + Tensor indices_W_2, + Tensor indices_output_1, + Tensor indices_output_2 +); + +template void mops::cuda::sparse_accumulation_scatter_add_with_weights( + Tensor output, + Tensor A, + Tensor B, + Tensor C, + Tensor W, + Tensor indices_A, + Tensor indices_W_1, + Tensor indices_W_2, + Tensor indices_output_1, + Tensor indices_output_2 +); + +template +void mops::cuda::sparse_accumulation_scatter_add_with_weights_vjp( + Tensor grad_A, + Tensor grad_B, + Tensor grad_W, + Tensor grad_output, + Tensor A, + Tensor B, + Tensor C, + Tensor W, + Tensor indices_A, + Tensor indices_W_1, + Tensor indices_W_2, + Tensor indices_output_1, + Tensor indices_output_2 +) { + // TODO +} + +template void mops::cuda::sparse_accumulation_scatter_add_with_weights_vjp( + Tensor grad_A, + Tensor grad_B, + Tensor grad_W, + Tensor grad_output, + Tensor A, + Tensor B, + Tensor C, + Tensor W, + Tensor indices_A, + Tensor indices_W_1, + Tensor indices_W_2, + Tensor indices_output_1, + Tensor indices_output_2 +); + +template void mops::cuda::sparse_accumulation_scatter_add_with_weights_vjp( + Tensor grad_A, + Tensor grad_B, + Tensor grad_W, + Tensor grad_output, + Tensor A, + Tensor B, + Tensor C, + Tensor W, + Tensor indices_A, + Tensor indices_W_1, + Tensor indices_W_2, + Tensor indices_output_1, + Tensor indices_output_2 +); diff --git a/python/mops/tests/sap_vjp.py b/python/mops/tests/sap_vjp.py index 76feff2b..e90e2802 100644 --- a/python/mops/tests/sap_vjp.py +++ b/python/mops/tests/sap_vjp.py @@ -102,7 +102,7 @@ def test_sap_vjp_cupy(valid_arguments): indices_B = cp.array(indices_B) indices_output = cp.array(indices_output) - ref_grad_A, ref_grad_B = ref_sap_vjp( # noqa: F841 + grad_A, grad_B = ref_sap_vjp( # noqa: F841 grad_output, A, B, @@ -111,16 +111,9 @@ def test_sap_vjp_cupy(valid_arguments): indices_B, indices_output, ) - - grad_A, grad_B = sap_vjp( - grad_output, - A, - B, - C, - indices_A, - indices_B, - indices_output, - ) - - assert cp.allclose(ref_grad_A, grad_A) - assert cp.allclose(ref_grad_B, grad_B) + with pytest.raises( + mops.status.MopsError, match="CUDA implementation does not exist yet" + ): + grad_A, grad_B = sap_vjp( # noqa: F841 + grad_output, A, B, C, indices_A, indices_B, indices_output + )