RVV and other mixed precission stuff :)

src/exo/API_cursors.py

@@ -581,6 +581,8 @@ def loopir_type_to_exotype(typ: LoopIR.Type) -> API.ExoType:
         return API.ExoType.F64
     elif isinstance(typ, LoopIR.INT8):
         return API.ExoType.I8
+    elif isinstance(typ, LoopIR.INT16):
+        return API.ExoType.I16
     elif isinstance(typ, LoopIR.INT32):
         return API.ExoType.I32
     elif isinstance(typ, LoopIR.Bool):

src/exo/API_scheduling.py

@@ -206,7 +206,6 @@ def __call__(self, bval, all_args):
             self.err("expected a bool")
         return bval
 
-
 class OptionalA(ArgumentProcessor):
     def __init__(self, arg_proc):
         if is_subclass_obj(arg_proc, ArgumentProcessor):
@@ -316,6 +315,7 @@ class TypeAbbrevA(ArgumentProcessor):
         "f32": T.f32,
         "f64": T.f64,
         "i8": T.int8,
+        "i16": T.int16,
         "ui8": T.uint8,
         "ui16": T.uint16,
         "i32": T.int32,
@@ -1518,8 +1518,8 @@ def stage_window(proc, expr_cursor, win_name, memory=None):
     return scheduling.DoStageWindow(proc, win_name, memory, e).result()
 
 
-@sched_op([BlockCursorA, CustomWindowExprA("block_cursor"), NameA, BoolA])
-def stage_mem(proc, block_cursor, win_expr, new_buf_name, accum=False):
+@sched_op([BlockCursorA, CustomWindowExprA("block_cursor"), NameA, BoolA, BoolA])
+def stage_mem(proc, block_cursor, win_expr, new_buf_name, accum=False, init_zero=False):
     """
     Stage the window of memory specified by `win_expr` into a new buffer
     before the indicated code block and move the memory back after the
@@ -1568,7 +1568,7 @@ def stage_mem(proc, block_cursor, win_expr, new_buf_name, accum=False):
     """
     buf_name, w_exprs = win_expr
     ir, fwd = scheduling.DoStageMem(
-        block_cursor._impl, buf_name, w_exprs, new_buf_name, use_accum_zero=accum
+        block_cursor._impl, buf_name, w_exprs, new_buf_name, use_accum_zero=accum, init_zero=init_zero
     )
     return Procedure(ir, _provenance_eq_Procedure=proc, _forward=fwd)
 

src/exo/LoopIR.py

@@ -107,6 +107,7 @@ def __new__(cls, op):
          | F64()
          | INT8()
          | UINT8()
+         | INT16()
          | UINT16()
          | INT32()
          | Bool()
@@ -140,6 +141,7 @@ def __new__(cls, op):
         "F32",
         "F64",
         "INT8",
+        "INT16",
         "UINT16",
         "UINT8",
         "INT32" "Bool",
@@ -203,6 +205,7 @@ def __new__(cls, op):
             | F64   ()
             | INT8  ()
             | UINT8  ()
+            | INT16  ()
             | UINT16 ()
             | INT32 ()
             | Bool  ()
@@ -228,6 +231,7 @@ def __new__(cls, op):
         "F32",
         "F64",
         "INT8",
+        "INT16",
         "UINT8",
         "UINT16",
         "INT32",
@@ -363,6 +367,7 @@ def __new__(cls, op):
 @extclass(UAST.F32)
 @extclass(UAST.F64)
 @extclass(UAST.INT8)
+@extclass(UAST.INT16)
 @extclass(UAST.UINT8)
 @extclass(UAST.UINT16)
 @extclass(UAST.INT32)
@@ -404,6 +409,7 @@ class T:
     F32 = LoopIR.F32
     F64 = LoopIR.F64
     INT8 = LoopIR.INT8
+    INT16 = LoopIR.INT16
     UINT8 = LoopIR.UINT8
     UINT16 = LoopIR.UINT16
     INT32 = LoopIR.INT32
@@ -420,9 +426,11 @@ class T:
     f16 = F16()
     f32 = F32()
     int8 = INT8()
+    int16 = INT16()
     uint8 = UINT8()
     uint16 = UINT16()
     i8 = INT8()
+    i16 = INT16()
     ui8 = UINT8()
     ui16 = UINT16()
     int32 = INT32()
@@ -447,6 +455,7 @@ class T:
 @extclass(T.F32)
 @extclass(T.F64)
 @extclass(T.INT8)
+@extclass(T.INT16)
 @extclass(T.UINT8)
 @extclass(T.UINT16)
 @extclass(T.INT32)
@@ -468,6 +477,7 @@ def shape(t):
 @extclass(T.F32)
 @extclass(T.F64)
 @extclass(T.INT8)
+@extclass(T.INT16)
 @extclass(T.UINT8)
 @extclass(T.UINT16)
 @extclass(T.INT32)
@@ -487,6 +497,8 @@ def ctype(t):
         return "double"
     elif isinstance(t, T.INT8):
         return "int8_t"
+    elif isinstance(t, T.INT16):
+        return "int16_t"
     elif isinstance(t, T.UINT8):
         return "uint8_t"
     elif isinstance(t, T.UINT16):
@@ -505,7 +517,7 @@ def ctype(t):
 @extclass(LoopIR.type)
 def is_real_scalar(t):
     return isinstance(
-        t, (T.Num, T.F16, T.F32, T.F64, T.INT8, T.UINT8, T.UINT16, T.INT32)
+        t, (T.Num, T.F16, T.F32, T.F64, T.INT8, T.INT16, T.UINT8, T.UINT16, T.INT32)
     )
 
 

src/exo/LoopIR_compiler.py

