From f97e0e7be79d04c0ab537f58500eb14bc6920be5 Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Wed, 12 Aug 2026 13:34:39 +0800 Subject: [PATCH 01/11] Add initial TileArray spatial abstraction --- python/synapse/language/__init__.py | 3 ++ python/synapse/language/spatial.py | 58 +++++++++++++++++++++++++++++ tests/language/test_spatial.py | 38 +++++++++++++++++++ 3 files changed, 99 insertions(+) create mode 100644 python/synapse/language/__init__.py create mode 100644 python/synapse/language/spatial.py create mode 100644 tests/language/test_spatial.py diff --git a/python/synapse/language/__init__.py b/python/synapse/language/__init__.py new file mode 100644 index 0000000..6897311 --- /dev/null +++ b/python/synapse/language/__init__.py @@ -0,0 +1,3 @@ +from .spatial import TileArray + +__all__ = ["TileArray"] \ No newline at end of file diff --git a/python/synapse/language/spatial.py b/python/synapse/language/spatial.py new file mode 100644 index 0000000..76797da --- /dev/null +++ b/python/synapse/language/spatial.py @@ -0,0 +1,58 @@ +"""Hardware spatial features exposed by the Synapse language. +This file exposes spatial structures that programmers can use to organize +computation and data movement according to the target hardware hierarchy. + +TileArray currently exposes the two-dimensional tile array of a CGRA. + +Future spatial abstractions may expose inter-core structures, such as the +core array of a multi-CGRA, AMD AIE/NPU, or Tenstorrent. +""" + + +class Tile: + """A hardware tile in a CGRA TileArray. + + A Tile is owned by a TileArray. Its row and column identify the + corresponding position in the target CGRA tile array. + """ + + def __init__(self, row: int, col: int): + self.row = row + self.col = col + + +class TileArray: + """A parameterized two-dimensional tile array. + + In the initial implementation, ``rows`` and ``cols`` must match the + dimensions of the target CGRA. A tile accessed as ``array[row, col]`` + corresponds directly to the tile at that hardware coordinate. + + Example: + array = TileArray(rows=4, cols=4) + tile = array[1, 2] + """ + + def __init__(self, rows: int, cols: int): + self.rows = rows + self.cols = cols + + self._tiles = [ + [Tile(row=row, col=col) for col in range(cols)] for row in range(rows) + ] + + def tiles(self): + """Iterate over all tiles in the array. + + The iteration order is a Python programming convenience and + does not specify sequential hardware execution. + """ + for row_tiles in self._tiles: + yield from row_tiles + + def __getitem__(self, coordinate: tuple[int, int]) -> Tile: + """ + Return the tile at the given ``(row, col)`` coordinate + """ + row, col = coordinate + return self._tiles[row][col] diff --git a/tests/language/test_spatial.py b/tests/language/test_spatial.py new file mode 100644 index 0000000..358b818 --- /dev/null +++ b/tests/language/test_spatial.py @@ -0,0 +1,38 @@ +import synapse.language as synl + + +def test_tile_array_is_parameterized(): + array_2x3 = synl.TileArray(rows=2, cols=3) + array_4x4 = synl.TileArray(rows=4, cols=4) + + assert array_2x3.rows == 2 + assert array_2x3.cols == 3 + + assert array_4x4.rows == 4 + assert array_4x4.cols == 4 + + +def test_tile_array_exposes_logical_tiles(): + array = synl.TileArray(2, 3) + + coordinates = {(tile.row, tile.col) for tile in array.tiles()} + + assert coordinates == {(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)} + + +def test_access_tile_array_by_coordinate(): + array = synl.TileArray(2, 3) + tile = array[1, 2] + + assert tile.row == 1 + assert tile.col == 2 + + assert tile is array[1, 2] + + enumerated_tile = next( + candidate + for candidate in array.tiles() + if candidate.row == 1 and candidate.col == 2 + ) + + assert tile is enumerated_tile From 3745cf3ad8807532bf026b25709bcba5b420954e Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Wed, 12 Aug 2026 13:57:17 +0800 Subject: [PATCH 02/11] Update MLIR dependencies for Amoeba-Neura split --- .gitmodules | 6 +++--- mlir/README.md | 47 +++++++++++++++++++++++++++++++++++------------ mlir/amoeba | 1 + mlir/neura | 1 - 4 files changed, 39 insertions(+), 16 deletions(-) create mode 160000 mlir/amoeba delete mode 160000 mlir/neura diff --git a/.gitmodules b/.gitmodules index 0fd26b1..7e9f9ca 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,3 +1,3 @@ -[submodule "mlir/neura"] - path = mlir/neura - url = https://github.com/coredac/neura +[submodule "mlir/amoeba"] + path = mlir/amoeba + url = https://github.com/coredac/amoeba diff --git a/mlir/README.md b/mlir/README.md index e49c9ba..87dbc4d 100644 --- a/mlir/README.md +++ b/mlir/README.md @@ -1,23 +1,46 @@ -# SYNAPSE MLIR Submodules +# SYNAPSE MLIR Dependencies -This directory is reserved for MLIR/compiler submodules used by SYNAPSE. +This directory contains the compiler projects used by SYNAPSE after its Python +frontend has captured a program. -Current layout: +## Dependency layout ```text mlir/ - neura/ # git submodule: https://github.com/coredac/neura + amoeba/ # git submodule: https://github.com/coredac/amoeba + thirdparty/neura/ # nested submodule managed by Amoeba ``` -SYNAPSE itself only defines the programming model, frontend, internal graph IR, -and lowering to Taskflow IR. Compilation from Taskflow IR to NEURA or lower -hardware/compiler targets belongs to the NEURA/Taskflow compiler stack in this -directory. +SYNAPSE owns the programming model, Python frontend, and lowering of a captured +program into compiler IR. Amoeba provides the backend-neutral Taskflow dialect +and backend integration. Neura is Amoeba's CGRA backend and provides the +single-CGRA dialect, mapping, routing, register allocation, and code generation. -The submodule is registered in `.gitmodules`: +SYNAPSE therefore depends on Amoeba rather than carrying a second, independent +Neura checkout. This keeps the Taskflow and Neura interfaces on the revisions +tested together by Amoeba. + +## Initial single-task compilation boundary + +The first CGRA path intentionally handles one task only. The frontend treats the +whole captured program as an implicit task and lowers it to the following +container hierarchy: ```text -[submodule "mlir/neura"] - path = mlir/neura - url = https://github.com/coredac/neura +taskflow.task + neura.kernel + mapped Neura operations +``` + +The operations inside `neura.kernel` are expected to carry post-mapping +information such as tile coordinates, time steps, links, and registers. The +single-task path does not yet define user-facing task syntax or perform +inter-task allocation, placement, scheduling, replication, or communication. + +## Checkout + +Initialize Amoeba and its Neura dependency recursively: + +```sh +git submodule update --init --recursive ``` diff --git a/mlir/amoeba b/mlir/amoeba new file mode 160000 index 0000000..f1599e3 --- /dev/null +++ b/mlir/amoeba @@ -0,0 +1 @@ +Subproject commit f1599e32b6eff5a41b06605528a2b6fcb563ff12 diff --git a/mlir/neura b/mlir/neura deleted file mode 160000 index 5763d85..0000000 --- a/mlir/neura +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 5763d850d2618392e14fe4bc4714d189ad906290 From 9a3f158fe3e039bce74b8f72be7b2e0b63d42b58 Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Thu, 13 Aug 2026 13:04:27 +0800 Subject: [PATCH 03/11] Add mapped neura IR for a single task --- README.md | 2 +- tests/mlir/single_task_mapped_neura.mlir | 23 +++++++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) create mode 100644 tests/mlir/single_task_mapped_neura.mlir diff --git a/README.md b/README.md index b03062d..5789d71 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,7 @@ python/synapse/ Python package source python/synapse/frontend Source capture and frontend parser utilities examples/ Small frontend examples tests/ Pytest tests -mlir/neura/ Downstream NEURA / Taskflow compiler stack +mlir/amoeba/ Downstream NEURA / Taskflow compiler stack ``` ## Setup diff --git a/tests/mlir/single_task_mapped_neura.mlir b/tests/mlir/single_task_mapped_neura.mlir new file mode 100644 index 0000000..9c9b525 --- /dev/null +++ b/tests/mlir/single_task_mapped_neura.mlir @@ -0,0 +1,23 @@ +// This file defines the expected post-mapping Neura IR for the +// Synapse single-task spatial program. + +module { + func.func @single_task() { + taskflow.task @Task_0 : () -> () { + neura.kernel attributes {accelerator = "neura", dataflow_mode = "predicate", mapping_info = {compiled_ii = 1 : i32, mapping_mode = "spatial-temporal", mapping_strategy = "manual", rec_mii = 1 : i32, res_mii = 1 : i32, x_tiles = 2 : i32, y_tiles = 1 : i32}} { + // Tile(row=0, col=0) produces one value at step 0. + %source = "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 + + // Move the value east to Tile(row=0, col=1) and keep it in local register 0. + %moved = "neura.data_mov"(%source) {dfg_id = 1 : i32, mapping_locs = [{id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 0 : i32}, {id = 32 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 1 : i32}]} : (!neura.data) -> !neura.data + + // Tile(row=0, col=1) consumes the transferred value at step 2. + %result = "neura.add"(%moved) {dfg_id = 2 : i32, mapping_locs = [{id = 1 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 2 : i32, x = 1 : i32, y = 0 : i32}], rhs_value = 1 : i32} : (!neura.data) -> !neura.data + + neura.yield + } + taskflow.yield + } + return + } +} \ No newline at end of file From 3cfd22ea4baf7f488ef0e559fa96a375fda77198 Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Mon, 17 Aug 2026 13:03:09 +0800 Subject: [PATCH 04/11] Uodate amoeba for neura template mapping --- mlir/amoeba | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/amoeba b/mlir/amoeba index f1599e3..89c451b 160000 --- a/mlir/amoeba +++ b/mlir/amoeba @@ -1 +1 @@ -Subproject commit f1599e32b6eff5a41b06605528a2b6fcb563ff12 +Subproject commit 89c451beadf325e8d996fae4acee03f3580ae364 From 2c50216d319960095d4c9e8ec481a4766ce26cf9 Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Mon, 17 Aug 2026 13:05:47 +0800 Subject: [PATCH 05/11] Reorganize python tests --- tests/mlir/single_task_mapped_neura.mlir | 23 --------------------- tests/{ => python}/frontend/test_parser.py | 0 tests/{ => python}/language/test_spatial.py | 0 3 files changed, 23 deletions(-) delete mode 100644 tests/mlir/single_task_mapped_neura.mlir rename tests/{ => python}/frontend/test_parser.py (100%) rename tests/{ => python}/language/test_spatial.py (100%) diff --git a/tests/mlir/single_task_mapped_neura.mlir b/tests/mlir/single_task_mapped_neura.mlir deleted file mode 100644 index 9c9b525..0000000 --- a/tests/mlir/single_task_mapped_neura.mlir +++ /dev/null @@ -1,23 +0,0 @@ -// This file defines the expected post-mapping Neura IR for the -// Synapse single-task spatial program. - -module { - func.func @single_task() { - taskflow.task @Task_0 : () -> () { - neura.kernel attributes {accelerator = "neura", dataflow_mode = "predicate", mapping_info = {compiled_ii = 1 : i32, mapping_mode = "spatial-temporal", mapping_strategy = "manual", rec_mii = 1 : i32, res_mii = 1 : i32, x_tiles = 2 : i32, y_tiles = 1 : i32}} { - // Tile(row=0, col=0) produces one value at step 0. - %source = "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 - - // Move the value east to Tile(row=0, col=1) and keep it in local register 0. - %moved = "neura.data_mov"(%source) {dfg_id = 1 : i32, mapping_locs = [{id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 0 : i32}, {id = 32 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 1 : i32}]} : (!neura.data) -> !neura.data - - // Tile(row=0, col=1) consumes the transferred value at step 2. - %result = "neura.add"(%moved) {dfg_id = 2 : i32, mapping_locs = [{id = 1 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 2 : i32, x = 1 : i32, y = 0 : i32}], rhs_value = 1 : i32} : (!neura.data) -> !neura.data - - neura.yield - } - taskflow.yield - } - return - } -} \ No newline at end of file diff --git a/tests/frontend/test_parser.py b/tests/python/frontend/test_parser.py similarity index 100% rename from tests/frontend/test_parser.py rename to tests/python/frontend/test_parser.py diff --git a/tests/language/test_spatial.py b/tests/python/language/test_spatial.py similarity index 100% rename from tests/language/test_spatial.py rename to tests/python/language/test_spatial.py From 694501a2c9211925b87fb2afd5ffe06b33276ada Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Wed, 19 Aug 2026 19:44:42 +0800 Subject: [PATCH 06/11] Add tile-array program representation --- python/synapse/__init__.py | 3 + python/synapse/language/__init__.py | 8 +- python/synapse/language/spatial.py | 6 +- python/synapse/language/tile_array_program.py | 261 ++++++++++++++++++ tests/python/language/test_spatial.py | 1 + .../language/test_tile_array_program.py | 48 ++++ 6 files changed, 324 insertions(+), 3 deletions(-) create mode 100644 python/synapse/language/tile_array_program.py create mode 100644 tests/python/language/test_tile_array_program.py diff --git a/python/synapse/__init__.py b/python/synapse/__init__.py index e69de29..236470b 100644 --- a/python/synapse/__init__.py +++ b/python/synapse/__init__.py @@ -0,0 +1,3 @@ +from .compiler import compile + +__all__ = ["compile"] diff --git a/python/synapse/language/__init__.py b/python/synapse/language/__init__.py index 6897311..eca34d4 100644 --- a/python/synapse/language/__init__.py +++ b/python/synapse/language/__init__.py @@ -1,3 +1,9 @@ +"""Public Synapse language API.""" + from .spatial import TileArray +from .tile_array_program import TileArrayScalarType, add, constant + +i32 = TileArrayScalarType.I32 +f32 = TileArrayScalarType.F32 -__all__ = ["TileArray"] \ No newline at end of file +__all__ = ["TileArray", "TileArrayScalarType", "add", "constant", "f32", "i32"] diff --git a/python/synapse/language/spatial.py b/python/synapse/language/spatial.py index 76797da..3d67180 100644 --- a/python/synapse/language/spatial.py +++ b/python/synapse/language/spatial.py @@ -16,7 +16,8 @@ class Tile: corresponding position in the target CGRA tile array. """ - def __init__(self, row: int, col: int): + def __init__(self, row: int, col: int, array: "TileArray"): + self.array = array self.row = row self.col = col @@ -38,7 +39,8 @@ def __init__(self, rows: int, cols: int): self.cols = cols self._tiles = [ - [Tile(row=row, col=col) for col in range(cols)] for row in range(rows) + [Tile(row=row, col=col, array=self) for col in range(cols)] + for row in range(rows) ] def tiles(self): diff --git a/python/synapse/language/tile_array_program.py b/python/synapse/language/tile_array_program.py new file mode 100644 index 0000000..12f0ff8 --- /dev/null +++ b/python/synapse/language/tile_array_program.py @@ -0,0 +1,261 @@ +"""Programming model for computations placed on a TileArray. + +This module defines the typed frontend representation of a tile-array program. +It records computation independently of MLIR. Compiler lowering later converts +the recorded program into Taskflow and Neura operations. +""" + +from __future__ import annotations + +from contextvars import ContextVar +from dataclasses import dataclass, field +from enum import Enum +from typing import TypeAlias + +from .spatial import Tile, TileArray + + +class TileArrayScalarType(str, Enum): + """The scalar types supported by the TileArray programming model. + This is intentionally independent of MLIR types. The lowering converts + these frontend types into the corresponding MLIR types. + + Additional scalar types can be added here as the language grows. + """ + + I32 = "i32" + F32 = "f32" + + +@dataclass(frozen=True) +class TileArrayValue: + """A typed value produced by one tile-array operation.""" + + id: int + dtype: TileArrayScalarType + _builder: TileArrayBuilder = field(repr=False) + + +# --------------------------------------------------------------- +# Typed tile-array operations +# --------------------------------------------------------------- + + +@dataclass(frozen=True) +class ConstantOp: + """A scalar constant produced by a tile-array operation.""" + + result: TileArrayValue + value: int | float + tile: Tile + + +@dataclass(frozen=True) +class AddOp: + """A scalar addition executed by a tile-array operation.""" + + result: TileArrayValue + lhs: TileArrayValue + rhs: TileArrayValue + tile: Tile + + +# This union explicitly lists every operation currently supported by the +# tile-array frontend. Future operations such as MacOp and GatherOp should be +# added here. +TileArrayOp: TypeAlias = ConstantOp | AddOp + + +@dataclass(frozen=True) +class TileArrayProgram: + """A tile-array program produced by TileArrayBuilder.""" + + array: TileArray + operations: tuple[TileArrayOp, ...] + + +# --------------------------------------------------------------- +# Internal program builder +# --------------------------------------------------------------- +class TileArrayBuilder: + """A tile-array program builder. + + The builder records typed operations in user-program order. Once build() + is called, it returns a TileArrayProgram. + """ + + def __init__(self): + self._array: TileArray | None = None + self._operations: list[TileArrayOp] = [] + self._next_value_id = 0 + self._token = None + self._is_built = False + + def __enter__(self): + """Make this builder active for tile-array DSL calls.""" + if _active_builder.get() is not None: + raise RuntimeError("Cannot enter a nested TileArrayBuilder context") + if self._token is not None: + raise RuntimeError("TileArrayBuilder context is already active") + self._token = _active_builder.set(self) + return self + + def __exit__(self, exc_type, exc_value, traceback) -> None: + """Restore the previously active builder.""" + if self._token is None: + raise RuntimeError("TileArrayBuilder is not active") + + _active_builder.reset(self._token) + self._token = None + + def _new_value(self, *, dtype: TileArrayScalarType) -> TileArrayValue: + """Allocate the next value identifier.""" + result = TileArrayValue(id=self._next_value_id, dtype=dtype, _builder=self) + self._next_value_id += 1 + return result + + def _bind_tile_array(self, tile: Tile) -> None: + """Bind the program to one TileArray.""" + if not isinstance(tile, Tile): + raise TypeError("tile must be a Tile") + if self._array is None: + self._array = tile.array + return + if tile.array is not self._array: + raise ValueError( + "all operations in a TileArrayProgram must use tiles from the same TileArray" + ) + + def _check_operand(self, operand: TileArrayValue) -> None: + """Verify that an operand was produced by this builder.""" + if not isinstance(operand, TileArrayValue): + raise TypeError("operation operand must be a TileArrayValue") + + if operand._builder is not self: + raise ValueError( + "operation operand belongs to a different TileArrayProgram" + ) + + def _check_can_emit(self) -> None: + """Reject operations emitted after the program has been finalized.""" + if self._is_built: + raise RuntimeError( + "cannot emit operations after building a TileArrayProgram" + ) + + def emit_constant( + self, value: int | float, *, dtype: TileArrayScalarType, tile: Tile + ) -> TileArrayValue: + """Record one scalar constant operation.""" + self._check_can_emit() + self._bind_tile_array(tile) + result = self._new_value(dtype=dtype) + + self._operations.append(ConstantOp(result=result, value=value, tile=tile)) + return result + + def emit_add( + self, lhs: TileArrayValue, rhs: TileArrayValue, *, tile: Tile + ) -> TileArrayValue: + """Record one scalar addition operation.""" + self._check_can_emit() + self._bind_tile_array(tile) + self._check_operand(lhs) + self._check_operand(rhs) + + if lhs.dtype != rhs.dtype: + raise TypeError("add operands must have the same tile-array value type") + + result = self._new_value(dtype=lhs.dtype) + + self._operations.append(AddOp(result=result, lhs=lhs, rhs=rhs, tile=tile)) + return result + + def build(self) -> TileArrayProgram: + """Finish recording and return a program.""" + if self._token is not None: + raise RuntimeError( + "cannot build a TileArrayProgram while its builder is active" + ) + if self._array is None: + raise RuntimeError( + "cannot build an empty TileArrayProgram without a TileArray" + ) + self._is_built = True + return TileArrayProgram(array=self._array, operations=tuple(self._operations)) + + +# The active builder is compiler-internal state. Public DSL calls use it to +# find the builder created by frontend lowering. +_active_builder: ContextVar[TileArrayBuilder | None] = ContextVar( + "active_tile_array_builder", + default=None, +) + + +def _require_active_builder() -> TileArrayBuilder: + """Return the active builder.""" + builder = _active_builder.get() + + if builder is None: + raise RuntimeError( + "tile-array operations must be called while lowering a Synapse program" + ) + + return builder + + +# --------------------------------------------------------------- +# User-facing tile-array program DSL +# --------------------------------------------------------------- +def constant( + value: int | float, *, tile: Tile, dtype: TileArrayScalarType | None = None +) -> TileArrayValue: + """Create a scalar constant on one hardware tile. + + Integer literals default to i32. Floating-point literals default to f32. + Use an explicit dtype when a different representation is required: + constant(1.0, tile=tile, dtype=TileArrayScalarType.F32) + """ + if isinstance(value, bool): + raise TypeError("boolean constants are not supported yet") + + if dtype is None: + if isinstance(value, int): + dtype = TileArrayScalarType.I32 + elif isinstance(value, float): + dtype = TileArrayScalarType.F32 + else: + raise TypeError( + "constant currently supports integer and floating-point values" + ) + + if not isinstance(dtype, TileArrayScalarType): + raise TypeError("dtype must be a TileArrayScalarType") + + if dtype == TileArrayScalarType.I32 and not isinstance(value, int): + raise TypeError("an i32 constant requires an integer value") + + if dtype in (TileArrayScalarType.F32,) and not isinstance(value, (int, float)): + raise TypeError("a floating-point constant requires a numeric value") + + return _require_active_builder().emit_constant( + value, + dtype=dtype, + tile=tile, + ) + + +def add( + lhs: TileArrayValue, + rhs: TileArrayValue, + *, + tile: Tile, +) -> TileArrayValue: + """Create a scalar addition on one hardware tile.""" + + return _require_active_builder().emit_add( + lhs, + rhs, + tile=tile, + ) diff --git a/tests/python/language/test_spatial.py b/tests/python/language/test_spatial.py index 358b818..0cace14 100644 --- a/tests/python/language/test_spatial.py +++ b/tests/python/language/test_spatial.py @@ -24,6 +24,7 @@ def test_access_tile_array_by_coordinate(): array = synl.TileArray(2, 3) tile = array[1, 2] + assert tile.array is array assert tile.row == 1 assert tile.col == 2 diff --git a/tests/python/language/test_tile_array_program.py b/tests/python/language/test_tile_array_program.py new file mode 100644 index 0000000..ccef8b7 --- /dev/null +++ b/tests/python/language/test_tile_array_program.py @@ -0,0 +1,48 @@ +import synapse.language as synl + + +def test_records_tile_array_program(): + array = synl.TileArray(4, 4) + builder = synl.tile_array_program.TileArrayBuilder() + + with builder: + lhs = synl.constant(1, tile=array[0, 0]) + rhs = synl.constant(2, tile=array[0, 2]) + result = synl.add(lhs, rhs, tile=array[0, 1]) + + program = builder.build() + lhs_op, rhs_op, add_op = program.operations + + assert program.array is array + assert isinstance(lhs_op, synl.tile_array_program.ConstantOp) + assert isinstance(rhs_op, synl.tile_array_program.ConstantOp) + assert isinstance(add_op, synl.tile_array_program.AddOp) + assert [(value.id, value.dtype) for value in (lhs, rhs, result)] == [ + (0, synl.i32), + (1, synl.i32), + (2, synl.i32), + ] + assert [op.tile for op in program.operations] == [ + array[0, 0], + array[0, 2], + array[0, 1], + ] + assert (add_op.lhs, add_op.rhs) == (lhs, rhs) + + +def test_infers_supported_scalar_types(): + array = synl.TileArray(1, 3) + builder = synl.tile_array_program.TileArrayBuilder() + + with builder: + integer = synl.constant(1, tile=array[0, 0]) + floating = synl.constant(1.0, tile=array[0, 1]) + explicit_f32 = synl.constant( + 1, + tile=array[0, 2], + dtype=synl.f32, + ) + + assert integer.dtype == synl.i32 + assert floating.dtype == synl.f32 + assert explicit_f32.dtype == synl.f32 From c87510f9e95c40e4367935bd7fcdab5858a842ea Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Wed, 19 Aug 2026 20:54:55 +0800 Subject: [PATCH 07/11] Add Python-to-Neura compiler flow --- examples/gemm.py | 2 + mlir/README.md | 20 +- mlir/amoeba | 2 +- pyproject.toml | 8 +- python/synapse/compiler/__init__.py | 3 + python/synapse/compiler/compiler.py | 70 +++++++ python/synapse/frontend/lowering.py | 187 ++++++++++++++++++ python/synapse/language/spatial.py | 42 ++-- .../compiler/test_add_constant_kernel.py | 61 ++++++ tests/python/language/test_spatial.py | 26 +-- .../language/test_tile_array_program.py | 4 +- 11 files changed, 383 insertions(+), 42 deletions(-) create mode 100644 python/synapse/compiler/__init__.py create mode 100644 python/synapse/compiler/compiler.py create mode 100644 python/synapse/frontend/lowering.py create mode 100644 tests/python/compiler/test_add_constant_kernel.py diff --git a/examples/gemm.py b/examples/gemm.py index b0bd60a..6fbba93 100644 --- a/examples/gemm.py +++ b/examples/gemm.py @@ -1,5 +1,6 @@ from synapse.frontend.parser import dump_ast + def gemm(A, B, C): for i in range(128): for j in range(128): @@ -8,5 +9,6 @@ def gemm(A, B, C): acc += A[i][k] * B[k][j] C[i][j] = acc + if __name__ == "__main__": print(dump_ast(gemm)) diff --git a/mlir/README.md b/mlir/README.md index 87dbc4d..a7e2ecc 100644 --- a/mlir/README.md +++ b/mlir/README.md @@ -29,13 +29,23 @@ container hierarchy: ```text taskflow.task neura.kernel - mapped Neura operations + placed Neura operations with ordinary MLIR value types ``` -The operations inside `neura.kernel` are expected to carry post-mapping -information such as tile coordinates, time steps, links, and registers. The -single-task path does not yet define user-facing task syntax or perform -inter-task allocation, placement, scheduling, replication, or communication. +The frontend fixes the spatial operation placement selected by the TileArray +program, but it does not construct Neura's predicated value type or final +mapping metadata. The backend compilation flow performs: + +```text +--leverage-predicated-value + -> --insert-data-mov + -> --map-to-accelerator="mapping-strategy=template mapping-mode=spatial-only" +``` + +The final mapped Neura IR contains tile coordinates, time steps, links, and +register information. The single-task path does not yet define user-facing +task syntax or perform inter-task allocation, placement, scheduling, +replication, or communication. ## Checkout diff --git a/mlir/amoeba b/mlir/amoeba index 89c451b..784cb87 160000 --- a/mlir/amoeba +++ b/mlir/amoeba @@ -1 +1 @@ -Subproject commit 89c451beadf325e8d996fae4acee03f3580ae364 +Subproject commit 784cb87a45f51574cb45088da25bd674bed03abb diff --git a/pyproject.toml b/pyproject.toml index e861013..f25db76 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,4 +11,10 @@ requires-python = ">=3.10" package-dir = {"" = "python"} [tool.setuptools.packages.find] -where = ["python"] \ No newline at end of file +where = ["python"] + +[tool.pytest.ini_options] +pythonpath = [ + "python", + "build/amoeba/python_packages/amoeba_core", +] diff --git a/python/synapse/compiler/__init__.py b/python/synapse/compiler/__init__.py new file mode 100644 index 0000000..236470b --- /dev/null +++ b/python/synapse/compiler/__init__.py @@ -0,0 +1,3 @@ +from .compiler import compile + +__all__ = ["compile"] diff --git a/python/synapse/compiler/compiler.py b/python/synapse/compiler/compiler.py new file mode 100644 index 0000000..1213525 --- /dev/null +++ b/python/synapse/compiler/compiler.py @@ -0,0 +1,70 @@ +"""Top-Level Synapse Compilation Flow.""" + +import subprocess +from collections.abc import Callable +from pathlib import Path +from tempfile import TemporaryDirectory + +from synapse.frontend.lowering import lower + + +def compile(program: Callable, *, target: str) -> str: + """Compile a Synapse program for the selected backend.""" + + # We only support the Neura backend for now, so we raise an error if the user tries to compile for any other target. + if target != "neura": + raise ValueError(f"unsupported compilation target: {target}") + # TODO: Support the amoeba backend. + + neura_ir = lower(program) + return _run_neura_backend(neura_ir) + + +def _run_neura_backend(neura_ir: str) -> str: + """Legalize Neura values, insert data movement, and run 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}") + + with TemporaryDirectory(prefix="synapse-") as temporary_directory: + output_path = Path(temporary_directory) / "mapped.mlir" + + command = [ + str(amoeba_opt), + "--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, + ) + + 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") diff --git a/python/synapse/frontend/lowering.py b/python/synapse/frontend/lowering.py new file mode 100644 index 0000000..74b3cfb --- /dev/null +++ b/python/synapse/frontend/lowering.py @@ -0,0 +1,187 @@ +"""Lower Synapse Python programs to compiler input IR.""" + +from collections.abc import Callable + +from synapse.language.spatial import Tile +from synapse.language.tile_array_program import ( + AddOp, + ConstantOp, + TileArrayBuilder, + TileArrayProgram, + TileArrayScalarType, +) + + +def lower(program_fn: Callable) -> str: + """Lower one tile-array program to pre-mapping Taskflow and Neura IR.""" + + builder = TileArrayBuilder() + + # Execute the user's tile-array DSL while recording its operations. + with builder: + program_fn() + + program = builder.build() + + return _lower_tile_array_program( + program_name=program_fn.__name__, + program=program, + ) + + +def _lower_tile_array_program( + *, + program_name: str, + program: TileArrayProgram, +) -> str: + """Convert a TileArrayProgram into pre-mapping Taskflow and Neura IR. + + MLIR imports remain local so users can import ``synapse.language`` without + requiring the compiled Amoeba Python bindings. + """ + + from taskflow_mlir.dialects import func, neura, taskflow + from taskflow_mlir.ir import ( + Context, + DictAttr, + F32Type, + FloatAttr, + InsertionPoint, + IntegerAttr, + IntegerType, + Location, + Module, + StringAttr, + ) + + with Context(), Location.unknown(): + taskflow.register_dialect() + neura.register_dialect() + + i32 = IntegerType.get_signless(32) + + def get_mlir_type(dtype: TileArrayScalarType): + """Translate a frontend scalar type into an MLIR type.""" + + if dtype == TileArrayScalarType.I32: + return i32 + + if dtype == TileArrayScalarType.F32: + return F32Type.get() + + raise NotImplementedError( + f"unsupported tile-array scalar type: {dtype.value}" + ) + + def get_constant_attribute( + operation: ConstantOp, + result_type, + ): + """Build the typed MLIR attribute for a constant value.""" + + if operation.result.dtype == TileArrayScalarType.I32: + if not isinstance(operation.value, int): + raise TypeError("an i32 constant requires an integer value") + return IntegerAttr.get(result_type, operation.value) + + if operation.result.dtype == TileArrayScalarType.F32: + if not isinstance(operation.value, (float, int)): + raise TypeError("an f32 constant requires a numeric value") + return FloatAttr.get(result_type, float(operation.value)) + + raise NotImplementedError( + f"unsupported constant type: {operation.result.dtype.value}" + ) + + def get_placement(tile: Tile) -> DictAttr: + """Build Neura placement directly from a Tile coordinate.""" + + return DictAttr.get( + { + "x": IntegerAttr.get(i32, tile.x), + "y": IntegerAttr.get(i32, tile.y), + } + ) + + module = Module.create() + + # This milestone lowers one Python function into one task containing + # one manually placed Neura kernel. + with InsertionPoint(module.body): + function = func.FuncOp(program_name, ([], [])) + function_block = function.add_entry_block() + + with InsertionPoint(function_block): + task = taskflow.TaskflowTaskOp( + done_reads=[], + done_writes=[], + value_outputs=[], + will_reads=[], + will_writes=[], + value_inputs=[], + task_name=program_name, + original_read_memrefs=[], + original_write_memrefs=[], + ) + task_block = task.body.blocks.append() + + func.ReturnOp([]) + + with InsertionPoint(task_block): + kernel = neura.KernelOp( + outputs=[], + inputs=[], + iter_args_init=[], + accelerator=StringAttr.get("neura"), + ) + kernel_block = kernel.body.blocks.append() + + taskflow.TaskflowYieldOp( + done_reads=[], + done_writes=[], + value_results=[], + ) + + # Map frontend value IDs to the MLIR SSA values produced while + # lowering the recorded operations. + values_by_id = {} + + with InsertionPoint(kernel_block): + for operation in program.operations: + result_type = get_mlir_type(operation.result.dtype) + + if isinstance(operation, ConstantOp): + mlir_operation = neura.ConstantOp( + result_type, + get_constant_attribute(operation, result_type), + ) + + elif isinstance(operation, AddOp): + lhs = values_by_id[operation.lhs.id] + rhs = values_by_id[operation.rhs.id] + + mlir_operation = neura.AddOp( + result_type, + lhs, + rhs=rhs, + ) + + else: + raise NotImplementedError( + f"unsupported tile-array operation: {type(operation).__name__}" + ) + + mlir_operation.operation.attributes["placement"] = get_placement( + operation.tile + ) + values_by_id[operation.result.id] = mlir_operation.result + + neura.YieldOp( + iter_args_next=[], + results_=[], + ) + + if not module.operation.verify(): + raise RuntimeError("generated Taskflow/Neura module is invalid") + + return str(module) diff --git a/python/synapse/language/spatial.py b/python/synapse/language/spatial.py index 3d67180..a18c886 100644 --- a/python/synapse/language/spatial.py +++ b/python/synapse/language/spatial.py @@ -12,35 +12,38 @@ class Tile: """A hardware tile in a CGRA TileArray. - A Tile is owned by a TileArray. Its row and column identify the - corresponding position in the target CGRA tile array. + A Tile is owned by a TileArray. Its ``x`` and ``y`` coordinates identify + the corresponding position in the target CGRA tile array. """ - def __init__(self, row: int, col: int, array: "TileArray"): + def __init__(self, x: int, y: int, array: "TileArray"): self.array = array - self.row = row - self.col = col + self.x = x + self.y = y class TileArray: """A parameterized two-dimensional tile array. - In the initial implementation, ``rows`` and ``cols`` must match the - dimensions of the target CGRA. A tile accessed as ``array[row, col]`` + In the initial implementation, ``x_tiles`` and ``y_tiles`` must match the + dimensions of the target CGRA. A tile accessed as ``array[x, y]`` corresponds directly to the tile at that hardware coordinate. + Coordinates follow Neura's convention: increasing ``x`` moves east + (right), and increasing ``y`` moves north (up). + Example: - array = TileArray(rows=4, cols=4) + array = TileArray(x_tiles=4, y_tiles=4) tile = array[1, 2] """ - def __init__(self, rows: int, cols: int): - self.rows = rows - self.cols = cols + def __init__(self, x_tiles: int, y_tiles: int): + self.x_tiles = x_tiles + self.y_tiles = y_tiles self._tiles = [ - [Tile(row=row, col=col, array=self) for col in range(cols)] - for row in range(rows) + [Tile(x=x, y=y, array=self) for x in range(x_tiles)] + for y in range(y_tiles) ] def tiles(self): @@ -49,12 +52,11 @@ def tiles(self): The iteration order is a Python programming convenience and does not specify sequential hardware execution. """ - for row_tiles in self._tiles: - yield from row_tiles + for y_row in self._tiles: + yield from y_row def __getitem__(self, coordinate: tuple[int, int]) -> Tile: - """ - Return the tile at the given ``(row, col)`` coordinate - """ - row, col = coordinate - return self._tiles[row][col] + """Return the tile at the given ``(x, y)`` coordinate.""" + + x, y = coordinate + return self._tiles[y][x] diff --git a/tests/python/compiler/test_add_constant_kernel.py b/tests/python/compiler/test_add_constant_kernel.py new file mode 100644 index 0000000..db551a9 --- /dev/null +++ b/tests/python/compiler/test_add_constant_kernel.py @@ -0,0 +1,61 @@ +import synapse +import synapse.language as synl +from synapse.frontend.lowering import lower + +PRE_MAPPING_IR = """ +module { + func.func @add_constant() { + taskflow.task @add_constant : () -> () { + neura.kernel attributes {accelerator = "neura"} { + %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 + neura.yield + } + taskflow.yield + } + return + } +} +""".strip() + + +MAPPED_IR = """ +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}} { + %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 + %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 + %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) -> !neura.data + %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) -> !neura.data + %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, !neura.data) -> !neura.data + neura.yield {dfg_id = 2 : i32} + } + taskflow.yield + } + return + } +} +""".strip() + + +def add_constant(): + array = synl.TileArray(x_tiles=4, y_tiles=4) + + lhs = synl.constant(1, tile=array[0, 0]) + rhs = synl.constant(2, tile=array[2, 0]) + + synl.add(lhs, rhs, tile=array[1, 0]) + + +def test_add_constant_lowers_to_placed_neura_ir(): + actual = lower(add_constant) + + assert actual.strip() == PRE_MAPPING_IR + + +def test_add_constant_compiles_to_mapped_neura_ir(): + actual = synapse.compile(add_constant, target="neura") + + assert actual.strip() == MAPPED_IR diff --git a/tests/python/language/test_spatial.py b/tests/python/language/test_spatial.py index 0cace14..1cab673 100644 --- a/tests/python/language/test_spatial.py +++ b/tests/python/language/test_spatial.py @@ -2,38 +2,38 @@ def test_tile_array_is_parameterized(): - array_2x3 = synl.TileArray(rows=2, cols=3) - array_4x4 = synl.TileArray(rows=4, cols=4) + array_2x3 = synl.TileArray(x_tiles=2, y_tiles=3) + array_4x4 = synl.TileArray(x_tiles=4, y_tiles=4) - assert array_2x3.rows == 2 - assert array_2x3.cols == 3 + assert array_2x3.x_tiles == 2 + assert array_2x3.y_tiles == 3 - assert array_4x4.rows == 4 - assert array_4x4.cols == 4 + assert array_4x4.x_tiles == 4 + assert array_4x4.y_tiles == 4 -def test_tile_array_exposes_logical_tiles(): - array = synl.TileArray(2, 3) +def test_tile_array_exposes_tiles(): + array = synl.TileArray(x_tiles=2, y_tiles=3) - coordinates = {(tile.row, tile.col) for tile in array.tiles()} + coordinates = {(tile.x, tile.y) for tile in array.tiles()} assert coordinates == {(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)} def test_access_tile_array_by_coordinate(): - array = synl.TileArray(2, 3) + array = synl.TileArray(x_tiles=2, y_tiles=3) tile = array[1, 2] assert tile.array is array - assert tile.row == 1 - assert tile.col == 2 + assert tile.x == 1 + assert tile.y == 2 assert tile is array[1, 2] enumerated_tile = next( candidate for candidate in array.tiles() - if candidate.row == 1 and candidate.col == 2 + if candidate.x == 1 and candidate.y == 2 ) assert tile is enumerated_tile diff --git a/tests/python/language/test_tile_array_program.py b/tests/python/language/test_tile_array_program.py index ccef8b7..847e7cb 100644 --- a/tests/python/language/test_tile_array_program.py +++ b/tests/python/language/test_tile_array_program.py @@ -2,7 +2,7 @@ def test_records_tile_array_program(): - array = synl.TileArray(4, 4) + array = synl.TileArray(x_tiles=4, y_tiles=4) builder = synl.tile_array_program.TileArrayBuilder() with builder: @@ -31,7 +31,7 @@ def test_records_tile_array_program(): def test_infers_supported_scalar_types(): - array = synl.TileArray(1, 3) + array = synl.TileArray(x_tiles=1, y_tiles=3) builder = synl.tile_array_program.TileArrayBuilder() with builder: From f2ac824f04e71153081e0e75165764178c51e102 Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Wed, 19 Aug 2026 21:03:56 +0800 Subject: [PATCH 08/11] Update README & workflow --- .github/workflows/test.yml | 168 +++++++++++++++++++++++++++++++++-- README.md | 175 +++++++++++++++++++++++++++++++++---- 2 files changed, 318 insertions(+), 25 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 3454d6c..98fedd3 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -3,9 +3,26 @@ name: tests on: push: pull_request: + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +env: + BUILD_TYPE: Release + LLVM_COMMIT: 6146a88f60492b520a36f8f8f3231e15f3cc6082 + LLVM_BUILD_DIR: ${{ github.workspace }}/llvm-project/build + AMOEBA_BUILD_DIR: ${{ github.workspace }}/build/amoeba + CCACHE_DIR: ${{ github.workspace }}/.ccache + CCACHE_COMPRESS: "true" + CCACHE_MAXSIZE: 4G jobs: - pytest: + python-unit-tests: runs-on: ubuntu-latest strategy: @@ -14,18 +31,157 @@ jobs: steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v7 - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: python-version: ${{ matrix.python-version }} - - name: Install package + - 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 tests - run: pytest -q \ No newline at end of file + - name: Run frontend and language tests + run: python -m pytest -q tests/python/frontend tests/python/language + + compiler-integration: + runs-on: ubuntu-22.04 + timeout-minutes: 240 + + steps: + - name: Checkout Synapse + uses: actions/checkout@v7 + + - name: Initialize Amoeba and Neura + run: | + git submodule update --init mlir/amoeba + git -C mlir/amoeba submodule update --init thirdparty/neura + + - name: Verify compiler submodules + run: | + test "$(git -C mlir/amoeba rev-parse HEAD)" = "$(git rev-parse HEAD:mlir/amoeba)" + test "$(git -C mlir/amoeba/thirdparty/neura rev-parse HEAD)" = "$(git -C mlir/amoeba rev-parse HEAD:thirdparty/neura)" + + - name: Install build dependencies + run: | + sudo apt-get update + sudo apt-get install --yes ccache clang lld ninja-build + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: 3.11.13 + + - name: Restore Python package cache + uses: actions/cache@v6 + with: + path: ~/.cache/pip + key: ${{ runner.os }}-pip-python-3.11.13-synapse + + - 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 + id: ccache + uses: actions/cache/restore@v6 + with: + path: ${{ env.CCACHE_DIR }} + key: ${{ runner.os }}-ccache-${{ env.LLVM_COMMIT }}-${{ github.sha }} + restore-keys: | + ${{ runner.os }}-ccache-${{ env.LLVM_COMMIT }}- + + - name: Restore LLVM build + id: llvm-cache + uses: actions/cache/restore@v6 + with: + path: llvm-project + key: ${{ runner.os }}-llvm-python-3.11.13-${{ env.LLVM_COMMIT }}-${{ env.BUILD_TYPE }}-v2 + + - name: Build LLVM and MLIR + if: steps.llvm-cache.outputs.cache-hit != 'true' + run: | + mkdir -p "${CCACHE_DIR}" + git init llvm-project + git -C llvm-project remote add origin https://github.com/llvm/llvm-project.git + git -C llvm-project fetch --depth=1 --filter=blob:none origin "${LLVM_COMMIT}" + git -C llvm-project checkout --detach FETCH_HEAD + + cmake -G Ninja \ + -S llvm-project/llvm \ + -B llvm-project/build \ + -DLLVM_ENABLE_PROJECTS="mlir;clang" \ + -DLLVM_BUILD_EXAMPLES=OFF \ + -DLLVM_TARGETS_TO_BUILD=Native \ + -DCMAKE_BUILD_TYPE="${BUILD_TYPE}" \ + -DLLVM_ENABLE_ASSERTIONS=ON \ + -DCMAKE_C_COMPILER=clang \ + -DCMAKE_CXX_COMPILER=clang++ \ + -DCMAKE_CXX_STANDARD=17 \ + -DCMAKE_CXX_FLAGS="-std=c++17 -frtti" \ + -DLLVM_ENABLE_LLD=ON \ + -DMLIR_INSTALL_AGGREGATE_OBJECTS=ON \ + -DLLVM_ENABLE_RTTI=ON \ + -DLLVM_CCACHE_BUILD=ON \ + -DMLIR_ENABLE_BINDINGS_PYTHON=ON \ + -DMLIR_BINDINGS_PYTHON_NB_DOMAIN=mlir \ + -DPython3_EXECUTABLE="$(which python)" \ + -DPython_EXECUTABLE="$(which python)" \ + -DCMAKE_C_COMPILER_LAUNCHER=ccache \ + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache + + cmake --build llvm-project/build --parallel 2 + + - name: Verify LLVM build + run: | + test -f llvm-project/build/lib/cmake/llvm/LLVMConfig.cmake + test -f llvm-project/build/lib/cmake/mlir/MLIRConfig.cmake + test -x llvm-project/build/bin/llvm-lit + + - name: Save LLVM build + if: steps.llvm-cache.outputs.cache-hit != 'true' + uses: actions/cache/save@v6 + with: + path: llvm-project + key: ${{ steps.llvm-cache.outputs.cache-primary-key }} + + - name: Configure Amoeba + run: | + cmake -G Ninja \ + -S mlir/amoeba \ + -B "${AMOEBA_BUILD_DIR}" \ + -DCMAKE_BUILD_TYPE="${BUILD_TYPE}" \ + -DPython3_EXECUTABLE="$(which python)" \ + -DPython_EXECUTABLE="$(which python)" \ + -DCMAKE_C_COMPILER_LAUNCHER=ccache \ + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache + + - name: Build Amoeba compiler and Python bindings + run: | + cmake --build "${AMOEBA_BUILD_DIR}" \ + --target mlir-amoeba-opt AmoebaPythonModules \ + --parallel 2 + + - name: Verify Amoeba build + run: | + 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: Save ccache + if: steps.ccache.outputs.cache-hit != 'true' + uses: actions/cache/save@v6 + with: + path: ${{ env.CCACHE_DIR }} + key: ${{ steps.ccache.outputs.cache-primary-key }} + + - name: Show ccache statistics + if: always() + run: ccache --show-stats diff --git a/README.md b/README.md index 5789d71..5289064 100644 --- a/README.md +++ b/README.md @@ -1,53 +1,190 @@ # SYNAPSE -SYNAPSE is a programming model and compiler frontend for -task-level dataflow systems. +SYNAPSE is a Python programming model and compiler frontend for spatial +dataflow systems. Its first backend target is a multi-CGRA architecture in +which each task may contain a program explicitly placed on a CGRA tile array. + +The current single-task compiler path lowers a Python TileArray program into: + +```text +func.func + taskflow.task + neura.kernel + placed Neura operations +``` + +The Neura backend then legalizes predicated values, inserts data movement, and +maps the placed operations to tiles, links, registers, and time steps. ## Repository Layout ```text -python/synapse/ Python package source -python/synapse/frontend Source capture and frontend parser utilities +python/synapse/ + language/ User-facing spatial and TileArray language APIs + frontend/ Python program capture and lowering to MLIR + compiler/ Backend compiler orchestration examples/ Small frontend examples -tests/ Pytest tests -mlir/amoeba/ Downstream NEURA / Taskflow compiler stack +tests/python/ Python unit and compiler-integration tests +mlir/amoeba/ Pinned Amoeba compiler dependency + thirdparty/neura/ Pinned Neura backend dependency managed by Amoeba ``` -## Setup +## Requirements + +Frontend and language development requires Python 3.10 or newer. + +The compiler-integration path additionally requires: + +- Python 3.11; +- CMake, Ninja, Clang, LLD, and ccache; +- `pybind11==2.13.6` and `nanobind==2.15.0`; and +- LLVM/MLIR at commit + [`6146a88f60492b520a36f8f8f3231e15f3cc6082`](https://github.com/llvm/llvm-project/commit/6146a88f60492b520a36f8f8f3231e15f3cc6082). + +This is the same LLVM revision and Python binding configuration used by the +pinned Amoeba workflow. + +## Checkout + +After cloning SYNAPSE, initialize Amoeba and its direct Neura dependency: + +```bash +git submodule update --init mlir/amoeba +git -C mlir/amoeba submodule update --init thirdparty/neura +``` + +These commands intentionally avoid downloading Neura's nested benchmark +submodules, which are not required to build the SYNAPSE compiler path. + +## Python Setup Create or activate a Python environment, then install SYNAPSE in editable mode from the repository root: ```bash -cd $PROJECT_PATH/synapse +python -m pip install --upgrade pip python -m pip install -e . +python -m pip install pytest ``` Editable install only needs to be done once per environment. After that, changes under `python/synapse/` are picked up directly. -For tests, install `pytest`: +## Build LLVM and MLIR + +Install the Python dependencies into the same Python 3.11 environment that +will configure LLVM and Amoeba: ```bash -python -m pip install pytest +python -m pip install pybind11==2.13.6 nanobind==2.15.0 ``` -## Run The GEMM Parser Example +Choose a location for LLVM, clone it, and check out the pinned revision: ```bash -cd $PROJECT_PATH/synapse -python examples/gemm.py +export SYNAPSE_LLVM_PROJECT=/absolute/path/to/llvm-project + +git clone https://github.com/llvm/llvm-project.git "${SYNAPSE_LLVM_PROJECT}" +git -C "${SYNAPSE_LLVM_PROJECT}" checkout \ + 6146a88f60492b520a36f8f8f3231e15f3cc6082 +``` + +Configure and build LLVM/MLIR with Python bindings enabled: + +```bash +cmake -G Ninja \ + -S "${SYNAPSE_LLVM_PROJECT}/llvm" \ + -B "${SYNAPSE_LLVM_PROJECT}/build" \ + -DLLVM_ENABLE_PROJECTS="mlir;clang" \ + -DLLVM_BUILD_EXAMPLES=OFF \ + -DLLVM_TARGETS_TO_BUILD=Native \ + -DCMAKE_BUILD_TYPE=Release \ + -DLLVM_ENABLE_ASSERTIONS=ON \ + -DCMAKE_C_COMPILER=clang \ + -DCMAKE_CXX_COMPILER=clang++ \ + -DCMAKE_CXX_STANDARD=17 \ + -DCMAKE_CXX_FLAGS="-std=c++17 -frtti" \ + -DLLVM_ENABLE_LLD=ON \ + -DMLIR_INSTALL_AGGREGATE_OBJECTS=ON \ + -DLLVM_ENABLE_RTTI=ON \ + -DLLVM_CCACHE_BUILD=ON \ + -DMLIR_ENABLE_BINDINGS_PYTHON=ON \ + -DMLIR_BINDINGS_PYTHON_NB_DOMAIN=mlir \ + -DPython3_EXECUTABLE="$(which python)" \ + -DPython_EXECUTABLE="$(which python)" \ + -DCMAKE_C_COMPILER_LAUNCHER=ccache \ + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache + +cmake --build "${SYNAPSE_LLVM_PROJECT}/build" --parallel 2 +``` + +## Build Amoeba and Neura + +From the SYNAPSE repository root, point Amoeba at the LLVM build and create the +compiler artifacts under `build/amoeba`: + +```bash +export LLVM_BUILD_DIR="${SYNAPSE_LLVM_PROJECT}/build" + +cmake -G Ninja \ + -S mlir/amoeba \ + -B build/amoeba \ + -DCMAKE_BUILD_TYPE=Release \ + -DPython3_EXECUTABLE="$(which python)" \ + -DPython_EXECUTABLE="$(which python)" + +cmake --build build/amoeba \ + --target mlir-amoeba-opt AmoebaPythonModules \ + --parallel 2 ``` -The example parses a plain Python GEMM function and prints its Python AST. At -this stage, the GEMM function is not executed; SYNAPSE only reads its source. +The build produces: + +```text +build/amoeba/tools/mlir-amoeba-opt/mlir-amoeba-opt +build/amoeba/python_packages/amoeba_core/taskflow_mlir/ +``` ## Run Tests +Run the frontend and language tests without building LLVM: + +```bash +python -m pytest -q tests/python/frontend tests/python/language +``` + +After building Amoeba, run the compiler-integration tests: + ```bash -cd $PROJECT_PATH/synapse -pytest -q +python -m pytest -q tests/python/compiler ``` -The current tests check that the frontend parser can capture a Python GEMM -function and expose the expected AST nodes. +Or run the complete Python suite: + +```bash +python -m pytest -q tests/python +``` + +## Run the GEMM Parser Example + +```bash +python examples/gemm.py +``` + +This example captures a plain Python GEMM function and prints its Python AST. +The TileArray compiler path is exercised by the tests under +`tests/python/compiler`. + +## Continuous Integration + +The GitHub Actions workflow contains two layers: + +- `python-unit-tests` runs frontend and language tests on Python 3.10 and 3.11 + without building LLVM. +- `compiler-integration` uses Python 3.11, checks out the pinned Amoeba and + Neura revisions, downloads the pinned LLVM commit, builds LLVM/MLIR with + Python bindings, builds `mlir-amoeba-opt` and `AmoebaPythonModules`, and runs + the compiler tests. + +LLVM and ccache artifacts are cached using the pinned LLVM revision so later +workflow runs do not rebuild the entire dependency stack unnecessarily. From f35fc4ac887ce058469fb406814c69408d54ab51 Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Wed, 19 Aug 2026 21:14:53 +0800 Subject: [PATCH 09/11] Update workflow --- .github/workflows/test.yml | 18 ++++++++++++------ README.md | 7 +++---- pyproject.toml | 2 +- 3 files changed, 16 insertions(+), 11 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 98fedd3..cf55a6c 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -2,6 +2,8 @@ name: tests on: push: + branches: + - main pull_request: workflow_dispatch: @@ -25,10 +27,6 @@ jobs: python-unit-tests: runs-on: ubuntu-latest - strategy: - matrix: - python-version: ["3.10", "3.11"] - steps: - name: Checkout uses: actions/checkout@v7 @@ -36,7 +34,7 @@ jobs: - name: Set up Python uses: actions/setup-python@v6 with: - python-version: ${{ matrix.python-version }} + python-version: 3.11.13 - name: Install Synapse and test dependencies run: | @@ -58,6 +56,9 @@ jobs: - name: Initialize Amoeba and Neura run: | git submodule update --init mlir/amoeba + git -C mlir/amoeba config \ + submodule.thirdparty/neura.url \ + https://github.com/coredac/neura.git git -C mlir/amoeba submodule update --init thirdparty/neura - name: Verify compiler submodules @@ -184,4 +185,9 @@ jobs: - name: Show ccache statistics if: always() - run: ccache --show-stats + run: | + if command -v ccache >/dev/null 2>&1; then + ccache --show-stats + else + echo "ccache was not installed because setup did not complete." + fi diff --git a/README.md b/README.md index 5289064..ac83b8b 100644 --- a/README.md +++ b/README.md @@ -31,11 +31,10 @@ mlir/amoeba/ Pinned Amoeba compiler dependency ## Requirements -Frontend and language development requires Python 3.10 or newer. +SYNAPSE currently requires Python 3.11. The compiler-integration path additionally requires: -- Python 3.11; - CMake, Ninja, Clang, LLD, and ccache; - `pybind11==2.13.6` and `nanobind==2.15.0`; and - LLVM/MLIR at commit @@ -179,8 +178,8 @@ The TileArray compiler path is exercised by the tests under The GitHub Actions workflow contains two layers: -- `python-unit-tests` runs frontend and language tests on Python 3.10 and 3.11 - without building LLVM. +- `python-unit-tests` runs frontend and language tests on Python 3.11 without + building LLVM. - `compiler-integration` uses Python 3.11, checks out the pinned Amoeba and Neura revisions, downloads the pinned LLVM commit, builds LLVM/MLIR with Python bindings, builds `mlir-amoeba-opt` and `AmoebaPythonModules`, and runs diff --git a/pyproject.toml b/pyproject.toml index f25db76..8c7c081 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta" [project] name = "synapse" version = "0.0.0" -requires-python = ">=3.10" +requires-python = ">=3.11" [tool.setuptools] package-dir = {"" = "python"} From c3878dff9ee48d5eb9b03921026c4244d12ac227 Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Tue, 25 Aug 2026 01:31:34 +0800 Subject: [PATCH 10/11] Update the code structure --- python/synapse/frontend/lowering.py | 61 +++++++++++++------ python/synapse/language/spatial.py | 10 ++- python/synapse/language/tile_array_program.py | 52 ++++++++-------- .../language/test_tile_array_program.py | 4 +- 4 files changed, 79 insertions(+), 48 deletions(-) diff --git a/python/synapse/frontend/lowering.py b/python/synapse/frontend/lowering.py index 74b3cfb..1e1da7f 100644 --- a/python/synapse/frontend/lowering.py +++ b/python/synapse/frontend/lowering.py @@ -1,12 +1,14 @@ """Lower Synapse Python programs to compiler input IR.""" from collections.abc import Callable +from functools import singledispatch from synapse.language.spatial import Tile from synapse.language.tile_array_program import ( AddOp, ConstantOp, TileArrayBuilder, + TileArrayOp, TileArrayProgram, TileArrayScalarType, ) @@ -103,6 +105,41 @@ def get_placement(tile: Tile) -> DictAttr: } ) + @singledispatch + def lower_operation(operation: TileArrayOp, operands, result_type): + """Lower one frontend TileArray operation to a Neura operation. + + The caller handles common lowering such as resolving operands, + attaching placement, and recording the resulting SSA value. + """ + raise NotImplementedError( + f"unsupported tile-array operation: {type(operation).__name__}" + ) + + @lower_operation.register + def lower_constant(operation: ConstantOp, operands, result_type): + """Lower a ConstantOp to neura.constant.""" + if operands: + raise ValueError( + f"ConstantOp requires zero operands, but got {len(operands)}" + ) + return neura.ConstantOp( + result_type, get_constant_attribute(operation, result_type) + ) + + @lower_operation.register + def lower_add(operation: AddOp, operands, result_type): + """Lower a frontend AddOp to neura.add.""" + + if len(operands) != 2: + raise ValueError( + f"AddOp requires two operands, but got {len(operands)}" + ) + + lhs, rhs = operands + + return neura.AddOp(result_type, lhs, rhs=rhs) + module = Module.create() # This milestone lowers one Python function into one task containing @@ -150,30 +187,16 @@ def get_placement(tile: Tile) -> DictAttr: for operation in program.operations: result_type = get_mlir_type(operation.result.dtype) - if isinstance(operation, ConstantOp): - mlir_operation = neura.ConstantOp( - result_type, - get_constant_attribute(operation, result_type), - ) - - elif isinstance(operation, AddOp): - lhs = values_by_id[operation.lhs.id] - rhs = values_by_id[operation.rhs.id] - - mlir_operation = neura.AddOp( - result_type, - lhs, - rhs=rhs, - ) + mlir_operands = tuple( + values_by_id[operand.id] for operand in operation.operands + ) - else: - raise NotImplementedError( - f"unsupported tile-array operation: {type(operation).__name__}" - ) + mlir_operation = lower_operation(operation, mlir_operands, result_type) mlir_operation.operation.attributes["placement"] = get_placement( operation.tile ) + values_by_id[operation.result.id] = mlir_operation.result neura.YieldOp( diff --git a/python/synapse/language/spatial.py b/python/synapse/language/spatial.py index a18c886..bb67ad5 100644 --- a/python/synapse/language/spatial.py +++ b/python/synapse/language/spatial.py @@ -42,8 +42,7 @@ def __init__(self, x_tiles: int, y_tiles: int): self.y_tiles = y_tiles self._tiles = [ - [Tile(x=x, y=y, array=self) for x in range(x_tiles)] - for y in range(y_tiles) + [Tile(x=x, y=y, array=self) for x in range(x_tiles)] for y in range(y_tiles) ] def tiles(self): @@ -59,4 +58,11 @@ def __getitem__(self, coordinate: tuple[int, int]) -> Tile: """Return the tile at the given ``(x, y)`` coordinate.""" x, y = coordinate + + if not (0 <= x < self.x_tiles and 0 <= y < self.y_tiles): + raise IndexError( + f"tile coordinate ({x}, {y}) is outside " + f"TileArray({self.x_tiles}, {self.y_tiles})" + ) + return self._tiles[y][x] diff --git a/python/synapse/language/tile_array_program.py b/python/synapse/language/tile_array_program.py index 12f0ff8..d7af970 100644 --- a/python/synapse/language/tile_array_program.py +++ b/python/synapse/language/tile_array_program.py @@ -10,7 +10,6 @@ from contextvars import ContextVar from dataclasses import dataclass, field from enum import Enum -from typing import TypeAlias from .spatial import Tile, TileArray @@ -39,31 +38,30 @@ class TileArrayValue: # --------------------------------------------------------------- # Typed tile-array operations # --------------------------------------------------------------- - - @dataclass(frozen=True) -class ConstantOp: - """A scalar constant produced by a tile-array operation.""" +class TileArrayOp: + """Base class for operations executed on a TileArray. + + ``operands`` may contain any number of input values. Constants + therefore use an empty tuple, while operations such as add and MAC + use two or more operands. + """ result: TileArrayValue - value: int | float + operands: tuple[TileArrayValue, ...] tile: Tile @dataclass(frozen=True) -class AddOp: - """A scalar addition executed by a tile-array operation.""" +class ConstantOp(TileArrayOp): + """A scalar constant produced by a tile-array operation.""" - result: TileArrayValue - lhs: TileArrayValue - rhs: TileArrayValue - tile: Tile + value: int | float -# This union explicitly lists every operation currently supported by the -# tile-array frontend. Future operations such as MacOp and GatherOp should be -# added here. -TileArrayOp: TypeAlias = ConstantOp | AddOp +@dataclass(frozen=True) +class AddOp(TileArrayOp): + """A scalar addition executed by a tile-array operation.""" @dataclass(frozen=True) @@ -126,8 +124,8 @@ def _bind_tile_array(self, tile: Tile) -> None: "all operations in a TileArrayProgram must use tiles from the same TileArray" ) - def _check_operand(self, operand: TileArrayValue) -> None: - """Verify that an operand was produced by this builder.""" + def _validate_operand_for_builder(self, operand: TileArrayValue) -> None: + """Validate that an operand was produced by this builder.""" if not isinstance(operand, TileArrayValue): raise TypeError("operation operand must be a TileArrayValue") @@ -136,8 +134,8 @@ def _check_operand(self, operand: TileArrayValue) -> None: "operation operand belongs to a different TileArrayProgram" ) - def _check_can_emit(self) -> None: - """Reject operations emitted after the program has been finalized.""" + def _ensure_not_built(self) -> None: + """Reject operations emitted after the program has been built.""" if self._is_built: raise RuntimeError( "cannot emit operations after building a TileArrayProgram" @@ -147,28 +145,30 @@ def emit_constant( self, value: int | float, *, dtype: TileArrayScalarType, tile: Tile ) -> TileArrayValue: """Record one scalar constant operation.""" - self._check_can_emit() + self._ensure_not_built() self._bind_tile_array(tile) result = self._new_value(dtype=dtype) - self._operations.append(ConstantOp(result=result, value=value, tile=tile)) + self._operations.append( + ConstantOp(result=result, operands=(), value=value, tile=tile) + ) return result def emit_add( self, lhs: TileArrayValue, rhs: TileArrayValue, *, tile: Tile ) -> TileArrayValue: """Record one scalar addition operation.""" - self._check_can_emit() + self._ensure_not_built() self._bind_tile_array(tile) - self._check_operand(lhs) - self._check_operand(rhs) + self._validate_operand_for_builder(lhs) + self._validate_operand_for_builder(rhs) if lhs.dtype != rhs.dtype: raise TypeError("add operands must have the same tile-array value type") result = self._new_value(dtype=lhs.dtype) - self._operations.append(AddOp(result=result, lhs=lhs, rhs=rhs, tile=tile)) + self._operations.append(AddOp(result=result, operands=(lhs, rhs), tile=tile)) return result def build(self) -> TileArrayProgram: diff --git a/tests/python/language/test_tile_array_program.py b/tests/python/language/test_tile_array_program.py index 847e7cb..29c8ae6 100644 --- a/tests/python/language/test_tile_array_program.py +++ b/tests/python/language/test_tile_array_program.py @@ -27,7 +27,9 @@ def test_records_tile_array_program(): array[0, 2], array[0, 1], ] - assert (add_op.lhs, add_op.rhs) == (lhs, rhs) + assert lhs_op.operands == () + assert rhs_op.operands == () + assert add_op.operands == (lhs, rhs) def test_infers_supported_scalar_types(): From 7002bcdfeb1a6c3bb54584561008f827989165ac Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Wed, 26 Aug 2026 00:06:47 +0800 Subject: [PATCH 11/11] Clean up the emit_xxOp/xxOp/xx functionality --- python/synapse/frontend/lowering.py | 17 +-- python/synapse/language/tile_array_program.py | 139 ++++++++++++------ 2 files changed, 93 insertions(+), 63 deletions(-) diff --git a/python/synapse/frontend/lowering.py b/python/synapse/frontend/lowering.py index 1e1da7f..93f7d25 100644 --- a/python/synapse/frontend/lowering.py +++ b/python/synapse/frontend/lowering.py @@ -2,6 +2,7 @@ from collections.abc import Callable from functools import singledispatch +from typing import cast from synapse.language.spatial import Tile from synapse.language.tile_array_program import ( @@ -82,13 +83,9 @@ def get_constant_attribute( """Build the typed MLIR attribute for a constant value.""" if operation.result.dtype == TileArrayScalarType.I32: - if not isinstance(operation.value, int): - raise TypeError("an i32 constant requires an integer value") - return IntegerAttr.get(result_type, operation.value) + return IntegerAttr.get(result_type, cast(int, operation.value)) if operation.result.dtype == TileArrayScalarType.F32: - if not isinstance(operation.value, (float, int)): - raise TypeError("an f32 constant requires a numeric value") return FloatAttr.get(result_type, float(operation.value)) raise NotImplementedError( @@ -119,10 +116,6 @@ def lower_operation(operation: TileArrayOp, operands, result_type): @lower_operation.register def lower_constant(operation: ConstantOp, operands, result_type): """Lower a ConstantOp to neura.constant.""" - if operands: - raise ValueError( - f"ConstantOp requires zero operands, but got {len(operands)}" - ) return neura.ConstantOp( result_type, get_constant_attribute(operation, result_type) ) @@ -130,12 +123,6 @@ def lower_constant(operation: ConstantOp, operands, result_type): @lower_operation.register def lower_add(operation: AddOp, operands, result_type): """Lower a frontend AddOp to neura.add.""" - - if len(operands) != 2: - raise ValueError( - f"AddOp requires two operands, but got {len(operands)}" - ) - lhs, rhs = operands return neura.AddOp(result_type, lhs, rhs=rhs) diff --git a/python/synapse/language/tile_array_program.py b/python/synapse/language/tile_array_program.py index d7af970..4e26598 100644 --- a/python/synapse/language/tile_array_program.py +++ b/python/synapse/language/tile_array_program.py @@ -7,6 +7,7 @@ from __future__ import annotations +from collections.abc import Callable from contextvars import ContextVar from dataclasses import dataclass, field from enum import Enum @@ -58,11 +59,48 @@ class ConstantOp(TileArrayOp): value: int | float + def __post_init__(self) -> None: + """Validate the operation-specific operands and scalar value.""" + if self.operands: + raise ValueError( + f"ConstantOp requires zero operands, but got {len(self.operands)}" + ) + + dtype = self.result.dtype + if not isinstance(dtype, TileArrayScalarType): + raise TypeError("ConstantOp result must use a TileArrayScalarType") + + if isinstance(self.value, bool): + raise TypeError("boolean constants are not supported yet") + + if dtype == TileArrayScalarType.I32 and not isinstance(self.value, int): + raise TypeError("an i32 constant requires an integer value") + + if dtype == TileArrayScalarType.F32 and not isinstance( + self.value, (int, float) + ): + raise TypeError("an f32 constant requires a numeric value") + @dataclass(frozen=True) class AddOp(TileArrayOp): """A scalar addition executed by a tile-array operation.""" + def __post_init__(self) -> None: + """Validate the operation-specific arity and scalar types.""" + if len(self.operands) != 2: + raise ValueError( + f"AddOp requires two operands, but got {len(self.operands)}" + ) + + lhs, rhs = self.operands + + if lhs.dtype != rhs.dtype: + raise TypeError("AddOp operands must have the same scalar type") + + if self.result.dtype != lhs.dtype: + raise TypeError("AddOp result type must match its operand type") + @dataclass(frozen=True) class TileArrayProgram: @@ -106,12 +144,6 @@ def __exit__(self, exc_type, exc_value, traceback) -> None: _active_builder.reset(self._token) self._token = None - def _new_value(self, *, dtype: TileArrayScalarType) -> TileArrayValue: - """Allocate the next value identifier.""" - result = TileArrayValue(id=self._next_value_id, dtype=dtype, _builder=self) - self._next_value_id += 1 - return result - def _bind_tile_array(self, tile: Tile) -> None: """Bind the program to one TileArray.""" if not isinstance(tile, Tile): @@ -141,34 +173,41 @@ def _ensure_not_built(self) -> None: "cannot emit operations after building a TileArrayProgram" ) - def emit_constant( - self, value: int | float, *, dtype: TileArrayScalarType, tile: Tile - ) -> TileArrayValue: - """Record one scalar constant operation.""" - self._ensure_not_built() - self._bind_tile_array(tile) - result = self._new_value(dtype=dtype) - - self._operations.append( - ConstantOp(result=result, operands=(), value=value, tile=tile) - ) - return result - - def emit_add( - self, lhs: TileArrayValue, rhs: TileArrayValue, *, tile: Tile + def emit( + self, + *, + operands: tuple[TileArrayValue, ...], + result_dtype: TileArrayScalarType, + tile: Tile, + create_operation: Callable[[TileArrayValue], TileArrayOp], ) -> TileArrayValue: - """Record one scalar addition operation.""" + """Create and record one tile-array operation.""" self._ensure_not_built() self._bind_tile_array(tile) - self._validate_operand_for_builder(lhs) - self._validate_operand_for_builder(rhs) - if lhs.dtype != rhs.dtype: - raise TypeError("add operands must have the same tile-array value type") - - result = self._new_value(dtype=lhs.dtype) + for operand in operands: + self._validate_operand_for_builder(operand) - self._operations.append(AddOp(result=result, operands=(lhs, rhs), tile=tile)) + # Commit the value ID only after operation construction and validation + # succeed, so a rejected operation does not consume a value ID. + result = TileArrayValue( + id=self._next_value_id, + dtype=result_dtype, + _builder=self, + ) + operation = create_operation(result) + + if not isinstance(operation, TileArrayOp): + raise TypeError("create_operation must return a TileArrayOp") + if operation.result is not result: + raise ValueError("create_operation must use the provided result value") + if operation.operands != operands: + raise ValueError("create_operation must use the provided operands") + if operation.tile is not tile: + raise ValueError("create_operation must use the provided tile") + + self._operations.append(operation) + self._next_value_id += 1 return result def build(self) -> TileArrayProgram: @@ -217,32 +256,28 @@ def constant( Use an explicit dtype when a different representation is required: constant(1.0, tile=tile, dtype=TileArrayScalarType.F32) """ - if isinstance(value, bool): - raise TypeError("boolean constants are not supported yet") - if dtype is None: - if isinstance(value, int): + if type(value) is int: dtype = TileArrayScalarType.I32 - elif isinstance(value, float): + elif type(value) is float: dtype = TileArrayScalarType.F32 else: raise TypeError( "constant currently supports integer and floating-point values" ) - if not isinstance(dtype, TileArrayScalarType): - raise TypeError("dtype must be a TileArrayScalarType") - - if dtype == TileArrayScalarType.I32 and not isinstance(value, int): - raise TypeError("an i32 constant requires an integer value") + builder = _require_active_builder() - if dtype in (TileArrayScalarType.F32,) and not isinstance(value, (int, float)): - raise TypeError("a floating-point constant requires a numeric value") - - return _require_active_builder().emit_constant( - value, - dtype=dtype, + return builder.emit( + operands=(), + result_dtype=dtype, tile=tile, + create_operation=lambda result: ConstantOp( + result=result, + operands=(), + tile=tile, + value=value, + ), ) @@ -254,8 +289,16 @@ def add( ) -> TileArrayValue: """Create a scalar addition on one hardware tile.""" - return _require_active_builder().emit_add( - lhs, - rhs, + builder = _require_active_builder() + operands = (lhs, rhs) + + return builder.emit( + operands=operands, + result_dtype=lhs.dtype, tile=tile, + create_operation=lambda result: AddOp( + result=result, + operands=operands, + tile=tile, + ), )