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

Initial commit of convolution builder.

Browse Source 
Creates a single instance with template metaprogramming. Many things are still hard-coded.
This commit is contained in:
John Shumway
2025-08-28 01:48:42 +00:00
parent 1c2078066b
commit 4117fcb36e
11 changed files with 574 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
CMakeLists.txt
View File

@@ -667,6 +667,10 @@ if (NOT MIOPEN_REQ_LIBS_ONLY)
add_subdirectory(profiler)
endif()
if (MIOPEN_REQ_LIBS_ONLY)
add_subdirectory(experimental/builder/test)
endif()
if(CK_USE_CODEGEN AND (SUPPORTED_GPU_TARGETS MATCHES "gfx9" OR GPU_ARCHS))
add_subdirectory(codegen)
endif()

95
experimental/builder/include/ck_tile/builder/conv_algorithm.hpp Normal file
View File

@@ -0,0 +1,95 @@
#pragma once
#include <type_traits>
namespace ck_tile::builder {
enum class GemmImplementationType
{
XDL,
WMMA,
DL
};
enum class ConvolutionDirection
{
Forward,
BackwardData,
BackwardWeight
};
enum class UniversalGemmSupport
{
Supported,
NotSupported
};
enum class SplitKSupport
{
Supported,
SupportedTwoStage,
NotSupported
};
enum class DepthwiseOptimization
{
X16,
X8,
X4,
X2,
NotSupported
};
enum class LargeTensorSupport
{
Supported,
SplitBatch,
NotSupported
};
enum class ImplementationType
{
ExplicitDefault,
ExplicitMPadding,
ExplicitNPadding,
ExplicitKPadding,
ExplicitMNPadding,
ExplicitMKPadding,
ExplicitNKPadding,
ExplicitMNKPadding,
Implicit
};
enum class GemmPipelineVersion
{
Naive,
ComputeFriendly,
MemFriendly,
ComputeFriendlyDoubleLDS,
ComputeFriendlyDoubleGlobalPrefetch
};
enum class GemmPipelineScheduler
{
Intrawave,
Interwave
};
enum class ConvolutionSpecialization
{
Default,
Filter1x1Pad0,
Filter1x1Stride1Pad0,
Filter3x3
};
enum class MFMAInstructionSize
{
M16N16,
M32N32
};
template <typename T>
concept ConvAlgorithm = std::is_class_v<T>;
} // namespace ck_tile::builder

30
experimental/builder/include/ck_tile/builder/conv_builder.hpp Normal file
View File

@@ -0,0 +1,30 @@
#pragma once
#include <concepts>
#include <type_traits>
#include "conv_builder_reference.hpp"
#include <ck_tile/builder/conv_algorithm.hpp>
#include <ck_tile/builder/conv_factory.hpp>
#include <ck_tile/builder/conv_signature.hpp>
#include <ck_tile/builder/versions.h>
namespace ck_tile::builder {
template <ConvSignature TSignature, ConvAlgorithm TAlgorithm, auto Version>
requires SupportedVersion<Version>
struct ConvBuilder
{
// Input: Signature describes the mathematical funcationality of the algorithm.
using Signature = TSignature;
// Input: Algorithm describes the implementation of the algorithm.
using Algorithm = TAlgorithm;
// Input: Version of the builder, exposed for testing.
static constexpr auto kVersion = Version;
// Implmentation: The factory handles the builder logic.
using builder = GroupedConvForwardXldCShuffleFactoryV3<Signature, Algorithm, Version>;
// Output: The kernel class.
using Instance = builder::Instance;
};
} // namespace ck_tile::builder

217
experimental/builder/include/ck_tile/builder/conv_factory.hpp Normal file
View File

