mulit-GPU and CUDA Stream Support

ci/docker/Dockerfile.base

@@ -0,0 +1,6 @@
+FROM nvcr.io/nvidia/pytorch:23.10-py3
+
+RUN apt-get update
+
+# install boost test framework
+RUN apt-get install -y libboost-test-dev

ci/pipeline.yml

@@ -0,0 +1,36 @@
+include:
+  - remote: 'https://gitlab.com/cscs-ci/recipes/-/raw/master/templates/v2/.ci-ext.yml'
+
+stages:
+  - build
+  - test
+
+build_base_image_job:
+  stage: build
+  extends: .container-builder-dynamic-name
+  timeout: 2h
+  variables:
+    DOCKERFILE: ci/docker/Dockerfile.base
+    WATCH_FILECHANGES: $DOCKERFILE
+    PERSIST_IMAGE_NAME: $CSCS_REGISTRY_PATH/base/public/mops
+
+test_job:
+  stage: test
+  extends: .container-runner-daint-gpu
+  image: $BASE_IMAGE
+  timeout: 2h
+  script:
+    - export CUDA_HOME="/usr/local/cuda"
+    - python3 -m pip install --upgrade pip
+    - echo "Install Tox"
+    - python3 -m pip install tox
+    - echo "Run the Tox Script"
+    - tox
+    - echo "Tox script completed"
+
+  variables:
+    SLURM_JOB_NUM_NODES: 1
+    SLURM_PARTITION: normal
+    SLURM_NTASKS: 1
+    SLURM_TIMELIMIT: '00:40:00'
+    GIT_STRATEGY: fetch

mops-torch/src/hpe.cpp

@@ -1,3 +1,8 @@
+#ifdef MOPS_CUDA_ENABLED
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+#endif
+
 #include "mops/torch/hpe.hpp"
 #include "mops/torch/utils.hpp"
 
