forked from microsoft/Multiverso
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathMatrixTableHandler.lua
More file actions
90 lines (83 loc) · 3.45 KB
/
MatrixTableHandler.lua
File metadata and controls
90 lines (83 loc) · 3.45 KB
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
local ffi = require 'ffi'
local util = require('multiverso.util')
local tbh = torch.class('MatrixTableHandler')
ffi.cdef[[
void MV_NewMatrixTable(int num_row, int num_col, TableHandler* out);
void MV_GetMatrixTableAll(TableHandler handler, float* data, int size);
void MV_AddMatrixTableAll(TableHandler handler, float* data, int size);
void MV_AddAsyncMatrixTableAll(TableHandler handler, float* data, int size);
void MV_GetMatrixTableByRows(TableHandler handler, float* data, int size, int row_ids[], int row_ids_n);
void MV_AddMatrixTableByRows(TableHandler handler, float* data, int size, int row_ids[], int row_ids_n);
void MV_AddAsyncMatrixTableByRows(TableHandler handler, float* data, int size, int row_ids[], int row_ids_n);
]]
function tbh:new(num_row, num_col, init_value)
tbh = {}
num_row = num_row or 0
num_col = num_col or 0
setmetatable(tbh, self)
self.__index = self
tbh._handler = ffi.new("TableHandler[1]")
tbh._num_row = ffi.new("int", num_row)
tbh._num_col = ffi.new("int", num_col)
tbh._size = ffi.new("int", num_row * num_col)
libmv.MV_NewMatrixTable(
tbh._num_row,
tbh._num_col,
tbh._handler
)
local init = require 'multiverso.init'
if init_value ~= nil then
init_value = init_value:float()
-- sync add is used because we want to make sure that the initial value
-- has taken effect when the call returns. No matter whether it is
-- master worker, we should call add to make sure it works in sync
-- mode
if init.worker_id() == 0 then
self.add(tbh, init_value, nil, true)
else
self.add(tbh, init_value:clone():zero(), nil, true)
end
end
return tbh
end
function tbh:get(row_ids)
if row_ids == nil then
cdata = ffi.new("float[?]", self._size)
libmv.MV_GetMatrixTableAll(self._handler[0], cdata, self._size)
data = util.cdata2tensor(cdata, tonumber(self._size))
return torch.reshape(data, tonumber(self._num_row), tonumber(self._num_col))
else
cdata = ffi.new("float[?]", #row_ids * self._num_col)
crow_ids = util.tensor2cdata(row_ids, 'int')
crow_ids_n = ffi.new("int", #row_ids)
libmv.MV_GetMatrixTableByRows(self._handler[0], cdata,
crow_ids_n * self._num_col,
crow_ids, crow_ids_n)
data = util.cdata2tensor(cdata, tonumber(#row_ids * self._num_col))
return torch.reshape(data, #row_ids, tonumber(self._num_col))
end
end
function tbh:add(data, row_ids, sync)
sync = sync or false
cdata = util.tensor2cdata(data)
if row_ids == nil then
if sync then
libmv.MV_AddMatrixTableAll(self._handler[0], cdata, self._size)
else
libmv.MV_AddAsyncMatrixTableAll(self._handler[0], cdata, self._size)
end
else
crow_ids = util.tensor2cdata(row_ids, 'int')
crow_ids_n = ffi.new("int", #row_ids)
if sync then
libmv.MV_AddMatrixTableByRows(self._handler[0], cdata,
crow_ids_n * self._num_col,
crow_ids, crow_ids_n)
else
libmv.MV_AddAsyncMatrixTableByRows(self._handler[0], cdata,
crow_ids_n * self._num_col,
crow_ids, crow_ids_n)
end
end
end
return tbh