tds/atomic.lua at master · torch/tds · 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
local ffi = require 'ffi'
local tds = require 'tds.env'
local C = tds.C

local mt = {}

local atomic = {}
local NULL = not jit and ffi.C.NULL or nil

function mt:inc()
  return tonumber(C.tds_atomic_inc(self))
end

function mt:get()
  return tonumber(C.tds_atomic_get(self))
end

function mt:set(value)
  return C.tds_atomic_set(self, value)
end

function atomic:__new(...)
  if C.tds_has_atomic()==0 then
    error('atomic counter not available (Torch not found)')
  end
  local self = C.tds_atomic_new()
  if self == NULL then
    error('unable to allocate atomic')
  end
  self = ffi.cast('tds_atomic_counter&', self)
  ffi.gc(self, C.tds_atomic_free)
  return self
end

function atomic:__tostring()
  return 'tds.AtomicCounter='.. tonumber(C.tds_atomic_get(self))
end

function atomic:__index(lkey)
  local method = rawget(mt, lkey)
  if method then
    return method
  else
    error('no such method')
  end
end

ffi.metatype('tds_atomic_counter', atomic)


if pcall(require, 'torch') and torch.metatype then

   function atomic:__write(f)
      f:writeLong(self:get())
   end

   function atomic:__read(f)
      self:set(f:readLong())
   end

   atomic.__factory = atomic.__new
   atomic.__version = 0

   torch.metatype('tds.AtomicCounter', atomic, 'tds_atomic_counter&')

end


local atomic_ctr = {}
setmetatable(
  atomic_ctr,
  {
    __call = atomic.__new
  }
)

tds.atomic_counter = atomic_ctr
tds.AtomicCounter = atomic_ctr

return atomic_ctr