Move stream-ordered memory allocations

Change-Id: Ief116e4e77c7579612b99cf552f3d8fc54c1d42a

internal/ceres/context_impl.cc

diff --git a/internal/ceres/context_impl.cc b/internal/ceres/context_impl.cc
index c39e098..c085644 100644
--- a/internal/ceres/context_impl.cc
+++ b/internal/ceres/context_impl.cc
@@ -171,40 +171,6 @@
   is_cuda_initialized_ = true;
   return true;
 }
-
-void* ContextImpl::CudaMalloc(size_t size, cudaStream_t stream) const {
-  void* data = nullptr;
-  // Stream-ordered alloaction API is available since CUDA 11.4, but might be
-  // not implemented by particular device
-#if CUDART_VERSION < 11040
-#warning \
-    "Stream-ordered allocations are unavailable, consider updating CUDA toolkit to version 11.4+"
-  CHECK_EQ(cudaSuccess, cudaMalloc(&data, size));
-#else
-  if (gpu_device_properties_.memoryPoolsSupported) {
-    CHECK_EQ(cudaSuccess, cudaMallocAsync(&data, size, stream));
-  } else {
-    CHECK_EQ(cudaSuccess, cudaMalloc(&data, size));
-  }
-#endif
-  return data;
-}
-
-void ContextImpl::CudaFree(void* data, cudaStream_t stream) const {
-  // Stream-ordered alloaction API is available since CUDA 11.4, but might be
-  // not implemented by particular device
-#if CUDART_VERSION < 11040
-#warning \
-    "Stream-ordered allocations are unavailable, consider updating CUDA toolkit to version 11.4+"
-  CHECK_EQ(cudaSuccess, cudaFree(data));
-#else
-  if (gpu_device_properties_.memoryPoolsSupported) {
-    CHECK_EQ(cudaSuccess, cudaFreeAsync(data, stream));
-  } else {
-    CHECK_EQ(cudaSuccess, cudaFree(data));
-  }
-#endif
-}
 #endif  // CERES_NO_CUDA
 
 ContextImpl::~ContextImpl() {

internal/ceres/context_impl.h

diff --git a/internal/ceres/context_impl.h b/internal/ceres/context_impl.h
index 6bdfd0f..508fd06 100644
--- a/internal/ceres/context_impl.h
+++ b/internal/ceres/context_impl.h
@@ -119,20 +119,6 @@
   // Returns the number of bytes of available global memory on the current CUDA
   // device. If it is called before InitCuda, it returns 0.
   size_t GpuMemoryAvailable() const;
-  // Allocate memory stream-synchronously, if cudaMallocAsync API is available
-  // in CUDA Toolkit and particular device supports it. Otherwise
-  // device-synchronous allocation is used
-  void* CudaMalloc(size_t size, cudaStream_t stream) const;
-  template <typename T>
-  T* CudaAllocate(size_t num_elements, cudaStream_t stream) const {
-    T* data = static_cast<T*>(CudaMalloc(num_elements * sizeof(T), stream));
-    return data;
-  }
-  // Free memory previously allocated by CudaMalloc. If cudaMallocAsync is
-  // supported by both CUDA Toolkit used for compilation and device used at
-  // runtime - deallocation will be performed stream-synchronously. Otherwise
-  // device-synchronous free is used.
-  void CudaFree(void* data, cudaStream_t stream) const;
 
   cusolverDnHandle_t cusolver_handle_ = nullptr;
   cublasHandle_t cublas_handle_ = nullptr;

internal/ceres/cuda_kernels_bsm_to_crs.cu.cc

diff --git a/internal/ceres/cuda_kernels_bsm_to_crs.cu.cc b/internal/ceres/cuda_kernels_bsm_to_crs.cu.cc
index 8dc68ca..05b52e7 100644
--- a/internal/ceres/cuda_kernels_bsm_to_crs.cu.cc
+++ b/internal/ceres/cuda_kernels_bsm_to_crs.cu.cc
@@ -35,7 +35,6 @@
 #include <thrust/scan.h>
 
 #include "ceres/block_structure.h"
-#include "ceres/context_impl.h"
 #include "ceres/cuda_kernels_utils.h"
 
 namespace ceres {
@@ -51,6 +50,50 @@
   return thrust::cuda::par_nosync.on(stream);
 #endif
 }
+
+void* CudaMalloc(size_t size,
+                 cudaStream_t stream,
+                 bool memory_pools_supported) {
+  void* data = nullptr;
+  // Stream-ordered alloaction API is available since CUDA 11.4, but might be
+  // not implemented by particular device
+#if CUDART_VERSION < 11040
+#warning \
+    "Stream-ordered allocations are unavailable, consider updating CUDA toolkit to version 11.4+"
+  cudaMalloc(&data, size);
+#else
+  if (memory_pools_supported) {
+    cudaMallocAsync(&data, size, stream);
+  } else {
+    cudaMalloc(&data, size);
+  }
+#endif
+  return data;
+}
+
+void CudaFree(void* data, cudaStream_t stream, bool memory_pools_supported) {
+  // Stream-ordered alloaction API is available since CUDA 11.4, but might be
+  // not implemented by particular device
+#if CUDART_VERSION < 11040
+#warning \
+    "Stream-ordered allocations are unavailable, consider updating CUDA toolkit to version 11.4+"
+  cudaSuccess, cudaFree(data);
+#else
+  if (memory_pools_supported) {
+    cudaFreeAsync(data, stream);
+  } else {
+    cudaFree(data);
+  }
+#endif
+}
+template <typename T>
+T* CudaAllocate(size_t num_elements,
+                cudaStream_t stream,
+                bool memory_pools_supported) {
+  T* data = static_cast<T*>(
+      CudaMalloc(num_elements * sizeof(T), stream, memory_pools_supported));
+  return data;
+}
 }  // namespace
 
 // Fill row block id and nnz for each row using block-sparse structure
