export IR and array allocation

code/code_gen/code_gen.py

@@ -500,6 +500,101 @@ class FortranProgram:
         return md
 
 
+    def print_graph (self):
+
+        dg = self.parser.dep_graph()
+
+        dot_shape = "{{ node [shape=box] {} }}"
+
+        dot_rank = "{{ rank=same {} }}"
+
+        dot_edge = "{} -> {}"
+
+
+
+        print "digraph G {"
+
+        print dot_shape.format(" ".join( self.pass1avg ))
+
+        print dot_rank.format(" ".join( self.pass1avg ))
+
+
+        for var in self.pass1list[::-1] :
+            for dv in dg[var]:
+                if dv in self.pass1set:
+                    print dot_edge.format(var, dv)
+
+        print "}"
+
+        print (self.pass1list)
+
+        pp.pprint (self.allocate_arr (self.pass1list))
+
+        print len(self.allocate_arr (self.pass1list))
+
+
+        print (self.pass2list)
+
+        pp.pprint (self.allocate_arr (self.pass2list))
+
+        print len(self.allocate_arr (self.pass2list))
+
+
+    def save_ir (self):
+
+        import json
+
+        dg = self.parser.dep_graph()
+
+        with open("ir.py", "w") as irf:
+
+            print >>irf, "g = ", json.dumps(dg, indent=4)
+
+            print >>irf, "l1 = ", json.dumps(self.pass1list, indent=4)
+
+            print >>irf, "l2 = ", json.dumps(self.pass2list, indent=4)
+
+            print >>irf, "avg1 = ", json.dumps(self.pass1avg, indent=4)
+
+            print >>irf, "avg2 = ", json.dumps(self.pass2avg, indent=4)
+
+
+    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.parser.dep_graph()
+
+        mask = self.liveness(l, dg)
+
+        narr = mask.astype(np.int).sum(axis=0).max()
+
+        livesets = [set(np.asarray(l)[row]) for row in mask.T]
+
+        array_pool = set(["xyzbuffer{}".format(i) for i in range(narr)])
+
+        var2arr = {}
+
+        var2arr[l[0]] = array_pool.pop()
+
+        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 var2arr
+
 
 def main():
     while True:
@@ -524,6 +619,9 @@ def test():
 
     print mod_form.format( fp.print_program() )
 
+    # fp.print_graph()
+
+    fp.save_ir()
 
 if __name__ == '__main__':
     test()