Update all rmm imports to use pylibrmm/librmm

docs/cudf/source/conf.py

@@ -342,7 +342,7 @@ def clean_all_xml_files(path):
     "cudf.Series": ("cudf.core.series.Series", "cudf.Series"),
     "cudf.Index": ("cudf.core.index.Index", "cudf.Index"),
     "cupy.core.core.ndarray": ("cupy.ndarray", "cupy.ndarray"),
-    "DeviceBuffer": ("rmm._lib.device_buffer.DeviceBuffer", "rmm.DeviceBuffer"),
+    "DeviceBuffer": ("rmm.pylibrmm.device_buffer.DeviceBuffer", "rmm.DeviceBuffer"),
 }
 
 

python/cudf/cudf/_lib/column.pxd

@@ -11,7 +11,7 @@ from pylibcudf.libcudf.column.column_view cimport (
     mutable_column_view,
 )
 from pylibcudf.libcudf.types cimport size_type
-from rmm._lib.device_buffer cimport device_buffer
+from rmm.librmm.device_buffer cimport device_buffer
 
 
 cdef class Column:

python/cudf/cudf/_lib/column.pyx

@@ -28,7 +28,7 @@ from libcpp.memory cimport make_unique, unique_ptr
 from libcpp.utility cimport move
 from libcpp.vector cimport vector
 