@@ -188,10 +231,11 @@
                       int* rows,
                       int* cols,
                       cudaStream_t stream,
-                      ContextImpl* context) {
+                      bool memory_pools_supported) {
   // Set number of non-zeros per row in rows array and row to row-block map in
   // row_block_ids array
-  int* row_block_ids = context->CudaAllocate<int>(num_rows, stream);
+  int* row_block_ids =
+      CudaAllocate<int>(num_rows, stream, memory_pools_supported);
   const int num_blocks_blockwise = NumBlocksInGrid(num_row_blocks + 1);
   RowBlockIdAndNNZ<false><<<num_blocks_blockwise, kCudaBlockSize, 0, stream>>>(
       num_row_blocks,
@@ -223,7 +267,7 @@
       nullptr,
       rows,
       cols);
-  context->CudaFree(row_block_ids, stream);
+  CudaFree(row_block_ids, stream, memory_pools_supported);
 }
 
 void FillCRSStructurePartitioned(const int num_row_blocks,
@@ -240,10 +284,11 @@
                                  int* rows_f,
                                  int* cols_f,
                                  cudaStream_t stream,
-                                 ContextImpl* context) {
+                                 bool memory_pools_supported) {
   // Set number of non-zeros per row in rows array and row to row-block map in
   // row_block_ids array
-  int* row_block_ids = context->CudaAllocate<int>(num_rows, stream);
+  int* row_block_ids =
+      CudaAllocate<int>(num_rows, stream, memory_pools_supported);
   const int num_blocks_blockwise = NumBlocksInGrid(num_row_blocks + 1);
   RowBlockIdAndNNZ<true><<<num_blocks_blockwise, kCudaBlockSize, 0, stream>>>(
       num_row_blocks,
@@ -281,7 +326,7 @@
       cols_e,
       rows_f,
       cols_f);
-  context->CudaFree(row_block_ids, stream);
+  CudaFree(row_block_ids, stream, memory_pools_supported);
 }
 
 template <typename T, typename Predicate>

internal/ceres/cuda_kernels_bsm_to_crs.h

diff --git a/internal/ceres/cuda_kernels_bsm_to_crs.h b/internal/ceres/cuda_kernels_bsm_to_crs.h
index ddfecfc..27f4a25 100644
--- a/internal/ceres/cuda_kernels_bsm_to_crs.h
+++ b/internal/ceres/cuda_kernels_bsm_to_crs.h
@@ -41,7 +41,6 @@
 namespace internal {
 struct Block;
 struct Cell;
-class ContextImpl;
 
 // Compute structure of CRS matrix using block-sparse structure.
 // Arrays corresponding to CRS matrix are to be allocated by caller
@@ -54,7 +53,7 @@
                       int* rows,
                       int* cols,
                       cudaStream_t stream,
-                      ContextImpl* context);
+                      bool memory_pools_supported);
 
 // Compute structure of partitioned CRS matrix using block-sparse structure.
 // Arrays corresponding to CRS matrices are to be allocated by caller
@@ -72,7 +71,7 @@
                                  int* rows_f,
                                  int* cols_f,
                                  cudaStream_t stream,
-                                 ContextImpl* context);
+                                 bool memory_pools_supported);
 
 // Permute segment of values from block-sparse matrix with sequential layout to
 // CRS order. Segment starts at block_sparse_offset and has length of num_values