torch7/test/test_sharedmem.lua at 5630866ff6d72f7d8a79ff0392e8058d665a30b9 · torch/torch7 · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
require 'torch'
local ffi = require 'ffi'

local tester = torch.Tester()
local tests = torch.TestSuite()

local function createSharedMemStorage(name, size, storageType)
  local storageType = storageType or 'FloatStorage'
  local shmName = name or os.tmpname():gsub('/','_')
  local isShared = true
  local isSharedMem = true
  local nElements = size or torch.random(10000, 20000)
  local storage = torch[storageType](shmName, isShared, nElements, isSharedMem)
  return storage, shmName
end

local function shmFilePath(shmName)
  return (ffi.os ~= 'Windows' and '/dev/shm/' or '') .. shmName
end

local function removeShmFile(shmFileName)
  if ffi.os == 'Windows' then
    os.remove(shmFileName)
  end
end

function tests.createSharedMemFile()
  local storage, shmName = createSharedMemStorage()
  local shmFileName = shmFilePath(shmName)

  -- check that file is at /dev/shm
  tester:assert(paths.filep(shmFileName),
                'Shared memory file exists')

  -- collect storage and make sure that file is gone
  storage = nil
  collectgarbage()
  collectgarbage()
  removeShmFile(shmFileName)
  tester:assert(not paths.filep(shmFileName),
                'Shared memory file does not exists')
end

function tests.checkContents()
  local storage, shmName = createSharedMemStorage()
  local shmFileName = shmFilePath(shmName)
  local tensor = torch.FloatTensor(storage, 1, torch.LongStorage{storage:size()})
  tensor:copy(torch.rand(storage:size()))

  local sharedFile = torch.DiskFile(shmFileName, 'r'):binary()
  for i = 1, storage:size() do
    tester:assert(sharedFile:readFloat() == storage[i], 'value is not correct')
  end
  sharedFile:close()
  removeShmFile(shmFileName)
end

function tests.testSharing()
  -- since we are going to cast numbers into double (lua default)
  -- we specifically generate double storage
  local storage, shmName = createSharedMemStorage(nil, nil, 'DoubleStorage')
  local shmFileName = shmFilePath(shmName)
  local tensor = torch.DoubleTensor(storage, 1, torch.LongStorage{storage:size()})
  tensor:copy(torch.rand(storage:size()))
  local tensorCopy = tensor.new():resizeAs(tensor):copy(tensor)

  -- access the same shared memory file as regular mapping from same process
  local storage2 = torch.DoubleStorage(shmFileName, true, storage:size())
  local tensor2 = torch.DoubleTensor(storage2, 1,
                                     torch.LongStorage{storage2:size()})
  local tensor2Copy = tensor2.new():resizeAs(tensor2):copy(tensor2)

  tester:assertTensorEq(tensorCopy, tensor2Copy, 0, 'contents don\'t match')

  -- fill tensor 1 with a random value and read from 2
  local rval = torch.uniform()
  tensor:fill(rval)
  for i = 1, tensor2:size(1) do
    tester:asserteq(tensor2[i], rval, 'content is wrong')
  end

  -- fill tensor 2 with a random value and read from 1
  local rval = torch.uniform()
  tensor2:fill(rval)
  for i = 1, tensor:size(1) do
    tester:asserteq(tensor[i], rval, 'content is wrong')
  end
  removeShmFile(shmFileName)
end

function tests.readWritePointer()
  local tensor = torch.rand(2,3)
  local addr1 = torch.pointer(tensor)
  tester:assert(type(addr1)=='number')

  local f = torch.MemoryFile()
  f:binary()
  f:writePointer(addr1)
  f:seek(1)
  local addr2 = f:readPointer()
  f:close()
  tester:assert(addr1 == addr2)

  local f = torch.MemoryFile()
  f:writePointer(addr1)
  f:seek(1)
  local addr2 = f:readPointer()
  f:close()
  tester:assert(addr1 == addr2)
end

tester:add(tests)
tester:run()