forked from terralang/terra
-
Notifications
You must be signed in to change notification settings - Fork 2
/
blocking.t
57 lines (49 loc) · 1.13 KB
/
blocking.t
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
function symmat(typ,name,I,...)
if not I then return symbol(typ,name) end
local r = {}
for i = 1,I do
r[i] = symmat(typ,name..tostring(i),...)
end
return r
end
terra min(a : int, b : int)
return terralib.select(a < b, a, b)
end
function blockedloop(bounds,sizes,bodyfn)
local indexes = symmat(int,"i",#sizes,#bounds)
--local makeloop --bug local function doesn't add to set of live variables...
local function makeloop(s,b)
if s > #sizes then
return bodyfn(unpack(indexes[#sizes]))
elseif b > #bounds then
return makeloop(s + 1, 1)
else
local topbound = bounds[b]
local blocksize = sizes[s]
local begin,bound
if s == 1 then
begin,bound = 0, topbound
else
begin,bound = indexes[s-1][b], sizes[s-1]
end
local step = sizes[s]
return quote
for [indexes[s][b]] = begin,min(begin+bound,topbound),step do
[ makeloop(s,b+1) ]
end
end
end
end
return makeloop(1,1)
end
IO = terralib.includec("stdio.h")
terra main()
var M,N = 30,40;
[blockedloop({M,N}, {10,1}, function(m,n)
return quote
IO.printf("%d %d\n",m,n)
end
end)]
end
main:printpretty()
main()