Move allocate_workspace to cuda/utils.h (#2923)

mlx/backend/cuda/cublas_utils.cpp

@@ -40,21 +40,6 @@ cublasLtMatmulPreference_t get_preference(cu::Device& device) {
   return pref.pref_;
 }
 
-void* allocate_workspace(cu::CommandEncoder& encoder, size_t workspace_size) {
-  if (workspace_size == 0) {
-    return nullptr;
-  }
-
-  // Ensure workspace is 256-byte aligned
-  int nbytes = cuda::ceil_div(workspace_size, 256) * 256;
-  array workspace(
-      cu::malloc_async(nbytes, encoder),
-      {static_cast<int>(workspace_size)},
-      int8);
-  encoder.add_temporary(workspace);
-  return gpu_ptr<void>(workspace);
-}
-
 cublasLtMatrixLayout_t create_matrix_layout(
     cudaDataType_t type,
     uint64_t rows,
@@ -193,8 +178,7 @@ void CublasMatmulBase::execute_matmul(
     }
   }
 
-  void* workspace_ptr =
-      cublas_utils::allocate_workspace(encoder, heuristic_.workspaceSize);
+  void* workspace_ptr = allocate_workspace(encoder, heuristic_.workspaceSize);
 
   // Execute matmul
   auto capture = encoder.capture_context();

mlx/backend/cuda/cublas_utils.h

@@ -12,8 +12,6 @@ namespace cublas_utils {
 // Get the shared cublas preference for a device
 cublasLtMatmulPreference_t get_preference(cu::Device& device);
 
-void* allocate_workspace(cu::CommandEncoder& encoder, size_t workspace_size);
-
 cublasLtMatrixLayout_t create_matrix_layout(
     cudaDataType_t type,
     uint64_t rows,

mlx/backend/cuda/cudnn_utils.cpp

@@ -94,15 +94,7 @@ fe::error_t DnnGraph::encode_capturing(
 void* DnnGraph::prepare_workspace(cu::CommandEncoder& encoder) {
   int64_t workspace_size = 0;
   CHECK_CUDNN_FE_ERROR(get_workspace_size(workspace_size));
-  if (workspace_size > 0) {
-    array workspace(
-        cu::malloc_async(workspace_size, encoder),
-        {static_cast<int>(workspace_size)},
-        uint8);
-    encoder.add_temporary(workspace);
-    return gpu_ptr<void>(workspace);
-  }
-  return nullptr;
+  return allocate_workspace(encoder, workspace_size);
 }
 
 void DnnGraph::set_tensor_attrs(

mlx/backend/cuda/utils.cpp

@@ -93,4 +93,23 @@ CudaStream::CudaStream(cu::Device& device) {
   CHECK_CUDA_ERROR(cudaStreamCreateWithFlags(&handle_, cudaStreamNonBlocking));
 }
 
+void* allocate_workspace(cu::CommandEncoder& encoder, size_t workspace_size) {
+  if (workspace_size == 0) {
+    return nullptr;
+  }
+
+  // Workspace allocation should not be captured.
+#ifndef NDEBUG
+  cudaStreamCaptureStatus status;
+  CHECK_CUDA_ERROR(cudaStreamIsCapturing(encoder.stream(), &status));
+  assert(status == cudaStreamCaptureStatusNone);
+#endif
+
+  // Ensure workspace is 256-byte aligned.
+  int nbytes = cuda::ceil_div(workspace_size, 256) * 256;
+  array workspace(cu::malloc_async(nbytes, encoder), {nbytes}, int8);
+  encoder.add_temporary(workspace);
+  return gpu_ptr<void>(workspace);
+}
+
 } // namespace mlx::core

mlx/backend/cuda/utils.h

@@ -43,4 +43,7 @@ struct Dtype;
 // Convert Dtype to CUDA C++ types.
 const char* dtype_to_cuda_type(const Dtype& dtype);
 
+// Allocate an empty array and add it as temporary.
+void* allocate_workspace(cu::CommandEncoder& encoder, size_t workspace_size);
+
 } // namespace mlx::core