Repository: nod-ai/shark-ai
Pull request #1407: [tuner] Improve support for convolutions
Author: login: rkayaith
Created: 2025-05-07T17:59:15Z | Merged: 2025-05-08T22:50:20Z

This improves support for convolutions in the tuner by:
- generalizing `ConvolutionOpInterfaceParser` to support parsing `linalg.generic` ops, as well as additional layouts. However since constraint generation for more layouts is untested, only `nhwc_hwcf` and the generic equivalent are considered valid.
- supporting generating the padding lowering attribute for convolutions


Diff:
diff --git a/sharktuner/sharktuner/candidate_gen.py b/sharktuner/sharktuner/candidate_gen.py
index ed27cd163d0..db95962f9ed 100644
--- a/sharktuner/sharktuner/candidate_gen.py
+++ b/sharktuner/sharktuner/candidate_gen.py
@@ -81,9 +81,6 @@ def get_td_spec(
         compilation_info: iree_codegen.CompilationInfoAttr,
     ) -> ir.Module:
         conv_op = self.get_root_op()
-        assert (
-            conv_op.name == "linalg.conv_2d_nhwc_hwcf"
-        ), "expected linalg.conv_2d_nhwc_hwcf"
         func_name = self.get_root_op_func_name()
         return build_td_spec(conv_op.context, conv_op, compilation_info, func_name)
 
diff --git a/sharktuner/sharktuner/common.py b/sharktuner/sharktuner/common.py
index ba8448da15a..cdc5344902a 100644
--- a/sharktuner/sharktuner/common.py
+++ b/sharktuner/sharktuner/common.py
@@ -252,27 +252,6 @@ def read_input_mlir(filename: str) -> list[str]:
         return f.readlines()
 
 
-@dataclass
-class ConvDimInfo:
-    n: int
-    oh: int
-    ow: int
-    oc: int
-    fh: int
-    fw: int
-    ic: int
-
-    @staticmethod
-    def from_rhs_res(rhs_shaped_type: ShapedType, res_shaped_type: ShapedType):
-        n, oh, ow, oc = res_shaped_type.shape
-        fh, fw, ic, _ = rhs_shaped_type.shape
-        return ConvDimInfo(n, oh, ow, oc, fh, fw, ic)
-
-    @staticmethod
-    def from_problem_size(problem_size: ProblemSize):
-        return ConvDimInfo.from_rhs_res(problem_size.rhs_type, problem_size.res_type)
-
-
 @dataclass
 class MLIRTransformation:
     """Transformation of MLIR context"""
