ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-14 19:18:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Create re-usable assets.

Browse Source
This commit is contained in:
Ville Pietilä
2025-10-03 10:43:28 +00:00
parent 0570ea48c0
commit e94d76de53
2 changed files with 52 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

59
example/ck_builder/01_conv_2d_fwd/ckb_example_conv_fwd_2d.cpp
View File

@@ -1,59 +1,26 @@
#include <iostream>
#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<ConvSignature>);
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<ConvAlgorithm>);
static_assert(ckb::SpecifiesThreadBlock<ConvAlgorithm>);
static_assert(ckb::SpecifiesConvTuning<ConvAlgorithm>);
static_assert(ckb::SpecifiesBlockATransfer<ConvAlgorithm>);
static_assert(ckb::SpecifiesBlockBTransfer<ConvAlgorithm>);
static_assert(ckb::SpecifiesBlockCTransfer<ConvAlgorithm>);
static_assert(ckb::SpecifiesGemmPipelineVersion<ConvAlgorithm>);
} // 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<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;
static_assert(ckb::ValidConvSignature<FwdConvSignature>);
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},

39
example/ck_builder/utils/types.hpp Normal file
View File

@@ -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<ConvSignature>);
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<ConvAlgorithm>);
static_assert(ckb::SpecifiesThreadBlock<ConvAlgorithm>);
static_assert(ckb::SpecifiesConvTuning<ConvAlgorithm>);
static_assert(ckb::SpecifiesBlockATransfer<ConvAlgorithm>);
static_assert(ckb::SpecifiesBlockBTransfer<ConvAlgorithm>);
static_assert(ckb::SpecifiesBlockCTransfer<ConvAlgorithm>);
static_assert(ckb::SpecifiesGemmPipelineVersion<ConvAlgorithm>);
} // namespace ck_tile::builder::examples