Implemented insert_noop_call scheduling operation (#661)

implement a new scheduling op to insert no-op function calls anywhere,
addressing #565. This scheduling op will be useful for inserting
prefetching and potentially synchronization primitives arbitrarily.

For now, it only works if the body of the proc is a single `pass`
statement, but we could potentially extend this in the future.

---------

Co-authored-by: Yuka Ikarashi <[email protected]>

Author: skeqiqevian

src/exo/API.py

@@ -11,7 +11,6 @@
 from .API_types import ProcedureBase, ExoType
 from . import LoopIR as LoopIR
 from .LoopIR_compiler import run_compile, compile_to_strings
-from .LoopIR_unification import DoReplace, UnificationError
 from .configs import Config
 from .boundscheck import CheckBounds
 from .memory import Memory

src/exo/API_scheduling.py

@@ -765,6 +765,14 @@ def parse_arg(a):
         return buf_name, args
 
 
+class NewExprOrCustomWindowExprA(NewExprA):
+    def __call__(self, expr_str, all_args):
+        try:
+            return NewExprA(self.cursor_arg, self.before)(expr_str, all_args)
+        except:
+            return CustomWindowExprA(self.cursor_arg, self.before)(expr_str, all_args)
+
+
 # --------------------------------------------------------------------------- #
 #  - * - * - * - * - * - * - * - * - * - * - * - * - * - * - * - * - * - * -
 # --------------------------------------------------------------------------- #
@@ -848,6 +856,15 @@ def insert_pass(proc, gap_cursor):
     return Procedure(ir, _provenance_eq_Procedure=proc, _forward=fwd)
 
 
+@sched_op([GapCursorA, ProcA, ListA(NewExprOrCustomWindowExprA("gap_cursor"))])
+def insert_noop_call(proc, gap_cursor, instr, args):
+    if len(args) != len(instr._loopir_proc.args):
+        raise TypeError("Function argument count mismatch")
+
+    ir, fwd = scheduling.DoInsertNoopCall(gap_cursor._impl, instr._loopir_proc, args)
+    return Procedure(ir, _provenance_eq_Procedure=proc, _forward=fwd)
+
+
 @sched_op([])
 def delete_pass(proc):
     """
@@ -1080,7 +1097,7 @@ def replace(proc, block_cursor, subproc, quiet=False):
     except UnificationError:
         if quiet:
             raise
-        print(f"Failed to unify the following:\nSubproc:\n{subproc}Statements:\n")
+        print(f"Failed to unify the following:\nSubproc:\n{subproc}\nStatements:")
         [print(sc._impl._node) for sc in block_cursor]
         raise
 

src/exo/LoopIR_scheduling.py

@@ -41,6 +41,7 @@
 import exo.API as api
 from .pattern_match import match_pattern
 from .memory import DRAM
+from .typecheck import check_call_types
 
 from functools import partial
 
@@ -2727,6 +2728,54 @@ def DoInsertPass(gap):
     return ir, fwd
 
 
+def DoInsertNoopCall(gap, proc, args):
+    srcinfo = gap.parent()._node.srcinfo
+
+    body = proc.body
+    if not (len(body) == 1 and isinstance(body[0], LoopIR.Pass)):
+        # TODO: We should allow for a more general case, e.g. loops of passes
+        raise SchedulingError("Cannot insert a proc whose body is not pass")
+
+    syms_env = extract_env(gap.anchor())
+
+    def get_typ_mem(buf_name):
+        for name, typ, mem in syms_env:
+            if str(name) == buf_name:
+                return name, typ, mem
+
+    def process_slice(idx):
+        if not isinstance(idx, tuple):
+            return idx
+
+        buf_name, w_exprs = idx
+        name, typ, _ = get_typ_mem(buf_name)
+
+        idxs = []
+        win_shape = []
+        for w_e in w_exprs:
+            if isinstance(w_e, tuple):
+                lo, hi = w_e
+                win_shape.append(LoopIR.BinOp("-", hi, lo, hi.type, srcinfo))
+                idxs.append(LoopIR.Interval(lo, hi, srcinfo))
+            else:
+                idxs.append(LoopIR.Point(w_e, srcinfo))
+
+        as_tensor = T.Tensor(win_shape, True, typ.basetype())
+        w_typ = T.Window(typ, as_tensor, name, idxs)
+        return LoopIR.WindowExpr(name, idxs, w_typ, srcinfo)
+
+    args = [process_slice(arg) for arg in args]
+    call_stmt = LoopIR.Call(proc, args, srcinfo)
+    ir, fwd = gap._insert([call_stmt])
+
+    def err_handler(_, msg):
+        raise SchedulingError(f"Function argument type mismatch:" + msg)
+
+    check_call_types(err_handler, args, proc.args)
+
+    return ir, fwd
+
+
 def DoDeleteConfig(proc_cursor, config_cursor):
     eq_mod_config = Check_DeleteConfigWrite(proc_cursor._node, [config_cursor._node])
     p, fwd = config_cursor._delete()

src/exo/platforms/x86.py

@@ -3,6 +3,17 @@
 from .. import instr, DRAM
 from ..libs.memories import AVX2, AVX512
 
+# --------------------------------------------------------------------------- #
+#   Prefetching
+# --------------------------------------------------------------------------- #
+
+
+@instr("_mm_prefetch(&{A_data}, {locality_hint});")
+def prefetch(A: [R][1] @ DRAM, locality_hint: size):
+    assert 0 <= locality_hint
+    assert locality_hint < 8
+    pass
+
 
 # --------------------------------------------------------------------------- #
 #   AVX2 intrinsics
@@ -127,7 +138,7 @@ def mm256_broadcast_sd_scalar(out: [f64][4] @ AVX2, val: f64):
         out[i] = val
 
 
-@instr("{dst_data} = _mm512_fmadd_ps({dst_data}, {lhs_data}, {rhs_data});")
+@instr("{dst_data} = _mm256_fmadd_ps({dst_data}, {lhs_data}, {rhs_data});")
 def mm256_fmadd_ps_broadcast(
     dst: [f32][8] @ AVX2, lhs: [f32][8] @ AVX2, rhs: [f32][1] @ DRAM
 ):
@@ -251,6 +262,39 @@ def mm256_add_epi16(out: [ui16][16] @ AVX2, x: [ui16][16] @ AVX2, y: [ui16][16]
 # --------------------------------------------------------------------------- #
 
 
+@instr("{dst_data} = _mm512_setzero_ps();")
+def mm512_setzero_ps(dst: [f32][16] @ AVX512):
+    assert stride(dst, 0) == 1
+
+    for i in seq(0, 16):
+        dst[i] = 0.0
+
+
+@instr("{out_data} = _mm512_add_ps({x_data}, {y_data});")
+def mm512_add_ps(out: [f32][16] @ AVX512, x: [f32][16] @ AVX512, y: [f32][16] @ AVX512):
+    assert stride(out, 0) == 1
+    assert stride(x, 0) == 1
+    assert stride(y, 0) == 1
+
+    for i in seq(0, 16):
+        out[i] = x[i] + y[i]
+
+
+@instr(
+    "{out_data} = _mm512_mask_add_ps({out_data}, ((1 << {N}) - 1), {x_data}, {y_data});"
+)
+def mm512_mask_add_ps(
+    N: size, out: [f32][16] @ AVX512, x: [f32][16] @ AVX512, y: [f32][16] @ AVX512
+):
+    assert stride(out, 0) == 1
+    assert stride(x, 0) == 1
+    assert stride(y, 0) == 1
+
+    for i in seq(0, 16):
+        if i < N:
+            out[i] = x[i] + y[i]
+
+
 @instr("{dst_data} = _mm512_loadu_ps(&{src_data});")
 def mm512_loadu_ps(dst: [f32][16] @ AVX512, src: [f32][16] @ DRAM):
     assert stride(src, 0) == 1

src/exo/pyparser.py

@@ -484,7 +484,7 @@ def parse_num_type(self, node, is_arg=False):
         elif isinstance(node, pyast.Name) and node.id in Parser._prim_types:
             return Parser._prim_types[node.id]
         else:
-            self.err(node, "unrecognized type: " + ast.dump(node))
+            self.err(node, "unrecognized type: " + pyast.dump(node))
 
     def parse_stmt_block(self, stmts):
         assert isinstance(stmts, list)

src/exo/stdlib/scheduling.py

@@ -35,6 +35,7 @@
     inline,
     replace,
     call_eqv,
+    insert_noop_call,
     #
     # precision, memory, and window annotation setting
     set_precision,

src/exo/typecheck.py

@@ -34,6 +34,63 @@
 # The typechecker
 
 
+def check_call_types(err_handler, args, call_args):
+    for call_a, sig_a in zip(args, call_args):
+        if call_a.type == T.err:
+            pass
+        elif sig_a.type is T.size or sig_a.type is T.index:
+            if not call_a.type.is_indexable():
+                err_handler(
+                    call_a,
+                    "expected size or index type "
+                    "expression, "
+                    f"but got type {call_a.type}",
+                )
+
+        elif sig_a.type is T.bool:
+            if not call_a.type is T.bool:
+                err_handler(
+                    call_a,
+                    "expected bool-type variable, " f"but got type {call_a.type}",
+                )
+
+        elif sig_a.type is T.stride:
+            if not call_a.type.is_stridable():
+                err_handler(
+                    call_a,
+                    "expected stride-type variable, " f"but got type {call_a.type}",
+                )
+
+        elif sig_a.type.is_numeric():
+            if call_a.type.is_numeric():
+                if len(call_a.type.shape()) != len(sig_a.type.shape()):
+                    err_handler(
+                        call_a,
+                        f"expected argument of type '{sig_a.type}', "
+                        f"but got type '{call_a.type}'",
+                    )
+
+                # ensure scalars are simply variable names
+                elif (
+                    call_a.type.is_real_scalar()
+                    and not isinstance(call_a, LoopIR.ReadConfig)
+                    and not (isinstance(call_a, LoopIR.Read) and len(call_a.idx) == 0)
+                ):
+                    err_handler(
+                        call_a,
+                        "expected scalar arguments "
+                        "to be simply variable names "
+                        "for now",
+                    )
+            else:
+                err_handler(
+                    call_a,
+                    "expected numeric type expression, " f"but got type {call_a.type}",
+                )
+        else:
+            assert False, "bad argument type case"
+
+
 class TypeChecker:
     def __init__(self, proc):
         self.uast_proc = proc
@@ -235,7 +292,6 @@ def check_single_stmt(self, stmt):
             return [LoopIR.WriteConfig(stmt.config, stmt.field, rhs, stmt.srcinfo)]
         elif isinstance(stmt, UAST.Pass):
             return [LoopIR.Pass(stmt.srcinfo)]
-
         elif isinstance(stmt, UAST.If):
             cond = self.check_e(stmt.cond, is_index=True)
             if cond.type != T.err and cond.type != T.bool:
@@ -290,60 +346,11 @@ def check_single_stmt(self, stmt):
                 for call_a, sig_a in zip(stmt.args, stmt.f.args)
             ]
 
-            for call_a, sig_a in zip(args, stmt.f.args):
-                if call_a.type == T.err:
-                    pass
-                elif sig_a.type is T.size or sig_a.type is T.index:
-                    if not call_a.type.is_indexable():
-                        self.err(
-                            call_a,
-                            "expected size or index type "
-                            "expression, "
-                            f"but got type {call_a.type}",
-                        )
-
-                elif sig_a.type is T.bool:
-                    if not call_a.type is T.bool:
-                        self.err(
-                            call_a,
-                            "expected bool-type variable, "
-                            f"but got type {call_a.type}",
-                        )
-
-                elif sig_a.type is T.stride:
-                    if not call_a.type.is_stridable():
-                        self.err(
-                            call_a,
-                            "expected stride-type variable, "
-                            f"but got type {call_a.type}",
-                        )
-
-                elif sig_a.type.is_numeric():
-                    if len(call_a.type.shape()) != len(sig_a.type.shape()):
-                        self.err(
-                            call_a,
-                            f"expected argument of type '{sig_a.type}', "
-                            f"but got '{call_a.type}'",
-                        )
-
-                    # ensure scalars are simply variable names
-                    elif call_a.type.is_real_scalar():
-                        if not isinstance(call_a, LoopIR.ReadConfig) and not (
-                            isinstance(call_a, LoopIR.Read) and len(call_a.idx) == 0
-                        ):
-                            self.err(
-                                call_a,
-                                "expected scalar arguments "
-                                "to be simply variable names "
-                                "for now",
-                            )
-
-                else:
-                    assert False, "bad argument type case"
+            check_call_types(self.err, args, stmt.f.args)
 
             return [LoopIR.Call(stmt.f, args, stmt.srcinfo)]
         else:
-            assert False, "not a loopir in check_stmts"
+            assert False, f"not a loopir in check_stmts {type(stmt)}"
 
     def check_w_access(self, e, orig_hi):
         if isinstance(e, UAST.Point):

tests/golden/test_schedules/test_insert_noop_call.txt

@@ -0,0 +1,5 @@
+def foo(n: size, x: i8[n] @ DRAM, locality_hint: size):
+    assert locality_hint >= 0
+    assert locality_hint < 8
+    prefetch(x[1:2], locality_hint)
+    pass

tests/test_schedules.py

@@ -2,8 +2,7 @@
 
 import pytest
 
-from exo import ParseFragmentError
-from exo import proc, DRAM, Procedure, config
+from exo import ParseFragmentError, proc, DRAM, Procedure, config
 from exo.libs.memories import GEMM_SCRATCH
 from exo.stdlib.scheduling import *
 from exo.platforms.x86 import *
@@ -4630,6 +4629,33 @@ def foo(x: f32[30], result: f32):
     assert str(simplify(foo)) == golden
 
 
+def test_insert_noop_call(golden):
+    @proc
+    def foo(n: size, x: i8[n], locality_hint: size):
+        assert locality_hint >= 0
+        assert locality_hint < 8
+        pass
+
+    foo = insert_noop_call(
+        foo, foo.find("pass").before(), prefetch, ["x[1:2]", "locality_hint"]
+    )
+    assert str(foo) == golden
+
+
+def test_insert_noop_call_bad_args():
+    @proc
+    def foo(n: size, x: i8[n], locality_hint: size):
+        pass
+
+    with pytest.raises(TypeError, match="Function argument count mismatch"):
+        insert_noop_call(foo, foo.find("pass").before(), prefetch, [])
+
+    with pytest.raises(SchedulingError, match="Function argument type mismatch"):
+        insert_noop_call(
+            foo, foo.find("pass").before(), prefetch, ["n", "locality_hint"]
+        )
+
+
 def test_old_lift_alloc_config(golden):
     @config
     class CFG:

tests/test_typecheck.py

@@ -497,3 +497,15 @@ def dot(m: size, x: R[1, 1], y: R[m]):
         @proc
         def proj(n: size, x: R[100, 10, 1], y: R[10, n]):
             dot(n, x[1], y[0])
+
+
+def test_numeric_type_mismatch():
+    with pytest.raises(TypeError, match="but got type size"):
+
+        @proc
+        def bar(n: R):
+            pass
+
+        @proc
+        def foo(n: size):
+            bar(n)