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
|
--[[
Copyright (c) 2022, Vsevolod Stakhov <vsevolod@rspamd.com>
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
]]--
--[[[
-- @module lua_ffi/linalg
-- This module contains ffi interfaces to linear algebra routines
--]]
local ffi = require 'ffi'
local exports = {}
ffi.cdef [[
void kad_sgemm_simple(int trans_A, int trans_B, int M, int N, int K, const float *A, const float *B, float *C);
bool kad_ssyev_simple (int N, float *A, float *output);
]]
local function table_to_ffi(a, m, n)
local a_conv = ffi.new("float[?]", m * n)
for i = 1, m or #a do
for j = 1, n or #a[1] do
a_conv[(i - 1) * n + (j - 1)] = a[i][j]
end
end
return a_conv
end
local function ffi_to_table(a, m, n)
local res = {}
for i = 0, m - 1 do
res[i + 1] = {}
for j = 0, n - 1 do
res[i + 1][j + 1] = a[i * n + j]
end
end
return res
end
exports.sgemm = function(a, m, b, n, k, trans_a, trans_b)
if type(a) == 'table' then
-- Need to convert, slow!
a = table_to_ffi(a, m, k)
end
if type(b) == 'table' then
b = table_to_ffi(b, k, n)
end
local res = ffi.new("float[?]", m * n)
ffi.C.kad_sgemm_simple(trans_a or 0, trans_b or 0, m, n, k, ffi.cast('const float*', a),
ffi.cast('const float*', b), ffi.cast('float*', res))
return res
end
exports.eigen = function(a, n)
if type(a) == 'table' then
-- Need to convert, slow!
n = n or #a
a = table_to_ffi(a, n, n)
end
local res = ffi.new("float[?]", n)
if ffi.C.kad_ssyev_simple(n, ffi.cast('float*', a), res) then
return res, a
end
return nil
end
exports.ffi_to_table = ffi_to_table
exports.table_to_ffi = table_to_ffi
return exports
|
