[tuner] add padding_conv attribute along IGEMM supprot for conv

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

Author: bangtianliu

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.
+    Used by get_padding_conv_sizes to calculate padding_conv attribute.
+
+    Corresponds to ConvToIgemmInfo struct in IREE:
+    https://github.com/iree-org/iree/blob/d3440737cc56a4d1b20c72181d9a37f194bd3ce5/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.cpp#L373-L379
+    """
+
+    conv_dims: linalg.ConvolutionDimensions
+    is_batch_dim_last: bool = False
+    is_spatial_dim_last: bool = False
+    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]
@@ -565,8 +600,103 @@ def get_dim_bounds(
     return result
 
 
-# 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 the 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]]:
+    """
+    Computes padding_conv by mapping padding from IGEMM space to convolution space.
+
+    Args:
+        bounds: Loop bounds for each dimension
+        padding_sizes: Padding sizes in IGEMM dimension space (M, N, K)
+        workgroup_tile_sizes: Workgroup tile sizes
+        reduction_tile_sizes: Reduction tile sizes
+        conv_to_igemm_info: Convolution to IGEMM transformation info
+
+    Returns:
+        Padding sizes in convolution dimension space, or None if no padding
+        is needed along original convolution dimensions.
+    """
+    # Skip padding convolution for NCHW layout (spatial dimensions are 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
+    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]
+        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
+
+    for conv_dim, igemm_pos in conv_to_igemm_map.items():
+        if reduction_tile_sizes[igemm_pos] != 0:
+            # Skip conv padding for reduction dims if already divisible by 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],

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,
@@ -259,6 +260,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]
@@ -271,6 +273,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
@@ -295,6 +309,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)))))
@@ -595,6 +610,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,
@@ -607,6 +624,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,
         )
 

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
 

amdsharktuner/amdsharktuner/dispatch_parser.py

@@ -40,6 +40,55 @@ def parse_mlir(mlir_text: str, ctx: common.TunerContext) -> ir.Module:
     return mlir_module
 
 
+def build_conv_to_igemm_info(
+    convolution_dims: linalg.ConvolutionDimensions,
+    input_type: ir.Type,
+    input_map: ir.AffineMap,
+    igemm_details,
+) -> common.ConvToIgemmInfo:
+    """
+    Builds ConvToIgemmInfo from convolution dimensions and IGEMM details.
+
+    Corresponds to IREE:
+    https://github.com/iree-org/iree/blob/d3440737cc56a4d1b20c72181d9a37f194bd3ce5/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.cpp#L872-L909
+    """
+    input_shape = input_type.shape
+    conv_to_igemm_info = common.ConvToIgemmInfo(conv_dims=convolution_dims)
+
+    # Map input channel dimensions to their sizes in the input tensor.
+    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)
+
+    input_image_pos = sorted(input_image_pos)
+    batch_pos = sorted(batch_pos)
+
+    conv_to_igemm_info.is_batch_dim_last = (
+        len(batch_pos) > 0 and batch_pos[-1] == len(input_shape) - 1
+    )
+    conv_to_igemm_info.is_spatial_dim_last = (
+        len(input_image_pos) > 0 and input_image_pos[-1] == len(input_shape) - 1
+    )
+
+    conv_to_igemm_info.conv_to_igemm_dim_map = dict(igemm_details.conv_to_igemm_dim_map)
+    return conv_to_igemm_info
+
+
 @dataclass
 class OpInfo:
     root_op: ir.Operation
@@ -74,6 +123,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 +326,13 @@ 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.
+        conv_to_igemm_info = None
+        if igemm_details:
+            conv_to_igemm_info = build_conv_to_igemm_info(
+                convolution_dims, lhs_type, indexing_maps[0], igemm_details
+            )
+
         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: