Remove mempool header in vtensor library (#26)

shengtsui · web-flow · commit c171251022b4 · 2025-01-27T15:23:38.000-08:00
diff --git a/docs/source/api/linalg/cholesky.rst b/docs/source/api/linalg/cholesky.rst
@@ -4,5 +4,5 @@ vt::linalg::cholesky
 .. toctree::
    :maxdepth: 1
 
-.. doxygenfunction:: vt::linalg::cholesky(Tensor<T, 2>& tensor, cusolverDnHandle_t handle = cuda::cusolver.get_handle())
-.. doxygenfunction:: vt::linalg::cholesky(Tensor<T, N>& tensor, cusolverDnHandle_t handle = cuda::cusolver.get_handle())
+.. doxygenfunction:: vt::linalg::cholesky(Tensor<T, 2>& tensor, cusolverDnHandle_t handle = cuda::CuSolver::get_instance().get_handle())
+.. doxygenfunction:: vt::linalg::cholesky(Tensor<T, N>& tensor, cusolverDnHandle_t handle = cuda::CuSolver::get_instance().get_handle())
diff --git a/docs/source/api/linalg/cublas.rst b/docs/source/api/linalg/cublas.rst
@@ -14,8 +14,6 @@ vt::cuda::CuBLAS
 .. doxygenclass:: vt::cuda::CuBLAS
    :members:
 
-.. doxygenvariable:: vt::cuda::cublas
-
 .. doxygenstruct:: vt::cuda::CuBLASFuncType
    :members:
 
diff --git a/docs/source/api/linalg/cusolver.rst b/docs/source/api/linalg/cusolver.rst
@@ -14,8 +14,6 @@ vt::cuda::CuSolver
 .. doxygenclass:: vt::cuda::CuSolver
    :members:
    
-.. doxygenvariable:: vt::cuda::cusolver
-
 .. doxygenstruct:: vt::cuda::CuSolverFuncType
    :members:
 
diff --git a/docs/source/api/linalg/inv.rst b/docs/source/api/linalg/inv.rst
@@ -4,5 +4,5 @@ vt::linalg::inv
 .. toctree::
    :maxdepth: 1
 
-.. doxygenfunction:: vt::linalg::inv(Tensor<T, 2>& tensor, cusolverDnHandle_t handle = cuda::cusolver.get_handle())
-.. doxygenfunction:: vt::linalg::inv(Tensor<T, N>& tensor, cublasHandle_t handle = cuda::cublas.get_handle())
+.. doxygenfunction:: vt::linalg::inv(Tensor<T, 2>& tensor, cusolverDnHandle_t handle = cuda::CuSolver::get_instance().get_handle())
+.. doxygenfunction:: vt::linalg::inv(Tensor<T, N>& tensor, cublasHandle_t handle = cuda::CuBLAS::get_instance().get_handle())
diff --git a/docs/source/api/linalg/matmul.rst b/docs/source/api/linalg/matmul.rst
@@ -4,7 +4,7 @@ vt::matmul
 .. toctree::
    :maxdepth: 1
 
-.. doxygenfunction:: vt::matmul(Tensor<T, 1>& tensor1, Tensor<T, 1>& tensor2, cublasHandle_t handle = cuda::cublas.get_handle())
-.. doxygenfunction:: vt::matmul(Tensor<T, 2>& tensor1, Tensor<T, 2>& tensor2, cublasHandle_t handle = cuda::cublas.get_handle())
-.. doxygenfunction:: vt::matmul(Tensor<T, N>& tensor1, Tensor<T, N>& tensor2, cublasHandle_t handle = cuda::cublas.get_handle())
+.. doxygenfunction:: vt::matmul(Tensor<T, 1>& tensor1, Tensor<T, 1>& tensor2, cublasHandle_t handle = cuda::CuBLAS::get_instance().get_handle())
+.. doxygenfunction:: vt::matmul(Tensor<T, 2>& tensor1, Tensor<T, 2>& tensor2, cublasHandle_t handle = cuda::CuBLAS::get_instance().get_handle())
+.. doxygenfunction:: vt::matmul(Tensor<T, N>& tensor1, Tensor<T, N>& tensor2, cublasHandle_t handle = cuda::CuBLAS::get_instance().get_handle())
 
diff --git a/docs/source/api/linalg/svd.rst b/docs/source/api/linalg/svd.rst
@@ -4,5 +4,5 @@ vt::linalg::svd
 .. toctree::
    :maxdepth: 1
 
-.. doxygenfunction:: vt::linalg::svd(Tensor<T, 2>& tensor, bool full_matrices = true, bool compute_uv = true, cusolverDnHandle_t handle = cuda::cusolver.get_handle())
-.. doxygenfunction:: vt::linalg::svd(Tensor<T, N>& tensor, bool full_matrices = true, bool compute_uv = true, cusolverDnHandle_t handle = cuda::cusolver.get_handle())
+.. doxygenfunction:: vt::linalg::svd(Tensor<T, 2>& tensor, bool full_matrices = true, bool compute_uv = true, cusolverDnHandle_t handle = cuda::CuSolver::get_instance().get_handle())
+.. doxygenfunction:: vt::linalg::svd(Tensor<T, N>& tensor, bool full_matrices = true, bool compute_uv = true, cusolverDnHandle_t handle = cuda::CuSolver::get_instance().get_handle())
diff --git a/docs/source/api/random/curand.rst b/docs/source/api/random/curand.rst
@@ -47,7 +47,5 @@ vt::cuda:CuRand
 .. doxygenclass:: vt::cuda::CuRand
    :members:
 
-.. doxygenvariable:: vt::cuda::curand
-
 .. doxygenfunction:: vt::cuda::set_seed(size_t seed, curandGenerator_t gen = curand.get_handle())
 .. doxygenfunction:: vt::cuda::set_offset(size_t offset, curandGenerator_t gen = curand.get_handle())
diff --git a/docs/source/api/random/normal.rst b/docs/source/api/random/normal.rst
@@ -4,6 +4,6 @@ vt::normal
 .. toctree::
    :maxdepth: 1
 
-.. doxygenfunction:: vt::random::normal(Shape<N> shape, T mean = T{0.0}, T stddev = T{1.0}, const Order order = Order::C, curandGenerator_t gen = cuda::curand.get_handle())
-.. doxygenfunction:: vt::random::normal(size_t m, T mean = T{0.0}, T stddev = T{1.0}, const Order order = Order::C, curandGenerator_t gen = cuda::curand.get_handle())
-.. doxygenfunction:: vt::random::normal(size_t m, size_t n, T mean = T{0.0}, T stddev = T{1.0}, const Order order = Order::C, curandGenerator_t gen = cuda::curand.get_handle())
+.. doxygenfunction:: vt::random::normal(Shape<N> shape, T mean = T{0.0}, T stddev = T{1.0}, const Order order = Order::C, curandGenerator_t gen = cuda::CuRand::get_instance().get_handle())
+.. doxygenfunction:: vt::random::normal(size_t m, T mean = T{0.0}, T stddev = T{1.0}, const Order order = Order::C, curandGenerator_t gen = cuda::CuRand::get_instance().get_handle())
+.. doxygenfunction:: vt::random::normal(size_t m, size_t n, T mean = T{0.0}, T stddev = T{1.0}, const Order order = Order::C, curandGenerator_t gen = cuda::CuRand::get_instance().get_handle())
diff --git a/docs/source/api/random/rand.rst b/docs/source/api/random/rand.rst
@@ -4,6 +4,6 @@ vt::rand
 .. toctree::
    :maxdepth: 1
 
-.. doxygenfunction:: vt::random::rand(Shape<N> shape, const Order order = Order::C, curandGenerator_t gen = cuda::curand.get_handle())
-.. doxygenfunction:: vt::random::rand(size_t m, const Order order = Order::C, curandGenerator_t gen = cuda::curand.get_handle())
-.. doxygenfunction:: vt::random::rand(size_t m, size_t n, const Order order = Order::C, curandGenerator_t gen = cuda::curand.get_handle())
+.. doxygenfunction:: vt::random::rand(Shape<N> shape, const Order order = Order::C, curandGenerator_t gen = cuda::CuRand::get_instance().get_handle())
+.. doxygenfunction:: vt::random::rand(size_t m, const Order order = Order::C, curandGenerator_t gen = cuda::CuRand::get_instance().get_handle())
+.. doxygenfunction:: vt::random::rand(size_t m, size_t n, const Order order = Order::C, curandGenerator_t gen = cuda::CuRand::get_instance().get_handle())
diff --git a/docs/source/conf.py b/docs/source/conf.py
@@ -7,11 +7,9 @@
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
 import pathlib
 import sys
-import warnings
-from sphinx.deprecation import RemovedInSphinx90Warning
 
 sys.path.insert(0, pathlib.Path(__file__).parents[2].resolve().as_posix())
-warnings.filterwarnings("ignore", category=RemovedInSphinx90Warning)
+
 project = 'VTensor'
 copyright = '2024, Vorticity Inc'
 author = 'Sheng-Yang Tsui'
diff --git a/docs/source/tutorials/memory.rst b/docs/source/tutorials/memory.rst
@@ -1,7 +1,7 @@
 Memory Management
 ==============
 
-VTensor leverages the `RMM device memory pool <https://github.com/rapidsai/rmm>`_ for memory allocation by default, offering two strategies for GPU memory management:
+User could leverage the `RMM device memory pool <https://github.com/rapidsai/rmm>`_ for memory allocation. Here are some ways to manage memory allocation in your program:
 
 1. **Global Memory Pool**: This approach involves configuring the compiler options to set `-DPOOL_SIZE` to a specific percentage of the total GPU memory. An example configuration in a Bazel BUILD file demonstrates this method. Unless altered, this setting reserves 50% of each GPU device's total memory for the global memory pool. Memory allocation logs are saved to `/tmp/vtensor/memory.log`.
 
@@ -11,32 +11,48 @@ VTensor leverages the `RMM device memory pool <https://github.com/rapidsai/rmm>`
        name = "my_program",
        srcs = ["my_program.cc"],
        copts = ["-DPOOL_SIZE=90", "-DLIBCUDACXX_ENABLE_EXPERIMENTAL_MEMORY_RESOURCE"],
-       deps = ["@vtensor"],
+       deps = [
+        "//lib:vtensor",
+        "//lib/core:mempool",
+       ],
    )
 
