[tuner] retire op_matchers.py through IREE python bindings (#1264)

This PR offloads the `op_matcher.py` logic to the IREE compiler through
exposed python bindings as shown below:

- `root_op_list = iree_codegen.get_tuner_root_ops(input_module)`
- `linalg::isaContractionOpInterface` and `linalg::inferContractionDims`
- `linalg::isaConvolutionOpInterface` and `linalg::inferConvolutionDims`


Issue: 
#1110 
#814 
#nod-ai/playbook#74

---------

Signed-off-by: Bangtian Liu <[email protected]>

Author: bangtianliu

tuner/tuner/candidate_gen.py

@@ -39,7 +39,6 @@ class DispatchTuner(DispatchParser):
     @abstractmethod
     def get_td_spec(
         self,
-        ir_module: ir.Module,
         compilation_info: iree_codegen.CompilationInfoAttr,
     ) -> ir.Module:
         """Generate a transform dialect spec that applies the compilation info attr."""
@@ -62,12 +61,14 @@ def find_handler(self, op_name: str) -> DispatchTuner:
 
 
 class ContractionOpInterfaceTuner(DispatchTuner, ContractionOpInterfaceParser):
+    def __init__(self, root_op: ir.Operation):
+        super().__init__(root_op)
+
     def get_td_spec(
         self,
-        ir_module: ir.Module,
         compilation_info: iree_codegen.CompilationInfoAttr,
     ) -> ir.Module:
-        contraction_op: ir.Operation = self.get_contraction_operation(ir_module)
+        contraction_op = self.get_root_op()
         lhs_type = ir.ShapedType(contraction_op.operands[0].type)
         rhs_type = ir.ShapedType(contraction_op.operands[1].type)
         acc_type = ir.ShapedType(contraction_op.operands[2].type)
@@ -77,17 +78,16 @@ def get_td_spec(
         # TODO(Max191): Get the function name from the func.func in the input module.
         func_name = f"match_contraction_{M}x{N}x{K}_{lhs_type.element_type}x{rhs_type.element_type}x{acc_type.element_type}"
         return build_td_spec(
-            ir_module.context, contraction_op, compilation_info, func_name
+            contraction_op.context, contraction_op, compilation_info, func_name
         )
 
 
 class ConvolutionOpInterfaceTuner(DispatchTuner, ConvolutionOpInterfaceParser):
     def get_td_spec(
         self,
-        ir_module: ir.Module,
         compilation_info: iree_codegen.CompilationInfoAttr,
     ) -> ir.Module:
-        conv_op: ir.Operation = self.get_conv_operation(ir_module)
+        conv_op = self.get_root_op()
         assert (
             conv_op.name == "linalg.conv_2d_nhwc_hwcf"
         ), "expected linalg.conv_2d_nhwc_hwcf"
@@ -104,7 +104,7 @@ def get_td_spec(
         conv_type = conv_op.name.split(".")[-1]
         # TODO(Max191): Get the function name from the func.func in the input module.
         func_name = f"match_{conv_type}_{N}x{H}x{W}x{C}x{P}x{Q}x{F}_{lhs_type.element_type}x{rhs_type.element_type}x{acc_type.element_type}"
-        return build_td_spec(ir_module.context, conv_op, compilation_info, func_name)
+        return build_td_spec(conv_op.context, conv_op, compilation_info, func_name)
 
 
 @dataclass
@@ -156,21 +156,33 @@ def generate_configs_and_td_specs(
     pipeline_options_search_space: PipelineOptionsSearchSpace = PipelineOptionsSearchSpace(),
     codegen_pipeline: iree_codegen.DispatchLoweringPassPipeline = iree_codegen.DispatchLoweringPassPipeline.LLVMGPUVectorDistribute,
 ) -> list[ir.Module]:
-    dispatch_tuner_registry = DispatchTunerRegistry()
-    dispatch_tuner_registry.register(
-        [
-            ContractionOpInterfaceTuner(),
-            ConvolutionOpInterfaceTuner(),
-        ]
-    )
+    dispatch_tuners: list[type[DispatchTuner]] = [
+        ContractionOpInterfaceTuner,
+        ConvolutionOpInterfaceTuner,
+    ]
+
+    root_op_list = iree_codegen.get_tuner_root_ops(input_module)
+    if len(root_op_list) == 0:
+        tune_logger.error(
+            "No root ops found. Did you forget to pass "
+            "--iree-config-add-tuner-attributes during compilation?"
+        )
+        return []
+    elif len(root_op_list) > 1:
+        tune_logger.error("Multiple root ops found. Only one is currently supported.")
+        return []
 
-    walk_result: OpWalkResult = walk_mlir_op(input_module, dispatch_tuner_registry)
+    root_op = root_op_list[0]
+
+    dispatch_tuner: Optional[DispatchTuner] = None
+    for tuner_class in dispatch_tuners:
+        tuner = tuner_class(root_op)
+        if tuner.has_valid_root_op():
+            dispatch_tuner = tuner
+            break
 
-    dispatch_tuner = walk_result.dispatch_tuner
     assert dispatch_tuner, "No suitable dispatch tuner found"
-    problem_size: ProblemSize = dispatch_tuner.get_shapes(
-        str(input_module).splitlines()
-    )
+    problem_size: ProblemSize = dispatch_tuner.get_problem_size()
     tune_logger.debug(str(problem_size))
 
     # Index 0 is reserved for default config, so it gets a placeholder spec.
@@ -196,7 +208,7 @@ def generate_configs_and_td_specs(
         if i >= limit:
             break
         tune_logger.debug(f"Solution #{i+1}: {config}")
-        td_spec_module = dispatch_tuner.get_td_spec(input_module, config)
+        td_spec_module = dispatch_tuner.get_td_spec(config)
         assert td_spec_module, "Failed to generate transform dialect spec"
         config_specs.append(td_spec_module)
 
@@ -263,7 +275,7 @@ def run_command(run_pack: RunPack) -> RunResult:
 # info makes the inputs to compilation consistent, and allows for overwriting
 # the compilation info with generated TD specs during codegen.
 def strip_root_op_attr(module: ir.Module):
-    root_ops: list[ir.Operation] = get_ops_from_module(module, is_root_op)
+    root_ops: list[ir.Operation] = iree_codegen.get_tuner_root_ops(module)
     for root_op in root_ops:
         assert (
             ROOT_OP_ATTR_NAME in root_op.opview.attributes

tuner/tuner/candidate_gen_test.py

@@ -19,11 +19,30 @@
 
 from . import candidate_gen
 from . import common
-from . import op_matchers
 
 from .test_utils import tuner_ctx
 
 
+def walk_collect_ops(
+    op: ir.Operation,
+    ops: list[ir.Operation],
+    fn,
+) -> ir.WalkResult:
+    if fn(op):
+        ops.append(op)
+    return ir.WalkResult.ADVANCE
+
+
+def get_ops_from_module(module: ir.Module, fn):
+    ops: list[ir.Operation] = []
+    for op in module.body.operations:
+        op.walk(
+            lambda op: walk_collect_ops(op, ops, fn),
+            ir.WalkOrder.POST_ORDER,
+        )
+    return ops
+
+
 def test_get_td_spec_contraction(tuner_ctx: common.TunerContext) -> None:
     context = tuner_ctx.mlir_ctx
     module_str = """
@@ -75,14 +94,15 @@ def test_get_td_spec_contraction(tuner_ctx: common.TunerContext) -> None:
     )
 
     ir_module = ir.Module.parse(module_str, context)
+    root_op_list = iree_codegen.get_tuner_root_ops(ir_module)
+    assert len(root_op_list) == 1
+    root_op = root_op_list[0]
 
-    tuner = candidate_gen.ContractionOpInterfaceTuner()
-    td_spec_module = tuner.get_td_spec(ir_module, compilation_info)
+    tuner = candidate_gen.ContractionOpInterfaceTuner(root_op)
+    td_spec_module = tuner.get_td_spec(compilation_info)
     assert td_spec_module
 
-    named_sequence_ops: list[
-        transform.NamedSequenceOp
-    ] = op_matchers.get_ops_from_module(
+    named_sequence_ops: list[transform.NamedSequenceOp] = get_ops_from_module(
         module=td_spec_module,
         fn=lambda op: isinstance(op.opview, transform.NamedSequenceOp),
     )
@@ -157,14 +177,14 @@ def test_get_td_spec_convolution(tuner_ctx: common.TunerContext) -> None:
     )
 
     ir_module = ir.Module.parse(module_str, context)
-
-    tuner = candidate_gen.ConvolutionOpInterfaceTuner()
-    td_spec_module = tuner.get_td_spec(ir_module, compilation_info)
+    root_op_list = iree_codegen.get_tuner_root_ops(ir_module)
+    assert len(root_op_list) == 1
+    root_op = root_op_list[0]
+    tuner = candidate_gen.ConvolutionOpInterfaceTuner(root_op)
+    td_spec_module = tuner.get_td_spec(compilation_info)
     assert td_spec_module
 
-    named_sequence_ops: list[
-        transform.NamedSequenceOp
-    ] = op_matchers.get_ops_from_module(
+    named_sequence_ops: list[transform.NamedSequenceOp] = get_ops_from_module(
         module=td_spec_module,
         fn=lambda op: isinstance(op.opview, transform.NamedSequenceOp),
     )

tuner/tuner/common.py

@@ -19,6 +19,7 @@
 from iree.compiler.dialects import iree_gpu  # type: ignore
 from iree.compiler.dialects import transform  # type: ignore
 import iree.compiler as ireec  # type: ignore
+from iree.compiler._mlir_libs._mlir import ir  # type: ignore
 
 
 class CommonTypes:
@@ -144,6 +145,13 @@ def MNK(self) -> tuple[list[int], list[int], list[int]]:
         return (self.matmul_size.M, self.matmul_size.N, self.matmul_size.K)
 
 
+def get_map_result_dim_positions(map: ir.AffineMap) -> Optional[list[int]]:
+    if not map.is_projected_permutation:
+        return None
+
+    return [ir.AffineDimExpr(expr).position for expr in map.results]
+
+
 def get_compatible_mfma_intrinsics(
     problem_size: ProblemSize,
     mma_intrinsics: list[iree_gpu.MMAIntrinsic],

tuner/tuner/common_test.py

@@ -63,6 +63,23 @@ def test_gpu_pipeline_options(tuner_ctx: common.TunerContext) -> None:
     )
 
 
+def test_get_map_result_dim_positions(tuner_ctx: common.TunerContext) -> None:
+    dim0 = ir.AffineDimExpr.get(0)
+    dim1 = ir.AffineDimExpr.get(1)
+    dim2 = ir.AffineDimExpr.get(2)
+
+    # Valid projected permutation: (d0, d1, d2) -> (d0, d2).
+    valid_map = ir.AffineMap.get(3, 0, [dim0, dim2])
+    result = common.get_map_result_dim_positions(valid_map)
+    assert result == [0, 2], f"Expected [0, 2], got {result}"
+
+    # Not a projected permutation: (d0, d1, d2) -> (d0 + d1).
+    sum_expr = dim0 + dim1
+    invalid_map = ir.AffineMap.get(3, 0, [sum_expr])
+    result = common.get_map_result_dim_positions(invalid_map)
+    assert result is None, "Expected None for non-projected permutation"
+
+
 def test_get_pipeline_config(tuner_ctx: common.TunerContext) -> None:
     mma_intrinsic = iree_gpu.MMAIntrinsic.MFMA_F32_16x16x16_F16
     mma_attr = iree_gpu.MMAAttr.get(mma_intrinsic)