diff --git a/sharktuner/sharktuner/dispatch_constraints.py b/sharktuner/sharktuner/dispatch_constraints.py
index 6abd9da21fc..31fc718048c 100644
--- a/sharktuner/sharktuner/dispatch_constraints.py
+++ b/sharktuner/sharktuner/dispatch_constraints.py
@@ -53,12 +53,15 @@ def get_dispatch_constraints(
     if problem_size.dispatch_kind != DispatchKind.conv:
         return []
 
-    dim_info = ConvDimInfo.from_problem_size(problem_size)
-    conv_constraints = []
-    # WARNING: This sometimes makes the constraints UNSAT for some reason.
-    conv_constraints += [tile_m <= dim_info.ow]
-    conv_constraints += [tile_n <= dim_info.oc]
-    conv_constraints += [tile_k <= dim_info.ic]
+    max_tile_m = problem_size.matmul_size.M[-1]
+    [max_tile_n] = problem_size.matmul_size.N
+    max_tile_k = problem_size.matmul_size.K[-1]
+    conv_constraints = [
+        # WARNING: This sometimes makes the constraints UNSAT for some reason.
+        tile_m <= max_tile_m,
+        tile_n <= max_tile_n,
+        tile_k <= max_tile_k,
+    ]
     return conv_constraints
 
 
@@ -523,13 +526,14 @@ def set_cdim_tile_sizes(tile_sizes, contraction_dims, csizes):
         if required_padding:
             # TODO: Remove promotion of operand 2 once codegen supports handling padded outputs without promotion.
             promote_operands = [0, 1, 2]
-            workgroup_tile_m, workgroup_tile_n, _ = workgroup_tile_sizes
-            _, _, reduction_tile_k = reduction_tile_sizes
             _, _, mma_intrinsic_k = mma_attr.mnk_shape
             padding = [
-                workgroup_tile_m,
-                workgroup_tile_n,
-                reduction_tile_k * mma_intrinsic_k,
+                *(workgroup_tile_sizes[d] for d in problem_size.contraction_dims.m),
+                *(workgroup_tile_sizes[d] for d in problem_size.contraction_dims.n),
+                *(
+                    reduction_tile_sizes[d] * mma_intrinsic_k
+                    for d in problem_size.contraction_dims.k
+                ),
             ]
 
         compilation_infos = generate_compilation_infos(
diff --git a/sharktuner/sharktuner/dispatch_parser.py b/sharktuner/sharktuner/dispatch_parser.py
index 074d670474d..dcc5c3cea1d 100644
--- a/sharktuner/sharktuner/dispatch_parser.py
+++ b/sharktuner/sharktuner/dispatch_parser.py
@@ -96,41 +96,66 @@ def get_problem_size(self) -> ProblemSize:
         )
 
 
-# TODO(Max191): Support more convolution types. Only NHWC convs are supported.
 class ConvolutionOpInterfaceParser(DispatchParser):
     def __init__(self, root_op: ir.Operation):
         super().__init__(root_op)
-        self.supported_ops = ["linalg.conv_2d_nhwc_hwcf"]
 
     def has_valid_root_op(self) -> bool:
         root_op = self.get_root_op()
         if not linalg.isa_convolution_op(root_op):
             return False
-
-        return root_op.name in self.supported_ops
+        convolution_dims = linalg.infer_convolution_dimensions(root_op)
+        assert convolution_dims, "no convolution dimensions"
+        # Only allow 'nhwc_hwcf' convs.
+        # TODO: This dispatch parser class supports more layouts, but constraint
+        #       generation is not tested. Relax this check as support is verified.
+        if (
+            list(convolution_dims.batch) != [0]
+            or list(convolution_dims.output_image) != [1, 2]
+            or list(convolution_dims.output_channel) != [3]
+            or list(convolution_dims.filter_loop) != [4, 5]
+            or list(convolution_dims.input_channel) != [6]
+            or list(convolution_dims.depth) != []
+        ):
+            return False
+        return True
 
     def get_problem_size(self) -> ProblemSize:
         root_op = self.get_root_op()
-        lhs_type = ir.RankedTensorType(root_op.operands[0].type)
-        rhs_type = ir.RankedTensorType(root_op.operands[1].type)
-        res_type = ir.RankedTensorType(root_op.operands[2].type)
-        dim_info = ConvDimInfo.from_rhs_res(rhs_type, res_type)
+
+        def find_iter_dim_size(iter_dim: int, operand: int):
+            operand_type = root_op.operands[operand].type
+            indexing_map = linalg.get_indexing_maps(root_op)[operand]
+            tensor_dim = list(indexing_map.value.results).index(
+                ir.AffineExpr.get_dim(iter_dim)
+            )
+            return operand_type.shape[tensor_dim]
+
         convolution_dims = linalg.infer_convolution_dimensions(root_op)
         assert convolution_dims, "no convolution dimensions"
+        contraction_dims = ContractionDimensions(
+            batch=list(convolution_dims.depth),
+            m=list(convolution_dims.batch) + list(convolution_dims.output_image),
+            n=list(convolution_dims.output_channel),
+            k=list(convolution_dims.filter_loop) + list(convolution_dims.input_channel),
+        )
+        # Parallel dimension sizes come from the output operand, and reduction
+        # dimension sizes come from filter.
+        matmul_size = ContractionSizes(
+            B=[find_iter_dim_size(d, operand=2) for d in contraction_dims.batch],
+            M=[find_iter_dim_size(d, operand=2) for d in contraction_dims.m],
+            N=[find_iter_dim_size(d, operand=2) for d in contraction_dims.n],
+            K=[find_iter_dim_size(d, operand=1) for d in contraction_dims.k],
+        )
+
+        lhs_type = root_op.operands[0].type
+        rhs_type = root_op.operands[1].type
+        res_type = root_op.operands[2].type
         return ProblemSize(
-            matmul_size=ContractionSizes(
-                M=[dim_info.n, dim_info.oh, dim_info.ow],
-                N=[dim_info.oc],
-                K=[dim_info.fh, dim_info.fw, dim_info.ic],
-            ),
+            matmul_size=matmul_size,
             lhs_type=ShapedType(lhs_type.shape, lhs_type.element_type),
             rhs_type=ShapedType(rhs_type.shape, rhs_type.element_type),
             res_type=ShapedType(res_type.shape, res_type.element_type),
             dispatch_kind=DispatchKind.conv,
-            contraction_dims=ContractionDimensions(
-                m=list(convolution_dims.batch) + list(convolution_dims.output_image),
-                n=list(convolution_dims.output_channel),
-                k=list(convolution_dims.filter_loop)
-                + list(convolution_dims.input_channel),
-            ),
+            contraction_dims=contraction_dims,
         )
diff --git a/sharktuner/tests/dispatch_constraints_test.py b/sharktuner/tests/dispatch_constraints_test.py
index 6917f929a83..99e7810aaba 100644
--- a/sharktuner/tests/dispatch_constraints_test.py
+++ b/sharktuner/tests/dispatch_constraints_test.py
@@ -51,7 +51,9 @@ def test_generate_solutions(tuner_ctx: common.TunerContext) -> None:
     assert configs is not None
 
 
-def test_generate_solutions_tile_and_fuse(tuner_ctx: common.TunerContext) -> None:
+def test_generate_solutions_tile_and_fuse_contraction_padding(
+    tuner_ctx: common.TunerContext,
+) -> None:
     matmul_size = common.ContractionSizes([5369], [112], [112])
     contraction_dims = common.ContractionDimensions([0], [1], [2])
     lhs_type = common.ShapedType([5369, 112], tuner_ctx.type.f16)
@@ -94,6 +96,47 @@ def test_generate_solutions_tile_and_fuse(tuner_ctx: common.TunerContext) -> Non
     ), "Not all lowering configs have promote_operands = [0, 1, 2]"
 
 
