ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-06-28 18:56:59 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
develop
composable_kernel/dispatcher/tests/test_grouped_conv_codegen.py
Ville Pietilä 60b276647b [rocm-libraries] ROCm/rocm-libraries#8157 (commit b0d9d39) 
[CK Tile] Rule-based configuration generation in CK
 Dispatcher codegen (#8157)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Motivation

The CK Tile Dispatcher code generation for CK Tile Profiler relies on
flat JSON files to list the generated configurations. This approach has
the following problems

- The JSON files are verbose
- The JSON files get easily out of sync with the CK Builder .config
files from which they were generated from.
- The JSON file based configuration make it hard to list explicitly the
rules that govern the instance generation.

## Technical Details

Replaced the JSON files with a rule based configuration. To preserve the
existing functionality, the `profiler` and the `tests` instance sets are
generated directly from the CK Builder config files. The JSON config
files are removed from source control, and the "on-the-fly" generation
guarantees that the Dispatcher codegen uses up to date configurations.

This is PR introduces six different rule sets for the CK Tile Dispatcher
code generation

1. `profiler`: matches with the old JSON set of profiler configurations.
2. `tests`: matches with the old JSON set of tests configurations.
3. `full`: full configuration set created from a rule-based config
selection
4. `full-tests`: a subset of `full` for generating configurations for
convolution integration tests.
5. `tiny`: a subset of `full-tests` to produce the minimal set of
configurations to test the Dispatcher codegen.
6. `default`: the default rules, which corresponds to the existing
heuristic rules for configuration selection. This ensures that ML based
kernel selection doesn't get broken.

The main use of the `full` rule set is to define a reasonable solution
space for the possible implicit GEMM configurations. We start from the
configurations that allowed by the device architecture. The `full` rule
set defines the relevant tile sizes for each convolution direction. From
the tile size we have a curated mapping to the number of waves over the
different GEMM axes, i.e., we describe how many waves each GEMM
dimensions corresponds to. The GEMM-K wave tile dimension can be
computed from the other parameters and does not need to be listed
explicitly.

An orthogonal axis to the tiling strategy is the vectorization strategy.
This mainly defined by the data type and hardware as in general, we want
to use the maximum possible load widths. The maximum sizes for each
convolution direction variant are defined by the implicit GEMM matrix
dimensions. For cases where have a low number of channels per
convolution group, we need smaller vector load sizes. These are captured
by the `VecStrategy` enumeration in the codegen rules.

The problem with the rule based configuration selection is that we "over
generate" configurations. The old JSON configurations compose
approximately 25% of all configuration that the `full` rule set creates.
The additional configurations are valid, but they many not provide any
performance benefits. Hence, we keep the `profiler` and `tests` rule set
for now to avoid building an excessive amount configurations by default.
The `full` rule set can be taken into use by specifying CMake
configuration flag `-D DISPATCHER_RULE_SET=full`. By default, the
`tests` rule set is used, i.e., we don't change the existing bahaviour.

## Test Plan

Added a new stage in the CI/CD pipeline that ensures the Dispatcher
codegen rules are up to date. Otherwise the functionality is covered by
the existing CI/CD tests. There are no functional changes to the
convolution kernels. Only how the different instances are generated.

## Test Result

If the CK Tile conv instances build without errors, the Dispatcher
codegen is generating valid code. If all tests in CI/CD pipeline are
passing, the Dispatcher codegen generates valid instances.

## Submission Checklist

- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
2026-06-18 01:22:50 +00:00

593 lines
21 KiB
Python
Raw Permalink Blame History

#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
TDD tests for codegen/unified_grouped_conv_codegen.py -- grouped convolution code generator.
These tests are written BEFORE the implementation exists.
Run: python3 -m pytest dispatcher/tests/test_grouped_conv_codegen.py -v
"""
import sys
import unittest
from pathlib import Path
from unittest.mock import patch
SCRIPT_DIR = Path(__file__).parent.resolve()
DISPATCHER_DIR = SCRIPT_DIR.parent
sys.path.insert(0, str(DISPATCHER_DIR / "codegen"))
sys.path.insert(0, str(DISPATCHER_DIR / "python"))
from codegen_common import TileConfig, TraitConfigBase # noqa: E402
from unified_grouped_conv_codegen import ( # noqa: E402
GroupedConvVariant,
GroupedConvLayout,
GroupedConvKernelConfig,
GroupedConvTypeMappings,
GroupedConvTraitConfig,
CKTileGroupedConvKernelGenerator,
GroupedConvDispatcherWrapperGenerator,
UnifiedGroupedConvCodegen,
)
# =============================================================================
# TestGroupedConvVariant
# =============================================================================
class TestGroupedConvVariant(unittest.TestCase):
"""Test GroupedConvVariant enum values."""
def test_forward_value(self):
self.assertEqual(GroupedConvVariant.FORWARD.value, "forward")
def test_backward_data_value(self):
self.assertEqual(GroupedConvVariant.BACKWARD_DATA.value, "bwd_data")
def test_backward_weight_value(self):
self.assertEqual(GroupedConvVariant.BACKWARD_WEIGHT.value, "bwd_weight")
def test_all_variants_exist(self):
self.assertIn(GroupedConvVariant.FORWARD, GroupedConvVariant)
self.assertIn(GroupedConvVariant.BACKWARD_DATA, GroupedConvVariant)
self.assertIn(GroupedConvVariant.BACKWARD_WEIGHT, GroupedConvVariant)
# =============================================================================
# TestGroupedConvLayout
# =============================================================================
class TestGroupedConvLayout(unittest.TestCase):
"""Test GroupedConvLayout enum for 1D/2D/3D layouts."""
def test_nhwgc_value(self):
self.assertEqual(GroupedConvLayout.NHWGC.value, "NHWGC")
def test_gkyxc_value(self):
self.assertEqual(GroupedConvLayout.GKYXC.value, "GKYXC")
def test_nhwgk_value(self):
self.assertEqual(GroupedConvLayout.NHWGK.value, "NHWGK")
def test_1d_layouts_exist(self):
"""1D conv layouts (e.g., NWGC, GYXC, NWGK)."""
layouts_1d = [
lay
for lay in GroupedConvLayout
if "W" in lay.value and "H" not in lay.value
]
self.assertGreater(len(layouts_1d), 0)
def test_2d_layouts_exist(self):
"""2D conv layouts (e.g., NHWGC, GKYXC, NHWGK)."""
layouts_2d = [lay for lay in GroupedConvLayout if "HW" in lay.value]
self.assertGreater(len(layouts_2d), 0)
def test_3d_layouts_exist(self):
"""3D conv layouts (e.g., NDHWGC, GDKYXC)."""
layouts_3d = [
lay for lay in GroupedConvLayout if "D" in lay.value or "DHW" in lay.value
]
self.assertGreater(len(layouts_3d), 0)
# =============================================================================
# TestGroupedConvKernelConfig
# =============================================================================
class TestGroupedConvKernelConfig(unittest.TestCase):
"""Test GroupedConvKernelConfig dataclass."""
def _make_tile(self):
return TileConfig(128, 128, 32, 2, 2, 1, 32, 32, 16)
def _make_trait(self):
return GroupedConvTraitConfig(
"mem",
"cshuffle",
"intrawave",
False,
False,
False,
double_smem_buffer=False,
num_groups_to_merge=1,
)
def test_name_contains_grouped_conv_fwd(self):
config = GroupedConvKernelConfig(
tile=self._make_tile(),
trait=self._make_trait(),
variant=GroupedConvVariant.FORWARD,
ndim_spatial=2,
arch="gfx942",
layout=GroupedConvLayout.NHWGC,
vector_sizes=(4, 4, 4),
)
name = config.name("fp16")
self.assertIn("grouped_conv_fwd", name)
def test_name_backward_data_contains_bwd_data(self):
config = GroupedConvKernelConfig(
tile=self._make_tile(),
trait=self._make_trait(),
variant=GroupedConvVariant.BACKWARD_DATA,
ndim_spatial=2,
arch="gfx942",
layout=GroupedConvLayout.NHWGC,
vector_sizes=(4, 4, 4),
)
name = config.name("fp16")
self.assertIn("bwd_data", name)
def test_is_valid_for_arch_supported(self):
config = GroupedConvKernelConfig(
tile=self._make_tile(),
trait=self._make_trait(),
variant=GroupedConvVariant.FORWARD,
ndim_spatial=2,
arch="gfx942",
layout=GroupedConvLayout.NHWGC,
vector_sizes=(4, 4, 4),
)
self.assertTrue(config.is_valid_for_arch("gfx942"))
def test_is_valid_for_arch_unsupported(self):
config = GroupedConvKernelConfig(
tile=self._make_tile(),
trait=self._make_trait(),
variant=GroupedConvVariant.FORWARD,
ndim_spatial=2,
arch="gfx942",
layout=GroupedConvLayout.NHWGC,
vector_sizes=(4, 4, 4),
)
self.assertFalse(config.is_valid_for_arch("gfx600"))
# =============================================================================
# TestGroupedConvTypeMappings
# =============================================================================
class TestGroupedConvTypeMappings(unittest.TestCase):
"""Test GroupedConvTypeMappings class."""
def test_dtype_to_ck_fp16(self):
self.assertEqual(GroupedConvTypeMappings.DTYPE_TO_CK["fp16"], "half_t")
def test_dtype_to_ck_bf16(self):
self.assertIn("bf16", GroupedConvTypeMappings.DTYPE_TO_CK)
def test_dtype_to_ck_fp32(self):
self.assertIn("fp32", GroupedConvTypeMappings.DTYPE_TO_CK)
def test_get_layouts_2d_has_in_wei_out_keys(self):
layouts = GroupedConvTypeMappings.get_layouts(2)
self.assertIn("in", layouts)
self.assertIn("wei", layouts)
self.assertIn("out", layouts)
def test_get_layouts_2d_returns_dict(self):
layouts = GroupedConvTypeMappings.get_layouts(2)
self.assertIsInstance(layouts, dict)
def test_get_layouts_1d(self):
layouts = GroupedConvTypeMappings.get_layouts(1)
self.assertIn("in", layouts)
self.assertIn("wei", layouts)
self.assertIn("out", layouts)
def test_get_layouts_3d(self):
layouts = GroupedConvTypeMappings.get_layouts(3)
self.assertIn("in", layouts)
self.assertIn("wei", layouts)
self.assertIn("out", layouts)
# =============================================================================
# TestCKTileGroupedConvKernelGenerator
# =============================================================================
class TestCKTileGroupedConvKernelGenerator(unittest.TestCase):
"""Test CKTileGroupedConvKernelGenerator.generate()."""
def _make_config(self):
tile = TileConfig(128, 128, 32, 2, 2, 1, 32, 32, 16)
trait = GroupedConvTraitConfig(
"mem",
"cshuffle",
"intrawave",
False,
False,
False,
double_smem_buffer=False,
num_groups_to_merge=1,
)
return GroupedConvKernelConfig(
tile=tile,
trait=trait,
variant=GroupedConvVariant.FORWARD,
ndim_spatial=2,
arch="gfx942",
layout=GroupedConvLayout.NHWGC,
vector_sizes=(4, 4, 4),
)
def test_generate_contains_pragma_once(self):
gen = CKTileGroupedConvKernelGenerator("fp16")
config = self._make_config()
result = gen.generate(config)
self.assertIn("#pragma once", result)
def test_generate_contains_forward_kernel_include(self):
gen = CKTileGroupedConvKernelGenerator("fp16")
config = self._make_config()
result = gen.generate(config)
self.assertIn("grouped_convolution_forward_kernel.hpp", result)
def test_generate_returns_non_empty_string(self):
gen = CKTileGroupedConvKernelGenerator("fp16")
config = self._make_config()
result = gen.generate(config)
self.assertIsInstance(result, str)
self.assertGreater(len(result), 100)
def test_generate_valid_cpp_structure(self):
gen = CKTileGroupedConvKernelGenerator("fp16")
config = self._make_config()
result = gen.generate(config)
self.assertIn("#include", result)
self.assertIn("ck_tile", result)
# =============================================================================
# TestGroupedConvDispatcherWrapperGenerator
# =============================================================================
class TestGroupedConvDispatcherWrapperGenerator(unittest.TestCase):
"""Test GroupedConvDispatcherWrapperGenerator.generate()."""
def _make_config(self):
tile = TileConfig(128, 128, 32, 2, 2, 1, 32, 32, 16)
trait = GroupedConvTraitConfig(
"mem",
"cshuffle",
"intrawave",
False,
False,
False,
double_smem_buffer=False,
num_groups_to_merge=1,
)
return GroupedConvKernelConfig(
tile=tile,
trait=trait,
variant=GroupedConvVariant.FORWARD,
ndim_spatial=2,
arch="gfx942",
layout=GroupedConvLayout.NHWGC,
vector_sizes=(4, 4, 4),
)
def test_generate_contains_dispatcher_registration(self):
gen = GroupedConvDispatcherWrapperGenerator("fp16")
config = self._make_config()
kernel_path = DISPATCHER_DIR / "build" / "generated" / "test_kernel.hpp"
output_dir = DISPATCHER_DIR / "build" / "generated"
result = gen.generate(config, kernel_path, output_dir)
self.assertIn("dispatcher", result)
self.assertIn("KernelKey", result)
self.assertIn("KernelInstancePtr", result)
def test_generate_contains_pragma_once(self):
gen = GroupedConvDispatcherWrapperGenerator("fp16")
config = self._make_config()
kernel_path = DISPATCHER_DIR / "build" / "generated" / "test_kernel.hpp"
output_dir = DISPATCHER_DIR / "build" / "generated"
result = gen.generate(config, kernel_path, output_dir)
self.assertIn("#pragma once", result)
def test_generate_valid_cpp(self):
gen = GroupedConvDispatcherWrapperGenerator("fp16")
config = self._make_config()
kernel_path = DISPATCHER_DIR / "build" / "generated" / "test_kernel.hpp"
output_dir = DISPATCHER_DIR / "build" / "generated"
result = gen.generate(config, kernel_path, output_dir)
self.assertIn("#include", result)
self.assertIn("namespace", result)
# =============================================================================
# TestUnifiedGroupedConvCodegen
# =============================================================================
class TestUnifiedGroupedConvCodegen(unittest.TestCase):
"""Test UnifiedGroupedConvCodegen.generate_all()."""
def test_generate_all_returns_dict_with_expected_keys(self):
output_dir = DISPATCHER_DIR / "build" / "generated" / "grouped_conv"
output_dir.mkdir(parents=True, exist_ok=True)
codegen = UnifiedGroupedConvCodegen(
output_dir=output_dir,
datatype="fp16",
ndim_spatial=2,
gpu_target="gfx942",
)
with patch.object(
codegen,
"_get_configs",
return_value=[], # Mock empty config list for fast test
):
results = codegen.generate_all(parallel=False)
self.assertIn("kernels", results)
self.assertIn("wrappers", results)
self.assertIn("failed", results)
self.assertIsInstance(results["kernels"], list)
self.assertIsInstance(results["wrappers"], list)
self.assertIsInstance(results["failed"], list)
def test_generate_all_with_mock_config_produces_output(self):
output_dir = DISPATCHER_DIR / "build" / "generated" / "grouped_conv_test"
output_dir.mkdir(parents=True, exist_ok=True)
codegen = UnifiedGroupedConvCodegen(
output_dir=output_dir,
datatype="fp16",
ndim_spatial=2,
gpu_target="gfx942",
)
# Use a real config - patch the config source to return one config
tile = TileConfig(128, 128, 32, 2, 2, 1, 32, 32, 16)
trait = GroupedConvTraitConfig(
"mem",
"cshuffle",
"intrawave",
False,
False,
False,
double_smem_buffer=False,
num_groups_to_merge=1,
)
config = GroupedConvKernelConfig(
tile=tile,
trait=trait,
variant=GroupedConvVariant.FORWARD,
ndim_spatial=2,
arch="gfx942",
layout=GroupedConvLayout.NHWGC,
vector_sizes=(4, 4, 4),
)
with patch.object(codegen, "_get_configs", return_value=[config]):
results = codegen.generate_all(parallel=False)
self.assertIsInstance(results, dict)
self.assertIn("kernels", results)
# =============================================================================
# TestSharedImports
# =============================================================================
class TestSharedImports(unittest.TestCase):
"""Verify TileConfig from codegen_common and GroupedConvTraitConfig extends TraitConfigBase."""
def test_tile_config_has_expected_fields(self):
"""TileConfig from codegen_common has tile_m, tile_n, tile_k, etc."""
tc = TileConfig(128, 128, 32, 2, 2, 1, 32, 32, 16)
self.assertEqual(tc.tile_m, 128)
self.assertEqual(tc.tile_n, 128)
self.assertEqual(tc.tile_k, 32)
self.assertEqual(tc.warp_m, 2)
self.assertEqual(tc.warp_n, 2)
self.assertEqual(tc.warp_k, 1)
self.assertEqual(tc.warp_tile_m, 32)
self.assertEqual(tc.warp_tile_n, 32)
self.assertEqual(tc.warp_tile_k, 16)
def test_tile_config_is_from_codegen_common(self):
"""TileConfig used by grouped conv is the same as codegen_common.TileConfig."""
tc = TileConfig(128, 128, 32, 2, 2, 1, 32, 32, 16)
self.assertTrue(tc.is_valid())
def test_grouped_conv_trait_config_extends_trait_config_base(self):
"""GroupedConvTraitConfig extends TraitConfigBase."""
self.assertTrue(issubclass(GroupedConvTraitConfig, TraitConfigBase))
def test_grouped_conv_trait_config_has_double_smem_buffer(self):
"""GroupedConvTraitConfig has double_smem_buffer field."""
trait = GroupedConvTraitConfig(
"mem",
"cshuffle",
"intrawave",
False,
False,
False,
double_smem_buffer=True,
num_groups_to_merge=2,
)
self.assertTrue(trait.double_smem_buffer)
self.assertEqual(trait.num_groups_to_merge, 2)
def test_grouped_conv_trait_config_has_num_groups_to_merge(self):
"""GroupedConvTraitConfig has num_groups_to_merge field."""
trait = GroupedConvTraitConfig(
"mem",
"cshuffle",
"intrawave",
False,
False,
False,
double_smem_buffer=False,
num_groups_to_merge=4,
)
self.assertEqual(trait.num_groups_to_merge, 4)
def test_grouped_conv_trait_config_inherits_base_fields(self):
"""GroupedConvTraitConfig inherits pipeline, epilogue, scheduler from base."""
trait = GroupedConvTraitConfig(
"compv4",
"cshuffle",
"intrawave",
True,
True,
True,
double_smem_buffer=False,
num_groups_to_merge=1,
)
self.assertEqual(trait.pipeline, "compv4")
self.assertEqual(trait.epilogue, "cshuffle")
self.assertEqual(trait.scheduler, "intrawave")
self.assertTrue(trait.pad_m)
self.assertTrue(trait.pad_n)
self.assertTrue(trait.pad_k)
# =============================================================================
# TestTwoStageBwdWeightCodegen
# =============================================================================
def _make_two_stage_config():
"""Helper: create a two-stage bwd_weight config."""
return GroupedConvKernelConfig(
tile=TileConfig(16, 64, 64, 1, 4, 1, 16, 16, 32),
trait=GroupedConvTraitConfig(
pipeline="compv3",
epilogue="cshuffle",
scheduler="intrawave",
pad_m=True,
pad_n=True,
pad_k=True,
two_stage=True,
),
variant=GroupedConvVariant.BACKWARD_WEIGHT,
ndim_spatial=2,
arch="gfx942",
)
class TestTwoStageBwdWeightCodegen(unittest.TestCase):
"""Tests for two-stage backward weight kernel generation."""
def test_kernel_name_contains_2stage(self):
config = _make_two_stage_config()
name = config.name("fp16")
self.assertIn("_2stage", name)
self.assertIn("bwd_weight", name)
def test_single_stage_name_has_no_2stage(self):
config = _make_two_stage_config()
config.trait.two_stage = False
name = config.name("fp16")
self.assertNotIn("_2stage", name)
def test_generate_contains_elementwise_include(self):
config = _make_two_stage_config()
gen = CKTileGroupedConvKernelGenerator(
"fp16", GroupedConvVariant.BACKWARD_WEIGHT
)
code = gen.generate(config)
self.assertIn("elementwise.hpp", code)
def test_generate_contains_workspace_type(self):
config = _make_two_stage_config()
gen = CKTileGroupedConvKernelGenerator(
"fp16", GroupedConvVariant.BACKWARD_WEIGHT
)
code = gen.generate(config)
self.assertIn("WorkspaceDataType", code)
def test_generate_contains_elementwise_kernel(self):
config = _make_two_stage_config()
gen = CKTileGroupedConvKernelGenerator(
"fp16", GroupedConvVariant.BACKWARD_WEIGHT
)
code = gen.generate(config)
self.assertIn("ElementWiseKernel", code)
def test_generate_contains_launch_kernel_time_mask(self):
config = _make_two_stage_config()
gen = CKTileGroupedConvKernelGenerator(
"fp16", GroupedConvVariant.BACKWARD_WEIGHT
)
code = gen.generate(config)
self.assertIn("launch_kernel_time_mask", code)
def test_two_stage_uses_fp32_workspace_vector_size_c(self):
# Two-stage writes the GEMM result to an fp32 workspace, so it uses the
# configured VectorSizeC directly instead of forcing it to 1.
config = _make_two_stage_config()
gen = CKTileGroupedConvKernelGenerator(
"fp16", GroupedConvVariant.BACKWARD_WEIGHT
)
code = gen.generate(config)
self.assertIn("WorkspaceDataType = float", code)
self.assertIn("Config::VectorSizeC", code)
self.assertNotIn("VectorSizeC_TwoStage", code)
def test_generate_contains_workspace_memset(self):
config = _make_two_stage_config()
gen = CKTileGroupedConvKernelGenerator(
"fp16", GroupedConvVariant.BACKWARD_WEIGHT
)
code = gen.generate(config)
self.assertIn("hipMemsetAsync", code)
def test_single_stage_does_not_contain_workspace(self):
config = _make_two_stage_config()
config.trait.two_stage = False
gen = CKTileGroupedConvKernelGenerator(
"fp16", GroupedConvVariant.BACKWARD_WEIGHT
)
code = gen.generate(config)
self.assertNotIn("WorkspaceDataType", code)
self.assertNotIn("ElementWiseKernel", code)
def test_default_configs_include_two_stage(self):
from unified_grouped_conv_codegen import get_default_configs
configs = get_default_configs(
arch="gfx942",
variants=[GroupedConvVariant.BACKWARD_WEIGHT],
ndims=[2],
)
two_stage = [c for c in configs if c.trait.two_stage]
single_stage = [c for c in configs if not c.trait.two_stage]
self.assertGreater(len(two_stage), 0, "Should have two-stage configs")
self.assertGreater(
len(single_stage), 0, "Should still have single-stage configs"
)
if __name__ == "__main__":
unittest.main()
Reference in New Issue View Git Blame Copy Permalink