Skip to content

Add TileArray pattern rewriting infrastructure - #5

Open
ShangkunLi wants to merge 1 commit into
mainfrom
split/v2-rewriter
Open

ShangkunLi wants to merge 1 commit into
mainfrom
split/v2-rewriter

Conversation

@ShangkunLi

Copy link
Copy Markdown
Collaborator

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.

Adds user-defined match-and-rewrite callbacks and task-preserving TileArray
replacement with shared type and memory checks. Exposes rewrite and IR
compilation with explicitly supplied patterns.

Split from #3, original commit 99d2def.
@tancheng

Copy link
Copy Markdown

"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?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants