[ShangkunLi]

Summary

This PR adds infrastructure for replacing computations inside existing Taskflow tasks with reusable TileArray programs.

It introduces:

• TileArrayRewritePattern, with an operation root and a match_and_rewrite() callback;
• recursive pattern application with root-type filtering;
• PatternRewriter utilities for replacing and erasing operations;
• lowering a selected TileArray program directly into the existing task;
• checks for argument types, task memory effects, and buffer aliasing;
• tests for successful replacement, rejected matches, invalid implementations, and memory safety.

Example

A pattern declares the source operation it handles and performs its checks and replacement in one callback:

class CopyPattern(TileArrayRewritePattern):
    root = linalg.CopyOp

    @classmethod
    def match_and_rewrite(cls, operation, rewriter):
        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),
        )

The pattern is applied to an existing Taskflow module with:

rewritten = synapse.rewrite(
    source,
    patterns=[CopyPattern],
)

When the pattern matches, Synapse builds and verifies the TileArray implementation, inserts its Neura kernel into the original task, and removes the matched operation. The task interface and dependency results remain unchanged.

When a pattern rejects the operation, the original IR is preserved.

[tancheng]

"rewrite" sounds confusing/ambiguous to me. Pattern match and rewrite has rewrite, what we are doing here is replace sth with existing kernel, rather than mlir pattern rewrite. Can you think about some other candidates for naming?

[ShangkunLi]

"rewrite" sounds confusing/ambiguous to me. Pattern match and rewrite has rewrite, what we are doing here is replace sth with existing kernel, rather than mlir pattern rewrite. Can you think about some other candidates for naming?

How about using "replace"?

e.g., TileArrayProgramPattern, match_and_replace(), PatternReplacer, apply_replacement_patterns(), etc.

[ShangkunLi]

"rewrite" sounds confusing/ambiguous to me. Pattern match and rewrite has rewrite, what we are doing here is replace sth with existing kernel, rather than mlir pattern rewrite. Can you think about some other candidates for naming?

How about using "replace"?

e.g., TileArrayProgramPattern, match_and_replace(), PatternReplacer, apply_replacement_patterns(), etc.

I have updated accordingly ~