提交 ce70df86 编写于 作者: L liaogang

Add gpu_allocator

上级 ce938ae5
---
Language: Cpp
BasedOnStyle: Google
Standard: Cpp11
...
cc_test(cpu_allocator_test SRCS cpu_allocator_test.cc)
nv_test(gpu_allocator_test SRCS gpu_allocator_test.cc)
......@@ -32,21 +32,21 @@ namespace detail {
// default, we should use CPUAllocator<staging=false>.
template <bool staging>
class CPUAllocator {
public:
public:
void* Alloc(size_t size);
void Free(void* p, size_t size);
};
template <>
class CPUAllocator<false> {
public:
public:
void* Alloc(size_t size) { return std::malloc(size); }
void Free(void* p, size_t size) { std::free(p); }
};
template <>
class CPUAllocator<true> {
public:
public:
void* Alloc(size_t size) {
void* p = std::malloc(size);
if (p == nullptr) {
......
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */
#pragma once
#include <stddef.h> // for size_t
#include <thrust/system_error.h>
#include <thrust/system/cuda/error.h>
namespace paddle {
namespace memory {
namespace detail {
inline void throw_on_error(cudaError_t e, const char* message) {
if (e) {
throw thrust::system_error(e, thrust::cuda_category(), message);
}
}
// GPUAllocator<staging=true> calls cudaHostMalloc, which returns
// pinned and locked memory as staging areas for data exchange
// between host and device. Allocates too much would reduce the
// amount of memory available to the system for paging. So, by
// default, we should use GPUAllocator<staging=false>.
template <bool staging>
class GPUAllocator {
public:
void* Alloc(size_t size);
void Free(void* p, size_t size);
};
template <>
class GPUAllocator<false> {
public:
void* Alloc(size_t size) {
void* p = 0;
cudaError_t result = cudaMalloc(&p, size);
if (result == cudaSuccess) {
return p;
}
// clear last error
cudaGetLastError();
return nullptr;
}
void Free(void* p, size_t size) {
// Purposefully allow cudaErrorCudartUnloading, because
// that is returned if you ever call cudaFree after the
// driver has already shutdown. This happens only if the
// process is terminating, in which case we don't care if
// cudaFree succeeds.
auto err = cudaFree(p);
if (err != cudaErrorCudartUnloading) {
throw_on_error(err, "cudaFree failed");
}
}
};
template <>
class GPUAllocator<true> {
public:
void* Alloc(size_t size) {
void* p = 0;
cudaError_t result = cudaMallocHost(&p, size);
if (result == cudaSuccess) {
return p;
}
// clear last error
cudaGetLastError();
return nullptr;
}
void Free(void* p, size_t size) {
throw_on_error(cudaFreeHost(p), "cudaFreeHost failed");
}
};
} // namespace detail
} // namespace memory
} // namespace paddle
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */
#include "paddle/memory/detail/gpu_allocator.h"
#include "gtest/gtest.h"
TEST(GPUAllocator, NonStaging) {
paddle::memory::detail::GPUAllocator<false> a;
void* p = a.Alloc(4096);
EXPECT_NE(p, nullptr);
a.Free(p, 4096);
}
TEST(GPUAllocator, Staging) {
paddle::memory::detail::GPUAllocator<true> a;
void* p = a.Alloc(4096);
EXPECT_NE(p, nullptr);
a.Free(p, 4096);
}
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册