@@ -38,15 +43,25 @@ torch::Tensor HomogeneousPolynomialEvaluation::forward(
         });
     } else if (A.device().is_cuda()) {
 
+#ifndef MOPS_CUDA_ENABLED
+        C10_THROW_ERROR(ValueError, "MOPS was not compiled with CUDA support " + A.device().str());
+#else
+        c10::cuda::CUDAGuard deviceGuard{A.device()};
+        cudaStream_t currstream = c10::cuda::getCurrentCUDAStream();
+        void* stream = reinterpret_cast<void*>(currstream);
+
         AT_DISPATCH_FLOATING_TYPES(A.scalar_type(), "homogeneous_polynomial_evaluation", [&]() {
             mops::cuda::homogeneous_polynomial_evaluation<scalar_t>(
                 details::torch_to_mops_1d<scalar_t>(output),
                 details::torch_to_mops_2d<scalar_t>(A),
                 details::torch_to_mops_1d<scalar_t>(C),
-                details::torch_to_mops_2d<int32_t>(indices_A)
+                details::torch_to_mops_2d<int32_t>(indices_A),
+                stream
             );
         });
 
+#endif
+
     } else {
         C10_THROW_ERROR(
             ValueError,
@@ -108,6 +123,12 @@ torch::Tensor HomogeneousPolynomialEvaluationBackward::forward(
             );
         });
     } else if (A.device().is_cuda()) {
+#ifndef MOPS_CUDA_ENABLED
+        C10_THROW_ERROR(ValueError, "MOPS was not compiled with CUDA support " + A.device().str());
+#else
+        c10::cuda::CUDAGuard deviceGuard{A.device()};
+        cudaStream_t currstream = c10::cuda::getCurrentCUDAStream();
+        void* stream = reinterpret_cast<void*>(currstream);
 
         AT_DISPATCH_FLOATING_TYPES(A.scalar_type(), "homogeneous_polynomial_evaluation_vjp", [&]() {
             auto mops_grad_A = mops::Tensor<scalar_t, 2>{nullptr, {0, 0}};
@@ -121,9 +142,11 @@ torch::Tensor HomogeneousPolynomialEvaluationBackward::forward(
                 details::torch_to_mops_1d<scalar_t>(grad_output),
                 details::torch_to_mops_2d<scalar_t>(A),
                 details::torch_to_mops_1d<scalar_t>(C),
-                details::torch_to_mops_2d<int32_t>(indices_A)
+                details::torch_to_mops_2d<int32_t>(indices_A),
+                stream
             );
         });
+#endif
     } else {
         C10_THROW_ERROR(
             ValueError,

mops-torch/src/opsa.cpp

@@ -1,3 +1,8 @@
+#ifdef MOPS_CUDA_ENABLED
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+#endif
+
 #include "mops/torch/opsa.hpp"
 #include "mops/torch/utils.hpp"
 
@@ -48,6 +53,10 @@ torch::Tensor OuterProductScatterAdd::forward(
 #ifndef MOPS_CUDA_ENABLED
         C10_THROW_ERROR(ValueError, "MOPS was not compiled with CUDA support " + A.device().str());
 #else
+        c10::cuda::CUDAGuard deviceGuard{A.device()};
+        cudaStream_t currstream = c10::cuda::getCurrentCUDAStream();
+        void* stream = reinterpret_cast<void*>(currstream);
+
         output = torch::empty(
             {output_size, A.size(1), B.size(1)},
             torch::TensorOptions().dtype(A.scalar_type()).device(A.device())
@@ -58,7 +67,8 @@ torch::Tensor OuterProductScatterAdd::forward(
                 details::torch_to_mops_3d<scalar_t>(output),
                 details::torch_to_mops_2d<scalar_t>(A),
                 details::torch_to_mops_2d<scalar_t>(B),
-                details::torch_to_mops_1d<int32_t>(indices_output)
+                details::torch_to_mops_1d<int32_t>(indices_output),
+                stream
             );
         });
 
@@ -130,6 +140,10 @@ std::vector<torch::Tensor> OuterProductScatterAddBackward::forward(
 #ifndef MOPS_CUDA_ENABLED
         C10_THROW_ERROR(ValueError, "MOPS was not compiled with CUDA support " + A.device().str());
 #else
+        c10::cuda::CUDAGuard deviceGuard{A.device()};
+        cudaStream_t currstream = c10::cuda::getCurrentCUDAStream();
+        void* stream = reinterpret_cast<void*>(currstream);
+
         AT_DISPATCH_FLOATING_TYPES(A.scalar_type(), "outer_product_scatter_add_vjp", [&]() {
             auto mops_grad_A = mops::Tensor<scalar_t, 2>{nullptr, {0, 0}};
 
@@ -150,7 +164,8 @@ std::vector<torch::Tensor> OuterProductScatterAddBackward::forward(
                 details::torch_to_mops_3d<scalar_t>(grad_output),
                 details::torch_to_mops_2d<scalar_t>(A),
                 details::torch_to_mops_2d<scalar_t>(B),
-                details::torch_to_mops_1d<int32_t>(indices_output)
+                details::torch_to_mops_1d<int32_t>(indices_output),
+                stream
             );
         });
 #endif

mops-torch/src/opsaw.cpp

@@ -1,3 +1,8 @@
+#ifdef MOPS_CUDA_ENABLED
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+#endif
+
 #include "mops/torch/opsaw.hpp"
 #include "mops/torch/utils.hpp"
 

mops-torch/src/sap.cpp

@@ -1,3 +1,8 @@
+#ifdef MOPS_CUDA_ENABLED
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+#endif
+
 #include "mops/torch/sap.hpp"
 #include "mops/torch/utils.hpp"
 
@@ -59,6 +64,14 @@ torch::Tensor SparseAccumulationOfProducts::forward(
             );
         });
     } else if (A.device().is_cuda()) {
+
+#ifndef MOPS_CUDA_ENABLED
+        C10_THROW_ERROR(ValueError, "MOPS was not compiled with CUDA support " + A.device().str());
+#else
+        c10::cuda::CUDAGuard deviceGuard{A.device()};
+        cudaStream_t currstream = c10::cuda::getCurrentCUDAStream();
+        void* stream = reinterpret_cast<void*>(currstream);
+
         output = torch::empty(
             {A.size(0), output_size},
             torch::TensorOptions().dtype(A.scalar_type()).device(A.device())
@@ -72,9 +85,11 @@ torch::Tensor SparseAccumulationOfProducts::forward(
                 details::torch_to_mops_1d<scalar_t>(C),
                 details::torch_to_mops_1d<int32_t>(indices_A),
                 details::torch_to_mops_1d<int32_t>(indices_B),
-                details::torch_to_mops_1d<int32_t>(indices_output)
+                details::torch_to_mops_1d<int32_t>(indices_output),
+                stream
             );
         });
