[tuner] add padding_conv attribute along IGEMM supprot for conv

bangtianliu · bangtianliu · commit ac90ffe8cad0 · 2025-12-01T17:22:09.000Z
Signed-off-by: Bangtian Liu &lt;liubangtian@gmail.com&gt;
diff --git a/amdsharktuner/amdsharktuner/common.py b/amdsharktuner/amdsharktuner/common.py
@@ -18,7 +18,7 @@
 import tempfile
 
 from iree.compiler import ir  # type: ignore
-from iree.compiler.dialects import iree_codegen, iree_gpu, transform  # type: ignore
+from iree.compiler.dialects import iree_codegen, iree_gpu, linalg, transform  # type: ignore
 import iree.compiler as ireec  # type: ignore
 from iree.compiler._mlir_libs._mlir import ir  # type: ignore
 
@@ -190,6 +190,23 @@ class ContractionDimensions:
     batch: list[int] = field(default_factory=list)
 
 
+@dataclass
+class ConvToIgemmInfo:
+    """
+    Stores information about convolution to IGEMM transformation.
+    This corresponds to the C++ ConvToIgemmInfo struct in IREE.
+
+    Note: In C++, conv_to_igemm_dim_map is DenseMap<int64_t, AffineExpr>,
+    but in Python bindings it's dict[int, int] mapping conv dim to IGEMM position.
+    """
+
+    is_batch_dim_last: bool = False
+    is_spatial_dim_last: bool = False
+    conv_dims: Optional[linalg.ConvolutionDimensions] = None
+    conv_to_igemm_dim_map: dict[int, int] = field(default_factory=dict)
+    input_channel_dim_to_size: dict[int, int] = field(default_factory=dict)
+
+
 @dataclass
 class MatmulShapeType:
     m: int
@@ -233,6 +250,24 @@ class AttentionKnobs(KnobAssignment):
     pass
 
 
+def is_affine_expr_function_of_dim(expr: ir.AffineExpr, position: int) -> bool:
+    """
+    Return True if the expression depends on the dimension at the given position.
+    """
+    if ir.AffineDimExpr.isinstance(expr):
+        dim_expr = ir.AffineDimExpr(expr)
+        return dim_expr.position == position
+
+    # Check if it's a binary operation and recursively check both sides.
+    if ir.AffineBinaryExpr.isinstance(expr):
+        binary_expr = ir.AffineBinaryExpr(expr)
+        return is_affine_expr_function_of_dim(
+            binary_expr.lhs, position
+        ) or is_affine_expr_function_of_dim(binary_expr.rhs, position)
+
+    return False
+
+
 def get_map_result_dim_positions(map: ir.AffineMap) -> Optional[list[int]]:
     if not map.is_projected_permutation:
         return None
