Correct signatures for torch allocator plug in

Since pytorch/pytorch#91398, the signature of
the pluggable allocate and deallocate functions must accept the device
id. The current version only accepts a device id for allocate, which
means that when using a stream ordered allocator with devices other
than device zero, we pass an invalid stream into the deallocation
function. To fix this, adapt the signature to match the one pytorch
expects.

Now, since we have the device available during allocation and
deallocation, we would like to use that device to obtain the
appropriate memory resource.

Unfortunately, since RMM's cuda_device_id does not have a nullary
constructor, we can't use it in Cython without some hacky workarounds.

However, since we don't actually need to build a Python module, but
rather just a single shared library that offers two extern "C"
functions, let's just write our allocator hooks directly in C++.

- Closes #1405

python/rmm/_lib/CMakeLists.txt

@@ -12,12 +12,15 @@
 # the License.
 # =============================================================================
 
-set(cython_sources device_buffer.pyx lib.pyx logger.pyx memory_resource.pyx cuda_stream.pyx
-                   torch_allocator.pyx)
+set(cython_sources device_buffer.pyx lib.pyx logger.pyx memory_resource.pyx cuda_stream.pyx)
 set(linked_libraries rmm::rmm)
 
 # Build all of the Cython targets
 rapids_cython_create_modules(SOURCE_FILES "${cython_sources}" LINKED_LIBRARIES "${linked_libraries}"
                                                                                CXX)
-# The cdef public functions in this file need to have a C ABI
-target_compile_definitions(torch_allocator PRIVATE CYTHON_EXTERN_C=extern\ "C")
+
+
+add_library(torch_allocator SHARED _torch_allocator.cpp)
+target_link_libraries(torch_allocator PRIVATE rmm::rmm)
+cmake_path(RELATIVE_PATH CMAKE_CURRENT_SOURCE_DIR BASE_DIRECTORY "${PROJECT_SOURCE_DIR}" OUTPUT_VARIABLE _torch_allocator_location)
+install(TARGETS torch_allocator DESTINATION "${_torch_allocator_location}")

python/rmm/allocators/torch.py

@@ -16,11 +16,14 @@
 except ImportError:
     rmm_torch_allocator = None
 else:
-    import rmm._lib.torch_allocator
+    import pathlib
 
-    _alloc_free_lib_path = rmm._lib.torch_allocator.__file__
     rmm_torch_allocator = CUDAPluggableAllocator(
-        _alloc_free_lib_path,
+        pathlib.Path(__file__).absolute().parent
+        / ".."
+        / "_lib"
+        / "libtorch_allocator.so",
         alloc_fn_name="allocate",
         free_fn_name="deallocate",
     )
+    del pathlib