Replace use of RMM provided CUDA bindings with CUDA Python (rapidsai#…

…4499) Authors: - Ashwin Srinath (https://github.com/shwina) - Divye Gala (https://github.com/divyegala) Approvers: - Jordan Jacobelli (https://github.com/Ethyling) - Dante Gama Dessavre (https://github.com/dantegd) URL: rapidsai#4499
vimarsh6739 · Jan 21, 2022 · dbbac47 · dbbac47
1 parent ac4db43
commit dbbac47
Show file tree

Hide file tree

Showing 8 changed files with 20 additions and 11 deletions.
diff --git a/conda/environments/cuml_dev_cuda11.0.yml b/conda/environments/cuml_dev_cuda11.0.yml
@@ -6,6 +6,7 @@ channels:
 - conda-forge
 dependencies:
 - cudatoolkit=11.0
+- cuda-python >=11.5,<12.0
 - rapids-build-env=22.02.*
 - rapids-notebook-env=22.02.*
 - rapids-doc-env=22.02.*

diff --git a/conda/environments/cuml_dev_cuda11.2.yml b/conda/environments/cuml_dev_cuda11.2.yml
@@ -6,6 +6,7 @@ channels:
 - conda-forge
 dependencies:
 - cudatoolkit=11.2
+- cuda-python >=11.5,<12.0
 - rapids-build-env=22.02.*
 - rapids-notebook-env=22.02.*
 - rapids-doc-env=22.02.*

diff --git a/conda/environments/cuml_dev_cuda11.4.yml b/conda/environments/cuml_dev_cuda11.4.yml
@@ -6,6 +6,7 @@ channels:
 - conda-forge
 dependencies:
 - cudatoolkit=11.4
+- cuda-python >=11.5,<12.0
 - rapids-build-env=22.02.*
 - rapids-notebook-env=22.02.*
 - rapids-doc-env=22.02.*

diff --git a/conda/environments/cuml_dev_cuda11.5.yml b/conda/environments/cuml_dev_cuda11.5.yml
@@ -6,6 +6,7 @@ channels:
 - conda-forge
 dependencies:
 - cudatoolkit=11.5
+- cuda-python >=11.5,<12.0
 - rapids-build-env=22.02.*
 - rapids-notebook-env=22.02.*
 - rapids-doc-env=22.02.*

diff --git a/conda/recipes/cuml/meta.yaml b/conda/recipes/cuml/meta.yaml
@@ -37,6 +37,7 @@ requirements:
     - cudatoolkit {{ cuda_version }}.*
     - ucx-py {{ ucx_py_version }}
     - ucx-proc=*=gpu
+    - cuda-python >=11.5,<12.0
   run:
     - python x.x
     - cudf {{ minor_version }}
@@ -52,6 +53,7 @@ requirements:
     - distributed>=2021.11.1
     - joblib >=0.11
     - {{ pin_compatible('cudatoolkit', max_pin='x', min_pin='x') }}
+    - cuda-python >=11.5,<12.0
 
 tests:                                 # [linux64]
   requirements:                        # [linux64]

diff --git a/python/cuml/svm/linear.pyx b/python/cuml/svm/linear.pyx
@@ -1,4 +1,4 @@
-# Copyright (c) 2021, NVIDIA CORPORATION.
+# Copyright (c) 2021-2022, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -33,12 +33,15 @@ from cuml.raft.common.handle cimport handle_t
 from cuml.common import input_to_cuml_array
 from libc.stdint cimport uintptr_t
 from libcpp cimport bool as cppbool
-cimport rmm._lib.lib as rmm
+from cuda.ccudart cimport(
+    cudaMemcpyAsync,
+    cudaMemcpyKind,
+    cudaMemcpyDeviceToDevice
+)
+
 
 __all__ = ['LinearSVM', 'LinearSVM_defaults']
 
-cdef extern from * nogil:
-    ctypedef void* _Stream "cudaStream_t"
 
 cdef extern from "cuml/svm/linear.hpp" namespace "ML::SVM":
 
@@ -236,12 +239,12 @@ cdef class LinearSVMWrapper:
             raise AttributeError(
                 f"Expected an array of type {target.dtype}, "
                 f"but got {source.dtype}")
-        rmm.cudaMemcpyAsync(
+        cudaMemcpyAsync(
             <void*><uintptr_t>target.ptr,
             <void*><uintptr_t>source.ptr,
             <size_t>(source.nbytes),
-            rmm.cudaMemcpyDeviceToDevice,
-            <_Stream> stream)
+            cudaMemcpyKind.cudaMemcpyDeviceToDevice,
+            stream.value())
         if synchronize:
             self.handle.sync_stream()
 

diff --git a/python/cuml/test/test_base.py b/python/cuml/test/test_base.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2019-2021, NVIDIA CORPORATION.
+# Copyright (c) 2019-2022, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -40,7 +40,7 @@ def test_base_class_usage():
 
 
 def test_base_class_usage_with_handle():
-    stream = cuml.cuda.Stream()
+    stream = cuml.raft.common.cuda.Stream()
     handle = cuml.Handle(stream=stream)
     base = cuml.Base(handle=handle)
     base.handle.sync()

diff --git a/python/cuml/test/utils.py b/python/cuml/test/utils.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2020-2021, NVIDIA CORPORATION.
+# Copyright (c) 2020-2022, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -171,7 +171,7 @@ def sqnorm(x):
 def get_handle(use_handle, n_streams=0):
     if not use_handle:
         return None, None
-    s = cuml.cuda.Stream()
+    s = cuml.raft.common.cuda.Stream()
     h = cuml.Handle(stream=s, n_streams=n_streams)
     return h, s