@@ -281,7 +316,7 @@ def get_lowering_config(
         # A local variable to hold the transformed value.
         promoted_value = value
         match key:
-            case "workgroup" | "reduction" | "subgroup" | "promote_operands" | "padding":
+            case "workgroup" | "reduction" | "subgroup" | "promote_operands" | "padding" | "padding_conv":
                 if isinstance(value, Sequence):
                     promoted_value = ir.ArrayAttr.get(
                         [tuner_ctx.type.getI64(x) for x in value]
@@ -562,8 +597,112 @@ def get_dim_bounds(
     return result, any_padding_applied
 
 
-# Use padding logic from IREE side:
-# https://github.com/iree-org/iree/blob/8ae91ebb0e555e660b8a6898f6071476f7a1f20b/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.cpp#L691-L703
+# Implemented padding logic from IREE side:
+# https://github.com/iree-org/iree/blob/8ae91ebb0e555e660b8a6898f6071476f7a1f20b/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.cpp#L382-L467
+def get_padding_conv_sizes(
+    bounds: list[int],
+    padding_sizes: list[int],
+    workgroup_tile_sizes: list[int],
+    reduction_tile_sizes: list[int],
+    conv_to_igemm_info: ConvToIgemmInfo,
+) -> Optional[list[int]]:
+    """
+    Calculate padding sizes for convolution dimensions when using IGEMM.
+    This corresponds to C++ getPaddingConvSizes function in IREE.
+
+    Args:
+        bounds: Loop bounds for each dimension
+        padding_sizes: Padding sizes for IGEMM dimensions
+        workgroup_tile_sizes: Workgroup tile sizes
+        reduction_tile_sizes: Reduction tile sizes
+        conv_to_igemm_info: Convolution to IGEMM transformation info (must not be None)
+
+    Returns:
+        List of padding sizes for convolution dimensions, or None if padding should be skipped.
+        Caller should convert to ArrayAttr if needed.
+    """
+    # Skip padding convolution for NCHW layout (spatial dimension last).
+    if conv_to_igemm_info.is_spatial_dim_last:
+        return None
+
+    conv_to_igemm_map = conv_to_igemm_info.conv_to_igemm_dim_map
+    padded_igemm_dims = set()
+    conv_dims = conv_to_igemm_info.conv_dims
+
+    assert conv_dims is not None, "Expected conv_dims to be set in ConvToIgemmInfo"
+
+    input_channel_dims = set(conv_dims.input_channel)
+
+    padding_conv_sizes = [0] * len(conv_to_igemm_map)
+
+    # For batch-last layout (e.g., CHWN), only pad the batch dimension to avoid
+    # introducing pad op as the producer of collapse_shape op which may cause fusion problem.
+    if conv_to_igemm_info.is_batch_dim_last:
+        last_batch_dim = conv_dims.batch[-1]
+        # The map stores integer positions, use them directly.
+        igemm_batch_pos = conv_to_igemm_map[last_batch_dim]
+
+        if (
+            padding_sizes[igemm_batch_pos]
+            and bounds[igemm_batch_pos] % padding_sizes[igemm_batch_pos] == 0
+        ):
+            return None
+
+        padding_conv_sizes[last_batch_dim] = padding_sizes[igemm_batch_pos]
+        return padding_conv_sizes
+
+    # Process each convolution dimension mapping.
+    for conv_dim, igemm_pos in conv_to_igemm_map.items():
+        # The map stores integer positions directly.
+
+        if reduction_tile_sizes[igemm_pos] != 0:
+            # For reduction dimensions, avoid setting padding on the convolution
+            # if the product of the corresponding conv sizes are already divisible by the padding size.
+            if (
+                padding_sizes[igemm_pos]
+                and bounds[igemm_pos] % padding_sizes[igemm_pos] == 0
+            ):
+                padded_igemm_dims.add(igemm_pos)
+                continue
+
+            # Only pad input channel dims. If we need to pad filter dims, then we
+            # would rather just do padding on the IGEMM instead.
+            if conv_dim in input_channel_dims:
+                # Multiple input channel dims for a single IGEMMPos is not supported.
+                if igemm_pos in padded_igemm_dims:
+                    return None
+
+                input_channel_size = conv_to_igemm_info.input_channel_dim_to_size.get(
+                    conv_dim, 0
+                )
+                is_input_channel_size_small = (
+                    padding_sizes[igemm_pos] // input_channel_size > 2
+                )
+
+                # If the input channel dimension is much smaller than the padding size,
+                # skip padding along that dimension while still padding the others.
+                if is_input_channel_size_small:
+                    padding_conv_sizes[conv_dim] = 0
+                else:
+                    padding_conv_sizes[conv_dim] = padding_sizes[igemm_pos]
+
+                padded_igemm_dims.add(igemm_pos)
+            continue
+
+        # Multiple padded parallel dims mapping to the same IGEMM dim is not supported.
+        if workgroup_tile_sizes[igemm_pos] != 0 and igemm_pos in padded_igemm_dims:
+            return None
+
+        padding_conv_sizes[conv_dim] = padding_sizes[igemm_pos]
+        padded_igemm_dims.add(igemm_pos)
+
+    # Ensure that all dimensions have been padded.
+    if len(padded_igemm_dims) != len(padding_sizes):
+        return None
+
+    return padding_conv_sizes
+
+
 def calculate_padded_dimensions(
     M: list[int],
     N: list[int],
diff --git a/amdsharktuner/amdsharktuner/constraint_generator.py b/amdsharktuner/amdsharktuner/constraint_generator.py
@@ -82,6 +82,7 @@ def generate_generic_contraction_solutions(
     allowed_waves_per_eu: list[int] = [2],
     pipeline_options_search_space: dispatch_constraints.PipelineOptionsSearchSpace = dispatch_constraints.PipelineOptionsSearchSpace(),
     igemm_details: Optional[iree_codegen.IGEMMGenericConvDetails] = None,
+    conv_to_igemm_info: Optional[common.ConvToIgemmInfo] = None,
 ) -> Iterator[list[common.TuningConfiguration]]:
     adjust_problem_size_for_pipeline(
         contraction_dims,
@@ -258,6 +259,7 @@ def set_cdim_tile_sizes(tile_sizes, contraction_dims, csizes):
 
         promote_operands = [0, 1]
         padding = None
+        padding_conv = None
         if padding_applied:
             # TODO: Remove promotion of operand 2 once codegen supports handling padded outputs without promotion.
             promote_operands = [0, 1, 2]
@@ -270,6 +272,18 @@ def set_cdim_tile_sizes(tile_sizes, contraction_dims, csizes):
             padding_tile_sizes[inner_k_dim] *= mma_intrinsic_k
 
             padding = padding_tile_sizes
+
+            # Calculate padding_conv sizes for convolutions when using IGEMM.
+            if conv_to_igemm_info and igemm_details:
+                # Use IGEMM loop bounds directly from igemm_details.
+                bounds = list(igemm_details.igemm_loop_bounds)
+                padding_conv = common.get_padding_conv_sizes(
+                    bounds,
+                    padding_tile_sizes,
+                    workgroup_tile_sizes,
+                    reduction_tile_sizes,
+                    conv_to_igemm_info,
+                )
         # Setting subgroup basis.
         # TODO(Bangtian): Sync changes from IREE PR: https://github.com/iree-org/iree/pull/22000.
         subgroup_basis_counts = [1] * num_loops
@@ -294,6 +308,7 @@ def set_cdim_tile_sizes(tile_sizes, contraction_dims, csizes):
             pipeline_options_search_space,
             allowed_waves_per_eu,
             padding=padding,
+            padding_conv=padding_conv,
         )
 
         solver.add(z3.simplify(z3.Not(z3.And(list(x == model[x] for x in all_vars)))))
@@ -594,6 +609,8 @@ def generate_solutions(
         codegen_pipeline: iree_codegen.DispatchLoweringPassPipeline,
         **pipeline_constraint_options,
     ) -> Iterator[list[common.TuningConfiguration]]:
+        # TODO(Bangtian): Simplify the function signature to accept op_info directly instead of
+        # unpacking all individual fields.
         return generate_generic_contraction_solutions(
             tuner_ctx=tuner_context,
             gpu_target_info=gpu_target_info,
@@ -606,6 +623,7 @@ def generate_solutions(
             indexing_maps=self.op_info.indexing_maps,
             codegen_pipeline=codegen_pipeline,
             igemm_details=self.op_info.igemm_details,
+            conv_to_igemm_info=self.op_info.conv_to_igemm_info,
             **pipeline_constraint_options,
         )
 
diff --git a/amdsharktuner/amdsharktuner/dispatch_constraints.py b/amdsharktuner/amdsharktuner/dispatch_constraints.py
@@ -672,6 +672,7 @@ def generate_compilation_infos(
     pipeline_options_search_space: PipelineOptionsSearchSpace,
     allowed_waves_per_eu: list[int],
     padding: Optional[list[int]] = None,
+    padding_conv: Optional[list[int]] = None,
 ) -> list[iree_codegen.CompilationInfoAttr]:
     subgroup_basis = [subgroup_basis_counts, subgroup_basis_mapping]
     # Create the LoweringConfigAttr.
@@ -688,6 +689,9 @@ def generate_compilation_infos(
     if padding is not None:
         lowering_config_args["padding"] = padding
 
+    if padding_conv is not None:
+        lowering_config_args["padding_conv"] = padding_conv
+
     if codegen_pipeline == iree_codegen.DispatchLoweringPassPipeline.LLVMGPUTileAndFuse:
         lowering_config_args["subgroup"] = subgroup_tile_sizes
 
diff --git a/amdsharktuner/amdsharktuner/dispatch_parser.py b/amdsharktuner/amdsharktuner/dispatch_parser.py
@@ -74,6 +74,8 @@ class ConvolutionOpInfo(OpInfo):
 
     # IGEMM details for TileAndFuse pipeline (None if not available).
     igemm_details: Optional[iree_codegen.IGEMMGenericConvDetails] = None
+    # Convolution to IGEMM transformation info (None if not available).
+    conv_to_igemm_info: Optional[common.ConvToIgemmInfo] = None
 
 
 @dataclass
@@ -275,6 +277,62 @@ def __init__(self, root_op: ir.Operation, tuner_ctx: common.TunerContext):
         # for any convolution layout (nhwc_hwcf, nchw_fchw, etc.).
         igemm_details = iree_codegen.get_igemm_generic_conv_details(root_op)
 
+        # Build ConvToIgemmInfo using convolution_dims (similar to C++ implementation).
+        conv_to_igemm_info = None
+        if igemm_details:
+            conv_to_igemm_info = common.ConvToIgemmInfo()
+
+            # Get input operand information (first operand).
+            input_type = lhs_type
+            input_shape = input_type.shape
+            input_map = indexing_maps[0]
+
+            # Store the convolution dimensions.
+            conv_to_igemm_info.conv_dims = convolution_dims
+
+            # Process input channel dimensions.
+            # Note: For convolutions with strides, expressions may be complex (e.g., d0*2 + d1).
+            # We use isFunctionOfDim to check if the expression depends on a specific dimension.
+            for dim in convolution_dims.input_channel:
+                for idx, expr in enumerate(input_map.results):
+                    if common.is_affine_expr_function_of_dim(expr, dim):
+                        conv_to_igemm_info.input_channel_dim_to_size[dim] = input_shape[
+                            idx
+                        ]
+
+            # Process output image dimensions to find input image positions.
+            input_image_pos = []
+            for dim in convolution_dims.output_image:
+                for idx, expr in enumerate(input_map.results):
+                    if common.is_affine_expr_function_of_dim(expr, dim):
+                        input_image_pos.append(idx)
+
+            # Process batch dimensions to find batch positions.
+            batch_pos = []
+            for dim in convolution_dims.batch:
+                for idx, expr in enumerate(input_map.results):
+                    if common.is_affine_expr_function_of_dim(expr, dim):
+                        batch_pos.append(idx)
+
+            # Sort positions.
+            input_image_pos = sorted(input_image_pos)
+            batch_pos = sorted(batch_pos)
+
+            # Determine if batch dimension is last.
+            conv_to_igemm_info.is_batch_dim_last = (
+                len(batch_pos) > 0 and batch_pos[-1] == len(input_shape) - 1
+            )
+
+            # Determine if spatial dimension is last.
+            conv_to_igemm_info.is_spatial_dim_last = (
+                len(input_image_pos) > 0 and input_image_pos[-1] == len(input_shape) - 1
+            )
+
+            # Store conv to IGEMM dimension mapping from IGEMM details.
+            conv_to_igemm_info.conv_to_igemm_dim_map = dict(
+                igemm_details.conv_to_igemm_dim_map
+            )
+
         self._op_info: ConvolutionOpInfo = ConvolutionOpInfo(
             root_op=root_op,
             indexing_maps=indexing_maps,
@@ -292,6 +350,7 @@ def __init__(self, root_op: ir.Operation, tuner_ctx: common.TunerContext):
             strides=strides,
             dilations=dilations,
             igemm_details=igemm_details,
+            conv_to_igemm_info=conv_to_igemm_info,
         )
 
     def has_valid_root_op(self) -> bool:
diff --git a/amdsharktuner/tests/common_test.py b/amdsharktuner/tests/common_test.py
diff --git a/amdsharktuner/tests/constraint_generator_test.py b/amdsharktuner/tests/constraint_generator_test.py
