src/caffe/test/test_syncedmem.cpp

   1 // Copyright 2013 Yangqing Jia
   2
   3 #include <cstring>
   4 #include <cuda_runtime.h>
   5
   6 #include "gtest/gtest.h"
   7 #include "caffe/common.hpp"
   8 #include "caffe/syncedmem.hpp"
   9
  10 #include "caffe/test/test_caffe_main.hpp"
  11
  12 namespace caffe {
  13
  14 class SyncedMemoryTest : public ::testing::Test {};
  15
  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 }
  24
  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 }
  32
  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 }
  65
  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);
  76
  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 }
  86
  87 }