Add tensor-based TileArray programs and GEMM pattern rewriting

.github/workflows/test.yml

@@ -24,28 +24,7 @@ env:
   CCACHE_MAXSIZE: 4G
 
 jobs:
-  python-unit-tests:
-    runs-on: ubuntu-latest
-
-    steps:
-      - name: Checkout
-        uses: actions/checkout@v7
-
-      - name: Set up Python
-        uses: actions/setup-python@v6
-        with:
-          python-version: 3.11.13
-
-      - name: Install Synapse and test dependencies
-        run: |
-          python -m pip install --upgrade pip
-          python -m pip install -e .
-          python -m pip install pytest
-
-      - name: Run frontend and language tests
-        run: python -m pytest -q tests/python/frontend tests/python/language
-
-  compiler-integration:
+  tests:
     runs-on: ubuntu-22.04
     timeout-minutes: 240
 
@@ -85,7 +64,6 @@ jobs:
       - name: Install Python build and test dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install pybind11==2.13.6 nanobind==2.15.0 pytest
           python -m pip install -e .
 
       - name: Restore ccache
@@ -173,8 +151,8 @@ jobs:
           test -x build/amoeba/tools/mlir-amoeba-opt/mlir-amoeba-opt
           test -d build/amoeba/python_packages/amoeba_core/taskflow_mlir
 
-      - name: Run compiler integration tests
-        run: python -m pytest -q tests/python/compiler
+      - name: Run all tests
+        run: python -m pytest -q tests
 
       - name: Save ccache
         if: steps.ccache.outputs.cache-hit != 'true'

.gitignore

@@ -1,5 +1,7 @@
 # Local docs / planning artifacts
 planbook/
+docs/
+pupa/
 
 # Python bytecode and caches
 __pycache__/
@@ -10,6 +12,7 @@ __pycache__/
 *.egg-info/
 build/
 dist/
+docs/
 .eggs/
 
 # Test / type-check / coverage caches

pyproject.toml

@@ -7,6 +7,16 @@ name = "synapse"
 version = "0.0.0"
 requires-python = ">=3.11"
 
+dependencies = [
+    "numpy==2.1.2",
+    "PyYAML==6.0.1",
+    "ml_dtypes==0.6.0",
+    "pybind11==2.13.6",
+    "nanobind==2.15.0",
+    "pytest==9.1.1",
+    "wheel"
+]
+
 [tool.setuptools]
 package-dir = {"" = "python"}
 

python/synapse/__init__.py

@@ -1,3 +1,3 @@
-from .compiler import compile
+from .compiler import compile, rewrite
 
-__all__ = ["compile"]
+__all__ = ["compile", "rewrite"]

python/synapse/compiler/__init__.py

@@ -1,3 +1,3 @@
-from .compiler import compile
+from .compiler import compile, rewrite
 
-__all__ = ["compile"]
+__all__ = ["compile", "rewrite"]

python/synapse/compiler/compiler.py

@@ -1,27 +1,68 @@
 """Top-Level Synapse Compilation Flow."""
 
 import subprocess
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from pathlib import Path
 from tempfile import TemporaryDirectory
 
 from synapse.frontend.lowering import lower
+from synapse.language.types import TensorType
+from synapse.patterns import TileArrayRewritePattern
 
 
-def compile(program: Callable, *, target: str) -> str:
-    """Compile a Synapse program for the selected backend."""
+def compile(
+    program: Callable | str,
+    *,
+    target: str,
+    argument_types: tuple[TensorType, ...] = (),
+    patterns: Sequence[type[TileArrayRewritePattern]] | None = None,
+) -> str:
+    """Compiles a TileArray function or bufferized task IR for the backend."""
 
-    # We only support the Neura backend for now, so we raise an error if the user tries to compile for any other target.
+    # The current compilation path targets Neura.
     if target != "neura":
         raise ValueError(f"unsupported compilation target: {target}")
-    # TODO: Support the amoeba backend.
-
-    neura_ir = lower(program)
+    if isinstance(program, str):
+        if argument_types:
+            raise ValueError("IR inputs already carry their argument types")
+        neura_ir = rewrite(program, patterns=patterns)
+    else:
+        if patterns is not None:
+            raise ValueError("rewrite patterns apply to IR inputs")
+        neura_ir = lower(program, argument_types=argument_types)
     return _run_neura_backend(neura_ir)
 
 
+def rewrite(
+    source: str,
+    *,
+    patterns: Sequence[type[TileArrayRewritePattern]] | None = None,
+) -> str:
+    """Applies patterns to task IR while preserving unmatched computations."""
+    from taskflow_mlir.dialects import neura, taskflow
+    from taskflow_mlir.ir import Context, Location, Module
+
+    from synapse.compiler.pattern_rewriter import apply_patterns
+    from synapse.patterns.gemm_pattern import (
+        AffineGemmPattern,
+        LinalgGemmPattern,
+        LinalgGenericGemmPattern,
+    )
+
+    if patterns is None:
+        patterns = [LinalgGemmPattern, LinalgGenericGemmPattern, AffineGemmPattern]
+    with Context(), Location.unknown():
+        taskflow.register_dialect()
+        neura.register_dialect()
+        module = Module.parse(source)
+        apply_patterns(module, patterns)
+        if not module.operation.verify():
+            raise ValueError("rewritten module failed verification")
+        return str(module)
+
+
 def _run_neura_backend(neura_ir: str) -> str:
-    """Legalize Neura values, insert data movement, and run template mapping."""
+    """Legalizes values, inserts data movement, and runs template mapping."""
 
     repository_root = Path(__file__).resolve().parents[3]
     amoeba_opt = (
@@ -36,11 +77,24 @@ def _run_neura_backend(neura_ir: str) -> str:
     if not amoeba_opt.is_file():
         raise FileNotFoundError(f"Amoeba compiler is not built: {amoeba_opt}")
 
+    architecture_spec = (
+        repository_root
+        / "mlir"
+        / "amoeba"
+        / "thirdparty"
+        / "neura"
+        / "test"
+        / "arch_spec"
+        / "architecture.yaml"
+    )
+
     with TemporaryDirectory(prefix="synapse-") as temporary_directory:
         output_path = Path(temporary_directory) / "mapped.mlir"
 
         command = [
             str(amoeba_opt),
+            f"--neura-architecture-spec={architecture_spec}",
+            "--promote-input-arg-to-const",
             "--leverage-predicated-value",
             "--insert-data-mov",
             (
@@ -58,6 +112,7 @@ def _run_neura_backend(neura_ir: str) -> str:
             capture_output=True,
             text=True,
             check=False,
+            cwd=temporary_directory,
         )
 
         if completed.returncode != 0: