未验证 提交 bdac244e 编辑于 作者: X Wang's avatar X Wang 提交者: GitHub
浏览文件

use cuda caching allocator from pytorch (#1180)

上级 2a559c51
......@@ -4,3 +4,5 @@ build
docs/build
*~
__pycache__
*.so
.vscode
\ No newline at end of file
#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <THC/THCNumerics.cuh>
#include <c10/cuda/CUDACachingAllocator.h>
#include "THC/THC.h"
......@@ -26,23 +27,20 @@ static size_t round_up_to_multiple(size_t x, int multiple) {
return ((x + multiple - 1) / multiple) * multiple;
}
// TODO: Stop manually allocating CUDA memory; allocate an ATen byte
// tensor instead.
struct Workspace {
Workspace(size_t size) : size(size), data(NULL) {
data = THCudaMalloc(at::globalContext().lazyInitCUDA(), size);
auto& allocator = *::c10::cuda::CUDACachingAllocator::get();
dataPtr = allocator.allocate(size);
data = dataPtr.get();
}
Workspace(const Workspace&) = delete;
Workspace(Workspace&&) = default;
Workspace& operator=(Workspace&&) = default;
~Workspace() {
if (data) {
THCudaFree(at::globalContext().lazyInitCUDA(), data);
}
}
~Workspace() = default;
size_t size;
void* data;
c10::DataPtr dataPtr;
};
// Return {y}
......
#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <THC/THCNumerics.cuh>
#include <c10/cuda/CUDACachingAllocator.h>
#include "THC/THC.h"
......@@ -27,23 +28,20 @@ static size_t round_up_to_multiple(size_t x, int multiple) {
return ((x + multiple - 1) / multiple) * multiple;
}
// TODO: Stop manually allocating CUDA memory; allocate an ATen byte
// tensor instead.
struct Workspace {
Workspace(size_t size) : size(size), data(NULL) {
data = THCudaMalloc(at::globalContext().lazyInitCUDA(), size);
auto& allocator = *::c10::cuda::CUDACachingAllocator::get();
dataPtr = allocator.allocate(size);
data = dataPtr.get();
}
Workspace(const Workspace&) = delete;
Workspace(Workspace&&) = default;
Workspace& operator=(Workspace&&) = default;
~Workspace() {
if (data) {
THCudaFree(at::globalContext().lazyInitCUDA(), data);
}
}
~Workspace() = default;
size_t size;
void* data;
c10::DataPtr dataPtr;
};
// Return {y}
......
......@@ -78,7 +78,6 @@
#include <THC/THC.h>
#include <THC/THCGeneral.h>
#include <THC/THCThrustAllocator.cuh>
#include "type_shim.h"
#include "compat.h"
......
Supports Markdown
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册