+#endif
     } else {
         C10_THROW_ERROR(
             ValueError,
@@ -170,6 +185,14 @@ std::vector<torch::Tensor> SparseAccumulationOfProductsBackward::forward(
             );
         });
     } else if (A.device().is_cuda()) {
+
+#ifndef MOPS_CUDA_ENABLED
+        C10_THROW_ERROR(ValueError, "MOPS was not compiled with CUDA support " + A.device().str());
+#else
+        c10::cuda::CUDAGuard deviceGuard{A.device()};
+        cudaStream_t currstream = c10::cuda::getCurrentCUDAStream();
+        void* stream = reinterpret_cast<void*>(currstream);
+
         AT_DISPATCH_FLOATING_TYPES(A.scalar_type(), "sparse_accumulation_of_products_vjp", [&]() {
             auto mops_grad_A = mops::Tensor<scalar_t, 2>{nullptr, {0, 0}};
             if (A.requires_grad()) {
@@ -192,9 +215,11 @@ std::vector<torch::Tensor> SparseAccumulationOfProductsBackward::forward(
                 details::torch_to_mops_1d<scalar_t>(C),
                 details::torch_to_mops_1d<int32_t>(indices_A),
                 details::torch_to_mops_1d<int32_t>(indices_B),
-                details::torch_to_mops_1d<int32_t>(indices_output)
+                details::torch_to_mops_1d<int32_t>(indices_output),
+                stream
             );
         });
+#endif
     } else {
         C10_THROW_ERROR(
             ValueError,
@@ -276,6 +301,10 @@ std::vector<torch::Tensor> SparseAccumulationOfProductsBackward::backward(
 #ifndef MOPS_CUDA_ENABLED
         C10_THROW_ERROR(ValueError, "MOPS was not compiled with CUDA support " + A.device().str());
 #else
+        c10::cuda::CUDAGuard deviceGuard{A.device()};
+        cudaStream_t currstream = c10::cuda::getCurrentCUDAStream();
+        void* stream = reinterpret_cast<void*>(currstream);
+
         AT_DISPATCH_FLOATING_TYPES(A.scalar_type(), "sparse_accumulation_of_products_vjp_vjp", [&]() {
             auto mops_grad_grad_output = mops::Tensor<scalar_t, 2>{nullptr, {0, 0}};
             if (grad_output.requires_grad()) {
@@ -317,7 +346,8 @@ std::vector<torch::Tensor> SparseAccumulationOfProductsBackward::backward(
                 details::torch_to_mops_1d<scalar_t>(C),
                 details::torch_to_mops_1d<int32_t>(indices_A),
                 details::torch_to_mops_1d<int32_t>(indices_B),
-                details::torch_to_mops_1d<int32_t>(indices_output)
+                details::torch_to_mops_1d<int32_t>(indices_output),
+                stream
             );
         });
 #endif

mops-torch/src/sasaw.cpp

@@ -1,3 +1,8 @@
+#ifdef MOPS_CUDA_ENABLED
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+#endif
+
 #include "mops/torch/sasaw.hpp"
 #include "mops/torch/utils.hpp"
 

mops/CMakeLists.txt

@@ -124,6 +124,8 @@ if(CMAKE_CUDA_COMPILER AND MOPS_CUDA)
         "src/opsa/opsa.cu"
         "src/hpe/hpe.cu"
         "src/sap/sap.cu"
+        "src/sasaw/sasaw.cu"
+        "src/opsaw/opsaw.cu"
     )
 
 endif()

mops/include/mops/hpe.h

@@ -69,15 +69,17 @@ int MOPS_EXPORT mops_cuda_homogeneous_polynomial_evaluation_f32(
     mops_tensor_1d_f32_t output,
     mops_tensor_2d_f32_t A,
     mops_tensor_1d_f32_t C,
-    mops_tensor_2d_i32_t indices_A
+    mops_tensor_2d_i32_t indices_A,
+    void* cuda_stream
 );
 
 /// CUDA version of mops::homogeneous_polynomial_evaluation for 64-bit floats
 int MOPS_EXPORT mops_cuda_homogeneous_polynomial_evaluation_f64(
     mops_tensor_1d_f64_t output,
     mops_tensor_2d_f64_t A,
     mops_tensor_1d_f64_t C,
-    mops_tensor_2d_i32_t indices_A
+    mops_tensor_2d_i32_t indices_A,
+    void* cuda_stream
 );
 
 /// CUDA version of mops::homogeneous_polynomial_evaluation_vjp for 32-bit floats
@@ -86,7 +88,8 @@ int MOPS_EXPORT mops_cuda_homogeneous_polynomial_evaluation_vjp_f32(
     mops_tensor_1d_f32_t grad_output,
     mops_tensor_2d_f32_t A,
     mops_tensor_1d_f32_t C,
-    mops_tensor_2d_i32_t indices_A
+    mops_tensor_2d_i32_t indices_A,
+    void* cuda_stream
 );
 
 /// CUDA version of mops::homogeneous_polynomial_evaluation_vjp for 64-bit floats
@@ -95,7 +98,8 @@ int MOPS_EXPORT mops_cuda_homogeneous_polynomial_evaluation_vjp_f64(
     mops_tensor_1d_f64_t grad_output,
     mops_tensor_2d_f64_t A,
     mops_tensor_1d_f64_t C,
-    mops_tensor_2d_i32_t indices_A
+    mops_tensor_2d_i32_t indices_A,
+    void* cuda_stream
 );
 
 /// CUDA version of mops::homogeneous_polynomial_evaluation_vjp_vjp for 32-bit floats
@@ -106,7 +110,8 @@ int MOPS_EXPORT mops_cuda_homogeneous_polynomial_evaluation_vjp_vjp_f32(
     mops_tensor_1d_f32_t grad_output,
     mops_tensor_2d_f32_t A,
     mops_tensor_1d_f32_t C,
-    mops_tensor_2d_i32_t indices_A
+    mops_tensor_2d_i32_t indices_A,
+    void* cuda_stream
 );
 
 /// CUDA version of mops::homogeneous_polynomial_evaluation_vjp_vjp for 64-bit floats
@@ -117,7 +122,8 @@ int MOPS_EXPORT mops_cuda_homogeneous_polynomial_evaluation_vjp_vjp_f64(
     mops_tensor_1d_f64_t grad_output,
     mops_tensor_2d_f64_t A,
     mops_tensor_1d_f64_t C,
-    mops_tensor_2d_i32_t indices_A
+    mops_tensor_2d_i32_t indices_A,
+    void* cuda_stream
 );
 
 #ifdef __cplusplus