Files
linguist/samples/Terra/dgemm.t
2016-01-28 11:22:27 +07:00

143 lines
3.2 KiB
Perl

local IO = terralib.includec("stdio.h")
local stdlib = terralib.includec("stdlib.h")
local function isinteger(x) return math.floor(x) == x end
local NB = 48
terra naivel1matmul(A : &double, B : &double, C : &double, lda : int, ldb : int, ldc : int, alpha : double)
for m = 0, NB do
for n = 0, NB do
C[m*ldc + n] = alpha * C[m*ldc + n]
for k = 0, NB do
C[m*ldc + n] = C[m*ldc + n] + A[m*lda + k] * B[k*ldb + n]
end
end
end
end
function symmat(name,I,...)
if not I then return symbol(name) end
local r = {}
for i = 0,I-1 do
r[i] = symmat(name..tostring(i),...)
end
return r
end
function genl1matmul(NB, NK, RM, RN, V,prefetch)
assert(isinteger(NB / (RN*V)))
assert(isinteger(NB / RM))
local VP = &vector(double,V)
local terra vecload(data : &double, idx : int)
var addr = &data[idx]
return @VP(addr)
end
local terra vecstore(data : &double, idx : int, v : vector(double,V))
var addr = &data[idx]
@VP(addr) = v
end
local A,B,C,mm,nn, alpha = symbol("A"),symbol("B"),symbol("C"),symbol("mn"),symbol("nn"),symbol("alpha")
local lda,ldb,ldc = symbol("lda"),symbol("ldb"), symbol("ldc")
local a,b,c = symmat("a",NB/V,RM), symmat("b",NB,RN), symmat("c",RM,RN)
local kk = symbol("kk")
local loadc,storec = terralib.newlist(),terralib.newlist()
for m = 0, RM-1 do
for n = 0, RN-1 do
loadc:insert(quote
var [c[m][n]] = alpha * vecload(C,(mm+m)*ldc + nn + n*V)
end)
storec:insert(quote
vecstore(C,(mm+m)*ldc + nn + n*V,[c[m][n]])
end)
end
end
local calcc = terralib.newlist()
for kb = 0, NK/V-1 do
local kbV = kb*V
for m = 0, RM-1 do
calcc:insert(quote
var [a[kb][m]] = vecload(A,(mm+m)*lda + kk + kbV)
end)
end
for v = 0, V-1 do
local k = kbV+v
if not prefetch or (v == 0 and kb == 0) then
for n = 0, RN-1 do
calcc:insert(quote
var [b[k][n]] = vecload(B,(kk + k)*ldb + nn + n*V)
end)
end
end
for m = 0, RM-1 do
for n = 0, RN-1 do
calcc:insert(quote
[c[m][n]] = [c[m][n]] + [a[kb][m]][v] * [b[k][n]]
end)
if prefetch and not (v == V-1 and kb == NK/V-1) and m == RM-1 then --prefetch the next b
calcc:insert(quote
var [b[k+1][n]] = vecload(B,(kk + k + 1)*ldb + nn + n*V)
end)
end
end
end
end
end
return terra([A] : &double, [B] : &double, [C] : &double, [lda] : int, [ldb] : int, [ldc] : int, [alpha] : double)
for [mm] = 0, NB, RM do
for [nn] = 0, NB,RN*V do
[loadc];
for [kk] = 0, NB, NK do
[calcc];
end
[storec];
end
end
end
end
local NB2 = 8 * NB
local l1matmul = genl1matmul(NB,4, 3, 2, 4)
terra min(a : int, b : int)
return terralib.select(a < b, a, b)
end
terra my_dgemm(gettime : {} -> double, M : int, N : int, K : int, alpha : double, A : &double, lda : int, B : &double, ldb : int,
beta : double, C : &double, ldc : int)
for mm = 0,M,NB2 do
for nn = 0,N,NB2 do
for kk = 0,K, NB2 do
for m = mm,min(mm+NB2,M),NB do
for n = nn,min(nn+NB2,N),NB do
for k = kk,min(kk+NB2,K),NB do
l1matmul(A + m*lda + k,
B + k*ldb + n,
C + m*ldc + n,
lda,ldb,ldc, terralib.select(k == 0,0,1))
end
end
end
end
end
end
end
terralib.saveobj("my_dgemm.o", {my_dgemm = my_dgemm})