1 #include <cstring>
2 #include <cuda_runtime.h>
4 #include "gtest/gtest.h"
5 #include "caffe/common.hpp"
6 #include "caffe/syncedmem.hpp"
8 #include "caffe/test/test_caffe_main.hpp"
10 namespace caffe {
12 class SyncedMemoryTest : public ::testing::Test {};
14 TEST_F(SyncedMemoryTest, TestInitialization) {
15 SyncedMemory mem(10);
16 EXPECT_EQ(mem.head(), SyncedMemory::UNINITIALIZED);
17 EXPECT_EQ(mem.size(), 10);
18 SyncedMemory* p_mem = new SyncedMemory(10 * sizeof(float));
19 EXPECT_EQ(p_mem->size(), 10 * sizeof(float));
20 delete p_mem;
21 }
23 TEST_F(SyncedMemoryTest, TestAllocation) {
24 SyncedMemory mem(10);
25 EXPECT_TRUE(mem.cpu_data());
26 EXPECT_TRUE(mem.gpu_data());
27 EXPECT_TRUE(mem.mutable_cpu_data());
28 EXPECT_TRUE(mem.mutable_gpu_data());
29 }
31 TEST_F(SyncedMemoryTest, TestCPUWrite) {
32 SyncedMemory mem(10);
33 void* cpu_data = mem.mutable_cpu_data();
34 EXPECT_EQ(mem.head(), SyncedMemory::HEAD_AT_CPU);
35 memset(cpu_data, 1, mem.size());
36 for (int i = 0; i < mem.size(); ++i) {
37 EXPECT_EQ(((char*)cpu_data)[i], 1);
38 }
39 const void* gpu_data = mem.gpu_data();
40 EXPECT_EQ(mem.head(), SyncedMemory::SYNCED);
41 // check if values are the same
42 char* recovered_value = new char[10];
43 cudaMemcpy((void*)recovered_value, gpu_data, 10, cudaMemcpyDeviceToHost);
44 for (int i = 0; i < mem.size(); ++i) {
45 EXPECT_EQ(((char*)recovered_value)[i], 1);
46 }
47 // do another round
48 cpu_data = mem.mutable_cpu_data();
49 EXPECT_EQ(mem.head(), SyncedMemory::HEAD_AT_CPU);
50 memset(cpu_data, 2, mem.size());
51 for (int i = 0; i < mem.size(); ++i) {
52 EXPECT_EQ(((char*)cpu_data)[i], 2);
53 }
54 gpu_data = mem.gpu_data();
55 EXPECT_EQ(mem.head(), SyncedMemory::SYNCED);
56 // check if values are the same
57 cudaMemcpy((void*)recovered_value, gpu_data, 10, cudaMemcpyDeviceToHost);
58 for (int i = 0; i < mem.size(); ++i) {
59 EXPECT_EQ(((char*)recovered_value)[i], 2);
60 }
61 delete[] recovered_value;
62 }
64 TEST_F(SyncedMemoryTest, TestGPUWrite) {
65 SyncedMemory mem(10);
66 void* gpu_data = mem.mutable_gpu_data();
67 EXPECT_EQ(mem.head(), SyncedMemory::HEAD_AT_GPU);
68 CUDA_CHECK(cudaMemset(gpu_data, 1, mem.size()));
69 const void* cpu_data = mem.cpu_data();
70 for (int i = 0; i < mem.size(); ++i) {
71 EXPECT_EQ(((char*)cpu_data)[i], 1);
72 }
73 EXPECT_EQ(mem.head(), SyncedMemory::SYNCED);
75 gpu_data = mem.mutable_gpu_data();
76 EXPECT_EQ(mem.head(), SyncedMemory::HEAD_AT_GPU);
77 CUDA_CHECK(cudaMemset(gpu_data, 2, mem.size()));
78 cpu_data = mem.cpu_data();
79 for (int i = 0; i < mem.size(); ++i) {
80 EXPECT_EQ(((char*)cpu_data)[i], 2);
81 }
82 EXPECT_EQ(mem.head(), SyncedMemory::SYNCED);
83 }
85 }