stage4 memory optimization based on liveness analysis
This commit is contained in:
ignis
parent
00caa1aa7e
commit
a6da2d9e2e
2 changed files with 219 additions and 49 deletions
|
|
@ -626,7 +626,13 @@ def test():
|
|||
|
||||
ir3 = Stage3(ir2)
|
||||
|
||||
print mod_form.format( ir3.print_program() )
|
||||
ir4 = Stage4(ir3)
|
||||
|
||||
# ir4.save_ir()
|
||||
|
||||
print mod_form.format( ir4.print_program() )
|
||||
|
||||
# print mod_form.format( ir3.print_program() )
|
||||
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,35 @@
|
|||
from lark import Lark, Visitor, Transformer, v_args, Token
|
||||
|
||||
|
||||
class CollectDefinitions(Visitor):
|
||||
|
||||
def __init__ (self, primary, derived, averaged):
|
||||
self.primary = primary
|
||||
self.derived = derived
|
||||
self.averaged = averaged
|
||||
|
||||
def varlist(self, tree):
|
||||
for v in tree.children:
|
||||
self.primary.add(v.value)
|
||||
self.derived[v.value] = PrimaryField(v.value, self.derived)
|
||||
|
||||
def assign_var (self, tree):
|
||||
lval, rval = tree.children
|
||||
|
||||
if lval.value in self.derived:
|
||||
raise StandardError("duplicate definition of " + lval)
|
||||
self.derived[lval.value] = Field(lval.value, rval, self.derived)
|
||||
|
||||
def assign_avg_var (self, tree):
|
||||
w = tree.children[0]
|
||||
targets = tree.children[1:]
|
||||
|
||||
try:
|
||||
self.averaged[w.children[0].value] = set([x.value for x in targets])
|
||||
except IndexError:
|
||||
self.averaged[None] = set([x.value for x in targets])
|
||||
|
||||
|
||||
class ExpInspector(Visitor):
|
||||
def __init__(self):
|
||||
self.fluctuation = False
|
||||
|
|
@ -108,12 +137,15 @@ class ExpToCode(Transformer):
|
|||
rxn_rate = lambda self : "rxn_rate"
|
||||
|
||||
|
||||
|
||||
|
||||
class FieldBase (object):
|
||||
def __init__ (self, name, fdict):
|
||||
self.name = name
|
||||
self.array = name
|
||||
self.dep = set([])
|
||||
self.fluc = False
|
||||
self.prime = False
|
||||
self.fdict = fdict
|
||||
self.shape = "nxp,nyp,nzp"
|
||||
self.dim = ":,:,:"
|
||||
|
|
@ -121,6 +153,12 @@ class FieldBase (object):
|
|||
def depends_on (self, a):
|
||||
return (a in self.dep)
|
||||
|
||||
def is_primary (self):
|
||||
return self.prime
|
||||
|
||||
def not_primary (self):
|
||||
return not self.prime
|
||||
|
||||
def is_fluctuation (self):
|
||||
return self.fluc
|
||||
|
||||
|
|
@ -176,7 +214,7 @@ class Field (FieldBase):
|
|||
raise StandardError(a + " is not defined")
|
||||
'''
|
||||
|
||||
def code (self):
|
||||
def code (self, alloc=None):
|
||||
real_array_loop = """
|
||||
! {1}
|
||||
do k = 1, nzp
|
||||
|
|
@ -187,6 +225,8 @@ end do
|
|||
end do
|
||||
end do
|
||||
"""
|
||||
self.array = alloc[self.name] if alloc else self.name
|
||||
|
||||
rhs = ExpToCode(self.fdict).transform(self.exp)
|
||||
return real_array_loop.format((self.array, rhs), self.comment)
|
||||
|
||||
|
|
@ -213,7 +253,7 @@ class FluctuationField (FieldBase):
|
|||
self.comment = self.comment.format(self.w)
|
||||
|
||||
|
||||
def code (self):
|
||||
def code (self, alloc=None):
|
||||
real_array_loop = """
|
||||
! {1}
|
||||
do k = 1, nzp
|
||||
|
|
@ -224,6 +264,8 @@ end do
|
|||
end do
|
||||
end do
|
||||
"""
|
||||
self.array = alloc[self.name] if alloc else self.name
|
||||
|
||||
rhs = ExpToCode(self.fdict).transform(self.field.exp)
|
||||
|
||||
if self.field.is_fluctuation():
|
||||
|
|
@ -239,13 +281,15 @@ end do
|
|||
name = "{}____avg".format(field)
|
||||
return name
|
||||
|
||||
|
||||
class PrimaryField (FieldBase):
|
||||
|
||||
def __init__ (self, name, fdict):
|
||||
super(PrimaryField,self).__init__(name, fdict)
|
||||
self.derivs = set([])
|
||||
self.prime = True
|
||||
|
||||
def code (self):
|
||||
def code (self, alloc=None):
|
||||
return "! {} is read from file".format(self.name)
|
||||
|
||||
def code_decl (self):
|
||||
|
|
@ -258,7 +302,6 @@ class PrimaryField (FieldBase):
|
|||
return "! {} is read from file".format(self.name)
|
||||
|
||||
|
||||
|
||||
class DerivedField (FieldBase):
|
||||
|
||||
def __init__ (self, op, v, fdict):
|
||||
|
|
@ -268,7 +311,8 @@ class DerivedField (FieldBase):
|
|||
self.v = v
|
||||
self.dep = set([v])
|
||||
|
||||
def code (self):
|
||||
def code (self, alloc=None):
|
||||
self.array = alloc[self.name] if alloc else self.name
|
||||
return "call {0} ( {2}, {1} )".format(self.op, self.v, self.array)
|
||||
|
||||
|
||||
|
|
@ -305,7 +349,7 @@ class AveragedField (FieldBase):
|
|||
self.w = fdict[w]
|
||||
self.dep.add(w)
|
||||
|
||||
def code (self):
|
||||
def code (self, alloc=None):
|
||||
|
||||
avg_array_sum = """
|
||||
do k = 1, nzp
|
||||
|
|
@ -344,33 +388,6 @@ call MPI_ALLREDUCE(MPI_IN_PLACE, {0}, nxp, MPI_REAL8, MPI_SUM, MPI_COMM_TASK, mp
|
|||
def pass2 (self):
|
||||
return len(self.fset) > 0
|
||||
|
||||
class CollectDefinitions(Visitor):
|
||||
|
||||
def __init__ (self, primary, derived, averaged):
|
||||
self.primary = primary
|
||||
self.derived = derived
|
||||
self.averaged = averaged
|
||||
|
||||
def varlist(self, tree):
|
||||
for v in tree.children:
|
||||
self.primary.add(v.value)
|
||||
self.derived[v.value] = PrimaryField(v.value, self.derived)
|
||||
|
||||
def assign_var (self, tree):
|
||||
lval, rval = tree.children
|
||||
|
||||
if lval.value in self.derived:
|
||||
raise StandardError("duplicate definition of " + lval)
|
||||
self.derived[lval.value] = Field(lval.value, rval, self.derived)
|
||||
|
||||
def assign_avg_var (self, tree):
|
||||
w = tree.children[0]
|
||||
targets = tree.children[1:]
|
||||
|
||||
try:
|
||||
self.averaged[w.children[0].value] = set([x.value for x in targets])
|
||||
except IndexError:
|
||||
self.averaged[None] = set([x.value for x in targets])
|
||||
|
||||
class Stage1():
|
||||
''' conversion from tree to python data '''
|
||||
|
|
@ -379,17 +396,20 @@ class Stage1():
|
|||
self.primary = set([])
|
||||
self.derived = {}
|
||||
self.averaged = {}
|
||||
|
||||
# Construct Field objects
|
||||
CollectDefinitions( self.primary, self.derived, self.averaged ).visit(raw_tree)
|
||||
|
||||
def __repr__ (self):
|
||||
return "\n".join(map(str, [self.primary, self.derived, self.averaged]))
|
||||
|
||||
|
||||
class Stage2():
|
||||
''' expand derivatives and averages'''
|
||||
|
||||
def __init__ (self, src):
|
||||
self.src = src
|
||||
self.primary = src.primary
|
||||
self.derived = src.derived
|
||||
self.derivative = {}
|
||||
self.averaged = {}
|
||||
|
|
@ -435,25 +455,26 @@ class Stage3():
|
|||
|
||||
def __init__ (self, src):
|
||||
self.src = src
|
||||
self.primary = src.primary
|
||||
self.derived = src.derived
|
||||
self.averaged = src.averaged
|
||||
self.dependency = src.dependency()
|
||||
|
||||
assert set(self.derived.keys()).isdisjoint(self.averaged.keys())
|
||||
|
||||
pass1set = set(filter(AveragedField.pass1, self.averaged.values()))
|
||||
pass2set = set(filter(AveragedField.pass2, self.averaged.values()))
|
||||
|
||||
self.avg1 = pass1set
|
||||
self.avg2 = pass2set
|
||||
self.avg1 = set(filter(AveragedField.pass1, self.averaged.values()))
|
||||
self.avg2 = set(filter(AveragedField.pass2, self.averaged.values()))
|
||||
|
||||
pass1calc = set(map(repr, pass1set))
|
||||
apply (pass1calc.update, map(AveragedField.depClosure, pass1set))
|
||||
self.pass1 = self.sort_vars(self.dependency, pass1calc)
|
||||
pass1calc = set(map(repr, self.avg1))
|
||||
map (pass1calc.update, map(AveragedField.depClosure, self.avg1))
|
||||
self.pass1 = self.sort_vars(self.dependency, pass1calc - self.primary)
|
||||
|
||||
pass2calc = set(map(repr, pass2set))
|
||||
apply (pass2calc.update, map(AveragedField.depClosure, pass2set))
|
||||
self.pass2 = self.sort_vars(self.dependency, pass2calc)
|
||||
# self.sort_vars(self.dependency, pass1calc - self.primary -set(map(repr, self.avg1)))
|
||||
|
||||
pass2calc = set(map(repr, self.avg2))
|
||||
map (pass2calc.update, map(AveragedField.depClosure, self.avg2))
|
||||
self.pass2 = self.sort_vars(self.dependency, pass2calc - self.primary)
|
||||
|
||||
|
||||
def __repr__ (self):
|
||||
|
|
@ -480,9 +501,7 @@ class Stage3():
|
|||
allvar.update(self.averaged)
|
||||
|
||||
decl = "\n".join(allvar[v].code_decl() for v in set(self.pass1+self.pass2))
|
||||
|
||||
alloc = "\n".join(allvar[v].code_alloc() for v in set(self.pass1+self.pass2))
|
||||
|
||||
free = "\n".join(allvar[v].code_free() for v in set(self.pass1+self.pass2))
|
||||
|
||||
calc1 = "\n".join(allvar[v].code() for v in self.pass1)
|
||||
|
|
@ -497,11 +516,9 @@ class Stage3():
|
|||
avg1 = "\n".join(allvar[v].code_avg() for v in set1)
|
||||
avg2 = "\n".join(allvar[v].code_avg() for v in set2)
|
||||
|
||||
|
||||
hfmt = 'character (len = *), parameter :: output_header="{}"'
|
||||
declh = hfmt.format(" ".join(["x"] + set1 + set2))
|
||||
|
||||
|
||||
avg_array_write = '''
|
||||
integer :: i
|
||||
|
||||
|
|
@ -522,7 +539,6 @@ close (200)
|
|||
", ".join(map(avgarr.format, set1+set2))
|
||||
)
|
||||
|
||||
|
||||
md = {}
|
||||
md["module_name"] = "terms"
|
||||
md["module_data"] = "\n".join((declh, decl))
|
||||
|
|
@ -539,6 +555,154 @@ close (200)
|
|||
|
||||
class Stage4():
|
||||
''' analyze liveness and allocate array '''
|
||||
def __init__ (self, src):
|
||||
self.src = src
|
||||
self.derived = src.derived
|
||||
self.averaged = src.averaged
|
||||
self.dependency = src.dependency
|
||||
self.avg1 = src.avg1
|
||||
self.avg2 = src.avg2
|
||||
self.pass1 = src.pass1
|
||||
self.pass2 = src.pass2
|
||||
|
||||
self.array_name = "xyzbuffer{}"
|
||||
|
||||
narr1, alloc1 = (self.allocate_arr(self.pass1))
|
||||
narr2, alloc2 = (self.allocate_arr(self.pass2))
|
||||
|
||||
self.narr = max(narr1, narr2)
|
||||
|
||||
self.alloc1 = alloc1
|
||||
self.alloc2 = alloc2
|
||||
|
||||
|
||||
def liveness (self, l1, g):
|
||||
import numpy as np
|
||||
img = np.zeros((len(l1), len(l1)))
|
||||
for i, v in enumerate(l1):
|
||||
for j in range(i, len(l1)):
|
||||
img[i,i:j] = img[i,i:j] + (1 if v in g[l1[j]] else 0)
|
||||
return img > 0
|
||||
|
||||
|
||||
def allocate_arr (self, l):
|
||||
|
||||
import numpy as np
|
||||
|
||||
dg = self.dependency
|
||||
|
||||
mask = self.liveness(l, dg)
|
||||
|
||||
narr = mask.astype(np.int).sum(axis=0).max()
|
||||
|
||||
array_pool = set([self.array_name.format(i) for i in range(narr)])
|
||||
|
||||
livesets = [set([])] + [set(np.asarray(l)[row]) for row in mask.T]
|
||||
|
||||
var2arr = {}
|
||||
|
||||
for i, (s0, s1) in enumerate(zip(livesets[:-1], livesets[1:])):
|
||||
|
||||
array_pool.update(map(var2arr.get, s0 - s1))
|
||||
|
||||
for new in s1 - s0:
|
||||
var2arr[new] = array_pool.pop()
|
||||
|
||||
return narr, var2arr
|
||||
|
||||
|
||||
def array_codes (self):
|
||||
|
||||
array_names = [self.array_name.format(i) for i in range(self.narr)]
|
||||
|
||||
real_array_decl = "real*8, allocatable, dimension(:,:,:) :: {0}"
|
||||
real_array_alloc = "allocate({0}(nxp,nyp,nzp), stat=ierr) ; {0} = 0."
|
||||
real_array_free = "deallocate({})"
|
||||
|
||||
decl = "\n".join([real_array_decl.format(v) for v in array_names])
|
||||
alloc = "\n".join([real_array_alloc.format(v) for v in array_names])
|
||||
free = "\n".join([real_array_free.format(v) for v in array_names])
|
||||
|
||||
return decl, alloc, free
|
||||
|
||||
|
||||
|
||||
def print_program (self):
|
||||
|
||||
allvar = dict(self.derived)
|
||||
allvar.update(self.averaged)
|
||||
|
||||
declarr, allocarr, freearr = self.array_codes()
|
||||
|
||||
declavg = "\n".join(v.code_decl() for v in self.averaged.itervalues())
|
||||
allocavg = "\n".join(v.code_alloc() for v in self.averaged.itervalues())
|
||||
freeavg = "\n".join(v.code_free() for v in self.averaged.itervalues())
|
||||
|
||||
calc1 = "\n".join(allvar[v].code(self.alloc1) for v in self.pass1)
|
||||
calc2 = "\n".join(allvar[v].code(self.alloc2) for v in self.pass2)
|
||||
|
||||
set1 = [a.name for a in filter(AveragedField.pass1, self.averaged.values())]
|
||||
set2 = [a.name for a in filter(AveragedField.pass2, self.averaged.values())]
|
||||
|
||||
set1.sort()
|
||||
set2.sort()
|
||||
|
||||
avg1 = "\n".join(allvar[v].code_avg() for v in set1)
|
||||
avg2 = "\n".join(allvar[v].code_avg() for v in set2)
|
||||
|
||||
hfmt = 'character (len = *), parameter :: output_header="{}"'
|
||||
declh = hfmt.format(" ".join(["x"] + set1 + set2))
|
||||
|
||||
avg_array_write = '''
|
||||
integer :: i
|
||||
|
||||
open (200, file="qEdge_X.dat")
|
||||
|
||||
write (200,*) output_header
|
||||
|
||||
do i=1,nxp
|
||||
write (200,'({0}e20.10)') real(i)*hxp, {1}
|
||||
end do
|
||||
|
||||
close (200)
|
||||
'''
|
||||
|
||||
avgarr = "{}(i)"
|
||||
write_avg = avg_array_write.format(
|
||||
len(self.averaged)+1,
|
||||
", ".join(map(avgarr.format, set1+set2))
|
||||
)
|
||||
|
||||
md = {}
|
||||
md["module_name"] = "terms"
|
||||
md["module_data"] = "\n".join((declh, declavg, declarr))
|
||||
md["module_init"] = "\n".join((allocavg, allocarr))
|
||||
md["module_finalize"] = "\n".join((freeavg, freearr))
|
||||
md["module_pass1"] = calc1
|
||||
md["module_pass1_avg"] = avg1
|
||||
md["module_pass2"] = calc2
|
||||
md["module_pass2_avg"] = avg2
|
||||
md["module_write_result"] = write_avg
|
||||
|
||||
return md
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
def save_ir (self):
|
||||
|
||||
import json
|
||||
|
||||
dg = {k:list(v) for k,v in self.dependency.items()}
|
||||
|
||||
with open("ir2.py", "w") as irf:
|
||||
|
||||
print >>irf, "g = ", json.dumps(dg, indent=4)
|
||||
print >>irf, "l1 = ", json.dumps(self.pass1, indent=4)
|
||||
print >>irf, "l2 = ", json.dumps(self.pass2, indent=4)
|
||||
print >>irf, "avg1 = ", json.dumps(map(repr,self.avg1), indent=4)
|
||||
print >>irf, "avg2 = ", json.dumps(map(repr,self.avg2), indent=4)
|
||||
|
||||
class Stage5():
|
||||
''' pass1 and pass2 seperation and calculation ordering '''
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue