forked from iree-org/iree
-
Notifications
You must be signed in to change notification settings - Fork 11
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[Codegen][Tuner] Make default and user-provided specs work (iree-org#…
…19449) This is a fixup to the tuning spec materialization that makes default and user-provided specs work e2e. As an example, a working spec for `linalg.matmul_transpose_b` is provided for gfx942. * Allow for tuning spec entry points that consume their argument op. This is so that tuning specs can use `transform.foreach_match`. * Require all tuning spec entry points to return `any_op`, so that we can chain includes. This works for both consumed and readonly args. * Add a test to show that user-provided tuning specs take precedence over default ones. * Work around a transform interpreter bug when multiple named sequenced across different modules share the same symbol name. Issue: iree-org#19214 --------- Signed-off-by: Jakub Kuderski <[email protected]>
- Loading branch information
Showing
7 changed files
with
240 additions
and
70 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
49 changes: 29 additions & 20 deletions
49
compiler/plugins/target/ROCM/test/tuning_spec_mmt_tile_and_fuse.mlir
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,24 +1,33 @@ | ||
| // RUN: iree-opt %s | ||
|
|
||
| module @mmt_tile_and_fuse_spec attributes { transform.with_named_sequence } { | ||
| transform.named_sequence @main(%arg0: !transform.any_op {transform.readonly}) -> () | ||
| attributes { iree_codegen.tuning_spec_entrypoint } { | ||
| %mmt = transform.structured.match ops{["linalg.generic"]} in %arg0 : (!transform.any_op) -> !transform.any_op | ||
| // transform.print %mmt {name="MMT"} : !transform.any_op | ||
| %config = transform.param.constant #iree_codegen.compilation_info< | ||
| lowering_config = #iree_gpu.lowering_config<{workgroup = [64, 64, 0], | ||
| reduction = [0, 0, 4], | ||
| thread = [8, 4], | ||
| promote_operands = [0, 1]}>, | ||
| translation_info = #iree_codegen.translation_info<pipeline = LLVMGPUTileAndFuse | ||
| workgroup_size = [128, 1, 1] subgroup_size = 64> | ||
| > -> !transform.any_param | ||
| transform.annotate %mmt "compilation_info" = %config : !transform.any_op, !transform.any_param | ||
| // Add a dummy unit attribute to be sure that the tuning spec applied. | ||
| // Otherwise it would be difficult to tell if the lowering config attribute | ||
| // comes from our tuning spec or if the compiler heuristic happened to produce | ||
| // the same config as this script. | ||
| transform.annotate %mmt "__tuning_spec_applied__" : !transform.any_op | ||
| transform.yield | ||
| } | ||
| transform.named_sequence @apply_op_config(%op: !transform.any_op {transform.readonly}, | ||
| %config: !transform.any_param {transform.readonly}) { | ||
| transform.annotate %op "compilation_info" = %config : !transform.any_op, !transform.any_param | ||
| transform.annotate %op "__custom_tuning_spec_applied__" : !transform.any_op | ||
| transform.yield | ||
| } | ||
|
|
||
| transform.named_sequence @match_mmt(%matmul: !transform.any_op {transform.readonly}) | ||
| -> (!transform.any_op, !transform.any_param) { | ||
| transform.match.operation_name %matmul ["linalg.generic"] : !transform.any_op | ||
| %config = transform.param.constant #iree_codegen.compilation_info< | ||
| lowering_config = #iree_gpu.lowering_config<{workgroup = [64, 64, 0], | ||
| reduction = [0, 0, 4], | ||
| thread = [8, 4], | ||
| promote_operands = [0, 1]}>, | ||
| translation_info = #iree_codegen.translation_info<pipeline = LLVMGPUTileAndFuse | ||
| workgroup_size = [128, 1, 1] subgroup_size = 64> | ||
| > -> !transform.any_param | ||
| transform.yield %matmul, %config : !transform.any_op, !transform.any_param | ||
| } | ||
|
|
||
| transform.named_sequence @main(%variant_op: !transform.any_op {transform.consumed}) -> (!transform.any_op) | ||
| attributes { iree_codegen.tuning_spec_entrypoint } { | ||
| transform.print %variant_op {name="Custom spec"} : !transform.any_op | ||
| %res = transform.foreach_match in %variant_op | ||
| @match_mmt -> @apply_op_config | ||
| : (!transform.any_op) -> !transform.any_op | ||
| transform.yield %res : !transform.any_op | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.