@@ -0,0 +1,217 @@
#pragma once
// #include "ck/library/tensor_operation_instance/gpu/grouped_conv_fwd/device_grouped_conv_fwd_xdl_comp_instance.hpp"
#include <ck/tensor_operation/gpu/device/impl/device_grouped_conv_fwd_multiple_abd_xdl_cshuffle_v3.hpp>
#include <ck_tile/builder/conv_signature.hpp>
#include <ck_tile/builder/conv_algorithm.hpp>
#include <ck_tile/builder/sequence_util.hpp>
#include <ck_tile/builder/versions.h>
namespace ck_tile::builder {
// Type mappings from the builder GroupConvLayout enum class to the CK tensor data types.
template <GroupConvLayout Layout>
struct ConvTensorLayouts;
template <>
struct ConvTensorLayouts<GroupConvLayout::NGCHW_GKCYX_NGKHW>
{
// Channels first convolution layout.
using ALayout = ck::tensor_layout::convolution::NHWGC;
using BLayout = ck::tensor_layout::convolution::GKCYX;
using DsLayout = ck::Tuple<>;
using ELayout = ck::tensor_layout::convolution::NGKHW;
};
template <>
struct ConvTensorLayouts<GroupConvLayout::NHWGC_GKYXC_NHWGK>
{
// Channels last convolution layout.
using ALayout = ck::tensor_layout::convolution::NHWGC;
using BLayout = ck::tensor_layout::convolution::GKYXC;
using DsLayout = ck::Tuple<>;
using ELayout = ck::tensor_layout::convolution::NHWGK;
};
// Type mappings from builder convolution data type to CK tensor types.
template <DataType T>
struct ConvTensorTypes;
template <>
struct ConvTensorTypes<DataType::FP16>
{
using ADataType = ck::bhalf_t;
using BDataType = ck::bhalf_t;
using CShuffleDataType = ck::bhalf_t;
using DsDataTypes = ck::Tuple<>;
using AccDataType = float;
using EDataTYpe = ck::bhalf_t;
};
// Hard-coded pass-through ops.
struct ConvPassThroughOps
{
using AElementwiseOp = ck::tensor_operation::element_wise::PassThrough;
using BElementwiseOp = ck::tensor_operation::element_wise::PassThrough;
using CDEElementwiseOp = ck::tensor_operation::element_wise::PassThrough;
};
// The specializations for the convolution and GEMM.
struct ConvSpec
{
ck::tensor_operation::device::ConvolutionForwardSpecialization conv_spec =
ck::tensor_operation::device::ConvolutionForwardSpecialization::Default;
ck::tensor_operation::device::GemmSpecialization gemm_spec =
ck::tensor_operation::device::GemmSpecialization::MNKPadding;
};
// Store M,N,K values.
template <typename T>
struct MNK
{
T m = 0;
T n = 0;
T k = 0;
};
// Block info for a convlution.
struct ConvBlock
{
int block_size = 0;
MNK<int> per_block;
};
// Convolution tuning parameters.
struct ConvTuning
{
int ak1 = 0;
int ak2 = 0;
int m_per_xdl = 0;
int n_per_dxl = 0;
int m_xdl_per_wave = 0;
int n_xdl_per_wave = 0;
};
// Block tranfser paramters for A or B tensor.
struct BlockTransfer
{
ck::Array<int, 3> thread_cluster_lengths = {0, 0, 0}; // k0, m, k1
ck::Array<int, 3> thread_cluster_order = {0, 0, 0};
ck::Array<int, 3> src_access_order = {0, 0, 0};
int src_vector_dim = 0;
int src_scaler_per_vector = 0;
int dest_scaler_per_vector_k1 = 0;
int add_extra = 0;
};
// Block transfer parameters for C tensor.
struct CBlockTransfer
{
int m_xdl_per_wave_per_shuffle = 0;
int n_xdl_per_wave_per_shuffle = 0;
ck::Array<int, 4> cluster_lengths = {0, 0, 0, 0};
int scaler_per_vector = 8;
};
// Factory builds an instance of a grouped convolution kernel.
template <ConvSignature Signature, ConvAlgorithm Algorithm, auto Version>
requires SupportedVersion<Version>
struct GroupedConvForwardXldCShuffleFactoryV3
{
static constexpr int SPATIAL_DIM = Signature::SPATIAL_DIM;
using Layouts = ConvTensorLayouts<Signature::LAYOUT>;
using Types = ConvTensorTypes<Signature::DATA_TYPE>;
using Ops = ConvPassThroughOps;
static constexpr ConvSpec SPECIALIZATION{
.conv_spec = ck::tensor_operation::device::ConvolutionForwardSpecialization::Default,
.gemm_spec = ck::tensor_operation::device::GemmSpecialization::MNKPadding,
};
static constexpr ConvBlock BLOCK{
.block_size = 256,
.per_block = {.m = 256, .n = 256, .k = 32},
};
static constexpr ConvTuning TUNING{
.ak1 = 8,
.ak2 = 8,
.m_per_xdl = 32,
.n_per_dxl = 32,
.m_xdl_per_wave = 4,
.n_xdl_per_wave = 4,
};
static constexpr BlockTransfer A_BLOCK_TRANSFER{
.thread_cluster_lengths = {4, 64, 1},
.thread_cluster_order = {1, 0, 2},
.src_access_order = {1, 0, 2},
.src_vector_dim = 2,
.src_scaler_per_vector = 8,
.dest_scaler_per_vector_k1 = 8,
.add_extra = 0,
};
static constexpr BlockTransfer B_BLOCK_TRANSFER{
.thread_cluster_lengths = {4, 64, 1},
.thread_cluster_order = {1, 0, 2},
.src_access_order = {1, 0, 2},
.src_vector_dim = 2,
.src_scaler_per_vector = 8,
.dest_scaler_per_vector_k1 = 8,
.add_extra = 0,
};
static constexpr CBlockTransfer C_BLOCK_TRANSFER{
.m_xdl_per_wave_per_shuffle = 1,
.n_xdl_per_wave_per_shuffle = 1,
.cluster_lengths = {1, 32, 1, 8},
.scaler_per_vector = 8,
};
static constexpr auto PIPELINE_SCHEDULER = ck::BlockGemmPipelineScheduler::Intrawave;
static constexpr auto PIPELINE_VERSION = ck::BlockGemmPipelineVersion::v4;
// The convlution kernel class instance.
using Instance = ck::tensor_operation::device::DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3< //
SPATIAL_DIM,
typename Layouts::ALayout,
typename Layouts::BLayout,
typename Layouts::DsLayout,
typename Layouts::ELayout,
typename Types::ADataType,
typename Types::BDataType,
typename Types::AccDataType,
typename Types::CShuffleDataType,
typename Types::DsDataTypes,
typename Types::EDataTYpe,
typename Ops::AElementwiseOp,
typename Ops::BElementwiseOp,
typename Ops::CDEElementwiseOp,
SPECIALIZATION.conv_spec,
SPECIALIZATION.gemm_spec,
BLOCK.block_size,
BLOCK.per_block.m,
BLOCK.per_block.n,
BLOCK.per_block.k,
TUNING.ak1,
TUNING.ak2,
TUNING.m_per_xdl,
TUNING.n_per_dxl,
TUNING.m_xdl_per_wave,
TUNING.n_xdl_per_wave,
ToSequence<A_BLOCK_TRANSFER.thread_cluster_lengths>,
ToSequence<A_BLOCK_TRANSFER.thread_cluster_order>,
ToSequence<A_BLOCK_TRANSFER.src_access_order>,
A_BLOCK_TRANSFER.src_vector_dim,
A_BLOCK_TRANSFER.src_scaler_per_vector,
A_BLOCK_TRANSFER.dest_scaler_per_vector_k1,
A_BLOCK_TRANSFER.add_extra,
ToSequence<B_BLOCK_TRANSFER.thread_cluster_lengths>,
ToSequence<B_BLOCK_TRANSFER.thread_cluster_order>,
ToSequence<B_BLOCK_TRANSFER.src_access_order>,
B_BLOCK_TRANSFER.src_vector_dim,
B_BLOCK_TRANSFER.src_scaler_per_vector,
B_BLOCK_TRANSFER.dest_scaler_per_vector_k1,
B_BLOCK_TRANSFER.add_extra,
C_BLOCK_TRANSFER.m_xdl_per_wave_per_shuffle,
C_BLOCK_TRANSFER.n_xdl_per_wave_per_shuffle,
ToSequence<C_BLOCK_TRANSFER.cluster_lengths>,
C_BLOCK_TRANSFER.scaler_per_vector,
PIPELINE_SCHEDULER,
PIPELINE_VERSION>;
};
} // namespace ck_tile::builder

