From 697b2cba10c985f68a0c4086302d11e7c4162227 Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Mon, 28 Sep 2026 14:10:27 +0800 Subject: [PATCH 1/2] Update the implementation of pattern replacement --- .../synapse/compiler/pattern_replacement.py | 82 +------------------ 1 file changed, 3 insertions(+), 79 deletions(-) diff --git a/python/synapse/compiler/pattern_replacement.py b/python/synapse/compiler/pattern_replacement.py index 55e73c5..8366baf 100644 --- a/python/synapse/compiler/pattern_replacement.py +++ b/python/synapse/compiler/pattern_replacement.py @@ -11,7 +11,6 @@ IntegerType, MemRefType, Module, - OpResult, OpView, Value, ) @@ -74,10 +73,9 @@ def replace_with_tile_array_program( ) -> 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. + The pattern establishes applicability and semantic preconditions. The + shared path adapts argument types and validates the generated kernel + before changing the source IR. """ from taskflow_mlir.dialects import func @@ -94,8 +92,6 @@ def replace_with_tile_array_program( 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): @@ -157,52 +153,6 @@ def _apply_patterns( return replace_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 @@ -230,29 +180,3 @@ def _task_argument_types(operation, arguments): 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 From 097a8339f2eea7a1ca7df072c067bb12e6e0a69e Mon Sep 17 00:00:00 2001 From: ShangkunLI Date: Mon, 28 Sep 2026 15:56:17 +0800 Subject: [PATCH 2/2] Modify the type and value systems in synapse --- python/synapse/compiler/compiler.py | 4 +- .../synapse/compiler/pattern_replacement.py | 18 ++--- python/synapse/frontend/lowering.py | 30 ++++---- python/synapse/language/__init__.py | 15 ++-- python/synapse/language/tensor.py | 70 ------------------ python/synapse/language/tile_array_program.py | 44 ++++++------ python/synapse/language/types.py | 23 +++--- python/synapse/language/values.py | 71 +++++++++++++++++++ python/synapse/library/gemm.py | 12 ++-- tests/python/compiler/test_gemm_patterns.py | 18 +---- .../compiler/test_pattern_replacement.py | 15 +--- tests/python/frontend/test_memory_lowering.py | 10 +-- tests/python/language/test_memory.py | 18 ++--- tests/python/language/test_types.py | 6 +- tests/python/language/test_values.py | 28 ++++++++ 15 files changed, 197 insertions(+), 185 deletions(-) delete mode 100644 python/synapse/language/tensor.py create mode 100644 python/synapse/language/values.py create mode 100644 tests/python/language/test_values.py diff --git a/python/synapse/compiler/compiler.py b/python/synapse/compiler/compiler.py index 6f7f958..9a237ac 100644 --- a/python/synapse/compiler/compiler.py +++ b/python/synapse/compiler/compiler.py @@ -6,7 +6,7 @@ from tempfile import TemporaryDirectory from synapse.frontend.lowering import lower -from synapse.language.types import TensorType +from synapse.language.types import BufferType from synapse.patterns import TileArrayProgramPattern @@ -14,7 +14,7 @@ def compile( program: Callable | str, *, target: str, - argument_types: tuple[TensorType, ...] = (), + argument_types: tuple[BufferType, ...] = (), patterns: Sequence[type[TileArrayProgramPattern]] | None = None, ) -> str: """Compiles a TileArray function or bufferized task IR for the backend.""" diff --git a/python/synapse/compiler/pattern_replacement.py b/python/synapse/compiler/pattern_replacement.py index 8366baf..1f1d43f 100644 --- a/python/synapse/compiler/pattern_replacement.py +++ b/python/synapse/compiler/pattern_replacement.py @@ -15,7 +15,7 @@ Value, ) -from synapse.language.types import DType, TensorType +from synapse.language.types import BufferType, DType from synapse.patterns import TileArrayProgramPattern @@ -86,11 +86,11 @@ def replace_with_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: + buffer_types = _task_buffer_types(operation, arguments) + if buffer_types is None: return False - tile_program = build_tile_array_program(program, argument_types=inferred_types) + tile_program = build_tile_array_program(program, argument_types=buffer_types) lowering = TileArrayProgramLowering(tile_program) staged = Module.create() @@ -153,13 +153,13 @@ def _apply_patterns( return replace_count -def _task_argument_types(operation, arguments): - """Returns supported capture types, or None when the task boundary is unsuitable.""" +def _task_buffer_types(operation, arguments): + """Returns supported buffer 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 = [] + buffer_types = [] for value in arguments: if ( not BlockArgument.isinstance(value) @@ -178,5 +178,5 @@ def _task_argument_types(operation, arguments): return None if memref != MemRefType.get(list(memref.shape), memref.element_type): return None - types.append(TensorType(tuple(memref.shape), dtype)) - return tuple(types) + buffer_types.append(BufferType(tuple(memref.shape), dtype)) + return tuple(buffer_types) diff --git a/python/synapse/frontend/lowering.py b/python/synapse/frontend/lowering.py index f55c774..c266888 100644 --- a/python/synapse/frontend/lowering.py +++ b/python/synapse/frontend/lowering.py @@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, cast from synapse.language.spatial import Tile -from synapse.language.tensor import Tensor, TensorAccess +from synapse.language.values import Buffer, BufferSlice from synapse.language.tile_array_program import ( AddOp, ConstantOp, @@ -21,27 +21,27 @@ TileArrayOp, TileArrayProgram, ) -from synapse.language.types import DType, TensorType +from synapse.language.types import BufferType, DType if TYPE_CHECKING: from taskflow_mlir.ir import AffineMapAttr, DenseI64ArrayAttr, DictAttr, Value -def lower(program_fn: Callable, *, argument_types: tuple[TensorType, ...] = ()) -> str: +def lower(program_fn: Callable, *, argument_types: tuple[BufferType, ...] = ()) -> str: """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, ...] = () + program_fn: Callable, *, argument_types: tuple[BufferType, ...] = () ) -> TileArrayProgram: - """Runs a TileArray function and records its operations and tensor accesses. + """Runs a TileArray function and records its operations and buffer slices. Direct compilation and pattern replacement share this construction step. It creates a TileArrayProgram without importing or constructing MLIR. """ - # Program arguments currently carry tensor types. Scalar argument capture + # Program arguments currently carry buffer types. Scalar argument capture # will use Taskflow value dependencies when that frontend path is added. function_signature = signature(program_fn) @@ -54,12 +54,12 @@ def build_tile_array_program( ) if any( - not isinstance(argument_type, TensorType) for argument_type in argument_types + not isinstance(argument_type, BufferType) for argument_type in argument_types ): - raise TypeError("program argument types must be TensorType values") + raise TypeError("program argument types must be BufferType values") arguments = [ - Tensor(name=name, type=argument_type) + Buffer(name=name, type=argument_type) for name, argument_type in zip(parameter_names, argument_types) ] @@ -199,7 +199,7 @@ def lower_to_single_task(self, program_name: str) -> str: return str(module) - def lower_to_kernel(self, argument_values: dict[Tensor, Value]) -> None: + def lower_to_kernel(self, argument_values: dict[Buffer, 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. @@ -217,7 +217,7 @@ def lower_to_kernel(self, argument_values: dict[Tensor, Value]) -> None: # Configuration validation precedes IR insertion. metadata = self.get_kernel_metadata() - memory_configs: dict[int, tuple[TensorAccess, DenseI64ArrayAttr]] = {} + memory_configs: dict[int, tuple[BufferSlice, DenseI64ArrayAttr]] = {} for index, operation in enumerate(self.program.operations): if isinstance(operation, LoadOp): access = operation.source @@ -277,16 +277,16 @@ def get_mlir_type(self, dtype: DType): raise NotImplementedError(f"unsupported tile-array data type: {dtype.value}") - def get_memref_type(self, tensor_type: TensorType): - """Translates a TensorType into a Taskflow MemRef type.""" + def get_memref_type(self, buffer_type: BufferType): + """Translates a BufferType into a Taskflow MemRef type.""" from taskflow_mlir.ir import MemRefType return MemRefType.get( - list(tensor_type.shape), self.get_mlir_type(tensor_type.dtype) + list(buffer_type.shape), self.get_mlir_type(buffer_type.dtype) ) - def get_memory_offsets(self, access: TensorAccess) -> DenseI64ArrayAttr: + def get_memory_offsets(self, access: BufferSlice) -> DenseI64ArrayAttr: """Converts a static access to Neura's constant element-offset array.""" from taskflow_mlir.ir import DenseI64ArrayAttr diff --git a/python/synapse/language/__init__.py b/python/synapse/language/__init__.py index 8ddb19c..7c42549 100644 --- a/python/synapse/language/__init__.py +++ b/python/synapse/language/__init__.py @@ -1,7 +1,6 @@ """Public Synapse language API.""" from .spatial import TileArray -from .tensor import Tensor from .tile_array_program import ( add, constant, @@ -9,12 +8,20 @@ mac, store, ) -from .types import DType, TensorType, f32, i32 +from .types import ( + BufferType, + DType, + f32, + i32, +) +from .values import Buffer, BufferSlice, SynapseValue __all__ = [ + "Buffer", + "BufferSlice", + "BufferType", "DType", - "Tensor", - "TensorType", + "SynapseValue", "TileArray", "add", "constant", diff --git a/python/synapse/language/tensor.py b/python/synapse/language/tensor.py deleted file mode 100644 index 2e97102..0000000 --- a/python/synapse/language/tensor.py +++ /dev/null @@ -1,70 +0,0 @@ -"""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 static 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 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 static 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 40d4694..79be5af 100644 --- a/python/synapse/language/tile_array_program.py +++ b/python/synapse/language/tile_array_program.py @@ -12,7 +12,7 @@ from dataclasses import dataclass, field from .spatial import Tile, TileArray -from .tensor import Tensor, TensorAccess +from .values import Buffer, BufferSlice from .types import DType @@ -43,9 +43,9 @@ class TileArrayOp: @dataclass(frozen=True) class LoadOp(TileArrayOp): - """Loads through an explicit address or a configured tensor access.""" + """Loads through an explicit address or a configured buffer slice.""" - source: TensorAccess | None = None + source: BufferSlice | None = None def __post_init__(self) -> None: """Validates the selected address form and result type.""" @@ -70,9 +70,9 @@ def addr(self) -> TileArrayValue | None: @dataclass(frozen=True) class StoreOp(TileArrayOp): - """Stores through an explicit address or a configured tensor access.""" + """Stores through an explicit address or a configured buffer slice.""" - target: TensorAccess | None = None + target: BufferSlice | None = None def __post_init__(self) -> None: """Validates the value and selected address form.""" @@ -154,7 +154,7 @@ def __post_init__(self) -> None: class MacOp(TileArrayOp): """A configured MAC using one stationary scalar value.""" - stationary_value: TensorAccess + stationary_value: BufferSlice @property def accumulated(self) -> TileArrayValue: @@ -189,8 +189,8 @@ def __post_init__(self) -> None: class StationaryBinding: """Stationary data assigned to the tiles of one template program.""" - source: Tensor - tile_values: tuple[tuple[Tile, TensorAccess], ...] + source: Buffer + tile_values: tuple[tuple[Tile, BufferSlice], ...] @dataclass(frozen=True) @@ -198,7 +198,7 @@ class TileArrayProgram: """A tile-array program produced by TileArrayBuilder.""" array: TileArray - arguments: tuple[Tensor, ...] + arguments: tuple[Buffer, ...] operations: tuple[TileArrayOp, ...] stationary: StationaryBinding | None @@ -213,7 +213,7 @@ class TileArrayBuilder: is called, it returns a TileArrayProgram. """ - def __init__(self, arguments: tuple[Tensor, ...] = ()): + def __init__(self, arguments: tuple[Buffer, ...] = ()): self._arguments = arguments self._array: TileArray | None = None self._operations: list[TileArrayOp] = [] @@ -428,23 +428,23 @@ def add( def load( - source: TensorAccess | None = None, + source: BufferSlice | None = None, *, addr: TileArrayValue | None = None, dtype: DType | None = None, tile: Tile, ) -> TileArrayValue: - """Loads from a tensor access or an explicit target address. + """Loads from a buffer slice or an explicit target address. 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. + already use the target address representation; they are not buffer 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 not isinstance(source, BufferSlice): + raise TypeError("load source must be a BufferSlice") if dtype is not None and dtype != source.dtype: raise TypeError("load dtype must match its source") dtype = source.dtype @@ -468,15 +468,15 @@ def load( def store( value: TileArrayValue, *, - target: TensorAccess | None = None, + target: BufferSlice | None = None, addr: TileArrayValue | None = None, tile: Tile, ) -> None: - """Stores a value through a tensor access or explicit target address.""" + """Stores a value through a buffer slice 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") + if target is not None and not isinstance(target, BufferSlice): + raise TypeError("store target must be a BufferSlice") operands = (value,) if addr is None else (value, addr) _require_active_builder().emit( operands=operands, @@ -495,13 +495,13 @@ def mac( input0: TileArrayValue, input1: TileArrayValue | None = None, *, - stationary: TensorAccess, + stationary: BufferSlice, tile: Tile, ) -> 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(stationary, BufferSlice): + raise TypeError("mac stationary data must be a buffer slice") operands = (input0,) if input1 is None else (input0, input1) diff --git a/python/synapse/language/types.py b/python/synapse/language/types.py index fc89d18..b060bc5 100644 --- a/python/synapse/language/types.py +++ b/python/synapse/language/types.py @@ -11,19 +11,23 @@ from enum import Enum -class DType(Enum): +class SynapseType: + """Base class for types in the Synapse language.""" + + +class DType(SynapseType, Enum): """A data type supported by Synapse.""" I32 = "i32" F32 = "f32" - def __getitem__(self, shape: int | tuple[int, ...]) -> TensorType: - """Create a tensor type with this data element type.""" + def __getitem__(self, shape: int | tuple[int, ...]) -> BufferType: + """Create a buffer type with this data element type.""" if not isinstance(shape, tuple): shape = (shape,) - return TensorType(shape=shape, dtype=self) + return BufferType(shape=shape, dtype=self) def __str__(self) -> str: """Return the source-level spelling of this type.""" @@ -32,11 +36,8 @@ def __str__(self) -> str: @dataclass(frozen=True) -class TensorType: - """The shape and scalar element type of a Synapse tensor. - - TensorType does not prescribe its later bufferized representation. - """ +class BufferType(SynapseType): + """The shape and element type of a buffer.""" shape: tuple[int, ...] dtype: DType @@ -45,12 +46,12 @@ def __post_init__(self) -> None: """Validates the dimensions and element type.""" if not self.shape: - raise TypeError("a tensor type requires at least one dimension") + raise TypeError("a buffer type requires at least one dimension") if not all( type(dimension) is int and dimension > 0 for dimension in self.shape ): - raise TypeError("tensor dimensions must be positive integers") + raise TypeError("buffer dimensions must be positive integers") if not isinstance(self.dtype, DType): raise TypeError("dtype must be a DType") diff --git a/python/synapse/language/values.py b/python/synapse/language/values.py new file mode 100644 index 0000000..caf196a --- /dev/null +++ b/python/synapse/language/values.py @@ -0,0 +1,71 @@ +"""Values exposed by the Synapse programming language.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from .types import BufferType, DType, SynapseType + + +class SynapseValue: + """Base class for values in the Synapse language.""" + + +@dataclass(frozen=True, eq=False) +class Buffer(SynapseValue): + """A symbolic buffer value passed to a TileArray program.""" + + name: str + type: BufferType + + def __getitem__( + self, + indices: int | slice | tuple[int | slice, ...], + ) -> BufferSlice: + """Selects an element or static slice of this buffer.""" + 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)}" + ) + + shape = [] + 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}" + ) + elif isinstance(index, slice): + shape.append(len(range(*index.indices(dimension)))) + else: + raise TypeError("indices must be integers or static slices") + + selected_type: SynapseType + if shape: + selected_type = BufferType(tuple(shape), self.type.dtype) + else: + selected_type = self.type.dtype + + return BufferSlice(self, indices, selected_type) + + +@dataclass(frozen=True) +class BufferSlice(SynapseValue): + """A static selection of one buffer value.""" + + source: Buffer + indices: tuple[int | slice, ...] + type: SynapseType + + @property + def dtype(self) -> DType: + """Returns the selected buffer element type.""" + return self.source.type.dtype + + @property + def is_scalar(self) -> bool: + """Reports whether the selection identifies one element.""" + return isinstance(self.type, DType) diff --git a/python/synapse/library/gemm.py b/python/synapse/library/gemm.py index 02704ee..373c936 100644 --- a/python/synapse/library/gemm.py +++ b/python/synapse/library/gemm.py @@ -4,14 +4,10 @@ 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") +def ws_gemm_3x3(A: synl.Buffer, B: synl.Buffer, C: synl.Buffer): + """Computes C = A @ B using nine MACs and six memory tiles.""" + if any(buffer.type != synl.i32[3, 3] for buffer in (A, B, C)): + raise ValueError("ws_gemm_3x3 requires three 3x3 i32 buffers") array = synl.TileArray(x_tiles=4, y_tiles=4) partial_sums: list[TileArrayValue] = [] diff --git a/tests/python/compiler/test_gemm_patterns.py b/tests/python/compiler/test_gemm_patterns.py index e82abd0..89b2cf0 100644 --- a/tests/python/compiler/test_gemm_patterns.py +++ b/tests/python/compiler/test_gemm_patterns.py @@ -125,9 +125,6 @@ def test_replaces_gemm_inside_existing_task(body, pattern): "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"), @@ -175,11 +172,11 @@ def test_affine_match_accepts_actual_linalg_lowering(): assert "arith.muli" not in replaced -def test_failed_tile_array_lowering_preserves_source_module(): +def test_invalid_buffer_slice_preserves_source_module(): import synapse.language as synl from synapse.compiler.pattern_replacement import PatternReplacer - def invalid_program(A: synl.Tensor, B: synl.Tensor, C: synl.Tensor): + def invalid_program(A: synl.Buffer, B: synl.Buffer, C: synl.Buffer): synl.load(A[0:0, 0], tile=synl.TileArray(4, 4)[0, 1]) with Context(), Location.unknown(): @@ -190,7 +187,7 @@ def invalid_program(A: synl.Tensor, B: synl.Tensor, C: synl.Tensor): root = task.regions[0].blocks[0].operations[2] assert isinstance(root, linalg.MatmulOp) before = str(module) - with pytest.raises(ValueError, match="cannot be empty"): + with pytest.raises(TypeError, match="positive integers"): PatternReplacer(root).replace_with_tile_array_program( root, program=invalid_program, @@ -202,15 +199,6 @@ def invalid_program(A: synl.Tensor, B: synl.Tensor, C: synl.Tensor): 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.replace(source) - - def test_intervening_memory_write_invalidates_zero_initialization(): body = ( """ diff --git a/tests/python/compiler/test_pattern_replacement.py b/tests/python/compiler/test_pattern_replacement.py index b1dfe34..ab74bae 100644 --- a/tests/python/compiler/test_pattern_replacement.py +++ b/tests/python/compiler/test_pattern_replacement.py @@ -67,8 +67,8 @@ def test_applies_pattern_to_matching_operation(): assert "arith.muli" in replaced -def copy_program(A: synl.Tensor, C: synl.Tensor): - """Copies a tensor through configured loads and stores.""" +def copy_program(A: synl.Buffer, C: synl.Buffer): + """Copies a buffer 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]) @@ -126,15 +126,6 @@ def test_custom_pattern_checks_and_replaces_in_one_callback(dtype): assert replaced.count('"neura.store"') == 3 -def test_alias_checks_apply_to_custom_repalces(): - source = ( - _copy_source() - .replace("%C = memref.alloc() : memref<3x3xi32>", "") - .replace("%C", "%A") - ) - assert "neura.kernel" not in synapse.replace(source, patterns=[CopyPattern]) - - def test_root_filter_runs_before_the_pattern_callback(): class UnrelatedPattern(TileArrayProgramPattern): root = arith.AddIOp @@ -162,7 +153,7 @@ def match_and_replace(cls, operation, replacer) -> bool: def test_invalid_user_implementation_remains_an_error(): - def invalid(A: synl.Tensor, C: synl.Tensor): + def invalid(A: synl.Buffer, C: synl.Buffer): raise ValueError("invalid user program") class InvalidPattern(TileArrayProgramPattern): diff --git a/tests/python/frontend/test_memory_lowering.py b/tests/python/frontend/test_memory_lowering.py index 36a3ff1..eebb953 100644 --- a/tests/python/frontend/test_memory_lowering.py +++ b/tests/python/frontend/test_memory_lowering.py @@ -12,7 +12,7 @@ from taskflow_mlir.ir import Context, DenseI64ArrayAttr, Location, Module -def configured_accesses(A: synl.Tensor): +def configured_accesses(A: synl.Buffer): array = synl.TileArray(4, 4) synl.load(A[1, 2], tile=array[0, 1]) synl.load(A[:, ::2], tile=array[0, 2]) @@ -61,11 +61,11 @@ def test_dynamic_ir_uses_address_operands_without_memory_configuration(): assert "memory_access" not in source -def test_empty_configured_queue_is_rejected(): - def empty(A: synl.Tensor): +def test_empty_buffer_slice_is_rejected(): + def empty(A: synl.Buffer): synl.load(A[0:0, 0], tile=synl.TileArray(4, 4)[0, 1]) - with pytest.raises(ValueError, match="cannot be empty"): + with pytest.raises(TypeError, match="positive integers"): lower(empty, argument_types=(synl.i32[3, 3],)) @@ -80,7 +80,7 @@ def test_dynamic_addresses_reach_backend_mapping(): def test_memory_bases_survive_task_argument_grouping(): - def program(A: synl.Tensor, unused: synl.Tensor, C: synl.Tensor): + def program(A: synl.Buffer, unused: synl.Buffer, C: synl.Buffer): array = synl.TileArray(4, 4) value = synl.load(A[:, 0], tile=array[0, 1]) synl.store(value, target=C[:, 0], tile=array[1, 0]) diff --git a/tests/python/language/test_memory.py b/tests/python/language/test_memory.py index d35960c..78956ac 100644 --- a/tests/python/language/test_memory.py +++ b/tests/python/language/test_memory.py @@ -26,12 +26,12 @@ def test_dynamic_addresses_are_explicit_operands(): def test_memory_configuration_retains_static_accesses(): - tensor = synl.Tensor("A", synl.i32[3, 4]) + buffer = synl.Buffer("A", synl.i32[3, 4]) array = synl.TileArray(4, 4) - builder = TileArrayBuilder((tensor,)) + builder = TileArrayBuilder((buffer,)) with builder: - value = synl.load(tensor[1, 2], tile=array[0, 1]) - synl.store(value, target=tensor[:, ::2], tile=array[1, 0]) + value = synl.load(buffer[1, 2], tile=array[0, 1]) + synl.store(value, target=buffer[:, ::2], tile=array[1, 0]) read, write = builder.build().operations assert isinstance(read, LoadOp) assert isinstance(write, StoreOp) @@ -44,21 +44,21 @@ def test_memory_configuration_retains_static_accesses(): def test_rejected_memory_forms_do_not_consume_value_ids(): - tensor = synl.Tensor("A", synl.i32[3, 3]) + buffer = synl.Buffer("A", synl.i32[3, 3]) array = synl.TileArray(4, 4) - builder = TileArrayBuilder((tensor,)) + builder = TileArrayBuilder((buffer,)) 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]) + synl.load(buffer[:, 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]) + synl.load(buffer[:, 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]) + synl.store(value, target=buffer[:, 0], addr=addr, tile=array[1, 0]) assert len(builder.build().operations) == 2 diff --git a/tests/python/language/test_types.py b/tests/python/language/test_types.py index acc0525..eda2dd0 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, TensorType +from synapse.language.types import BufferType, DType 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 == TensorType( + assert matrix == BufferType( shape=(4, 4), dtype=synl.i32, ) - assert vector == TensorType( + assert vector == BufferType( shape=(8,), dtype=synl.f32, ) diff --git a/tests/python/language/test_values.py b/tests/python/language/test_values.py new file mode 100644 index 0000000..0276f27 --- /dev/null +++ b/tests/python/language/test_values.py @@ -0,0 +1,28 @@ +import synapse.language as synl +from synapse.language.values import Buffer, SynapseValue + + +def test_buffer_value_has_a_separate_type(): + buffer_type = synl.BufferType(shape=(3, 4), dtype=synl.i32) + buffer = Buffer(name="A", type=buffer_type) + + assert isinstance(buffer, SynapseValue) + assert buffer.type is buffer_type + assert buffer.type.shape == (3, 4) + assert buffer.type.dtype == synl.i32 + + +def test_buffer_subscripts_produce_typed_slices(): + buffer = Buffer("A", synl.i32[3, 4]) + + scalar = buffer[1, 2] + column = buffer[:, 1] + + assert scalar.source is buffer + assert scalar.indices == (1, 2) + assert scalar.type == synl.i32 + assert scalar.is_scalar + + assert column.source is buffer + assert column.indices == (slice(None), 1) + assert column.type == synl.i32[3]