From a6da2d9e2e8a8aa9f7d4a0b23f670959a27fffdf Mon Sep 17 00:00:00 2001 From: ignis Date: Thu, 9 May 2019 04:13:48 +0900 Subject: [PATCH] stage4 memory optimization based on liveness analysis --- code/code_gen/code_gen.py | 8 +- code/code_gen/post.py | 260 +++++++++++++++++++++++++++++++------- 2 files changed, 219 insertions(+), 49 deletions(-) diff --git a/code/code_gen/code_gen.py b/code/code_gen/code_gen.py index 04d5f92..8f4bf02 100644 --- a/code/code_gen/code_gen.py +++ b/code/code_gen/code_gen.py @@ -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() ) diff --git a/code/code_gen/post.py b/code/code_gen/post.py index 47b19d9..b183629 100644 --- a/code/code_gen/post.py +++ b/code/code_gen/post.py @@ -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 '''