alg.lua

--------------------------------------------------------------------------------
-- Matrix and vector algebra module.
--
-- Copyright (C) 2011-2016 Stefano Peluchetti. All rights reserved.
--------------------------------------------------------------------------------

-- TODO: Stack.
-- TODO: Use only the algorithms in OpenBLAS: support only the supported types.
-- TODO: Views: sub, row, col, diag
-- TODO: Custom allocator
-- TODO: Can bound checks be optimized?
-- TODO: Can access be optimized (contiguous memory + $** / no shifts) ?
-- TODO: Can access to BLAS functions be optimized?
-- TODO: For vectors: r = c = 0 or r = 1, c = n ? (totable then changes)
-- TODO: Minimize dimension and type checks in BLAS operations.
-- TODO: Can the request to the stack be optimized (is the pointer sunk?) ?
-- TODO: better dimension reporting.
-- TODO: Use column vectors: think of matrix vector ,multiplication.
-- TODO: Just mul() instead of mulmv() and mulmm().
-- TODO: Consider avoiding type checks via calls like x:_method_element_type().
-- TODO: Remove _new.
-- TODO: Think of removing _gem*.

-- Notes:
-- + BLAS requires contiguous memory and allows only for aliasing between inputs
-- + to decide faster way a reasonable number of benchmarks must be available, 
--   including ones that perform allocations.

local ffi  = require 'ffi'
local bit  = require 'bit'
local xsys = require 'xsys'

local STACK_BUFFER = 10e6
assert(STACK_BUFFER >= 0)
local STACK_ELEMENT = ffi.typeof('double') -- TODO: Fix casting!
local JOIN_UNROLL = 5
assert(JOIN_UNROLL > 0)

local type, setmetatable, rawequal = type, setmetatable, rawequal
local band = bit.band
local floor, ceil = math.floor, math.ceil

local template = xsys.template
local width = xsys.string.width

-- Array memory ops ------------------------------------------------------------
-- Invariant: n == r*c.
local function array_alloc(ct, n, r, c)
  local a = ffi.new(ct, n) -- Default initialization of VLS, compiled.
  a._n, a._r, a._c = n, r, c
  a._p = a._v
  return a -- VLS are automatically zero-filled for default initializer case.
end

local function array_map(ct, n, r, c, p)
  local a = ffi.new(ct, 0) -- Default initialization of VLS, compiled.
  a._n, a._r, a._c = n, r, c
  a._p = p
  return a -- VLS are automatically zero-filled for default initializer case.
end

local function array_copy_data(dest, source)
  local source_size = ffi.sizeof(source:elementct())*source._n
  ffi.copy(dest._p, source._p, source_size)
end

local function array_copy_data_offset(dest, source, offset)
  local source_size = ffi.sizeof(source:elementct())*source._n
  ffi.copy(dest._p + offset, source._p, source_size)
end

local function array_clear(x)
  local x_size = ffi.sizeof(x:elementct())*x._n
  ffi.fill(x._p, x_size)
end

-- Stack -----------------------------------------------------------------------
-- TODO: Use malloc.
-- TODO: Allow growth.

local stack_struct = 'struct { int32_t _max, _n; $ _p[?]; }'

local function mem_stack_data(self, n)
  self._n = self._n + n
  return self._p + (self._n - n)
end

local function new_mem_stack_ct(element_ct)
  local stack_mt = {
    __new = function(ct, maxsize)
      local o = ffi.new(ct, maxsize)
      o._max = maxsize
      return o
    end,
    clear = function(self)
      array_clear(self)
      self._n = 0
    end,
    request = function(self, n)
      return self._n + n <= self._max and mem_stack_data(self, n)
    end,
    elementct = function()
      return element_ct
    end,
  }
  stack_mt.__index = stack_mt 

  local stack_ct = ffi.typeof(stack_struct, element_ct) 
  return ffi.metatype(stack_ct, stack_mt)
end

local stack_element_size = ffi.sizeof(STACK_ELEMENT)
local stack_elements = STACK_BUFFER/stack_element_size
local stack = new_mem_stack_ct(STACK_ELEMENT)(stack_elements)

