incomp-flame-post/code/code_gen/post.py

from lark import Lark, Visitor, Transformer, v_args, Token


class ExpInspector(Visitor):
    def __init__(self):
        self.fluctuation = False
        self.dep = set([])
        self.deriv = set([])

    @classmethod
    def inspect(cls, tree):
        self = cls()
        return self(tree)

    def __call__(self, tree):
        self.visit(tree)
        return self.fluctuation, self.dep, self.deriv

    def fluc(self, tree):
        self.fluctuation = True
        self.dep.add(tree.children[0].value)

    def var(self, tree):
        self.dep.add(tree.children[0].value)

    def dnx (self, tree):
        op, v = tree.children
        deriv = "{}_{}".format(op.data, v.value)
        self.dep.add(deriv)
        self.deriv.add((op.data, v.value))


class FieldBase (object):
    def __init__ (self, name, fdict):
        self.name = name
        self.array = name
        self.dep = set([])
        self.fluc = False
        self.fdict = fdict

    def depends_on (self, a):
        return (a in self.dep)

    def is_fluctuation (self):
        return False

    def __repr__ (self):
        return self.name

    def checkFluctuation (self):

        fset = set([])

        if self.is_fluctuation():
            fset.add(self.name)

        for d in map(self.fdict.get, self.dep):
            fset.update(d.checkFluctuation())

        if len(fset) > 0:
            fset.add(self.name)

        return fset


class Field (FieldBase):

    def __init__ (self, name, exp, fdict):
        super(Field,self).__init__(name, fdict)
        self.exp = exp
        self.fluc, self.dep, self.derivs = ExpInspector.inspect(exp)

        '''
        for a in self.dep:
            if a not in self.fdict:
                raise StandardError(a + " is not defined")
        '''

    def is_fluctuation (self):
        return self.fluc

    def code (self, alloc):
        transform(self.exp)
        return (a in self.dep)


class PrimaryField (FieldBase):

    def __init__ (self, name, fdict):
        super(PrimaryField,self).__init__(name, fdict)
        self.derivs = set([])

    def code (self):
        return ""


class DerivedField (FieldBase):

    def __init__ (self, op, v, fdict):
        name = "{}_{}".format(op, v)
        super(DerivedField,self).__init__(name, fdict)
        self.op = op
        self.v = v
        self.dep = set([v])

    def code (self):
        return "call {0} ( {2}, {1} )".format(self.op, self.v, self.array)


class AveragedField (FieldBase):

    def __init__ (self, w, tgt, fdict):

        if w:
            name = "{}_avg_{}".format(w, tgt)
        else:
            name = "avg_{}".format(tgt)

        super(AveragedField,self).__init__(name, fdict)

        self.tgt = fdict[tgt]
        self.dep.add(tgt)
        self.weighted = False

        if w:
            self.weighted = True
            self.w = fdict[w]
            self.dep.add(w)

        self.fset = self.checkFluctuation() - set([self.name])

    def isWeighted (self):
        return self.weighted

    def pass1 (self):
        return not self.pass2()

    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 '''

    def __init__ (self, raw_tree):
        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 '''

    def __init__ (self, src):
        self.src = src
        self.derived = src.derived
        self.derivative = {}
        self.averaged = {}

        # Construct Derivative Field objects

        dset = set([])
        for k, v in self.derived.iteritems():
            dset.update(v.derivs)

        for tup in dset:
            a = DerivedField(tup[0], tup[1], self.derived)
            self.derived[a.name] = a
            self.derivative[tup] = a

        # Construct Derivative Field objects

        for w, tgts in src.averaged.iteritems():
            for t in tgts:
                a = AveragedField(w, t, self.derived)
                self.averaged[a.name] = a

    def __repr__ (self):
        return "\n".join(map(str, [self.derived, self.derivative, self.averaged]))

    def dependency (self):
        dgraph = {}

        for k,v in self.derived.iteritems():
            dgraph[k] = v.dep

        for k,v in self.averaged.iteritems():
            dgraph[k] = v.dep

        return dgraph


class Stage3():
    ''' calculate execution order '''

    def __init__ (self, src):
        self.src = src
        self.derived = src.derived
        self.averaged = src.averaged
        self.dependency = src.dependency()

        pass1set = filter(AveragedField.pass1, self.averaged.values())
        pass2set = filter(AveragedField.pass2, self.averaged.values())

        print pass1set
        print pass2set

    def __repr__ (self):
        return "\n".join(map(str, [self.derived, self.derivative]))




class Stage4():
    ''' analyze liveness and allocate array '''

class Stage5():
    ''' pass1 and pass2 seperation and calculation ordering '''

class Stage6():
    ''' generate fortran code '''