Add tensor-based TileArray programs and GEMM pattern rewriting

[ShangkunLi]

Summary

This PR adds two complementary capabilities:

• TileArray programming: users describe computation, memory access, and data movement on a CGRA tile array.
• Pattern matching and rewriting: the compiler reuses the same TileArray implementation to replace matching computations inside an existing Taskflow task.

Example: programming a GEMM

The following program computes C = A @ B for 3×3 i32 tensors on a physical 4×4 TileArray. The west column performs loads, nine tiles perform MACs, and the south row performs stores.

import synapse
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):
    array = synl.TileArray(x_tiles=4, y_tiles=4)
    partial_sums: list[TileArrayValue] = []

    # Physical rows descend 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] = []

        # A flows east; partial sums flow south.
        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

    for x, accumulated in enumerate(partial_sums, start=1):
        synl.store(
            accumulated,
            target=C[:, x - 1],
            tile=array[x, 0],
        )


mapped_ir = synapse.compile(
    ws_gemm_3x3,
    target="neura",
    argument_types=(synl.i32[3, 3],) * 3,
)

x and y identify physical tiles; k identifies the reduction position assigned to each MAC row. B[k, x - 1] remains stationary, while A[:, k] supplies successive output rows.

Standalone compilation records a TileArrayProgram, creates a Taskflow task, and lowers its body to a Neura kernel. The backend handles routing and scheduling. C is overwritten and must not overlap A or B.

Pattern matching and rewriting

Each pattern directly inherits TileArrayRewritePattern and defines:

• root: the operation type considered by the driver;
• match_and_rewrite(): computation-specific checks and the replacement call.

This PR provides patterns for named linalg.matmul, canonical linalg.generic contraction, and canonical Affine GEMM. All three reuse ws_gemm_3x3.

flowchart TD
    A["Existing Taskflow task<br/>Linalg or Affine computation"]
    B["Filter by pattern.root"]
    C["match_and_rewrite()<br/>Check computation and implementation conditions"]
    D["Build TileArrayProgram<br/>from the selected Python function"]
    E["Check task and memory compatibility"]
    F["lower_to_kernel()<br/>Generate and verify the replacement"]
    G["Replace the original computation<br/>Preserve the task, results, and completion dependencies"]
    U["Keep original IR"]

    A --> B
    B -->|Matching root| C
    C -->|Match| D
    C -->|Not applicable| U
    D --> E
    E -->|Compatible| F
    E -->|Not compatible| U
    F --> G

Loading

The current GEMM implementation overwrites C, whereas source matmul accumulates into C. Its patterns therefore require zero-initialized C. Generic and Affine patterns additionally check the indexing and multiply-add dataflow.

The rewriter constructs and verifies the replacement before changing the original IR. Both standalone compilation and pattern rewriting share lower_to_kernel().

from synapse.patterns.gemm_pattern import LinalgGemmPattern

rewritten_ir = synapse.rewrite(
    source,
    patterns=[LinalgGemmPattern],
)

Unmatched computations remain available for later lowering. Automatic Counter/General fallback is not connected in this PR.