local function stack_array(self, n, r, c)
  local nbuff = ceil(ffi.sizeof(self:elementct())*n/stack_element_size)
  local p = stack:request(nbuff)
  return p and array_map(self, n, r, c, ffi.cast(self._p, p)) or array_alloc(self, n, r, c)
end

local function stack_clear()
  stack:clear()
end

-- BLAS ------------------------------------------------------------------------
local blas_element_code = template([[
local ffi     = require 'ffi'
local cblas_h = require 'sci._cblas_h'

ffi.cdef(cblas_h)
local blas = ffi.load('libopenblas')

local complex_a1 = ffi.typeof('complex[1]')
local compflo_a1 = ffi.typeof('complex float[1]')

local compfloa, compflob = compflo_a1(), compflo_a1()
local complexa, complexb = complex_a1(), complex_a1()

return {
| for ELEMENT_NAME, BLAS in pairs{
|   float             = { PREFIX = 's' },
|   double            = { PREFIX = 'd' },
|   ['complex float'] = { PREFIX = 'c', ALPHA = 'compfloa', BETA = 'compflob' },
|   complex           = { PREFIX = 'z', ALPHA = 'complexa', BETA = 'complexb' },
| } do
  [tonumber(ffi.typeof('${ELEMENT_NAME}'))] = {
    gemm = function(C, A, B, At, Bt, alpha, beta)
      ${BLAS.ALPHA and BLAS.ALPHA..'[0] = alpha'}
      ${BLAS.BETA and BLAS.BETA..'[0] = beta'}
      blas.cblas_${BLAS.PREFIX}gemm(
        blas.CblasRowMajor, 
        At and blas.CblasTrans or blas.CblasNoTrans, 
        Bt and blas.CblasTrans or blas.CblasNoTrans, 
        C:nrow(), 
        C:ncol(), 
        At and A:nrow() or A:ncol(), 
        ${BLAS.ALPHA and BLAS.ALPHA or 'alpha'},
        A:data(), 
        A:ncol(), 
        B:data(), 
        B:ncol(), 
        ${BLAS.BETA and BLAS.BETA or 'beta'},
        C:data(), 
        C:ncol()
      )
    end,
    gemv = function(y, A, x, At, alpha, beta)
      ${BLAS.ALPHA and BLAS.ALPHA..'[0] = alpha'}
      ${BLAS.BETA and BLAS.BETA..'[0] = beta'}
      blas.cblas_${BLAS.PREFIX}gemv(
        blas.CblasRowMajor, 
        At and blas.CblasTrans or blas.CblasNoTrans, 
        A:nrow(), 
        A:ncol(),
        ${BLAS.ALPHA and BLAS.ALPHA or 'alpha'},
        A:data(), 
        A:ncol(),
        x:data(), 
        1,
        ${BLAS.BETA and BLAS.BETA or 'beta'},
        y:data(),
        1
      )
    end,
  },
| end
}
]])()

local blas_element_ct = assert(loadstring(blas_element_code))()

local function same_type_check_2(x, y)
  if x:elementct() ~= y:elementct() then
    error('constant element type required')
  end
end

local function same_type_check_3(x, y, z)
  local ct = x:elementct()
  if ct ~= y:elementct() or ct ~= z:elementct() then
    error('constant element type required')
  end
end

local function dimensions_mat(A, At)
  local Ar, Ac = A:nrow(), A:ncol()
  if At then
    Ar, Ac = Ac, Ar
  end
  return Ar*Ac, Ar, Ac
end

local function dimensions_mat_same_check(A, At, Br, Bc)
  local _, Ar, Ac = dimensions_mat(A, At)
  if Ar ~= Br or Ac ~= Bc then
    error('matrix dimensions disagree')
  end
end

local function dimensions_mat_square_check(Ar, Ac)
  if Ar ~= Ac then
    error('square matrix expected')
  end
end

local function dimensions_mul_check_2(A, B, At, Bt)
  local _, Ar, Ac = dimensions_mat(A, At)
  local _, Br, Bc = dimensions_mat(B, Bt)
  if Ac ~= Br then
    error("incompatible dimensions in matrix-matrix multiplication")
  end
  return Ar*Bc, Ar, Bc
end

local function dimensions_mul_check_3(C, A, B, At, Bt)
  local Cn, Cr, Cc = dimensions_mul_check_2(A, B, At, Bt)
  dimensions_mat_same_check(C, false, Cr, Cc)
  return Cn, Cr, Cc