@@ -285,7 +285,7 @@ def _window_struct(typename, ctype, n_dims, is_const) -> WindowStruct:
         f"    const int_fast32_t strides[{n_dims}];\n"
         f"}};"
     )
-
+    #sdef = ("") #ADRIAN
     return WindowStruct(sname, sdef)
 
 
@@ -297,6 +297,7 @@ def window_struct(base_type, n_dims, is_const) -> WindowStruct:
         T.f32: "f32",
         T.f64: "f64",
         T.i8: "i8",
+        T.i16: "i16",
         T.ui8: "ui8",
         T.ui16: "ui16",
         T.i32: "i32",
@@ -740,6 +741,7 @@ def access_str(self, nm, idx_list) -> str:
         if not type.is_win():
             return f"{buf}[{idx_expr_s}]"
         else:
+            #return f"{buf}[{idx_expr_s}]"
             return f"{buf}.data[{idx_expr_s}]"
 
     def shape_strs(self, shape, prec=100) -> str:
@@ -962,6 +964,7 @@ def comp_fnarg(self, e, fn, i, *, prec=0):
             win_struct = self.get_window_type(e.type, is_const)
             data, strides = self.window_struct_fields(e)
             return f"(struct {win_struct}){{ &{data}, {{ {strides} }} }}"
+            #return f"&{data}, {strides}"
         else:
             return self.comp_e(e, prec)
 
@@ -990,6 +993,7 @@ def comp_e(self, e, prec=0):
             win_struct = self.get_window_type(e.type)
             data, strides = self.window_struct_fields(e)
             return f"(struct {win_struct}){{ &{data}, {{ {strides} }} }}"
+            #return f"&{data}, {strides}"
 
         elif isinstance(e, LoopIR.Const):
             if isinstance(e.val, bool):

src/exo/LoopIR_pprint.py

@@ -292,6 +292,8 @@ def ptype(self, t):
             return "f64"
         elif isinstance(t, UAST.INT8):
             return "i8"
+        elif isinstance(t, UAST.INT16):
+            return "i16"
         elif isinstance(t, UAST.UINT8):
             return "ui8"
         elif isinstance(t, UAST.UINT16):
@@ -530,6 +532,8 @@ def _print_type(t, env: PrintEnv) -> str:
         return "f64"
     elif isinstance(t, T.INT8):
         return "i8"
+    elif isinstance(t, T.INT16):
+        return "i16"
     elif isinstance(t, T.UINT8):
         return "ui8"
     elif isinstance(t, T.UINT16):

src/exo/LoopIR_scheduling.py

@@ -3667,7 +3667,7 @@ def do_e(self, e):
         pass
 
 
-def DoStageMem(block_cursor, buf_name, w_exprs, new_name, use_accum_zero=False):
+def DoStageMem(block_cursor, buf_name, w_exprs, new_name, use_accum_zero=False, init_zero=False):
     proc = block_cursor.get_root()
     new_name = Sym(new_name)
 
@@ -3799,8 +3799,10 @@ def guard_wrapper(body):
                     )
                 else:
                     load_ridx.append(w)
-            load_rhs = LoopIR.Read(buf_name, load_ridx, basetyp, srcinfo)
-
+            if init_zero == False:
+                load_rhs = LoopIR.Read(buf_name, load_ridx, basetyp, srcinfo)
+            else:
+                load_rhs = LoopIR.Const(0.0, basetyp, srcinfo)
         load_nest = [
             LoopIR.Assign(new_name, basetyp, None, load_widx, load_rhs, None, srcinfo)
         ]
@@ -3822,9 +3824,10 @@ def guard_wrapper(body):
 
         if not use_accum_zero:
             load_nest_c = fwd(block_cursor[0]).prev()
-            ir, fwd = insert_safety_guards(
-                ir, fwd, get_inner_stmt(load_nest_c), load_rhs, buf_typ
-            )
+            if init_zero == False:
+                ir, fwd = insert_safety_guards(
+                    ir, fwd, get_inner_stmt(load_nest_c), load_rhs, buf_typ
+                )
     if isW:
         store_iter = [Sym(f"i{i}") for i, _ in enumerate(shape)]
         store_ridx = [LoopIR.Read(s, [], T.index, srcinfo) for s in store_iter]

src/exo/configs.py

@@ -28,6 +28,7 @@ def new_config(name, fields, disable_rw=False):
         'f32'       : LoopIR.T.f32,
         'f64'       : LoopIR.T.f64,
         'i8'        : LoopIR.T.i8,
+        'i16'       : LoopIR.T.i16,
         'i32'       : LoopIR.T.i32,
     }
     good_args = (isinstance(name, str) and
@@ -61,6 +62,8 @@ def ctyp(typ):
         return "double"
     elif isinstance(typ, LoopIR.T.INT8):
         return "int8_t"
+    elif isinstance(typ, LoopIR.T.INT16):
+        return "int16_t"
     elif isinstance(typ, LoopIR.T.INT32):
         return "int32_t"
     elif isinstance(typ, LoopIR.T.Bool):
@@ -97,6 +100,7 @@ def __init__(self, name, fields, disable_rw):
             LoopIR.UAST.F32(): LoopIR.T.f32,
             LoopIR.UAST.F64(): LoopIR.T.f64,
             LoopIR.UAST.INT8(): LoopIR.T.i8,
+            LoopIR.UAST.INT16(): LoopIR.T.i16,
             LoopIR.UAST.INT32(): LoopIR.T.i32,
         }
 

src/exo/memory.py

@@ -109,6 +109,7 @@ def window(cls, basetyp, baseptr, indices, strides, srcinfo):
 
         if basetyp.is_win():
             baseptr = f"{baseptr}.data"
+            #baseptr = f"{baseptr}"
 
         return f"{baseptr}[{offset}]"