somewhat simplified the symtab code

1 files changed, 49 insertions, 27 deletions

diff --git a/lib/c.py b/lib/c.py
index 25c02eb..08203cb 100644
--- a/lib/c.py
+++ b/lib/c.py
@@ -166,6 +166,9 @@ def IsVectorType(kind):
 
     return False
 
+def IsArrayType(kind):
+    return isinstance(kind, Array)
+
 # TODO: Make this ARCH dependent
 BitDepth= {
     Void:    8,
@@ -1118,7 +1121,7 @@ def _(op: BinaryOp, cond, results: List[Expr]):
 
 @singledispatch
 def Eval(x: Expr) -> object:
-    raise TypeError(f"{type(s)} not supported by Eval operation")
+    raise TypeError(f"{type(x)} not supported by Eval operation")
 
 @Eval.register
 def _(x: Empty):
@@ -1129,6 +1132,8 @@ def _(x: Empty):
 def _(x):
     return sympy.symbols(f"{x}")
 
+@Eval.register(float)
+@Eval.register(int)
 @Eval.register(I)
 @Eval.register(F)
 def _(x):
@@ -1156,6 +1161,18 @@ def _(op: Sub):
     return sympy.simplify(l - r)
 
 @Eval.register
+def _(op: Mul):
+    l = Eval(op.l)
+    r = Eval(op.r)
+    return sympy.simplify(l * r)
+
+@Eval.register
+def _(op: Div):
+    l = Eval(op.l)
+    r = Eval(op.r)
+    return sympy.simplify(l / r)
+
+@Eval.register
 def _(s: Assign):
     return Eval(s.rhs)
 
@@ -1189,6 +1206,7 @@ def AddrAccessed(stmt: Stmt):
         if vars[0] in scalars:
             scalars.remove(vars[0])
 
+    print([(key.name, val) for key, val in vectors.items()])
     return set(scalars), vectors
 
 # ------------------------------------------
@@ -1259,6 +1277,15 @@ def Unroll(loop: For, times: int, name: str, vars: List[vars] = []) -> (For, Fun
 
 # Replaces all vectorizable loops inside Func with vectorized variants
 # Returns the new vectorized function (tagged by the SIMD chosen)
+
+class SymTab(object):
+    def __init__(self):
+        self.stack = {}
+        self.addrs = {}
+
+    def have(self, name):
+        return name in self.stack or name in self.addrs
+
 def Vectorize(func: Func, isa: SIMD) -> Func:
     if isa != SIMD.AVX2:
         raise ValueError(f"ISA '{isa}' not currently implemented")
@@ -1280,11 +1307,6 @@ def Vectorize(func: Func, isa: SIMD) -> Func:
                 if numel(instr) % 4 != 0:
                     raise ValueError("loop can not be vectorized, instructions can not be globbed equally")
 
-                class symtab(object):
-                    scalar = {}
-                    stack  = {}
-                    addrs  = {}
-
                 # TODO: Allow for non-sequential accesses?
                 for i in range(0, numel(instr), 4):
                     scalars = []
@@ -1306,36 +1328,36 @@ def Vectorize(func: Func, isa: SIMD) -> Func:
                                 raise ValueError("non uniform vector accesses in consecutive line. can not vectorize")
 
                     # If we made it to here, we have passed all checks. vectorize!
-                    # Create symbol table
+                    vecs = vectors[0]
                     if i == 0:
-                        for s in scalars[0]:
-                            if type(s.var) == Param:
-                                intermediate          = vfunc.declare(Var(Float64x2, f"{s.name}128"))
-                                symtab.scalar[s.name] = vfunc.declare(Var(Float64x4, f"{s.name}256"))
-                                vfunc.execute(Set(intermediate, Ref(s)))
-                                vfunc.execute(Set(symtab.scalar[s.name], intermediate))
-                            else:
-                                symtab.scalar[s.name] = vfunc.declare(Var(Float64x4, f"{s.name}256"))
-
-                        for v in vectors[0].keys():
+                        syms = SymTab()
+                        for v in vecs.keys():
+                            # All params are treated AS addresses to load into vectorized registers
                             if type(v.var) == Param:
-                                symtab.addrs[v.name] = Var(Ptr(Float64x4), f"{v.name}256")
+                                syms.addrs[v.name] = Var(Ptr(Float64x4), f"{v.name}256")
+                            # All stack declared parameters MUST be moved to vectorized registers
                             else:
-                                symtab.stack[v.name] = vfunc.declare(Var(Float64x4, f"{v.name}256"))
+                                assert IsArrayType(v.var.type), f"must be an array type, instead got {v.var.type}"
+                                nreg = v.var.type.len // 4
+                                if nreg > 1:
+                                    syms.stack[v.name] = vfunc.declare(Var(Array(Float64x4, nreg), f"{v.name}256"))
+                                else:
+                                    syms.stack[v.name] = vfunc.declare(Var(Float64x4, f"{v.name}256"))
 
                     # IMPORTANT: We do a post-order traversal.
                     # We transforms leaves (identifiers) first and then move back up the root
                     def translate(x):
                         if isinstance(x, Ident):
-                            if x.name in symtab.scalar:
-                                return symtab.scalar[x.name]
-                            elif x.name in symtab.stack:
-                                return symtab.stack[x.name]
-                            elif x.name in symtab.addrs:
-                                x.var = symtab.addrs[x.name]
+                            if x.name in syms.stack:
+                                return syms.stack[x.name]
+                            elif x.name in syms.addrs:
+                                x.var = syms.addrs[x.name]
                         if isinstance(x, Index):
-                            if type(x.x) == Ident and (x.x.name in symtab.addrs or x.x.name in symtab.stack):
-                                return Add(x.x, x.i)
+                            if type(x.x) == Ident:
+                                if x.x.name in syms.addrs:
+                                    return Add(x.x, x.i)
+                                elif f"{x.x.name[:-3]}" in syms.stack:
+                                    return Index(x.x, I(Eval(Div(x.i, 4))))
                         if isinstance(x, BinaryOp):
                             l, r = GetType(x.l), GetType(x.r)
                             if IsVectorType(l) or IsVectorType(r):