+.. code-block:: cpp
+
+    #include "lib/core/mempool.hpp"
+    #include "lib/vtensor.hpp"
+
+   int main(){
+      // Make sure you include mempool.hpp in your program. A global memory pool is initialized for every GPU device.
+   }
+
+
 2. **Custom Memory Pool**: Alternatively, a custom memory pool can be instantiated by specifying its size and the log file's path. This pool is dedicated to the GPU device in use. 
 
 .. code-block:: python
 
    cc_binary(
        name = "my_program",
        srcs = ["my_program.cc"],
-       copts = ["-DPOOL_SIZE=90", "-DLIBCUDACXX_ENABLE_EXPERIMENTAL_MEMORY_RESOURCE", "-DDISABLE_GLOBAL_MEMPOOL=true"],
-       deps = ["@vtensor"],
+       copts = ["-DLIBCUDACXX_ENABLE_EXPERIMENTAL_MEMORY_RESOURCE"],
+       deps = [
+        "//lib:vtensor",
+        "//lib/core:rmm_utils",
+    ],
    )
 
 .. code-block:: cpp
 
-   #include <lib/vtensor.hpp>
+   #include "lib/core/rmm_utils.hpp"
 
    int main(){
        auto pool = vt::Mempool(50, "/tmp/vtensor/memory.log"); // Current GPU device's memory pool
    }
 
 .. code-block:: cpp
 
