Uvm backend #1174

dllehr-amd · 2023-02-07T16:26:01Z

Add CUDAMallocManagedAllocator Backend

With the new CUDAAllocator class, we have created a new
CUDAMallocManagedAllocator, which will handle allocator requests from both
cpu and cuda device types when the backend is enabled
You can enable the backend using PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocManaged
And view inside PyTorch using torch.cuda.get_allocator_backend()

This allocator is initially rudimentary as the performance implications of a
managed allocator are still being worked out. However, the goal is to be able to
swap out the backend when running without any code change required.

With the new CUDAAllocator class, we have created a new CUDAMallocManagedAllocator, which will handle allocator requests from both cpu and cuda device types when the backend is enabled You can enable the backend using PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocManaged And view inside PyTorch using torch.cuda.get_allocator_backend() This allocator is initially rudimentary as the performance implications of a managed allocator are still being worked out. However, the goal is to be able to swap out the backend when running without any code change required.

jayfurmanek · 2023-02-07T17:14:36Z

aten/src/ATen/cuda/EmptyTensor.cpp

@@ -1,6 +1,7 @@
 #define TORCH_ASSERT_NO_OPERATORS
 #include <ATen/cuda/EmptyTensor.h>
 #include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/UvmMemoryAllocator.h>


Header no longer needed

jayfurmanek and others added 8 commits December 21, 2022 17:56

[UVM] Allocation and move enablement front end

e5a95d5

[UVM] Add Caching Managed Allocator

c19db86

[UVM] Add CUDA Hooks, isManagedPtr, tensor.is_managed()

5d2b294

[UVM] Enable UVM stats module

dcebdca

[UVM] Add Move kernel and associated functions

d1cfce1

Pytorch 2.0 updates for uvm

189be36

Clean up front end to use default DeviceAllocator

ad82c20

dllehr-amd marked this pull request as draft February 7, 2023 16:26

dllehr-amd requested review from jayfurmanek and zstreet87 February 7, 2023 16:26

jayfurmanek reviewed Feb 7, 2023

View reviewed changes

Cleanup commit

12328b5

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uvm backend #1174

Uvm backend #1174

dllehr-amd commented Feb 7, 2023

jayfurmanek Feb 7, 2023

Uvm backend #1174

Are you sure you want to change the base?

Uvm backend #1174

Conversation

dllehr-amd commented Feb 7, 2023

jayfurmanek Feb 7, 2023

Choose a reason for hiding this comment