mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-04-19 22:39:03 +00:00
[CK_BUILDER] Refactor convolution signature to provide data type/layout/elementwise op per tensor (#3331)
* Separate layouts into separate entities for input, weight, and output tensors. * Add test for handling bias tensor layouts. * Use instance string in builder tests. * Add handling of output bias data types and layouts. * Generalize handling of the elementwise ops. * Test fix. * Create builder for layouts. * Layout builder improvements. * Improve layout builder. * Simplify bias layout handling. * Code clean-up. * Move layout utils into separate file. * Remove hard-coded layout combinations. * Small code clean-up. * Move data type utils into a separate file. * Add data types, layouts, and elementwise ops per conv tensor. * Builder bug fixes after refactoring. * Working baseline. * Make signature definition look nice in the test code. * Move TensorConfig into test implementations. * Fix all fwd conv builder tests. * Fix conv traits and descriptors tests. * More factory assets under a separate directory. * Fix building conv traits. * Fix clang-format. * Add Readme doc to describe the design. * Add link to main Readme. Fix links in the builder design doc. * Clean-up data type/layout/elementwise op conversions. * Switch from dimension and tensor type specific layouts to a flat list of tensor layouts. * Fix clang-formatting. * Fix clang-format for test code. * Simplify fwd conv signature definitions in the test code. * Remove accidental edits. * Fix comment string. * Fix instance factory after rebase. * Fix tests after rebase. * Unify layout handling. * Add more conv layout unit tests. * Clang-format. * Fix merge conflicts. * Improve elementwise op handling. --------- Co-authored-by: Ville Pietilä <>
This commit is contained in:
Ville Pietilä
@@ -4,116 +4,481 @@
|
||||
#include <gtest/gtest.h>
|
||||
#include <type_traits>
|
||||
|
||||
// Include the helper file we're testing
|
||||
#include "ck_tile/builder/factory/helpers/conv_tensor_layout.hpp"
|
||||
#include "impl/conv_signature_types.hpp"
|
||||
|
||||
namespace {
|
||||
|
||||
namespace ckb = ::ck_tile::builder;
|
||||
using ::ck_tile::builder::DataType;
|
||||
using ::ck_tile::builder::ElementwiseOperation;
|
||||
using ::ck_tile::builder::TensorLayout;
|
||||
using ::ck_tile::builder::factory::internal::AuxiliaryTensorLayouts;
|
||||
using ::ck_tile::builder::factory::internal::ConvTensorLayouts;
|
||||
using ::ck_tile::builder::factory::internal::GetTensorLayout;
|
||||
using ::ck_tile::builder::factory::internal::LayoutToCK;
|
||||
|
||||
using namespace ::ck_tile::builder::test;
|
||||
using enum ::ck_tile::builder::ConvDirection;
|
||||
|
||||
TEST(ConvTensorLayout, AssignsLayoutsFor1D_NWGC_GKXC_NWGK)
|
||||
{
|
||||
using TensorLayouts = ConvTensorLayouts<ckb::GroupConvLayout1D::NWGC_GKXC_NWGK, 1, FORWARD>;
|
||||
static constexpr auto sig =
|
||||
ConvSignature<>{.spatial_dim = 1,
|
||||
.direction = FORWARD,
|
||||
.data_type = DataType::FP16,
|
||||
.accumulation_data_type = DataType::FP32,
|
||||
.input = {.config = {.layout = TensorLayout::NWGC}},
|
||||
.weight = {.config = {.layout = TensorLayout::GKXC}},
|
||||
.output = {.config = {.layout = TensorLayout::NWGK}}};
|
||||
|
||||
using TensorLayouts = ConvTensorLayouts<sig, 1, FORWARD>;
|
||||
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ALayout, ck::tensor_layout::convolution::NWGC>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::BLayout, ck::tensor_layout::convolution::GKXC>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ELayout, ck::tensor_layout::convolution::NWGK>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::DsLayout, ck::Tuple<>>));
|
||||
}
|
||||
|
||||
TEST(ConvTensorLayout, AssignsLayoutsFor1D_NGCW_GKXC_NGKW)
|
||||
{
|
||||
using TensorLayouts = ConvTensorLayouts<ckb::GroupConvLayout1D::NGCW_GKXC_NGKW, 1, FORWARD>;
|
||||
static constexpr auto sig =
|
||||
ConvSignature<>{.spatial_dim = 1,
|
||||
.direction = FORWARD,
|
||||
.data_type = DataType::FP16,
|
||||
.accumulation_data_type = DataType::FP32,
|
||||
.input = {.config = {.layout = TensorLayout::NGCW}},
|
||||
.weight = {.config = {.layout = TensorLayout::GKXC}},
|
||||
.output = {.config = {.layout = TensorLayout::NGKW}}};
|
||||
|
||||
using TensorLayouts = ConvTensorLayouts<sig, 1, FORWARD>;
|
||||
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ALayout, ck::tensor_layout::convolution::NGCW>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::BLayout, ck::tensor_layout::convolution::GKXC>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ELayout, ck::tensor_layout::convolution::NGKW>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::DsLayout, ck::Tuple<>>));
|
||||
}
|
||||
|
||||
TEST(ConvTensorLayout, AssignsLayoutsFor1D_GNWC_GKXC_GNWK)
|
||||
{
|
||||
using TensorLayouts = ConvTensorLayouts<ckb::GroupConvLayout1D::GNWC_GKXC_GNWK, 1, FORWARD>;
|
||||
static constexpr auto sig =
|
||||
ConvSignature<>{.spatial_dim = 1,
|
||||
.direction = FORWARD,
|
||||
.data_type = DataType::FP16,
|
||||
.accumulation_data_type = DataType::FP32,
|
||||
.input = {.config = {.layout = TensorLayout::GNWC}},
|
||||
.weight = {.config = {.layout = TensorLayout::GKXC}},
|
||||
.output = {.config = {.layout = TensorLayout::GNWK}}};
|
||||
|
||||
using TensorLayouts = ConvTensorLayouts<sig, 1, FORWARD>;
|
||||
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ALayout, ck::tensor_layout::convolution::GNWC>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::BLayout, ck::tensor_layout::convolution::GKXC>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ELayout, ck::tensor_layout::convolution::GNWK>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::DsLayout, ck::Tuple<>>));
|
||||
}
|
||||
|
||||
TEST(ConvTensorLayout, AssignsLayoutsFor1D_NGCW_GKCX_NGKW)
|
||||
{
|
||||
using TensorLayouts = ConvTensorLayouts<ckb::GroupConvLayout1D::NGCW_GKCX_NGKW, 1, FORWARD>;
|
||||
static constexpr auto sig =
|
||||
ConvSignature<>{.spatial_dim = 1,
|
||||
.direction = FORWARD,
|
||||
.data_type = DataType::FP16,
|
||||
.accumulation_data_type = DataType::FP32,
|
||||
.input = {.config = {.layout = TensorLayout::NGCW}},
|
||||
.weight = {.config = {.layout = TensorLayout::GKCX}},
|
||||
.output = {.config = {.layout = TensorLayout::NGKW}}};
|
||||
|
||||
using TensorLayouts = ConvTensorLayouts<sig, 1, FORWARD>;
|
||||
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ALayout, ck::tensor_layout::convolution::NGCW>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::BLayout, ck::tensor_layout::convolution::GKCX>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ELayout, ck::tensor_layout::convolution::NGKW>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::DsLayout, ck::Tuple<>>));
|
||||
}
|
||||
|
||||
TEST(ConvTensorLayout, AssignsLayoutsFor2D_NGCHW_GKYXC_NGKHW)
|
||||
{
|
||||
using TensorLayouts = ConvTensorLayouts<ckb::GroupConvLayout2D::NGCHW_GKYXC_NGKHW, 2, FORWARD>;
|
||||
static constexpr auto sig =
|
||||
ConvSignature<>{.spatial_dim = 2,
|
||||
.direction = FORWARD,
|
||||
.data_type = DataType::FP16,
|
||||
.accumulation_data_type = DataType::FP32,
|
||||
.input = {.config = {.layout = TensorLayout::NGCHW}},
|
||||
.weight = {.config = {.layout = TensorLayout::GKYXC}},
|
||||
.output = {.config = {.layout = TensorLayout::NGKHW}}};
|
||||
|
||||
using TensorLayouts = ConvTensorLayouts<sig, 2, FORWARD>;
|
||||
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ALayout, ck::tensor_layout::convolution::NGCHW>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::BLayout, ck::tensor_layout::convolution::GKYXC>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ELayout, ck::tensor_layout::convolution::NGKHW>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::DsLayout, ck::Tuple<>>));
|
||||
}
|
||||
|
||||
TEST(ConvTensorLayout, AssignsLayoutsFor2D_NHWGC_GKYXC_NHWGK)
|
||||
{
|
||||
using TensorLayouts = ConvTensorLayouts<ckb::GroupConvLayout2D::NHWGC_GKYXC_NHWGK, 2, FORWARD>;
|
||||
static constexpr auto sig =
|
||||
ConvSignature<>{.spatial_dim = 2,
|
||||
.direction = FORWARD,
|
||||
.data_type = DataType::FP16,
|
||||
.accumulation_data_type = DataType::FP32,
|
||||
.input = {.config = {.layout = TensorLayout::NHWGC}},
|
||||
.weight = {.config = {.layout = TensorLayout::GKYXC}},
|
||||
.output = {.config = {.layout = TensorLayout::NHWGK}}};
|
||||
|
||||
using TensorLayouts = ConvTensorLayouts<sig, 2, FORWARD>;
|
||||
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ALayout, ck::tensor_layout::convolution::NHWGC>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::BLayout, ck::tensor_layout::convolution::GKYXC>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ELayout, ck::tensor_layout::convolution::NHWGK>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::DsLayout, ck::Tuple<>>));
|
||||
}
|
||||
|
||||
TEST(ConvTensorLayout, AssignsLayoutsFor2D_GNHWC_GKYXC_GNHWK)
|
||||
{
|
||||
using TensorLayouts = ConvTensorLayouts<ckb::GroupConvLayout2D::GNHWC_GKYXC_GNHWK, 2, FORWARD>;
|
||||
static constexpr auto sig =
|
||||
ConvSignature<>{.spatial_dim = 2,
|
||||
.direction = FORWARD,
|
||||
.data_type = DataType::FP16,
|
||||
.accumulation_data_type = DataType::FP32,
|
||||
.input = {.config = {.layout = TensorLayout::GNHWC}},
|
||||
.weight = {.config = {.layout = TensorLayout::GKYXC}},
|
||||
.output = {.config = {.layout = TensorLayout::GNHWK}}};
|
||||
|
||||
using TensorLayouts = ConvTensorLayouts<sig, 2, FORWARD>;
|
||||
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ALayout, ck::tensor_layout::convolution::GNHWC>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::BLayout, ck::tensor_layout::convolution::GKYXC>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ELayout, ck::tensor_layout::convolution::GNHWK>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::DsLayout, ck::Tuple<>>));
|
||||
}
|
||||
|
||||
TEST(ConvTensorLayout, AssignsLayoutsFor2D_NGCHW_GKCYX_NGKHW)
|
||||
{
|
||||
using TensorLayouts = ConvTensorLayouts<ckb::GroupConvLayout2D::NGCHW_GKCYX_NGKHW, 2, FORWARD>;
|
||||
static constexpr auto sig =
|
||||
ConvSignature<>{.spatial_dim = 2,
|
||||
.direction = FORWARD,
|
||||
.data_type = DataType::FP16,
|
||||
.accumulation_data_type = DataType::FP32,
|
||||
.input = {.config = {.layout = TensorLayout::NGCHW}},
|
||||
.weight = {.config = {.layout = TensorLayout::GKCYX}},
|
||||
.output = {.config = {.layout = TensorLayout::NGKHW}}};
|
||||
|
||||
using TensorLayouts = ConvTensorLayouts<sig, 2, FORWARD>;
|
||||
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ALayout, ck::tensor_layout::convolution::NGCHW>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::BLayout, ck::tensor_layout::convolution::GKCYX>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ELayout, ck::tensor_layout::convolution::NGKHW>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::DsLayout, ck::Tuple<>>));
|
||||
}
|
||||
|
||||
TEST(ConvTensorLayout, AssignsLayoutsFor3D_NGCDHW_GKCZYX_NGKDHW)
|
||||
{
|
||||
using TensorLayouts =
|
||||
ConvTensorLayouts<ckb::GroupConvLayout3D::NGCDHW_GKCZYX_NGKDHW, 3, FORWARD>;
|
||||
static constexpr auto sig =
|
||||
ConvSignature<>{.spatial_dim = 3,
|
||||
.direction = FORWARD,
|
||||
.data_type = DataType::FP16,
|
||||
.accumulation_data_type = DataType::FP32,
|
||||
.input = {.config = {.layout = TensorLayout::NGCDHW}},
|
||||
.weight = {.config = {.layout = TensorLayout::GKCZYX}},
|
||||
.output = {.config = {.layout = TensorLayout::NGKDHW}}};
|
||||
|
||||
using TensorLayouts = ConvTensorLayouts<sig, 3, FORWARD>;
|
||||
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ALayout, ck::tensor_layout::convolution::NGCDHW>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::BLayout, ck::tensor_layout::convolution::GKCZYX>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ELayout, ck::tensor_layout::convolution::NGKDHW>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::DsLayout, ck::Tuple<>>));
|
||||
}
|
||||
|
||||
TEST(ConvTensorLayout, AssignsLayoutsFor3D_NDHWGC_GKZYXC_NDHWGK)
|
||||
{
|
||||
using TensorLayouts =
|
||||
ConvTensorLayouts<ckb::GroupConvLayout3D::NDHWGC_GKZYXC_NDHWGK, 3, FORWARD>;
|
||||
static constexpr auto sig =
|
||||
ConvSignature<>{.spatial_dim = 3,
|
||||
.direction = FORWARD,
|
||||
.data_type = DataType::FP16,
|
||||
.accumulation_data_type = DataType::FP32,
|
||||
.input = {.config = {.layout = TensorLayout::NDHWGC}},
|
||||
.weight = {.config = {.layout = TensorLayout::GKZYXC}},
|
||||
.output = {.config = {.layout = TensorLayout::NDHWGK}}};
|
||||
|
||||
using TensorLayouts = ConvTensorLayouts<sig, 3, FORWARD>;
|
||||
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ALayout, ck::tensor_layout::convolution::NDHWGC>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::BLayout, ck::tensor_layout::convolution::GKZYXC>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ELayout, ck::tensor_layout::convolution::NDHWGK>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::DsLayout, ck::Tuple<>>));
|
||||
}
|
||||
|
||||
TEST(ConvTensorLayout, AssignsLayoutsFor3D_GNDHWC_GKZYXC_GNDHWK)
|
||||
{
|
||||
using TensorLayouts =
|
||||
ConvTensorLayouts<ckb::GroupConvLayout3D::GNDHWC_GKZYXC_GNDHWK, 3, FORWARD>;
|
||||
static constexpr auto sig =
|
||||
ConvSignature<>{.spatial_dim = 3,
|
||||
.direction = FORWARD,
|
||||
.data_type = DataType::FP16,
|
||||
.accumulation_data_type = DataType::FP32,
|
||||
.input = {.config = {.layout = TensorLayout::GNDHWC}},
|
||||
.weight = {.config = {.layout = TensorLayout::GKZYXC}},
|
||||
.output = {.config = {.layout = TensorLayout::GNDHWK}}};
|
||||
|
||||
using TensorLayouts = ConvTensorLayouts<sig, 3, FORWARD>;
|
||||
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ALayout, ck::tensor_layout::convolution::GNDHWC>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::BLayout, ck::tensor_layout::convolution::GKZYXC>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ELayout, ck::tensor_layout::convolution::GNDHWK>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::DsLayout, ck::Tuple<>>));
|
||||
}
|
||||
|
||||
TEST(AuxiliaryTensorLayout, AssignsLayoutForG_K_strided)
|
||||
{
|
||||
using CKLayout = LayoutToCK<TensorLayout::G_K_strided>::type;
|
||||
EXPECT_TRUE((std::is_same_v<CKLayout, ck::tensor_layout::convolution::G_K>));
|
||||
}
|
||||
|
||||
TEST(AuxiliaryTensorLayout, AssignsLayoutForGC)
|
||||
{
|
||||
using CKLayout = LayoutToCK<TensorLayout::GC>::type;
|
||||
EXPECT_TRUE((std::is_same_v<CKLayout, ck::tensor_layout::convolution::GC>));
|
||||
}
|
||||
|
||||
TEST(AuxiliaryTensorLayout, AssignsLayoutForG_C_strided)
|
||||
{
|
||||
using CKLayout = LayoutToCK<TensorLayout::G_C_strided>::type;
|
||||
EXPECT_TRUE((std::is_same_v<CKLayout, ck::tensor_layout::convolution::G_C>));
|
||||
}
|
||||
|
||||
TEST(AuxiliaryTensorLayout, EmptyAuxiliaryTensorLayoutIsEmptyTuple)
|
||||
{
|
||||
using ::ck_tile::builder::factory::internal::EmptyAuxiliaryTensorLayout;
|
||||
using EmptyLayout = EmptyAuxiliaryTensorLayout::type;
|
||||
EXPECT_TRUE((std::is_same_v<EmptyLayout, ck::Tuple<>>));
|
||||
}
|
||||
|
||||
struct MockAuxiliaryTensorConfig
|
||||
{
|
||||
TensorLayout layout;
|
||||
};
|
||||
|
||||
TEST(AuxiliaryTensorLayoutIntegration, SingleBiasTensorWithG_K_Layout)
|
||||
{
|
||||
static constexpr std::array<MockAuxiliaryTensorConfig, 1> aux_configs = {
|
||||
MockAuxiliaryTensorConfig{.layout = TensorLayout::G_K_strided}};
|
||||
|
||||
using AuxLayouts = AuxiliaryTensorLayouts<aux_configs, 2, FORWARD>;
|
||||
|
||||
EXPECT_EQ(AuxLayouts::Size, 1);
|
||||
using ExpectedType = ck::Tuple<ck::tensor_layout::convolution::G_K>;
|
||||
EXPECT_TRUE((std::is_same_v<AuxLayouts::type, ExpectedType>));
|
||||
}
|
||||
|
||||
TEST(AuxiliaryTensorLayoutIntegration, SingleBiasTensorWithGC_Layout)
|
||||
{
|
||||
static constexpr std::array<MockAuxiliaryTensorConfig, 1> aux_configs = {
|
||||
MockAuxiliaryTensorConfig{.layout = TensorLayout::GC}};
|
||||
|
||||
using AuxLayouts = AuxiliaryTensorLayouts<aux_configs, 2, FORWARD>;
|
||||
|
||||
EXPECT_EQ(AuxLayouts::Size, 1);
|
||||
using ExpectedType = ck::Tuple<ck::tensor_layout::convolution::GC>;
|
||||
EXPECT_TRUE((std::is_same_v<AuxLayouts::type, ExpectedType>));
|
||||
}
|
||||
|
||||
TEST(AuxiliaryTensorLayoutIntegration, SingleBiasTensorWithG_C_Layout)
|
||||
{
|
||||
static constexpr std::array<MockAuxiliaryTensorConfig, 1> aux_configs = {
|
||||
MockAuxiliaryTensorConfig{.layout = TensorLayout::G_C_strided}};
|
||||
|
||||
using AuxLayouts = AuxiliaryTensorLayouts<aux_configs, 2, FORWARD>;
|
||||
|
||||
EXPECT_EQ(AuxLayouts::Size, 1);
|
||||
using ExpectedType = ck::Tuple<ck::tensor_layout::convolution::G_C>;
|
||||
EXPECT_TRUE((std::is_same_v<AuxLayouts::type, ExpectedType>));
|
||||
}
|
||||
|
||||
TEST(AuxiliaryTensorLayoutIntegration, TwoAuxiliaryTensors)
|
||||
{
|
||||
static constexpr std::array<MockAuxiliaryTensorConfig, 2> aux_configs = {
|
||||
MockAuxiliaryTensorConfig{.layout = TensorLayout::G_K_strided},
|
||||
MockAuxiliaryTensorConfig{.layout = TensorLayout::GC}};
|
||||
|
||||
using AuxLayouts = AuxiliaryTensorLayouts<aux_configs, 2, FORWARD>;
|
||||
|
||||
EXPECT_EQ(AuxLayouts::Size, 2);
|
||||
using ExpectedType =
|
||||
ck::Tuple<ck::tensor_layout::convolution::G_K, ck::tensor_layout::convolution::GC>;
|
||||
EXPECT_TRUE((std::is_same_v<AuxLayouts::type, ExpectedType>));
|
||||
}
|
||||
|
||||
TEST(AuxiliaryTensorLayoutIntegration, ThreeAuxiliaryTensors)
|
||||
{
|
||||
static constexpr std::array<MockAuxiliaryTensorConfig, 3> aux_configs = {
|
||||
MockAuxiliaryTensorConfig{.layout = TensorLayout::G_K_strided},
|
||||
MockAuxiliaryTensorConfig{.layout = TensorLayout::GC},
|
||||
MockAuxiliaryTensorConfig{.layout = TensorLayout::G_C_strided}};
|
||||
|
||||
using AuxLayouts = AuxiliaryTensorLayouts<aux_configs, 2, FORWARD>;
|
||||
|
||||
EXPECT_EQ(AuxLayouts::Size, 3);
|
||||
using ExpectedType = ck::Tuple<ck::tensor_layout::convolution::G_K,
|
||||
ck::tensor_layout::convolution::GC,
|
||||
ck::tensor_layout::convolution::G_C>;
|
||||
EXPECT_TRUE((std::is_same_v<AuxLayouts::type, ExpectedType>));
|
||||
}
|
||||
|
||||
TEST(AuxiliaryTensorLayoutIntegration, WorksWith1DConvolution)
|
||||
{
|
||||
static constexpr std::array<MockAuxiliaryTensorConfig, 1> aux_configs = {
|
||||
MockAuxiliaryTensorConfig{.layout = TensorLayout::G_K_strided}};
|
||||
|
||||
using AuxLayouts = AuxiliaryTensorLayouts<aux_configs, 1, FORWARD>;
|
||||
|
||||
EXPECT_EQ(AuxLayouts::Size, 1);
|
||||
using ExpectedType = ck::Tuple<ck::tensor_layout::convolution::G_K>;
|
||||
EXPECT_TRUE((std::is_same_v<AuxLayouts::type, ExpectedType>));
|
||||
}
|
||||
|
||||
TEST(AuxiliaryTensorLayoutIntegration, WorksWith3DConvolution)
|
||||
{
|
||||
static constexpr std::array<MockAuxiliaryTensorConfig, 1> aux_configs = {
|
||||
MockAuxiliaryTensorConfig{.layout = TensorLayout::GC}};
|
||||
|
||||
using AuxLayouts = AuxiliaryTensorLayouts<aux_configs, 3, FORWARD>;
|
||||
|
||||
EXPECT_EQ(AuxLayouts::Size, 1);
|
||||
using ExpectedType = ck::Tuple<ck::tensor_layout::convolution::GC>;
|
||||
EXPECT_TRUE((std::is_same_v<AuxLayouts::type, ExpectedType>));
|
||||
}
|
||||
|
||||
TEST(ConvTensorLayoutsWithAuxiliary, Conv2DWithSingleBiasG_K)
|
||||
{
|
||||
using OutputOp = TensorOperation<TensorConfig{.layout = TensorLayout::G_K_strided}>;
|
||||
|
||||
static constexpr auto sig =
|
||||
ConvSignature<ConvolutionTensor<>, ConvolutionTensor<>, ConvolutionTensor<OutputOp>>{
|
||||
.spatial_dim = 2,
|
||||
.direction = FORWARD,
|
||||
.data_type = DataType::FP16,
|
||||
.accumulation_data_type = DataType::FP32,
|
||||
.input = {.config = {.layout = TensorLayout::NGCHW}},
|
||||
.weight = {.config = {.layout = TensorLayout::GKYXC}},
|
||||
.output = {.config = {.layout = TensorLayout::NGKHW},
|
||||
.operation =
|
||||
OutputOp{.elementwise_operation = ElementwiseOperation::SCALE}}};
|
||||
|
||||
using TensorLayouts = ConvTensorLayouts<sig, 2, FORWARD>;
|
||||
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ALayout, ck::tensor_layout::convolution::NGCHW>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::BLayout, ck::tensor_layout::convolution::GKYXC>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ELayout, ck::tensor_layout::convolution::NGKHW>));
|
||||
|
||||
using ExpectedDsLayout = ck::Tuple<ck::tensor_layout::convolution::G_K>;
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::DsLayout, ExpectedDsLayout>));
|
||||
}
|
||||
|
||||
TEST(ConvTensorLayoutsWithAuxiliary, Conv2DWithSingleBiasGC)
|
||||
{
|
||||
using OutputOp = TensorOperation<TensorConfig{.layout = TensorLayout::GC}>;
|
||||
|
||||
static constexpr auto sig =
|
||||
ConvSignature<ConvolutionTensor<>, ConvolutionTensor<>, ConvolutionTensor<OutputOp>>{
|
||||
.spatial_dim = 2,
|
||||
.direction = FORWARD,
|
||||
.data_type = DataType::BF16,
|
||||
.accumulation_data_type = DataType::FP32,
|
||||
.input = {.config = {.layout = TensorLayout::NHWGC}},
|
||||
.weight = {.config = {.layout = TensorLayout::GKYXC}},
|
||||
.output = {.config = {.layout = TensorLayout::NHWGK},
|
||||
.operation =
|
||||
OutputOp{.elementwise_operation = ElementwiseOperation::SCALE}}};
|
||||
|
||||
using TensorLayouts = ConvTensorLayouts<sig, 2, FORWARD>;
|
||||
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ALayout, ck::tensor_layout::convolution::NHWGC>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::BLayout, ck::tensor_layout::convolution::GKYXC>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ELayout, ck::tensor_layout::convolution::NHWGK>));
|
||||
|
||||
using ExpectedDsLayout = ck::Tuple<ck::tensor_layout::convolution::GC>;
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::DsLayout, ExpectedDsLayout>));
|
||||
}
|
||||
|
||||
TEST(ConvTensorLayoutsWithAuxiliary, Conv2DWithTwoAuxiliaryTensors)
|
||||
{
|
||||
using OutputOp = TensorOperation<TensorConfig{.layout = TensorLayout::G_K_strided},
|
||||
TensorConfig{.layout = TensorLayout::GC}>;
|
||||
|
||||
static constexpr auto sig =
|
||||
ConvSignature<ConvolutionTensor<>, ConvolutionTensor<>, ConvolutionTensor<OutputOp>>{
|
||||
.spatial_dim = 2,
|
||||
.direction = FORWARD,
|
||||
.data_type = DataType::FP16,
|
||||
.accumulation_data_type = DataType::FP32,
|
||||
.input = {.config = {.layout = TensorLayout::GNHWC}},
|
||||
.weight = {.config = {.layout = TensorLayout::GKYXC}},
|
||||
.output = {.config = {.layout = TensorLayout::GNHWK},
|
||||
.operation = OutputOp{.elementwise_operation =
|
||||
ElementwiseOperation::SCALEADD_SCALEADD_RELU}}};
|
||||
|
||||
using TensorLayouts = ConvTensorLayouts<sig, 2, FORWARD>;
|
||||
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ALayout, ck::tensor_layout::convolution::GNHWC>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::BLayout, ck::tensor_layout::convolution::GKYXC>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ELayout, ck::tensor_layout::convolution::GNHWK>));
|
||||
|
||||
using ExpectedDsLayout =
|
||||
ck::Tuple<ck::tensor_layout::convolution::G_K, ck::tensor_layout::convolution::GC>;
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::DsLayout, ExpectedDsLayout>));
|
||||
}
|
||||
|
||||
TEST(ConvTensorLayoutsWithAuxiliary, Conv1DWithBias)
|
||||
{
|
||||
using OutputOp = TensorOperation<TensorConfig{.layout = TensorLayout::G_K_strided}>;
|
||||
|
||||
static constexpr auto sig =
|
||||
ConvSignature<ConvolutionTensor<>, ConvolutionTensor<>, ConvolutionTensor<OutputOp>>{
|
||||
.spatial_dim = 1,
|
||||
.direction = FORWARD,
|
||||
.data_type = DataType::FP32,
|
||||
.accumulation_data_type = DataType::FP32,
|
||||
.input = {.config = {.layout = TensorLayout::NWGC}},
|
||||
.weight = {.config = {.layout = TensorLayout::GKXC}},
|
||||
.output = {.config = {.layout = TensorLayout::NWGK},
|
||||
.operation =
|
||||
OutputOp{.elementwise_operation = ElementwiseOperation::SCALE}}};
|
||||
|
||||
using TensorLayouts = ConvTensorLayouts<sig, 1, FORWARD>;
|
||||
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ALayout, ck::tensor_layout::convolution::NWGC>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::BLayout, ck::tensor_layout::convolution::GKXC>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ELayout, ck::tensor_layout::convolution::NWGK>));
|
||||
|
||||
using ExpectedDsLayout = ck::Tuple<ck::tensor_layout::convolution::G_K>;
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::DsLayout, ExpectedDsLayout>));
|
||||
}
|
||||
|
||||
TEST(ConvTensorLayoutsWithAuxiliary, Conv3DWithBias)
|
||||
{
|
||||
using OutputOp = TensorOperation<TensorConfig{.layout = TensorLayout::G_C_strided}>;
|
||||
|
||||
static constexpr auto sig =
|
||||
ConvSignature<ConvolutionTensor<>, ConvolutionTensor<>, ConvolutionTensor<OutputOp>>{
|
||||
.spatial_dim = 3,
|
||||
.direction = FORWARD,
|
||||
.data_type = DataType::FP16,
|
||||
.accumulation_data_type = DataType::FP32,
|
||||
.input = {.config = {.layout = TensorLayout::NDHWGC}},
|
||||
.weight = {.config = {.layout = TensorLayout::GKZYXC}},
|
||||
.output = {.config = {.layout = TensorLayout::NDHWGK},
|
||||
.operation = OutputOp{.elementwise_operation =
|
||||
ElementwiseOperation::BIAS_BNORM_CLAMP}}};
|
||||
|
||||
using TensorLayouts = ConvTensorLayouts<sig, 3, FORWARD>;
|
||||
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ALayout, ck::tensor_layout::convolution::NDHWGC>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::BLayout, ck::tensor_layout::convolution::GKZYXC>));
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::ELayout, ck::tensor_layout::convolution::NDHWGK>));
|
||||
|
||||
using ExpectedDsLayout = ck::Tuple<ck::tensor_layout::convolution::G_C>;
|
||||
EXPECT_TRUE((std::is_same_v<TensorLayouts::DsLayout, ExpectedDsLayout>));
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
Reference in New Issue
Block a user