end

local function dimensions_pow_check_1(A)
  local An, Ar, Ac = dimensions_mat(A)
  dimensions_mat_square_check(Ar, Ac)
  return An, Ar, Ac
end

local function dimensions_pow_check_2(B, A)
  local An, Ar, Ac = dimensions_pow_check_1(A)
  dimensions_mat_same_check(B, false, Ar, Ac)
  return An, Ar, Ac
end

local function __mul(C, A, B, At, Bt)
  same_type_check_3(C, A, B)
  local Cn, Cr, Cc = dimensions_mat(C)
  local alias = rawequal(C, A) or rawequal(C, B)
  local T = alias and stack_array(C, Cn, Cr, Cc) or C
  if Cc == 1 then
    T:_gemv(A, B, At, 1, 0)
  else
    T:_gemm(A, B, At, Bt, 1, 0)
  end
  if alias then
    array_copy_data(C, T)
  end
end

local function mul(C, A, B, At, Bt)
  dimensions_mul_check_3(C, A, B, At, Bt)
  __mul(C, A, B, At, Bt)
  stack_clear()
end

-- Exponentiation by squaring algorithm:
local function pow_recursive(A, s, n)
  local T = stack_array(A, n*n, n, n)
  if s == 1 then
    -- Cannot return A because could generate aliasing between R and T below.
    array_copy_data(T, A)
    return T
  elseif s == 2 then
    T:_gemm(A, A, false, false, 1, 0)
    return T
  elseif band(s, 1) == 0 then -- Even.
    T:_gemm(A, A, false, false, 1, 0)
    return pow_recursive(T, s/2, n)
  else
    T:_gemm(A, A, false, false, 1, 0)
    local R = pow_recursive(T, (s - 1)/2, n) -- R cannot alias T.
    T:_gemm(R, A, false, false, 1, 0)
    return T
  end
end

local function pow_dispatch(B, A, s)
  local n = B:nrow()
  if s == 0 then
    array_clear(B)
    for i=1,n do B[{i,i}] = 1 end
  elseif s == 1 then
    array_copy_data(B, A)
  else
    local T = pow_recursive(A, s, n)
    array_copy_data(B, T)
  end
end

-- TODO: Use SVD decomposition for large s and allow positive real s.
local function __pow(B, A, s)
  same_type_check_2(B, A)
  if s < 0 or floor(s) ~= s then
    error('NYI: matrix exponentiation supported only for nonnegative integers')
  end
  pow_dispatch(B, A, s)
end

local function pow(B, A, s)
  dimensions_pow_check_2(B, A)
  __pow(B, A, s)
  stack_clear()
end

--------------------------------------------------------------------------------

local function sum(x)
  local v = 0
  for i=0,#x-1 do v = v + x._p[i] end
  return v
end

local function prod(x)
  local v = 1
  for i=0,#x-1 do v = v * x._p[i] end
  return v
end

local function trace(A)
  local _, Ar, Ac = dimensions_mat(A)
  dimensions_mat_square_check(Ar, Ac)
  local v = 0
  for i=1,Ar do
    v = v + A[{i,i}]
  end
  return v
end

