Integrate Counter and General lowering into Synapse

python/synapse/__init__.py

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

python/synapse/compiler/__init__.py

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

python/synapse/compiler/compiler.py

@@ -1,21 +1,45 @@
 """Top-Level Synapse Compilation Flow."""
 
-import subprocess
 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 BufferType
 from synapse.patterns import TileArrayProgramPattern
 
+_NEURA_BACKEND_PIPELINE = """builtin.module(
+  convert-affine-to-taskflow,
+  func.func(construct-hyperblock-from-task),
+  classify-task-and-counter,
+  convert-taskflow-to-neura,
+  lower-affine,
+  convert-scf-to-cf,
+  convert-cf-to-llvm,
+  assign-accelerator,
+  lower-memref-to-neura,
+  lower-arith-to-neura,
+  lower-builtin-to-neura,
+  lower-llvm-to-neura,
+  canonicalize-return,
+  canonicalize-cast,
+  promote-input-arg-to-const,
+  fold-constant,
+  canonicalize-live-in,
+  leverage-predicated-value,
+  transform-ctrl-to-data-flow,
+  fold-constant,
+  insert-data-mov,
+  map-to-accelerator
+)"""
+
 
 def compile(
     program: Callable | str,
     *,
     target: str,
     argument_types: tuple[BufferType, ...] = (),
     patterns: Sequence[type[TileArrayProgramPattern]] | None = None,
+    architecture_spec: str | Path | None = None,
 ) -> str:
     """Compiles a TileArray function or bufferized task IR for the backend."""
 
