From 096592eb995f819da20caa1efc7a874e55b34220 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ville=20Pietil=C3=A4?= <> Date: Wed, 14 Jan 2026 08:29:49 -0500 Subject: [PATCH] Refactor conv algorithms into more categorized form. --- .../include/ck_tile/builder/factory/README.md | 228 ++++++++++++++++++ .../builder/factory/conv_algorithms.hpp | 89 +++---- .../test/impl/conv_algorithm_types.hpp | 4 +- 3 files changed, 268 insertions(+), 53 deletions(-) diff --git a/experimental/builder/include/ck_tile/builder/factory/README.md b/experimental/builder/include/ck_tile/builder/factory/README.md index d1794349ab..4c30348f9a 100644 --- a/experimental/builder/include/ck_tile/builder/factory/README.md +++ b/experimental/builder/include/ck_tile/builder/factory/README.md @@ -29,3 +29,231 @@ using Factory = decltype(make_conv_instance()); ``` The dispatcher automatically selects the appropriate factory following explicit logic. + +# Convolution Algorithm Hierarchy + +This section illustrates the hierarchy of convolution algorithm concepts defined in `conv_algorithms.hpp`. + +## Overview + +The convolution algorithms are organized into three main categories: + +1. **XDL Algorithms** - GPU matrix multiplication using XDL (matrix core instructions) +2. **WMMA Algorithms** - GPU matrix multiplication using WMMA (Wave Matrix Multiply-Accumulate) +3. **DL Algorithms** - Special vectorized dot-product kernels optimized for specific data layouts with separate implementation. + +XDL and WMMA algorithms share a common base, while DL algorithms have their own independent base. + +## Common Base Hierarchy (XDL & WMMA) + +Both XDL and WMMA algorithms share the following foundational concepts: + +``` +ConvAlgorithm (Base Concept) +│ +│ Requirements: +│ • ConvAlgorithmDescriptor +│ • SpecifiesThreadBlock +│ • SpecifiesTileTransferParameters (ThreadClusters, LdsTransfer, AccessOrders) +│ • SpecifiesWarpGemm +│ • SpecifiesGemmPipeline +│ +├─── FwdAlgorithm (Forward Convolution) +│ │ +│ │ Additional: SpecifiesFwdConvSpecialization +│ │ +│ └─── FwdAlgorithmV3 +│ │ +│ │ Additional: SpecifiesPipelineV3 +│ │ +│ +└─── BwdAlgorithm (Backward Weight Convolution) + │ + │ Additional: SpecifiesBwdWeightConvSpecialization + │ + └─── BwdAlgorithmV3 + │ + │ Additional: SpecifiesPipelineV3 + │ +``` + +--- + +## XDL Algorithm Hierarchy + +### Forward XDL Algorithms + +``` +FwdAlgorithm + SpecifiesXdl +│ +├─── FwdXdlAlgorithmBase + │ + ├─── FwdXdlAlgorithm + │ │ + │ └─ Requirements: Base + SpecifiesGenericInstance + │ + ├─── LargeTensorAlgorithm + │ │ + │ └─ Requirements: Base + SpecifiesLargeTensorSupport + │ + └─── FwdXdlV3Algorithm + │ + └─ Based on: FwdAlgorithmV3 + SpecifiesXdl +``` + +### Backward XDL Algorithms + +``` +BwdAlgorithm + SpecifiesXdl +│ +├─── BwdXdlAlgorithmBase (ThreadClusterRank=4) +│ │ +│ ├─── BwdXdlAlgorithm +│ │ │ +│ │ └─ Requirements: Base + SpecifiesTransposeTransfer + SpecifiesGenericInstance +│ │ +│ └─── BwdMultiDXdlAlgorithm +│ │ +│ └─ Requirements: Base + SpecifiesMultipleDSupport +│ +└─── BwdXdlV3AlgorithmBase + │ + ├─── BwdXdlV3Algorithm + │ │ + │ └─ Requirements: Base + │ + └─── BwdTwoStageXdlAlgorithm + │ + └─ Requirements: Base + SpecifiesTransposeTransfer + SpecifiesTwoStageSupport +``` + +**Valid XDL Algorithms:** +- FwdXdlAlgorithm +- FwdXdlV3Algorithm +- LargeTensorAlgorithm +- BwdXdlAlgorithm +- BwdXdlV3Algorithm +- BwdTwoStageXdlAlgorithm +- BwdMultiDXdlAlgorithm + +--- + +## WMMA Algorithm Hierarchy + +### Forward WMMA Algorithms + +``` +FwdAlgorithm + SpecifiesWmma +│ +└─── FwdWmmaAlgorithm + │ + └─ Requirements: Base + SpecifiesWmma +``` + +### Backward WMMA Algorithms + +``` +BwdAlgorithm + SpecifiesWmma +│ +├─── BwdWmmaAlgorithmBase (ThreadClusterRank=3) +│ │ +│ └─── BwdWmmaAlgorithm +│ │ +│ └─ Requirements: Base + SpecifiesNumPrefetchStages + SpecifiesGemmPipeline + SpecifiesGenericInstance +│ +└─── BwdWmmaV3AlgorithmBase (Based on BwdAlgorithmV3) + │ + ├─── BwdMultiDWmmaV3Algorithm + │ │ + │ └─ Requirements: Base + SpecifiesMultipleDSupport + │ + ├─── BwdWmmaV3Algorithm + │ │ + │ └─ Requirements: Base + SpecifiesTransposeTransfer + SpecifiesGenericInstance + │ + └─── BwdTwoStageWmmaV3Algorithm + │ + └─ Requirements: Base + SpecifiesTransposeTransfer + SpecifiesTwoStageSupport +``` + +**Valid WMMA Algorithms:** +- FwdWmmaAlgorithm +- BwdWmmaAlgorithm +- BwdWmmaV3Algorithm +- BwdTwoStageWmmaV3Algorithm +- BwdMultiDWmmaV3Algorithm + +--- + +## DL Algorithm Hierarchy + +DL algorithms have a separate base and do not share the common hierarchy with XDL/WMMA algorithms. + +``` +DlAlgorithm +│ +│ Requirements: +│ • ConvAlgorithmDescriptor +│ • SpecifiesThreadBlock +│ • SpecifiesDlThreadConfig +│ • SpecifiesDlThreadCluster +│ • SpecifiesDlEpilogue +│ +├─── FwdDlAlgorithmBase +│ │ +│ │ Requirements: Base + SpecifiesFwdConvSpecialization + SpecifiesDlFwdBlockTransfer + SpecifiesGemmSpecialization +│ │ +│ └─── FwdDlAlgorithm +│ +└─── BwdDlAlgorithm + │ + └─ Requirements: Base + SpecifiesBwdWeightConvSpecialization + SpecifiesDlBwdBlockTransfer +``` + +**Valid DL Algorithms:** +- FwdDlAlgorithm +- BwdDlAlgorithm + +--- + + +## Reference Algorithms + +``` +ReferenceAlgorithm +│ +└─ Requirements: ConvAlgorithmDescriptor + + SpecifiesReferenceAlgorithm +``` + +Used for reference implementations and testing. + +## CK Tile Algorithms + +``` +TileAlgorithm +│ +└─ Requirements: ConvAlgorithmDescriptor + + SpecifiesTileThreadBlock + + SpecifiesTileTransfer + + SpecifiesTileConvSpecialization + + SpecifiesTileBlockGemm + + SpecifiesTileOptimizations +``` + +The CK Tile algorithms are applicable to foward convolution as well as backwards convolution (weight and data). + +--- + +## Summary for XDL/WMMA/DL algorithms + +| Category | Algorithm Type | Forward Variants | Backward Variants | +|----------|---------------|------------------|-------------------| +| **XDL** | Base | FwdXdlAlgorithmBase | BwdXdlAlgorithmBase, BwdXdlV3AlgorithmBase | +| | Concrete | • FwdXdlAlgorithm
• FwdXdlV3Algorithm
• LargeTensorAlgorithm | • BwdXdlAlgorithm
• BwdXdlV3Algorithm
• BwdTwoStageXdlAlgorithm
• BwdMultiDXdlAlgorithm | +| **WMMA** | Base | FwdAlgorithm | BwdWmmaAlgorithmBase, BwdWmmaV3AlgorithmBase | +| | Concrete | • FwdWmmaAlgorithm | • BwdWmmaAlgorithm
• BwdWmmaV3Algorithm
• BwdTwoStageWmmaV3Algorithm
• BwdMultiDWmmaV3Algorithm | +| **DL** | Base | FwdDlAlgorithmBase | DlAlgorithm | +| | Concrete | • FwdDlAlgorithm | • BwdDlAlgorithm | + +--- diff --git a/experimental/builder/include/ck_tile/builder/factory/conv_algorithms.hpp b/experimental/builder/include/ck_tile/builder/factory/conv_algorithms.hpp index 3b2f7f8c51..5c169d96a7 100644 --- a/experimental/builder/include/ck_tile/builder/factory/conv_algorithms.hpp +++ b/experimental/builder/include/ck_tile/builder/factory/conv_algorithms.hpp @@ -7,52 +7,53 @@ namespace ck_tile::builder::factory { -// Base algorithm concepts template -concept TileTransferParameters = +concept SpecifiesTileTransferParameters = SpecifiesThreadClusters && SpecifiesLdsTransfer && SpecifiesThreadClusterAccessOrder && SpecifiesSourceAccessOrder; -template -concept SpecifiesTileTransferParameters3D = TileTransferParameters; +// Base algorithm concepts +template +concept ConvAlgorithm = + ConvAlgorithmDescriptor && SpecifiesThreadBlock && + SpecifiesTileTransferParameters && + SpecifiesWarpGemm && SpecifiesGemmPipeline;; template -concept SpecifiesTileTransferParameters4D = TileTransferParameters; +concept FwdAlgorithm = ConvAlgorithm && SpecifiesFwdConvSpecialization; + +template +concept FwdAlgorithmV3 = FwdAlgorithm && SpecifiesPipelineV3; + +template +concept BwdAlgorithm = ConvAlgorithm && SpecifiesBwdWeightConvSpecialization; + +template +concept BwdAlgorithmV3 = BwdAlgorithm && SpecifiesPipelineV3; + +template +concept DlAlgorithm = + ConvAlgorithmDescriptor && SpecifiesThreadBlock && + SpecifiesDlThreadConfig && SpecifiesDlThreadCluster && SpecifiesDlEpilogue; + +template +concept FwdDlAlgorithmBase = DlAlgorithm && SpecifiesFwdConvSpecialization && + SpecifiesDlFwdBlockTransfer && SpecifiesGemmSpecialization; template -concept FwdXdlAlgorithmBase = - ConvAlgorithmDescriptor && SpecifiesThreadBlock && - SpecifiesTileTransferParameters3D && - SpecifiesWarpGemm && SpecifiesFwdConvSpecialization && - SpecifiesGemmPipeline && SpecifiesXdl; +concept FwdXdlAlgorithmBase = FwdAlgorithm && SpecifiesXdl; template -concept BwdXdlAlgorithmBase = - ConvAlgorithmDescriptor && SpecifiesThreadBlock && - SpecifiesTileTransferParameters4D && - SpecifiesWarpGemm && SpecifiesBwdWeightConvSpecialization && - SpecifiesXdl; +concept BwdXdlAlgorithmBase = BwdAlgorithm && SpecifiesXdl; template -concept BwdXdlV3AlgorithmBase = - ConvAlgorithmDescriptor && SpecifiesThreadBlock && - SpecifiesTileTransferParameters3D && - SpecifiesWarpGemm && SpecifiesBwdWeightConvSpecialization && - SpecifiesGemmPipeline && SpecifiesXdl && SpecifiesPipelineV3; +concept BwdXdlV3AlgorithmBase = BwdAlgorithmV3 && SpecifiesXdl; template -concept BwdWmmaAlgorithmBase = - ConvAlgorithmDescriptor && SpecifiesThreadBlock && - SpecifiesTileTransferParameters3D && - SpecifiesWarpGemm && SpecifiesBwdWeightConvSpecialization && - SpecifiesWmma; +concept BwdWmmaAlgorithmBase = BwdAlgorithm && SpecifiesWmma; template -concept BwdWmmaV3AlgorithmBase = - ConvAlgorithmDescriptor && SpecifiesThreadBlock && - SpecifiesTileTransferParameters3D && - SpecifiesWarpGemm && SpecifiesBwdWeightConvSpecialization && - SpecifiesGemmPipeline && SpecifiesWmma && SpecifiesPipelineV3; +concept BwdWmmaV3AlgorithmBase = BwdAlgorithmV3 && SpecifiesWmma; // Reference algorithm concept template @@ -72,28 +73,16 @@ template concept LargeTensorAlgorithm = FwdXdlAlgorithmBase && SpecifiesLargeTensorSupport; template -concept FwdXdlV3Algorithm = - ConvAlgorithmDescriptor && SpecifiesThreadBlock && - SpecifiesTileTransferParameters3D && - SpecifiesWarpGemm && SpecifiesFwdConvSpecialization && - SpecifiesGemmPipeline && SpecifiesXdl && SpecifiesPipelineV3; +concept FwdXdlV3Algorithm = FwdAlgorithmV3 && SpecifiesXdl; // FWD WMMA algorithm concepts template -concept FwdWmmaAlgorithm = - ConvAlgorithmDescriptor && SpecifiesThreadBlock && - SpecifiesTileTransferParameters3D && - SpecifiesWarpGemm && SpecifiesFwdConvSpecialization && - SpecifiesGemmPipeline && SpecifiesWmma; +concept FwdWmmaAlgorithm = FwdAlgorithm && SpecifiesWmma; // FWD DL algorithms template -concept FwdDlAlgorithm = - ConvAlgorithmDescriptor && SpecifiesThreadBlock && - SpecifiesFwdConvSpecialization && SpecifiesGemmSpecialization && - SpecifiesDlThreadConfig && SpecifiesDlThreadCluster && - SpecifiesDlFwdBlockTransfer && SpecifiesDlEpilogue; - +concept FwdDlAlgorithm = FwdDlAlgorithmBase; + // BWD weight XDL algorithm concepts template concept BwdXdlAlgorithm = @@ -130,11 +119,8 @@ concept BwdTwoStageWmmaV3Algorithm = BwdWmmaV3AlgorithmBase && SpecifiesT // BWD weight DL algorithms template -concept BwdDlAlgorithm = - ConvAlgorithmDescriptor && SpecifiesThreadBlock && - SpecifiesBwdWeightConvSpecialization && SpecifiesDlThreadConfig && - SpecifiesDlThreadCluster && SpecifiesDlBwdBlockTransfer && - SpecifiesDlEpilogue; +concept BwdDlAlgorithm = DlAlgorithm && SpecifiesBwdWeightConvSpecialization && + SpecifiesDlBwdBlockTransfer; // Concepts for valid XDL/WMMA algorithms template @@ -154,7 +140,6 @@ concept SpecifiesValidBwdWmmaAlgorithm = BwdWmmaAlgorithm || BwdWmmaV3Algorithm || BwdTwoStageWmmaV3Algorithm || BwdMultiDWmmaV3Algorithm; - template concept FwdWarpGemmOrDL = SpecifiesValidWarpGemm || FwdDlAlgorithm; diff --git a/experimental/builder/test/impl/conv_algorithm_types.hpp b/experimental/builder/test/impl/conv_algorithm_types.hpp index 8d315f69b8..264f69ad81 100644 --- a/experimental/builder/test/impl/conv_algorithm_types.hpp +++ b/experimental/builder/test/impl/conv_algorithm_types.hpp @@ -42,7 +42,7 @@ struct GemmPipeline { size_t num_gemm_k_prefetch_stages{1}; size_t num_conv_groups_to_merge{1}; - PipelineVersion pipeline_version; + PipelineVersion pipeline_version{PipelineVersion::V1}; PipelineScheduler scheduler{PipelineScheduler::DEFAULT}; }; static_assert(ckb::GemmPipelineDescriptor); @@ -543,6 +543,7 @@ using ConvAlgorithm_DeviceGroupedConvBwdWeight_Xdl_CShuffle = WarpGemm_, InputOutputTileTransfer_<4>, ConvSpecializationBwdWeight_, + GemmPipeline_, // Not needed, but we need this to meet the ConvAlgorithm concept. TransposeParams_, AlgorithmSpecialization_<>>; @@ -592,6 +593,7 @@ using ConvAlgorithm_DeviceGroupedConvBwdWeightMultipleD_Xdl_CShuffle = ConvAlgorithmTemplate, + GemmPipeline_, // Not needed, but we need this to meet the ConvAlgorithm concept. ConvSpecializationBwdWeight_, AlgorithmSpecialization_>;