WIP: Generalize Memory to MemWin, SpecialWindow; WindowStmt upgrade

src/exo/API.py

@@ -13,7 +13,7 @@
 from .backend.LoopIR_compiler import run_compile, compile_to_strings
 from .core.configs import Config
 from .frontend.boundscheck import CheckBounds
-from .core.memory import Memory
+from .core.memory import MemWin, Memory, SpecialWindow
 from .frontend.parse_fragment import parse_fragment
 from .frontend.pattern_match import match_pattern
 from .core.prelude import *

src/exo/API_cursors.py

@@ -3,13 +3,13 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 
-from typing import List, Any
+from typing import List, Any, Optional, Type
 
 from . import API  # TODO: remove this circular import
 from .API_types import ExoType, loopir_type_to_exotype
 from .core.LoopIR import LoopIR
 from .core.configs import Config
-from .core.memory import Memory
+from .core.memory import MemWin, Memory, SpecialWindow
 
 from .core import internal_cursors as C
 from .frontend.pattern_match import match_pattern
@@ -65,7 +65,7 @@ class Cursor(ABC):
                | For( name : str, hi : Expr, body : Block )
                | Alloc( name : str, mem : Memory? )
                | Call( subproc : Procedure, args : ExprList )
-               | WindowStmt( name : str, winexpr : WindowExpr )
+               | WindowStmt( name : str, winexpr : WindowExpr, special_window : SpecialWindow? )
 
         Expr ::= Read( name : str, idx : ExprList )
                | ReadConfig( config : Config, field : str )
@@ -203,13 +203,13 @@ def name(self) -> str:
 
         return self._impl._node.name.name()
 
-    def mem(self) -> Memory:
+    def mem(self) -> MemWin:
         assert isinstance(self._impl, C.Node)
         assert isinstance(self._impl._node, LoopIR.fnarg)
         assert not self._impl._node.type.is_indexable()
 
         mem = self._impl._node.mem
-        assert issubclass(mem, Memory)
+        assert issubclass(mem, MemWin)
         return mem
 
     def is_tensor(self) -> bool:
@@ -657,7 +657,7 @@ class WindowStmtCursor(StmtCursor):
     """
     Cursor pointing to a window declaration statement:
         ```
-        name = winexpr
+        name = winexpr @ special_window
         ```
     """
 
@@ -673,6 +673,13 @@ def winexpr(self) -> ExprCursor:
 
         return WindowExprCursor(self._impl._child_node("rhs"), self._proc)
 
+    def special_window(self) -> Optional[Type[SpecialWindow]]:
+        assert isinstance(self._impl, C.Node)
+        assert isinstance(self._impl._node, LoopIR.WindowStmt)
+        special_window = self._impl._node.special_window
+        assert issubclass(special_window, SpecialWindow)
+        return special_window
+
 
 # --------------------------------------------------------------------------- #
 # --------------------------------------------------------------------------- #

src/exo/API_scheduling.py

@@ -1161,6 +1161,7 @@ def set_window(proc, cursor, is_window=True):
     return Procedure(ir, _provenance_eq_Procedure=proc, _forward=fwd)
 
 
+# TODO support SpecialWindow for arg cursor (but not alloc)
 @sched_op([ArgOrAllocCursorA, MemoryA])
 def set_memory(proc, cursor, memory_type):
     """

src/exo/__init__.py

@@ -10,7 +10,15 @@
 from .rewrite.LoopIR_scheduling import SchedulingError
 from .frontend.parse_fragment import ParseFragmentError
 from .core.configs import Config
-from .core.memory import Memory, DRAM
+from .core.memory import (
+    MemWin,
+    Memory,
+    SpecialWindow,
+    DRAM,
+    WindowStructCtx,
+    WindowFromDenseCtx,
+    memwin_template,
+)
 from .core.extern import Extern
 
 from . import stdlib
@@ -25,9 +33,14 @@
     "instr",
     "config",
     "Config",
+    "MemWin",
     "Memory",
-    "Extern",
+    "SpecialWindow",
+    "WindowStructCtx",
+    "WindowFromDenseCtx",
+    "memwin_template",
     "DRAM",
+    "Extern",
     "SchedulingError",
     "ParseFragmentError",
     #