-from rmm._lib.device_buffer cimport DeviceBuffer
+from rmm.pylibrmm.device_buffer cimport DeviceBuffer
 
 from cudf._lib.types cimport (
     dtype_from_column_view,

python/cudf/cudf/_lib/copying.pyx

@@ -8,7 +8,7 @@ from libcpp.memory cimport make_shared, shared_ptr, unique_ptr
 from libcpp.utility cimport move
 from libcpp.vector cimport vector
 
-from rmm._lib.device_buffer cimport DeviceBuffer
+from rmm.pylibrmm.device_buffer cimport DeviceBuffer
 
 import pylibcudf
 

python/cudf/cudf/_lib/scalar.pxd

@@ -4,7 +4,7 @@ from libcpp cimport bool
 from libcpp.memory cimport unique_ptr
 
 from pylibcudf.libcudf.scalar.scalar cimport scalar
-from rmm._lib.memory_resource cimport DeviceMemoryResource
+from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource
 
 
 cdef class DeviceScalar:

python/cudf/cudf/_lib/string_casting.pyx

@@ -6,7 +6,6 @@ from cudf._lib.scalar import as_device_scalar
 from cudf._lib.types import SUPPORTED_NUMPY_TO_LIBCUDF_TYPES
 
 from libcpp.memory cimport unique_ptr
-from libcpp.string cimport string
 from libcpp.utility cimport move
 
 from pylibcudf.libcudf.column.column cimport column

python/cudf/cudf/_lib/strings_udf.pyx

@@ -23,7 +23,7 @@ from pylibcudf.libcudf.strings_udf cimport (
     to_string_view_array as cpp_to_string_view_array,
     udf_string,
 )
-from rmm._lib.device_buffer cimport DeviceBuffer, device_buffer
+from rmm.pylibrmm.device_buffer cimport DeviceBuffer, device_buffer
 
 from cudf._lib.column cimport Column
 

python/cudf/cudf/core/buffer/buffer.py

@@ -284,7 +284,7 @@ def memoryview(
         """Read-only access to the buffer through host memory."""
         size = self._size if size is None else size
         host_buf = host_memory_allocation(size)
-        rmm._lib.device_buffer.copy_ptr_to_host(
+        rmm.pylibrmm.device_buffer.copy_ptr_to_host(
             self.get_ptr(mode="read") + offset, host_buf
         )
         return memoryview(host_buf).toreadonly()

python/cudf/cudf/core/buffer/spillable_buffer.py

@@ -207,7 +207,7 @@ def spill(self, target: str = "cpu") -> None:
                     domain="cudf_python-spill",
                 ):
                     host_mem = host_memory_allocation(self.size)
-                    rmm._lib.device_buffer.copy_ptr_to_host(
+                    rmm.pylibrmm.device_buffer.copy_ptr_to_host(
                         self._ptr, host_mem
                     )
                 self._ptr_desc["memoryview"] = host_mem
@@ -352,7 +352,7 @@ def memoryview(
             else:
                 assert self._ptr_desc["type"] == "gpu"
                 ret = host_memory_allocation(size)
-                rmm._lib.device_buffer.copy_ptr_to_host(
+                rmm.pylibrmm.device_buffer.copy_ptr_to_host(
                     self._ptr + offset, ret
                 )
                 return ret

python/cudf/cudf/core/udf/strings_typing.py

@@ -99,7 +99,7 @@ def prepare_args(self, ty, val, **kwargs):
             ty.dtype, (StringView, UDFString)
         ):
             return types.uint64, val.ptr if isinstance(
-                val, rmm._lib.device_buffer.DeviceBuffer
+                val, rmm.pylibrmm.device_buffer.DeviceBuffer
             ) else val.get_ptr(mode="read")
         else:
             return ty, val

python/pylibcudf/pylibcudf/column.pyx

@@ -8,7 +8,7 @@ from pylibcudf.libcudf.column.column_factories cimport make_column_from_scalar
 from pylibcudf.libcudf.scalar.scalar cimport scalar
 from pylibcudf.libcudf.types cimport size_type
 
-from rmm._lib.device_buffer cimport DeviceBuffer
+from rmm.pylibrmm.device_buffer cimport DeviceBuffer
 
 from .gpumemoryview cimport gpumemoryview
 from .scalar cimport Scalar

python/pylibcudf/pylibcudf/join.pyx

@@ -9,7 +9,7 @@ from pylibcudf.libcudf.column.column cimport column
 from pylibcudf.libcudf.table.table cimport table
 from pylibcudf.libcudf.types cimport null_equality
 
-from rmm._lib.device_buffer cimport device_buffer
+from rmm.librmm.device_buffer cimport device_buffer
 
 from .column cimport Column
 from .table cimport Table

python/pylibcudf/pylibcudf/libcudf/column/column.pxd

@@ -9,7 +9,7 @@ from pylibcudf.libcudf.column.column_view cimport (
 )
 from pylibcudf.libcudf.types cimport data_type, size_type
 
-from rmm._lib.device_buffer cimport device_buffer
+from rmm.librmm.device_buffer cimport device_buffer
 
 
 cdef extern from "cudf/column/column.hpp" namespace "cudf" nogil:

python/pylibcudf/pylibcudf/libcudf/column/column_factories.pxd

@@ -11,7 +11,7 @@ from pylibcudf.libcudf.types cimport (
     type_id,
 )
 
-from rmm._lib.device_buffer cimport device_buffer
+from rmm.librmm.device_buffer cimport device_buffer
 
 
 cdef extern from "cudf/column/column_factories.hpp" namespace "cudf" nogil:

python/pylibcudf/pylibcudf/libcudf/concatenate.pxd

@@ -6,7 +6,7 @@ from pylibcudf.libcudf.column.column cimport column, column_view
 from pylibcudf.libcudf.table.table cimport table, table_view
 from pylibcudf.libcudf.utilities.host_span cimport host_span
 
-from rmm._lib.device_buffer cimport device_buffer
+from rmm.librmm.device_buffer cimport device_buffer
 
 
 cdef extern from "cudf/concatenate.hpp" namespace "cudf" nogil:

python/pylibcudf/pylibcudf/libcudf/contiguous_split.pxd

@@ -6,7 +6,7 @@ from libcpp.vector cimport vector
 from pylibcudf.libcudf.table.table_view cimport table_view
 from pylibcudf.libcudf.types cimport size_type
 
-from rmm._lib.device_buffer cimport device_buffer
+from rmm.librmm.device_buffer cimport device_buffer
 
 
 cdef extern from "cudf/contiguous_split.hpp" namespace "cudf" nogil:

python/pylibcudf/pylibcudf/libcudf/copying.pxd

@@ -16,7 +16,7 @@ from pylibcudf.libcudf.table.table cimport table
 from pylibcudf.libcudf.table.table_view cimport table_view
 from pylibcudf.libcudf.types cimport size_type
 
-from rmm._lib.device_buffer cimport device_buffer
+from rmm.librmm.device_buffer cimport device_buffer
 
 ctypedef const scalar constscalar
 

python/pylibcudf/pylibcudf/libcudf/join.pxd

@@ -9,7 +9,7 @@ from pylibcudf.libcudf.table.table cimport table
 from pylibcudf.libcudf.table.table_view cimport table_view
 from pylibcudf.libcudf.types cimport null_equality, size_type
 
-from rmm._lib.device_uvector cimport device_uvector
+from rmm.librmm.device_uvector cimport device_uvector
 
 ctypedef unique_ptr[device_uvector[size_type]] gather_map_type
 ctypedef pair[gather_map_type, gather_map_type] gather_map_pair_type

python/pylibcudf/pylibcudf/libcudf/null_mask.pxd

@@ -6,7 +6,7 @@ from pylibcudf.libcudf.column.column_view cimport column_view
 from pylibcudf.libcudf.table.table_view cimport table_view
 from pylibcudf.libcudf.types cimport bitmask_type, mask_state, size_type
 
-from rmm._lib.device_buffer cimport device_buffer
+from rmm.librmm.device_buffer cimport device_buffer
 
 
 cdef extern from "cudf/null_mask.hpp" namespace "cudf" nogil:

python/pylibcudf/pylibcudf/libcudf/strings_udf.pxd

@@ -8,7 +8,7 @@ from pylibcudf.libcudf.column.column cimport column
 from pylibcudf.libcudf.column.column_view cimport column_view
 from pylibcudf.libcudf.types cimport size_type
 
-from rmm._lib.device_buffer cimport DeviceBuffer, device_buffer
+from rmm.pylibrmm.device_buffer cimport DeviceBuffer, device_buffer
 
 
 cdef extern from "cudf/strings/udf/udf_string.hpp" namespace \

python/pylibcudf/pylibcudf/libcudf/transform.pxd

@@ -11,7 +11,7 @@ from pylibcudf.libcudf.table.table cimport table
 from pylibcudf.libcudf.table.table_view cimport table_view
 from pylibcudf.libcudf.types cimport bitmask_type, data_type, size_type
 
-from rmm._lib.device_buffer cimport device_buffer
+from rmm.librmm.device_buffer cimport device_buffer
 
 
 cdef extern from "cudf/transform.hpp" namespace "cudf" nogil:

python/pylibcudf/pylibcudf/null_mask.pxd

@@ -2,7 +2,7 @@
 
 from pylibcudf.libcudf.types cimport mask_state, size_type
 
-from rmm._lib.device_buffer cimport DeviceBuffer
+from rmm.pylibrmm.device_buffer cimport DeviceBuffer
 
 from .column cimport Column
 

python/pylibcudf/pylibcudf/null_mask.pyx

@@ -6,10 +6,12 @@ from libcpp.utility cimport move
 from pylibcudf.libcudf cimport null_mask as cpp_null_mask
 from pylibcudf.libcudf.types cimport mask_state, size_type
 
-from rmm._lib.device_buffer cimport DeviceBuffer, device_buffer
+from rmm.pylibrmm.device_buffer cimport DeviceBuffer, device_buffer
 
 from pylibcudf.libcudf.types import mask_state as MaskState  # no-cython-lint
 
+import rmm  # no-cython-lint
+
 from .column cimport Column
 from .table cimport Table
 
@@ -19,7 +21,7 @@ cdef DeviceBuffer buffer_to_python(device_buffer buf):
 
 
 cpdef DeviceBuffer copy_bitmask(Column col):
-    """Copies ``col``'s bitmask into a ``DeviceBuffer``.
+    """Copies ``col``'s bitmask into a ``rmm.DeviceBuffer``.
 
     For details, see :cpp:func:`copy_bitmask`.
 
@@ -31,8 +33,8 @@ cpdef DeviceBuffer copy_bitmask(Column col):
     Returns
     -------
     rmm.DeviceBuffer
-        A ``DeviceBuffer`` containing ``col``'s bitmask, or an empty ``DeviceBuffer``
-        if ``col`` is not nullable
+        A ``rmm.DeviceBuffer`` containing ``col``'s bitmask, or an empty
+        ``rmm.DeviceBuffer`` if ``col`` is not nullable
     """
     cdef device_buffer db
 
@@ -66,7 +68,7 @@ cpdef DeviceBuffer create_null_mask(
     size_type size,
     mask_state state = mask_state.UNINITIALIZED
 ):
-    """Creates a ``DeviceBuffer`` for use as a null value indicator bitmask of a
+    """Creates a ``rmm.DeviceBuffer`` for use as a null value indicator bitmask of a
     ``Column``.
 
     For details, see :cpp:func:`create_null_mask`.
@@ -83,7 +85,7 @@ cpdef DeviceBuffer create_null_mask(
     Returns
     -------
     rmm.DeviceBuffer
-        A ``DeviceBuffer`` for use as a null bitmask satisfying the desired size and
+        A ``rmm.DeviceBuffer`` for use as a null bitmask satisfying the desired size and
         state
     """
     cdef device_buffer db

python/pylibcudf/pylibcudf/scalar.pxd

@@ -4,7 +4,7 @@ from libcpp cimport bool
 from libcpp.memory cimport unique_ptr
 from pylibcudf.libcudf.scalar.scalar cimport scalar
 
-from rmm._lib.memory_resource cimport DeviceMemoryResource
+from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource
 
 from .column cimport Column
 from .types cimport DataType

python/pylibcudf/pylibcudf/scalar.pyx

@@ -6,7 +6,7 @@ from libcpp.utility cimport move
 from pylibcudf.libcudf.scalar.scalar cimport scalar
 from pylibcudf.libcudf.scalar.scalar_factories cimport make_empty_scalar_like
 
-from rmm._lib.memory_resource cimport get_current_device_resource
+from rmm.pylibrmm.memory_resource cimport get_current_device_resource
 
 from .column cimport Column
 from .types cimport DataType

python/pylibcudf/pylibcudf/transform.pyx

@@ -9,7 +9,7 @@ from pylibcudf.libcudf.table.table cimport table
 from pylibcudf.libcudf.table.table_view cimport table_view
 from pylibcudf.libcudf.types cimport bitmask_type, size_type
 
-from rmm._lib.device_buffer cimport DeviceBuffer, device_buffer
+from rmm.pylibrmm.device_buffer cimport DeviceBuffer, device_buffer
 
 from .column cimport Column
 from .gpumemoryview cimport gpumemoryview