[python] Set up binding for preprocessing transform ops by bangtianliu · Pull Request #22227 · iree-org/iree

bangtianliu · 2025-10-07T06:10:19Z

This PR adds Python bindings for the Preprocessing Transform Extension ops.
It exposes the transform ops defined in that extension, such as transform.iree.match.contraction, as well as other commonly used ops like transform.iree.match.cast_compatible_dag_from_root and transform.iree.match.dim_is_multiple_of, etc.

The bindings are exposed using declare_mlir_dialect_extension_python_bindings. The logic mainly follows this llvm upstream PR: llvm/llvm-project#159450.

kuhar

Looks good overall. Maybe reference the mlir PR you used as a reference in the PR description, so that we can track the overall design/mechanics to that original PR?

bangtianliu · 2025-10-08T04:54:45Z

When updating the test, I noticed that the printed operation from python looks like:

 %batch_dims, %m_dims, %n_dims, %k_dims = transform.iree.match.contraction %arg0, lhs_type = f32, rhs_type = f32, output_type = f32 indexing_maps [#map, #map1, #map2] : !transform.any_op -> !transform.param<i64>

determined by assembly format

let assemblyFormat = [{
    $operand_handle
    `,` `lhs_type` `=` $lhs_type
    `,` `rhs_type` `=` $rhs_type
    `,` `output_type` `=` $output_type
    (`indexing_maps` $indexing_maps^)?   // ← Prints WITHOUT braces!
    attr-dict `:` type($operand_handle) `->` type($batch_dims)
}];

It turns out the parser accepts both of the following formats:

// Format 1: Custom format (what the printer outputs)
transform.iree.match.contraction %op, lhs_type = f32, rhs_type = f32, output_type = f32
  indexing_maps [#map1, #map2, #map3]
  : !transform.any_op -> !transform.param<i64>

// Format 2: Attribute dictionary format (also accepted!)
transform.iree.match.contraction %op, lhs_type = f32, rhs_type = f32, output_type = f32
  {indexing_maps = [#map1, #map2, #map3]}
  : !transform.any_op -> !transform.param<i64>

To make the printed and test formats consistent, I updated the assembly format of transform.iree.match.contraction to :

let assemblyFormat = [{
    $operand_handle
    `,` `lhs_type` `=` $lhs_type
    `,` `rhs_type` `=` $rhs_type
    `,` `output_type` `=` $output_type
    attr-dict                              // ← indexing_maps goes here now
    `:` type($operand_handle) `->` type($batch_dims)
}];

This change doesn’t alter semantics or parsing behavior. It only ensures that the Python-printed module matches the format used in tests.

kuhar · 2025-10-08T12:09:02Z

Can we print it as:

`(` `,` `indexing_maps` `=` $indexing_maps^)?

with a leading comma and equals? It's weird to have some attributes printed inline and some in the attribute dictionary

bangtianliu · 2025-10-08T15:56:01Z

Can we print it as:
`(` `,` `indexing_maps` `=` $indexing_maps^)?
with a leading comma and equals? It's weird to have some attributes printed inline and some in the attribute dictionary

Updated! The assembly format now prints all attributes inline:

%batch_dims, %m_dims, %n_dims, %k_dims = transform.iree.match.contraction %arg0, lhs_type = f32, rhs_type = f32, output_type = f32, indexing_maps = [#map, #map1, #map2] : !transform.any_op -> !transform.param<i64>

Should I update the tests in preprocessing_match_ops.mlir for consistency, if so I prefer to do it in another separate PR.

kuhar · 2025-10-08T16:12:41Z

    `,` `lhs_type` `=` $lhs_type
    `,` `rhs_type` `=` $rhs_type
    `,` `output_type` `=` $output_type
+    (`,` `indexing_maps` `=` $indexing_maps^)?


Can we also use this syntax in all tests and add an op example in let description = above?

Here is my comment: #22227 (comment)

Should we do it in this PR?

bangtianliu · 2025-10-09T14:22:34Z

cc @kuhar for another look, since other PRs are also on top of this.

kuhar · 2025-10-10T00:34:36Z

cc @kuhar for another look, since other PRs are also on top of this.

Don't stack PRs when not necessary -- AFAICT, the other PRs don't need these python bindings, and we don't need syntax changes in this PR.

bangtianliu · 2025-10-10T00:39:17Z

cc @kuhar for another look, since other PRs are also on top of this.

Don't stack PRs when not necessary -- AFAICT, the other PRs don't need these python bindings, and we don't need syntax changes in this PR.

For python testing, we updated the assembly format for match contraction op in this PR.

kuhar · 2025-10-10T00:42:15Z

You can either update the assembly format first and then come back and update this PR, or land this first and update the assembly format later. We should avoid stacking PRs when not necessary.

bangtianliu · 2025-10-10T00:52:35Z

You can either update the assembly format first and then come back and update this PR, or land this first and update the assembly format later. We should avoid stacking PRs when not necessary.

Ok, could I send a new PR that focuses only on updating the assembly format? The existing PR (#22247) should be closed since it depends on the current one.

kuhar · 2025-10-10T00:53:36Z

You can force-push to an existing PR branch to avoid closing/opening PRs

bangtianliu · 2025-10-10T01:41:05Z

You can force-push to an existing PR branch to avoid closing/opening PRs

Ok the PR is closed automatically when reset --hard and push it to no change. I sent a PR which is independent and pure about updating assembly format: #22270.

…cher ops (#22270) According to this comment: #22227 (comment), this PR updates the assembly format for contraction and convolution matcher ops. The attention matcher op's indexing map is not optional and also consistent with what is suggested in the comment, so no need to update. Signed-off-by: Bangtian Liu <liubangtian@gmail.com>

Signed-off-by: Bangtian Liu <liubangtian@gmail.com> address the comment and also update the assembly format Signed-off-by: Bangtian Liu <liubangtian@gmail.com> update assembly format Signed-off-by: Bangtian Liu <liubangtian@gmail.com>

…op (#22311) This PR extends the preprocessing transform dialect Python bindings by adding `MatchConvolutionOp `and `MatchAttentionOp`, completing the matcher operation API suite started in #22227. **Motivation** Completes the Python binding API for all preprocessing transform matcher operations, enabling SHARK Tuner to programmatically generate Transform Dialect specs for convolutions and attention operations in addition to contractions. This removes the need for string-based spec generation and improves type safety and maintainability. Signed-off-by: Bangtian Liu <liubangtian@gmail.com>

This PR uses python binding to build td specs for contraction op. The python binding is exposed from IREE PR iree-org/iree#22227 (MatchContractionOp binding). Right now, the code still contains two paths for building TD specs: - Contraction ops: Use new Python bindings via `build_contraction_td_spec()` - Convolution/Attention ops: Still use old string template approach via `build_td_spec()` After this PR, I will also send PRs for convolution and attention ops to migrate them to Python bindings. Once all operations are supported, some code relevant to temporary solutions and string-based method will be fully cleaned up. --------- Signed-off-by: Bangtian Liu <liubangtian@gmail.com>

…cher ops (iree-org#22270) According to this comment: iree-org#22227 (comment), this PR updates the assembly format for contraction and convolution matcher ops. The attention matcher op's indexing map is not optional and also consistent with what is suggested in the comment, so no need to update. Signed-off-by: Bangtian Liu <liubangtian@gmail.com> Signed-off-by: Philipp <philipp.weidel@intel.com>

This PR adds Python bindings for the Preprocessing Transform Extension ops. It exposes the transform ops defined in that extension, such as `transform.iree.match.contraction`, as well as other commonly used ops like `transform.iree.match.cast_compatible_dag_from_root `and `transform.iree.match.dim_is_multiple_of`, etc. The bindings are exposed using `declare_mlir_dialect_extension_python_bindings`. The logic mainly follows this llvm upstream PR: llvm/llvm-project#159450. Signed-off-by: Bangtian Liu <liubangtian@gmail.com> Signed-off-by: Philipp <philipp.weidel@intel.com>

…op (iree-org#22311) This PR extends the preprocessing transform dialect Python bindings by adding `MatchConvolutionOp `and `MatchAttentionOp`, completing the matcher operation API suite started in iree-org#22227. **Motivation** Completes the Python binding API for all preprocessing transform matcher operations, enabling SHARK Tuner to programmatically generate Transform Dialect specs for convolutions and attention operations in addition to contractions. This removes the need for string-based spec generation and improves type safety and maintainability. Signed-off-by: Bangtian Liu <liubangtian@gmail.com> Signed-off-by: Philipp <philipp.weidel@intel.com>

…cher ops (#22270) According to this comment: iree-org/iree#22227 (comment), this PR updates the assembly format for contraction and convolution matcher ops. The attention matcher op's indexing map is not optional and also consistent with what is suggested in the comment, so no need to update. Signed-off-by: Bangtian Liu <liubangtian@gmail.com>

This PR adds Python bindings for the Preprocessing Transform Extension ops. It exposes the transform ops defined in that extension, such as `transform.iree.match.contraction`, as well as other commonly used ops like `transform.iree.match.cast_compatible_dag_from_root `and `transform.iree.match.dim_is_multiple_of`, etc. The bindings are exposed using `declare_mlir_dialect_extension_python_bindings`. The logic mainly follows this llvm upstream PR: llvm/llvm-project#159450. Signed-off-by: Bangtian Liu <liubangtian@gmail.com>

…op (iree-org#22311) This PR extends the preprocessing transform dialect Python bindings by adding `MatchConvolutionOp `and `MatchAttentionOp`, completing the matcher operation API suite started in iree-org#22227. **Motivation** Completes the Python binding API for all preprocessing transform matcher operations, enabling SHARK Tuner to programmatically generate Transform Dialect specs for convolutions and attention operations in addition to contractions. This removes the need for string-based spec generation and improves type safety and maintainability. Signed-off-by: Bangtian Liu <liubangtian@gmail.com>

This PR uses python binding to build td specs for contraction op. The python binding is exposed from IREE PR iree-org/iree#22227 (MatchContractionOp binding). Right now, the code still contains two paths for building TD specs: - Contraction ops: Use new Python bindings via `build_contraction_td_spec()` - Convolution/Attention ops: Still use old string template approach via `build_td_spec()` After this PR, I will also send PRs for convolution and attention ops to migrate them to Python bindings. Once all operations are supported, some code relevant to temporary solutions and string-based method will be fully cleaned up. --------- Signed-off-by: Bangtian Liu <liubangtian@gmail.com>

bangtianliu requested a review from benvanik as a code owner October 7, 2025 06:10

bangtianliu marked this pull request as draft October 7, 2025 06:10

bangtianliu force-pushed the transform_python_binding branch 10 times, most recently from 3419a6f to 9f3f729 Compare October 7, 2025 22:46

bangtianliu marked this pull request as ready for review October 7, 2025 23:13

bangtianliu requested review from Max191 and kuhar October 7, 2025 23:13

kuhar reviewed Oct 7, 2025

View reviewed changes

Comment thread compiler/bindings/python/test/ir/dialects_test.py Outdated

Comment thread compiler/bindings/python/test/ir/dialects_test.py Outdated

bangtianliu requested review from MaheshRavishankar and qedawkins as code owners October 8, 2025 04:57

bangtianliu requested a review from kuhar October 8, 2025 15:56

kuhar reviewed Oct 8, 2025

View reviewed changes

bangtianliu requested a review from kuhar October 8, 2025 17:09

bangtianliu force-pushed the transform_python_binding branch from be0e02c to 942931b Compare October 8, 2025 17:39

This was referenced Oct 8, 2025

[Codegen] Update assembly format for match transform ops #22247

Closed

[Codegen] Update the td spec using the contraction matcher op #22249

Merged

bangtianliu mentioned this pull request Oct 9, 2025

[sharktuner] v3.8 release reediness checklist nod-ai/amd-shark-ai#2453

Closed

7 tasks

bangtianliu changed the title ~~[python] Set up binding for preprocessing transform ops~~ [python] Set up binding for preprocessing transform ops (1/4) Oct 10, 2025

bangtianliu changed the title ~~[python] Set up binding for preprocessing transform ops (1/4)~~ [python] Set up binding for preprocessing transform ops Oct 10, 2025

bangtianliu mentioned this pull request Oct 10, 2025

[Codegen] Update the assembly formats and corresponding tests for matcher ops #22270

Merged

bangtianliu force-pushed the transform_python_binding branch from 942931b to a3d441c Compare October 10, 2025 03:49

bangtianliu force-pushed the transform_python_binding branch from a3d441c to 0e595e3 Compare October 10, 2025 05:51

kuhar approved these changes Oct 10, 2025

View reviewed changes

bangtianliu merged commit aea660c into iree-org:main Oct 10, 2025
45 checks passed

This was referenced Oct 15, 2025

[python] Set up python binding for matcher convolution and attention op #22311

Merged

[tuner] use python bindings to build td specs nod-ai/amd-shark-ai#2515

Closed

bangtianliu mentioned this pull request Oct 17, 2025

[tuner] use python binding to build td specs for contraction nod-ai/amd-shark-ai#2543

Merged

bangtianliu mentioned this pull request Dec 15, 2025

[Tuner] Improving ease of use for the tuner nod-ai/amd-shark-ai#814

Open

30 tasks

Uh oh!

Conversation

bangtianliu commented Oct 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kuhar left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

bangtianliu commented Oct 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kuhar commented Oct 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

bangtianliu commented Oct 8, 2025

Uh oh!

kuhar Oct 8, 2025

Choose a reason for hiding this comment

Uh oh!

bangtianliu Oct 8, 2025

Choose a reason for hiding this comment

Uh oh!

bangtianliu commented Oct 9, 2025

Uh oh!

kuhar commented Oct 10, 2025

Uh oh!

bangtianliu commented Oct 10, 2025

Uh oh!

kuhar commented Oct 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

bangtianliu commented Oct 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kuhar commented Oct 10, 2025

Uh oh!

bangtianliu commented Oct 10, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

bangtianliu commented Oct 7, 2025 •

edited

Loading

bangtianliu commented Oct 8, 2025 •

edited

Loading

kuhar commented Oct 8, 2025 •

edited

Loading

kuhar commented Oct 10, 2025 •

edited

Loading

bangtianliu commented Oct 10, 2025 •

edited

Loading