62
experimental/builder/include/ck_tile/builder/conv_signature.hpp Normal file
View File

@@ -0,0 +1,62 @@
#pragma once
#include <concepts>
#include <type_traits>
#include <ck_tile/builder/types.hpp>
namespace ck_tile::builder {
// Layouts for grouped convolutions.
enum class GroupConvLayout
{
NHWGC_GKYXC_NHWGK, // Channels-last
NGCHW_GKCYX_NGKHW // Channels-first
};
// Spatial dimensionalities of grouped convolutions.
// N represents the number of spatial dimensions (e.g., 1 for 1D, 2 for 2D, 3 for 3D).
template <auto N>
concept ConvSpatialDim = std::is_integral_v<decltype(N)> && (N == 1 || N == 2 || N == 3);
// Allowed datatypes for grouped convolutions.
// Currently limited to floating-point types commonly accelerated on GPUs.
template <DataType T>
concept ConvDataType = (T == DataType::FP32) || (T == DataType::FP16);
// Direction of the convolution operation.
enum class ConvDirection
{
Forward,
BackwardData,
BackwardWeight
};
// Elementwise operation to fuse to convolution.
enum class ElementwiseOperation
{
Bias,
BiasClamp,
Bilinear,
Clamp,
Scale,
PassThrough
};
// Operational signature of a convolution.
template <typename T>
concept ConvSignature = requires {
// Dimensionality of the convolution (e.g., 1, 2, or 3).
requires ConvSpatialDim<T::SPATIAL_DIM>;
// Direction of the convolition (fwd, bwd, or weights).
{ T::DIRECTION } -> std::same_as<const ConvDirection&>;
// Memory layout of the tensors.
{ T::LAYOUT } -> std::same_as<const GroupConvLayout&>;
// Tensor datatype for input and output.
requires ConvDataType<T::DATA_TYPE>;
};
} // namespace ck_tile::builder