-- Join ------------------------------------------------------------------------
local function rep(what, first, last, sep)
  sep = sep or ', '
  local increment = last >= first and 1 or -1
  local o = { }
  for i=first,last,increment do
    o[#o + 1] = what:gsub('@', i)
  end
  return table.concat(o, sep)
end

local concat_code = template([[
local setmetatable = setmetatable

local concat_n_mt = {
  _new = function(self, n, r, c)
    return self[1]:_new(n, r, c)
  end,
  nrow = function(self)
    return self[1]:nrow()
  end,
  ncol = function(self)
    local nc = 0
    for i=1,self[0] do
      nc = nc + self[i]:ncol()
    end
    return nc
  end,
  elementct = function(self)
    return self[1]:elementct()
  end,
  _concat_dispatch = function(self, lhs)
    local na = self[0]
    self[na + 1] = lhs
    self[0] = na + 1
    return self
  end,
  _copy_into = function(self, out, offset)
    local na, nr = self[0], self[1]._r
    for r=1,nr do
      for a=na,1,-1 do
        local nc = self[a]._c
        for c=1,nc do
          out._p[offset + c - 1] = self[a]._p[(r-1)*nc + c - 1]
        end
        offset = offset + nc
      end
    end
    return offset
  end,
}
concat_n_mt.__index = concat_n_mt

| for N=JOIN_UNROLL,2,-1 do
local concat_${N}_mt = {
  _new = function(self, n, r, c)
    return self[1]:_new(n, r, c)
  end,
  nrow = function(self)
    return self[1]:nrow()
  end,
  ncol = function(self)
    return ${R('self[@]:ncol()', 1, N, ' + ')}
  end,
  elementct = function(self)
    return self[1]:elementct()
  end,
  _concat_dispatch = function(self, lhs)
| if N == JOIN_UNROLL then
    self[0] = ${N + 1}
    return setmetatable({ ${R('self[@]', 1, N)}, lhs }, concat_n_mt)
| else
    return setmetatable({ ${R('self[@]', 1, N)}, lhs }, concat_${N + 1}_mt)
| end
  end,
  _copy_into = function(self, out, offset)
    for r=1,self[1]:nrow() do
| for I=N,1,-1 do
      local nc = self[${I}]._c
      for c=1,nc do
        out._p[offset + c - 1] = self[${I}]._p[(r-1)*nc + c - 1]
      end
      offset = offset + nc
| end
    end
    return offset
  end,
}
concat_${N}_mt.__index = concat_${N}_mt

| end
return concat_2_mt
]])({ JOIN_UNROLL = JOIN_UNROLL, R = rep })

local concat_2_mt = assert(loadstring(concat_code))()

local join_code = template([[
local select = select
local error = error

local function join_1(x1)
  local nr, nc = x1:nrow(), x1:ncol()
  local a = x1:_new(nr*nc, nr, nc)
  x1:_copy_into(a, 0)
  return a
end

| for N=2,JOIN_UNROLL do
local function join_${N}(${R('x@', 1, N)})
  local nr, nc, ct = x1:nrow(), x1:ncol(), x1:elementct()
  if ${R('x@:elementct() ~= ct', 2, N, ' or ')} then
    error('constant element type required')
  end
  if ${R('x@:ncol() ~= nc', 2, N, ' or ')} then
    error('constant number of columns required')
  end
  nr = nr + ${R('x@:nrow()', 2, N, ' + ')}
  local a = x1:_new(nr*nc, nr, nc)
  local offset = 0
| for I=1,N do
  offset = x${I}:_copy_into(a, offset)
| end
  return a
end

| end
local function join_n(n, ...)
  local arg = { ... }
  local nr, nc, ct = arg[1]:nrow(), arg[1]:ncol(), arg[1]:elementct()
  for i=2,n do
    if arg[i]:elementct() ~= ct then
      error('constant element type required')
    end
    if arg[i]:ncol() ~= nc then
      error('constant number of columns required')
    end
    nr = nr + arg[i]:nrow()
  end
  local a = arg[1]:_new(nr*nc, nr, nc)
  local offset = 0
  for i=1,n do
    offset = arg[i]:_copy_into(a, offset)
  end
  return a
end

return function(...)
  local n = select('#', ...)
  if n == 1 then
    return join_1(...)
| for I=2,JOIN_UNROLL do
  elseif n == ${I} then
    return join_${I}(...)
| end
  else
    return join_n(n, ...)
  end
end
]])({ JOIN_UNROLL = JOIN_UNROLL, R = rep })

local join = assert(loadstring(join_code))()

-- Array -----------------------------------------------------------------------
local array_struct = 'struct { int32_t _n, _r, _c; $* _p; $ _v[?]; }'

local function unsupported_element_ct(self)
  error('operation not supported for element type '..tostring(self:elementct()))
end

local function new_array_ct(element_ct, element_copy)
  local array_mt
  array_mt = {
    new = function(self)
      return array_alloc(self, self._n, self._r, self._c)
    end,
    copy = function(self)
      local a = self:new()
      array_copy_data(a, self)
      return a
    end,
    _new = function(self, n, r, c)
      return array_alloc(self, n, r, c)
    end,
    _copy_into = function(self, out, offset)
      array_copy_data_offset(out, self, offset)
      return offset + self._n
    end,
    _concat_dispatch = function(self, lhs) -- Concatenating two array_ct.
      return setmetatable({ [0] = 2, self, lhs }, concat_2_mt)
    end,
    __concat = function(lhs, rhs)
      if lhs:nrow() ~= rhs:nrow() then
        error('constant number of rows required')
      end
      same_type_check_2(lhs, rhs)
      return rhs:_concat_dispatch(lhs)
    end,
    sub = function(self, f, l)
      if f < 1 or f - 1 > l or l > self._n then
        error('out of bounds first: '..f..', last: '..l..', length: '..self._n)
      end
      if self._n ~= 0 and self._c ~= 1 then
        error('single-column array required')
      end
      local a = array_alloc(self, l - f + 1, l - f + 1, 1)
      array_copy_data_offset(a, self, f - 1)
      return a
    end,
    __len = function(self)
      return self._n
    end,
    nrow = function(self)
      return self._r
    end,
    ncol = function(self)
      return self._c
    end,
    __index = element_copy and function(self, k)
      if type(k) == 'number' then
        if k < 1 or k > self._n then
          error('out of bounds index: '..k..', length: '..self._n)
        end
        return element_copy(self._p[k-1])
      elseif type(k) == 'table' then
        local r, c = k[1], k[2]
        if r < 1 or r > self._r then
          error('out of bounds row: '..r..', number of rows: '..self._r)
        end
        if c < 1 or c > self._c then
          error('out of bounds column: '..c..', number of columns: '..self._c)
        end
        return element_copy(self._p[(r-1)*self._c + (c-1)])
      else    
        return array_mt[k]
      end
    end or function(self, k)
      if type(k) == 'number' then
        if k < 1 or k > self._n then
          error('out of bounds index: '..k..', length: '..self._n)
        end
        return self._p[k-1]
      elseif type(k) == 'table' then
        local r, c = k[1], k[2]
        if r < 1 or r > self._r then
          error('out of bounds row: '..r..', number of rows: '..self._r)
        end
        if c < 1 or c > self._c then
          error('out of bounds column: '..c..', number of columns: '..self._c)
        end
        return self._p[(r-1)*self._c + (c-1)]
      else    
        return array_mt[k]
      end
    end,
    __newindex = element_copy and function(self, k, v)
      if type(k) == 'number' then
        if k < 1 or k > self._n then
          error('out of bounds index: '..k..', length: '..self._n)
        end
        self._p[k-1] =  element_copy(v)
      elseif type(k) == 'table' then
        local r, c = k[1], k[2]
        if r < 1 or r > self._r then
          error('out of bounds row: '..r..', number of rows: '..self._r)
        end
        if c < 1 or c > self._c then
          error('out of bounds column: '..c..', number of columns: '..self._c)
        end
        self._p[(r-1)*self._c + (c-1)] = element_copy(v)
      end
    end or function(self, k, v)
      if type(k) == 'number' then
        if k < 1 or k > self._n then
          error('out of bounds index: '..k..', length: '..self._n)
        end
        self._p[k-1] = v
      elseif type(k) == 'table' then
        local r, c = k[1], k[2]
        if r < 1 or r > self._r then
          error('out of bounds row: '..r..', number of rows: '..self._r)
        end
        if c < 1 or c > self._c then
          error('out of bounds column: '..c..', number of columns: '..self._c)
        end
        self._p[(r-1)*self._c + (c-1)] = v
      end
    end,
    totable = function(self)
      local o = { }
      for i=1,self:nrow() do
        o[i] = { }
        for j=1,self:ncol() do
          o[i][j] = self[{i, j}]
        end
      end
      return o
    end,
    __tostring = function(self)
      local o = { }
      for i=1,self:nrow() do
        o[i] = { }
        for j=1,self:ncol() do
          o[i][j] = width(self[{i, j}])
        end
        o[i] = table.concat(o[i], ",")
      end
      return table.concat(o, "\n")
    end,
    elementct = function()
      return element_ct
    end,
    data = function(self)
      return self._p
    end,
  }

  local element_ct_id = tonumber(element_ct)

  local blas_algo = blas_element_ct[element_ct_id]
  if blas_algo then
    array_mt._gemm = blas_algo.gemm
    array_mt._gemv = blas_algo.gemv
  else
    array_mt._gemm = unsupported_element_ct
    array_mt._gemv = unsupported_element_ct
  end
  
  local ct = ffi.typeof(array_struct, element_ct, element_ct)
  return ffi.metatype(ct, array_mt)
end

-- Typeof ----------------------------------------------------------------------

-- To preserve value semantics.
local allowed_element_ct = { }

local diff = require 'sci.diff'

for ct_name in pairs{
  bool              = true,
  char              = true,
  int8_t            = true,
  int16_t           = true,
  int32_t           = true,
  int64_t           = true,
  uint8_t           = true,
  uint16_t          = true,
  uint32_t          = true,
  uint64_t          = true,
  float             = true,
  double            = true,
  ['complex float'] = true,
  complex           = true,
  [diff.dn]         = true,
} do
  local ct_id = tonumber(ffi.typeof(ct_name))
  allowed_element_ct[ct_id] = true
end

local alg_element_ct = { }

local function alg_typeof(element_ct)
  element_ct = ffi.typeof(element_ct) -- Allow for strings, now it's ctype.
  local element_ct_id = tonumber(element_ct)
  if not allowed_element_ct[element_ct_id] then
    error('element type "'..tostring(element_ct)..'" not allowed')
  end
  
  if alg_element_ct[element_ct_id] then
    return alg_element_ct[element_ct_id]
  end

  local is_diff_dn = element_ct == diff.dn
  local array_ct = new_array_ct(element_ct, is_diff_dn and element_ct)

  local function vec(n)
    if n < 0 then
      error('length '..n..' is negative')
    end
    return array_alloc(array_ct, n, n, 1)
  end

  local function mat(r, c)
    if r < 0 then
      error('number of rows '..r..' is negative')
    end
    if c < 0 then
      error('number of columns '..c..' is negative')
    end
    return array_alloc(array_ct, r*c, r, c)
  end

  local function tovec(t)
    if type(t) ~= 'table' then
      error('table argument expected, got '..type(t))
    end
    local n = #t
    local a = vec(n)
    for i=1,n do
      a[i] = t[i]
    end
    return a
  end

  local function tomat(t)
    if type(t) ~= 'table' then
      error('table argument expected, got '..type(t))
    end
    local r, c = #t, #t > 0 and #t[1] or 0
    local a = mat(r, c)
    for i=1,r do
      for j=1,c do
        if #t[i] ~= c then
          error('all rows of the table must have the same number of elements')
        end
        a[{i, j}] = t[i][j]
      end
    end
    return a
  end

  local alg = { 
    vec = vec, 
    mat = mat,
    tovec = tovec,
    tomat = tomat,
    arrayct = array_ct,
  }

  alg_element_ct[element_ct_id] = alg
  return alg_element_ct[element_ct_id]
end

--------------------------------------------------------------------------------
local __code = template([[
return {
| for NEL = 1,10 do
  dim_elw_${NEL} = function(${R('__x@', 1, NEL)})
    local n, r, c = __x1._n, __x1._r, __x1._c
| for N=2,NEL do
    if ${R('__x@._r ~= r or __x@._c ~= c', 2, N, ' or ')} then
      error('incompatible dimensions in element-wise operation')
    end
| end
    return n, r, c
  end,
| end
}
]])({ R = rep })

local __ = assert(loadstring(__code))()

__.array_alloc = array_alloc
__.stack_array = stack_array
__.stack_clear = stack_clear
__.mul = __mul
__.pow = __pow
__.dim_pow_1 = dimensions_pow_check_1
__.dim_pow_2 = dimensions_pow_check_2
__.dim_mul_2 = dimensions_mul_check_2
__.dim_mul_3 = dimensions_mul_check_3

--------------------------------------------------------------------------------
local alg_double = alg_typeof('double')

return {
  typeof = alg_typeof,

  vec = alg_double.vec,
  mat = alg_double.mat,
  tovec = alg_double.tovec,
  tomat = alg_double.tomat,
  arrayct = alg_double.arrayct,

  join = join,

  mul = mul,
  pow = pow,

  sum   = sum,
  prod  = prod,
  trace = trace,

  __ = __,
}