From 9e8f5dea9174dce37535cd9c0abaf044448d7b8f Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Thu, 27 Aug 2026 20:55:31 +0800 Subject: [PATCH 01/19] Add mac operation in language --- python/synapse/language/__init__.py | 12 +++++- python/synapse/language/tile_array_program.py | 42 +++++++++++++++++++ .../language/test_tile_array_program.py | 29 +++++++++++++ 3 files changed, 81 insertions(+), 2 deletions(-) diff --git a/python/synapse/language/__init__.py b/python/synapse/language/__init__.py index eca34d4..b8df3f9 100644 --- a/python/synapse/language/__init__.py +++ b/python/synapse/language/__init__.py @@ -1,9 +1,17 @@ """Public Synapse language API.""" from .spatial import TileArray -from .tile_array_program import TileArrayScalarType, add, constant +from .tile_array_program import TileArrayScalarType, add, constant, mac i32 = TileArrayScalarType.I32 f32 = TileArrayScalarType.F32 -__all__ = ["TileArray", "TileArrayScalarType", "add", "constant", "f32", "i32"] +__all__ = [ + "TileArray", + "TileArrayScalarType", + "add", + "constant", + "f32", + "i32", + "mac", +] diff --git a/python/synapse/language/tile_array_program.py b/python/synapse/language/tile_array_program.py index 4e26598..4cb7e8e 100644 --- a/python/synapse/language/tile_array_program.py +++ b/python/synapse/language/tile_array_program.py @@ -102,6 +102,26 @@ def __post_init__(self) -> None: raise TypeError("AddOp result type must match its operand type") +@dataclass(frozen=True) +class MacOp(TileArrayOp): + """A scalar multiply-accumulate operation executed on a TileArray.""" + + def __post_init__(self) -> None: + """Validate the operation-specific arity and scalar types.""" + + if len(self.operands) != 3: + raise ValueError( + f"MacOp requires three operands, but got {len(self.operands)}" + ) + lhs, rhs, accumulator = self.operands + + if not (lhs.dtype == rhs.dtype == accumulator.dtype): + raise TypeError("MacOp operands must have the same scalar type") + + if self.result.dtype != lhs.dtype: + raise TypeError("MacOp result type must match its operand type") + + @dataclass(frozen=True) class TileArrayProgram: """A tile-array program produced by TileArrayBuilder.""" @@ -302,3 +322,25 @@ def add( tile=tile, ), ) + + +def mac( + lhs: TileArrayValue, + rhs: TileArrayValue, + accumulator: TileArrayValue, + *, + tile: Tile, +) -> TileArrayValue: + """Create a scalar multiply-accumulate operation on one hardware tile.""" + builder = _require_active_builder() + operands = (lhs, rhs, accumulator) + return builder.emit( + operands=operands, + result_dtype=lhs.dtype, + tile=tile, + create_operation=lambda result: MacOp( + result=result, + operands=operands, + tile=tile, + ), + ) diff --git a/tests/python/language/test_tile_array_program.py b/tests/python/language/test_tile_array_program.py index 29c8ae6..f1ce1f4 100644 --- a/tests/python/language/test_tile_array_program.py +++ b/tests/python/language/test_tile_array_program.py @@ -48,3 +48,32 @@ def test_infers_supported_scalar_types(): assert integer.dtype == synl.i32 assert floating.dtype == synl.f32 assert explicit_f32.dtype == synl.f32 + + +def test_records_mac_operation(): + array = synl.TileArray(x_tiles=1, y_tiles=1) + builder = synl.tile_array_program.TileArrayBuilder() + + with builder: + lhs = synl.constant(2.0, tile=array[0, 0]) + rhs = synl.constant(3.0, tile=array[0, 0]) + accumulator = synl.constant(0.0, tile=array[0, 0]) + + result = synl.mac( + lhs, + rhs, + accumulator, + tile=array[0, 0], + ) + + program = builder.build() + mac_op = program.operations[-1] + + assert isinstance( + mac_op, + synl.tile_array_program.MacOp, + ) + assert mac_op.operands == (lhs, rhs, accumulator) + assert mac_op.result is result + assert mac_op.tile is array[0, 0] + assert result.dtype == synl.f32 From bff3fdc529e007c077258146e5f4d2f6abfdd6db Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Thu, 27 Aug 2026 21:29:21 +0800 Subject: [PATCH 02/19] Lower MAC operations to Neura --- python/synapse/frontend/lowering.py | 26 +++++++++ tests/python/compiler/test_mac_kernel.py | 69 ++++++++++++++++++++++++ 2 files changed, 95 insertions(+) create mode 100644 tests/python/compiler/test_mac_kernel.py diff --git a/python/synapse/frontend/lowering.py b/python/synapse/frontend/lowering.py index 93f7d25..6018634 100644 --- a/python/synapse/frontend/lowering.py +++ b/python/synapse/frontend/lowering.py @@ -8,6 +8,7 @@ from synapse.language.tile_array_program import ( AddOp, ConstantOp, + MacOp, TileArrayBuilder, TileArrayOp, TileArrayProgram, @@ -127,6 +128,31 @@ def lower_add(operation: AddOp, operands, result_type): return neura.AddOp(result_type, lhs, rhs=rhs) + @lower_operation.register + def lower_mac(operation: MacOp, operands, result_type): + """Lower a frontend MacOp to the matching Neura fused operation.""" + lhs, rhs, accumulator = operands + + if operation.result.dtype == TileArrayScalarType.I32: + return neura.MulAddOp( + result_type, + lhs, + rhs, + accumulator, + ) + + if operation.result.dtype == TileArrayScalarType.F32: + return neura.FMulFAddOp( + result_type, + lhs, + rhs, + accumulator, + ) + + raise NotImplementedError( + f"unsupported MacOp scalar type: {operation.result.dtype.value}" + ) + module = Module.create() # This milestone lowers one Python function into one task containing diff --git a/tests/python/compiler/test_mac_kernel.py b/tests/python/compiler/test_mac_kernel.py new file mode 100644 index 0000000..9052260 --- /dev/null +++ b/tests/python/compiler/test_mac_kernel.py @@ -0,0 +1,69 @@ +import synapse +import synapse.language as synl +from synapse.frontend.lowering import lower + +EXPECTED_IR = """ +module { + func.func @f32_mac() { + taskflow.task @f32_mac : () -> () { + neura.kernel attributes {accelerator = "neura"} { + %0 = "neura.constant"() <{value = 2.000000e+00 : f32}> {placement = {x = 0 : i32, y = 0 : i32}} : () -> f32 + %1 = "neura.constant"() <{value = 3.000000e+00 : f32}> {placement = {x = 0 : i32, y = 0 : i32}} : () -> f32 + %2 = "neura.constant"() <{value = 0.000000e+00 : f32}> {placement = {x = 0 : i32, y = 0 : i32}} : () -> f32 + %3 = "neura.fmul_fadd"(%0, %1, %2) {placement = {x = 0 : i32, y = 0 : i32}} : (f32, f32, f32) -> f32 + neura.yield + } + taskflow.yield + } + return + } +} +""".strip() + + +def f32_mac(): + array = synl.TileArray(x_tiles=1, y_tiles=1) + + lhs = synl.constant(2.0, tile=array[0, 0]) + rhs = synl.constant(3.0, tile=array[0, 0]) + accumulator = synl.constant(0.0, tile=array[0, 0]) + + synl.mac( + lhs, + rhs, + accumulator, + tile=array[0, 0], + ) + + +def mapped_f32_mac(): + array = synl.TileArray(x_tiles=2, y_tiles=2) + + lhs = synl.constant(2.0, tile=array[0, 0]) + rhs = synl.constant(3.0, tile=array[1, 0]) + accumulator = synl.constant(0.0, tile=array[0, 1]) + + synl.mac( + lhs, + rhs, + accumulator, + tile=array[1, 1], + ) + + +def test_f32_mac_lowers_to_neura(): + actual = lower(f32_mac) + + assert actual.strip() == EXPECTED_IR + + +def test_f32_mac_compiles_to_mapped_neura(): + actual = synapse.compile( + mapped_f32_mac, + target="neura", + ) + + assert '"neura.fmul_fadd"' in actual + assert 'mapping_strategy = "template"' in actual + assert 'resource = "link"' in actual + assert "placement =" not in actual From 2eb4f5f3b8d19c1d9bd1a96641ea0df1833f93e0 Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Sun, 30 Aug 2026 13:27:09 +0800 Subject: [PATCH 03/19] Update Amoeba for systolic template mapping --- mlir/amoeba | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/amoeba b/mlir/amoeba index 784cb87..35e7f98 160000 --- a/mlir/amoeba +++ b/mlir/amoeba @@ -1 +1 @@ -Subproject commit 784cb87a45f51574cb45088da25bd674bed03abb +Subproject commit 35e7f98c9dbfe513bde0bc1f0da3e747fe794adf From b777ff7d6bb63b33ab4336d68c3248e0bf662e46 Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Sun, 30 Aug 2026 15:56:13 +0800 Subject: [PATCH 04/19] Update Amoeba for Taskflow stream adapters --- mlir/amoeba | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/amoeba b/mlir/amoeba index 35e7f98..b514f8d 160000 --- a/mlir/amoeba +++ b/mlir/amoeba @@ -1 +1 @@ -Subproject commit 35e7f98c9dbfe513bde0bc1f0da3e747fe794adf +Subproject commit b514f8d0202d5f6ed6efabc1849d3f891642c805 From f24dd043be06029dd6527a39cabd396e2e0dd5ad Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Sun, 30 Aug 2026 16:21:59 +0800 Subject: [PATCH 05/19] Remove generic MAC from TileArray frontend --- python/synapse/frontend/lowering.py | 26 ------ python/synapse/language/__init__.py | 3 +- python/synapse/language/tile_array_program.py | 42 --------- tests/python/compiler/test_mac_kernel.py | 69 --------------- tests/python/compiler/test_systolic_gemm.py | 85 +++++++++++++++++++ .../language/test_tile_array_program.py | 29 ------- 6 files changed, 86 insertions(+), 168 deletions(-) delete mode 100644 tests/python/compiler/test_mac_kernel.py create mode 100644 tests/python/compiler/test_systolic_gemm.py diff --git a/python/synapse/frontend/lowering.py b/python/synapse/frontend/lowering.py index 6018634..93f7d25 100644 --- a/python/synapse/frontend/lowering.py +++ b/python/synapse/frontend/lowering.py @@ -8,7 +8,6 @@ from synapse.language.tile_array_program import ( AddOp, ConstantOp, - MacOp, TileArrayBuilder, TileArrayOp, TileArrayProgram, @@ -128,31 +127,6 @@ def lower_add(operation: AddOp, operands, result_type): return neura.AddOp(result_type, lhs, rhs=rhs) - @lower_operation.register - def lower_mac(operation: MacOp, operands, result_type): - """Lower a frontend MacOp to the matching Neura fused operation.""" - lhs, rhs, accumulator = operands - - if operation.result.dtype == TileArrayScalarType.I32: - return neura.MulAddOp( - result_type, - lhs, - rhs, - accumulator, - ) - - if operation.result.dtype == TileArrayScalarType.F32: - return neura.FMulFAddOp( - result_type, - lhs, - rhs, - accumulator, - ) - - raise NotImplementedError( - f"unsupported MacOp scalar type: {operation.result.dtype.value}" - ) - module = Module.create() # This milestone lowers one Python function into one task containing diff --git a/python/synapse/language/__init__.py b/python/synapse/language/__init__.py index b8df3f9..463d6f9 100644 --- a/python/synapse/language/__init__.py +++ b/python/synapse/language/__init__.py @@ -1,7 +1,7 @@ """Public Synapse language API.""" from .spatial import TileArray -from .tile_array_program import TileArrayScalarType, add, constant, mac +from .tile_array_program import TileArrayScalarType, add, constant i32 = TileArrayScalarType.I32 f32 = TileArrayScalarType.F32 @@ -13,5 +13,4 @@ "constant", "f32", "i32", - "mac", ] diff --git a/python/synapse/language/tile_array_program.py b/python/synapse/language/tile_array_program.py index 4cb7e8e..4e26598 100644 --- a/python/synapse/language/tile_array_program.py +++ b/python/synapse/language/tile_array_program.py @@ -102,26 +102,6 @@ def __post_init__(self) -> None: raise TypeError("AddOp result type must match its operand type") -@dataclass(frozen=True) -class MacOp(TileArrayOp): - """A scalar multiply-accumulate operation executed on a TileArray.""" - - def __post_init__(self) -> None: - """Validate the operation-specific arity and scalar types.""" - - if len(self.operands) != 3: - raise ValueError( - f"MacOp requires three operands, but got {len(self.operands)}" - ) - lhs, rhs, accumulator = self.operands - - if not (lhs.dtype == rhs.dtype == accumulator.dtype): - raise TypeError("MacOp operands must have the same scalar type") - - if self.result.dtype != lhs.dtype: - raise TypeError("MacOp result type must match its operand type") - - @dataclass(frozen=True) class TileArrayProgram: """A tile-array program produced by TileArrayBuilder.""" @@ -322,25 +302,3 @@ def add( tile=tile, ), ) - - -def mac( - lhs: TileArrayValue, - rhs: TileArrayValue, - accumulator: TileArrayValue, - *, - tile: Tile, -) -> TileArrayValue: - """Create a scalar multiply-accumulate operation on one hardware tile.""" - builder = _require_active_builder() - operands = (lhs, rhs, accumulator) - return builder.emit( - operands=operands, - result_dtype=lhs.dtype, - tile=tile, - create_operation=lambda result: MacOp( - result=result, - operands=operands, - tile=tile, - ), - ) diff --git a/tests/python/compiler/test_mac_kernel.py b/tests/python/compiler/test_mac_kernel.py deleted file mode 100644 index 9052260..0000000 --- a/tests/python/compiler/test_mac_kernel.py +++ /dev/null @@ -1,69 +0,0 @@ -import synapse -import synapse.language as synl -from synapse.frontend.lowering import lower - -EXPECTED_IR = """ -module { - func.func @f32_mac() { - taskflow.task @f32_mac : () -> () { - neura.kernel attributes {accelerator = "neura"} { - %0 = "neura.constant"() <{value = 2.000000e+00 : f32}> {placement = {x = 0 : i32, y = 0 : i32}} : () -> f32 - %1 = "neura.constant"() <{value = 3.000000e+00 : f32}> {placement = {x = 0 : i32, y = 0 : i32}} : () -> f32 - %2 = "neura.constant"() <{value = 0.000000e+00 : f32}> {placement = {x = 0 : i32, y = 0 : i32}} : () -> f32 - %3 = "neura.fmul_fadd"(%0, %1, %2) {placement = {x = 0 : i32, y = 0 : i32}} : (f32, f32, f32) -> f32 - neura.yield - } - taskflow.yield - } - return - } -} -""".strip() - - -def f32_mac(): - array = synl.TileArray(x_tiles=1, y_tiles=1) - - lhs = synl.constant(2.0, tile=array[0, 0]) - rhs = synl.constant(3.0, tile=array[0, 0]) - accumulator = synl.constant(0.0, tile=array[0, 0]) - - synl.mac( - lhs, - rhs, - accumulator, - tile=array[0, 0], - ) - - -def mapped_f32_mac(): - array = synl.TileArray(x_tiles=2, y_tiles=2) - - lhs = synl.constant(2.0, tile=array[0, 0]) - rhs = synl.constant(3.0, tile=array[1, 0]) - accumulator = synl.constant(0.0, tile=array[0, 1]) - - synl.mac( - lhs, - rhs, - accumulator, - tile=array[1, 1], - ) - - -def test_f32_mac_lowers_to_neura(): - actual = lower(f32_mac) - - assert actual.strip() == EXPECTED_IR - - -def test_f32_mac_compiles_to_mapped_neura(): - actual = synapse.compile( - mapped_f32_mac, - target="neura", - ) - - assert '"neura.fmul_fadd"' in actual - assert 'mapping_strategy = "template"' in actual - assert 'resource = "link"' in actual - assert "placement =" not in actual diff --git a/tests/python/compiler/test_systolic_gemm.py b/tests/python/compiler/test_systolic_gemm.py new file mode 100644 index 0000000..f46ea7d --- /dev/null +++ b/tests/python/compiler/test_systolic_gemm.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +import synapse +import synapse.language as synl +from synapse.frontend import lowering + +PRE_MAPPING_IR = """ +#map = affine_map<(d0) -> (d0, 0)> +#map1 = affine_map<(d0) -> (d0, 1)> +#map2 = affine_map<(d0) -> (d0, 2)> +#map3 = affine_map<(d0) -> (d0, 3)> +#map4 = affine_map<(d0, d1) -> (-d1 + 3, d0)> +module { + func.func @ws_gemm_4x4(%arg0: memref<4x4xi32>, %arg1: memref<4x4xi32>, %arg2: memref<4x4xi32>) -> memref<4x4xi32> { + %done_writes = taskflow.task @ws_gemm_4x4 will_reads(%arg0, %arg1 : memref<4x4xi32>, memref<4x4xi32>) will_writes(%arg2 : memref<4x4xi32>) [original_read_memrefs(%arg0, %arg1 : memref<4x4xi32>, memref<4x4xi32>), original_write_memrefs(%arg2 : memref<4x4xi32>)] : (memref<4x4xi32>, memref<4x4xi32>, memref<4x4xi32>) -> (memref<4x4xi32>) { + ^bb0(%arg3: memref<4x4xi32>, %arg4: memref<4x4xi32>, %arg5: memref<4x4xi32>): + %0:4 = taskflow.stream_read %arg3 maps [#map, #map1, #map2, #map3] : memref<4x4xi32> -> (i32, i32, i32, i32) + %1:4 = neura.kernel inputs(%0#0, %0#1, %0#2, %0#3, %arg4 : i32, i32, i32, i32, memref<4x4xi32>) attributes {accelerator = "neura", kernel_metadata = {kind = "template", template = {input_ports = [{direction = "west", kernel_input = 0 : i32, x = 0 : i32, y = 3 : i32}, {direction = "west", kernel_input = 1 : i32, x = 0 : i32, y = 2 : i32}, {direction = "west", kernel_input = 2 : i32, x = 0 : i32, y = 1 : i32}, {direction = "west", kernel_input = 3 : i32, x = 0 : i32, y = 0 : i32}], name = "systolic_array", output_ports = [{direction = "south", kernel_result = 0 : i32, x = 0 : i32, y = 0 : i32}, {direction = "south", kernel_result = 1 : i32, x = 1 : i32, y = 0 : i32}, {direction = "south", kernel_result = 2 : i32, x = 2 : i32, y = 0 : i32}, {direction = "south", kernel_result = 3 : i32, x = 3 : i32, y = 0 : i32}], stationary = {kernel_input = 4 : i32, map = #map4, mode = "weight"}}}} { + ^bb0(%arg6: i32, %arg7: i32, %arg8: i32, %arg9: i32, %arg10: memref<4x4xi32>): + %2 = "neura.mac"(%arg6) <{stationary = "weight"}> {placement = {x = 0 : i32, y = 3 : i32}} : (i32) -> i32 + %3 = "neura.mac"(%arg6) <{stationary = "weight"}> {placement = {x = 1 : i32, y = 3 : i32}} : (i32) -> i32 + %4 = "neura.mac"(%arg6) <{stationary = "weight"}> {placement = {x = 2 : i32, y = 3 : i32}} : (i32) -> i32 + %5 = "neura.mac"(%arg6) <{stationary = "weight"}> {placement = {x = 3 : i32, y = 3 : i32}} : (i32) -> i32 + %6 = "neura.mac"(%arg7, %2) <{stationary = "weight"}> {placement = {x = 0 : i32, y = 2 : i32}} : (i32, i32) -> i32 + %7 = "neura.mac"(%arg7, %3) <{stationary = "weight"}> {placement = {x = 1 : i32, y = 2 : i32}} : (i32, i32) -> i32 + %8 = "neura.mac"(%arg7, %4) <{stationary = "weight"}> {placement = {x = 2 : i32, y = 2 : i32}} : (i32, i32) -> i32 + %9 = "neura.mac"(%arg7, %5) <{stationary = "weight"}> {placement = {x = 3 : i32, y = 2 : i32}} : (i32, i32) -> i32 + %10 = "neura.mac"(%arg8, %6) <{stationary = "weight"}> {placement = {x = 0 : i32, y = 1 : i32}} : (i32, i32) -> i32 + %11 = "neura.mac"(%arg8, %7) <{stationary = "weight"}> {placement = {x = 1 : i32, y = 1 : i32}} : (i32, i32) -> i32 + %12 = "neura.mac"(%arg8, %8) <{stationary = "weight"}> {placement = {x = 2 : i32, y = 1 : i32}} : (i32, i32) -> i32 + %13 = "neura.mac"(%arg8, %9) <{stationary = "weight"}> {placement = {x = 3 : i32, y = 1 : i32}} : (i32, i32) -> i32 + %14 = "neura.mac"(%arg9, %10) <{stationary = "weight"}> {placement = {x = 0 : i32, y = 0 : i32}} : (i32, i32) -> i32 + %15 = "neura.mac"(%arg9, %11) <{stationary = "weight"}> {placement = {x = 1 : i32, y = 0 : i32}} : (i32, i32) -> i32 + %16 = "neura.mac"(%arg9, %12) <{stationary = "weight"}> {placement = {x = 2 : i32, y = 0 : i32}} : (i32, i32) -> i32 + %17 = "neura.mac"(%arg9, %13) <{stationary = "weight"}> {placement = {x = 3 : i32, y = 0 : i32}} : (i32, i32) -> i32 + neura.yield results(%14, %15, %16, %17 : i32, i32, i32, i32) + } : i32, i32, i32, i32 + taskflow.stream_write(%1#0, %1#1, %1#2, %1#3 : i32, i32, i32, i32) to %arg5 maps [#map, #map1, #map2, #map3] : memref<4x4xi32> + taskflow.yield done_writes(%arg5 : memref<4x4xi32>) + } + return %done_writes : memref<4x4xi32> + } +} +""".strip() + + +def ws_gemm_4x4( + A: synl.MemRef[4, 4, synl.i32], + B: synl.MemRef[4, 4, synl.i32], + C: synl.MemRef[4, 4, synl.i32], +): + array = synl.TileArray(x_tiles=4, y_tiles=4) + + partial_sums = [None] * array.x_tiles + + for k in range(array.y_tiles): + y = array.y_tiles - 1 - k + + # A[:, k] is streamed through one west boundary Port. + activation = synl.input_port( + A[:, k], + port=array.west[y], + ) + + for x in range(array.x_tiles): + partial_sums[x] = synl.mac( + activation, + partial_sums[x], + weight=B[k, x], + tile=array[x, y], + ) + + for x, result in enumerate(partial_sums): + # Each south Port writes one column of C. + synl.output_port( + result, + target=C[:, x], + port=array.south[x], + ) + + +def test_lowers_systolic_gemm_to_exact_pre_mapping_ir(): + actual = lowering.lower(ws_gemm_4x4) + + assert actual.strip() == PRE_MAPPING_IR diff --git a/tests/python/language/test_tile_array_program.py b/tests/python/language/test_tile_array_program.py index f1ce1f4..29c8ae6 100644 --- a/tests/python/language/test_tile_array_program.py +++ b/tests/python/language/test_tile_array_program.py @@ -48,32 +48,3 @@ def test_infers_supported_scalar_types(): assert integer.dtype == synl.i32 assert floating.dtype == synl.f32 assert explicit_f32.dtype == synl.f32 - - -def test_records_mac_operation(): - array = synl.TileArray(x_tiles=1, y_tiles=1) - builder = synl.tile_array_program.TileArrayBuilder() - - with builder: - lhs = synl.constant(2.0, tile=array[0, 0]) - rhs = synl.constant(3.0, tile=array[0, 0]) - accumulator = synl.constant(0.0, tile=array[0, 0]) - - result = synl.mac( - lhs, - rhs, - accumulator, - tile=array[0, 0], - ) - - program = builder.build() - mac_op = program.operations[-1] - - assert isinstance( - mac_op, - synl.tile_array_program.MacOp, - ) - assert mac_op.operands == (lhs, rhs, accumulator) - assert mac_op.result is result - assert mac_op.tile is array[0, 0] - assert result.dtype == synl.f32 From 5cd9f244c7d229d7d9bc2d0ac83167794634a405 Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Sun, 30 Aug 2026 19:39:15 +0800 Subject: [PATCH 06/19] Expose TileArray boundary ports --- python/synapse/language/spatial.py | 66 +++++++++++++++++---- tests/python/compiler/test_systolic_gemm.py | 5 +- tests/python/language/test_spatial.py | 36 +++++++++-- 3 files changed, 89 insertions(+), 18 deletions(-) diff --git a/python/synapse/language/spatial.py b/python/synapse/language/spatial.py index bb67ad5..1739304 100644 --- a/python/synapse/language/spatial.py +++ b/python/synapse/language/spatial.py @@ -8,6 +8,8 @@ core array of a multi-CGRA, AMD AIE/NPU, or Tenstorrent. """ +from typing import Literal + class Tile: """A hardware tile in a CGRA TileArray. @@ -22,6 +24,39 @@ def __init__(self, x: int, y: int, array: "TileArray"): self.y = y +PortDirection = Literal["west", "east", "north", "south"] + + +class Port: + """A boundary data port of a CGRA TileArray. + + ``direction`` identifies the array boundary. The ``x`` and ``y`` + coordinates identify the boundary tile attached to this Port. + + Ports expose hardware connectivity to the programming model. They do not + prescribe when data is transferred or introduce clock-based scheduling. + """ + + def __init__( + self, + *, + direction: PortDirection, + x: int, + y: int, + array: "TileArray", + ): + self.direction = direction + self.x = x + self.y = y + self.array = array + + @property + def tile(self) -> Tile: + """Return the boundary tile attached to this Port.""" + + return self.array[self.x, self.y] + + class TileArray: """A parameterized two-dimensional tile array. @@ -41,18 +76,27 @@ def __init__(self, x_tiles: int, y_tiles: int): self.x_tiles = x_tiles 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) - ] + self.tiles = tuple( + Tile(x=x, y=y, array=self) for y in range(y_tiles) for x in range(x_tiles) + ) + + self.west_ports = tuple( + Port(direction="west", x=0, y=y, array=self) for y in range(y_tiles) + ) + + self.east_ports = tuple( + Port(direction="east", x=x_tiles - 1, y=y, array=self) + for y in range(y_tiles) + ) - def tiles(self): - """Iterate over all tiles in the array. + self.north_ports = tuple( + Port(direction="north", x=x, y=y_tiles - 1, array=self) + for x in range(x_tiles) + ) - The iteration order is a Python programming convenience and - does not specify sequential hardware execution. - """ - for y_row in self._tiles: - yield from y_row + self.south_ports = tuple( + Port(direction="south", x=x, y=0, array=self) for x in range(x_tiles) + ) def __getitem__(self, coordinate: tuple[int, int]) -> Tile: """Return the tile at the given ``(x, y)`` coordinate.""" @@ -65,4 +109,4 @@ def __getitem__(self, coordinate: tuple[int, int]) -> Tile: f"TileArray({self.x_tiles}, {self.y_tiles})" ) - return self._tiles[y][x] + return self.tiles[y * self.x_tiles + x] diff --git a/tests/python/compiler/test_systolic_gemm.py b/tests/python/compiler/test_systolic_gemm.py index f46ea7d..83b7ac6 100644 --- a/tests/python/compiler/test_systolic_gemm.py +++ b/tests/python/compiler/test_systolic_gemm.py @@ -1,6 +1,5 @@ from __future__ import annotations -import synapse import synapse.language as synl from synapse.frontend import lowering @@ -59,7 +58,7 @@ def ws_gemm_4x4( # A[:, k] is streamed through one west boundary Port. activation = synl.input_port( A[:, k], - port=array.west[y], + port=array.west_ports[y], ) for x in range(array.x_tiles): @@ -75,7 +74,7 @@ def ws_gemm_4x4( synl.output_port( result, target=C[:, x], - port=array.south[x], + port=array.south_ports[x], ) diff --git a/tests/python/language/test_spatial.py b/tests/python/language/test_spatial.py index 1cab673..52cbc7a 100644 --- a/tests/python/language/test_spatial.py +++ b/tests/python/language/test_spatial.py @@ -15,7 +15,7 @@ def test_tile_array_is_parameterized(): def test_tile_array_exposes_tiles(): array = synl.TileArray(x_tiles=2, y_tiles=3) - coordinates = {(tile.x, tile.y) 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)} @@ -31,9 +31,37 @@ def test_access_tile_array_by_coordinate(): assert tile is array[1, 2] enumerated_tile = next( - candidate - for candidate in array.tiles() - if candidate.x == 1 and candidate.y == 2 + candidate for candidate in array.tiles if candidate.x == 1 and candidate.y == 2 ) assert tile is enumerated_tile + + +def test_tile_array_exposes_boundary_ports(): + array = synl.TileArray(x_tiles=2, y_tiles=3) + + assert [(port.direction, port.x, port.y) for port in array.west_ports] == [ + ("west", 0, 0), + ("west", 0, 1), + ("west", 0, 2), + ] + + assert [(port.direction, port.x, port.y) for port in array.east_ports] == [ + ("east", 1, 0), + ("east", 1, 1), + ("east", 1, 2), + ] + + assert [(port.direction, port.x, port.y) for port in array.north_ports] == [ + ("north", 0, 2), + ("north", 1, 2), + ] + + assert [(port.direction, port.x, port.y) for port in array.south_ports] == [ + ("south", 0, 0), + ("south", 1, 0), + ] + + assert array.west_ports[2].array is array + assert array.west_ports[2].tile is array[0, 2] + assert array.south_ports[1].tile is array[1, 0] From 7c8028482651936b5ed92b598e541a71ad2e9614 Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Sun, 30 Aug 2026 21:00:11 +0800 Subject: [PATCH 07/19] Generalize Synapse language types --- python/synapse/frontend/lowering.py | 24 +++---- python/synapse/language/__init__.py | 8 +-- python/synapse/language/tile_array_program.py | 41 ++++-------- python/synapse/language/types.py | 64 +++++++++++++++++++ tests/python/compiler/test_systolic_gemm.py | 6 +- tests/python/language/test_types.py | 20 ++++++ 6 files changed, 115 insertions(+), 48 deletions(-) create mode 100644 python/synapse/language/types.py create mode 100644 tests/python/language/test_types.py diff --git a/python/synapse/frontend/lowering.py b/python/synapse/frontend/lowering.py index 93f7d25..38831bc 100644 --- a/python/synapse/frontend/lowering.py +++ b/python/synapse/frontend/lowering.py @@ -11,8 +11,8 @@ TileArrayBuilder, TileArrayOp, TileArrayProgram, - TileArrayScalarType, ) +from synapse.language.types import DType def lower(program_fn: Callable) -> str: @@ -61,19 +61,19 @@ def _lower_tile_array_program( taskflow.register_dialect() neura.register_dialect() - i32 = IntegerType.get_signless(32) + i32_type = IntegerType.get_signless(32) - def get_mlir_type(dtype: TileArrayScalarType): - """Translate a frontend scalar type into an MLIR type.""" + def get_mlir_type(dtype: DType): + """Translate a frontend data type into an MLIR type.""" - if dtype == TileArrayScalarType.I32: - return i32 + if dtype == DType.I32: + return i32_type - if dtype == TileArrayScalarType.F32: + if dtype == DType.F32: return F32Type.get() raise NotImplementedError( - f"unsupported tile-array scalar type: {dtype.value}" + f"unsupported tile-array data type: {dtype.value}" ) def get_constant_attribute( @@ -82,10 +82,10 @@ def get_constant_attribute( ): """Build the typed MLIR attribute for a constant value.""" - if operation.result.dtype == TileArrayScalarType.I32: + if operation.result.dtype == DType.I32: return IntegerAttr.get(result_type, cast(int, operation.value)) - if operation.result.dtype == TileArrayScalarType.F32: + if operation.result.dtype == DType.F32: return FloatAttr.get(result_type, float(operation.value)) raise NotImplementedError( @@ -97,8 +97,8 @@ def get_placement(tile: Tile) -> DictAttr: return DictAttr.get( { - "x": IntegerAttr.get(i32, tile.x), - "y": IntegerAttr.get(i32, tile.y), + "x": IntegerAttr.get(i32_type, tile.x), + "y": IntegerAttr.get(i32_type, tile.y), } ) diff --git a/python/synapse/language/__init__.py b/python/synapse/language/__init__.py index 463d6f9..e75b889 100644 --- a/python/synapse/language/__init__.py +++ b/python/synapse/language/__init__.py @@ -1,14 +1,12 @@ """Public Synapse language API.""" from .spatial import TileArray -from .tile_array_program import TileArrayScalarType, add, constant - -i32 = TileArrayScalarType.I32 -f32 = TileArrayScalarType.F32 +from .tile_array_program import add, constant +from .types import DType, f32, i32 __all__ = [ + "DType", "TileArray", - "TileArrayScalarType", "add", "constant", "f32", diff --git a/python/synapse/language/tile_array_program.py b/python/synapse/language/tile_array_program.py index 4e26598..d74552f 100644 --- a/python/synapse/language/tile_array_program.py +++ b/python/synapse/language/tile_array_program.py @@ -10,21 +10,9 @@ from collections.abc import Callable from contextvars import ContextVar from dataclasses import dataclass, field -from enum import Enum 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" +from .types import DType, f32, i32 @dataclass(frozen=True) @@ -32,7 +20,7 @@ class TileArrayValue: """A typed value produced by one tile-array operation.""" id: int - dtype: TileArrayScalarType + dtype: DType _builder: TileArrayBuilder = field(repr=False) @@ -43,9 +31,8 @@ class TileArrayValue: 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. + ``operands`` may contain any number of input values. Constants use an + empty tuple, while operations such as add consume input values. """ result: TileArrayValue @@ -67,18 +54,16 @@ def __post_init__(self) -> None: ) dtype = self.result.dtype - if not isinstance(dtype, TileArrayScalarType): - raise TypeError("ConstantOp result must use a TileArrayScalarType") + if not isinstance(dtype, DType): + raise TypeError("ConstantOp result must use a DType") if isinstance(self.value, bool): raise TypeError("boolean constants are not supported yet") - if dtype == TileArrayScalarType.I32 and not isinstance(self.value, int): + if dtype == DType.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) - ): + if dtype == DType.F32 and not isinstance(self.value, (int, float)): raise TypeError("an f32 constant requires a numeric value") @@ -177,7 +162,7 @@ def emit( self, *, operands: tuple[TileArrayValue, ...], - result_dtype: TileArrayScalarType, + result_dtype: DType, tile: Tile, create_operation: Callable[[TileArrayValue], TileArrayOp], ) -> TileArrayValue: @@ -248,19 +233,19 @@ def _require_active_builder() -> TileArrayBuilder: # User-facing tile-array program DSL # --------------------------------------------------------------- def constant( - value: int | float, *, tile: Tile, dtype: TileArrayScalarType | None = None + value: int | float, *, tile: Tile, dtype: DType | 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) + constant(1.0, tile=tile, dtype=DType.F32) """ if dtype is None: if type(value) is int: - dtype = TileArrayScalarType.I32 + dtype = DType.I32 elif type(value) is float: - dtype = TileArrayScalarType.F32 + dtype = DType.F32 else: raise TypeError( "constant currently supports integer and floating-point values" diff --git a/python/synapse/language/types.py b/python/synapse/language/types.py new file mode 100644 index 0000000..9ddec17 --- /dev/null +++ b/python/synapse/language/types.py @@ -0,0 +1,64 @@ +"""Types exposed by the Synapse programming language. + +These types describe program values independently of a particular hardware +hierarchy or compiler IR. Compiler lowering later decides whether a shaped +value becomes a tensor, MemRef, stream source, or another backend type. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum + + +class DType(Enum): + """A data type supported by Synapse.""" + + I32 = "i32" + F32 = "f32" + + def __getitem__(self, shape: int | tuple[int, ...]) -> ShapedType: + """Create a shaped type with this data element type.""" + + if not isinstance(shape, tuple): + shape = (shape,) + + return ShapedType( + shape=shape, + dtype=self, + ) + + def __str__(self) -> str: + """Return the source-level spelling of this type.""" + + return self.value + + +@dataclass(frozen=True) +class ShapedType: + """The shape and scalar element type of a Synapse program value. + + ShapedType does not prescribe whether compiler lowering uses a tensor, + MemRef, or another backend representation. + """ + + shape: tuple[int, ...] + dtype: DType + + def __post_init__(self) -> None: + """Validate the dimensions and scalar element type.""" + + if not self.shape: + raise TypeError("a shaped type requires at least one dimension") + + if not all( + type(dimension) is int and dimension > 0 for dimension in self.shape + ): + raise TypeError("shape dimensions must be positive integers") + + if not isinstance(self.dtype, DType): + raise TypeError("dtype must be a DType") + + +i32 = DType.I32 +f32 = DType.F32 diff --git a/tests/python/compiler/test_systolic_gemm.py b/tests/python/compiler/test_systolic_gemm.py index 83b7ac6..a9303a3 100644 --- a/tests/python/compiler/test_systolic_gemm.py +++ b/tests/python/compiler/test_systolic_gemm.py @@ -44,9 +44,9 @@ def ws_gemm_4x4( - A: synl.MemRef[4, 4, synl.i32], - B: synl.MemRef[4, 4, synl.i32], - C: synl.MemRef[4, 4, synl.i32], + A: synl.i32[4, 4], + B: synl.i32[4, 4], + C: synl.i32[4, 4], ): array = synl.TileArray(x_tiles=4, y_tiles=4) diff --git a/tests/python/language/test_types.py b/tests/python/language/test_types.py new file mode 100644 index 0000000..3f040d1 --- /dev/null +++ b/tests/python/language/test_types.py @@ -0,0 +1,20 @@ +import synapse.language as synl +from synapse.language.types import DType, ShapedType + + +def test_dtypes_create_shaped_types(): + matrix = synl.i32[4, 4] + vector = synl.f32[8] + + assert synl.i32 == DType.I32 + assert synl.f32 == DType.F32 + + assert matrix == ShapedType( + shape=(4, 4), + dtype=synl.i32, + ) + + assert vector == ShapedType( + shape=(8,), + dtype=synl.f32, + ) From 8e6f834e29a340e8c1f45642c7b3993d442a89f8 Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Tue, 1 Sep 2026 13:31:02 +0800 Subject: [PATCH 08/19] Add tensor-based TileArray program representation --- python/synapse/language/__init__.py | 18 +- python/synapse/language/tensor.py | 65 +++++ python/synapse/language/tile_array_program.py | 252 ++++++++++++++++-- python/synapse/language/types.py | 22 +- tests/python/language/test_types.py | 6 +- 5 files changed, 329 insertions(+), 34 deletions(-) create mode 100644 python/synapse/language/tensor.py diff --git a/python/synapse/language/__init__.py b/python/synapse/language/__init__.py index e75b889..64b7b97 100644 --- a/python/synapse/language/__init__.py +++ b/python/synapse/language/__init__.py @@ -1,14 +1,28 @@ """Public Synapse language API.""" from .spatial import TileArray -from .tile_array_program import add, constant -from .types import DType, f32, i32 +from .tensor import Tensor +from .tile_array_program import ( + StationaryMode, + add, + constant, + input_port, + mac, + output_port, +) +from .types import DType, TensorType, f32, i32 __all__ = [ "DType", + "StationaryMode", + "Tensor", + "TensorType", "TileArray", "add", "constant", "f32", "i32", + "input_port", + "mac", + "output_port", ] diff --git a/python/synapse/language/tensor.py b/python/synapse/language/tensor.py new file mode 100644 index 0000000..f9aca45 --- /dev/null +++ b/python/synapse/language/tensor.py @@ -0,0 +1,65 @@ +"""Tensor values used by the Synapse programming language.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from .types import DType, TensorType + +TensorIndex = int | slice + + +@dataclass(frozen=True) +class Tensor: + """A symbolic tensor value passed to a Synapse program.""" + + name: str + type: TensorType + + def __getitem__( + self, indices: TensorIndex | tuple[TensorIndex, ...] + ) -> TensorAccess: + """Describes a scalar element or full-dimensional slice of this tensor.""" + if not isinstance(indices, tuple): + indices = (indices,) + + if len(indices) != len(self.type.shape): + raise IndexError( + f"expected {len(self.type.shape)} indices, but got {len(indices)}" + ) + + for dimension, index in zip(self.type.shape, indices): + if type(index) is int: + if not (0 <= index < dimension): + raise IndexError( + f"index {index} is outside dimension size {dimension}" + ) + continue + + if isinstance(index, slice): + if index != slice(None): + raise ValueError("only full slices are supported initially") + continue + + raise TypeError("indices must be integers or full slices") + + return TensorAccess(source=self, indices=indices) + + +@dataclass(frozen=True) +class TensorAccess: + """A symbolic element or slice selected from a Tensor.""" + + source: Tensor + indices: tuple[TensorIndex, ...] + + @property + def dtype(self) -> DType: + """Returns the scalar element type of the tensor access.""" + return self.source.type.dtype + + @property + def is_scalar(self) -> bool: + """Returns whether this access identifies one scalar element.""" + + return all(type(index) is int for index in self.indices) diff --git a/python/synapse/language/tile_array_program.py b/python/synapse/language/tile_array_program.py index d74552f..bcdb828 100644 --- a/python/synapse/language/tile_array_program.py +++ b/python/synapse/language/tile_array_program.py @@ -10,9 +10,19 @@ from collections.abc import Callable from contextvars import ContextVar from dataclasses import dataclass, field +from enum import Enum -from .spatial import Tile, TileArray -from .types import DType, f32, i32 +from .spatial import Port, Tile, TileArray +from .tensor import Tensor, TensorAccess +from .types import DType + + +class StationaryMode(Enum): + """The dataflow operand retained locally by configured MACs.""" + + WEIGHT = "weight" + INPUT = "input" + OUTPUT = "output" @dataclass(frozen=True) @@ -24,6 +34,24 @@ class TileArrayValue: _builder: TileArrayBuilder = field(repr=False) +@dataclass(frozen=True) +class InputPortBinding: + """A tensor slice bound to one TileArray input Port.""" + + input_des: TileArrayValue + input_src: TensorAccess + port: Port + + +@dataclass(frozen=True) +class OutputPortBinding: + """A TileArray value bound to a tensor output slice and Port.""" + + output_src: TileArrayValue + output_des: TensorAccess + port: Port + + # --------------------------------------------------------------- # Typed tile-array operations # --------------------------------------------------------------- @@ -87,12 +115,56 @@ def __post_init__(self) -> None: raise TypeError("AddOp result type must match its operand type") +@dataclass(frozen=True) +class MacOp(TileArrayOp): + """A configured MAC using one stationary scalar value.""" + + stationary_value: TensorAccess + stationary_mode: StationaryMode + + def __post_init__(self) -> None: + """Validates configured MAC operands and stationary data.""" + + if self.stationary_mode in ( + StationaryMode.WEIGHT, + StationaryMode.INPUT, + ) and len(self.operands) not in (1, 2): + raise ValueError( + "weight/input-stationary MacOp requires " + "an input and optional partial sum" + ) + + if self.stationary_mode == StationaryMode.OUTPUT and len(self.operands) != 2: + raise ValueError("output-stationary MacOp requires two multiplicands") + + if not self.stationary_value.is_scalar: + raise ValueError("stationary data must identify one scalar") + if any(operand.dtype != self.result.dtype for operand in self.operands): + raise TypeError("MacOp operands and result must have the same dtype") + if self.stationary_value.dtype != self.result.dtype: + raise TypeError("MacOp stationary data and result must have the same dtype") + + +@dataclass(frozen=True) +class StationaryBinding: + """Stationary data assigned to the tiles of one template program.""" + + mode: StationaryMode + source: Tensor + tile_values: tuple[tuple[Tile, TensorAccess], ...] + + @dataclass(frozen=True) class TileArrayProgram: """A tile-array program produced by TileArrayBuilder.""" array: TileArray + arguments: tuple[Tensor, ...] + input_ports: tuple[InputPortBinding, ...] + output_ports: tuple[OutputPortBinding, ...] operations: tuple[TileArrayOp, ...] + template_name: str | None + stationary: StationaryBinding | None # --------------------------------------------------------------- @@ -105,15 +177,18 @@ class TileArrayBuilder: is called, it returns a TileArrayProgram. """ - def __init__(self): + def __init__(self, arguments: tuple[Tensor, ...] = ()): + self._arguments = arguments self._array: TileArray | None = None + self._input_ports: list[InputPortBinding] = [] + self._output_ports: list[OutputPortBinding] = [] 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.""" + """Makes 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: @@ -122,27 +197,25 @@ def __enter__(self): return self def __exit__(self, exc_type, exc_value, traceback) -> None: - """Restore the previously active builder.""" + """Restores the previously active builder.""" if self._token is None: raise RuntimeError("TileArrayBuilder is not active") _active_builder.reset(self._token) self._token = None - 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") + def _bind_tile_array(self, array: TileArray) -> None: + """Binds the program to one TileArray.""" if self._array is None: - self._array = tile.array + self._array = array return - if tile.array is not self._array: + if array is not self._array: raise ValueError( "all operations in a TileArrayProgram must use tiles from the same TileArray" ) def _validate_operand_for_builder(self, operand: TileArrayValue) -> None: - """Validate that an operand was produced by this builder.""" + """Validates that an operand was produced by this builder.""" if not isinstance(operand, TileArrayValue): raise TypeError("operation operand must be a TileArrayValue") @@ -152,7 +225,7 @@ def _validate_operand_for_builder(self, operand: TileArrayValue) -> None: ) def _ensure_not_built(self) -> None: - """Reject operations emitted after the program has been built.""" + """Rejects operations emitted after the program has been built.""" if self._is_built: raise RuntimeError( "cannot emit operations after building a TileArrayProgram" @@ -166,9 +239,9 @@ def emit( tile: Tile, create_operation: Callable[[TileArrayValue], TileArrayOp], ) -> TileArrayValue: - """Create and record one tile-array operation.""" + """Creates and record one tile-array operation.""" self._ensure_not_built() - self._bind_tile_array(tile) + self._bind_tile_array(tile.array) for operand in operands: self._validate_operand_for_builder(operand) @@ -195,6 +268,47 @@ def emit( self._next_value_id += 1 return result + def add_input_port(self, *, source: TensorAccess, port: Port) -> TileArrayValue: + """Creates one scalar kernel input and bind it to a Port.""" + self._ensure_not_built() + if source.is_scalar: + raise ValueError("an input Port requires a tensor slice") + + if sum(isinstance(index, slice) for index in source.indices) != 1: + raise ValueError("an input Port currently supports one varying dimension") + + self._bind_tile_array(port.array) + + result = TileArrayValue( + id=self._next_value_id, dtype=source.dtype, _builder=self + ) + self._next_value_id += 1 + self._input_ports.append( + InputPortBinding(input_src=source, input_des=result, port=port) + ) + return result + + def add_output_port( + self, *, value: TileArrayValue, target: TensorAccess, port: Port + ) -> None: + """Bind one kernel result to an output slice and Port.""" + self._ensure_not_built() + self._validate_operand_for_builder(value) + + if target.is_scalar: + raise ValueError("an output Port requires a tensor slice") + if sum(isinstance(index, slice) for index in target.indices) != 1: + raise ValueError("an output Port currently supports one varying dimension") + + if target.dtype != value.dtype: + raise TypeError("output value and target must have the same dtype") + + self._bind_tile_array(port.array) + + self._output_ports.append( + OutputPortBinding(output_src=value, output_des=target, port=port) + ) + def build(self) -> TileArrayProgram: """Finish recording and return a program.""" if self._token is not None: @@ -205,8 +319,52 @@ def build(self) -> TileArrayProgram: raise RuntimeError( "cannot build an empty TileArrayProgram without a TileArray" ) + + mac_operations = [ + operation for operation in self._operations if isinstance(operation, MacOp) + ] + + stationary = None + template_name = None + + if mac_operations: + stationary_mode = mac_operations[0].stationary_mode + stationary_source = mac_operations[0].stationary_value.source + + if any( + operation.stationary_mode != stationary_mode + for operation in mac_operations + ): + raise ValueError( + "all configured MACs in one template must use one stationary mode" + ) + if any( + operation.stationary_value.source is not stationary_source + for operation in mac_operations + ): + raise ValueError("all configured MACs must use one stationary source") + stationary = StationaryBinding( + mode=stationary_mode, + source=stationary_source, + tile_values=tuple( + (operation.tile, operation.stationary_value) + for operation in mac_operations + ), + ) + # Configured MAC networks currently use the systolic template. + # Template registration will replace this inference later. + template_name = "systolic_array" + self._is_built = True - return TileArrayProgram(array=self._array, operations=tuple(self._operations)) + return TileArrayProgram( + array=self._array, + arguments=self._arguments, + input_ports=tuple(self._input_ports), + output_ports=tuple(self._output_ports), + operations=tuple(self._operations), + template_name=template_name, + stationary=stationary, + ) # The active builder is compiler-internal state. Public DSL calls use it to @@ -287,3 +445,65 @@ def add( tile=tile, ), ) + + +def input_port(source: TensorAccess, *, port: Port) -> TileArrayValue: + """Reads one tensor slice through a TileArray input port.""" + if not isinstance(source, TensorAccess): + raise TypeError("input_port source must be a tensor access") + if not isinstance(port, Port): + raise TypeError("input_port port must be a Port") + return _require_active_builder().add_input_port(source=source, port=port) + + +def output_port(value: TileArrayValue, *, target: TensorAccess, port: Port) -> None: + """Write one TileArray value stream through an output Port.""" + + if not isinstance(target, TensorAccess): + raise TypeError("output_port target must be a tensor access") + + if not isinstance(port, Port): + raise TypeError("output_port port must be a Port") + + _require_active_builder().add_output_port(value=value, target=target, port=port) + + +def mac( + input0: TileArrayValue, + input1: TileArrayValue | None, + *, + stationary: TensorAccess, + mode: StationaryMode, + tile: Tile, +) -> TileArrayValue: + """Create one configured MAC operation. + + Operand meanings depend on the selected stationary mode: + + - WEIGHT: input0 is activation, input1 is an optional partial sum. + - INPUT: input0 is weight, input1 is an optional partial sum. + - OUTPUT: input0 and input1 are the two multiplicands. + """ + + if not isinstance(stationary, TensorAccess): + raise TypeError("mac stationary data must be a tensor access") + + if not isinstance(mode, StationaryMode): + raise TypeError("mac mode must be a StationaryMode") + + operands = (input0,) if input1 is None else (input0, input1) + + builder = _require_active_builder() + + return builder.emit( + operands=operands, + result_dtype=input0.dtype, + tile=tile, + create_operation=lambda result: MacOp( + result=result, + operands=operands, + tile=tile, + stationary_value=stationary, + stationary_mode=mode, + ), + ) diff --git a/python/synapse/language/types.py b/python/synapse/language/types.py index 9ddec17..fc89d18 100644 --- a/python/synapse/language/types.py +++ b/python/synapse/language/types.py @@ -17,16 +17,13 @@ class DType(Enum): I32 = "i32" F32 = "f32" - def __getitem__(self, shape: int | tuple[int, ...]) -> ShapedType: - """Create a shaped type with this data element type.""" + def __getitem__(self, shape: int | tuple[int, ...]) -> TensorType: + """Create a tensor type with this data element type.""" if not isinstance(shape, tuple): shape = (shape,) - return ShapedType( - shape=shape, - dtype=self, - ) + return TensorType(shape=shape, dtype=self) def __str__(self) -> str: """Return the source-level spelling of this type.""" @@ -35,26 +32,25 @@ def __str__(self) -> str: @dataclass(frozen=True) -class ShapedType: - """The shape and scalar element type of a Synapse program value. +class TensorType: + """The shape and scalar element type of a Synapse tensor. - ShapedType does not prescribe whether compiler lowering uses a tensor, - MemRef, or another backend representation. + TensorType does not prescribe its later bufferized representation. """ shape: tuple[int, ...] dtype: DType def __post_init__(self) -> None: - """Validate the dimensions and scalar element type.""" + """Validates the dimensions and element type.""" if not self.shape: - raise TypeError("a shaped type requires at least one dimension") + raise TypeError("a tensor type requires at least one dimension") if not all( type(dimension) is int and dimension > 0 for dimension in self.shape ): - raise TypeError("shape dimensions must be positive integers") + raise TypeError("tensor dimensions must be positive integers") if not isinstance(self.dtype, DType): raise TypeError("dtype must be a DType") diff --git a/tests/python/language/test_types.py b/tests/python/language/test_types.py index 3f040d1..acc0525 100644 --- a/tests/python/language/test_types.py +++ b/tests/python/language/test_types.py @@ -1,5 +1,5 @@ import synapse.language as synl -from synapse.language.types import DType, ShapedType +from synapse.language.types import DType, TensorType def test_dtypes_create_shaped_types(): @@ -9,12 +9,12 @@ def test_dtypes_create_shaped_types(): assert synl.i32 == DType.I32 assert synl.f32 == DType.F32 - assert matrix == ShapedType( + assert matrix == TensorType( shape=(4, 4), dtype=synl.i32, ) - assert vector == ShapedType( + assert vector == TensorType( shape=(8,), dtype=synl.f32, ) From 4774a9e2152b1901bca3a79414e2cbf58aedc927 Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Tue, 1 Sep 2026 13:31:22 +0800 Subject: [PATCH 09/19] Lower systolic GEMM to Taskflow and Neura --- python/synapse/frontend/lowering.py | 307 ++++++++++++++++++-- tests/python/compiler/test_systolic_gemm.py | 30 +- 2 files changed, 296 insertions(+), 41 deletions(-) diff --git a/python/synapse/frontend/lowering.py b/python/synapse/frontend/lowering.py index 38831bc..a04c933 100644 --- a/python/synapse/frontend/lowering.py +++ b/python/synapse/frontend/lowering.py @@ -2,34 +2,57 @@ from collections.abc import Callable from functools import singledispatch +from inspect import signature from typing import cast from synapse.language.spatial import Tile +from synapse.language.tensor import Tensor, TensorAccess from synapse.language.tile_array_program import ( AddOp, ConstantOp, + MacOp, + StationaryMode, TileArrayBuilder, TileArrayOp, TileArrayProgram, ) -from synapse.language.types import DType +from synapse.language.types import DType, TensorType -def lower(program_fn: Callable) -> str: - """Lower one tile-array program to pre-mapping Taskflow and Neura IR.""" +def lower(program_fn: Callable, *, argument_types: tuple[TensorType, ...] = ()) -> str: + """Lowers one tile-array program to pre-mapping Taskflow and Neura IR.""" + # TODO: Generalize program arguments beyond Tensor. Accept both DType and + # TensorType, create Scalar or Tensor symbolic values accordingly, and + # lower scalar dependencies through Taskflow value_inputs/value_outputs. - builder = TileArrayBuilder() + function_signature = signature(program_fn) + parameter_names = tuple(function_signature.parameters) - # Execute the user's tile-array DSL while recording its operations. + if len(parameter_names) != len(argument_types): + raise TypeError( + f"program {program_fn.__name__} expects {len(parameter_names)} arguments, " + f"but got {len(argument_types)} argument types" + ) + + if any( + not isinstance(argument_type, TensorType) for argument_type in argument_types + ): + raise TypeError("program argument types must be TensorType values") + + arguments = [ + Tensor(name=name, type=argument_type) + for name, argument_type in zip(parameter_names, argument_types) + ] + + builder = TileArrayBuilder(arguments=tuple(arguments)) + + # Executes the user's tile-array DSL while recording its operations. with builder: - program_fn() + program_fn(*arguments) program = builder.build() - return _lower_tile_array_program( - program_name=program_fn.__name__, - program=program, - ) + return _lower_tile_array_program(program_name=program_fn.__name__, program=program) def _lower_tile_array_program( @@ -45,6 +68,10 @@ def _lower_tile_array_program( from taskflow_mlir.dialects import func, neura, taskflow from taskflow_mlir.ir import ( + AffineExpr, + AffineMap, + AffineMapAttr, + ArrayAttr, Context, DictAttr, F32Type, @@ -53,6 +80,7 @@ def _lower_tile_array_program( IntegerAttr, IntegerType, Location, + MemRefType, Module, StringAttr, ) @@ -76,6 +104,54 @@ def get_mlir_type(dtype: DType): f"unsupported tile-array data type: {dtype.value}" ) + def get_memref_type(tensor_type: TensorType): + """Translates a TensorType into a Taskflow MemRef type.""" + return MemRefType.get( + list(tensor_type.shape), get_mlir_type(tensor_type.dtype) + ) + + def get_access_map(access: TensorAccess) -> AffineMapAttr: + """Builds the affine map that enumerates one tensor access.""" + + results = [] + next_dimension = 0 + + for index in access.indices: + if isinstance(index, slice): + results.append(AffineExpr.get_dim(next_dimension)) + next_dimension += 1 + else: + results.append(AffineExpr.get_constant(index)) + + return AffineMapAttr.get(AffineMap.get(next_dimension, 0, results)) + + def get_stationary_map() -> AffineMapAttr: + """Builds and validates the stationary Tile-to-data map.""" + if program.stationary is None: + raise RuntimeError("template program requires stationary data") + + if program.stationary.mode != StationaryMode.WEIGHT: + raise NotImplementedError( + "only weight-stationary template lowering is implemented initially" + ) + + for tile, access in program.stationary.tile_values: + expected_indices = (program.array.y_tiles - 1 - tile.y, tile.x) + + if access.indices != expected_indices: + raise ValueError("weight-stationary GEMM requires B[K - 1 - y, x]") + + x = AffineExpr.get_dim(0) + y = AffineExpr.get_dim(1) + + last_row = AffineExpr.get_constant(program.array.y_tiles - 1) + + negative_y = AffineExpr.get_mul(AffineExpr.get_constant(-1), y) + + weight_row = AffineExpr.get_add(last_row, negative_y) + + return AffineMapAttr.get(AffineMap.get(2, 0, [weight_row, x])) + def get_constant_attribute( operation: ConstantOp, result_type, @@ -102,6 +178,67 @@ def get_placement(tile: Tile) -> DictAttr: } ) + def get_kernel_metadata() -> DictAttr | None: + """Materialize template, stationary, and Port metadata.""" + + if program.template_name is None: + return None + + if program.stationary is None: + raise RuntimeError("template program requires stationary metadata") + + input_ports = ArrayAttr.get( + [ + DictAttr.get( + { + "kernel_input": IntegerAttr.get(i32_type, index), + "direction": StringAttr.get(binding.port.direction), + "x": IntegerAttr.get(i32_type, binding.port.x), + "y": IntegerAttr.get(i32_type, binding.port.y), + } + ) + for index, binding in enumerate(program.input_ports) + ] + ) + + output_ports = ArrayAttr.get( + [ + DictAttr.get( + { + "kernel_result": IntegerAttr.get(i32_type, index), + "direction": StringAttr.get(binding.port.direction), + "x": IntegerAttr.get(i32_type, binding.port.x), + "y": IntegerAttr.get(i32_type, binding.port.y), + } + ) + for index, binding in enumerate(program.output_ports) + ] + ) + + stationary = DictAttr.get( + { + "kernel_input": IntegerAttr.get(i32_type, len(program.input_ports)), + "map": get_stationary_map(), + "mode": StringAttr.get(program.stationary.mode.value), + } + ) + + template = DictAttr.get( + { + "input_ports": input_ports, + "name": StringAttr.get(program.template_name), + "output_ports": output_ports, + "stationary": stationary, + } + ) + + return DictAttr.get( + { + "kind": StringAttr.get("template"), + "template": template, + } + ) + @singledispatch def lower_operation(operation: TileArrayOp, operands, result_type): """Lower one frontend TileArray operation to a Neura operation. @@ -127,48 +264,159 @@ def lower_add(operation: AddOp, operands, result_type): return neura.AddOp(result_type, lhs, rhs=rhs) + @lower_operation.register + def lower_mac(operation: MacOp, operands, result_type): + """Lower a configured MacOp to neura.mac.""" + + input0 = operands[0] + + input1 = operands[1] if len(operands) == 2 else None + + return neura.MacOp( + result_type, + input0, + StringAttr.get(operation.stationary_mode.value), + input1=input1, + ) + + read_sources = {binding.input_src.source for binding in program.input_ports} + + if program.stationary is not None: + read_sources.add(program.stationary.source) + + write_sources = {binding.output_des.source for binding in program.output_ports} + + read_arguments = tuple( + argument for argument in program.arguments if argument in read_sources + ) + + write_arguments = tuple( + argument for argument in program.arguments if argument in write_sources + ) + + function_argument_types = [ + get_memref_type(argument.type) for argument in program.arguments + ] + + function_result_types = [ + get_memref_type(argument.type) for argument in write_arguments + ] + 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 = func.FuncOp( + program_name, + (function_argument_types, function_result_types), + ) + function_block = function.add_entry_block() + function_values = dict(zip(program.arguments, function_block.arguments)) + + read_values = [function_values[argument] for argument in read_arguments] + + write_values = [function_values[argument] for argument in write_arguments] + with InsertionPoint(function_block): task = taskflow.TaskflowTaskOp( done_reads=[], - done_writes=[], + done_writes=[value.type for value in write_values], value_outputs=[], - will_reads=[], - will_writes=[], + will_reads=read_values, + will_writes=write_values, value_inputs=[], task_name=program_name, - original_read_memrefs=[], - original_write_memrefs=[], + original_read_memrefs=read_values, + original_write_memrefs=write_values, + ) + + task_block = task.body.blocks.append( + *[value.type for value in read_values], + *[value.type for value in write_values], ) - task_block = task.body.blocks.append() - func.ReturnOp([]) + func.ReturnOp(task.done_writes) + + task_arguments = dict( + zip(read_arguments + write_arguments, task_block.arguments) + ) with InsertionPoint(task_block): + stream_values_by_id = {} + + for argument in read_arguments: + bindings = [ + binding + for binding in program.input_ports + if binding.input_src.source is argument + ] + + if not bindings: + continue + + stream_read = taskflow.TaskflowStreamReadOp( + [get_mlir_type(binding.input_des.dtype) for binding in bindings], + task_arguments[argument], + ArrayAttr.get( + [get_access_map(binding.input_src) for binding in bindings] + ), + ) + + for binding, value in zip(bindings, stream_read.values): + stream_values_by_id[binding.input_des.id] = value + + kernel_input_values = [ + stream_values_by_id[binding.input_des.id] + for binding in program.input_ports + ] + + kernel_input_types = [value.type for value in kernel_input_values] + + if program.stationary is not None: + stationary_value = task_arguments[program.stationary.source] + + kernel_input_values.append(stationary_value) + + kernel_input_types.append(stationary_value.type) + + kernel_output_types = [ + get_mlir_type(binding.output_src.dtype) + for binding in program.output_ports + ] + kernel = neura.KernelOp( - outputs=[], - inputs=[], + outputs=kernel_output_types, + inputs=kernel_input_values, iter_args_init=[], accelerator=StringAttr.get("neura"), + kernel_metadata=get_kernel_metadata(), ) - kernel_block = kernel.body.blocks.append() + + kernel_block = kernel.body.blocks.append(*kernel_input_types) + + if program.output_ports: + taskflow.TaskflowStreamWriteOp( + list(kernel.results), + task_arguments[write_arguments[0]], + ArrayAttr.get( + [ + get_access_map(binding.output_des) + for binding in program.output_ports + ] + ), + ) taskflow.TaskflowYieldOp( done_reads=[], - done_writes=[], + done_writes=[task_arguments[argument] for argument in write_arguments], value_results=[], ) - # Map frontend value IDs to the MLIR SSA values produced while - # lowering the recorded operations. - values_by_id = {} + values_by_id = { + binding.input_des.id: kernel_block.arguments[index] + for index, binding in enumerate(program.input_ports) + } with InsertionPoint(kernel_block): for operation in program.operations: @@ -188,7 +436,10 @@ def lower_add(operation: AddOp, operands, result_type): neura.YieldOp( iter_args_next=[], - results_=[], + results_=[ + values_by_id[binding.output_src.id] + for binding in program.output_ports + ], ) if not module.operation.verify(): diff --git a/tests/python/compiler/test_systolic_gemm.py b/tests/python/compiler/test_systolic_gemm.py index a9303a3..77c229f 100644 --- a/tests/python/compiler/test_systolic_gemm.py +++ b/tests/python/compiler/test_systolic_gemm.py @@ -1,5 +1,3 @@ -from __future__ import annotations - import synapse.language as synl from synapse.frontend import lowering @@ -43,14 +41,10 @@ """.strip() -def ws_gemm_4x4( - A: synl.i32[4, 4], - B: synl.i32[4, 4], - C: synl.i32[4, 4], -): +def ws_gemm_4x4(A: synl.Tensor, B: synl.Tensor, C: synl.Tensor): array = synl.TileArray(x_tiles=4, y_tiles=4) - partial_sums = [None] * array.x_tiles + partial_sums = [] for k in range(array.y_tiles): y = array.y_tiles - 1 - k @@ -61,13 +55,16 @@ def ws_gemm_4x4( port=array.west_ports[y], ) - for x in range(array.x_tiles): - partial_sums[x] = synl.mac( + partial_sums = [ + synl.mac( activation, - partial_sums[x], - weight=B[k, x], + partial_sums[x] if partial_sums else None, + stationary=B[k, x], + mode=synl.StationaryMode.WEIGHT, tile=array[x, y], ) + for x in range(array.x_tiles) + ] for x, result in enumerate(partial_sums): # Each south Port writes one column of C. @@ -79,6 +76,13 @@ def ws_gemm_4x4( def test_lowers_systolic_gemm_to_exact_pre_mapping_ir(): - actual = lowering.lower(ws_gemm_4x4) + actual = lowering.lower( + ws_gemm_4x4, + argument_types=( + synl.i32[4, 4], + synl.i32[4, 4], + synl.i32[4, 4], + ), + ) assert actual.strip() == PRE_MAPPING_IR From 6c58830b28e790e8190111c0be4e5a8d3b2212ad Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Tue, 1 Sep 2026 13:31:38 +0800 Subject: [PATCH 10/19] Clean up frontend parser typing imports --- python/synapse/frontend/parser.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/python/synapse/frontend/parser.py b/python/synapse/frontend/parser.py index 43db617..8e43246 100644 --- a/python/synapse/frontend/parser.py +++ b/python/synapse/frontend/parser.py @@ -1,10 +1,9 @@ -from __future__ import annotations - import ast import inspect import textwrap +from collections.abc import Callable from dataclasses import dataclass -from typing import Callable + @dataclass(frozen=True) class ParsedFunction: From 5c0fe13c82173f22216fa0f393c5744a01543e1c Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Tue, 1 Sep 2026 17:27:58 +0800 Subject: [PATCH 11/19] Update tests --- python/synapse/compiler/compiler.py | 11 ++- tests/python/compiler/test_systolic_gemm.py | 87 +++++++++++++++++++++ 2 files changed, 96 insertions(+), 2 deletions(-) diff --git a/python/synapse/compiler/compiler.py b/python/synapse/compiler/compiler.py index 1213525..5492a41 100644 --- a/python/synapse/compiler/compiler.py +++ b/python/synapse/compiler/compiler.py @@ -6,9 +6,15 @@ from tempfile import TemporaryDirectory from synapse.frontend.lowering import lower +from synapse.language.types import TensorType -def compile(program: Callable, *, target: str) -> str: +def compile( + program: Callable, + *, + target: str, + argument_types: tuple[TensorType, ...] = (), +) -> 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. @@ -16,7 +22,7 @@ def compile(program: Callable, *, target: str) -> str: raise ValueError(f"unsupported compilation target: {target}") # TODO: Support the amoeba backend. - neura_ir = lower(program) + neura_ir = lower(program, argument_types=argument_types) return _run_neura_backend(neura_ir) @@ -41,6 +47,7 @@ def _run_neura_backend(neura_ir: str) -> str: command = [ str(amoeba_opt), + "--promote-input-arg-to-const", "--leverage-predicated-value", "--insert-data-mov", ( diff --git a/tests/python/compiler/test_systolic_gemm.py b/tests/python/compiler/test_systolic_gemm.py index 77c229f..5c709f5 100644 --- a/tests/python/compiler/test_systolic_gemm.py +++ b/tests/python/compiler/test_systolic_gemm.py @@ -1,3 +1,4 @@ +import synapse import synapse.language as synl from synapse.frontend import lowering @@ -41,6 +42,78 @@ """.strip() +MAPPED_IR = """ +#map = affine_map<(d0) -> (d0, 0)> +#map1 = affine_map<(d0) -> (d0, 1)> +#map2 = affine_map<(d0) -> (d0, 2)> +#map3 = affine_map<(d0) -> (d0, 3)> +#map4 = affine_map<(d0, d1) -> (-d1 + 3, d0)> +module { + func.func @ws_gemm_4x4(%arg0: memref<4x4xi32>, %arg1: memref<4x4xi32>, %arg2: memref<4x4xi32>) -> memref<4x4xi32> { + %done_writes = taskflow.task @ws_gemm_4x4 will_reads(%arg0, %arg1 : memref<4x4xi32>, memref<4x4xi32>) will_writes(%arg2 : memref<4x4xi32>) [original_read_memrefs(%arg0, %arg1 : memref<4x4xi32>, memref<4x4xi32>), original_write_memrefs(%arg2 : memref<4x4xi32>)] : (memref<4x4xi32>, memref<4x4xi32>, memref<4x4xi32>) -> (memref<4x4xi32>) { + ^bb0(%arg3: memref<4x4xi32>, %arg4: memref<4x4xi32>, %arg5: memref<4x4xi32>): + %0:4 = taskflow.stream_read %arg3 maps [#map, #map1, #map2, #map3] : memref<4x4xi32> -> (i32, i32, i32, i32) + %1:4 = neura.kernel inputs(%0#0, %0#1, %0#2, %0#3, %arg4 : i32, i32, i32, i32, memref<4x4xi32>) attributes {accelerator = "neura", kernel_metadata = {kind = "template", template = {input_ports = [{direction = "west", kernel_input = 0 : i32, x = 0 : i32, y = 3 : i32}, {direction = "west", kernel_input = 1 : i32, x = 0 : i32, y = 2 : i32}, {direction = "west", kernel_input = 2 : i32, x = 0 : i32, y = 1 : i32}, {direction = "west", kernel_input = 3 : i32, x = 0 : i32, y = 0 : i32}], name = "systolic_array", output_ports = [{direction = "south", kernel_result = 0 : i32, x = 0 : i32, y = 0 : i32}, {direction = "south", kernel_result = 1 : i32, x = 1 : i32, y = 0 : i32}, {direction = "south", kernel_result = 2 : i32, x = 2 : i32, y = 0 : i32}, {direction = "south", kernel_result = 3 : i32, x = 3 : i32, y = 0 : i32}], stationary = {kernel_input = 4 : i32, map = #map4, mode = "weight"}}}, 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}} { + ^bb0(%arg6: !neura.data, %arg7: !neura.data, %arg8: !neura.data, %arg9: !neura.data, %arg10: !neura.data, i1>): + %2 = "neura.data_mov"(%arg6) {dfg_id = 0 : i32, mapping_locs = [{direction = "west", id = 20 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, io = "input", resource = "port", time_step = 0 : i32, x = 0 : i32, y = 3 : i32}]} : (!neura.data) -> !neura.data + %3 = "neura.mac"(%2) <{stationary = "weight"}> {dfg_id = 16 : i32, mapping_locs = [{id = 12 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 0 : i32, x = 0 : i32, y = 3 : i32}]} : (!neura.data) -> !neura.data + %4 = "neura.data_mov"(%arg6) {dfg_id = 1 : i32, mapping_locs = [{direction = "west", id = 20 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, io = "input", resource = "port", time_step = 0 : i32, x = 0 : i32, y = 3 : i32}, {id = 38 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 0 : i32}]} : (!neura.data) -> !neura.data + %5 = "neura.mac"(%4) <{stationary = "weight"}> {dfg_id = 17 : i32, mapping_locs = [{id = 13 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 1 : i32, x = 1 : i32, y = 3 : i32}]} : (!neura.data) -> !neura.data + %6 = "neura.data_mov"(%arg6) {dfg_id = 2 : i32, mapping_locs = [{direction = "west", id = 20 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, io = "input", resource = "port", time_step = 0 : i32, x = 0 : i32, y = 3 : i32}, {id = 38 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 0 : i32}, {id = 41 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data) -> !neura.data + %7 = "neura.mac"(%6) <{stationary = "weight"}> {dfg_id = 18 : i32, mapping_locs = [{id = 14 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 2 : i32, x = 2 : i32, y = 3 : i32}]} : (!neura.data) -> !neura.data + %8 = "neura.data_mov"(%arg6) {dfg_id = 3 : i32, mapping_locs = [{direction = "west", id = 20 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, io = "input", resource = "port", time_step = 0 : i32, x = 0 : i32, y = 3 : i32}, {id = 38 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 0 : i32}, {id = 41 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 1 : i32}, {id = 44 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data + %9 = "neura.mac"(%8) <{stationary = "weight"}> {dfg_id = 19 : i32, mapping_locs = [{id = 15 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "tile", time_step = 3 : i32, x = 3 : i32, y = 3 : i32}]} : (!neura.data) -> !neura.data + %10 = "neura.data_mov"(%arg7) {dfg_id = 4 : i32, mapping_locs = [{direction = "west", id = 16 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, io = "input", resource = "port", time_step = 1 : i32, x = 0 : i32, y = 2 : i32}]} : (!neura.data) -> !neura.data + %11 = "neura.data_mov"(%3) {dfg_id = 20 : i32, mapping_locs = [{id = 39 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 0 : i32}]} : (!neura.data) -> !neura.data + %12 = "neura.mac"(%10, %11) <{stationary = "weight"}> {dfg_id = 24 : i32, mapping_locs = [{id = 8 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 1 : i32, x = 0 : i32, y = 2 : i32}]} : (!neura.data, !neura.data) -> !neura.data + %13 = "neura.data_mov"(%arg7) {dfg_id = 5 : i32, mapping_locs = [{direction = "west", id = 16 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, io = "input", resource = "port", time_step = 1 : i32, x = 0 : i32, y = 2 : i32}, {id = 24 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data) -> !neura.data + %14 = "neura.data_mov"(%5) {dfg_id = 21 : i32, mapping_locs = [{id = 42 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data) -> !neura.data + %15 = "neura.mac"(%13, %14) <{stationary = "weight"}> {dfg_id = 25 : i32, mapping_locs = [{id = 9 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 2 : i32, x = 1 : i32, y = 2 : i32}]} : (!neura.data, !neura.data) -> !neura.data + %16 = "neura.data_mov"(%arg7) {dfg_id = 6 : i32, mapping_locs = [{direction = "west", id = 16 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, io = "input", resource = "port", time_step = 1 : i32, x = 0 : i32, y = 2 : i32}, {id = 24 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 1 : i32}, {id = 28 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data + %17 = "neura.data_mov"(%7) {dfg_id = 22 : i32, mapping_locs = [{id = 45 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data + %18 = "neura.mac"(%16, %17) <{stationary = "weight"}> {dfg_id = 26 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "tile", time_step = 3 : i32, x = 2 : i32, y = 2 : i32}]} : (!neura.data, !neura.data) -> !neura.data + %19 = "neura.data_mov"(%arg7) {dfg_id = 7 : i32, mapping_locs = [{direction = "west", id = 16 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, io = "input", resource = "port", time_step = 1 : i32, x = 0 : i32, y = 2 : i32}, {id = 24 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 1 : i32}, {id = 28 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}, {id = 32 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data + %20 = "neura.data_mov"(%9) {dfg_id = 23 : i32, mapping_locs = [{id = 47 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data + %21 = "neura.mac"(%19, %20) <{stationary = "weight"}> {dfg_id = 27 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "tile", time_step = 4 : i32, x = 3 : i32, y = 2 : i32}]} : (!neura.data, !neura.data) -> !neura.data + %22 = "neura.data_mov"(%arg8) {dfg_id = 8 : i32, mapping_locs = [{direction = "west", id = 12 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, io = "input", resource = "port", time_step = 2 : i32, x = 0 : i32, y = 1 : i32}]} : (!neura.data) -> !neura.data + %23 = "neura.data_mov"(%12) {dfg_id = 28 : i32, mapping_locs = [{id = 25 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data) -> !neura.data + %24 = "neura.mac"(%22, %23) <{stationary = "weight"}> {dfg_id = 32 : i32, mapping_locs = [{id = 4 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 2 : i32, x = 0 : i32, y = 1 : i32}]} : (!neura.data, !neura.data) -> !neura.data + %25 = "neura.data_mov"(%arg8) {dfg_id = 9 : i32, mapping_locs = [{direction = "west", id = 12 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, io = "input", resource = "port", time_step = 2 : i32, x = 0 : i32, y = 1 : i32}, {id = 10 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data + %26 = "neura.data_mov"(%15) {dfg_id = 29 : i32, mapping_locs = [{id = 29 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data + %27 = "neura.mac"(%25, %26) <{stationary = "weight"}> {dfg_id = 33 : i32, mapping_locs = [{id = 5 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "tile", time_step = 3 : i32, x = 1 : i32, y = 1 : i32}]} : (!neura.data, !neura.data) -> !neura.data + %28 = "neura.data_mov"(%arg8) {dfg_id = 10 : i32, mapping_locs = [{direction = "west", id = 12 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, io = "input", resource = "port", time_step = 2 : i32, x = 0 : i32, y = 1 : i32}, {id = 10 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}, {id = 14 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data + %29 = "neura.data_mov"(%18) {dfg_id = 30 : i32, mapping_locs = [{id = 33 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data + %30 = "neura.mac"(%28, %29) <{stationary = "weight"}> {dfg_id = 34 : i32, mapping_locs = [{id = 6 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "tile", time_step = 4 : i32, x = 2 : i32, y = 1 : i32}]} : (!neura.data, !neura.data) -> !neura.data + %31 = "neura.data_mov"(%arg8) {dfg_id = 11 : i32, mapping_locs = [{direction = "west", id = 12 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, io = "input", resource = "port", time_step = 2 : i32, x = 0 : i32, y = 1 : i32}, {id = 10 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}, {id = 14 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}, {id = 18 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "link", time_step = 4 : i32}]} : (!neura.data) -> !neura.data + %32 = "neura.data_mov"(%21) {dfg_id = 31 : i32, mapping_locs = [{id = 36 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "link", time_step = 4 : i32}]} : (!neura.data) -> !neura.data + %33 = "neura.mac"(%31, %32) <{stationary = "weight"}> {dfg_id = 35 : i32, mapping_locs = [{id = 7 : i32, index_per_ii = 0 : i32, invalid_iterations = 5 : i32, resource = "tile", time_step = 5 : i32, x = 3 : i32, y = 1 : i32}]} : (!neura.data, !neura.data) -> !neura.data + %34 = "neura.data_mov"(%arg9) {dfg_id = 12 : i32, mapping_locs = [{direction = "west", id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, io = "input", resource = "port", time_step = 3 : i32, x = 0 : i32, y = 0 : i32}]} : (!neura.data) -> !neura.data + %35 = "neura.data_mov"(%24) {dfg_id = 36 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data + %36 = "neura.mac"(%34, %35) <{stationary = "weight"}> {dfg_id = 40 : i32, mapping_locs = [{id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "tile", time_step = 3 : i32, x = 0 : i32, y = 0 : i32}]} : (!neura.data, !neura.data) -> !neura.data + %37 = "neura.data_mov"(%arg9) {dfg_id = 13 : i32, mapping_locs = [{direction = "west", id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, io = "input", resource = "port", time_step = 3 : i32, x = 0 : i32, y = 0 : i32}, {id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data + %38 = "neura.data_mov"(%27) {dfg_id = 37 : i32, mapping_locs = [{id = 15 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data + %39 = "neura.mac"(%37, %38) <{stationary = "weight"}> {dfg_id = 41 : i32, mapping_locs = [{id = 1 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "tile", time_step = 4 : i32, x = 1 : i32, y = 0 : i32}]} : (!neura.data, !neura.data) -> !neura.data + %40 = "neura.data_mov"(%arg9) {dfg_id = 14 : i32, mapping_locs = [{direction = "west", id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, io = "input", resource = "port", time_step = 3 : i32, x = 0 : i32, y = 0 : i32}, {id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}, {id = 3 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "link", time_step = 4 : i32}]} : (!neura.data) -> !neura.data + %41 = "neura.data_mov"(%30) {dfg_id = 38 : i32, mapping_locs = [{id = 19 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "link", time_step = 4 : i32}]} : (!neura.data) -> !neura.data + %42 = "neura.mac"(%40, %41) <{stationary = "weight"}> {dfg_id = 42 : i32, mapping_locs = [{id = 2 : i32, index_per_ii = 0 : i32, invalid_iterations = 5 : i32, resource = "tile", time_step = 5 : i32, x = 2 : i32, y = 0 : i32}]} : (!neura.data, !neura.data) -> !neura.data + %43 = "neura.data_mov"(%arg9) {dfg_id = 15 : i32, mapping_locs = [{direction = "west", id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, io = "input", resource = "port", time_step = 3 : i32, x = 0 : i32, y = 0 : i32}, {id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}, {id = 3 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "link", time_step = 4 : i32}, {id = 6 : i32, index_per_ii = 0 : i32, invalid_iterations = 5 : i32, resource = "link", time_step = 5 : i32}]} : (!neura.data) -> !neura.data + %44 = "neura.data_mov"(%33) {dfg_id = 39 : i32, mapping_locs = [{id = 22 : i32, index_per_ii = 0 : i32, invalid_iterations = 5 : i32, resource = "link", time_step = 5 : i32}]} : (!neura.data) -> !neura.data + %45 = "neura.mac"(%43, %44) <{stationary = "weight"}> {dfg_id = 43 : i32, mapping_locs = [{id = 3 : i32, index_per_ii = 0 : i32, invalid_iterations = 6 : i32, resource = "tile", time_step = 6 : i32, x = 3 : i32, y = 0 : i32}]} : (!neura.data, !neura.data) -> !neura.data + %46 = "neura.data_mov"(%36) {dfg_id = 44 : i32, mapping_locs = [{direction = "south", id = 3 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, io = "output", resource = "port", time_step = 3 : i32, x = 0 : i32, y = 0 : i32}]} : (!neura.data) -> !neura.data + %47 = "neura.data_mov"(%39) {dfg_id = 45 : i32, mapping_locs = [{direction = "south", id = 5 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, io = "output", resource = "port", time_step = 4 : i32, x = 1 : i32, y = 0 : i32}]} : (!neura.data) -> !neura.data + %48 = "neura.data_mov"(%42) {dfg_id = 46 : i32, mapping_locs = [{direction = "south", id = 7 : i32, index_per_ii = 0 : i32, invalid_iterations = 5 : i32, io = "output", resource = "port", time_step = 5 : i32, x = 2 : i32, y = 0 : i32}]} : (!neura.data) -> !neura.data + %49 = "neura.data_mov"(%45) {dfg_id = 47 : i32, mapping_locs = [{direction = "south", id = 11 : i32, index_per_ii = 0 : i32, invalid_iterations = 6 : i32, io = "output", resource = "port", time_step = 6 : i32, x = 3 : i32, y = 0 : i32}]} : (!neura.data) -> !neura.data + neura.yield results(%46, %47, %48, %49 : !neura.data, !neura.data, !neura.data, !neura.data) {dfg_id = 48 : i32} + } : i32, i32, i32, i32 + taskflow.stream_write(%1#0, %1#1, %1#2, %1#3 : i32, i32, i32, i32) to %arg5 maps [#map, #map1, #map2, #map3] : memref<4x4xi32> + taskflow.yield done_writes(%arg5 : memref<4x4xi32>) + } + return %done_writes : memref<4x4xi32> + } +} +""".strip() + + def ws_gemm_4x4(A: synl.Tensor, B: synl.Tensor, C: synl.Tensor): array = synl.TileArray(x_tiles=4, y_tiles=4) @@ -86,3 +159,17 @@ def test_lowers_systolic_gemm_to_exact_pre_mapping_ir(): ) assert actual.strip() == PRE_MAPPING_IR + + +def test_compiles_systolic_gemm_to_exact_mapped_ir(): + actual = synapse.compile( + ws_gemm_4x4, + target="neura", + argument_types=( + synl.i32[4, 4], + synl.i32[4, 4], + synl.i32[4, 4], + ), + ) + + assert actual.strip() == MAPPED_IR From ddc61a99e750f492647908fc70583b39663f20fc Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Wed, 2 Sep 2026 20:28:23 +0800 Subject: [PATCH 12/19] Extract systolic GEMM into TileArray templates --- python/synapse/templates/__init__.py | 1 + .../synapse/templates/tile_array/__init__.py | 7 +++ python/synapse/templates/tile_array/gemm.py | 45 +++++++++++++++++++ tests/python/compiler/test_systolic_gemm.py | 35 +-------------- 4 files changed, 54 insertions(+), 34 deletions(-) create mode 100644 python/synapse/templates/__init__.py create mode 100644 python/synapse/templates/tile_array/__init__.py create mode 100644 python/synapse/templates/tile_array/gemm.py diff --git a/python/synapse/templates/__init__.py b/python/synapse/templates/__init__.py new file mode 100644 index 0000000..3affb48 --- /dev/null +++ b/python/synapse/templates/__init__.py @@ -0,0 +1 @@ +"""Reusable spatial templates provided by Synapse.""" diff --git a/python/synapse/templates/tile_array/__init__.py b/python/synapse/templates/tile_array/__init__.py new file mode 100644 index 0000000..66eb077 --- /dev/null +++ b/python/synapse/templates/tile_array/__init__.py @@ -0,0 +1,7 @@ +"""Templates for computation within one TileArray.""" + +from .gemm import ws_gemm_4x4 + +__all__ = [ + "ws_gemm_4x4", +] diff --git a/python/synapse/templates/tile_array/gemm.py b/python/synapse/templates/tile_array/gemm.py new file mode 100644 index 0000000..b174b79 --- /dev/null +++ b/python/synapse/templates/tile_array/gemm.py @@ -0,0 +1,45 @@ +"""Reusable GEMM implementations authored with the Synapse language.""" + +import synapse.language as synl + + +def ws_gemm_4x4( + A: synl.Tensor, + B: synl.Tensor, + C: synl.Tensor, +): + """Describe a fixed 4x4 weight-stationary GEMM.""" + + array = synl.TileArray(x_tiles=4, y_tiles=4) + + partial_sums = [] + + for k in range(array.y_tiles): + y = array.y_tiles - 1 - k + + # Stream one column of A through the corresponding west Port. + activation = synl.input_port( + A[:, k], + port=array.west_ports[y], + ) + + # Each Tile keeps one value from B stationary while partial sums + # propagate toward the south boundary. + partial_sums = [ + synl.mac( + activation, + partial_sums[x] if partial_sums else None, + stationary=B[k, x], + mode=synl.StationaryMode.WEIGHT, + tile=array[x, y], + ) + for x in range(array.x_tiles) + ] + + for x, result in enumerate(partial_sums): + # Each south Port writes one result column into C. + synl.output_port( + result, + target=C[:, x], + port=array.south_ports[x], + ) diff --git a/tests/python/compiler/test_systolic_gemm.py b/tests/python/compiler/test_systolic_gemm.py index 5c709f5..df2456d 100644 --- a/tests/python/compiler/test_systolic_gemm.py +++ b/tests/python/compiler/test_systolic_gemm.py @@ -1,6 +1,7 @@ import synapse import synapse.language as synl from synapse.frontend import lowering +from synapse.templates.tile_array import ws_gemm_4x4 PRE_MAPPING_IR = """ #map = affine_map<(d0) -> (d0, 0)> @@ -114,40 +115,6 @@ """.strip() -def ws_gemm_4x4(A: synl.Tensor, B: synl.Tensor, C: synl.Tensor): - array = synl.TileArray(x_tiles=4, y_tiles=4) - - partial_sums = [] - - for k in range(array.y_tiles): - y = array.y_tiles - 1 - k - - # A[:, k] is streamed through one west boundary Port. - activation = synl.input_port( - A[:, k], - port=array.west_ports[y], - ) - - partial_sums = [ - synl.mac( - activation, - partial_sums[x] if partial_sums else None, - stationary=B[k, x], - mode=synl.StationaryMode.WEIGHT, - tile=array[x, y], - ) - for x in range(array.x_tiles) - ] - - for x, result in enumerate(partial_sums): - # Each south Port writes one column of C. - synl.output_port( - result, - target=C[:, x], - port=array.south_ports[x], - ) - - def test_lowers_systolic_gemm_to_exact_pre_mapping_ir(): actual = lowering.lower( ws_gemm_4x4, From 074ce1c5eea52f03cc668446efba67fc1b1f8ede Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Fri, 4 Sep 2026 18:11:55 +0800 Subject: [PATCH 13/19] Align TileArray MAC lowering with Neura forwarding --- mlir/amoeba | 2 +- python/synapse/frontend/lowering.py | 49 +++--- python/synapse/language/__init__.py | 2 - python/synapse/language/tile_array_program.py | 148 ++++++++---------- python/synapse/templates/tile_array/gemm.py | 27 ++-- tests/python/compiler/test_systolic_gemm.py | 136 ++++++++-------- 6 files changed, 180 insertions(+), 184 deletions(-) diff --git a/mlir/amoeba b/mlir/amoeba index b514f8d..60a8bed 160000 --- a/mlir/amoeba +++ b/mlir/amoeba @@ -1 +1 @@ -Subproject commit b514f8d0202d5f6ed6efabc1849d3f891642c805 +Subproject commit 60a8bed1b657b4a0c036e3df018d5bc945075b41 diff --git a/python/synapse/frontend/lowering.py b/python/synapse/frontend/lowering.py index a04c933..0f5e7e7 100644 --- a/python/synapse/frontend/lowering.py +++ b/python/synapse/frontend/lowering.py @@ -11,7 +11,6 @@ AddOp, ConstantOp, MacOp, - StationaryMode, TileArrayBuilder, TileArrayOp, TileArrayProgram, @@ -130,11 +129,6 @@ def get_stationary_map() -> AffineMapAttr: if program.stationary is None: raise RuntimeError("template program requires stationary data") - if program.stationary.mode != StationaryMode.WEIGHT: - raise NotImplementedError( - "only weight-stationary template lowering is implemented initially" - ) - for tile, access in program.stationary.tile_values: expected_indices = (program.array.y_tiles - 1 - tile.y, tile.x) @@ -157,15 +151,16 @@ def get_constant_attribute( result_type, ): """Build the typed MLIR attribute for a constant value.""" + result = operation.results[0] - if operation.result.dtype == DType.I32: + if result.dtype == DType.I32: return IntegerAttr.get(result_type, cast(int, operation.value)) - if operation.result.dtype == DType.F32: + if result.dtype == DType.F32: return FloatAttr.get(result_type, float(operation.value)) raise NotImplementedError( - f"unsupported constant type: {operation.result.dtype.value}" + f"unsupported constant type: {result.dtype.value}" ) def get_placement(tile: Tile) -> DictAttr: @@ -219,7 +214,6 @@ def get_kernel_metadata() -> DictAttr | None: { "kernel_input": IntegerAttr.get(i32_type, len(program.input_ports)), "map": get_stationary_map(), - "mode": StringAttr.get(program.stationary.mode.value), } ) @@ -240,7 +234,7 @@ def get_kernel_metadata() -> DictAttr | None: ) @singledispatch - def lower_operation(operation: TileArrayOp, operands, result_type): + def lower_operation(operation: TileArrayOp, operands, result_types): """Lower one frontend TileArray operation to a Neura operation. The caller handles common lowering such as resolving operands, @@ -251,31 +245,32 @@ def lower_operation(operation: TileArrayOp, operands, result_type): ) @lower_operation.register - def lower_constant(operation: ConstantOp, operands, result_type): + def lower_constant(operation: ConstantOp, operands, result_types): """Lower a ConstantOp to neura.constant.""" return neura.ConstantOp( - result_type, get_constant_attribute(operation, result_type) + result_types[0], get_constant_attribute(operation, result_types[0]) ) @lower_operation.register - def lower_add(operation: AddOp, operands, result_type): + def lower_add(operation: AddOp, operands, result_types): """Lower a frontend AddOp to neura.add.""" lhs, rhs = operands - return neura.AddOp(result_type, lhs, rhs=rhs) + return neura.AddOp(result_types[0], lhs, rhs=rhs) @lower_operation.register - def lower_mac(operation: MacOp, operands, result_type): + def lower_mac(operation: MacOp, operands, result_types): """Lower a configured MacOp to neura.mac.""" input0 = operands[0] input1 = operands[1] if len(operands) == 2 else None + accumulated_type, forwarded_type = result_types return neura.MacOp( - result_type, + accumulated_type, + forwarded_type, input0, - StringAttr.get(operation.stationary_mode.value), input1=input1, ) @@ -420,19 +415,31 @@ def lower_mac(operation: MacOp, operands, result_type): with InsertionPoint(kernel_block): for operation in program.operations: - result_type = get_mlir_type(operation.result.dtype) + result_types = tuple( + get_mlir_type(result.dtype) for result in operation.results + ) mlir_operands = tuple( values_by_id[operand.id] for operand in operation.operands ) - mlir_operation = lower_operation(operation, mlir_operands, result_type) + mlir_operation = lower_operation(operation, mlir_operands, result_types) mlir_operation.operation.attributes["placement"] = get_placement( operation.tile ) - values_by_id[operation.result.id] = mlir_operation.result + mlir_results = tuple(mlir_operation.results) + + if len(mlir_results) != len(operation.results): + raise RuntimeError( + "frontend and MLIR operation result counts differ" + ) + + for frontend_result, mlir_result in zip( + operation.results, mlir_results + ): + values_by_id[frontend_result.id] = mlir_result neura.YieldOp( iter_args_next=[], diff --git a/python/synapse/language/__init__.py b/python/synapse/language/__init__.py index 64b7b97..4e97894 100644 --- a/python/synapse/language/__init__.py +++ b/python/synapse/language/__init__.py @@ -3,7 +3,6 @@ from .spatial import TileArray from .tensor import Tensor from .tile_array_program import ( - StationaryMode, add, constant, input_port, @@ -14,7 +13,6 @@ __all__ = [ "DType", - "StationaryMode", "Tensor", "TensorType", "TileArray", diff --git a/python/synapse/language/tile_array_program.py b/python/synapse/language/tile_array_program.py index bcdb828..ba9cffa 100644 --- a/python/synapse/language/tile_array_program.py +++ b/python/synapse/language/tile_array_program.py @@ -10,21 +10,12 @@ from collections.abc import Callable from contextvars import ContextVar from dataclasses import dataclass, field -from enum import Enum from .spatial import Port, Tile, TileArray from .tensor import Tensor, TensorAccess from .types import DType -class StationaryMode(Enum): - """The dataflow operand retained locally by configured MACs.""" - - WEIGHT = "weight" - INPUT = "input" - OUTPUT = "output" - - @dataclass(frozen=True) class TileArrayValue: """A typed value produced by one tile-array operation.""" @@ -63,7 +54,7 @@ class TileArrayOp: empty tuple, while operations such as add consume input values. """ - result: TileArrayValue + results: tuple[TileArrayValue, ...] operands: tuple[TileArrayValue, ...] tile: Tile @@ -75,13 +66,17 @@ class ConstantOp(TileArrayOp): value: int | float def __post_init__(self) -> None: - """Validate the operation-specific operands and scalar value.""" + """Validates the operation-specific operands and scalar value.""" + if len(self.results) != 1: + raise ValueError("ConstantOp requires exactly one result") + if self.operands: raise ValueError( f"ConstantOp requires zero operands, but got {len(self.operands)}" ) - dtype = self.result.dtype + result = self.results[0] + dtype = result.dtype if not isinstance(dtype, DType): raise TypeError("ConstantOp result must use a DType") @@ -100,18 +95,22 @@ 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.""" + """Validates the operation-specific arity and scalar types.""" + if len(self.results) != 1: + raise ValueError("AddOp requires exactly one result") + if len(self.operands) != 2: raise ValueError( f"AddOp requires two operands, but got {len(self.operands)}" ) + result = self.results[0] 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: + if result.dtype != lhs.dtype: raise TypeError("AddOp result type must match its operand type") @@ -120,36 +119,40 @@ class MacOp(TileArrayOp): """A configured MAC using one stationary scalar value.""" stationary_value: TensorAccess - stationary_mode: StationaryMode - def __post_init__(self) -> None: - """Validates configured MAC operands and stationary data.""" + @property + def accumulated(self) -> TileArrayValue: + """Returns the accumulated output value.""" - if self.stationary_mode in ( - StationaryMode.WEIGHT, - StationaryMode.INPUT, - ) and len(self.operands) not in (1, 2): - raise ValueError( - "weight/input-stationary MacOp requires " - "an input and optional partial sum" - ) + return self.results[0] + + @property + def forwarded(self) -> TileArrayValue: + """Returns the value forwarded from input0.""" - if self.stationary_mode == StationaryMode.OUTPUT and len(self.operands) != 2: - raise ValueError("output-stationary MacOp requires two multiplicands") + return self.results[1] + def __post_init__(self) -> None: + """Validates configured MAC operands and stationary data.""" + + if len(self.results) != 2: + raise ValueError("MacOp requires accumulated and forwarded results") + if len(self.operands) not in (1, 2): + raise ValueError("MacOp requires a flowing input and optional partial sum") if not self.stationary_value.is_scalar: raise ValueError("stationary data must identify one scalar") - if any(operand.dtype != self.result.dtype for operand in self.operands): + if any(operand.dtype != self.accumulated.dtype for operand in self.operands): raise TypeError("MacOp operands and result must have the same dtype") - if self.stationary_value.dtype != self.result.dtype: + if self.stationary_value.dtype != self.accumulated.dtype: raise TypeError("MacOp stationary data and result must have the same dtype") + if self.forwarded.dtype != self.accumulated.dtype: + raise TypeError("MacOp results must have the same dtype") @dataclass(frozen=True) class StationaryBinding: """Stationary data assigned to the tiles of one template program.""" - mode: StationaryMode source: Tensor tile_values: tuple[tuple[Tile, TensorAccess], ...] @@ -235,38 +238,44 @@ def emit( self, *, operands: tuple[TileArrayValue, ...], - result_dtype: DType, + result_dtypes: tuple[DType, ...], tile: Tile, - create_operation: Callable[[TileArrayValue], TileArrayOp], - ) -> TileArrayValue: + create_operation: Callable[[tuple[TileArrayValue, ...]], TileArrayOp], + ) -> tuple[TileArrayValue, ...]: """Creates and record one tile-array operation.""" self._ensure_not_built() self._bind_tile_array(tile.array) + if not result_dtypes: + raise ValueError("an operation must produce at least one result") + for operand in operands: self._validate_operand_for_builder(operand) - # Commit the value ID only after operation construction and validation + # Commits 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, + results = tuple( + TileArrayValue(id=self._next_value_id + index, dtype=dtype, _builder=self) + for index, dtype in enumerate(result_dtypes) ) - operation = create_operation(result) + + operation = create_operation(results) 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 len(operation.results) != len(results) or any( + actual is not expected + for actual, expected in zip(operation.results, results) + ): + raise ValueError("create_operation must use the provided result values") 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 + self._next_value_id += len(results) + return results def add_input_port(self, *, source: TensorAccess, port: Port) -> TileArrayValue: """Creates one scalar kernel input and bind it to a Port.""" @@ -328,23 +337,14 @@ def build(self) -> TileArrayProgram: template_name = None if mac_operations: - stationary_mode = mac_operations[0].stationary_mode stationary_source = mac_operations[0].stationary_value.source - if any( - operation.stationary_mode != stationary_mode - for operation in mac_operations - ): - raise ValueError( - "all configured MACs in one template must use one stationary mode" - ) if any( operation.stationary_value.source is not stationary_source for operation in mac_operations ): raise ValueError("all configured MACs must use one stationary source") stationary = StationaryBinding( - mode=stationary_mode, source=stationary_source, tile_values=tuple( (operation.tile, operation.stationary_value) @@ -413,15 +413,15 @@ def constant( return builder.emit( operands=(), - result_dtype=dtype, + result_dtypes=(dtype,), tile=tile, - create_operation=lambda result: ConstantOp( - result=result, + create_operation=lambda results: ConstantOp( + results=results, operands=(), tile=tile, value=value, ), - ) + )[0] def add( @@ -437,14 +437,14 @@ def add( return builder.emit( operands=operands, - result_dtype=lhs.dtype, + result_dtypes=(lhs.dtype,), tile=tile, - create_operation=lambda result: AddOp( - result=result, + create_operation=lambda results: AddOp( + results=results, operands=operands, tile=tile, ), - ) + )[0] def input_port(source: TensorAccess, *, port: Port) -> TileArrayValue: @@ -470,40 +470,30 @@ def output_port(value: TileArrayValue, *, target: TensorAccess, port: Port) -> N def mac( input0: TileArrayValue, - input1: TileArrayValue | None, + input1: TileArrayValue | None = None, *, stationary: TensorAccess, - mode: StationaryMode, tile: Tile, -) -> TileArrayValue: - """Create one configured MAC operation. - - Operand meanings depend on the selected stationary mode: - - - WEIGHT: input0 is activation, input1 is an optional partial sum. - - INPUT: input0 is weight, input1 is an optional partial sum. - - OUTPUT: input0 and input1 are the two multiplicands. - """ +) -> tuple[TileArrayValue, TileArrayValue]: + """Creates one configured MAC operation.""" if not isinstance(stationary, TensorAccess): raise TypeError("mac stationary data must be a tensor access") - if not isinstance(mode, StationaryMode): - raise TypeError("mac mode must be a StationaryMode") - operands = (input0,) if input1 is None else (input0, input1) builder = _require_active_builder() - return builder.emit( + accumulated, forwarded = builder.emit( operands=operands, - result_dtype=input0.dtype, + result_dtypes=(input0.dtype, input0.dtype), tile=tile, - create_operation=lambda result: MacOp( - result=result, + create_operation=lambda results: MacOp( + results=results, operands=operands, tile=tile, stationary_value=stationary, - stationary_mode=mode, ), ) + + return accumulated, forwarded diff --git a/python/synapse/templates/tile_array/gemm.py b/python/synapse/templates/tile_array/gemm.py index b174b79..2520463 100644 --- a/python/synapse/templates/tile_array/gemm.py +++ b/python/synapse/templates/tile_array/gemm.py @@ -8,7 +8,7 @@ def ws_gemm_4x4( B: synl.Tensor, C: synl.Tensor, ): - """Describe a fixed 4x4 weight-stationary GEMM.""" + """Describes a fixed 4x4 weight-stationary GEMM.""" array = synl.TileArray(x_tiles=4, y_tiles=4) @@ -17,29 +17,30 @@ def ws_gemm_4x4( for k in range(array.y_tiles): y = array.y_tiles - 1 - k - # Stream one column of A through the corresponding west Port. - activation = synl.input_port( + # A[:, k] enters from the west boundary ports and flows east across this row. + flowing = synl.input_port( A[:, k], port=array.west_ports[y], ) - # Each Tile keeps one value from B stationary while partial sums - # propagate toward the south boundary. - partial_sums = [ - synl.mac( - activation, + next_partial_sums = [] + + for x in range(array.x_tiles): + accumulated, flowing = synl.mac( + flowing, partial_sums[x] if partial_sums else None, stationary=B[k, x], - mode=synl.StationaryMode.WEIGHT, tile=array[x, y], ) - for x in range(array.x_tiles) - ] - for x, result in enumerate(partial_sums): + next_partial_sums.append(accumulated) + + partial_sums = next_partial_sums + + for x, accumulated in enumerate(partial_sums): # Each south Port writes one result column into C. synl.output_port( - result, + accumulated, target=C[:, x], port=array.south_ports[x], ) diff --git a/tests/python/compiler/test_systolic_gemm.py b/tests/python/compiler/test_systolic_gemm.py index df2456d..b6da5f9 100644 --- a/tests/python/compiler/test_systolic_gemm.py +++ b/tests/python/compiler/test_systolic_gemm.py @@ -14,25 +14,25 @@ %done_writes = taskflow.task @ws_gemm_4x4 will_reads(%arg0, %arg1 : memref<4x4xi32>, memref<4x4xi32>) will_writes(%arg2 : memref<4x4xi32>) [original_read_memrefs(%arg0, %arg1 : memref<4x4xi32>, memref<4x4xi32>), original_write_memrefs(%arg2 : memref<4x4xi32>)] : (memref<4x4xi32>, memref<4x4xi32>, memref<4x4xi32>) -> (memref<4x4xi32>) { ^bb0(%arg3: memref<4x4xi32>, %arg4: memref<4x4xi32>, %arg5: memref<4x4xi32>): %0:4 = taskflow.stream_read %arg3 maps [#map, #map1, #map2, #map3] : memref<4x4xi32> -> (i32, i32, i32, i32) - %1:4 = neura.kernel inputs(%0#0, %0#1, %0#2, %0#3, %arg4 : i32, i32, i32, i32, memref<4x4xi32>) attributes {accelerator = "neura", kernel_metadata = {kind = "template", template = {input_ports = [{direction = "west", kernel_input = 0 : i32, x = 0 : i32, y = 3 : i32}, {direction = "west", kernel_input = 1 : i32, x = 0 : i32, y = 2 : i32}, {direction = "west", kernel_input = 2 : i32, x = 0 : i32, y = 1 : i32}, {direction = "west", kernel_input = 3 : i32, x = 0 : i32, y = 0 : i32}], name = "systolic_array", output_ports = [{direction = "south", kernel_result = 0 : i32, x = 0 : i32, y = 0 : i32}, {direction = "south", kernel_result = 1 : i32, x = 1 : i32, y = 0 : i32}, {direction = "south", kernel_result = 2 : i32, x = 2 : i32, y = 0 : i32}, {direction = "south", kernel_result = 3 : i32, x = 3 : i32, y = 0 : i32}], stationary = {kernel_input = 4 : i32, map = #map4, mode = "weight"}}}} { + %1:4 = neura.kernel inputs(%0#0, %0#1, %0#2, %0#3, %arg4 : i32, i32, i32, i32, memref<4x4xi32>) attributes {accelerator = "neura", kernel_metadata = {kind = "template", template = {input_ports = [{direction = "west", kernel_input = 0 : i32, x = 0 : i32, y = 3 : i32}, {direction = "west", kernel_input = 1 : i32, x = 0 : i32, y = 2 : i32}, {direction = "west", kernel_input = 2 : i32, x = 0 : i32, y = 1 : i32}, {direction = "west", kernel_input = 3 : i32, x = 0 : i32, y = 0 : i32}], name = "systolic_array", output_ports = [{direction = "south", kernel_result = 0 : i32, x = 0 : i32, y = 0 : i32}, {direction = "south", kernel_result = 1 : i32, x = 1 : i32, y = 0 : i32}, {direction = "south", kernel_result = 2 : i32, x = 2 : i32, y = 0 : i32}, {direction = "south", kernel_result = 3 : i32, x = 3 : i32, y = 0 : i32}], stationary = {kernel_input = 4 : i32, map = #map4}}}} { ^bb0(%arg6: i32, %arg7: i32, %arg8: i32, %arg9: i32, %arg10: memref<4x4xi32>): - %2 = "neura.mac"(%arg6) <{stationary = "weight"}> {placement = {x = 0 : i32, y = 3 : i32}} : (i32) -> i32 - %3 = "neura.mac"(%arg6) <{stationary = "weight"}> {placement = {x = 1 : i32, y = 3 : i32}} : (i32) -> i32 - %4 = "neura.mac"(%arg6) <{stationary = "weight"}> {placement = {x = 2 : i32, y = 3 : i32}} : (i32) -> i32 - %5 = "neura.mac"(%arg6) <{stationary = "weight"}> {placement = {x = 3 : i32, y = 3 : i32}} : (i32) -> i32 - %6 = "neura.mac"(%arg7, %2) <{stationary = "weight"}> {placement = {x = 0 : i32, y = 2 : i32}} : (i32, i32) -> i32 - %7 = "neura.mac"(%arg7, %3) <{stationary = "weight"}> {placement = {x = 1 : i32, y = 2 : i32}} : (i32, i32) -> i32 - %8 = "neura.mac"(%arg7, %4) <{stationary = "weight"}> {placement = {x = 2 : i32, y = 2 : i32}} : (i32, i32) -> i32 - %9 = "neura.mac"(%arg7, %5) <{stationary = "weight"}> {placement = {x = 3 : i32, y = 2 : i32}} : (i32, i32) -> i32 - %10 = "neura.mac"(%arg8, %6) <{stationary = "weight"}> {placement = {x = 0 : i32, y = 1 : i32}} : (i32, i32) -> i32 - %11 = "neura.mac"(%arg8, %7) <{stationary = "weight"}> {placement = {x = 1 : i32, y = 1 : i32}} : (i32, i32) -> i32 - %12 = "neura.mac"(%arg8, %8) <{stationary = "weight"}> {placement = {x = 2 : i32, y = 1 : i32}} : (i32, i32) -> i32 - %13 = "neura.mac"(%arg8, %9) <{stationary = "weight"}> {placement = {x = 3 : i32, y = 1 : i32}} : (i32, i32) -> i32 - %14 = "neura.mac"(%arg9, %10) <{stationary = "weight"}> {placement = {x = 0 : i32, y = 0 : i32}} : (i32, i32) -> i32 - %15 = "neura.mac"(%arg9, %11) <{stationary = "weight"}> {placement = {x = 1 : i32, y = 0 : i32}} : (i32, i32) -> i32 - %16 = "neura.mac"(%arg9, %12) <{stationary = "weight"}> {placement = {x = 2 : i32, y = 0 : i32}} : (i32, i32) -> i32 - %17 = "neura.mac"(%arg9, %13) <{stationary = "weight"}> {placement = {x = 3 : i32, y = 0 : i32}} : (i32, i32) -> i32 - neura.yield results(%14, %15, %16, %17 : i32, i32, i32, i32) + %result, %forwarded = "neura.mac"(%arg6) {placement = {x = 0 : i32, y = 3 : i32}} : (i32) -> (i32, i32) + %result_0, %forwarded_1 = "neura.mac"(%forwarded) {placement = {x = 1 : i32, y = 3 : i32}} : (i32) -> (i32, i32) + %result_2, %forwarded_3 = "neura.mac"(%forwarded_1) {placement = {x = 2 : i32, y = 3 : i32}} : (i32) -> (i32, i32) + %result_4, %forwarded_5 = "neura.mac"(%forwarded_3) {placement = {x = 3 : i32, y = 3 : i32}} : (i32) -> (i32, i32) + %result_6, %forwarded_7 = "neura.mac"(%arg7, %result) {placement = {x = 0 : i32, y = 2 : i32}} : (i32, i32) -> (i32, i32) + %result_8, %forwarded_9 = "neura.mac"(%forwarded_7, %result_0) {placement = {x = 1 : i32, y = 2 : i32}} : (i32, i32) -> (i32, i32) + %result_10, %forwarded_11 = "neura.mac"(%forwarded_9, %result_2) {placement = {x = 2 : i32, y = 2 : i32}} : (i32, i32) -> (i32, i32) + %result_12, %forwarded_13 = "neura.mac"(%forwarded_11, %result_4) {placement = {x = 3 : i32, y = 2 : i32}} : (i32, i32) -> (i32, i32) + %result_14, %forwarded_15 = "neura.mac"(%arg8, %result_6) {placement = {x = 0 : i32, y = 1 : i32}} : (i32, i32) -> (i32, i32) + %result_16, %forwarded_17 = "neura.mac"(%forwarded_15, %result_8) {placement = {x = 1 : i32, y = 1 : i32}} : (i32, i32) -> (i32, i32) + %result_18, %forwarded_19 = "neura.mac"(%forwarded_17, %result_10) {placement = {x = 2 : i32, y = 1 : i32}} : (i32, i32) -> (i32, i32) + %result_20, %forwarded_21 = "neura.mac"(%forwarded_19, %result_12) {placement = {x = 3 : i32, y = 1 : i32}} : (i32, i32) -> (i32, i32) + %result_22, %forwarded_23 = "neura.mac"(%arg9, %result_14) {placement = {x = 0 : i32, y = 0 : i32}} : (i32, i32) -> (i32, i32) + %result_24, %forwarded_25 = "neura.mac"(%forwarded_23, %result_16) {placement = {x = 1 : i32, y = 0 : i32}} : (i32, i32) -> (i32, i32) + %result_26, %forwarded_27 = "neura.mac"(%forwarded_25, %result_18) {placement = {x = 2 : i32, y = 0 : i32}} : (i32, i32) -> (i32, i32) + %result_28, %forwarded_29 = "neura.mac"(%forwarded_27, %result_20) {placement = {x = 3 : i32, y = 0 : i32}} : (i32, i32) -> (i32, i32) + neura.yield results(%result_22, %result_24, %result_26, %result_28 : i32, i32, i32, i32) } : i32, i32, i32, i32 taskflow.stream_write(%1#0, %1#1, %1#2, %1#3 : i32, i32, i32, i32) to %arg5 maps [#map, #map1, #map2, #map3] : memref<4x4xi32> taskflow.yield done_writes(%arg5 : memref<4x4xi32>) @@ -54,57 +54,57 @@ %done_writes = taskflow.task @ws_gemm_4x4 will_reads(%arg0, %arg1 : memref<4x4xi32>, memref<4x4xi32>) will_writes(%arg2 : memref<4x4xi32>) [original_read_memrefs(%arg0, %arg1 : memref<4x4xi32>, memref<4x4xi32>), original_write_memrefs(%arg2 : memref<4x4xi32>)] : (memref<4x4xi32>, memref<4x4xi32>, memref<4x4xi32>) -> (memref<4x4xi32>) { ^bb0(%arg3: memref<4x4xi32>, %arg4: memref<4x4xi32>, %arg5: memref<4x4xi32>): %0:4 = taskflow.stream_read %arg3 maps [#map, #map1, #map2, #map3] : memref<4x4xi32> -> (i32, i32, i32, i32) - %1:4 = neura.kernel inputs(%0#0, %0#1, %0#2, %0#3, %arg4 : i32, i32, i32, i32, memref<4x4xi32>) attributes {accelerator = "neura", kernel_metadata = {kind = "template", template = {input_ports = [{direction = "west", kernel_input = 0 : i32, x = 0 : i32, y = 3 : i32}, {direction = "west", kernel_input = 1 : i32, x = 0 : i32, y = 2 : i32}, {direction = "west", kernel_input = 2 : i32, x = 0 : i32, y = 1 : i32}, {direction = "west", kernel_input = 3 : i32, x = 0 : i32, y = 0 : i32}], name = "systolic_array", output_ports = [{direction = "south", kernel_result = 0 : i32, x = 0 : i32, y = 0 : i32}, {direction = "south", kernel_result = 1 : i32, x = 1 : i32, y = 0 : i32}, {direction = "south", kernel_result = 2 : i32, x = 2 : i32, y = 0 : i32}, {direction = "south", kernel_result = 3 : i32, x = 3 : i32, y = 0 : i32}], stationary = {kernel_input = 4 : i32, map = #map4, mode = "weight"}}}, 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}} { + %1:4 = neura.kernel inputs(%0#0, %0#1, %0#2, %0#3, %arg4 : i32, i32, i32, i32, memref<4x4xi32>) attributes {accelerator = "neura", kernel_metadata = {kind = "template", template = {input_ports = [{direction = "west", kernel_input = 0 : i32, x = 0 : i32, y = 3 : i32}, {direction = "west", kernel_input = 1 : i32, x = 0 : i32, y = 2 : i32}, {direction = "west", kernel_input = 2 : i32, x = 0 : i32, y = 1 : i32}, {direction = "west", kernel_input = 3 : i32, x = 0 : i32, y = 0 : i32}], name = "systolic_array", output_ports = [{direction = "south", kernel_result = 0 : i32, x = 0 : i32, y = 0 : i32}, {direction = "south", kernel_result = 1 : i32, x = 1 : i32, y = 0 : i32}, {direction = "south", kernel_result = 2 : i32, x = 2 : i32, y = 0 : i32}, {direction = "south", kernel_result = 3 : i32, x = 3 : i32, y = 0 : i32}], stationary = {kernel_input = 4 : i32, map = #map4}}}, 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}} { ^bb0(%arg6: !neura.data, %arg7: !neura.data, %arg8: !neura.data, %arg9: !neura.data, %arg10: !neura.data, i1>): - %2 = "neura.data_mov"(%arg6) {dfg_id = 0 : i32, mapping_locs = [{direction = "west", id = 20 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, io = "input", resource = "port", time_step = 0 : i32, x = 0 : i32, y = 3 : i32}]} : (!neura.data) -> !neura.data - %3 = "neura.mac"(%2) <{stationary = "weight"}> {dfg_id = 16 : i32, mapping_locs = [{id = 12 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 0 : i32, x = 0 : i32, y = 3 : i32}]} : (!neura.data) -> !neura.data - %4 = "neura.data_mov"(%arg6) {dfg_id = 1 : i32, mapping_locs = [{direction = "west", id = 20 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, io = "input", resource = "port", time_step = 0 : i32, x = 0 : i32, y = 3 : i32}, {id = 38 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 0 : i32}]} : (!neura.data) -> !neura.data - %5 = "neura.mac"(%4) <{stationary = "weight"}> {dfg_id = 17 : i32, mapping_locs = [{id = 13 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 1 : i32, x = 1 : i32, y = 3 : i32}]} : (!neura.data) -> !neura.data - %6 = "neura.data_mov"(%arg6) {dfg_id = 2 : i32, mapping_locs = [{direction = "west", id = 20 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, io = "input", resource = "port", time_step = 0 : i32, x = 0 : i32, y = 3 : i32}, {id = 38 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 0 : i32}, {id = 41 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data) -> !neura.data - %7 = "neura.mac"(%6) <{stationary = "weight"}> {dfg_id = 18 : i32, mapping_locs = [{id = 14 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 2 : i32, x = 2 : i32, y = 3 : i32}]} : (!neura.data) -> !neura.data - %8 = "neura.data_mov"(%arg6) {dfg_id = 3 : i32, mapping_locs = [{direction = "west", id = 20 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, io = "input", resource = "port", time_step = 0 : i32, x = 0 : i32, y = 3 : i32}, {id = 38 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 0 : i32}, {id = 41 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 1 : i32}, {id = 44 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data - %9 = "neura.mac"(%8) <{stationary = "weight"}> {dfg_id = 19 : i32, mapping_locs = [{id = 15 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "tile", time_step = 3 : i32, x = 3 : i32, y = 3 : i32}]} : (!neura.data) -> !neura.data - %10 = "neura.data_mov"(%arg7) {dfg_id = 4 : i32, mapping_locs = [{direction = "west", id = 16 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, io = "input", resource = "port", time_step = 1 : i32, x = 0 : i32, y = 2 : i32}]} : (!neura.data) -> !neura.data - %11 = "neura.data_mov"(%3) {dfg_id = 20 : i32, mapping_locs = [{id = 39 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 0 : i32}]} : (!neura.data) -> !neura.data - %12 = "neura.mac"(%10, %11) <{stationary = "weight"}> {dfg_id = 24 : i32, mapping_locs = [{id = 8 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 1 : i32, x = 0 : i32, y = 2 : i32}]} : (!neura.data, !neura.data) -> !neura.data - %13 = "neura.data_mov"(%arg7) {dfg_id = 5 : i32, mapping_locs = [{direction = "west", id = 16 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, io = "input", resource = "port", time_step = 1 : i32, x = 0 : i32, y = 2 : i32}, {id = 24 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data) -> !neura.data - %14 = "neura.data_mov"(%5) {dfg_id = 21 : i32, mapping_locs = [{id = 42 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data) -> !neura.data - %15 = "neura.mac"(%13, %14) <{stationary = "weight"}> {dfg_id = 25 : i32, mapping_locs = [{id = 9 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 2 : i32, x = 1 : i32, y = 2 : i32}]} : (!neura.data, !neura.data) -> !neura.data - %16 = "neura.data_mov"(%arg7) {dfg_id = 6 : i32, mapping_locs = [{direction = "west", id = 16 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, io = "input", resource = "port", time_step = 1 : i32, x = 0 : i32, y = 2 : i32}, {id = 24 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 1 : i32}, {id = 28 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data - %17 = "neura.data_mov"(%7) {dfg_id = 22 : i32, mapping_locs = [{id = 45 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data - %18 = "neura.mac"(%16, %17) <{stationary = "weight"}> {dfg_id = 26 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "tile", time_step = 3 : i32, x = 2 : i32, y = 2 : i32}]} : (!neura.data, !neura.data) -> !neura.data - %19 = "neura.data_mov"(%arg7) {dfg_id = 7 : i32, mapping_locs = [{direction = "west", id = 16 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, io = "input", resource = "port", time_step = 1 : i32, x = 0 : i32, y = 2 : i32}, {id = 24 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 1 : i32}, {id = 28 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}, {id = 32 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data - %20 = "neura.data_mov"(%9) {dfg_id = 23 : i32, mapping_locs = [{id = 47 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data - %21 = "neura.mac"(%19, %20) <{stationary = "weight"}> {dfg_id = 27 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "tile", time_step = 4 : i32, x = 3 : i32, y = 2 : i32}]} : (!neura.data, !neura.data) -> !neura.data - %22 = "neura.data_mov"(%arg8) {dfg_id = 8 : i32, mapping_locs = [{direction = "west", id = 12 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, io = "input", resource = "port", time_step = 2 : i32, x = 0 : i32, y = 1 : i32}]} : (!neura.data) -> !neura.data - %23 = "neura.data_mov"(%12) {dfg_id = 28 : i32, mapping_locs = [{id = 25 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data) -> !neura.data - %24 = "neura.mac"(%22, %23) <{stationary = "weight"}> {dfg_id = 32 : i32, mapping_locs = [{id = 4 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 2 : i32, x = 0 : i32, y = 1 : i32}]} : (!neura.data, !neura.data) -> !neura.data - %25 = "neura.data_mov"(%arg8) {dfg_id = 9 : i32, mapping_locs = [{direction = "west", id = 12 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, io = "input", resource = "port", time_step = 2 : i32, x = 0 : i32, y = 1 : i32}, {id = 10 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data - %26 = "neura.data_mov"(%15) {dfg_id = 29 : i32, mapping_locs = [{id = 29 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data - %27 = "neura.mac"(%25, %26) <{stationary = "weight"}> {dfg_id = 33 : i32, mapping_locs = [{id = 5 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "tile", time_step = 3 : i32, x = 1 : i32, y = 1 : i32}]} : (!neura.data, !neura.data) -> !neura.data - %28 = "neura.data_mov"(%arg8) {dfg_id = 10 : i32, mapping_locs = [{direction = "west", id = 12 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, io = "input", resource = "port", time_step = 2 : i32, x = 0 : i32, y = 1 : i32}, {id = 10 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}, {id = 14 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data - %29 = "neura.data_mov"(%18) {dfg_id = 30 : i32, mapping_locs = [{id = 33 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data - %30 = "neura.mac"(%28, %29) <{stationary = "weight"}> {dfg_id = 34 : i32, mapping_locs = [{id = 6 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "tile", time_step = 4 : i32, x = 2 : i32, y = 1 : i32}]} : (!neura.data, !neura.data) -> !neura.data - %31 = "neura.data_mov"(%arg8) {dfg_id = 11 : i32, mapping_locs = [{direction = "west", id = 12 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, io = "input", resource = "port", time_step = 2 : i32, x = 0 : i32, y = 1 : i32}, {id = 10 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}, {id = 14 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}, {id = 18 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "link", time_step = 4 : i32}]} : (!neura.data) -> !neura.data - %32 = "neura.data_mov"(%21) {dfg_id = 31 : i32, mapping_locs = [{id = 36 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "link", time_step = 4 : i32}]} : (!neura.data) -> !neura.data - %33 = "neura.mac"(%31, %32) <{stationary = "weight"}> {dfg_id = 35 : i32, mapping_locs = [{id = 7 : i32, index_per_ii = 0 : i32, invalid_iterations = 5 : i32, resource = "tile", time_step = 5 : i32, x = 3 : i32, y = 1 : i32}]} : (!neura.data, !neura.data) -> !neura.data - %34 = "neura.data_mov"(%arg9) {dfg_id = 12 : i32, mapping_locs = [{direction = "west", id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, io = "input", resource = "port", time_step = 3 : i32, x = 0 : i32, y = 0 : i32}]} : (!neura.data) -> !neura.data - %35 = "neura.data_mov"(%24) {dfg_id = 36 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data - %36 = "neura.mac"(%34, %35) <{stationary = "weight"}> {dfg_id = 40 : i32, mapping_locs = [{id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "tile", time_step = 3 : i32, x = 0 : i32, y = 0 : i32}]} : (!neura.data, !neura.data) -> !neura.data - %37 = "neura.data_mov"(%arg9) {dfg_id = 13 : i32, mapping_locs = [{direction = "west", id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, io = "input", resource = "port", time_step = 3 : i32, x = 0 : i32, y = 0 : i32}, {id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data - %38 = "neura.data_mov"(%27) {dfg_id = 37 : i32, mapping_locs = [{id = 15 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data - %39 = "neura.mac"(%37, %38) <{stationary = "weight"}> {dfg_id = 41 : i32, mapping_locs = [{id = 1 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "tile", time_step = 4 : i32, x = 1 : i32, y = 0 : i32}]} : (!neura.data, !neura.data) -> !neura.data - %40 = "neura.data_mov"(%arg9) {dfg_id = 14 : i32, mapping_locs = [{direction = "west", id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, io = "input", resource = "port", time_step = 3 : i32, x = 0 : i32, y = 0 : i32}, {id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}, {id = 3 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "link", time_step = 4 : i32}]} : (!neura.data) -> !neura.data - %41 = "neura.data_mov"(%30) {dfg_id = 38 : i32, mapping_locs = [{id = 19 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "link", time_step = 4 : i32}]} : (!neura.data) -> !neura.data - %42 = "neura.mac"(%40, %41) <{stationary = "weight"}> {dfg_id = 42 : i32, mapping_locs = [{id = 2 : i32, index_per_ii = 0 : i32, invalid_iterations = 5 : i32, resource = "tile", time_step = 5 : i32, x = 2 : i32, y = 0 : i32}]} : (!neura.data, !neura.data) -> !neura.data - %43 = "neura.data_mov"(%arg9) {dfg_id = 15 : i32, mapping_locs = [{direction = "west", id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, io = "input", resource = "port", time_step = 3 : i32, x = 0 : i32, y = 0 : i32}, {id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}, {id = 3 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "link", time_step = 4 : i32}, {id = 6 : i32, index_per_ii = 0 : i32, invalid_iterations = 5 : i32, resource = "link", time_step = 5 : i32}]} : (!neura.data) -> !neura.data - %44 = "neura.data_mov"(%33) {dfg_id = 39 : i32, mapping_locs = [{id = 22 : i32, index_per_ii = 0 : i32, invalid_iterations = 5 : i32, resource = "link", time_step = 5 : i32}]} : (!neura.data) -> !neura.data - %45 = "neura.mac"(%43, %44) <{stationary = "weight"}> {dfg_id = 43 : i32, mapping_locs = [{id = 3 : i32, index_per_ii = 0 : i32, invalid_iterations = 6 : i32, resource = "tile", time_step = 6 : i32, x = 3 : i32, y = 0 : i32}]} : (!neura.data, !neura.data) -> !neura.data - %46 = "neura.data_mov"(%36) {dfg_id = 44 : i32, mapping_locs = [{direction = "south", id = 3 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, io = "output", resource = "port", time_step = 3 : i32, x = 0 : i32, y = 0 : i32}]} : (!neura.data) -> !neura.data - %47 = "neura.data_mov"(%39) {dfg_id = 45 : i32, mapping_locs = [{direction = "south", id = 5 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, io = "output", resource = "port", time_step = 4 : i32, x = 1 : i32, y = 0 : i32}]} : (!neura.data) -> !neura.data - %48 = "neura.data_mov"(%42) {dfg_id = 46 : i32, mapping_locs = [{direction = "south", id = 7 : i32, index_per_ii = 0 : i32, invalid_iterations = 5 : i32, io = "output", resource = "port", time_step = 5 : i32, x = 2 : i32, y = 0 : i32}]} : (!neura.data) -> !neura.data - %49 = "neura.data_mov"(%45) {dfg_id = 47 : i32, mapping_locs = [{direction = "south", id = 11 : i32, index_per_ii = 0 : i32, invalid_iterations = 6 : i32, io = "output", resource = "port", time_step = 6 : i32, x = 3 : i32, y = 0 : i32}]} : (!neura.data) -> !neura.data - neura.yield results(%46, %47, %48, %49 : !neura.data, !neura.data, !neura.data, !neura.data) {dfg_id = 48 : i32} + %2 = "neura.data_mov"(%arg6) {dfg_id = 0 : i32, mapping_locs = [{direction = "west", id = 20 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, io = "input", resource = "boundary_port", time_step = 0 : i32, x = 0 : i32, y = 3 : i32}]} : (!neura.data) -> !neura.data + %result, %forwarded = "neura.mac"(%2) {dfg_id = 4 : i32, mapping_locs = [{id = 12 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 0 : i32, x = 0 : i32, y = 3 : i32}]} : (!neura.data) -> (!neura.data, !neura.data) + %3 = "neura.data_mov"(%forwarded) {dfg_id = 6 : i32, mapping_locs = [{id = 38 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 0 : i32}]} : (!neura.data) -> !neura.data + %result_0, %forwarded_1 = "neura.mac"(%3) {dfg_id = 8 : i32, mapping_locs = [{id = 13 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 1 : i32, x = 1 : i32, y = 3 : i32}]} : (!neura.data) -> (!neura.data, !neura.data) + %4 = "neura.data_mov"(%forwarded_1) {dfg_id = 12 : i32, mapping_locs = [{id = 41 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data) -> !neura.data + %result_2, %forwarded_3 = "neura.mac"(%4) {dfg_id = 15 : i32, mapping_locs = [{id = 14 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 2 : i32, x = 2 : i32, y = 3 : i32}]} : (!neura.data) -> (!neura.data, !neura.data) + %5 = "neura.data_mov"(%forwarded_3) {dfg_id = 21 : i32, mapping_locs = [{id = 44 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data + %result_4, %forwarded_5 = "neura.mac"(%5) {dfg_id = 25 : i32, mapping_locs = [{id = 15 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "tile", time_step = 3 : i32, x = 3 : i32, y = 3 : i32}]} : (!neura.data) -> (!neura.data, !neura.data) + %6 = "neura.data_mov"(%arg7) {dfg_id = 1 : i32, mapping_locs = [{direction = "west", id = 16 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, io = "input", resource = "boundary_port", time_step = 1 : i32, x = 0 : i32, y = 2 : i32}]} : (!neura.data) -> !neura.data + %7 = "neura.data_mov"(%result) {dfg_id = 5 : i32, mapping_locs = [{id = 39 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 0 : i32}]} : (!neura.data) -> !neura.data + %result_6, %forwarded_7 = "neura.mac"(%6, %7) {dfg_id = 7 : i32, mapping_locs = [{id = 8 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 1 : i32, x = 0 : i32, y = 2 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) + %8 = "neura.data_mov"(%forwarded_7) {dfg_id = 10 : i32, mapping_locs = [{id = 24 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data) -> !neura.data + %9 = "neura.data_mov"(%result_0) {dfg_id = 11 : i32, mapping_locs = [{id = 42 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data) -> !neura.data + %result_8, %forwarded_9 = "neura.mac"(%8, %9) {dfg_id = 14 : i32, mapping_locs = [{id = 9 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 2 : i32, x = 1 : i32, y = 2 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) + %10 = "neura.data_mov"(%forwarded_9) {dfg_id = 19 : i32, mapping_locs = [{id = 28 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data + %11 = "neura.data_mov"(%result_2) {dfg_id = 20 : i32, mapping_locs = [{id = 45 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data + %result_10, %forwarded_11 = "neura.mac"(%10, %11) {dfg_id = 24 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "tile", time_step = 3 : i32, x = 2 : i32, y = 2 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) + %12 = "neura.data_mov"(%forwarded_11) {dfg_id = 31 : i32, mapping_locs = [{id = 32 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data + %13 = "neura.data_mov"(%result_4) {dfg_id = 32 : i32, mapping_locs = [{id = 47 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data + %result_12, %forwarded_13 = "neura.mac"(%12, %13) {dfg_id = 35 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "tile", time_step = 4 : i32, x = 3 : i32, y = 2 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) + %14 = "neura.data_mov"(%arg8) {dfg_id = 2 : i32, mapping_locs = [{direction = "west", id = 12 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, io = "input", resource = "boundary_port", time_step = 2 : i32, x = 0 : i32, y = 1 : i32}]} : (!neura.data) -> !neura.data + %15 = "neura.data_mov"(%result_6) {dfg_id = 9 : i32, mapping_locs = [{id = 25 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data) -> !neura.data + %result_14, %forwarded_15 = "neura.mac"(%14, %15) {dfg_id = 13 : i32, mapping_locs = [{id = 4 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 2 : i32, x = 0 : i32, y = 1 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) + %16 = "neura.data_mov"(%forwarded_15) {dfg_id = 17 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data + %17 = "neura.data_mov"(%result_8) {dfg_id = 18 : i32, mapping_locs = [{id = 29 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data + %result_16, %forwarded_17 = "neura.mac"(%16, %17) {dfg_id = 23 : i32, mapping_locs = [{id = 5 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "tile", time_step = 3 : i32, x = 1 : i32, y = 1 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) + %18 = "neura.data_mov"(%forwarded_17) {dfg_id = 29 : i32, mapping_locs = [{id = 14 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data + %19 = "neura.data_mov"(%result_10) {dfg_id = 30 : i32, mapping_locs = [{id = 33 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data + %result_18, %forwarded_19 = "neura.mac"(%18, %19) {dfg_id = 34 : i32, mapping_locs = [{id = 6 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "tile", time_step = 4 : i32, x = 2 : i32, y = 1 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) + %20 = "neura.data_mov"(%forwarded_19) {dfg_id = 39 : i32, mapping_locs = [{id = 18 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "link", time_step = 4 : i32}]} : (!neura.data) -> !neura.data + %21 = "neura.data_mov"(%result_12) {dfg_id = 40 : i32, mapping_locs = [{id = 36 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "link", time_step = 4 : i32}]} : (!neura.data) -> !neura.data + %result_20, %forwarded_21 = "neura.mac"(%20, %21) {dfg_id = 42 : i32, mapping_locs = [{id = 7 : i32, index_per_ii = 0 : i32, invalid_iterations = 5 : i32, resource = "tile", time_step = 5 : i32, x = 3 : i32, y = 1 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) + %22 = "neura.data_mov"(%arg9) {dfg_id = 3 : i32, mapping_locs = [{direction = "west", id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, io = "input", resource = "boundary_port", time_step = 3 : i32, x = 0 : i32, y = 0 : i32}]} : (!neura.data) -> !neura.data + %23 = "neura.data_mov"(%result_14) {dfg_id = 16 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data + %result_22, %forwarded_23 = "neura.mac"(%22, %23) {dfg_id = 22 : i32, mapping_locs = [{id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "tile", time_step = 3 : i32, x = 0 : i32, y = 0 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) + %24 = "neura.data_mov"(%forwarded_23) {dfg_id = 27 : i32, mapping_locs = [{id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data + %25 = "neura.data_mov"(%result_16) {dfg_id = 28 : i32, mapping_locs = [{id = 15 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data + %result_24, %forwarded_25 = "neura.mac"(%24, %25) {dfg_id = 33 : i32, mapping_locs = [{id = 1 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "tile", time_step = 4 : i32, x = 1 : i32, y = 0 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) + %26 = "neura.data_mov"(%forwarded_25) {dfg_id = 37 : i32, mapping_locs = [{id = 3 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "link", time_step = 4 : i32}]} : (!neura.data) -> !neura.data + %27 = "neura.data_mov"(%result_18) {dfg_id = 38 : i32, mapping_locs = [{id = 19 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "link", time_step = 4 : i32}]} : (!neura.data) -> !neura.data + %result_26, %forwarded_27 = "neura.mac"(%26, %27) {dfg_id = 41 : i32, mapping_locs = [{id = 2 : i32, index_per_ii = 0 : i32, invalid_iterations = 5 : i32, resource = "tile", time_step = 5 : i32, x = 2 : i32, y = 0 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) + %28 = "neura.data_mov"(%forwarded_27) {dfg_id = 44 : i32, mapping_locs = [{id = 6 : i32, index_per_ii = 0 : i32, invalid_iterations = 5 : i32, resource = "link", time_step = 5 : i32}]} : (!neura.data) -> !neura.data + %29 = "neura.data_mov"(%result_20) {dfg_id = 45 : i32, mapping_locs = [{id = 22 : i32, index_per_ii = 0 : i32, invalid_iterations = 5 : i32, resource = "link", time_step = 5 : i32}]} : (!neura.data) -> !neura.data + %result_28, %forwarded_29 = "neura.mac"(%28, %29) {dfg_id = 46 : i32, mapping_locs = [{id = 3 : i32, index_per_ii = 0 : i32, invalid_iterations = 6 : i32, resource = "tile", time_step = 6 : i32, x = 3 : i32, y = 0 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) + %30 = "neura.data_mov"(%result_22) {dfg_id = 26 : i32, mapping_locs = [{direction = "south", id = 3 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, io = "output", resource = "boundary_port", time_step = 3 : i32, x = 0 : i32, y = 0 : i32}]} : (!neura.data) -> !neura.data + %31 = "neura.data_mov"(%result_24) {dfg_id = 36 : i32, mapping_locs = [{direction = "south", id = 5 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, io = "output", resource = "boundary_port", time_step = 4 : i32, x = 1 : i32, y = 0 : i32}]} : (!neura.data) -> !neura.data + %32 = "neura.data_mov"(%result_26) {dfg_id = 43 : i32, mapping_locs = [{direction = "south", id = 7 : i32, index_per_ii = 0 : i32, invalid_iterations = 5 : i32, io = "output", resource = "boundary_port", time_step = 5 : i32, x = 2 : i32, y = 0 : i32}]} : (!neura.data) -> !neura.data + %33 = "neura.data_mov"(%result_28) {dfg_id = 47 : i32, mapping_locs = [{direction = "south", id = 11 : i32, index_per_ii = 0 : i32, invalid_iterations = 6 : i32, io = "output", resource = "boundary_port", time_step = 6 : i32, x = 3 : i32, y = 0 : i32}]} : (!neura.data) -> !neura.data + neura.yield results(%30, %31, %32, %33 : !neura.data, !neura.data, !neura.data, !neura.data) {dfg_id = 48 : i32} } : i32, i32, i32, i32 taskflow.stream_write(%1#0, %1#1, %1#2, %1#3 : i32, i32, i32, i32) to %arg5 maps [#map, #map1, #map2, #map3] : memref<4x4xi32> taskflow.yield done_writes(%arg5 : memref<4x4xi32>) From 90100e6bcff6b754aef9413664ccacd1822e21f2 Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Sat, 12 Sep 2026 00:50:49 +0800 Subject: [PATCH 14/19] Add TileArray memory lowering --- .gitignore | 1 + mlir/amoeba | 2 +- python/synapse/__init__.py | 4 +- python/synapse/compiler/__init__.py | 4 +- python/synapse/frontend/lowering.py | 706 +++++++++--------- python/synapse/language/__init__.py | 8 +- python/synapse/language/spatial.py | 55 +- python/synapse/language/tensor.py | 13 +- python/synapse/language/tile_array_program.py | 217 +++--- python/synapse/library/__init__.py | 5 + python/synapse/library/gemm.py | 38 + python/synapse/patterns/__init__.py | 5 + python/synapse/templates/__init__.py | 1 - .../synapse/templates/tile_array/__init__.py | 7 - python/synapse/templates/tile_array/gemm.py | 46 -- tests/python/compiler/test_systolic_gemm.py | 225 +++--- tests/python/frontend/test_memory_lowering.py | 124 +++ tests/python/language/test_memory.py | 73 ++ tests/python/language/test_spatial.py | 30 - 19 files changed, 851 insertions(+), 713 deletions(-) create mode 100644 python/synapse/library/__init__.py create mode 100644 python/synapse/library/gemm.py create mode 100644 python/synapse/patterns/__init__.py delete mode 100644 python/synapse/templates/__init__.py delete mode 100644 python/synapse/templates/tile_array/__init__.py delete mode 100644 python/synapse/templates/tile_array/gemm.py create mode 100644 tests/python/frontend/test_memory_lowering.py create mode 100644 tests/python/language/test_memory.py diff --git a/.gitignore b/.gitignore index 7117fdc..03d3467 100644 --- a/.gitignore +++ b/.gitignore @@ -10,6 +10,7 @@ __pycache__/ *.egg-info/ build/ dist/ +docs/ .eggs/ # Test / type-check / coverage caches diff --git a/mlir/amoeba b/mlir/amoeba index 60a8bed..a6fb32a 160000 --- a/mlir/amoeba +++ b/mlir/amoeba @@ -1 +1 @@ -Subproject commit 60a8bed1b657b4a0c036e3df018d5bc945075b41 +Subproject commit a6fb32a2473f8b616e1044bdc3bed178a2dcc692 diff --git a/python/synapse/__init__.py b/python/synapse/__init__.py index 236470b..441fb11 100644 --- a/python/synapse/__init__.py +++ b/python/synapse/__init__.py @@ -1,3 +1,3 @@ -from .compiler import compile +from .compiler import compile, rewrite -__all__ = ["compile"] +__all__ = ["compile", "rewrite"] diff --git a/python/synapse/compiler/__init__.py b/python/synapse/compiler/__init__.py index 236470b..441fb11 100644 --- a/python/synapse/compiler/__init__.py +++ b/python/synapse/compiler/__init__.py @@ -1,3 +1,3 @@ -from .compiler import compile +from .compiler import compile, rewrite -__all__ = ["compile"] +__all__ = ["compile", "rewrite"] diff --git a/python/synapse/frontend/lowering.py b/python/synapse/frontend/lowering.py index 0f5e7e7..ed724ae 100644 --- a/python/synapse/frontend/lowering.py +++ b/python/synapse/frontend/lowering.py @@ -1,28 +1,48 @@ -"""Lower Synapse Python programs to compiler input IR.""" +"""Lowers Synapse Python programs to compiler input IR.""" + +from __future__ import annotations from collections.abc import Callable -from functools import singledispatch +from functools import singledispatchmethod from inspect import signature -from typing import cast +from itertools import product +from math import prod +from typing import TYPE_CHECKING, cast from synapse.language.spatial import Tile from synapse.language.tensor import Tensor, TensorAccess from synapse.language.tile_array_program import ( AddOp, ConstantOp, + LoadOp, MacOp, + StoreOp, TileArrayBuilder, TileArrayOp, TileArrayProgram, ) from synapse.language.types import DType, TensorType +if TYPE_CHECKING: + from taskflow_mlir.ir import AffineMapAttr, DenseI64ArrayAttr, DictAttr, Value + def lower(program_fn: Callable, *, argument_types: tuple[TensorType, ...] = ()) -> str: - """Lowers one tile-array program to pre-mapping Taskflow and Neura IR.""" - # TODO: Generalize program arguments beyond Tensor. Accept both DType and - # TensorType, create Scalar or Tensor symbolic values accordingly, and - # lower scalar dependencies through Taskflow value_inputs/value_outputs. + """Lowers a standalone TileArray program to pre-mapping Taskflow and Neura IR.""" + program = build_tile_array_program(program_fn, argument_types=argument_types) + return TileArrayProgramLowering(program).lower_to_single_task(program_fn.__name__) + + +def build_tile_array_program( + program_fn: Callable, *, argument_types: tuple[TensorType, ...] = () +) -> TileArrayProgram: + """Runs a TileArray function and records its operations and tensor accesses. + + Direct compilation and pattern rewriting share this construction step. + It creates a TileArrayProgram without importing or constructing MLIR. + """ + # Program arguments currently carry tensor types. Scalar argument capture + # will use Taskflow value dependencies when that frontend path is added. function_signature = signature(program_fn) parameter_names = tuple(function_signature.parameters) @@ -49,407 +69,391 @@ def lower(program_fn: Callable, *, argument_types: tuple[TensorType, ...] = ()) with builder: program_fn(*arguments) - 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 ( - AffineExpr, - AffineMap, - AffineMapAttr, - ArrayAttr, - Context, - DictAttr, - F32Type, - FloatAttr, - InsertionPoint, - IntegerAttr, - IntegerType, - Location, - MemRefType, - Module, - StringAttr, - ) - - with Context(), Location.unknown(): - taskflow.register_dialect() - neura.register_dialect() - - i32_type = IntegerType.get_signless(32) - - def get_mlir_type(dtype: DType): - """Translate a frontend data type into an MLIR type.""" - - if dtype == DType.I32: - return i32_type - - if dtype == DType.F32: - return F32Type.get() - - raise NotImplementedError( - f"unsupported tile-array data type: {dtype.value}" - ) - - def get_memref_type(tensor_type: TensorType): - """Translates a TensorType into a Taskflow MemRef type.""" - return MemRefType.get( - list(tensor_type.shape), get_mlir_type(tensor_type.dtype) - ) - - def get_access_map(access: TensorAccess) -> AffineMapAttr: - """Builds the affine map that enumerates one tensor access.""" - - results = [] - next_dimension = 0 - - for index in access.indices: - if isinstance(index, slice): - results.append(AffineExpr.get_dim(next_dimension)) - next_dimension += 1 - else: - results.append(AffineExpr.get_constant(index)) - - return AffineMapAttr.get(AffineMap.get(next_dimension, 0, results)) - - def get_stationary_map() -> AffineMapAttr: - """Builds and validates the stationary Tile-to-data map.""" - if program.stationary is None: - raise RuntimeError("template program requires stationary data") - - for tile, access in program.stationary.tile_values: - expected_indices = (program.array.y_tiles - 1 - tile.y, tile.x) - - if access.indices != expected_indices: - raise ValueError("weight-stationary GEMM requires B[K - 1 - y, x]") - - x = AffineExpr.get_dim(0) - y = AffineExpr.get_dim(1) - - last_row = AffineExpr.get_constant(program.array.y_tiles - 1) - - negative_y = AffineExpr.get_mul(AffineExpr.get_constant(-1), y) - - weight_row = AffineExpr.get_add(last_row, negative_y) + return builder.build() - return AffineMapAttr.get(AffineMap.get(2, 0, [weight_row, x])) - def get_constant_attribute( - operation: ConstantOp, - result_type, - ): - """Build the typed MLIR attribute for a constant value.""" - result = operation.results[0] +class TileArrayProgramLowering: + """Lowers a TileArrayProgram into the current MLIR context.""" - if result.dtype == DType.I32: - return IntegerAttr.get(result_type, cast(int, operation.value)) - - if result.dtype == DType.F32: - return FloatAttr.get(result_type, float(operation.value)) - - raise NotImplementedError( - f"unsupported constant type: {result.dtype.value}" - ) - - def get_placement(tile: Tile) -> DictAttr: - """Build Neura placement directly from a Tile coordinate.""" - - return DictAttr.get( - { - "x": IntegerAttr.get(i32_type, tile.x), - "y": IntegerAttr.get(i32_type, tile.y), - } - ) - - def get_kernel_metadata() -> DictAttr | None: - """Materialize template, stationary, and Port metadata.""" + def __init__(self, program: TileArrayProgram): + self.program = program + self.kernel_input_indices = { + argument: index for index, argument in enumerate(program.arguments) + } + read_sources = { + operation.source.source + for operation in program.operations + if isinstance(operation, LoadOp) and operation.source is not None + } + write_sources = { + operation.target.source + for operation in program.operations + if isinstance(operation, StoreOp) and operation.target is not None + } + if program.stationary is not None: + read_sources.add(program.stationary.source) + if not (read_sources | write_sources).issubset(self.kernel_input_indices): + raise ValueError("memory accesses must reference program arguments") + self.has_dynamic_memory = any( + isinstance(operation, (LoadOp, StoreOp)) and operation.addr is not None + for operation in program.operations + ) + # Raw addresses may alias any captured buffer. Task dependencies remain + # conservative until address provenance becomes available. + if self.has_dynamic_memory: + read_sources.update(program.arguments) + write_sources.update(program.arguments) + self.read_arguments = tuple( + argument for argument in program.arguments if argument in read_sources + ) + self.write_arguments = tuple( + argument for argument in program.arguments if argument in write_sources + ) - if program.template_name is None: - return None + # The standalone path is used only when a TileArray program is lowered + # without an existing Taskflow graph. + def lower_to_single_task(self, program_name: str) -> str: + """Creates a complete module containing one TileArray task.""" - if program.stationary is None: - raise RuntimeError("template program requires stationary metadata") + from taskflow_mlir.dialects import func, neura, taskflow + from taskflow_mlir.ir import Context, InsertionPoint, Location, Module - input_ports = ArrayAttr.get( - [ - DictAttr.get( - { - "kernel_input": IntegerAttr.get(i32_type, index), - "direction": StringAttr.get(binding.port.direction), - "x": IntegerAttr.get(i32_type, binding.port.x), - "y": IntegerAttr.get(i32_type, binding.port.y), - } - ) - for index, binding in enumerate(program.input_ports) - ] - ) + with Context(), Location.unknown(): + taskflow.register_dialect() + neura.register_dialect() - output_ports = ArrayAttr.get( - [ - DictAttr.get( - { - "kernel_result": IntegerAttr.get(i32_type, index), - "direction": StringAttr.get(binding.port.direction), - "x": IntegerAttr.get(i32_type, binding.port.x), - "y": IntegerAttr.get(i32_type, binding.port.y), - } - ) - for index, binding in enumerate(program.output_ports) - ] - ) + module = Module.create() + argument_types = [ + self.get_memref_type(argument.type) + for argument in self.program.arguments + ] + result_types = [ + self.get_memref_type(argument.type) for argument in self.write_arguments + ] - stationary = DictAttr.get( - { - "kernel_input": IntegerAttr.get(i32_type, len(program.input_ports)), - "map": get_stationary_map(), - } - ) + with InsertionPoint(module.body): + function = func.FuncOp( + program_name, + (argument_types, result_types), + ) + function_block = function.add_entry_block() - template = DictAttr.get( - { - "input_ports": input_ports, - "name": StringAttr.get(program.template_name), - "output_ports": output_ports, - "stationary": stationary, - } + function_values = dict( + zip(self.program.arguments, function_block.arguments) ) + read_values = [ + function_values[argument] for argument in self.read_arguments + ] + write_values = [ + function_values[argument] for argument in self.write_arguments + ] - return DictAttr.get( - { - "kind": StringAttr.get("template"), - "template": template, - } + used_arguments = set(self.read_arguments) | set(self.write_arguments) + other_arguments = tuple( + argument + for argument in self.program.arguments + if argument not in used_arguments ) + other_values = [function_values[argument] for argument in other_arguments] + + with InsertionPoint(function_block): + task = taskflow.TaskflowTaskOp( + done_reads=[], + done_writes=[value.type for value in write_values], + value_outputs=[], + will_reads=read_values, + will_writes=write_values, + value_inputs=other_values, + task_name=program_name, + original_read_memrefs=read_values, + original_write_memrefs=write_values, + ) + task_block = task.body.blocks.append( + *[value.type for value in read_values], + *[value.type for value in write_values], + *[value.type for value in other_values], + ) - @singledispatch - def lower_operation(operation: TileArrayOp, operands, result_types): - """Lower one frontend TileArray operation to a Neura operation. + func.ReturnOp(task.done_writes) - 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__}" + task_arguments = dict( + zip( + self.read_arguments + self.write_arguments + other_arguments, + task_block.arguments, + ) ) - @lower_operation.register - def lower_constant(operation: ConstantOp, operands, result_types): - """Lower a ConstantOp to neura.constant.""" - return neura.ConstantOp( - result_types[0], get_constant_attribute(operation, result_types[0]) - ) + with InsertionPoint(task_block): + self.lower_to_kernel(task_arguments) - @lower_operation.register - def lower_add(operation: AddOp, operands, result_types): - """Lower a frontend AddOp to neura.add.""" - lhs, rhs = operands + taskflow.TaskflowYieldOp( + done_reads=[], + done_writes=[ + task_arguments[argument] for argument in self.write_arguments + ], + value_results=[], + ) - return neura.AddOp(result_types[0], lhs, rhs=rhs) + if not module.operation.verify(): + raise RuntimeError("generated Taskflow/Neura module is invalid") + + return str(module) + + def lower_to_kernel(self, argument_values: dict[Tensor, Value]) -> None: + """Creates one kernel at the caller's insertion point in an existing task. + + The caller owns the current context, location, task, and terminator. + Buffer arguments retain their program order for configuration indices. + """ + from taskflow_mlir.dialects import neura + from taskflow_mlir.ir import InsertionPoint, StringAttr + + inputs = [] + for argument in self.program.arguments: + value = argument_values[argument] + if value.type != self.get_memref_type(argument.type): + raise TypeError(f"unexpected memref type for {argument.name}") + inputs.append(value) + + # Configuration validation precedes IR insertion. + metadata = self.get_kernel_metadata() + memory_configs: dict[int, tuple[TensorAccess, DenseI64ArrayAttr]] = {} + for index, operation in enumerate(self.program.operations): + if isinstance(operation, LoadOp): + access = operation.source + elif isinstance(operation, StoreOp): + access = operation.target + else: + continue + if access is not None: + memory_configs[index] = (access, self.get_memory_offsets(access)) + kernel = neura.KernelOp( + outputs=[], + inputs=inputs, + iter_args_init=[], + accelerator=StringAttr.get("neura"), + kernel_metadata=metadata, + ) + block = kernel.body.blocks.append(*[value.type for value in inputs]) + values_by_id: dict[int, Value] = {} + with InsertionPoint(block): + for index, operation in enumerate(self.program.operations): + operands = tuple(values_by_id[value.id] for value in operation.operands) + memory_config = memory_configs.get(index) + if memory_config is not None: + access, _ = memory_config + # The memref supplies a launch-time base rather than a + # routed address value. Its SSA use stays inside the kernel. + base = block.arguments[self.kernel_input_indices[access.source]] + operands += (base,) + result_types = tuple( + self.get_mlir_type(value.dtype) for value in operation.results + ) + lowered = self.lower_operation(operation, operands, result_types) + lowered.operation.attributes["placement"] = self.get_placement( + operation.tile + ) + if memory_config is not None: + lowered.operation.attributes["constants"] = memory_config[1] + results = tuple(lowered.results) + if len(results) != len(operation.results): + raise RuntimeError( + "frontend and MLIR operation result counts differ" + ) + for source, result in zip(operation.results, results): + values_by_id[source.id] = result + neura.YieldOp(iter_args_next=[], results_=[]) - @lower_operation.register - def lower_mac(operation: MacOp, operands, result_types): - """Lower a configured MacOp to neura.mac.""" + def get_mlir_type(self, dtype: DType): + """Translates a frontend data type into an MLIR type.""" - input0 = operands[0] + from taskflow_mlir.ir import F32Type, IntegerType - input1 = operands[1] if len(operands) == 2 else None - accumulated_type, forwarded_type = result_types + if dtype == DType.I32: + return IntegerType.get_signless(32) - return neura.MacOp( - accumulated_type, - forwarded_type, - input0, - input1=input1, - ) + if dtype == DType.F32: + return F32Type.get() - read_sources = {binding.input_src.source for binding in program.input_ports} + raise NotImplementedError(f"unsupported tile-array data type: {dtype.value}") - if program.stationary is not None: - read_sources.add(program.stationary.source) + def get_memref_type(self, tensor_type: TensorType): + """Translates a TensorType into a Taskflow MemRef type.""" - write_sources = {binding.output_des.source for binding in program.output_ports} + from taskflow_mlir.ir import MemRefType - read_arguments = tuple( - argument for argument in program.arguments if argument in read_sources + return MemRefType.get( + list(tensor_type.shape), self.get_mlir_type(tensor_type.dtype) ) - write_arguments = tuple( - argument for argument in program.arguments if argument in write_sources - ) + def get_memory_offsets(self, access: TensorAccess) -> DenseI64ArrayAttr: + """Converts a static access to Neura's constant element-offset array.""" + from taskflow_mlir.ir import DenseI64ArrayAttr - function_argument_types = [ - get_memref_type(argument.type) for argument in program.arguments + shape = access.source.type.shape + dimensions = [ + range(*index.indices(size)) if isinstance(index, slice) else (index,) + for size, index in zip(shape, access.indices) ] - - function_result_types = [ - get_memref_type(argument.type) for argument in write_arguments + strides = [prod(shape[index + 1 :]) for index in range(len(shape))] + offsets = [ + sum(index * stride for index, stride in zip(indices, strides)) + for indices in product(*dimensions) ] - - module = Module.create() - - with InsertionPoint(module.body): - function = func.FuncOp( - program_name, - (function_argument_types, function_result_types), + if not offsets: + raise ValueError("configured address queues cannot be empty") + return DenseI64ArrayAttr.get(offsets) + + def get_stationary_map(self) -> AffineMapAttr: + """Infers a translated weight-stationary map from tile bindings.""" + from taskflow_mlir.ir import AffineExpr, AffineMap, AffineMapAttr + + stationary = self.program.stationary + if stationary is None or not stationary.tile_values: + raise ValueError("stationary bindings are required") + tile, access = stationary.tile_values[0] + if len(access.indices) != 2 or not access.is_scalar: + raise ValueError( + "stationary mapping currently requires static rank-2 elements" ) + k_index, column_index = access.indices + if not isinstance(k_index, int) or not isinstance(column_index, int): + raise TypeError("stationary indices must be static integers") + row_origin = k_index + tile.y + column_shift = column_index - tile.x + for tile, access in stationary.tile_values: + if access.indices != (row_origin - tile.y, tile.x + column_shift): + raise ValueError("stationary accesses do not form a translated WS map") + x = AffineExpr.get_dim(0) + y = AffineExpr.get_dim(1) + row = AffineExpr.get_add( + AffineExpr.get_constant(row_origin), + AffineExpr.get_mul(AffineExpr.get_constant(-1), y), + ) + column = AffineExpr.get_add(x, AffineExpr.get_constant(column_shift)) + return AffineMapAttr.get(AffineMap.get(2, 0, [row, column])) + + def get_kernel_metadata(self) -> DictAttr | None: + """Materializes the supported Neura stationary implementation metadata.""" + from taskflow_mlir.ir import DictAttr, IntegerAttr, StringAttr + + stationary = self.program.stationary + if stationary is None: + return None + return DictAttr.get( + { + "kind": StringAttr.get("template"), + "template": DictAttr.get( + { + "name": StringAttr.get("systolic_array"), + "stationary": DictAttr.get( + { + "kernel_input": IntegerAttr.get( + self.get_mlir_type(DType.I32), + self.kernel_input_indices[stationary.source], + ), + "map": self.get_stationary_map(), + } + ), + } + ), + } + ) - function_block = function.add_entry_block() - - function_values = dict(zip(program.arguments, function_block.arguments)) + def get_placement(self, tile: Tile) -> DictAttr: + """Builds Neura placement directly from a Tile coordinate.""" - read_values = [function_values[argument] for argument in read_arguments] + from taskflow_mlir.ir import DictAttr, IntegerAttr - write_values = [function_values[argument] for argument in write_arguments] + i32_type = self.get_mlir_type(DType.I32) - with InsertionPoint(function_block): - task = taskflow.TaskflowTaskOp( - done_reads=[], - done_writes=[value.type for value in write_values], - value_outputs=[], - will_reads=read_values, - will_writes=write_values, - value_inputs=[], - task_name=program_name, - original_read_memrefs=read_values, - original_write_memrefs=write_values, - ) + return DictAttr.get( + { + "x": IntegerAttr.get(i32_type, tile.x), + "y": IntegerAttr.get(i32_type, tile.y), + } + ) - task_block = task.body.blocks.append( - *[value.type for value in read_values], - *[value.type for value in write_values], - ) + def get_constant_attribute( + self, + operation: ConstantOp, + result_type, + ): + """Builds the typed MLIR attribute for a constant value.""" - func.ReturnOp(task.done_writes) + from taskflow_mlir.ir import FloatAttr, IntegerAttr - task_arguments = dict( - zip(read_arguments + write_arguments, task_block.arguments) - ) + result = operation.results[0] - with InsertionPoint(task_block): - stream_values_by_id = {} - - for argument in read_arguments: - bindings = [ - binding - for binding in program.input_ports - if binding.input_src.source is argument - ] - - if not bindings: - continue - - stream_read = taskflow.TaskflowStreamReadOp( - [get_mlir_type(binding.input_des.dtype) for binding in bindings], - task_arguments[argument], - ArrayAttr.get( - [get_access_map(binding.input_src) for binding in bindings] - ), - ) + if result.dtype == DType.I32: + return IntegerAttr.get(result_type, cast(int, operation.value)) - for binding, value in zip(bindings, stream_read.values): - stream_values_by_id[binding.input_des.id] = value + if result.dtype == DType.F32: + return FloatAttr.get(result_type, float(operation.value)) - kernel_input_values = [ - stream_values_by_id[binding.input_des.id] - for binding in program.input_ports - ] + raise NotImplementedError(f"unsupported constant type: {result.dtype.value}") - kernel_input_types = [value.type for value in kernel_input_values] + @singledispatchmethod + def lower_operation(self, operation: TileArrayOp, operands, result_types): + """Lowers one frontend TileArray operation to a Neura operation. - if program.stationary is not None: - stationary_value = task_arguments[program.stationary.source] + 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__}" + ) - kernel_input_values.append(stationary_value) + @lower_operation.register + def lower_constant(self, operation: ConstantOp, operands, result_types): + """Lowers a ConstantOp to neura.constant.""" - kernel_input_types.append(stationary_value.type) + from taskflow_mlir.dialects import neura - kernel_output_types = [ - get_mlir_type(binding.output_src.dtype) - for binding in program.output_ports - ] + return neura.ConstantOp( + result_types[0], self.get_constant_attribute(operation, result_types[0]) + ) - kernel = neura.KernelOp( - outputs=kernel_output_types, - inputs=kernel_input_values, - iter_args_init=[], - accelerator=StringAttr.get("neura"), - kernel_metadata=get_kernel_metadata(), - ) + @lower_operation.register + def lower_add(self, operation: AddOp, operands, result_types): + """Lowers a frontend AddOp to neura.add.""" - kernel_block = kernel.body.blocks.append(*kernel_input_types) - - if program.output_ports: - taskflow.TaskflowStreamWriteOp( - list(kernel.results), - task_arguments[write_arguments[0]], - ArrayAttr.get( - [ - get_access_map(binding.output_des) - for binding in program.output_ports - ] - ), - ) + from taskflow_mlir.dialects import neura - taskflow.TaskflowYieldOp( - done_reads=[], - done_writes=[task_arguments[argument] for argument in write_arguments], - value_results=[], - ) + lhs, rhs = operands - values_by_id = { - binding.input_des.id: kernel_block.arguments[index] - for index, binding in enumerate(program.input_ports) - } + return neura.AddOp(result_types[0], lhs, rhs=rhs) - with InsertionPoint(kernel_block): - for operation in program.operations: - result_types = tuple( - get_mlir_type(result.dtype) for result in operation.results - ) + @lower_operation.register + def lower_mac(self, operation: MacOp, operands, result_types): + """Lowers a configured MacOp to neura.mac.""" - mlir_operands = tuple( - values_by_id[operand.id] for operand in operation.operands - ) + from taskflow_mlir.dialects import neura - mlir_operation = lower_operation(operation, mlir_operands, result_types) + input0 = operands[0] - mlir_operation.operation.attributes["placement"] = get_placement( - operation.tile - ) + input1 = operands[1] if len(operands) == 2 else None + accumulated_type, forwarded_type = result_types - mlir_results = tuple(mlir_operation.results) + return neura.MacOp( + accumulated_type, + forwarded_type, + input0, + input1=input1, + ) - if len(mlir_results) != len(operation.results): - raise RuntimeError( - "frontend and MLIR operation result counts differ" - ) + @lower_operation.register + def lower_load(self, operation: LoadOp, operands, result_types): + """Lowers a configured or explicitly addressed load.""" + from taskflow_mlir.dialects import neura - for frontend_result, mlir_result in zip( - operation.results, mlir_results - ): - values_by_id[frontend_result.id] = mlir_result - - neura.YieldOp( - iter_args_next=[], - results_=[ - values_by_id[binding.output_src.id] - for binding in program.output_ports - ], - ) + return neura.LoadOp(result_types[0], addr=operands[0] if operands else None) - if not module.operation.verify(): - raise RuntimeError("generated Taskflow/Neura module is invalid") + @lower_operation.register + def lower_store(self, operation: StoreOp, operands, result_types): + """Lowers a configured or explicitly addressed store.""" + from taskflow_mlir.dialects import neura - return str(module) + return neura.StoreOp( + operands[0], addr=operands[1] if len(operands) == 2 else None + ) diff --git a/python/synapse/language/__init__.py b/python/synapse/language/__init__.py index 4e97894..8ddb19c 100644 --- a/python/synapse/language/__init__.py +++ b/python/synapse/language/__init__.py @@ -5,9 +5,9 @@ from .tile_array_program import ( add, constant, - input_port, + load, mac, - output_port, + store, ) from .types import DType, TensorType, f32, i32 @@ -20,7 +20,7 @@ "constant", "f32", "i32", - "input_port", + "load", "mac", - "output_port", + "store", ] diff --git a/python/synapse/language/spatial.py b/python/synapse/language/spatial.py index 1739304..4b8ab73 100644 --- a/python/synapse/language/spatial.py +++ b/python/synapse/language/spatial.py @@ -8,8 +8,6 @@ core array of a multi-CGRA, AMD AIE/NPU, or Tenstorrent. """ -from typing import Literal - class Tile: """A hardware tile in a CGRA TileArray. @@ -24,39 +22,6 @@ def __init__(self, x: int, y: int, array: "TileArray"): self.y = y -PortDirection = Literal["west", "east", "north", "south"] - - -class Port: - """A boundary data port of a CGRA TileArray. - - ``direction`` identifies the array boundary. The ``x`` and ``y`` - coordinates identify the boundary tile attached to this Port. - - Ports expose hardware connectivity to the programming model. They do not - prescribe when data is transferred or introduce clock-based scheduling. - """ - - def __init__( - self, - *, - direction: PortDirection, - x: int, - y: int, - array: "TileArray", - ): - self.direction = direction - self.x = x - self.y = y - self.array = array - - @property - def tile(self) -> Tile: - """Return the boundary tile attached to this Port.""" - - return self.array[self.x, self.y] - - class TileArray: """A parameterized two-dimensional tile array. @@ -80,26 +45,8 @@ def __init__(self, x_tiles: int, y_tiles: int): Tile(x=x, y=y, array=self) for y in range(y_tiles) for x in range(x_tiles) ) - self.west_ports = tuple( - Port(direction="west", x=0, y=y, array=self) for y in range(y_tiles) - ) - - self.east_ports = tuple( - Port(direction="east", x=x_tiles - 1, y=y, array=self) - for y in range(y_tiles) - ) - - self.north_ports = tuple( - Port(direction="north", x=x, y=y_tiles - 1, array=self) - for x in range(x_tiles) - ) - - self.south_ports = tuple( - Port(direction="south", x=x, y=0, array=self) for x in range(x_tiles) - ) - def __getitem__(self, coordinate: tuple[int, int]) -> Tile: - """Return the tile at the given ``(x, y)`` coordinate.""" + """Returns the tile at the given ``(x, y)`` coordinate.""" x, y = coordinate diff --git a/python/synapse/language/tensor.py b/python/synapse/language/tensor.py index f9aca45..2e97102 100644 --- a/python/synapse/language/tensor.py +++ b/python/synapse/language/tensor.py @@ -19,7 +19,7 @@ class Tensor: def __getitem__( self, indices: TensorIndex | tuple[TensorIndex, ...] ) -> TensorAccess: - """Describes a scalar element or full-dimensional slice of this tensor.""" + """Describes a scalar element or static slice of this tensor.""" if not isinstance(indices, tuple): indices = (indices,) @@ -37,11 +37,16 @@ def __getitem__( continue if isinstance(index, slice): - if index != slice(None): - raise ValueError("only full slices are supported initially") + if any( + value is not None and type(value) is not int + for value in (index.start, index.stop, index.step) + ): + raise TypeError("slice bounds and steps must be integers") + if index.step == 0: + raise ValueError("slice step cannot be zero") continue - raise TypeError("indices must be integers or full slices") + raise TypeError("indices must be integers or static slices") return TensorAccess(source=self, indices=indices) diff --git a/python/synapse/language/tile_array_program.py b/python/synapse/language/tile_array_program.py index ba9cffa..40d4694 100644 --- a/python/synapse/language/tile_array_program.py +++ b/python/synapse/language/tile_array_program.py @@ -11,7 +11,7 @@ from contextvars import ContextVar from dataclasses import dataclass, field -from .spatial import Port, Tile, TileArray +from .spatial import Tile, TileArray from .tensor import Tensor, TensorAccess from .types import DType @@ -25,24 +25,6 @@ class TileArrayValue: _builder: TileArrayBuilder = field(repr=False) -@dataclass(frozen=True) -class InputPortBinding: - """A tensor slice bound to one TileArray input Port.""" - - input_des: TileArrayValue - input_src: TensorAccess - port: Port - - -@dataclass(frozen=True) -class OutputPortBinding: - """A TileArray value bound to a tensor output slice and Port.""" - - output_src: TileArrayValue - output_des: TensorAccess - port: Port - - # --------------------------------------------------------------- # Typed tile-array operations # --------------------------------------------------------------- @@ -59,6 +41,60 @@ class TileArrayOp: tile: Tile +@dataclass(frozen=True) +class LoadOp(TileArrayOp): + """Loads through an explicit address or a configured tensor access.""" + + source: TensorAccess | None = None + + def __post_init__(self) -> None: + """Validates the selected address form and result type.""" + if len(self.results) != 1: + raise ValueError("LoadOp requires exactly one result") + if self.source is None: + if len(self.operands) != 1: + raise ValueError("dynamic LoadOp requires one address operand") + if self.operands[0].dtype != DType.I32: + raise TypeError("dynamic addresses currently require i32") + else: + if self.operands: + raise ValueError("configured LoadOp has no address operand") + if self.results[0].dtype != self.source.dtype: + raise TypeError("LoadOp result type must match its source") + + @property + def addr(self) -> TileArrayValue | None: + """Returns the explicit address when the load is dynamic.""" + return self.operands[0] if self.operands else None + + +@dataclass(frozen=True) +class StoreOp(TileArrayOp): + """Stores through an explicit address or a configured tensor access.""" + + target: TensorAccess | None = None + + def __post_init__(self) -> None: + """Validates the value and selected address form.""" + if self.results: + raise ValueError("StoreOp produces no results") + if self.target is None: + if len(self.operands) != 2: + raise ValueError("dynamic StoreOp requires value and address operands") + if self.operands[1].dtype != DType.I32: + raise TypeError("dynamic addresses currently require i32") + else: + if len(self.operands) != 1: + raise ValueError("configured StoreOp requires one value operand") + if self.operands[0].dtype != self.target.dtype: + raise TypeError("StoreOp value type must match its target") + + @property + def addr(self) -> TileArrayValue | None: + """Returns the explicit address when the store is dynamic.""" + return self.operands[1] if len(self.operands) == 2 else None + + @dataclass(frozen=True) class ConstantOp(TileArrayOp): """A scalar constant produced by a tile-array operation.""" @@ -163,10 +199,7 @@ class TileArrayProgram: array: TileArray arguments: tuple[Tensor, ...] - input_ports: tuple[InputPortBinding, ...] - output_ports: tuple[OutputPortBinding, ...] operations: tuple[TileArrayOp, ...] - template_name: str | None stationary: StationaryBinding | None @@ -183,8 +216,6 @@ class TileArrayBuilder: def __init__(self, arguments: tuple[Tensor, ...] = ()): self._arguments = arguments self._array: TileArray | None = None - self._input_ports: list[InputPortBinding] = [] - self._output_ports: list[OutputPortBinding] = [] self._operations: list[TileArrayOp] = [] self._next_value_id = 0 self._token = None @@ -242,13 +273,10 @@ def emit( tile: Tile, create_operation: Callable[[tuple[TileArrayValue, ...]], TileArrayOp], ) -> tuple[TileArrayValue, ...]: - """Creates and record one tile-array operation.""" + """Creates and records one tile-array operation.""" self._ensure_not_built() self._bind_tile_array(tile.array) - if not result_dtypes: - raise ValueError("an operation must produce at least one result") - for operand in operands: self._validate_operand_for_builder(operand) @@ -277,49 +305,8 @@ def emit( self._next_value_id += len(results) return results - def add_input_port(self, *, source: TensorAccess, port: Port) -> TileArrayValue: - """Creates one scalar kernel input and bind it to a Port.""" - self._ensure_not_built() - if source.is_scalar: - raise ValueError("an input Port requires a tensor slice") - - if sum(isinstance(index, slice) for index in source.indices) != 1: - raise ValueError("an input Port currently supports one varying dimension") - - self._bind_tile_array(port.array) - - result = TileArrayValue( - id=self._next_value_id, dtype=source.dtype, _builder=self - ) - self._next_value_id += 1 - self._input_ports.append( - InputPortBinding(input_src=source, input_des=result, port=port) - ) - return result - - def add_output_port( - self, *, value: TileArrayValue, target: TensorAccess, port: Port - ) -> None: - """Bind one kernel result to an output slice and Port.""" - self._ensure_not_built() - self._validate_operand_for_builder(value) - - if target.is_scalar: - raise ValueError("an output Port requires a tensor slice") - if sum(isinstance(index, slice) for index in target.indices) != 1: - raise ValueError("an output Port currently supports one varying dimension") - - if target.dtype != value.dtype: - raise TypeError("output value and target must have the same dtype") - - self._bind_tile_array(port.array) - - self._output_ports.append( - OutputPortBinding(output_src=value, output_des=target, port=port) - ) - def build(self) -> TileArrayProgram: - """Finish recording and return a program.""" + """Finishes recording and returns a program.""" if self._token is not None: raise RuntimeError( "cannot build a TileArrayProgram while its builder is active" @@ -334,7 +321,6 @@ def build(self) -> TileArrayProgram: ] stationary = None - template_name = None if mac_operations: stationary_source = mac_operations[0].stationary_value.source @@ -351,18 +337,12 @@ def build(self) -> TileArrayProgram: for operation in mac_operations ), ) - # Configured MAC networks currently use the systolic template. - # Template registration will replace this inference later. - template_name = "systolic_array" self._is_built = True return TileArrayProgram( array=self._array, arguments=self._arguments, - input_ports=tuple(self._input_ports), - output_ports=tuple(self._output_ports), operations=tuple(self._operations), - template_name=template_name, stationary=stationary, ) @@ -376,7 +356,7 @@ def build(self) -> TileArrayProgram: def _require_active_builder() -> TileArrayBuilder: - """Return the active builder.""" + """Returns the active builder.""" builder = _active_builder.get() if builder is None: @@ -393,10 +373,10 @@ def _require_active_builder() -> TileArrayBuilder: def constant( value: int | float, *, tile: Tile, dtype: DType | None = None ) -> TileArrayValue: - """Create a scalar constant on one hardware tile. + """Creates 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: + An explicit dtype selects a different representation: constant(1.0, tile=tile, dtype=DType.F32) """ if dtype is None: @@ -430,7 +410,7 @@ def add( *, tile: Tile, ) -> TileArrayValue: - """Create a scalar addition on one hardware tile.""" + """Creates a scalar addition on one hardware tile.""" builder = _require_active_builder() operands = (lhs, rhs) @@ -447,25 +427,68 @@ def add( )[0] -def input_port(source: TensorAccess, *, port: Port) -> TileArrayValue: - """Reads one tensor slice through a TileArray input port.""" - if not isinstance(source, TensorAccess): - raise TypeError("input_port source must be a tensor access") - if not isinstance(port, Port): - raise TypeError("input_port port must be a Port") - return _require_active_builder().add_input_port(source=source, port=port) - - -def output_port(value: TileArrayValue, *, target: TensorAccess, port: Port) -> None: - """Write one TileArray value stream through an output Port.""" +def load( + source: TensorAccess | None = None, + *, + addr: TileArrayValue | None = None, + dtype: DType | None = None, + tile: Tile, +) -> TileArrayValue: + """Loads from a tensor access or an explicit target address. - if not isinstance(target, TensorAccess): - raise TypeError("output_port target must be a tensor access") + Configured accesses enumerate logical indices in lexicographic order, + with the last varying dimension advancing fastest. Dynamic addresses + already use the target address representation; they are not tensor indices. + """ + if (source is None) == (addr is None): + raise ValueError("load requires exactly one of source and addr") + if source is not None: + if not isinstance(source, TensorAccess): + raise TypeError("load source must be a TensorAccess") + if dtype is not None and dtype != source.dtype: + raise TypeError("load dtype must match its source") + dtype = source.dtype + if not isinstance(dtype, DType): + raise TypeError("dynamic load requires an explicit DType") + + operands = () if addr is None else (addr,) + return _require_active_builder().emit( + operands=operands, + result_dtypes=(dtype,), + tile=tile, + create_operation=lambda results: LoadOp( + results=results, + operands=operands, + tile=tile, + source=source, + ), + )[0] - if not isinstance(port, Port): - raise TypeError("output_port port must be a Port") - _require_active_builder().add_output_port(value=value, target=target, port=port) +def store( + value: TileArrayValue, + *, + target: TensorAccess | None = None, + addr: TileArrayValue | None = None, + tile: Tile, +) -> None: + """Stores a value through a tensor access or explicit target address.""" + if (target is None) == (addr is None): + raise ValueError("store requires exactly one of target and addr") + if target is not None and not isinstance(target, TensorAccess): + raise TypeError("store target must be a TensorAccess") + operands = (value,) if addr is None else (value, addr) + _require_active_builder().emit( + operands=operands, + result_dtypes=(), + tile=tile, + create_operation=lambda results: StoreOp( + results=results, + operands=operands, + tile=tile, + target=target, + ), + ) def mac( diff --git a/python/synapse/library/__init__.py b/python/synapse/library/__init__.py new file mode 100644 index 0000000..da19a4d --- /dev/null +++ b/python/synapse/library/__init__.py @@ -0,0 +1,5 @@ +"""Reusable programs written in the Synapse language.""" + +from .gemm import ws_gemm_3x3 + +__all__ = ["ws_gemm_3x3"] diff --git a/python/synapse/library/gemm.py b/python/synapse/library/gemm.py new file mode 100644 index 0000000..02704ee --- /dev/null +++ b/python/synapse/library/gemm.py @@ -0,0 +1,38 @@ +"""Explicit spatial implementations shared by direct programs and patterns.""" + +import synapse.language as synl +from synapse.language.tile_array_program import TileArrayValue + + +def ws_gemm_3x3(A: synl.Tensor, B: synl.Tensor, C: synl.Tensor): + """Computes C = A @ B using nine MACs and six memory tiles. + + C is overwritten and must not overlap either input buffer. The matching + patterns prove this precondition before selecting the implementation. + """ + if any(tensor.type != synl.i32[3, 3] for tensor in (A, B, C)): + raise ValueError("ws_gemm_3x3 requires three 3x3 i32 tensors") + + array = synl.TileArray(x_tiles=4, y_tiles=4) + partial_sums: list[TileArrayValue] = [] + + # Physical MAC rows run north to south while reduction indices increase. + for y in range(3, 0, -1): + k = 3 - y + flowing = synl.load(A[:, k], tile=array[0, y]) + next_partial_sums: list[TileArrayValue] = [] + + # The west column contains loads; MAC columns start at x=1. + for x in range(1, 4): + accumulated, flowing = synl.mac( + flowing, + partial_sums[x - 1] if partial_sums else None, + stationary=B[k, x - 1], + tile=array[x, y], + ) + next_partial_sums.append(accumulated) + partial_sums = next_partial_sums + + # Each south memory tile writes one output column. + for x, accumulated in enumerate(partial_sums, start=1): + synl.store(accumulated, target=C[:, x - 1], tile=array[x, 0]) diff --git a/python/synapse/patterns/__init__.py b/python/synapse/patterns/__init__.py new file mode 100644 index 0000000..79b0f3b --- /dev/null +++ b/python/synapse/patterns/__init__.py @@ -0,0 +1,5 @@ +"""User-defined rewrite patterns for Synapse programs.""" + +from .pattern import TileArrayRewritePattern + +__all__ = ["TileArrayRewritePattern"] diff --git a/python/synapse/templates/__init__.py b/python/synapse/templates/__init__.py deleted file mode 100644 index 3affb48..0000000 --- a/python/synapse/templates/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Reusable spatial templates provided by Synapse.""" diff --git a/python/synapse/templates/tile_array/__init__.py b/python/synapse/templates/tile_array/__init__.py deleted file mode 100644 index 66eb077..0000000 --- a/python/synapse/templates/tile_array/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -"""Templates for computation within one TileArray.""" - -from .gemm import ws_gemm_4x4 - -__all__ = [ - "ws_gemm_4x4", -] diff --git a/python/synapse/templates/tile_array/gemm.py b/python/synapse/templates/tile_array/gemm.py deleted file mode 100644 index 2520463..0000000 --- a/python/synapse/templates/tile_array/gemm.py +++ /dev/null @@ -1,46 +0,0 @@ -"""Reusable GEMM implementations authored with the Synapse language.""" - -import synapse.language as synl - - -def ws_gemm_4x4( - A: synl.Tensor, - B: synl.Tensor, - C: synl.Tensor, -): - """Describes a fixed 4x4 weight-stationary GEMM.""" - - array = synl.TileArray(x_tiles=4, y_tiles=4) - - partial_sums = [] - - for k in range(array.y_tiles): - y = array.y_tiles - 1 - k - - # A[:, k] enters from the west boundary ports and flows east across this row. - flowing = synl.input_port( - A[:, k], - port=array.west_ports[y], - ) - - next_partial_sums = [] - - for x in range(array.x_tiles): - accumulated, flowing = synl.mac( - flowing, - partial_sums[x] if partial_sums else None, - stationary=B[k, x], - tile=array[x, y], - ) - - next_partial_sums.append(accumulated) - - partial_sums = next_partial_sums - - for x, accumulated in enumerate(partial_sums): - # Each south Port writes one result column into C. - synl.output_port( - accumulated, - target=C[:, x], - port=array.south_ports[x], - ) diff --git a/tests/python/compiler/test_systolic_gemm.py b/tests/python/compiler/test_systolic_gemm.py index b6da5f9..46766ac 100644 --- a/tests/python/compiler/test_systolic_gemm.py +++ b/tests/python/compiler/test_systolic_gemm.py @@ -1,142 +1,139 @@ +"""Checks complete standalone GEMM IR and code generation.""" + +import subprocess +from pathlib import Path + import synapse import synapse.language as synl from synapse.frontend import lowering -from synapse.templates.tile_array import ws_gemm_4x4 +from synapse.library import ws_gemm_3x3 + +REPOSITORY_ROOT = Path(__file__).resolve().parents[3] +NEURA_ROOT = REPOSITORY_ROOT / "mlir/amoeba/thirdparty/neura" PRE_MAPPING_IR = """ -#map = affine_map<(d0) -> (d0, 0)> -#map1 = affine_map<(d0) -> (d0, 1)> -#map2 = affine_map<(d0) -> (d0, 2)> -#map3 = affine_map<(d0) -> (d0, 3)> -#map4 = affine_map<(d0, d1) -> (-d1 + 3, d0)> +#map = affine_map<(d0, d1) -> (-d1 + 3, d0 - 1)> module { - func.func @ws_gemm_4x4(%arg0: memref<4x4xi32>, %arg1: memref<4x4xi32>, %arg2: memref<4x4xi32>) -> memref<4x4xi32> { - %done_writes = taskflow.task @ws_gemm_4x4 will_reads(%arg0, %arg1 : memref<4x4xi32>, memref<4x4xi32>) will_writes(%arg2 : memref<4x4xi32>) [original_read_memrefs(%arg0, %arg1 : memref<4x4xi32>, memref<4x4xi32>), original_write_memrefs(%arg2 : memref<4x4xi32>)] : (memref<4x4xi32>, memref<4x4xi32>, memref<4x4xi32>) -> (memref<4x4xi32>) { - ^bb0(%arg3: memref<4x4xi32>, %arg4: memref<4x4xi32>, %arg5: memref<4x4xi32>): - %0:4 = taskflow.stream_read %arg3 maps [#map, #map1, #map2, #map3] : memref<4x4xi32> -> (i32, i32, i32, i32) - %1:4 = neura.kernel inputs(%0#0, %0#1, %0#2, %0#3, %arg4 : i32, i32, i32, i32, memref<4x4xi32>) attributes {accelerator = "neura", kernel_metadata = {kind = "template", template = {input_ports = [{direction = "west", kernel_input = 0 : i32, x = 0 : i32, y = 3 : i32}, {direction = "west", kernel_input = 1 : i32, x = 0 : i32, y = 2 : i32}, {direction = "west", kernel_input = 2 : i32, x = 0 : i32, y = 1 : i32}, {direction = "west", kernel_input = 3 : i32, x = 0 : i32, y = 0 : i32}], name = "systolic_array", output_ports = [{direction = "south", kernel_result = 0 : i32, x = 0 : i32, y = 0 : i32}, {direction = "south", kernel_result = 1 : i32, x = 1 : i32, y = 0 : i32}, {direction = "south", kernel_result = 2 : i32, x = 2 : i32, y = 0 : i32}, {direction = "south", kernel_result = 3 : i32, x = 3 : i32, y = 0 : i32}], stationary = {kernel_input = 4 : i32, map = #map4}}}} { - ^bb0(%arg6: i32, %arg7: i32, %arg8: i32, %arg9: i32, %arg10: memref<4x4xi32>): - %result, %forwarded = "neura.mac"(%arg6) {placement = {x = 0 : i32, y = 3 : i32}} : (i32) -> (i32, i32) - %result_0, %forwarded_1 = "neura.mac"(%forwarded) {placement = {x = 1 : i32, y = 3 : i32}} : (i32) -> (i32, i32) - %result_2, %forwarded_3 = "neura.mac"(%forwarded_1) {placement = {x = 2 : i32, y = 3 : i32}} : (i32) -> (i32, i32) - %result_4, %forwarded_5 = "neura.mac"(%forwarded_3) {placement = {x = 3 : i32, y = 3 : i32}} : (i32) -> (i32, i32) - %result_6, %forwarded_7 = "neura.mac"(%arg7, %result) {placement = {x = 0 : i32, y = 2 : i32}} : (i32, i32) -> (i32, i32) - %result_8, %forwarded_9 = "neura.mac"(%forwarded_7, %result_0) {placement = {x = 1 : i32, y = 2 : i32}} : (i32, i32) -> (i32, i32) - %result_10, %forwarded_11 = "neura.mac"(%forwarded_9, %result_2) {placement = {x = 2 : i32, y = 2 : i32}} : (i32, i32) -> (i32, i32) - %result_12, %forwarded_13 = "neura.mac"(%forwarded_11, %result_4) {placement = {x = 3 : i32, y = 2 : i32}} : (i32, i32) -> (i32, i32) - %result_14, %forwarded_15 = "neura.mac"(%arg8, %result_6) {placement = {x = 0 : i32, y = 1 : i32}} : (i32, i32) -> (i32, i32) - %result_16, %forwarded_17 = "neura.mac"(%forwarded_15, %result_8) {placement = {x = 1 : i32, y = 1 : i32}} : (i32, i32) -> (i32, i32) - %result_18, %forwarded_19 = "neura.mac"(%forwarded_17, %result_10) {placement = {x = 2 : i32, y = 1 : i32}} : (i32, i32) -> (i32, i32) - %result_20, %forwarded_21 = "neura.mac"(%forwarded_19, %result_12) {placement = {x = 3 : i32, y = 1 : i32}} : (i32, i32) -> (i32, i32) - %result_22, %forwarded_23 = "neura.mac"(%arg9, %result_14) {placement = {x = 0 : i32, y = 0 : i32}} : (i32, i32) -> (i32, i32) - %result_24, %forwarded_25 = "neura.mac"(%forwarded_23, %result_16) {placement = {x = 1 : i32, y = 0 : i32}} : (i32, i32) -> (i32, i32) - %result_26, %forwarded_27 = "neura.mac"(%forwarded_25, %result_18) {placement = {x = 2 : i32, y = 0 : i32}} : (i32, i32) -> (i32, i32) - %result_28, %forwarded_29 = "neura.mac"(%forwarded_27, %result_20) {placement = {x = 3 : i32, y = 0 : i32}} : (i32, i32) -> (i32, i32) - neura.yield results(%result_22, %result_24, %result_26, %result_28 : i32, i32, i32, i32) - } : i32, i32, i32, i32 - taskflow.stream_write(%1#0, %1#1, %1#2, %1#3 : i32, i32, i32, i32) to %arg5 maps [#map, #map1, #map2, #map3] : memref<4x4xi32> - taskflow.yield done_writes(%arg5 : memref<4x4xi32>) + func.func @ws_gemm_3x3(%arg0: memref<3x3xi32>, %arg1: memref<3x3xi32>, %arg2: memref<3x3xi32>) -> memref<3x3xi32> { + %done_writes = taskflow.task @ws_gemm_3x3 will_reads(%arg0, %arg1 : memref<3x3xi32>, memref<3x3xi32>) will_writes(%arg2 : memref<3x3xi32>) [original_read_memrefs(%arg0, %arg1 : memref<3x3xi32>, memref<3x3xi32>), original_write_memrefs(%arg2 : memref<3x3xi32>)] : (memref<3x3xi32>, memref<3x3xi32>, memref<3x3xi32>) -> (memref<3x3xi32>) { + ^bb0(%arg3: memref<3x3xi32>, %arg4: memref<3x3xi32>, %arg5: memref<3x3xi32>): + neura.kernel inputs(%arg3, %arg4, %arg5 : memref<3x3xi32>, memref<3x3xi32>, memref<3x3xi32>) attributes {accelerator = "neura", kernel_metadata = {kind = "template", template = {name = "systolic_array", stationary = {kernel_input = 1 : i32, map = #map}}}} { + ^bb0(%arg6: memref<3x3xi32>, %arg7: memref<3x3xi32>, %arg8: memref<3x3xi32>): + %0 = "neura.load"(%arg6) {constants = array, placement = {x = 0 : i32, y = 3 : i32}} : (memref<3x3xi32>) -> i32 + %result, %forwarded = "neura.mac"(%0) {placement = {x = 1 : i32, y = 3 : i32}} : (i32) -> (i32, i32) + %result_0, %forwarded_1 = "neura.mac"(%forwarded) {placement = {x = 2 : i32, y = 3 : i32}} : (i32) -> (i32, i32) + %result_2, %forwarded_3 = "neura.mac"(%forwarded_1) {placement = {x = 3 : i32, y = 3 : i32}} : (i32) -> (i32, i32) + %1 = "neura.load"(%arg6) {constants = array, placement = {x = 0 : i32, y = 2 : i32}} : (memref<3x3xi32>) -> i32 + %result_4, %forwarded_5 = "neura.mac"(%1, %result) {placement = {x = 1 : i32, y = 2 : i32}} : (i32, i32) -> (i32, i32) + %result_6, %forwarded_7 = "neura.mac"(%forwarded_5, %result_0) {placement = {x = 2 : i32, y = 2 : i32}} : (i32, i32) -> (i32, i32) + %result_8, %forwarded_9 = "neura.mac"(%forwarded_7, %result_2) {placement = {x = 3 : i32, y = 2 : i32}} : (i32, i32) -> (i32, i32) + %2 = "neura.load"(%arg6) {constants = array, placement = {x = 0 : i32, y = 1 : i32}} : (memref<3x3xi32>) -> i32 + %result_10, %forwarded_11 = "neura.mac"(%2, %result_4) {placement = {x = 1 : i32, y = 1 : i32}} : (i32, i32) -> (i32, i32) + %result_12, %forwarded_13 = "neura.mac"(%forwarded_11, %result_6) {placement = {x = 2 : i32, y = 1 : i32}} : (i32, i32) -> (i32, i32) + %result_14, %forwarded_15 = "neura.mac"(%forwarded_13, %result_8) {placement = {x = 3 : i32, y = 1 : i32}} : (i32, i32) -> (i32, i32) + "neura.store"(%result_10, %arg8) {constants = array, placement = {x = 1 : i32, y = 0 : i32}} : (i32, memref<3x3xi32>) -> () + "neura.store"(%result_12, %arg8) {constants = array, placement = {x = 2 : i32, y = 0 : i32}} : (i32, memref<3x3xi32>) -> () + "neura.store"(%result_14, %arg8) {constants = array, placement = {x = 3 : i32, y = 0 : i32}} : (i32, memref<3x3xi32>) -> () + neura.yield + } + taskflow.yield done_writes(%arg5 : memref<3x3xi32>) } - return %done_writes : memref<4x4xi32> + return %done_writes : memref<3x3xi32> } } """.strip() MAPPED_IR = """ -#map = affine_map<(d0) -> (d0, 0)> -#map1 = affine_map<(d0) -> (d0, 1)> -#map2 = affine_map<(d0) -> (d0, 2)> -#map3 = affine_map<(d0) -> (d0, 3)> -#map4 = affine_map<(d0, d1) -> (-d1 + 3, d0)> +#map = affine_map<(d0, d1) -> (-d1 + 3, d0 - 1)> module { - func.func @ws_gemm_4x4(%arg0: memref<4x4xi32>, %arg1: memref<4x4xi32>, %arg2: memref<4x4xi32>) -> memref<4x4xi32> { - %done_writes = taskflow.task @ws_gemm_4x4 will_reads(%arg0, %arg1 : memref<4x4xi32>, memref<4x4xi32>) will_writes(%arg2 : memref<4x4xi32>) [original_read_memrefs(%arg0, %arg1 : memref<4x4xi32>, memref<4x4xi32>), original_write_memrefs(%arg2 : memref<4x4xi32>)] : (memref<4x4xi32>, memref<4x4xi32>, memref<4x4xi32>) -> (memref<4x4xi32>) { - ^bb0(%arg3: memref<4x4xi32>, %arg4: memref<4x4xi32>, %arg5: memref<4x4xi32>): - %0:4 = taskflow.stream_read %arg3 maps [#map, #map1, #map2, #map3] : memref<4x4xi32> -> (i32, i32, i32, i32) - %1:4 = neura.kernel inputs(%0#0, %0#1, %0#2, %0#3, %arg4 : i32, i32, i32, i32, memref<4x4xi32>) attributes {accelerator = "neura", kernel_metadata = {kind = "template", template = {input_ports = [{direction = "west", kernel_input = 0 : i32, x = 0 : i32, y = 3 : i32}, {direction = "west", kernel_input = 1 : i32, x = 0 : i32, y = 2 : i32}, {direction = "west", kernel_input = 2 : i32, x = 0 : i32, y = 1 : i32}, {direction = "west", kernel_input = 3 : i32, x = 0 : i32, y = 0 : i32}], name = "systolic_array", output_ports = [{direction = "south", kernel_result = 0 : i32, x = 0 : i32, y = 0 : i32}, {direction = "south", kernel_result = 1 : i32, x = 1 : i32, y = 0 : i32}, {direction = "south", kernel_result = 2 : i32, x = 2 : i32, y = 0 : i32}, {direction = "south", kernel_result = 3 : i32, x = 3 : i32, y = 0 : i32}], stationary = {kernel_input = 4 : i32, map = #map4}}}, 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}} { - ^bb0(%arg6: !neura.data, %arg7: !neura.data, %arg8: !neura.data, %arg9: !neura.data, %arg10: !neura.data, i1>): - %2 = "neura.data_mov"(%arg6) {dfg_id = 0 : i32, mapping_locs = [{direction = "west", id = 20 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, io = "input", resource = "boundary_port", time_step = 0 : i32, x = 0 : i32, y = 3 : i32}]} : (!neura.data) -> !neura.data - %result, %forwarded = "neura.mac"(%2) {dfg_id = 4 : i32, mapping_locs = [{id = 12 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 0 : i32, x = 0 : i32, y = 3 : i32}]} : (!neura.data) -> (!neura.data, !neura.data) - %3 = "neura.data_mov"(%forwarded) {dfg_id = 6 : i32, mapping_locs = [{id = 38 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 0 : i32}]} : (!neura.data) -> !neura.data - %result_0, %forwarded_1 = "neura.mac"(%3) {dfg_id = 8 : i32, mapping_locs = [{id = 13 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 1 : i32, x = 1 : i32, y = 3 : i32}]} : (!neura.data) -> (!neura.data, !neura.data) - %4 = "neura.data_mov"(%forwarded_1) {dfg_id = 12 : i32, mapping_locs = [{id = 41 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data) -> !neura.data - %result_2, %forwarded_3 = "neura.mac"(%4) {dfg_id = 15 : i32, mapping_locs = [{id = 14 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 2 : i32, x = 2 : i32, y = 3 : i32}]} : (!neura.data) -> (!neura.data, !neura.data) - %5 = "neura.data_mov"(%forwarded_3) {dfg_id = 21 : i32, mapping_locs = [{id = 44 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data - %result_4, %forwarded_5 = "neura.mac"(%5) {dfg_id = 25 : i32, mapping_locs = [{id = 15 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "tile", time_step = 3 : i32, x = 3 : i32, y = 3 : i32}]} : (!neura.data) -> (!neura.data, !neura.data) - %6 = "neura.data_mov"(%arg7) {dfg_id = 1 : i32, mapping_locs = [{direction = "west", id = 16 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, io = "input", resource = "boundary_port", time_step = 1 : i32, x = 0 : i32, y = 2 : i32}]} : (!neura.data) -> !neura.data - %7 = "neura.data_mov"(%result) {dfg_id = 5 : i32, mapping_locs = [{id = 39 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 0 : i32}]} : (!neura.data) -> !neura.data - %result_6, %forwarded_7 = "neura.mac"(%6, %7) {dfg_id = 7 : i32, mapping_locs = [{id = 8 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 1 : i32, x = 0 : i32, y = 2 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) - %8 = "neura.data_mov"(%forwarded_7) {dfg_id = 10 : i32, mapping_locs = [{id = 24 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data) -> !neura.data - %9 = "neura.data_mov"(%result_0) {dfg_id = 11 : i32, mapping_locs = [{id = 42 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data) -> !neura.data - %result_8, %forwarded_9 = "neura.mac"(%8, %9) {dfg_id = 14 : i32, mapping_locs = [{id = 9 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 2 : i32, x = 1 : i32, y = 2 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) - %10 = "neura.data_mov"(%forwarded_9) {dfg_id = 19 : i32, mapping_locs = [{id = 28 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data - %11 = "neura.data_mov"(%result_2) {dfg_id = 20 : i32, mapping_locs = [{id = 45 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data - %result_10, %forwarded_11 = "neura.mac"(%10, %11) {dfg_id = 24 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "tile", time_step = 3 : i32, x = 2 : i32, y = 2 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) - %12 = "neura.data_mov"(%forwarded_11) {dfg_id = 31 : i32, mapping_locs = [{id = 32 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data - %13 = "neura.data_mov"(%result_4) {dfg_id = 32 : i32, mapping_locs = [{id = 47 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data - %result_12, %forwarded_13 = "neura.mac"(%12, %13) {dfg_id = 35 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "tile", time_step = 4 : i32, x = 3 : i32, y = 2 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) - %14 = "neura.data_mov"(%arg8) {dfg_id = 2 : i32, mapping_locs = [{direction = "west", id = 12 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, io = "input", resource = "boundary_port", time_step = 2 : i32, x = 0 : i32, y = 1 : i32}]} : (!neura.data) -> !neura.data - %15 = "neura.data_mov"(%result_6) {dfg_id = 9 : i32, mapping_locs = [{id = 25 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data) -> !neura.data - %result_14, %forwarded_15 = "neura.mac"(%14, %15) {dfg_id = 13 : i32, mapping_locs = [{id = 4 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 2 : i32, x = 0 : i32, y = 1 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) - %16 = "neura.data_mov"(%forwarded_15) {dfg_id = 17 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data - %17 = "neura.data_mov"(%result_8) {dfg_id = 18 : i32, mapping_locs = [{id = 29 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data - %result_16, %forwarded_17 = "neura.mac"(%16, %17) {dfg_id = 23 : i32, mapping_locs = [{id = 5 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "tile", time_step = 3 : i32, x = 1 : i32, y = 1 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) - %18 = "neura.data_mov"(%forwarded_17) {dfg_id = 29 : i32, mapping_locs = [{id = 14 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data - %19 = "neura.data_mov"(%result_10) {dfg_id = 30 : i32, mapping_locs = [{id = 33 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data - %result_18, %forwarded_19 = "neura.mac"(%18, %19) {dfg_id = 34 : i32, mapping_locs = [{id = 6 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "tile", time_step = 4 : i32, x = 2 : i32, y = 1 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) - %20 = "neura.data_mov"(%forwarded_19) {dfg_id = 39 : i32, mapping_locs = [{id = 18 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "link", time_step = 4 : i32}]} : (!neura.data) -> !neura.data - %21 = "neura.data_mov"(%result_12) {dfg_id = 40 : i32, mapping_locs = [{id = 36 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "link", time_step = 4 : i32}]} : (!neura.data) -> !neura.data - %result_20, %forwarded_21 = "neura.mac"(%20, %21) {dfg_id = 42 : i32, mapping_locs = [{id = 7 : i32, index_per_ii = 0 : i32, invalid_iterations = 5 : i32, resource = "tile", time_step = 5 : i32, x = 3 : i32, y = 1 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) - %22 = "neura.data_mov"(%arg9) {dfg_id = 3 : i32, mapping_locs = [{direction = "west", id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, io = "input", resource = "boundary_port", time_step = 3 : i32, x = 0 : i32, y = 0 : i32}]} : (!neura.data) -> !neura.data - %23 = "neura.data_mov"(%result_14) {dfg_id = 16 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data - %result_22, %forwarded_23 = "neura.mac"(%22, %23) {dfg_id = 22 : i32, mapping_locs = [{id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "tile", time_step = 3 : i32, x = 0 : i32, y = 0 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) - %24 = "neura.data_mov"(%forwarded_23) {dfg_id = 27 : i32, mapping_locs = [{id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data - %25 = "neura.data_mov"(%result_16) {dfg_id = 28 : i32, mapping_locs = [{id = 15 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data - %result_24, %forwarded_25 = "neura.mac"(%24, %25) {dfg_id = 33 : i32, mapping_locs = [{id = 1 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "tile", time_step = 4 : i32, x = 1 : i32, y = 0 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) - %26 = "neura.data_mov"(%forwarded_25) {dfg_id = 37 : i32, mapping_locs = [{id = 3 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "link", time_step = 4 : i32}]} : (!neura.data) -> !neura.data - %27 = "neura.data_mov"(%result_18) {dfg_id = 38 : i32, mapping_locs = [{id = 19 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "link", time_step = 4 : i32}]} : (!neura.data) -> !neura.data - %result_26, %forwarded_27 = "neura.mac"(%26, %27) {dfg_id = 41 : i32, mapping_locs = [{id = 2 : i32, index_per_ii = 0 : i32, invalid_iterations = 5 : i32, resource = "tile", time_step = 5 : i32, x = 2 : i32, y = 0 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) - %28 = "neura.data_mov"(%forwarded_27) {dfg_id = 44 : i32, mapping_locs = [{id = 6 : i32, index_per_ii = 0 : i32, invalid_iterations = 5 : i32, resource = "link", time_step = 5 : i32}]} : (!neura.data) -> !neura.data - %29 = "neura.data_mov"(%result_20) {dfg_id = 45 : i32, mapping_locs = [{id = 22 : i32, index_per_ii = 0 : i32, invalid_iterations = 5 : i32, resource = "link", time_step = 5 : i32}]} : (!neura.data) -> !neura.data - %result_28, %forwarded_29 = "neura.mac"(%28, %29) {dfg_id = 46 : i32, mapping_locs = [{id = 3 : i32, index_per_ii = 0 : i32, invalid_iterations = 6 : i32, resource = "tile", time_step = 6 : i32, x = 3 : i32, y = 0 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) - %30 = "neura.data_mov"(%result_22) {dfg_id = 26 : i32, mapping_locs = [{direction = "south", id = 3 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, io = "output", resource = "boundary_port", time_step = 3 : i32, x = 0 : i32, y = 0 : i32}]} : (!neura.data) -> !neura.data - %31 = "neura.data_mov"(%result_24) {dfg_id = 36 : i32, mapping_locs = [{direction = "south", id = 5 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, io = "output", resource = "boundary_port", time_step = 4 : i32, x = 1 : i32, y = 0 : i32}]} : (!neura.data) -> !neura.data - %32 = "neura.data_mov"(%result_26) {dfg_id = 43 : i32, mapping_locs = [{direction = "south", id = 7 : i32, index_per_ii = 0 : i32, invalid_iterations = 5 : i32, io = "output", resource = "boundary_port", time_step = 5 : i32, x = 2 : i32, y = 0 : i32}]} : (!neura.data) -> !neura.data - %33 = "neura.data_mov"(%result_28) {dfg_id = 47 : i32, mapping_locs = [{direction = "south", id = 11 : i32, index_per_ii = 0 : i32, invalid_iterations = 6 : i32, io = "output", resource = "boundary_port", time_step = 6 : i32, x = 3 : i32, y = 0 : i32}]} : (!neura.data) -> !neura.data - neura.yield results(%30, %31, %32, %33 : !neura.data, !neura.data, !neura.data, !neura.data) {dfg_id = 48 : i32} - } : i32, i32, i32, i32 - taskflow.stream_write(%1#0, %1#1, %1#2, %1#3 : i32, i32, i32, i32) to %arg5 maps [#map, #map1, #map2, #map3] : memref<4x4xi32> - taskflow.yield done_writes(%arg5 : memref<4x4xi32>) + func.func @ws_gemm_3x3(%arg0: memref<3x3xi32>, %arg1: memref<3x3xi32>, %arg2: memref<3x3xi32>) -> memref<3x3xi32> { + %done_writes = taskflow.task @ws_gemm_3x3 will_reads(%arg0, %arg1 : memref<3x3xi32>, memref<3x3xi32>) will_writes(%arg2 : memref<3x3xi32>) [original_read_memrefs(%arg0, %arg1 : memref<3x3xi32>, memref<3x3xi32>), original_write_memrefs(%arg2 : memref<3x3xi32>)] : (memref<3x3xi32>, memref<3x3xi32>, memref<3x3xi32>) -> (memref<3x3xi32>) { + ^bb0(%arg3: memref<3x3xi32>, %arg4: memref<3x3xi32>, %arg5: memref<3x3xi32>): + neura.kernel inputs(%arg3, %arg4, %arg5 : memref<3x3xi32>, memref<3x3xi32>, memref<3x3xi32>) attributes {accelerator = "neura", kernel_metadata = {kind = "template", template = {name = "systolic_array", stationary = {kernel_input = 1 : i32, map = #map}}}, 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}} { + ^bb0(%arg6: !neura.data, i1>, %arg7: !neura.data, i1>, %arg8: !neura.data, i1>): + %0 = "neura.load"(%arg6) {constants = array, dfg_id = 0 : i32, mapping_locs = [{id = 12 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 0 : i32, x = 0 : i32, y = 3 : i32}]} : (!neura.data, i1>) -> !neura.data + %1 = "neura.data_mov"(%0) {dfg_id = 4 : i32, mapping_locs = [{id = 38 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 0 : i32}]} : (!neura.data) -> !neura.data + %result, %forwarded = "neura.mac"(%1) {dfg_id = 7 : i32, mapping_locs = [{id = 13 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 1 : i32, x = 1 : i32, y = 3 : i32}]} : (!neura.data) -> (!neura.data, !neura.data) + %2 = "neura.data_mov"(%forwarded) {dfg_id = 9 : i32, mapping_locs = [{id = 41 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data) -> !neura.data + %result_0, %forwarded_1 = "neura.mac"(%2) {dfg_id = 11 : i32, mapping_locs = [{id = 14 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 2 : i32, x = 2 : i32, y = 3 : i32}]} : (!neura.data) -> (!neura.data, !neura.data) + %3 = "neura.data_mov"(%forwarded_1) {dfg_id = 15 : i32, mapping_locs = [{id = 44 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data + %result_2, %forwarded_3 = "neura.mac"(%3) {dfg_id = 18 : i32, mapping_locs = [{id = 15 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "tile", time_step = 3 : i32, x = 3 : i32, y = 3 : i32}]} : (!neura.data) -> (!neura.data, !neura.data) + %4 = "neura.load"(%arg6) {constants = array, dfg_id = 1 : i32, mapping_locs = [{id = 8 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 1 : i32, x = 0 : i32, y = 2 : i32}]} : (!neura.data, i1>) -> !neura.data + %5 = "neura.data_mov"(%4) {dfg_id = 5 : i32, mapping_locs = [{id = 24 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data) -> !neura.data + %6 = "neura.data_mov"(%result) {dfg_id = 8 : i32, mapping_locs = [{id = 42 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data) -> !neura.data + %result_4, %forwarded_5 = "neura.mac"(%5, %6) {dfg_id = 10 : i32, mapping_locs = [{id = 9 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 2 : i32, x = 1 : i32, y = 2 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) + %7 = "neura.data_mov"(%forwarded_5) {dfg_id = 13 : i32, mapping_locs = [{id = 28 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data + %8 = "neura.data_mov"(%result_0) {dfg_id = 14 : i32, mapping_locs = [{id = 45 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data + %result_6, %forwarded_7 = "neura.mac"(%7, %8) {dfg_id = 17 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "tile", time_step = 3 : i32, x = 2 : i32, y = 2 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) + %9 = "neura.data_mov"(%forwarded_7) {dfg_id = 22 : i32, mapping_locs = [{id = 32 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data + %10 = "neura.data_mov"(%result_2) {dfg_id = 23 : i32, mapping_locs = [{id = 47 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data + %result_8, %forwarded_9 = "neura.mac"(%9, %10) {dfg_id = 26 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "tile", time_step = 4 : i32, x = 3 : i32, y = 2 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) + %11 = "neura.load"(%arg6) {constants = array, dfg_id = 2 : i32, mapping_locs = [{id = 4 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 2 : i32, x = 0 : i32, y = 1 : i32}]} : (!neura.data, i1>) -> !neura.data + %12 = "neura.data_mov"(%11) {dfg_id = 6 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data + %13 = "neura.data_mov"(%result_4) {dfg_id = 12 : i32, mapping_locs = [{id = 29 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "link", time_step = 2 : i32}]} : (!neura.data) -> !neura.data + %result_10, %forwarded_11 = "neura.mac"(%12, %13) {dfg_id = 16 : i32, mapping_locs = [{id = 5 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "tile", time_step = 3 : i32, x = 1 : i32, y = 1 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) + %14 = "neura.data_mov"(%forwarded_11) {dfg_id = 20 : i32, mapping_locs = [{id = 14 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data + %15 = "neura.data_mov"(%result_6) {dfg_id = 21 : i32, mapping_locs = [{id = 33 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data + %result_12, %forwarded_13 = "neura.mac"(%14, %15) {dfg_id = 25 : i32, mapping_locs = [{id = 6 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "tile", time_step = 4 : i32, x = 2 : i32, y = 1 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) + %16 = "neura.data_mov"(%forwarded_13) {dfg_id = 28 : i32, mapping_locs = [{id = 18 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "link", time_step = 4 : i32}]} : (!neura.data) -> !neura.data + %17 = "neura.data_mov"(%result_8) {dfg_id = 29 : i32, mapping_locs = [{id = 36 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "link", time_step = 4 : i32}]} : (!neura.data) -> !neura.data + %result_14, %forwarded_15 = "neura.mac"(%16, %17) {dfg_id = 31 : i32, mapping_locs = [{id = 7 : i32, index_per_ii = 0 : i32, invalid_iterations = 5 : i32, resource = "tile", time_step = 5 : i32, x = 3 : i32, y = 1 : i32}]} : (!neura.data, !neura.data) -> (!neura.data, !neura.data) + %18 = "neura.data_mov"(%result_10) {dfg_id = 19 : i32, mapping_locs = [{id = 15 : i32, index_per_ii = 0 : i32, invalid_iterations = 3 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data) -> !neura.data + "neura.store"(%18, %arg8) {constants = array, dfg_id = 24 : i32, mapping_locs = [{id = 1 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "tile", time_step = 4 : i32, x = 1 : i32, y = 0 : i32}]} : (!neura.data, !neura.data, i1>) -> () + %19 = "neura.data_mov"(%result_12) {dfg_id = 27 : i32, mapping_locs = [{id = 19 : i32, index_per_ii = 0 : i32, invalid_iterations = 4 : i32, resource = "link", time_step = 4 : i32}]} : (!neura.data) -> !neura.data + "neura.store"(%19, %arg8) {constants = array, dfg_id = 30 : i32, mapping_locs = [{id = 2 : i32, index_per_ii = 0 : i32, invalid_iterations = 5 : i32, resource = "tile", time_step = 5 : i32, x = 2 : i32, y = 0 : i32}]} : (!neura.data, !neura.data, i1>) -> () + %20 = "neura.data_mov"(%result_14) {dfg_id = 32 : i32, mapping_locs = [{id = 22 : i32, index_per_ii = 0 : i32, invalid_iterations = 5 : i32, resource = "link", time_step = 5 : i32}]} : (!neura.data) -> !neura.data + "neura.store"(%20, %arg8) {constants = array, dfg_id = 33 : i32, mapping_locs = [{id = 3 : i32, index_per_ii = 0 : i32, invalid_iterations = 6 : i32, resource = "tile", time_step = 6 : i32, x = 3 : i32, y = 0 : i32}]} : (!neura.data, !neura.data, i1>) -> () + neura.yield {dfg_id = 3 : i32} + } + taskflow.yield done_writes(%arg5 : memref<3x3xi32>) } - return %done_writes : memref<4x4xi32> + return %done_writes : memref<3x3xi32> } } """.strip() def test_lowers_systolic_gemm_to_exact_pre_mapping_ir(): - actual = lowering.lower( - ws_gemm_4x4, - argument_types=( - synl.i32[4, 4], - synl.i32[4, 4], - synl.i32[4, 4], - ), - ) - + actual = lowering.lower(ws_gemm_3x3, argument_types=(synl.i32[3, 3],) * 3) assert actual.strip() == PRE_MAPPING_IR def test_compiles_systolic_gemm_to_exact_mapped_ir(): actual = synapse.compile( - ws_gemm_4x4, - target="neura", - argument_types=( - synl.i32[4, 4], - synl.i32[4, 4], - synl.i32[4, 4], - ), + ws_gemm_3x3, target="neura", argument_types=(synl.i32[3, 3],) * 3 ) - assert actual.strip() == MAPPED_IR + + +def test_generates_configured_memory_and_mac_instructions(tmp_path): + mapped = synapse.compile( + ws_gemm_3x3, target="neura", argument_types=(synl.i32[3, 3],) * 3 + ) + command = [ + str(REPOSITORY_ROOT / "build/amoeba/tools/mlir-amoeba-opt/mlir-amoeba-opt"), + f"--neura-architecture-spec={NEURA_ROOT / 'test/arch_spec/architecture.yaml'}", + "--generate-code", + "-o", + str(tmp_path / "mapped.mlir"), + ] + completed = subprocess.run( + command, input=mapped, text=True, capture_output=True, cwd=tmp_path + ) + assert completed.returncode == 0, completed.stderr + assembly = (tmp_path / "tmp-generated-instructions.asm").read_text() + assert assembly.count(" LOAD,") == 3 + assert assembly.count(" STORE,") == 3 + assert assembly.count(" MUL_ADD,") == 6 + assert assembly.count(" MUL,") == 3 + assert "address(arg0, 0, 3, 6)" in assembly + assert "address(arg2, 2, 5, 8)" in assembly + assert "value(arg1, 8)" in assembly + + invalid = mapped.replace("array", "array") + completed = subprocess.run( + command, input=invalid, text=True, capture_output=True, cwd=tmp_path + ) + assert completed.returncode != 0 + assert "constant element offset is out of bounds" in completed.stderr diff --git a/tests/python/frontend/test_memory_lowering.py b/tests/python/frontend/test_memory_lowering.py new file mode 100644 index 0000000..36a3ff1 --- /dev/null +++ b/tests/python/frontend/test_memory_lowering.py @@ -0,0 +1,124 @@ +"""Checks configured offset semantics and explicit address lowering.""" + +import pytest +import synapse.language as synl +from synapse.frontend.lowering import ( + TileArrayProgramLowering, + build_tile_array_program, + lower, +) +from synapse.language.tile_array_program import LoadOp +from taskflow_mlir.dialects import neura, taskflow +from taskflow_mlir.ir import Context, DenseI64ArrayAttr, Location, Module + + +def configured_accesses(A: synl.Tensor): + array = synl.TileArray(4, 4) + synl.load(A[1, 2], tile=array[0, 1]) + synl.load(A[:, ::2], tile=array[0, 2]) + synl.load(A[::-1, 1], tile=array[0, 3]) + + +def dynamic_accesses(): + array = synl.TileArray(4, 4) + addr = synl.constant(16, tile=array[1, 1]) + value = synl.load(addr=addr, dtype=synl.i32, tile=array[0, 1]) + synl.store(value, addr=addr, tile=array[1, 0]) + + +def test_static_offsets_cover_scalar_strided_and_multidimensional_accesses(): + program = build_tile_array_program( + configured_accesses, argument_types=(synl.i32[3, 4],) + ) + with Context(), Location.unknown(): + lowering = TileArrayProgramLowering(program) + offsets = [] + for operation in program.operations: + assert isinstance(operation, LoadOp) + assert operation.source is not None + offsets.append(tuple(lowering.get_memory_offsets(operation.source))) + assert offsets == [(6,), (0, 2, 4, 6, 8, 10), (9, 5, 1)] + + +def test_dynamic_ir_uses_address_operands_without_memory_configuration(): + source = lower(dynamic_accesses) + with Context(), Location.unknown(): + taskflow.register_dialect() + neura.register_dialect() + module = Module.parse(source) + task = module.body.operations[0].regions[0].blocks[0].operations[0] + kernel = task.regions[0].blocks[0].operations[0] + constant, read, write, _ = tuple(kernel.regions[0].blocks[0].operations) + assert tuple(read.operands[index] for index in range(len(read.operands))) == ( + constant.results[0], + ) + assert tuple(write.operands[index] for index in range(len(write.operands))) == ( + read.results[0], + constant.results[0], + ) + assert len(kernel.results) == 0 + assert module.operation.verify() + assert "memory_access" not in source + + +def test_empty_configured_queue_is_rejected(): + def empty(A: synl.Tensor): + synl.load(A[0:0, 0], tile=synl.TileArray(4, 4)[0, 1]) + + with pytest.raises(ValueError, match="cannot be empty"): + lower(empty, argument_types=(synl.i32[3, 3],)) + + +def test_dynamic_addresses_reach_backend_mapping(): + import synapse + + mapped = synapse.compile(dynamic_accesses, target="neura") + assert '"neura.load"' in mapped + assert '"neura.store"' in mapped + assert "mapping_locs" in mapped + assert "memory_access" not in mapped + + +def test_memory_bases_survive_task_argument_grouping(): + def program(A: synl.Tensor, unused: synl.Tensor, C: synl.Tensor): + array = synl.TileArray(4, 4) + value = synl.load(A[:, 0], tile=array[0, 1]) + synl.store(value, target=C[:, 0], tile=array[1, 0]) + + source = lower(program, argument_types=(synl.i32[3, 3],) * 3) + with Context(), Location.unknown(): + taskflow.register_dialect() + neura.register_dialect() + module = Module.parse(source) + task = module.body.operations[0].regions[0].blocks[0].operations[0] + block = task.regions[0].blocks[0] + kernel = block.operations[0] + assert isinstance(kernel, neura.KernelOp) + assert tuple( + kernel.operands[index] for index in range(len(kernel.operands)) + ) == ( + block.arguments[0], + block.arguments[2], + block.arguments[1], + ) + kernel_block = kernel.regions[0].blocks[0] + read, store, _ = tuple(kernel_block.operations) + assert tuple(read.operands[index] for index in range(len(read.operands))) == ( + kernel_block.arguments[0], + ) + assert tuple(store.operands[index] for index in range(len(store.operands))) == ( + read.results[0], + kernel_block.arguments[2], + ) + assert tuple(DenseI64ArrayAttr(read.operation.attributes["constants"])) == ( + 0, + 3, + 6, + ) + assert tuple(DenseI64ArrayAttr(store.operation.attributes["constants"])) == ( + 0, + 3, + 6, + ) + assert "memory_access" not in str(kernel) + assert module.operation.verify() diff --git a/tests/python/language/test_memory.py b/tests/python/language/test_memory.py new file mode 100644 index 0000000..d35960c --- /dev/null +++ b/tests/python/language/test_memory.py @@ -0,0 +1,73 @@ +"""Checks general memory operations independently of Neura bindings.""" + +import pytest +import synapse.language as synl +from synapse.frontend.lowering import build_tile_array_program +from synapse.language.tile_array_program import LoadOp, StoreOp, TileArrayBuilder +from synapse.library import ws_gemm_3x3 + + +def test_dynamic_addresses_are_explicit_operands(): + array = synl.TileArray(4, 4) + builder = TileArrayBuilder() + with builder: + addr = synl.constant(16, tile=array[1, 1]) + value = synl.load(addr=addr, dtype=synl.f32, tile=array[0, 1]) + synl.store(value, addr=addr, tile=array[1, 0]) + _, read, write = builder.build().operations + assert isinstance(read, LoadOp) + assert read.addr is addr + assert read.operands == (addr,) + assert read.source is None + assert isinstance(write, StoreOp) + assert write.operands == (value, addr) + assert write.addr is addr + assert write.results == () + + +def test_memory_configuration_retains_static_accesses(): + tensor = synl.Tensor("A", synl.i32[3, 4]) + array = synl.TileArray(4, 4) + builder = TileArrayBuilder((tensor,)) + with builder: + value = synl.load(tensor[1, 2], tile=array[0, 1]) + synl.store(value, target=tensor[:, ::2], tile=array[1, 0]) + read, write = builder.build().operations + assert isinstance(read, LoadOp) + assert isinstance(write, StoreOp) + assert read.source is not None + assert write.target is not None + assert read.source.indices == (1, 2) + assert read.operands == () + assert write.target.indices == (slice(None), slice(None, None, 2)) + assert write.operands == (value,) + + +def test_rejected_memory_forms_do_not_consume_value_ids(): + tensor = synl.Tensor("A", synl.i32[3, 3]) + array = synl.TileArray(4, 4) + builder = TileArrayBuilder((tensor,)) + with builder: + addr = synl.constant(0, tile=array[1, 1]) + with pytest.raises(ValueError, match="exactly one"): + synl.load(tensor[:, 0], addr=addr, tile=array[0, 1]) + with pytest.raises(TypeError, match="explicit DType"): + synl.load(addr=addr, tile=array[0, 1]) + with pytest.raises(TypeError, match="must match"): + synl.load(tensor[:, 0], dtype=synl.f32, tile=array[0, 1]) + value = synl.load(addr=addr, dtype=synl.i32, tile=array[0, 1]) + assert value.id == 1 + with pytest.raises(ValueError, match="exactly one"): + synl.store(value, target=tensor[:, 0], addr=addr, tile=array[1, 0]) + assert len(builder.build().operations) == 2 + + +def test_gemm_program_records_load_mac_store_network(): + program = build_tile_array_program( + ws_gemm_3x3, argument_types=(synl.i32[3, 3],) * 3 + ) + assert sum(isinstance(op, LoadOp) for op in program.operations) == 3 + assert sum(isinstance(op, StoreOp) for op in program.operations) == 3 + assert len(program.operations) == 15 + assert program.stationary is not None + assert program.stationary.source is program.arguments[1] diff --git a/tests/python/language/test_spatial.py b/tests/python/language/test_spatial.py index 52cbc7a..d22bbbb 100644 --- a/tests/python/language/test_spatial.py +++ b/tests/python/language/test_spatial.py @@ -35,33 +35,3 @@ def test_access_tile_array_by_coordinate(): ) assert tile is enumerated_tile - - -def test_tile_array_exposes_boundary_ports(): - array = synl.TileArray(x_tiles=2, y_tiles=3) - - assert [(port.direction, port.x, port.y) for port in array.west_ports] == [ - ("west", 0, 0), - ("west", 0, 1), - ("west", 0, 2), - ] - - assert [(port.direction, port.x, port.y) for port in array.east_ports] == [ - ("east", 1, 0), - ("east", 1, 1), - ("east", 1, 2), - ] - - assert [(port.direction, port.x, port.y) for port in array.north_ports] == [ - ("north", 0, 2), - ("north", 1, 2), - ] - - assert [(port.direction, port.x, port.y) for port in array.south_ports] == [ - ("south", 0, 0), - ("south", 1, 0), - ] - - assert array.west_ports[2].array is array - assert array.west_ports[2].tile is array[0, 2] - assert array.south_ports[1].tile is array[1, 0] From 99d2def95f906c7888c7a2350f691c75b3476a91 Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Sat, 12 Sep 2026 00:51:14 +0800 Subject: [PATCH 15/19] Add TileArray pattern rewriting --- python/synapse/compiler/compiler.py | 64 +++- python/synapse/compiler/pattern_rewriter.py | 258 ++++++++++++++++ python/synapse/patterns/gemm_pattern.py | 278 ++++++++++++++++++ python/synapse/patterns/pattern.py | 28 ++ tests/python/compiler/test_gemm_patterns.py | 223 ++++++++++++++ .../python/compiler/test_pattern_rewriter.py | 182 ++++++++++++ 6 files changed, 1025 insertions(+), 8 deletions(-) create mode 100644 python/synapse/compiler/pattern_rewriter.py create mode 100644 python/synapse/patterns/gemm_pattern.py create mode 100644 python/synapse/patterns/pattern.py create mode 100644 tests/python/compiler/test_gemm_patterns.py create mode 100644 tests/python/compiler/test_pattern_rewriter.py diff --git a/python/synapse/compiler/compiler.py b/python/synapse/compiler/compiler.py index 5492a41..142ee28 100644 --- a/python/synapse/compiler/compiler.py +++ b/python/synapse/compiler/compiler.py @@ -1,33 +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, + program: Callable | str, *, target: str, argument_types: tuple[TensorType, ...] = (), + patterns: Sequence[type[TileArrayRewritePattern]] | None = None, ) -> str: - """Compile a Synapse program for the selected backend.""" + """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, argument_types=argument_types) + 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 = ( @@ -42,11 +77,23 @@ 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", @@ -65,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: diff --git a/python/synapse/compiler/pattern_rewriter.py b/python/synapse/compiler/pattern_rewriter.py new file mode 100644 index 0000000..d0e6a43 --- /dev/null +++ b/python/synapse/compiler/pattern_rewriter.py @@ -0,0 +1,258 @@ +"""Applies user-defined rewrite patterns to MLIR modules.""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence + +from taskflow_mlir.ir import ( + BlockArgument, + F32Type, + InsertionPoint, + IntegerType, + MemRefType, + Module, + OpResult, + OpView, + Value, +) + +from synapse.language.types import DType, TensorType +from synapse.patterns import TileArrayRewritePattern + + +class PatternRewriter: + """Mutates IR on behalf of one successfully matched pattern.""" + + def __init__(self, root: OpView): + self._root = root + + @property + def ip(self) -> InsertionPoint: + """Returns an insertion point immediately before the pattern root.""" + + return InsertionPoint(self._root) + + def erase_op(self, operation: OpView) -> None: + """Erases an operation that has no live results.""" + + for index in range(len(operation.results)): + if any(operation.results[index].uses): + raise ValueError("cannot erase an operation with live results") + + operation.erase() + + def replace_op(self, operation: OpView, replacement: OpView) -> None: + """Replaces an operation and redirects its SSA results.""" + + old_results = tuple( + operation.results[index] for index in range(len(operation.results)) + ) + new_results = tuple( + replacement.results[index] for index in range(len(replacement.results)) + ) + + if len(old_results) != len(new_results): + raise ValueError("replacement must produce the same number of results") + + if any(old.type != new.type for old, new in zip(old_results, new_results)): + raise TypeError("replacement must preserve result types") + + for old_result, new_result in zip( + old_results, + new_results, + ): + old_result.replace_all_uses_with(new_result) + + operation.erase() + + def replace_with_tile_array( + self, + operation: OpView, + *, + program: Callable, + arguments: tuple[Value, ...], + ) -> bool: + """Replaces a compatible buffer computation with a TileArray kernel. + + The pattern establishes computation semantics. Shared compiler checks + establish type, task, and memory compatibility. An unsuitable candidate + returns False without mutation; invalid implementations raise errors. + Staging validates the generated kernel before the source IR is changed. + """ + from taskflow_mlir.dialects import func + + from synapse.frontend.lowering import ( + TileArrayProgramLowering, + build_tile_array_program, + ) + + if operation.operation != self._root.operation: + raise ValueError("TileArray replacement requires the pattern root") + inferred_types = _task_argument_types(operation, arguments) + if inferred_types is None: + return False + + tile_program = build_tile_array_program(program, argument_types=inferred_types) + lowering = TileArrayProgramLowering(tile_program) + if not _replacement_memory_is_legal(operation, arguments, lowering): + return False + + staged = Module.create() + with InsertionPoint(staged.body): + function = func.FuncOp( + "replacement", ([value.type for value in arguments], []) + ) + block = function.add_entry_block() + with InsertionPoint(block): + lowering.lower_to_kernel(dict(zip(tile_program.arguments, block.arguments))) + func.ReturnOp([]) + if not staged.operation.verify(): + raise ValueError("replacement kernel failed verification") + + kernel = block.operations[0] + for index, value in enumerate(arguments): + kernel.operation.operands[index] = value + kernel.operation.move_before(operation.operation) + operation.erase() + return True + + +def apply_patterns( + module: Module, + patterns: Sequence[type[TileArrayRewritePattern]], +) -> int: + """Walks a module and applies the first matching pattern at each operation.""" + + return _apply_patterns( + module.operation.opview, + patterns, + ) + + +def _apply_patterns( + operation: OpView, + patterns: Sequence[type[TileArrayRewritePattern]], +) -> int: + """Applies patterns recursively, stopping below a replaced operation.""" + + for pattern in patterns: + if not isinstance(operation, pattern.root): + continue + + rewriter = PatternRewriter(operation) + + if pattern.match_and_rewrite(operation, rewriter): + return 1 + + rewrite_count = 0 + + for region in operation.regions: + for block in region.blocks: + for nested_operation in tuple(block.operations): + rewrite_count += _apply_patterns( + nested_operation, + patterns, + ) + + return rewrite_count + + +def _base_buffer(value): + """Traces task captures and view-like operations to their memory origin.""" + while True: + if BlockArgument.isinstance(value): + argument = BlockArgument(value) + parent = argument.owner.owner.operation + if parent.name == "taskflow.task": + value = parent.operands[argument.arg_number] + continue + return value + if not OpResult.isinstance(value): + return value + producer = OpResult(value).owner + if producer.name in ( + "memref.cast", + "memref.subview", + "memref.reinterpret_cast", + ): + value = producer.operands[0] + continue + return value + + +def _disjoint_buffers(lhs, rhs): + """Proves disjointness for fresh allocations and incoming function buffers.""" + lhs, rhs = _base_buffer(lhs), _base_buffer(rhs) + if lhs == rhs: + return False + + def is_allocation(value): + return OpResult.isinstance(value) and OpResult(value).owner.name in ( + "memref.alloc", + "memref.alloca", + ) + + def is_function_argument(value): + return ( + BlockArgument.isinstance(value) + and BlockArgument(value).owner.owner.operation.name == "func.func" + ) + + return ( + is_allocation(lhs) and (is_allocation(rhs) or is_function_argument(rhs)) + ) or (is_allocation(rhs) and is_function_argument(lhs)) + + +def _task_argument_types(operation, arguments): + """Returns supported capture types, or None when the task boundary is unsuitable.""" + parent = operation.operation.parent + if parent is None or parent.name != "taskflow.task" or len(operation.results): + return None + block = parent.regions[0].blocks[0] + types = [] + for value in arguments: + if ( + not BlockArgument.isinstance(value) + or BlockArgument(value).owner != block + or not MemRefType.isinstance(value.type) + ): + return None + memref = MemRefType(value.type) + if not memref.has_static_shape or any(size <= 0 for size in memref.shape): + return None + if memref.element_type == IntegerType.get_signless(32): + dtype = DType.I32 + elif memref.element_type == F32Type.get(): + dtype = DType.F32 + else: + return None + if memref != MemRefType.get(list(memref.shape), memref.element_type): + return None + types.append(TensorType(tuple(memref.shape), dtype)) + return tuple(types) + + +def _replacement_memory_is_legal(operation, arguments, lowering): + """Checks inferred implementation effects against task declarations and aliasing.""" + if lowering.has_dynamic_memory: + return False + task = operation.operation.parent.opview + block_arguments = tuple(task.body.blocks[0].arguments) + read_count = len(task.will_reads) + write_count = len(task.will_writes) + declared_reads = block_arguments[:read_count] + declared_writes = block_arguments[read_count : read_count + write_count] + values = dict(zip(lowering.program.arguments, arguments)) + reads = lowering.read_arguments + writes = lowering.write_arguments + if any(values[argument] not in declared_reads for argument in reads): + return False + if any(values[argument] not in declared_writes for argument in writes): + return False + for output in writes: + for other in reads + writes: + if output is other: + continue + if not _disjoint_buffers(values[output], values[other]): + return False + return True diff --git a/python/synapse/patterns/gemm_pattern.py b/python/synapse/patterns/gemm_pattern.py new file mode 100644 index 0000000..3a58859 --- /dev/null +++ b/python/synapse/patterns/gemm_pattern.py @@ -0,0 +1,278 @@ +"""Matches GEMM source forms to the weight-stationary implementation.""" + +from typing import Protocol, cast + +from taskflow_mlir.dialects import affine, linalg +from taskflow_mlir.ir import ( + AffineDimExpr, + AffineMap, + AffineMapAttr, + ArrayAttr, + Attribute, + IntegerAttr, + IntegerType, + MemRefType, + OpResult, +) + +from synapse.library import ws_gemm_3x3 +from synapse.patterns import TileArrayRewritePattern + + +class _AffineMapAttrValue(Protocol): + """Describes the value property omitted by the pinned MLIR attribute stub.""" + + @property + def value(self) -> AffineMap: ... + + +def _loop(operation, extent): + """Recognizes one canonical GEMM loop over the required extent.""" + if operation.operation.name != "affine.for" or len(operation.operands): + return None + attrs = operation.operation.attributes + if ( + attrs["lowerBoundMap"] != AffineMapAttr.get(AffineMap.get_constant(0)) + or attrs["upperBoundMap"] != AffineMapAttr.get(AffineMap.get_constant(extent)) + or IntegerAttr(attrs["step"]).value != 1 + ): + return None + block = operation.regions[0].blocks[0] + return block if len(block.arguments) == 1 and not len(operation.results) else None + + +def _indices(operation, operands): + """Resolves the projected indices of a GEMM load or store.""" + access_map = cast( + _AffineMapAttrValue, AffineMapAttr(operation.operation.attributes["map"]) + ).value + if access_map.n_symbols or access_map.n_dims != len(operands): + return None + if any(not AffineDimExpr.isinstance(expr) for expr in access_map.results): + return None + return tuple(operands[AffineDimExpr(expr).position] for expr in access_map.results) + + +def _zero(value): + """Recognizes an integer zero used by GEMM initialization.""" + if not OpResult.isinstance(value): + return False + owner = OpResult(value).owner + return ( + owner.name == "arith.constant" + and IntegerAttr.isinstance(owner.attributes["value"]) + and IntegerAttr(owner.attributes["value"]).value == 0 + ) + + +def _has_zero_output(operation, output): + """Checks the overwrite template's precondition on the initial accumulator. + + 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 + previous = None + found = False + for region in parent.regions: + for block in region.blocks: + previous = None + for sibling in block.operations: + if sibling.operation == operation.operation: + found = True + break + if sibling.operation.name != "arith.constant": + previous = sibling + if found: + break + if found: + break + if previous is None: + return False + if isinstance(previous, linalg.FillOp): + return tuple(previous.outputs) == (output,) and _zero(previous.inputs[0]) + + rows, columns = MemRefType(output.type).shape + outer = _loop(previous, rows) + if outer is None or len(tuple(outer.operations)) != 2: + return False + inner = _loop(outer.operations[0], columns) + if inner is None or len(tuple(inner.operations)) != 2: + return False + store = inner.operations[0] + return ( + store.operation.name == "affine.store" + and store.operands[1] == output + and _zero(store.operands[0]) + and _indices( + store, + tuple(store.operands[index] for index in range(len(store.operands)))[2:], + ) + == (outer.arguments[0], inner.arguments[0]) + ) + + +class LinalgGemmPattern(TileArrayRewritePattern): + """Matches named matmul using its declared semantics and operand types.""" + + root = linalg.MatmulOp + + @classmethod + def match_and_rewrite(cls, operation, rewriter) -> bool: + """Checks the selected matmul and replaces it with the WS kernel.""" + # The driver has already checked the declared root operation type. + operation = cast(linalg.MatmulOp, operation) + arguments = tuple(operation.inputs) + tuple(operation.outputs) + expected = MemRefType.get([3, 3], IntegerType.get_signless(32)) + if len(operation.results) or len(arguments) != 3: + return False + if any(value.type != expected for value in arguments): + return False + if not _has_zero_output(operation, arguments[2]): + return False + return rewriter.replace_with_tile_array( + operation, + program=ws_gemm_3x3, + arguments=arguments, + ) + + +class LinalgGenericGemmPattern(TileArrayRewritePattern): + """Matches contraction maps and the multiply-add dataflow in a generic op.""" + + root = linalg.GenericOp + + @classmethod + def match_and_rewrite(cls, operation, rewriter) -> bool: + """Checks the contraction dataflow and replaces it with the WS kernel.""" + # The driver has already checked the declared root operation type. + operation = cast(linalg.GenericOp, operation) + attrs = operation.operation.attributes + maps = ( + "affine_map<(i,j,k)->(i,k)>", + "affine_map<(i,j,k)->(k,j)>", + "affine_map<(i,j,k)->(i,j)>", + ) + if tuple(ArrayAttr(attrs["indexing_maps"])) != tuple( + Attribute.parse(text) for text in maps + ): + return False + if tuple(ArrayAttr(attrs["iterator_types"])) != tuple( + Attribute.parse(f"#linalg.iterator_type<{kind}>") + for kind in ("parallel", "parallel", "reduction") + ): + return False + block = operation.regions[0].blocks[0] + body = tuple(block.operations) + if len(block.arguments) != 3 or len(body) != 3: + return False + multiply, add, terminator = body + if ( + multiply.operation.name != "arith.muli" + or tuple( + multiply.operands[index] for index in range(len(multiply.operands)) + ) + != tuple(block.arguments[index] for index in range(2)) + or add.operation.name != "arith.addi" + or set(add.operands[index] for index in range(len(add.operands))) + != {multiply.results[0], block.arguments[2]} + or terminator.operation.name != "linalg.yield" + or tuple( + terminator.operands[index] for index in range(len(terminator.operands)) + ) + != (add.results[0],) + ): + return False + arguments = tuple(operation.inputs) + tuple(operation.outputs) + expected = MemRefType.get([3, 3], IntegerType.get_signless(32)) + if len(operation.results) or len(arguments) != 3: + return False + if any(value.type != expected for value in arguments): + return False + if not _has_zero_output(operation, arguments[2]): + return False + return rewriter.replace_with_tile_array( + operation, + program=ws_gemm_3x3, + arguments=arguments, + ) + + +class AffineGemmPattern(TileArrayRewritePattern): + """Matches matrix accesses and multiply-add dataflow in a canonical loop nest.""" + + root = affine.AffineForOp + + @classmethod + def match_and_rewrite(cls, operation, rewriter) -> bool: + """Checks the loop dataflow and replaces it with the WS kernel.""" + outer = _loop(operation, 3) + if outer is None or len(tuple(outer.operations)) != 2: + return False + middle = _loop(outer.operations[0], 3) + if middle is None or len(tuple(middle.operations)) != 2: + return False + inner = _loop(middle.operations[0], 3) + if inner is None: + return False + body = tuple(inner.operations) + if tuple(op.operation.name for op in body) != ( + "affine.load", + "affine.load", + "affine.load", + "arith.muli", + "arith.addi", + "affine.store", + "affine.yield", + ): + return False + lhs, rhs, current, multiply, add, store, _ = body + i, j, k = outer.arguments[0], middle.arguments[0], inner.arguments[0] + if ( + _indices( + lhs, + tuple(lhs.operands[index] for index in range(len(lhs.operands)))[1:], + ) + != (i, k) + or _indices( + rhs, + tuple(rhs.operands[index] for index in range(len(rhs.operands)))[1:], + ) + != (k, j) + or _indices( + current, + tuple( + current.operands[index] for index in range(len(current.operands)) + )[1:], + ) + != (i, j) + or _indices( + store, + tuple(store.operands[index] for index in range(len(store.operands)))[ + 2: + ], + ) + != (i, j) + or tuple( + multiply.operands[index] for index in range(len(multiply.operands)) + ) + != (lhs.results[0], rhs.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] + != (add.results[0], current.operands[0]) + ): + return False + arguments = (lhs.operands[0], rhs.operands[0], current.operands[0]) + 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]): + return False + return rewriter.replace_with_tile_array( + operation, + program=ws_gemm_3x3, + arguments=arguments, + ) diff --git a/python/synapse/patterns/pattern.py b/python/synapse/patterns/pattern.py new file mode 100644 index 0000000..4f0dfd2 --- /dev/null +++ b/python/synapse/patterns/pattern.py @@ -0,0 +1,28 @@ +"""Public API for Synapse rewrite patterns.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, ClassVar + +if TYPE_CHECKING: + from taskflow_mlir.ir import OpView + + from synapse.compiler.pattern_rewriter import PatternRewriter + + +class TileArrayRewritePattern(ABC): + """Matches source IR and rewrites it into a TileArray implementation.""" + + root: ClassVar[type[OpView] | tuple[type[OpView], ...]] + + @classmethod + @abstractmethod + def match_and_rewrite( + cls, + operation: OpView, + rewriter: PatternRewriter, + ) -> bool: + """Matches one root operation and rewrites it when supported.""" + + raise NotImplementedError diff --git a/tests/python/compiler/test_gemm_patterns.py b/tests/python/compiler/test_gemm_patterns.py new file mode 100644 index 0000000..ec686ce --- /dev/null +++ b/tests/python/compiler/test_gemm_patterns.py @@ -0,0 +1,223 @@ +"""Checks GEMM matching, semantic rejection, and task-preserving replacement.""" + +import pytest +import synapse +from synapse.compiler.pattern_rewriter import apply_patterns +from synapse.patterns.gemm_pattern import ( + AffineGemmPattern, + LinalgGemmPattern, + LinalgGenericGemmPattern, +) +from taskflow_mlir.dialects import linalg, neura, taskflow +from taskflow_mlir.ir import Context, Location, Module + +LINALG_GEMM = """ + linalg.matmul ins(%a, %b : memref<3x3xi32>, memref<3x3xi32>) + outs(%c : memref<3x3xi32>) +""" + +GENERIC_GEMM = """ + linalg.generic { + indexing_maps = [affine_map<(i,j,k)->(i,k)>, + affine_map<(i,j,k)->(k,j)>, + affine_map<(i,j,k)->(i,j)>], + iterator_types = ["parallel", "parallel", "reduction"] + } ins(%a, %b : memref<3x3xi32>, memref<3x3xi32>) + outs(%c : memref<3x3xi32>) { + ^bb0(%lhs: i32, %rhs: i32, %acc: i32): + %product = arith.muli %lhs, %rhs : i32 + %sum = arith.addi %product, %acc : i32 + linalg.yield %sum : i32 + } +""" + +AFFINE_GEMM = """ + affine.for %i = 0 to 3 { + affine.for %j = 0 to 3 { + affine.for %k = 0 to 3 { + %lhs = affine.load %a[%i, %k] : memref<3x3xi32> + %rhs = affine.load %b[%k, %j] : memref<3x3xi32> + %acc = affine.load %c[%i, %j] : memref<3x3xi32> + %product = arith.muli %lhs, %rhs : i32 + %sum = arith.addi %acc, %product : i32 + affine.store %sum, %c[%i, %j] : memref<3x3xi32> + } + } + } +""" + + +def task_source(body=LINALG_GEMM): + """Wraps a semantic kernel with an independently allocated output buffer.""" + return """ +module { + func.func @gemm(%A: memref<3x3xi32>, %B: memref<3x3xi32>) -> memref<3x3xi32> { + %C = memref.alloc() : memref<3x3xi32> + %done = taskflow.task @gemm + will_reads(%A, %B : memref<3x3xi32>, memref<3x3xi32>) + will_writes(%C : memref<3x3xi32>) + [original_read_memrefs(%A, %B : memref<3x3xi32>, memref<3x3xi32>), + original_write_memrefs(%C : memref<3x3xi32>)] + : (memref<3x3xi32>, memref<3x3xi32>, memref<3x3xi32>) -> memref<3x3xi32> { + ^bb0(%a: memref<3x3xi32>, %b: memref<3x3xi32>, %c: memref<3x3xi32>): + %zero = arith.constant 0 : i32 + linalg.fill ins(%zero : i32) outs(%c : memref<3x3xi32>) + BODY + taskflow.yield done_writes(%c : memref<3x3xi32>) + } + return %done : memref<3x3xi32> + } +} +""".replace("BODY", body) + + +@pytest.mark.parametrize( + "body,pattern", + [ + (LINALG_GEMM, LinalgGemmPattern), + (GENERIC_GEMM, LinalgGenericGemmPattern), + (AFFINE_GEMM, AffineGemmPattern), + ], +) +def test_rewrites_gemm_inside_existing_task(body, pattern): + import synapse.language as synl + from synapse.frontend.lowering import lower + from synapse.library import ws_gemm_3x3 + + direct_source = lower(ws_gemm_3x3, argument_types=(synl.i32[3, 3],) * 3) + with Context(), Location.unknown(): + taskflow.register_dialect() + neura.register_dialect() + module = Module.parse(task_source(body)) + function = module.body.operations[0] + task = function.regions[0].blocks[0].operations[1] + block = task.regions[0].blocks[0] + terminator = tuple(block.operations)[-1] + result = task.results[0] + assert apply_patterns(module, [pattern]) == 1 + assert module.operation.verify() + assert task.results[0] == result + assert tuple(block.operations)[-1] == terminator + kernel = next( + op for op in block.operations if op.operation.name == "neura.kernel" + ) + assert isinstance(kernel, neura.KernelOp) + assert len(kernel.results) == 0 + assert tuple( + kernel.operands[index] for index in range(len(kernel.operands)) + ) == tuple(block.arguments[index] for index in range(len(block.arguments))) + direct = Module.parse(direct_source) + direct_task = direct.body.operations[0].regions[0].blocks[0].operations[0] + direct_kernel = direct_task.regions[0].blocks[0].operations[0] + assert kernel.operation.get_asm( + use_local_scope=True + ) == direct_kernel.operation.get_asm(use_local_scope=True) + assert str(module).count('"neura.load"') == 3 + assert str(module).count('"neura.mac"') == 9 + assert str(module).count('"neura.store"') == 3 + + +@pytest.mark.parametrize( + "source", + [ + task_source().replace("constant 0", "constant 1"), + task_source().replace( + "linalg.fill ins(%zero : i32) outs(%c : memref<3x3xi32>)", "" + ), + task_source().replace("3x3xi32", "2x2xi32"), + task_source() + .replace("%C = memref.alloc() : memref<3x3xi32>", "") + .replace("%C", "%A"), + task_source(GENERIC_GEMM).replace("arith.muli", "arith.subi"), + task_source(AFFINE_GEMM).replace("%a[%i, %k]", "%a[%k, %i]"), + task_source(AFFINE_GEMM).replace("%k = 0 to 3", "%k = 0 to 2"), + ], +) +def test_rejected_candidates_preserve_original_ir(source): + with Context(), Location.unknown(): + taskflow.register_dialect() + neura.register_dialect() + module = Module.parse(source) + before = str(module) + assert ( + apply_patterns( + module, [LinalgGemmPattern, LinalgGenericGemmPattern, AffineGemmPattern] + ) + == 0 + ) + assert module.operation.verify() + assert str(module) == before + + +def test_public_compile_accepts_task_ir(): + mapped = synapse.compile(task_source(), target="neura") + assert "compiled_ii = 1" in mapped + assert "linalg.matmul" not in mapped + + +def test_affine_match_accepts_actual_linalg_lowering(): + import subprocess + from pathlib import Path + + executable = ( + Path(__file__).resolve().parents[3] + / "build/amoeba/tools/mlir-amoeba-opt/mlir-amoeba-opt" + ) + lowered = subprocess.run( + [str(executable), "--convert-linalg-to-affine-loops"], + input=task_source(), + text=True, + capture_output=True, + check=True, + ).stdout + rewritten = synapse.rewrite(lowered, patterns=[AffineGemmPattern]) + assert "neura.kernel" in rewritten + assert "arith.muli" not in rewritten + + +def test_failed_tile_array_lowering_preserves_source_module(): + import synapse.language as synl + from synapse.compiler.pattern_rewriter import PatternRewriter + + def invalid_program(A: synl.Tensor, B: synl.Tensor, C: synl.Tensor): + synl.load(A[0:0, 0], tile=synl.TileArray(4, 4)[0, 1]) + + with Context(), Location.unknown(): + taskflow.register_dialect() + neura.register_dialect() + module = Module.parse(task_source()) + task = module.body.operations[0].regions[0].blocks[0].operations[1] + root = task.regions[0].blocks[0].operations[2] + assert isinstance(root, linalg.MatmulOp) + before = str(module) + with pytest.raises(ValueError, match="cannot be empty"): + PatternRewriter(root).replace_with_tile_array( + root, + program=invalid_program, + arguments=tuple( + root.operands[index] for index in range(len(root.operands)) + ), + ) + assert str(module) == before + assert module.operation.verify() + + +def test_unknown_output_aliasing_is_not_assumed_safe(): + source = ( + task_source() + .replace("%B: memref<3x3xi32>)", "%B: memref<3x3xi32>, %C: memref<3x3xi32>)") + .replace("%C = memref.alloc() : memref<3x3xi32>", "") + ) + assert "neura.kernel" not in synapse.rewrite(source) + + +def test_intervening_memory_write_invalidates_zero_initialization(): + body = ( + """ + %one = arith.constant 1 : i32 + %index = arith.constant 0 : index + memref.store %one, %c[%index, %index] : memref<3x3xi32> + """ + + LINALG_GEMM + ) + assert "neura.kernel" not in synapse.rewrite(task_source(body)) diff --git a/tests/python/compiler/test_pattern_rewriter.py b/tests/python/compiler/test_pattern_rewriter.py new file mode 100644 index 0000000..e8fb188 --- /dev/null +++ b/tests/python/compiler/test_pattern_rewriter.py @@ -0,0 +1,182 @@ +"""Tests for the Synapse MLIR pattern driver.""" + +from typing import cast + +import pytest +import synapse +import synapse.language as synl +from synapse.compiler.pattern_rewriter import ( + PatternRewriter, + apply_patterns, +) +from synapse.patterns import TileArrayRewritePattern +from taskflow_mlir.dialects import arith, func, linalg, neura, taskflow +from taskflow_mlir.ir import Context, Location, Module + + +class AddToMultiplyPattern(TileArrayRewritePattern): + """Replaces integer addition to verify the rewrite mechanism.""" + + root = arith.AddIOp + + @classmethod + def match_and_rewrite( + cls, + operation: arith.AddIOp, + rewriter: PatternRewriter, + ) -> bool: + """Replaces arith.addi with arith.muli.""" + + with rewriter.ip: + replacement = arith.MulIOp( + operation.lhs, + operation.rhs, + ) + + rewriter.replace_op( + operation, + replacement, + ) + + return True + + +def test_applies_pattern_to_matching_operation(): + source = """ + module { + func.func @compute(%lhs: i32, %rhs: i32) -> i32 { + %result = arith.addi %lhs, %rhs : i32 + return %result : i32 + } + } + """ + + with Context(), Location.unknown(): + module = Module.parse(source) + + rewrite_count = apply_patterns( + module, + patterns=[AddToMultiplyPattern], + ) + + assert module.operation.verify() + rewritten = str(module) + + assert rewrite_count == 1 + assert "arith.addi" not in rewritten + assert "arith.muli" in rewritten + + +def copy_program(A: synl.Tensor, C: synl.Tensor): + """Copies a tensor through configured loads and stores.""" + array = synl.TileArray(4, 4) + for x in range(1, A.type.shape[1] + 1): + value = synl.load(A[:, x - 1], tile=array[0, x]) + synl.store(value, target=C[:, x - 1], tile=array[x, 0]) + + +class CopyPattern(TileArrayRewritePattern): + """Checks and replaces a copy in one user-defined callback.""" + + root = linalg.CopyOp + + @classmethod + def match_and_rewrite(cls, operation, rewriter) -> bool: + """Checks the buffer types and inserts the copy implementation.""" + operation = cast(linalg.CopyOp, operation) + (source,) = operation.inputs + (target,) = operation.outputs + if source.type != target.type: + return False + return rewriter.replace_with_tile_array( + operation, + program=copy_program, + arguments=(source, target), + ) + + +def _copy_source(): + """Provides a bufferized copy with an independently allocated destination.""" + return """ + module { + func.func @copy(%A: memref<3x3xi32>) -> memref<3x3xi32> { + %C = memref.alloc() : memref<3x3xi32> + %done = taskflow.task @copy + will_reads(%A : memref<3x3xi32>) + will_writes(%C : memref<3x3xi32>) + [original_read_memrefs(%A : memref<3x3xi32>), + original_write_memrefs(%C : memref<3x3xi32>)] + : (memref<3x3xi32>, memref<3x3xi32>) -> memref<3x3xi32> { + ^bb0(%a: memref<3x3xi32>, %c: memref<3x3xi32>): + linalg.copy ins(%a : memref<3x3xi32>) outs(%c : memref<3x3xi32>) + taskflow.yield done_writes(%c : memref<3x3xi32>) + } + return %done : memref<3x3xi32> + } + } + """ + + +@pytest.mark.parametrize("dtype", ["i32", "f32"]) +def test_custom_pattern_checks_and_rewrites_in_one_callback(dtype): + source = _copy_source().replace("3x3xi32", f"3x3x{dtype}") + rewritten = synapse.rewrite(source, patterns=[CopyPattern]) + assert "linalg.copy" not in rewritten + assert rewritten.count('"neura.load"') == 3 + assert rewritten.count('"neura.store"') == 3 + + +def test_alias_checks_apply_to_custom_rewrites(): + source = ( + _copy_source() + .replace("%C = memref.alloc() : memref<3x3xi32>", "") + .replace("%C", "%A") + ) + assert "neura.kernel" not in synapse.rewrite(source, patterns=[CopyPattern]) + + +def test_root_filter_runs_before_the_pattern_callback(): + class UnrelatedPattern(TileArrayRewritePattern): + root = arith.AddIOp + + @classmethod + def match_and_rewrite(cls, operation, rewriter) -> bool: + """Detects a callback invoked for an unrelated root type.""" + raise AssertionError("the root filter must reject this operation") + + assert "linalg.copy" in synapse.rewrite(_copy_source(), patterns=[UnrelatedPattern]) + + +def test_failed_user_check_preserves_the_original_operation(): + class RejectingPattern(TileArrayRewritePattern): + root = linalg.CopyOp + + @classmethod + def match_and_rewrite(cls, operation, rewriter) -> bool: + """Declines a candidate after the root type has matched.""" + return False + + rewritten = synapse.rewrite(_copy_source(), patterns=[RejectingPattern]) + assert "linalg.copy" in rewritten + assert "neura.kernel" not in rewritten + + +def test_invalid_user_implementation_remains_an_error(): + def invalid(A: synl.Tensor, C: synl.Tensor): + raise ValueError("invalid user program") + + class InvalidPattern(TileArrayRewritePattern): + root = linalg.CopyOp + + @classmethod + def match_and_rewrite(cls, operation, rewriter) -> bool: + """Passes an invalid implementation to the staged lowering path.""" + operation = cast(linalg.CopyOp, operation) + return rewriter.replace_with_tile_array( + operation, + program=invalid, + arguments=tuple(operation.inputs) + tuple(operation.outputs), + ) + + with pytest.raises(ValueError, match="invalid user program"): + synapse.rewrite(_copy_source(), patterns=[InvalidPattern]) From 845ff167044d268b1b9a1a213ee4094e39670474 Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Sat, 12 Sep 2026 01:15:11 +0800 Subject: [PATCH 16/19] Update CI --- .github/workflows/test.yml | 27 +++------------------------ 1 file changed, 3 insertions(+), 24 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index cf55a6c..e660417 100644 --- a/.github/workflows/test.yml +++ b/.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 @@ -173,8 +152,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' From 8f8cbd9089941e24d2ddad0804af52195616732d Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Sat, 12 Sep 2026 14:22:19 +0800 Subject: [PATCH 17/19] Update project dependencies --- pyproject.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 8c7c081..82fe143 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,6 +7,8 @@ name = "synapse" version = "0.0.0" requires-python = ">=3.11" +dependencies = ["PyYAML"] + [tool.setuptools] package-dir = {"" = "python"} From 748f7b7314a6039ee959df8b8cbdb9c97c4c7a61 Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Sat, 12 Sep 2026 14:40:35 +0800 Subject: [PATCH 18/19] Update dependencies --- .github/workflows/test.yml | 1 - pyproject.toml | 10 +++++++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index e660417..4b06aa4 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -64,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 diff --git a/pyproject.toml b/pyproject.toml index 82fe143..c537a96 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,15 @@ name = "synapse" version = "0.0.0" requires-python = ">=3.11" -dependencies = ["PyYAML"] +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"} From b0571ebb4d4a468f014d59580c60ff5698ae33a9 Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Wed, 16 Sep 2026 10:17:32 +0800 Subject: [PATCH 19/19] Update dependencies & ignore file --- .gitignore | 2 ++ mlir/amoeba | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 03d3467..9341e02 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,7 @@ # Local docs / planning artifacts planbook/ +docs/ +pupa/ # Python bytecode and caches __pycache__/ diff --git a/mlir/amoeba b/mlir/amoeba index a6fb32a..50cb16a 160000 --- a/mlir/amoeba +++ b/mlir/amoeba @@ -1 +1 @@ -Subproject commit a6fb32a2473f8b616e1044bdc3bed178a2dcc692 +Subproject commit 50cb16a8ef7c07da19dc06359810d35908a39589