23
experimental/builder/include/ck_tile/builder/sequence_util.hpp Normal file
View File

@@ -0,0 +1,23 @@
#pragma once
#include <ck/utility/sequence.hpp>
namespace ck_tile::builder {
// Helper struct to get the Sequence type from a constexpr ck::Array.
template <typename T, const T& Arr, typename>
struct ToSequenceHelper;
template <typename T, const T& Arr, std::size_t... Is>
struct ToSequenceHelper<T, Arr, std::index_sequence<Is...>>
{
using type = ck::Sequence<Arr[Is]...>;
};
// The main interface to get the type
template <auto& Arr>
using ToSequence = typename ToSequenceHelper<
std::remove_cvref_t<decltype(Arr)>,
Arr,
std::make_index_sequence<std::remove_reference_t<decltype(Arr)>::Size()>>::type;
} // namespace ck_tile::builder

14
experimental/builder/include/ck_tile/builder/types.hpp Normal file
View File

@@ -0,0 +1,14 @@
#pragma once
namespace ck_tile::builder {
enum class DataType
{
FP64,
FP32,
FP16,
BF16,
S16,
S8,
S4,
};
} // namespace ck_tile::builder

15
experimental/builder/include/ck_tile/builder/versions.h Normal file
View File

@@ -0,0 +1,15 @@
#pragma once
#include <concepts>
#include <string_view>
namespace ck_tile::builder {
static constexpr char V0_0_0[] = "0.0.0";
static constexpr char V0_1_0[] = "0.1.0";
template <const char* V>
concept SupportedVersion = (std::string_view{V} == std::string_view{V0_0_0}) ||
(std::string_view{V} == std::string_view{V0_1_0});
} // namespace ck_tile::builder

17
experimental/builder/test/CMakeLists.txt Normal file
View File

@@ -0,0 +1,17 @@
set(CMAKE_CXX_STANDARD 20)
add_executable(gemm_example gemm_example.cpp)
target_include_directories(gemm_example PRIVATE
"${PROJECT_SOURCE_DIR}/include"
)
include(gtest)
add_executable(test_conv_builder builder.cpp)
target_include_directories(test_conv_builder PRIVATE
"${PROJECT_SOURCE_DIR}/experimental/builder/include"
"${PROJECT_SOURCE_DIR}/include"
)
target_compile_options(test_conv_builder PRIVATE -Wno-global-constructors -Wno-c++20-compat)
target_link_libraries(test_conv_builder PRIVATE GTest::gtest GTest::gtest_main)

88
experimental/builder/test/builder.cpp Normal file
View File