+def test_generate_solutions_tile_and_fuse_conv_padding(
+    tuner_ctx: common.TunerContext,
+) -> None:
+    matmul_size = common.ContractionSizes(B=[], M=[2, 5, 5], N=[64], K=[3, 3, 32])
+    lhs_type = common.ShapedType([2, 7, 7, 32], tuner_ctx.type.f16)
+    rhs_type = common.ShapedType([64, 3, 3, 32], tuner_ctx.type.f16)
+    res_type = common.ShapedType([2, 5, 5, 64], tuner_ctx.type.f32)
+    contraction_dims = common.ContractionDimensions(
+        batch=[], m=[0, 1, 2], n=[3], k=[4, 5, 6]
+    )
+    problem_size = common.ProblemSize(
+        matmul_size,
+        lhs_type,
+        rhs_type,
+        res_type,
+        common.DispatchKind.conv,
+        contraction_dims,
+    )
+
+    solutions = list(
+        dispatch_constraints.generate_solutions(
+            tuner_ctx=tuner_ctx,
+            problem_size=problem_size,
+            num_subgrups=4,
+            mma_intrinsics=[iree_gpu.MMAIntrinsic.MFMA_F32_16x16x16_F16],
+            codegen_pipeline=iree_codegen.DispatchLoweringPassPipeline.LLVMGPUTileAndFuse,
+        )
+    )
+
+    assert len(solutions) > 0, "No solutions generated with TileAndFuse pipeline."
+    assert all(isinstance(sol, iree_codegen.CompilationInfoAttr) for sol in solutions)
+    assert all(
+        "padding =" in str(sol.lowering_config) for sol in solutions
+    ), "Not all lowering configs have padding option"
+    assert all(
+        [int(x) for x in sol.lowering_config.attributes["promote_operands"]]
+        == [0, 1, 2]
+        for sol in solutions
+    ), "Not all lowering configs have promote_operands = [0, 1, 2]"
+
+
 def test_calculate_shared_memory_usage_in_bytes(tuner_ctx: common.TunerContext) -> None:
     matmul_size = common.ContractionSizes([1024], [1024], [1024])
     contraction_dims = common.ContractionDimensions([0], [1], [2])
diff --git a/sharktuner/tests/dispatch_parser_test.py b/sharktuner/tests/dispatch_parser_test.py
index 8c0a0f57cf3..044cf6a6a28 100644
--- a/sharktuner/tests/dispatch_parser_test.py
+++ b/sharktuner/tests/dispatch_parser_test.py
@@ -12,6 +12,9 @@
 
 from typing import Generator
 
+# TODO: remove after https://github.com/llvm/llvm-project/pull/117918 is resolved.
+import sharktuner
+
 from iree.compiler import ir  # type: ignore
 from iree.compiler.dialects import func  # type: ignore
 from iree.compiler.dialects import iree_gpu  # type: ignore
@@ -24,7 +27,7 @@
 from sharktuner.test_utils import tuner_ctx
 
 
-CONTRACTION_TEMPLATE = r"""
+GENERIC_TEMPLATE = r"""
 builtin.module{{
     func.func @test(%arg0: {lhs_type}, %arg1: {rhs_type}) -> {res_type} {{
         %cst = arith.constant 0.000000e+00 : f32
@@ -55,7 +58,7 @@ def test_get_contraction_operation(tuner_ctx: common.TunerContext) -> None:
     context = tuner_ctx.mlir_ctx
 
     with ir.Location.unknown():
-        transpose_b_str = CONTRACTION_TEMPLATE.format(
+        transpose_b_str = GENERIC_TEMPLATE.format(
             lhs_type=ir.RankedTensorType.get([16, 64], ir.F16Type.get()),
             rhs_type=ir.RankedTensorType.get([32, 64], ir.F16Type.get()),
             res_type=ir.RankedTensorType.get([16, 32], ir.F32Type.get()),
@@ -82,7 +85,7 @@ def test_get_contraction_operation(tuner_ctx: common.TunerContext) -> None:
     assert isinstance(shapes.res_type.element_type, ir.F32Type)
 
     with ir.Location.unknown():
-        bmm_transposed_inputs_str = CONTRACTION_TEMPLATE.format(
+        bmm_transposed_inputs_str = GENERIC_TEMPLATE.format(
             lhs_type=ir.RankedTensorType.get([5, 8, 128], ir.F16Type.get()),
             rhs_type=ir.RankedTensorType.get([128, 40, 5], ir.F16Type.get()),
             res_type=ir.RankedTensorType.get([5, 40, 8], ir.F32Type.get()),
@@ -103,7 +106,7 @@ def test_get_contraction_operation(tuner_ctx: common.TunerContext) -> None:
     assert shapes.matmul_size.K == [128]
 
     with ir.Location.unknown():
-        bmm_transposed_inputs_str = CONTRACTION_TEMPLATE.format(
+        bmm_transposed_inputs_str = GENERIC_TEMPLATE.format(
             lhs_type=ir.RankedTensorType.get(
                 [16, 8, 15, 16, 64, 256], ir.F16Type.get()
             ),
@@ -223,7 +226,7 @@ def named_contraction(lhs, rhs, res):
         assert shape.res_type.shape == [5, 7]
 
 
-def test_get_conv_operation(tuner_ctx: common.TunerContext) -> None:
+def test_get_conv_nhwc_hwcf_operation(tuner_ctx: common.TunerContext) -> None:
     context = tuner_ctx.mlir_ctx
     module_str = """
         builtin.module{
@@ -247,6 +250,81 @@ def test_get_conv_operation(tuner_ctx: common.TunerContext) -> None:
         parser.has_valid_root_op()
     ), f"ConvolutionOpInterfaceParser does not support the op: {root_op.name}"
 
+    problem_size = parser.get_problem_size()
+    assert problem_size.contraction_dims.batch == []
+    assert problem_size.contraction_dims.m == [0, 1, 2]
+    assert problem_size.contraction_dims.n == [3]
+    assert problem_size.contraction_dims.k == [4, 5, 6]
+    assert problem_size.matmul_size.B == []
+    assert problem_size.matmul_size.M == [2, 32, 32]
+    assert problem_size.matmul_size.N == [16]
+    assert problem_size.matmul_size.K == [3, 3, 16]
+
+
+def test_get_group_conv_operation(tuner_ctx: common.TunerContext) -> None:
+    context = tuner_ctx.mlir_ctx
+    module_str = """
+    module {
+      func.func @test(%arg0: tensor<2x10x10x7x4xf32>, %arg1: tensor<7x16x3x3x4xf32>, %arg2: tensor<2x8x8x7x16xf32>) -> tensor<2x8x8x7x16xf32> {
+        %0 = linalg.conv_2d_nhwgc_gfhwc {
+           root_op,
+           dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>
+        } ins(%arg0, %arg1: tensor<2x10x10x7x4xf32>, tensor<7x16x3x3x4xf32>)
+          outs(%arg2: tensor<2x8x8x7x16xf32>) -> tensor<2x8x8x7x16xf32>
+        return %0 : tensor<2x8x8x7x16xf32>
+      }
+    }
+    """
+    module = ir.Module.parse(module_str, context)
+    root_op_list = iree_codegen.get_tuner_root_ops(module)
+    assert len(root_op_list) == 1
+    root_op = root_op_list[0]
+    parser = dispatch_parser.ConvolutionOpInterfaceParser(root_op)
+    assert parser.get_root_op_func_name() == "match_test"
+    assert parser.has_valid_root_op() is False, "group convs aren't supported yet"
+
+    problem_size = parser.get_problem_size()
+    assert problem_size.contraction_dims.batch == [3]
+    assert problem_size.contraction_dims.m == [0, 1, 2]
+    assert problem_size.contraction_dims.n == [4]
+    assert problem_size.contraction_dims.k == [5, 6, 7]
+    assert problem_size.matmul_size.B == [7]
+    assert problem_size.matmul_size.M == [2, 8, 8]
+    assert problem_size.matmul_size.N == [16]
+    assert problem_size.matmul_size.K == [3, 3, 4]
+
+
+def test_get_generic_conv_operation(tuner_ctx: common.TunerContext) -> None:
+    context = tuner_ctx.mlir_ctx
+    with ir.Location.name("generic_conv"):
+        # nhwc_hwcf
+        module_str = GENERIC_TEMPLATE.format(
+            lhs_type=ir.RankedTensorType.get([2, 7, 7, 32], ir.F16Type.get()),
+            rhs_type=ir.RankedTensorType.get([3, 3, 32, 64], ir.F16Type.get()),
+            res_type=ir.RankedTensorType.get([2, 5, 5, 64], ir.F32Type.get()),
+            lhs_map="affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 + d4, d2 + d5, d6)>",
+            rhs_map="affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d4, d5, d6, d3)>",
+            res_map="affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)>",
+            iterator_types='["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction"]',
+        )
+    module = ir.Module.parse(module_str, context)
+    root_op_list = iree_codegen.get_tuner_root_ops(module)
+    assert len(root_op_list) == 1
+    root_op = root_op_list[0]
+    parser = dispatch_parser.ConvolutionOpInterfaceParser(root_op)
+    assert parser.get_root_op_func_name() == "match_test"
+    assert parser.has_valid_root_op()
+
+    problem_size = parser.get_problem_size()
+    assert problem_size.contraction_dims.batch == []
+    assert problem_size.contraction_dims.m == [0, 1, 2]
+    assert problem_size.contraction_dims.n == [3]
+    assert problem_size.contraction_dims.k == [4, 5, 6]
+    assert problem_size.matmul_size.B == []
+    assert problem_size.matmul_size.M == [2, 5, 5]
+    assert problem_size.matmul_size.N == [64]
+    assert problem_size.matmul_size.K == [3, 3, 32]
+
 
 def test_get_mmt_tile_sizes(tuner_ctx: common.TunerContext) -> None:
     mma_intrinsic = iree_gpu.MMAIntrinsic.MFMA_F32_16x16x16_F16


Reviews:
user: login: bangtianliu
body: Overall, it looks good—I’ll review it more thoroughly later.

For now, I have one question: which convolution ops does the tuner support besides `conv_2d_nhwc_hwcf`?
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T18:28:22Z
user: login: rkayaith
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-08T18:20:53Z
user: login: bangtianliu
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-08T21:20:17Z
user: login: bangtianliu
body: Thanks for improving the tuner.
LGTM, it seems that you need to do a rebase. 
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-08T22:07:49Z

Review comments:
diff_hunk: @@ -247,6 +250,82 @@ def test_get_conv_operation(tuner_ctx: common.TunerContext) -> None:
         parser.has_valid_root_op()
     ), f"ConvolutionOpInterfaceParser does not support the op: {root_op.name}"
 
+    problem_size = parser.get_problem_size()
+    assert problem_size.contraction_dims.batch == []
+    assert problem_size.contraction_dims.m == [0, 1, 2]
+    assert problem_size.contraction_dims.n == [3]
+    assert problem_size.contraction_dims.k == [4, 5, 6]
+    assert problem_size.matmul_size.B == []
+    assert problem_size.matmul_size.M == [2, 32, 32]
+    assert problem_size.matmul_size.N == [16]
+    assert problem_size.matmul_size.K == [3, 3, 16]
+
+
+def test_get_group_conv_operation(tuner_ctx: common.TunerContext) -> None:
+    context = tuner_ctx.mlir_ctx
+    module_str = """
+    module {
+      func.func @test(%arg0: tensor<2x10x10x7x4xf32>, %arg1: tensor<7x16x3x3x4xf32>, %arg2: tensor<2x8x8x7x16xf32>) -> tensor<2x8x8x7x16xf32> {
+        %0 = linalg.conv_2d_nhwgc_gfhwc {
+           root_op,
+           dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>
+        } ins(%arg0, %arg1: tensor<2x10x10x7x4xf32>, tensor<7x16x3x3x4xf32>)
+          outs(%arg2: tensor<2x8x8x7x16xf32>) -> tensor<2x8x8x7x16xf32>
+        return %0 : tensor<2x8x8x7x16xf32>
+      }
+    }
+    """
+    module = ir.Module.parse(module_str, context)
+    module = ir.Module.parse(module_str, context)
path: sharktuner/tests/dispatch_parser_test.py
user: login: rkayaith
body: note to self: redundant line
created_at: 2025-05-08T18:20:53Z
author_association: MEMBER
diff_hunk: @@ -96,41 +96,64 @@ def get_problem_size(self) -> ProblemSize:
         )
 
 
-# TODO(Max191): Support more convolution types. Only NHWC convs are supported.
 class ConvolutionOpInterfaceParser(DispatchParser):
     def __init__(self, root_op: ir.Operation):
         super().__init__(root_op)
-        self.supported_ops = ["linalg.conv_2d_nhwc_hwcf"]
 
     def has_valid_root_op(self) -> bool:
         root_op = self.get_root_op()
         if not linalg.isa_convolution_op(root_op):
             return False
-
-        return root_op.name in self.supported_ops
+        convolution_dims = linalg.infer_convolution_dimensions(root_op)
+        assert convolution_dims, "no convolution dimensions"
+        # Only allow 'nhwc_hwcf' convs.
path: sharktuner/sharktuner/dispatch_parser.py
user: login: bangtianliu
body: Maybe also add a comment like "TODO: support more conv ops" in case we miss this. 
created_at: 2025-05-08T21:20:17Z
author_association: CONTRIBUTOR

PR comments:
user: login: rkayaith
created_at: 2025-05-07T20:27:19Z
body: > For now, I have one question: which convolution ops does the tuner support besides `conv_2d_nhwc_hwcf`?

Currently I'm only using the tuner on `nhwc_hwcf`, but represented as `linalg.generic`, so only the parser needs changes to support that. Next I'm looking at group convolutions, which needs more changes in the tuner (I'll be asking for help with that soon 😄). And after that, probably 3D convs.
author_association: MEMBER
user: login: bangtianliu
created_at: 2025-05-07T22:55:21Z
body: > > For now, I have one question: which convolution ops does the tuner support besides `conv_2d_nhwc_hwcf`?
> 
> Currently I'm only using the tuner on `nhwc_hwcf`, but represented as `linalg.generic`, so only the parser needs changes to support that. Next I'm looking at group convolutions, which needs more changes in the tuner (I'll be asking for help with that soon 😄). And after that, probably 3D convs.

Then how should we reflect that in the tuner? The current PR gives the impression that the tuner supports all convolution operations with different layouts.
author_association: CONTRIBUTOR
user: login: rkayaith
created_at: 2025-05-07T23:08:08Z
body: > Then how should we reflect that in the tuner? The current PR gives the impression that the tuner supports all convolution operations with different layouts.

If the configuration generator doesn't support something, I suppose it would either error out or produce no configs. I can add an error for the latter case.
(imo it would be best for any assumptions being made in the constraints generation to also be validated there, but I'm not familiar enough with the code to identify those places).
author_association: MEMBER
user: login: bangtianliu
created_at: 2025-05-07T23:17:18Z
body: > If the configuration generator doesn't support something, I suppose it would either error out or produce no configs. I can add an error for the latter case. (imo it would be best for any assumptions being made in the constraints generation to also be validated there, but I'm not familiar enough with the code to identify those places).

The tuner might either crash or proceed with tuning unnecessarily, which wastes time (I'm not sure, as I haven't tested it with other conv ops). But ideally, we should catch this case early to terminate the tuner. The contraction dimension information can often indicate the layout, so it can be used for filtering. Alternatively, you can just add a TODO for now — I can follow up with a separate PR to handle it properly.
author_association: CONTRIBUTOR
user: login: Max191
created_at: 2025-05-08T14:27:14Z
body: Fly by comment for right now. I'll do a full review later.

> Then how should we reflect that in the tuner?

We should be testing everything that we support, so whatever we decide on, we need to also add tests for it. The simplest thing for now would be to add some additional artificial constraints in `has_valid_root_op`, and we can slowly relax those as we test and support more types of layouts. This way we can reduce the size of each PR, which makes reviewing and landing easier.

We don't need to test every single permutation of layouts (a few key layouts that we regularly see would be good), but we want a couple representative tests each for things like 2D convs, 3D convs, group convs, etc. as we support each type.

For this PR, I think let's just support the 2D layouts. A good place to start for which layouts to test is in the IREE tests for conv_to_im2col (I see you already have some tests, which is great!): https://github.com/iree-org/iree/blob/main/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/test/conv_to_im2col.mlir
author_association: CONTRIBUTOR
user: login: bangtianliu
created_at: 2025-05-08T22:08:05Z
body: nvm lol
author_association: CONTRIBUTOR
user: login: rkayaith
created_at: 2025-05-08T22:50:07Z
body: Thanks for the reviews!
author_association: MEMBER