Merge pull request #8 from coredac/feature/type-system-modification

Simplify TileArray replacement and model typed buffers

Author: ShangkunLi

python/synapse/compiler/compiler.py

@@ -6,15 +6,15 @@
 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
 
 
 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."""

python/synapse/compiler/pattern_replacement.py

@@ -11,12 +11,11 @@
     IntegerType,
     MemRefType,
     Module,
-    OpResult,
     OpView,
     Value,
 )
 
-from synapse.language.types import DType, TensorType
+from synapse.language.types import BufferType, DType
 from synapse.patterns import TileArrayProgramPattern
 
 
@@ -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
 
@@ -88,14 +86,12 @@ 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)
-        if not _replacement_memory_is_legal(operation, arguments, lowering):
-            return False
 
         staged = Module.create()
         with InsertionPoint(staged.body):
@@ -157,59 +153,13 @@ 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."""
+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)
@@ -228,31 +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)
-
-
-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
+        buffer_types.append(BufferType(tuple(memref.shape), dtype))
+    return tuple(buffer_types)

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
 

python/synapse/language/__init__.py

@@ -1,20 +1,27 @@
 """Public Synapse language API."""
 
 from .spatial import TileArray
-from .tensor import Tensor
 from .tile_array_program import (
     add,
     constant,
     load,
     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",