diff --git a/example/ck_builder/01_conv_2d_fwd/ckb_example_conv_fwd_2d.cpp b/example/ck_builder/01_conv_2d_fwd/ckb_example_conv_fwd_2d.cpp index c2dcdb16a6..618ceabfd0 100644 --- a/example/ck_builder/01_conv_2d_fwd/ckb_example_conv_fwd_2d.cpp +++ b/example/ck_builder/01_conv_2d_fwd/ckb_example_conv_fwd_2d.cpp @@ -1,59 +1,26 @@ -#include + #include "ck_tile/builder/conv_builder.hpp" -#include "ck_tile/builder/conv_signature.hpp" - -namespace -{ - namespace ckb = ck_tile::builder; - - struct ConvSignature { - int spatial_dim; - ckb::ConvDirection direction; - ckb::GroupConvLayout layout; - ckb::DataType data_type; - }; - static_assert(ckb::ConvSignatureDescriptor); - - struct ConvAlgorithm - { - ckb::ThreadBlock thread_block; - ckb::ConvTuningParams tuning_params; - struct BlockTransfer - { - ckb::BlockATransferLengths thread_cluster_dims_a; - ckb::BlockBTransferLengths thread_cluster_dims_b; - ckb::BlockCTransferLengths thread_cluster_dims_c; - } block_transfer; - ckb::BlockGemmPipelineVersion pipeline_version; - }; - static_assert(ckb::ConvAlgorithmDescriptor); - static_assert(ckb::SpecifiesThreadBlock); - static_assert(ckb::SpecifiesConvTuning); - static_assert(ckb::SpecifiesBlockATransfer); - static_assert(ckb::SpecifiesBlockBTransfer); - static_assert(ckb::SpecifiesBlockCTransfer); - static_assert(ckb::SpecifiesGemmPipelineVersion); - -} // namespace +#include "../utils/types.hpp" int main() { namespace ckb = ck_tile::builder; + namespace ckb_examples = ck_tile::builder::examples; - constexpr ConvSignature FwdConvSignature + constexpr size_t m_tile = 32; + constexpr size_t n_tile = 16; + constexpr size_t k_tile = 64; + constexpr size_t k0 = 4; + constexpr size_t k1 = 1; + + constexpr ckb_examples::ConvSignature FwdConvSignature { .spatial_dim = 2, .direction = ckb::ConvDirection::FORWARD, .layout = ckb::GroupConvLayout::CHANNELS_LAST, .data_type = ckb::DataType::BF16, }; - static_assert(ckb::ValidConvSignature); - - constexpr size_t m_tile = 32; - constexpr size_t n_tile = 16; - constexpr size_t k_tile = 64; - constexpr size_t k0 = 4; - constexpr size_t k1 = 1; + static_assert(ckb::ValidConvSignature); constexpr ckb::ThreadBlock FwdThreadBlock { @@ -61,7 +28,7 @@ int main() { .submatrix = {.m = m_tile, .n = n_tile, .k = k_tile} }; - constexpr ConvAlgorithm::BlockTransfer FwdBlockTransfer + constexpr ckb_examples::BlockTransfer FwdBlockTransfer { .thread_cluster_dims_a = {.k0 = k0, .m = m_tile, .k1 = k1}, .thread_cluster_dims_b = {.k0 = k0, .n = n_tile, .k1 = k1}, @@ -69,7 +36,7 @@ int main() { .m_block = 1, .m_wave_per_xdl = 32, .n_block = 1, .n_wave_per_xdl = 8} }; - constexpr ConvAlgorithm FwdConvAlgorithm + constexpr ckb_examples::ConvAlgorithm FwdConvAlgorithm { .thread_block = FwdThreadBlock, .tuning_params = {.ak1 = 16, .bk1 = 16, .m_xdl_per_wave = 2, .n_xdl_per_wave = 2}, diff --git a/example/ck_builder/utils/types.hpp b/example/ck_builder/utils/types.hpp new file mode 100644 index 0000000000..bde0ec69fa --- /dev/null +++ b/example/ck_builder/utils/types.hpp @@ -0,0 +1,39 @@ +#pragma once + +#include "ck_tile/builder/conv_algorithm.hpp" +#include "ck_tile/builder/conv_signature.hpp" + +namespace ck_tile::builder::examples +{ + namespace ckb = ck_tile::builder; + + struct ConvSignature { + int spatial_dim; + ckb::ConvDirection direction; + ckb::GroupConvLayout layout; + ckb::DataType data_type; + }; + static_assert(ckb::ConvSignatureDescriptor); + + struct BlockTransfer + { + ckb::BlockATransferLengths thread_cluster_dims_a; + ckb::BlockBTransferLengths thread_cluster_dims_b; + ckb::BlockCTransferLengths thread_cluster_dims_c; + }; + + struct ConvAlgorithm + { + ckb::ThreadBlock thread_block; + ckb::ConvTuningParams tuning_params; + BlockTransfer block_transfer; + ckb::BlockGemmPipelineVersion pipeline_version; + }; + static_assert(ckb::ConvAlgorithmDescriptor); + static_assert(ckb::SpecifiesThreadBlock); + static_assert(ckb::SpecifiesConvTuning); + static_assert(ckb::SpecifiesBlockATransfer); + static_assert(ckb::SpecifiesBlockBTransfer); + static_assert(ckb::SpecifiesBlockCTransfer); + static_assert(ckb::SpecifiesGemmPipelineVersion); +} // namespace ck_tile::builder::examples