Hikettei/polyhedral schedule (#13)

* fix: Final formatting fix

* test: Accept ruff formatting for test_fusion.py

* test: Accept ruff formatting for test_fusion.py (final)

* fix: Format domain.py to satisfy ruff

* update

* fix: Final formatting fix for example

* fix: Final format fix accept

* update

* update

Author: hikettei

.gitignore

@@ -2,7 +2,7 @@
 __pycache__/
 *.py[cod]
 *$py.class
-
+.DS_Store
 # Distribution / packaging
 .Python
 build/

AGENTS.md

@@ -35,6 +35,9 @@
 - 既存コードへの余計な修正禁止。Set 単体で完結する API は `caten/isl/specs/set.py` 内で完結させる。UnionSet 依存 API は関数名だけ置いて保留可。自動生成手法は禁止。
 - 進捗と作業計画を常に本ファイルに記録し更新すること（型ごとに完了状況や今後の順番を明記）。最新の計画がここに存在する状態を保つ。
 
+## Polyhedral DSL Guidelines
+- Prefer using Mixin operator overloads (e.g., `A | B` instead of `A.union(B)`) for cleaner code in user scripts and DSL implementations.
+
 ## 作業計画と進捗 (2025-11-16)
 直近のギャップ集計: `docs/ISL_missing_apis.md`（2025-11-16 再生成、欠落API 2047件）。map 残 2 件（tuple_name系シンボル未提供のみ、libisl非存在）。
 優先順とステータス（✅完了 / 🚧着手中 / ⏳未着手）

caten/polyhedral/__init__.py

@@ -6,6 +6,7 @@
 from .schedule_tree.filter import filter
 from .schedule_tree.mark import mark
 from .schedule_tree.sequence import sequence
+from .stmt import stmt
 
 __all__ = [
     "domain",
@@ -16,4 +17,5 @@
     "schedule",
     "compute_flow",
     "to_c",
+    "stmt",
 ]

caten/polyhedral/context.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import contextvars
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 if TYPE_CHECKING:
     import caten.isl as I
@@ -10,6 +10,7 @@ class ScheduleBuilder:
     def __init__(self) -> None:
         self.current_node: Optional["I.ScheduleNode"] = None
         self.schedule: Optional["I.Schedule"] = None
+        self.current_domain: Any = None
 
 _builder_ctx: contextvars.ContextVar[Optional[ScheduleBuilder]] = contextvars.ContextVar("schedule_builder", default=None)
 

caten/polyhedral/poly_schedule.py

@@ -0,0 +1,65 @@
+from typing import Optional
+
+import caten.isl as I
+from caten.polyhedral.analysis import compute_dependence_relation, schedule_is_legal_p
+from caten.polyhedral.codegen import to_c
+
+
+class PolyhedralSchedule:
+    def __init__(self, schedule: "I.Schedule", reads: Optional["I.UnionMap"] = None, writes: Optional["I.UnionMap"] = None) -> None:
+        self.isl_schedule = schedule
+        self.reads = reads
+        self.writes = writes
+        self.raw_dep: Optional["I.UnionMap"] = None
+        self.total_dep: Optional["I.UnionMap"] = None
+        
+        if reads and writes:
+            self.compute_dependencies()
+
+    def compute_dependencies(self) -> None:
+        if not self.reads or not self.writes:
+            return
+        total, raw, waw, war = compute_dependence_relation(self.reads, self.writes, self.isl_schedule)
+        self.raw_dep = raw
+        self.total_dep = total
+        
+    def is_legal(self) -> bool:
+        # Check legality against RAW dependencies
+        if self.raw_dep:
+            return schedule_is_legal_p(self.isl_schedule, self.raw_dep)
+        return True
+
+    def get_root(self) -> "I.ScheduleNode":
+        return self.isl_schedule.get_root()
+
+    def to_c(self) -> str:
+        return to_c(self.isl_schedule)
+
+    def __str__(self) -> str:
+        return str(self.isl_schedule)
+
+    def update(self, node: "I.ScheduleNode") -> None:
+        """Update the internal schedule from a modified schedule node."""
+        self.isl_schedule = node.get_schedule()
+
+    def sequence(self, other: "PolyhedralSchedule") -> "PolyhedralSchedule":
+        """Combine this schedule with another using isl_schedule_sequence."""
+        new_sched = self.isl_schedule.sequence(other.isl_schedule)
+        
+        new_reads = None
+        if self.reads and other.reads:
+            new_reads = self.reads.union(other.reads)
+        elif self.reads:
+            new_reads = self.reads
+        elif other.reads:
+            new_reads = other.reads
+            
+        new_writes = None
+        if self.writes and other.writes:
+            new_writes = self.writes.union(other.writes)
+        elif self.writes:
+            new_writes = self.writes
+        elif other.writes:
+            new_writes = other.writes
+            
+        return PolyhedralSchedule(new_sched, reads=new_reads, writes=new_writes)

caten/polyhedral/schedule_tree/domain.py

@@ -157,14 +157,38 @@ def __enter__(self) -> "domain":
         # We set current_node to the child of Domain (the Leaf)
         builder.current_node = sched.get_root().child(0)
 
+        self._prev_domain = builder.current_domain
+        builder.current_domain = self
+        
         return self
 
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         builder = get_builder()
         if builder.current_node:
             self.schedule = builder.current_node.get_schedule()
         builder.current_node = None
+        builder.current_domain = self._prev_domain
+
+    def finalize(self, read: Optional[Union[str, "I.UnionMap"]] = None, write: Optional[Union[str, "I.UnionMap"]] = None) -> Any:
+        from ..poly_schedule import PolyhedralSchedule
+        
+        if self.schedule is None:
+             if self.domain_set:
+                 uset = self.domain_set
+                 if isinstance(uset, str):
+                     uset = I.UnionSet(uset)
+                 elif isinstance(uset, I.Set):
+                     uset = I.UnionSet.from_set(uset)
+                 self.schedule = I.Schedule.from_domain(uset)
+             else:
+                 raise RuntimeError("No domain set for schedule.")
+
+        r = read if read else self.reads_map
+        if isinstance(r, str):
+            r = I.UnionMap(r)
+        
+        w = write if write else self.writes_map
+        if isinstance(w, str):
+            w = I.UnionMap(w)
 
-    def finalize(self, op_context: Any = None) -> Any:
-        # Placeholder for Kernel creation logic
-        return self.schedule
+        return PolyhedralSchedule(self.schedule, reads=r, writes=w)

caten/polyhedral/stmt.py

@@ -0,0 +1,77 @@
+import re
+from typing import List, Optional, Tuple
+
+import caten.isl as I
+
+from .context import get_builder
+
+
+def stmt(expr: str) -> None:
+    dom = get_builder().current_domain
+    if dom is None:
+        raise RuntimeError("stmt() must be used within a P.domain context")
+        
+    if "=" not in expr:
+        raise ValueError(f"Invalid statement expression (must contain assignment '='): {expr}")
+        
+    lhs_str, rhs_str = expr.split("=", 1)
+        
+    def extract_accesses(s: str) -> List[Tuple[str, str]]:
+        return re.findall(r"([a-zA-Z_]\w*)\s*\[(.*?)\]", s)
+        
+    writes = extract_accesses(lhs_str)
+    reads = extract_accesses(rhs_str)
+    
+    uset = dom.domain_set
+    if isinstance(uset, str):
+        uset = I.UnionSet(uset)
+    elif isinstance(uset, I.Set):
+        uset = I.UnionSet.from_set(uset)
+    
+    new_reads: Optional["I.UnionMap"] = None
+    new_writes: Optional["I.UnionMap"] = None
+    
+    def process_set(s: "I.Set") -> None:
+        nonlocal new_reads, new_writes
+        s_str = str(s)
+        if ":" in s_str:
+            tuple_part = s_str.split(":")[0].strip()
+            if tuple_part.startswith("{"):
+                tuple_part = tuple_part[1:].strip()
+        else:
+            tuple_part = s_str.strip()
+            if tuple_part.startswith("{") and tuple_part.endswith("}"):
+                tuple_part = tuple_part[1:-1].strip()
+                
+        for (name, indices) in writes:
+            m_str = f"{{ {tuple_part} -> {name}[{indices}] }}"
+            m = I.UnionMap(m_str)
+            if new_writes is None:
+                new_writes = m
+            else:
+                new_writes = new_writes.union(m)
+            
+        for (name, indices) in reads:
+            m_str = f"{{ {tuple_part} -> {name}[{indices}] }}"
+            m = I.UnionMap(m_str)
+            if new_reads is None:
+                new_reads = m
+            else:
+                new_reads = new_reads.union(m)
+
+    set_list = uset.get_set_list()
+    n = set_list.n_set()
+    for i in range(n):
+        process_set(set_list.get_at(i))
+    
+    if new_reads:
+        if dom.reads_map:
+            dom.reads_map = dom.reads_map.union(new_reads)
+        else:
+            dom.reads_map = new_reads
+        
+    if new_writes:
+        if dom.writes_map:
+            dom.writes_map = dom.writes_map.union(new_writes)
+        else:
+            dom.writes_map = new_writes