@@ -25,22 +49,45 @@ def compile(
     if isinstance(program, str):
         if argument_types:
             raise ValueError("IR inputs already carry their argument types")
-        neura_ir = replace(program, patterns=patterns)
+        source = program
     else:
         if patterns is not None:
             raise ValueError("replacement patterns apply to IR inputs")
-        neura_ir = lower(program, argument_types=argument_types)
-    return _run_neura_backend(neura_ir)
+        source = lower(program, argument_types=argument_types)
+
+    if architecture_spec is None:
+        architecture_spec = (
+            Path(__file__).resolve().parents[3]
+            / "mlir"
+            / "amoeba"
+            / "thirdparty"
+            / "neura"
+            / "test"
+            / "arch_spec"
+            / "architecture.yaml"
+        )
+    architecture_spec = Path(architecture_spec)
+    if not architecture_spec.is_file():
+        raise FileNotFoundError(
+            f"Neura architecture specification does not exist: {architecture_spec}"
+        )
+    return _compile_task_graph(
+        source,
+        patterns=patterns,
+        architecture_spec=architecture_spec,
+    )
 
 
-def replace(
+def _compile_task_graph(
     source: str,
     *,
-    patterns: Sequence[type[TileArrayProgramPattern]] | None = None,
+    patterns: Sequence[type[TileArrayProgramPattern]] | None,
+    architecture_spec: Path,
 ) -> str:
-    """Applies patterns to task IR while preserving unmatched computations."""
+    """Applies patterns before and after Linalg-to-Affine conversion."""
     from taskflow_mlir.dialects import neura, taskflow
     from taskflow_mlir.ir import Context, Location, Module
+    from taskflow_mlir.passmanager import PassManager
 
     from synapse.compiler.pattern_replacement import apply_patterns
     from synapse.patterns.gemm_pattern import (
@@ -50,76 +97,26 @@ def replace(
     )
 
     if patterns is None:
-        patterns = [LinalgGemmPattern, LinalgGenericGemmPattern, AffineGemmPattern]
+        replacement_patterns = (
+            LinalgGemmPattern,
+            LinalgGenericGemmPattern,
+            AffineGemmPattern,
+        )
+    else:
+        replacement_patterns = patterns
+
+    taskflow.set_neura_architecture_spec(str(architecture_spec))
+
     with Context(), Location.unknown():
         taskflow.register_dialect()
         neura.register_dialect()
         module = Module.parse(source)
-        apply_patterns(module, patterns)
+        apply_patterns(module, replacement_patterns)
+        PassManager.parse(
+            "builtin.module(func.func(convert-linalg-to-affine-loops))"
+        ).run(module.operation)
+        apply_patterns(module, replacement_patterns)
+        PassManager.parse(_NEURA_BACKEND_PIPELINE).run(module.operation)
         if not module.operation.verify():
-            raise ValueError("replaced module failed verification")
+            raise ValueError("compiled task graph failed verification")
         return str(module)
-
-
-def _run_neura_backend(neura_ir: str) -> str:
-    """Legalizes values, inserts data movement, and runs template mapping."""
-
-    repository_root = Path(__file__).resolve().parents[3]
-    amoeba_opt = (
-        repository_root
-        / "build"
-        / "amoeba"
-        / "tools"
-        / "mlir-amoeba-opt"
-        / "mlir-amoeba-opt"
-    )
-
-    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",
-            (
-                "--map-to-accelerator="
-                "mapping-strategy=template "
-                "mapping-mode=spatial-only"
-            ),
-            "-o",
-            str(output_path),
-        ]
-
-        completed = subprocess.run(
-            command,
-            input=neura_ir,
-            capture_output=True,
-            text=True,
-            check=False,
-            cwd=temporary_directory,
-        )
-
-        if completed.returncode != 0:
-            diagnostics = "\n".join(
-                output for output in (completed.stdout, completed.stderr) if output
-            )
-
-            raise RuntimeError(f"Neura backend compilation failed:\n{diagnostics}")
-
-        return output_path.read_text(encoding="utf-8")

python/synapse/frontend/lowering.py

@@ -339,7 +339,7 @@ def get_kernel_metadata(self) -> DictAttr | None:
 
         stationary = self.program.stationary
         if stationary is None:
-            return None
+            return DictAttr.get({"kind": StringAttr.get("template")})
         return DictAttr.get(
             {
                 "kind": StringAttr.get("template"),

python/synapse/patterns/gemm_pattern.py

@@ -13,6 +13,7 @@
     IntegerType,
     MemRefType,
     OpResult,
+    OpView,
 )
 
 from synapse.library import ws_gemm_3x3
@@ -65,15 +66,15 @@ def _zero(value):
     )
 
 
-def _has_zero_output(operation, output):
-    """Checks the overwrite template's precondition on the initial accumulator.
+def _zero_output_initialization(operation, output) -> OpView | None:
+    """Finds the zero initialization consumed by the overwrite template.
 
     Named matmul and the recognized loop bodies compute C += A @ B, whereas
     the WS program computes C = A @ B. A preceding zero fill makes them agree.
     """
     parent = operation.operation.parent
     if parent is None:
-        return False
+        return None
     previous = None
     found = False
     for region in parent.regions:
@@ -90,19 +91,23 @@ def _has_zero_output(operation, output):
         if found:
             break
     if previous is None:
-        return False
+        return None
     if isinstance(previous, linalg.FillOp):
-        return tuple(previous.outputs) == (output,) and _zero(previous.inputs[0])
+        return (
+            previous
+            if tuple(previous.outputs) == (output,) and _zero(previous.inputs[0])
+            else None
+        )
 
     rows, columns = MemRefType(output.type).shape
     outer = _loop(previous, rows)
     if outer is None or len(tuple(outer.operations)) != 2:
-        return False
+        return None
     inner = _loop(outer.operations[0], columns)
     if inner is None or len(tuple(inner.operations)) != 2:
-        return False
+        return None
     store = inner.operations[0]
-    return (
+    matches = (
         store.operation.name == "affine.store"
         and store.operands[1] == output
         and _zero(store.operands[0])
@@ -112,6 +117,7 @@ def _has_zero_output(operation, output):
         )
         == (outer.arguments[0], inner.arguments[0])
     )
+    return previous if matches else None
 
 
 class LinalgGemmPattern(TileArrayProgramPattern):
@@ -130,13 +136,17 @@ def match_and_replace(cls, operation, replacer) -> bool:
             return False
         if any(value.type != expected for value in arguments):
             return False
-        if not _has_zero_output(operation, arguments[2]):
+        initialization = _zero_output_initialization(operation, arguments[2])
+        if initialization is None:
             return False
-        return replacer.replace_with_tile_array_program(
+        if not replacer.replace_with_tile_array_program(
             operation,
             program=ws_gemm_3x3,
             arguments=arguments,
-        )
+        ):
+            return False
+        replacer.erase_op(initialization)
+        return True
 
 
 class LinalgGenericGemmPattern(TileArrayProgramPattern):
@@ -191,13 +201,17 @@ def match_and_replace(cls, operation, replacer) -> bool:
             return False
         if any(value.type != expected for value in arguments):
             return False
-        if not _has_zero_output(operation, arguments[2]):
+        initialization = _zero_output_initialization(operation, arguments[2])
+        if initialization is None:
             return False
-        return replacer.replace_with_tile_array_program(
+        if not replacer.replace_with_tile_array_program(
             operation,
             program=ws_gemm_3x3,
             arguments=arguments,
-        )
+        ):
+            return False
+        replacer.erase_op(initialization)
+        return True
 
 
 class AffineGemmPattern(TileArrayProgramPattern):
@@ -258,7 +272,10 @@ def match_and_replace(cls, operation, replacer) -> bool:
             or tuple(
                 multiply.operands[index] for index in range(len(multiply.operands))
             )
-            != (lhs.results[0], rhs.results[0])
+            not in (
+                (lhs.results[0], rhs.results[0]),
+                (rhs.results[0], lhs.results[0]),
+            )
             or set(add.operands[index] for index in range(len(add.operands)))
             != {current.results[0], multiply.results[0]}
             or tuple(store.operands[index] for index in range(len(store.operands)))[:2]
@@ -269,10 +286,14 @@ def match_and_replace(cls, operation, replacer) -> bool:
         expected = MemRefType.get([3, 3], IntegerType.get_signless(32))
         if any(value.type != expected for value in arguments):
             return False
-        if not _has_zero_output(operation, arguments[2]):
+        initialization = _zero_output_initialization(operation, arguments[2])
+        if initialization is None:
             return False
-        return replacer.replace_with_tile_array_program(
+        if not replacer.replace_with_tile_array_program(
             operation,
             program=ws_gemm_3x3,
             arguments=arguments,
-        )
+        ):
+            return False
+        replacer.erase_op(initialization)
+        return True

tests/python/compiler/test_add_constant_kernel.py

@@ -6,7 +6,7 @@
 module {
   func.func @add_constant() {
     taskflow.task @add_constant : () -> () {
-      neura.kernel attributes {accelerator = "neura"} {
+      neura.kernel attributes {accelerator = "neura", kernel_metadata = {kind = "template"}} {
         %0 = "neura.constant"() <{value = 1 : i32}> {placement = {x = 0 : i32, y = 0 : i32}} : () -> i32
         %1 = "neura.constant"() <{value = 2 : i32}> {placement = {x = 2 : i32, y = 0 : i32}} : () -> i32
         %2 = "neura.add"(%0, %1) {placement = {x = 1 : i32, y = 0 : i32}} : (i32, i32) -> i32
@@ -24,13 +24,11 @@
 module {
   func.func @add_constant() {
     taskflow.task @add_constant : () -> () {
-      neura.kernel attributes {accelerator = "neura", mapping_info = {compiled_ii = 1 : i32, mapping_mode = "spatial-only", mapping_strategy = "template", rec_mii = 1 : i32, res_mii = 1 : i32, x_tiles = 4 : i32, y_tiles = 4 : i32}} {
+      neura.kernel attributes {accelerator = "neura", kernel_metadata = {kind = "template"}, mapping_info = {compiled_ii = 1 : i32, mapping_mode = "spatial-only", mapping_strategy = "template", rec_mii = 1 : i32, res_mii = 1 : i32, x_tiles = 4 : i32, y_tiles = 4 : i32}} {
         %0 = "neura.constant"() <{value = 1 : i32}> {dfg_id = 0 : i32, mapping_locs = [{id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 0 : i32, x = 0 : i32, y = 0 : i32}]} : () -> !neura.data<i32, i1>
-        %1 = "neura.constant"() <{value = 2 : i32}> {dfg_id = 1 : i32, mapping_locs = [{id = 2 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 0 : i32, x = 2 : i32, y = 0 : i32}]} : () -> !neura.data<i32, i1>
-        %2 = "neura.data_mov"(%0) {dfg_id = 3 : i32, mapping_locs = [{id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 0 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-        %3 = "neura.data_mov"(%1) {dfg_id = 4 : i32, mapping_locs = [{id = 5 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 0 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-        %4 = "neura.add"(%2, %3) {dfg_id = 5 : i32, mapping_locs = [{id = 1 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 1 : i32, x = 1 : i32, y = 0 : i32}]} : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-        neura.yield {dfg_id = 2 : i32}
+        %1 = "neura.data_mov"(%0) {dfg_id = 2 : i32, mapping_locs = [{id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 0 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+        %2 = "neura.add"(%1) {dfg_id = 3 : i32, mapping_locs = [{id = 1 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 1 : i32, x = 1 : i32, y = 0 : i32}], rhs_value = 2 : i32} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+        neura.yield {dfg_id = 1 : i32}
       }
       taskflow.yield
     }