-   #include <lib/vtensor.hpp>
+   #include "lib/core/rmm_utils.hpp"
 
    int main(){
-       auto pool = vt::GlobalMempool::get_instance(); // Initialize global memory pool for every GPU device
+       auto pool = vt::GlobalMempool(50); // Initialize global memory pool for every GPU device
    }
diff --git a/examples/BUILD b/examples/BUILD
@@ -7,6 +7,7 @@ cuda_binary(
     copts = ["-DLIBCUDACXX_ENABLE_EXPERIMENTAL_MEMORY_RESOURCE"],
     deps = [
         "//lib:vtensor",
+        "//lib/core:mempool",
     ]
 )
 
@@ -17,5 +18,6 @@ cuda_binary(
     copts = ["-DLIBCUDACXX_ENABLE_EXPERIMENTAL_MEMORY_RESOURCE"],
     deps = [
         "//lib:vtensor",
+        "//lib/core:mempool",
     ],
 )
diff --git a/examples/cutensor.cc b/examples/cutensor.cc
@@ -1,4 +1,5 @@
 #include <iostream>
+#include <lib/core/mempool.hpp>
 #include <lib/vtensor.hpp>
 
 __global__ void kernel1(vt::CuTensor<float, 2> tensor) {
diff --git a/examples/monte_carlo.cc b/examples/monte_carlo.cc
@@ -1,6 +1,7 @@
 
 #include <iostream>
 #include <vector>
+#include <lib/core/mempool.hpp>
 #include <lib/vtensor.hpp>
 
 int main() {
diff --git a/lib/core/BUILD b/lib/core/BUILD
@@ -9,19 +9,28 @@ cuda_library(
         "broadcast.hpp",
         "cutensor.hpp",
         "iterator.hpp",
-        "mempool.hpp",
         "operator.hpp",
         "print.hpp",
         "slice.hpp",
         "tensor.hpp",
     ],
     deps = [
-        "//lib/core:rmm_libs",
         "@local_cuda//:thrust",
         "@rmm",
     ],
 )
 
+cuda_library(
+    name = "mempool",
+    visibility = ["//visibility:public"],
+    hdrs = ["mempool.hpp"],
+    deps = [
+        ":rmm_libs",
+        "@rmm",
+    ],
+
+)
+
 cc_library(
     name = "rmm_libs",
     visibility = ["//visibility:public"],
@@ -44,3 +53,8 @@ cuda_library(
     hdrs = ["rmm_utils.hpp"],
     deps = ["@rmm"],
 )
+
+exports_files(
+    ["rmm_utils.hpp", "rmm_utils.cc"],
+    visibility = ["//visibility:public"],
+)
diff --git a/lib/core/mempool.hpp b/lib/core/mempool.hpp
@@ -4,10 +4,6 @@
 #define POOL_SIZE 50
 #endif
 
-#ifndef DISABLE_GLOBAL_MEMPOOL
-#define DISABLE_GLOBAL_MEMPOOL false
-#endif
-
 #include "lib/core/rmm_utils.hpp"
 
 namespace vt {
@@ -17,11 +13,7 @@ namespace vt {
  */
 class GlobalMempoolInitializer {
    public:
-    GlobalMempoolInitializer() {
-        if (!DISABLE_GLOBAL_MEMPOOL) {
-            GlobalMempool::get_instance(POOL_SIZE);
-        }
-    }
+    GlobalMempoolInitializer() { GlobalMempool::get_instance(POOL_SIZE); }
 };
 
 static GlobalMempoolInitializer mempool_initializer;
diff --git a/lib/core/tests/BUILD b/lib/core/tests/BUILD
@@ -15,6 +15,7 @@ cuda_test(
     copts = ["-DLIBCUDACXX_ENABLE_EXPERIMENTAL_MEMORY_RESOURCE"],
     deps = [
         "//lib:vtensor",
+        "//lib/core:mempool",
         "@googletest//:gtest_main",
     ],
 )
diff --git a/lib/core/tests/test_mempool.cc b/lib/core/tests/test_mempool.cc
@@ -1,6 +1,6 @@
 #include <gtest/gtest.h>
 
-#include <lib/vtensor.hpp>
+#include "lib/core/mempool.hpp"
 
 TEST(GlobalMempoolTest, BasicAssertions) {
     vt::GlobalMempool& instance1 = vt::GlobalMempool::get_instance(1);
diff --git a/lib/linalg/cholesky.hpp b/lib/linalg/cholesky.hpp
@@ -19,7 +19,7 @@ namespace linalg {
  * @return Tensor: The lower triangle of the Cholesky decomposition.
  */
 template <typename T>
-Tensor<T, 2> cholesky(Tensor<T, 2>& tensor, cusolverDnHandle_t handle = cuda::cusolver.get_handle()) {
+Tensor<T, 2> cholesky(Tensor<T, 2>& tensor, cusolverDnHandle_t handle = cuda::CuSolver::get_instance().get_handle()) {
     assert(tensor.order() == vt::Order::C);
     auto [n, m] = tensor.shape();
     assert(n == m);
@@ -52,7 +52,7 @@ Tensor<T, 2> cholesky(Tensor<T, 2>& tensor, cusolverDnHandle_t handle = cuda::cu
  * @return Tensor: The lower triangle of the Cholesky decomposition.
  */
 template <typename T, size_t N>
-Tensor<T, N> cholesky(Tensor<T, N>& tensor, cusolverDnHandle_t handle = cuda::cusolver.get_handle()) {
+Tensor<T, N> cholesky(Tensor<T, N>& tensor, cusolverDnHandle_t handle = cuda::CuSolver::get_instance().get_handle()) {
     assert(tensor.order() == vt::Order::C);
     assert_at_least_3d_tensor<N>();
     auto shape = tensor.shape();
diff --git a/lib/linalg/cublas.hpp b/lib/linalg/cublas.hpp
@@ -77,9 +77,6 @@ class CuBLAS {
     CuBLASHandleT handle;
 };
 
-// This is global CuBLAS instance
-static CuBLAS& cublas = CuBLAS::get_instance();
-
 /**
  * @brief The type of the CuBLAS functions.
  *
diff --git a/lib/linalg/cusolver.hpp b/lib/linalg/cusolver.hpp
@@ -78,9 +78,6 @@ class CuSolver {
     CuSolverHandleT handle;
 };
 
-// This is global CuSolver instance
-static CuSolver& cusolver = CuSolver::get_instance();
-
 /**
  * @brief The type of the CuSolver functions.
  *
diff --git a/lib/linalg/inv.hpp b/lib/linalg/inv.hpp
@@ -22,7 +22,7 @@ namespace linalg {
  * @return Tensor: The inverse tensor object.
  */
 template <typename T>
-Tensor<T, 2> inv(Tensor<T, 2>& tensor, cusolverDnHandle_t handle = cuda::cusolver.get_handle()) {
+Tensor<T, 2> inv(Tensor<T, 2>& tensor, cusolverDnHandle_t handle = cuda::CuSolver::get_instance().get_handle()) {
     assert(tensor.order() == vt::Order::C);
     auto shape = tensor.shape();
     assert(shape[0] == shape[1]);
@@ -65,7 +65,7 @@ Tensor<T, 2> inv(Tensor<T, 2>& tensor, cusolverDnHandle_t handle = cuda::cusolve
  * @return Tensor<T, N>: The inverse tensor object.
  */
 template <typename T, size_t N>
-Tensor<T, N> inv(Tensor<T, N>& tensor, cublasHandle_t handle = cuda::cublas.get_handle()) {
+Tensor<T, N> inv(Tensor<T, N>& tensor, cublasHandle_t handle = cuda::CuBLAS::get_instance().get_handle()) {
     assert(tensor.order() == vt::Order::C);
     assert_at_least_3d_tensor<N>();
     auto shape = tensor.shape();
diff --git a/lib/linalg/matmul.hpp b/lib/linalg/matmul.hpp
@@ -18,7 +18,7 @@ namespace vt {
  * @return Tensor: The result tensor.
  */
 template <typename T>
-Tensor<T, 1> matmul(Tensor<T, 1>& tensor1, Tensor<T, 1>& tensor2, cublasHandle_t handle = cuda::cublas.get_handle()) {
+Tensor<T, 1> matmul(Tensor<T, 1>& tensor1, Tensor<T, 1>& tensor2, cublasHandle_t handle = cuda::CuBLAS::get_instance().get_handle()) {
     assert(tensor1.order() == vt::Order::C);
     assert(tensor2.order() == vt::Order::C);
     assert(tensor1.size() == tensor2.size());
@@ -40,7 +40,7 @@ Tensor<T, 1> matmul(Tensor<T, 1>& tensor1, Tensor<T, 1>& tensor2, cublasHandle_t
  * @return Tensor: The result tensor.
  */
 template <typename T>
-Tensor<T, 2> matmul(Tensor<T, 2>& tensor1, Tensor<T, 2>& tensor2, cublasHandle_t handle = cuda::cublas.get_handle()) {
+Tensor<T, 2> matmul(Tensor<T, 2>& tensor1, Tensor<T, 2>& tensor2, cublasHandle_t handle = cuda::CuBLAS::get_instance().get_handle()) {
     assert(tensor1.order() == vt::Order::C);
     assert(tensor2.order() == vt::Order::C);
     auto shape1 = tensor1.shape();
@@ -75,7 +75,7 @@ Tensor<T, 2> matmul(Tensor<T, 2>& tensor1, Tensor<T, 2>& tensor2, cublasHandle_t
  * @return Tensor: The result tensor.
  */
 template <typename T, size_t N>
-Tensor<T, N> matmul(Tensor<T, N>& tensor1, Tensor<T, N>& tensor2, cublasHandle_t handle = cuda::cublas.get_handle()) {
+Tensor<T, N> matmul(Tensor<T, N>& tensor1, Tensor<T, N>& tensor2, cublasHandle_t handle = cuda::CuBLAS::get_instance().get_handle()) {
     assert(tensor1.order() == vt::Order::C);
     assert(tensor2.order() == vt::Order::C);
     assert_at_least_3d_tensor<N>();
diff --git a/lib/linalg/svd.hpp b/lib/linalg/svd.hpp
@@ -26,7 +26,7 @@ namespace linalg {
  */
 template <typename T>
 std::tuple<Tensor<T, 2>, Tensor<T, 1>, Tensor<T, 2>> svd(Tensor<T, 2>& tensor, bool full_matrices = true, bool compute_uv = true,
-                                                         cusolverDnHandle_t handle = cuda::cusolver.get_handle()) {
+                                                         cusolverDnHandle_t handle = cuda::CuSolver::get_instance().get_handle()) {
     assert(tensor.order() == vt::Order::C);
     auto [n, m] = tensor.shape();
     Tensor<T, 2> x;
@@ -109,7 +109,7 @@ std::tuple<Tensor<T, 2>, Tensor<T, 1>, Tensor<T, 2>> svd(Tensor<T, 2>& tensor, b
  */
 template <typename T, size_t N>
 std::tuple<Tensor<T, N>, Tensor<T, N - 1>, Tensor<T, N>> svd(Tensor<T, N>& tensor, bool full_matrices = true, bool compute_uv = true,
-                                                             cusolverDnHandle_t handle = cuda::cusolver.get_handle()) {
+                                                             cusolverDnHandle_t handle = cuda::CuSolver::get_instance().get_handle()) {
     assert(tensor.order() == vt::Order::C);
     assert_at_least_3d_tensor<N>();
     auto shape = tensor.shape();
diff --git a/lib/linalg/tests/test_cublas.cc b/lib/linalg/tests/test_cublas.cc
@@ -4,13 +4,13 @@
 #include <lib/vtensor.hpp>
 
 TEST(GetGlobalCuBLASHandle, BasicAssertions) {
-    auto handle1 = vt::cuda::cublas.get_handle();
-    auto handle2 = vt::cuda::cublas.get_handle();
+    auto handle1 = vt::cuda::CuBLAS::get_instance().get_handle();
+    auto handle2 = vt::cuda::CuBLAS::get_instance().get_handle();
     EXPECT_EQ(handle1, handle2);
 }
 
 TEST(CreateNewCuBLASHandle, BasicAssertions) {
-    auto handle1 = vt::cuda::cublas.get_handle();
+    auto handle1 = vt::cuda::CuBLAS::get_instance().get_handle();
     auto handle2 = vt::cuda::create_cublas_handle();
     EXPECT_NE(handle1, *handle2.get());
 }
diff --git a/lib/linalg/tests/test_cusolver.cc b/lib/linalg/tests/test_cusolver.cc
diff --git a/lib/linalg/tests/test_inv.cc b/lib/linalg/tests/test_inv.cc
diff --git a/lib/linalg/tests/test_svd.cc b/lib/linalg/tests/test_svd.cc
diff --git a/lib/memory/asfortrantensor.hpp b/lib/memory/asfortrantensor.hpp
diff --git a/lib/random/curand.hpp b/lib/random/curand.hpp
diff --git a/lib/random/normal.hpp b/lib/random/normal.hpp
diff --git a/lib/random/rand.hpp b/lib/random/rand.hpp
diff --git a/lib/random/tests/test_curand.cc b/lib/random/tests/test_curand.cc
diff --git a/lib/random/tests/test_normal.cc b/lib/random/tests/test_normal.cc
diff --git a/lib/random/tests/test_rand.cc b/lib/random/tests/test_rand.cc
diff --git a/lib/vtensor.hpp b/lib/vtensor.hpp
