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