@@ -0,0 +1,88 @@
#include <gtest/gtest.h>
#include <ck_tile/builder/conv_builder.hpp>
namespace {
namespace ckb = ck_tile::builder;
// Example of kernel description for Forward Conv with default settings
struct GroupedConvFwdXdlImplicitGemm : public GroupedConvBaseXdlV1
{
static constexpr ConvolutionDirection ConvolutionDirection_ = ConvolutionDirection::Forward;
static constexpr ElementwiseOperation ElementwiseOperation_ = ElementwiseOperation::Bias;
};
// Example of kernel description for Backward Weight Conv with default settings and Split K Two
// Stage
struct GroupedConvBwdWeightXdlImplicitGemmTwoStage : public GroupedConvBaseXdlV1
{
[[maybe_unused]] static constexpr ConvolutionDirection ConvolutionDirection_ =
ConvolutionDirection::BackwardWeight;
[[maybe_unused]] static constexpr SplitKSupport SplitKSupport_ =
SplitKSupport::SupportedTwoStage;
};
struct Implementation16x16 : ImplementationDefaultV1
{
static constexpr ck::index_t BlockSize_ = 64;
static constexpr auto TileSizes_ = std::make_tuple(16, 16, 32);
static constexpr ck::index_t K1_ = 8;
static constexpr MFMAInstructionSize MFMAInstructionSize_ = MFMAInstructionSize::M16N16;
static constexpr auto XdlPerWave_ = std::make_tuple(16, 16);
static constexpr auto GlobalTransferVectorSize_ = std::make_tuple(1, 1, 1);
static constexpr auto LDSStoreVectorSize_ = std::make_tuple(4, 4);
};
struct ProblemBF16NHWGC : public BF16ProblemBaseV1, public NHWGCProblemBaseV1
{
};
TEST(ConvBuilderTest, TestBuilderV0_0_0)
{
ConvolutionBuilder<GroupedConvFwdXdlImplicitGemm, ProblemBF16NHWGC, Implementation16x16>
builder_fwd;
EXPECT_EQ(builder_fwd.GetInstanceName(),
"GroupedConvFwdMultipleABD_Xdl_CShuffle<64, 16, 16, 32, Default, 8, 16x16, 16, 16, "
"1, 4, 1, 4, 1, Intrawave, v1, 1>");
// It would be nice if this worked, but it fails.
// [[maybe_unused]] auto instance = builder_fwd.GetInstance();
}
struct FwdConvSignature
{
static constexpr int SPATIAL_DIM = 2;
static constexpr auto DIRECTION = ckb::ConvDirection::Forward;
static constexpr auto LAYOUT = ckb::GroupConvLayout::NHWGC_GKYXC_NHWGK;
static constexpr auto DATA_TYPE = ckb::DataType::FP16;
};
TEST(ConvBuilderTest, TestSignature)
{
static_assert(ckb::ConvSignature<FwdConvSignature>);
SUCCEED();
}
struct FwdConvAlgorithm
{
//
};
TEST(ConvBuilderTest, TestAlgorithm)
{
static_assert(ckb::ConvAlgorithm<FwdConvAlgorithm>);
SUCCEED();
}
static constexpr char API_VERSION[] = "0.1.0";
using FwdConvBuilder = ckb::ConvBuilder<FwdConvSignature, FwdConvAlgorithm, API_VERSION>;
TEST(ConvBuilderTest, TestKernel)
{
EXPECT_EQ(
FwdConvBuilder::Instance::TypeString(),
"DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3<256, 256, 256, 32, Default, 32, 32, 4, 4, "
"8, 8, 8, 1, 1, BlkGemmPipelineScheduler: Intrawave, BlkGemmPipelineVersion: v4>");
}
} // namespace

11
include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_fwd_multiple_abd_xdl_cshuffle_v3.hpp
View File

@@ -1889,8 +1889,9 @@ struct DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3
return std::make_unique<Invoker>(Invoker{});
}
std::string GetTypeString() const override
{
static std::string TypeString() {
// Make this a static function on the class so we don't have to instantiate the
// object to get the classes type string.
auto str = std::stringstream();
std::map<BlockGemmPipelineScheduler, std::string> BlkGemmPipelineSchedulerToString{
@@ -1931,6 +1932,12 @@ struct DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3
return str.str();
}
std::string GetTypeString() const override
{
// Make the static type string available through the base class.
return TypeString();
}
size_t GetWorkSpaceSize(const BaseArgument* p_arg) const override
{
auto arg = dynamic_cast<const Argument*>(p_arg);