Skip to content

Commit b2f57c4

Browse files
peterlau123peterlau123
committed
test: add unit tests for amp_buffer_manager and arena
1 parent 0ba44e0 commit b2f57c4

2 files changed

Lines changed: 684 additions & 0 deletions

File tree

‎test/source/amp_buffer_manager_test.cpp‎

Lines changed: 332 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,332 @@
1+
#include "NovaLLM/memory/amp_buffer_manager.h"
2+
#include "NovaLLM/memory/allocator.h"
3+
4+
#include <gtest/gtest.h>
5+
#include <thread>
6+
#include <vector>
7+
#include <atomic>
8+
9+
using namespace nova_llm;
10+
11+
class AMPBufferManagerTest : public ::testing::Test {
12+
protected:
13+
void SetUp() override {
14+
// Note: AMPBufferManager uses singleton pattern, tests should be careful
15+
// about global state. In a real implementation, we'd want better isolation.
16+
}
17+
18+
void TearDown() override {
19+
// Cleanup is handled by the singleton's lifetime
20+
}
21+
};
22+
23+
// Test AMPBufferManager construction and initialization
24+
TEST_F(AMPBufferManagerTest, Construction) {
25+
AMPBufferManager::Config config;
26+
config.amp_config.thread_cache_size_kb = 512;
27+
config.device_flags.set(DeviceType::CPU);
28+
29+
// Add CPU allocator
30+
config.allocators[DeviceType::CPU] =
31+
nova_llm::amp::AllocatorFactory::Create(nova_llm::amp::AllocatorType::STANDARD);
32+
33+
EXPECT_NO_THROW({
34+
AMPBufferManager manager(config);
35+
EXPECT_TRUE(manager.IsInitialized());
36+
});
37+
}
38+
39+
// Test Builder::Build method
40+
TEST_F(AMPBufferManagerTest, BuilderBuild) {
41+
AMPBufferManager::Config config;
42+
config.amp_config.thread_cache_size_kb = 512;
43+
config.device_flags.set(DeviceType::CPU);
44+
config.allocators[DeviceType::CPU] =
45+
nova_llm::amp::AllocatorFactory::Create(nova_llm::amp::AllocatorType::STANDARD);
46+
47+
auto manager = AMPBufferManager::Builder::Build(config);
48+
EXPECT_NE(manager, nullptr);
49+
EXPECT_TRUE(manager->IsInitialized());
50+
}
51+
52+
// Test basic CPU allocation
53+
TEST_F(AMPBufferManagerTest, FetchCpuSmall) {
54+
AMPBufferManager::Config config;
55+
config.amp_config.thread_cache_size_kb = 512;
56+
config.device_flags.set(DeviceType::CPU);
57+
config.allocators[DeviceType::CPU] =
58+
nova_llm::amp::AllocatorFactory::Create(nova_llm::amp::AllocatorType::STANDARD);
59+
60+
auto manager = AMPBufferManager::Builder::Build(config);
61+
62+
Buffer buffer = manager->Fetch(64, DeviceType::CPU);
63+
EXPECT_NE(buffer.data, nullptr);
64+
EXPECT_GE(buffer.size, 64);
65+
EXPECT_EQ(buffer.device_type, DeviceType::CPU);
66+
67+
manager->Put(buffer);
68+
EXPECT_EQ(buffer.data, nullptr);
69+
EXPECT_EQ(buffer.size, 0);
70+
}
71+
72+
// Test CPU allocation with different sizes
73+
TEST_F(AMPBufferManagerTest, FetchCpuVariousSizes) {
74+
AMPBufferManager::Config config;
75+
config.amp_config.thread_cache_size_kb = 512;
76+
config.device_flags.set(DeviceType::CPU);
77+
config.allocators[DeviceType::CPU] =
78+
nova_llm::amp::AllocatorFactory::Create(nova_llm::amp::AllocatorType::STANDARD);
79+
80+
auto manager = AMPBufferManager::Builder::Build(config);
81+
82+
std::vector<size_t> sizes = {1, 64, 512, 4096, 65536};
83+
84+
for (size_t size : sizes) {
85+
Buffer buffer = manager->Fetch(size, DeviceType::CPU);
86+
EXPECT_NE(buffer.data, nullptr);
87+
EXPECT_GE(buffer.size, size);
88+
EXPECT_EQ(buffer.device_type, DeviceType::CPU);
89+
90+
// Verify we can write to the memory
91+
if (buffer.data) {
92+
memset(buffer.data, 0xAA, std::min(size, buffer.size));
93+
}
94+
95+
manager->Put(buffer);
96+
}
97+
}
98+
99+
// Test zero size allocation
100+
TEST_F(AMPBufferManagerTest, FetchZeroSize) {
101+
AMPBufferManager::Config config;
102+
config.amp_config.thread_cache_size_kb = 512;
103+
config.device_flags.set(DeviceType::CPU);
104+
config.allocators[DeviceType::CPU] =
105+
nova_llm::amp::AllocatorFactory::Create(nova_llm::amp::AllocatorType::STANDARD);
106+
107+
auto manager = AMPBufferManager::Builder::Build(config);
108+
109+
Buffer buffer = manager->Fetch(0, DeviceType::CPU);
110+
EXPECT_EQ(buffer.data, nullptr);
111+
EXPECT_EQ(buffer.size, 0);
112+
EXPECT_EQ(buffer.device_type, DeviceType::CPU);
113+
}
114+
115+
// Test Put with invalid buffer
116+
TEST_F(AMPBufferManagerTest, PutInvalidBuffer) {
117+
AMPBufferManager::Config config;
118+
config.amp_config.thread_cache_size_kb = 512;
119+
config.device_flags.set(DeviceType::CPU);
120+
config.allocators[DeviceType::CPU] =
121+
nova_llm::amp::AllocatorFactory::Create(nova_llm::amp::AllocatorType::STANDARD);
122+
123+
auto manager = AMPBufferManager::Builder::Build(config);
124+
125+
Buffer invalid_buffer{nullptr, 0, DeviceType::CPU};
126+
EXPECT_NO_THROW(manager->Put(invalid_buffer));
127+
}
128+
129+
// Test multiple allocations and deallocations
130+
TEST_F(AMPBufferManagerTest, MultipleOperations) {
131+
AMPBufferManager::Config config;
132+
config.amp_config.thread_cache_size_kb = 512;
133+
config.device_flags.set(DeviceType::CPU);
134+
config.allocators[DeviceType::CPU] =
135+
nova_llm::amp::AllocatorFactory::Create(nova_llm::amp::AllocatorType::STANDARD);
136+
137+
auto manager = AMPBufferManager::Builder::Build(config);
138+
139+
const int num_operations = 100;
140+
std::vector<Buffer> buffers;
141+
142+
// Allocate buffers
143+
for (int i = 0; i < num_operations; ++i) {
144+
Buffer buffer = manager->Fetch(128, DeviceType::CPU);
145+
EXPECT_NE(buffer.data, nullptr);
146+
buffers.push_back(buffer);
147+
}
148+
149+
// Deallocate all buffers
150+
for (auto& buffer : buffers) {
151+
manager->Put(buffer);
152+
}
153+
154+
// Verify all buffers are cleared
155+
for (const auto& buffer : buffers) {
156+
EXPECT_EQ(buffer.data, nullptr);
157+
EXPECT_EQ(buffer.size, 0);
158+
}
159+
}
160+
161+
// Test concurrent access
162+
TEST_F(AMPBufferManagerTest, ConcurrentAccess) {
163+
AMPBufferManager::Config config;
164+
config.amp_config.thread_cache_size_kb = 1024;
165+
config.device_flags.set(DeviceType::CPU);
166+
config.allocators[DeviceType::CPU] =
167+
nova_llm::amp::AllocatorFactory::Create(nova_llm::amp::AllocatorType::STANDARD);
168+
169+
auto manager = AMPBufferManager::Builder::Build(config);
170+
171+
const int num_threads = 4;
172+
const int operations_per_thread = 50;
173+
174+
auto thread_func = [&manager]() {
175+
for (int i = 0; i < operations_per_thread; ++i) {
176+
Buffer buffer = manager->Fetch(256, DeviceType::CPU);
177+
EXPECT_NE(buffer.data, nullptr);
178+
EXPECT_GE(buffer.size, 256);
179+
180+
// Simulate some work
181+
std::this_thread::sleep_for(std::chrono::microseconds(10));
182+
183+
manager->Put(buffer);
184+
}
185+
};
186+
187+
std::vector<std::thread> threads;
188+
for (int i = 0; i < num_threads; ++i) {
189+
threads.emplace_back(thread_func);
190+
}
191+
192+
for (auto& thread : threads) {
193+
thread.join();
194+
}
195+
}
196+
197+
// Test GetStats functionality
198+
TEST_F(AMPBufferManagerTest, GetStats) {
199+
AMPBufferManager::Config config;
200+
config.amp_config.thread_cache_size_kb = 512;
201+
config.device_flags.set(DeviceType::CPU);
202+
config.allocators[DeviceType::CPU] =
203+
nova_llm::amp::AllocatorFactory::Create(nova_llm::amp::AllocatorType::STANDARD);
204+
205+
auto manager = AMPBufferManager::Builder::Build(config);
206+
207+
// Initially should have some stats
208+
auto initial_stats = manager->GetStats();
209+
EXPECT_GE(initial_stats.total_allocated, 0);
210+
211+
// Allocate some memory
212+
Buffer buffer = manager->Fetch(1024, DeviceType::CPU);
213+
auto after_alloc_stats = manager->GetStats();
214+
EXPECT_GE(after_alloc_stats.total_allocated, initial_stats.total_allocated);
215+
216+
manager->Put(buffer);
217+
}
218+
219+
// Test IsHealthy functionality
220+
TEST_F(AMPBufferManagerTest, IsHealthy) {
221+
AMPBufferManager::Config config;
222+
config.amp_config.thread_cache_size_kb = 512;
223+
config.device_flags.set(DeviceType::CPU);
224+
config.allocators[DeviceType::CPU] =
225+
nova_llm::amp::AllocatorFactory::Create(nova_llm::amp::AllocatorType::STANDARD);
226+
227+
auto manager = AMPBufferManager::Builder::Build(config);
228+
EXPECT_TRUE(manager->IsHealthy());
229+
}
230+
231+
// Test GetArenaRouter
232+
TEST_F(AMPBufferManagerTest, GetArenaRouter) {
233+
AMPBufferManager::Config config;
234+
config.amp_config.thread_cache_size_kb = 512;
235+
config.device_flags.set(DeviceType::CPU);
236+
config.allocators[DeviceType::CPU] =
237+
nova_llm::amp::AllocatorFactory::Create(nova_llm::amp::AllocatorType::STANDARD);
238+
239+
auto manager = AMPBufferManager::Builder::Build(config);
240+
EXPECT_NE(manager->GetArenaRouter(), nullptr);
241+
}
242+
243+
// Test different configurations
244+
TEST_F(AMPBufferManagerTest, DifferentConfigurations) {
245+
std::vector<size_t> cache_sizes = {0, 64, 512, 2048};
246+
247+
for (size_t cache_size : cache_sizes) {
248+
AMPBufferManager::Config config;
249+
config.amp_config.thread_cache_size_kb = cache_size;
250+
config.device_flags.set(DeviceType::CPU);
251+
config.allocators[DeviceType::CPU] =
252+
nova_llm::amp::AllocatorFactory::Create(nova_llm::amp::AllocatorType::STANDARD);
253+
254+
auto manager = AMPBufferManager::Builder::Build(config);
255+
EXPECT_TRUE(manager->IsInitialized());
256+
257+
// Test basic functionality
258+
Buffer buffer = manager->Fetch(128, DeviceType::CPU);
259+
EXPECT_NE(buffer.data, nullptr);
260+
manager->Put(buffer);
261+
}
262+
}
263+
264+
// Test edge cases
265+
TEST_F(AMPBufferManagerTest, EdgeCases) {
266+
AMPBufferManager::Config config;
267+
config.amp_config.thread_cache_size_kb = 512;
268+
config.device_flags.set(DeviceType::CPU);
269+
config.allocators[DeviceType::CPU] =
270+
nova_llm::amp::AllocatorFactory::Create(nova_llm::amp::AllocatorType::STANDARD);
271+
272+
auto manager = AMPBufferManager::Builder::Build(config);
273+
274+
// Test very small allocation
275+
Buffer tiny = manager->Fetch(1, DeviceType::CPU);
276+
EXPECT_NE(tiny.data, nullptr);
277+
EXPECT_GE(tiny.size, 1);
278+
manager->Put(tiny);
279+
280+
// Test larger allocation
281+
Buffer large = manager->Fetch(1024 * 1024, DeviceType::CPU); // 1MB
282+
if (large.data != nullptr) {
283+
EXPECT_GE(large.size, 1024 * 1024);
284+
manager->Put(large);
285+
}
286+
}
287+
288+
// Test buffer reuse patterns
289+
TEST_F(AMPBufferManagerTest, BufferReuse) {
290+
AMPBufferManager::Config config;
291+
config.amp_config.thread_cache_size_kb = 1024;
292+
config.device_flags.set(DeviceType::CPU);
293+
config.allocators[DeviceType::CPU] =
294+
nova_llm::amp::AllocatorFactory::Create(nova_llm::amp::AllocatorType::STANDARD);
295+
296+
auto manager = AMPBufferManager::Builder::Build(config);
297+
298+
// Allocate and deallocate same size multiple times
299+
for (int i = 0; i < 10; ++i) {
300+
Buffer buffer = manager->Fetch(256, DeviceType::CPU);
301+
EXPECT_NE(buffer.data, nullptr);
302+
303+
// Fill with pattern
304+
memset(buffer.data, static_cast<uint8_t>(i), 256);
305+
306+
manager->Put(buffer);
307+
}
308+
}
309+
310+
// Test destructor cleanup
311+
TEST_F(AMPBufferManagerTest, DestructorCleanup) {
312+
// Create manager in scope
313+
{
314+
AMPBufferManager::Config config;
315+
config.amp_config.thread_cache_size_kb = 512;
316+
config.device_flags.set(DeviceType::CPU);
317+
config.allocators[DeviceType::CPU] =
318+
nova_llm::amp::AllocatorFactory::Create(nova_llm::amp::AllocatorType::STANDARD);
319+
320+
auto manager = AMPBufferManager::Builder::Build(config);
321+
322+
// Allocate some buffers
323+
std::vector<Buffer> buffers;
324+
for (int i = 0; i < 5; ++i) {
325+
buffers.push_back(manager->Fetch(128, DeviceType::CPU));
326+
}
327+
328+
// Don't explicitly deallocate - destructor should handle cleanup
329+
}
330+
// Should not crash on destruction
331+
SUCCEED();
332+
}

0 commit comments

Comments
 (0)