Commit 3d4f98e2 authored by Ayan Moitra's avatar Ayan Moitra Committed by Scott Cyphers

Handle workspace reservation when workspace size is zero (#1797)

* return nullptr when workspace size is zero+modify insert method to accept const lvalue ref

* Unit test added
parent 93a5cf92
...@@ -152,6 +152,11 @@ size_t runtime::gpu::GPUAllocator::reserve_argspace(const void* data, size_t siz ...@@ -152,6 +152,11 @@ size_t runtime::gpu::GPUAllocator::reserve_argspace(const void* data, size_t siz
size_t runtime::gpu::GPUAllocator::reserve_workspace(size_t size, bool zero_initialize) size_t runtime::gpu::GPUAllocator::reserve_workspace(size_t size, bool zero_initialize)
{ {
if (size == 0)
{
return m_manager->m_primitive_emitter->insert([]() { return nullptr; });
}
size_t offset = m_manager->m_workspace_manager->allocate(size); size_t offset = m_manager->m_workspace_manager->allocate(size);
m_active.push(offset); m_active.push(offset);
auto local = std::prev(m_manager->m_workspace_mem.end()); auto local = std::prev(m_manager->m_workspace_mem.end());
...@@ -171,7 +176,7 @@ size_t runtime::gpu::GPUAllocator::reserve_workspace(size_t size, bool zero_init ...@@ -171,7 +176,7 @@ size_t runtime::gpu::GPUAllocator::reserve_workspace(size_t size, bool zero_init
} }
return workspace_ptr; return workspace_ptr;
}; };
return m_manager->m_primitive_emitter->insert(mem_primitive); return m_manager->m_primitive_emitter->insert(std::move(mem_primitive));
} }
void runtime::gpu::GPUAllocator::close() void runtime::gpu::GPUAllocator::close()
......
...@@ -59,7 +59,7 @@ size_t GPUPrimitiveEmitter::insert(std::unique_ptr<gpu::primitive>&& f) ...@@ -59,7 +59,7 @@ size_t GPUPrimitiveEmitter::insert(std::unique_ptr<gpu::primitive>&& f)
m_gpu_primitives.push_back(m_managed_primitives.back().get()); m_gpu_primitives.push_back(m_managed_primitives.back().get());
return m_gpu_primitives.size() - 1; return m_gpu_primitives.size() - 1;
} }
size_t GPUPrimitiveEmitter::insert(gpu::memory_primitive& f) size_t GPUPrimitiveEmitter::insert(const gpu::memory_primitive& f)
{ {
m_gpu_mem_primitives.push_back(f); m_gpu_mem_primitives.push_back(f);
return m_gpu_mem_primitives.size() - 1; return m_gpu_mem_primitives.size() - 1;
......
...@@ -43,7 +43,7 @@ namespace ngraph ...@@ -43,7 +43,7 @@ namespace ngraph
return m_gpu_mem_primitives; return m_gpu_mem_primitives;
} }
size_t insert(std::unique_ptr<gpu::primitive>&& f); size_t insert(std::unique_ptr<gpu::primitive>&& f);
size_t insert(gpu::memory_primitive& f); size_t insert(const gpu::memory_primitive& f);
size_t lookup(std::string hash); size_t lookup(std::string hash);
void cache(const std::string& hash, const size_t& index); void cache(const std::string& hash, const size_t& index);
GPUAllocator get_memory_allocator() { return m_memory_manager.build_allocator(); } GPUAllocator get_memory_allocator() { return m_memory_manager.build_allocator(); }
......
...@@ -40,6 +40,20 @@ TEST(gpu_test, memory_manager_unallocated) ...@@ -40,6 +40,20 @@ TEST(gpu_test, memory_manager_unallocated)
ASSERT_ANY_THROW(mem_primitive()); ASSERT_ANY_THROW(mem_primitive());
} }
TEST(gpu_test, memory_manager_zero_workspace)
{
runtime::gpu::GPUPrimitiveEmitter emitter;
size_t idx_null, idx_not_null;
{
auto allocator = emitter.get_memory_allocator();
idx_null = allocator.reserve_workspace(0);
idx_not_null = allocator.reserve_workspace(10);
}
emitter.allocate_primitive_memory();
EXPECT_EQ(emitter.get_memory_primitives()[idx_null](), nullptr);
EXPECT_NE(emitter.get_memory_primitives()[idx_not_null](), nullptr);
}
TEST(gpu_test, memory_manager_allocated) TEST(gpu_test, memory_manager_allocated)
{ {
runtime::gpu::GPUPrimitiveEmitter emitter; runtime::gpu::GPUPrimitiveEmitter emitter;
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment