Extract GPU code out of SyncedMemoryTest::TestCPUWrite

diff --git a/src/caffe/test/test_syncedmem.cpp b/src/caffe/test/test_syncedmem.cpp
index cd74758..d0d6c31 100644 (file)
--- a/src/caffe/test/test_syncedmem.cpp
+++ b/src/caffe/test/test_syncedmem.cpp
@@ -39,6 +39,20 @@ TEST_F(SyncedMemoryTest, TestCPUWrite) {
   for (int i = 0; i < mem.size(); ++i) {
     EXPECT_EQ((reinterpret_cast<char*>(cpu_data))[i], 1);
   }
+  // do another round
+  cpu_data = mem.mutable_cpu_data();
+  EXPECT_EQ(mem.head(), SyncedMemory::HEAD_AT_CPU);
+  memset(cpu_data, 2, mem.size());
+  for (int i = 0; i < mem.size(); ++i) {
+    EXPECT_EQ((reinterpret_cast<char*>(cpu_data))[i], 2);
+  }
+}
+
+TEST_F(SyncedMemoryTest, TestGPURead) {
+  SyncedMemory mem(10);
+  void* cpu_data = mem.mutable_cpu_data();
+  EXPECT_EQ(mem.head(), SyncedMemory::HEAD_AT_CPU);
+  memset(cpu_data, 1, mem.size());
   const void* gpu_data = mem.gpu_data();
   EXPECT_EQ(mem.head(), SyncedMemory::SYNCED);
   // check if values are the same
@@ -52,9 +66,6 @@ TEST_F(SyncedMemoryTest, TestCPUWrite) {
   cpu_data = mem.mutable_cpu_data();
   EXPECT_EQ(mem.head(), SyncedMemory::HEAD_AT_CPU);
   memset(cpu_data, 2, mem.size());
-  for (int i = 0; i < mem.size(); ++i) {
-    EXPECT_EQ((reinterpret_cast<char*>(cpu_data))[i], 2);
-  }
   gpu_data = mem.gpu_data();
   EXPECT_EQ(mem.head(), SyncedMemory::SYNCED);
   // check if values are the same