mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-22 19:23:30 +00:00
11b369a864
Applies to all public headers and macros, plus many internal ones. There are still some internal things with OnnxRuntime in the name, but this fixes all public functions & macros.
71 lines
2.6 KiB
C++
71 lines
2.6 KiB
C++
// Copyright (c) Microsoft Corporation. All rights reserved.
|
|
// Licensed under the MIT License.
|
|
|
|
#include "core/framework/allocatormgr.h"
|
|
#include "test/framework/test_utils.h"
|
|
#include "gtest/gtest.h"
|
|
#include "cuda_runtime.h"
|
|
#include "core/providers/cuda/cuda_allocator.h"
|
|
|
|
namespace onnxruntime {
|
|
namespace test {
|
|
TEST(AllocatorTest, CUDAAllocatorTest) {
|
|
int cuda_device_id = 0;
|
|
DeviceAllocatorRegistrationInfo default_allocator_info({OrtMemTypeDefault,
|
|
[](int id) { return std::make_unique<CUDAAllocator>(id); }, std::numeric_limits<size_t>::max()});
|
|
|
|
auto cuda_arena = CreateAllocator(default_allocator_info, cuda_device_id);
|
|
|
|
size_t size = 1024;
|
|
|
|
EXPECT_STREQ(cuda_arena->Info().name, CUDA);
|
|
EXPECT_EQ(cuda_arena->Info().id, cuda_device_id);
|
|
EXPECT_EQ(cuda_arena->Info().mem_type, OrtMemTypeDefault);
|
|
EXPECT_EQ(cuda_arena->Info().type, OrtArenaAllocator);
|
|
|
|
//test cuda allocation
|
|
auto cuda_addr = cuda_arena->Alloc(size);
|
|
EXPECT_TRUE(cuda_addr);
|
|
|
|
DeviceAllocatorRegistrationInfo pinned_allocator_info({OrtMemTypeCPUOutput,
|
|
[](int) { return std::make_unique<CUDAPinnedAllocator>(); }, std::numeric_limits<size_t>::max()});
|
|
|
|
auto pinned_allocator = CreateAllocator(pinned_allocator_info);
|
|
|
|
EXPECT_STREQ(pinned_allocator->Info().name, CUDA_PINNED);
|
|
EXPECT_EQ(pinned_allocator->Info().id, 0);
|
|
EXPECT_EQ(pinned_allocator->Info().mem_type, OrtMemTypeCPUOutput);
|
|
EXPECT_EQ(pinned_allocator->Info().type, OrtArenaAllocator);
|
|
|
|
//test pinned allocation
|
|
auto pinned_addr = pinned_allocator->Alloc(size);
|
|
EXPECT_TRUE(pinned_addr);
|
|
|
|
const auto& cpu_arena = TestCPUExecutionProvider()->GetAllocator(0, OrtMemTypeDefault);
|
|
EXPECT_STREQ(cpu_arena->Info().name, CPU);
|
|
EXPECT_EQ(cpu_arena->Info().id, 0);
|
|
EXPECT_EQ(cpu_arena->Info().mem_type, OrtMemTypeDefault);
|
|
EXPECT_EQ(cpu_arena->Info().type, OrtArenaAllocator);
|
|
|
|
auto cpu_addr_a = cpu_arena->Alloc(size);
|
|
EXPECT_TRUE(cpu_addr_a);
|
|
auto cpu_addr_b = cpu_arena->Alloc(size);
|
|
EXPECT_TRUE(cpu_addr_b);
|
|
memset(cpu_addr_a, -1, 1024);
|
|
|
|
//test host-device memory copy
|
|
cudaMemcpy(cuda_addr, cpu_addr_a, size, cudaMemcpyHostToDevice);
|
|
cudaMemcpy(cpu_addr_b, cuda_addr, size, cudaMemcpyDeviceToHost);
|
|
EXPECT_EQ(*((int*)cpu_addr_b), -1);
|
|
|
|
cudaMemcpyAsync(pinned_addr, cuda_addr, size, cudaMemcpyDeviceToHost);
|
|
cudaDeviceSynchronize();
|
|
EXPECT_EQ(*((int*)pinned_addr), -1);
|
|
|
|
cpu_arena->Free(cpu_addr_a);
|
|
cpu_arena->Free(cpu_addr_b);
|
|
cuda_arena->Free(cuda_addr);
|
|
pinned_allocator->Free(pinned_addr);
|
|
}
|
|
} // namespace test
|
|
} // namespace onnxruntime
|