From fe6140353fe6a97e4d2056bb75ed614e1aa673f2 Mon Sep 17 00:00:00 2001 From: John Shumway Date: Sun, 14 Sep 2025 14:03:35 +0000 Subject: [PATCH] Rename layouts to channels first or channels last. Also add a spatial dimension parameter to layout mapping in the factory to handle different layout enums for 2D and 3D convolutions --- .../include/ck_tile/builder/conv_factory.hpp | 20 ++++++++++++------- .../ck_tile/builder/conv_signature.hpp | 5 ++--- .../builder/test/test_conv_builder.cpp | 2 +- .../builder/test/test_conv_grp_fwd_2d.cpp | 2 +- .../builder/test/test_conv_grp_fwd_3d.cpp | 2 +- .../builder/test/test_conv_instances.cpp | 2 +- 6 files changed, 19 insertions(+), 14 deletions(-) diff --git a/experimental/builder/include/ck_tile/builder/conv_factory.hpp b/experimental/builder/include/ck_tile/builder/conv_factory.hpp index 97428c4dbe..7f683273e8 100644 --- a/experimental/builder/include/ck_tile/builder/conv_factory.hpp +++ b/experimental/builder/include/ck_tile/builder/conv_factory.hpp @@ -12,11 +12,18 @@ namespace ck_tile::builder { // Type mappings from the builder GroupConvLayout enum class to the CK tensor data types. -template -struct ConvTensorLayouts; +template + requires(ConvSpatialDim) +struct ConvTensorLayouts +{ + // This will trigger if a specialization for the given layout is not found. + // We should always catch this in an earlier validation check. + static_assert(sizeof(Layout) == 0, + "Internal error. Unsupported layout for convolution factory."); +}; template <> -struct ConvTensorLayouts +struct ConvTensorLayouts { // Channels first convolution layout. using ALayout = ck::tensor_layout::convolution::NHWGC; @@ -26,7 +33,7 @@ struct ConvTensorLayouts }; template <> -struct ConvTensorLayouts +struct ConvTensorLayouts { // Channels last convolution layout. using ALayout = ck::tensor_layout::convolution::NHWGC; @@ -36,7 +43,7 @@ struct ConvTensorLayouts }; template <> -struct ConvTensorLayouts +struct ConvTensorLayouts { // Channels last convolution layout. using ALayout = ck::tensor_layout::convolution::NDHWGC; @@ -45,7 +52,6 @@ struct ConvTensorLayouts using ELayout = ck::tensor_layout::convolution::NDHWGK; }; - // Type mappings from builder convolution data type to CK tensor types. template struct ConvTensorTypes @@ -283,7 +289,7 @@ template ; + using Layouts = ConvTensorLayouts; using Types = ConvTensorTypes; using Ops = ConvPassThroughOps; static constexpr ConvSpec SPECIALIZATION{ diff --git a/experimental/builder/include/ck_tile/builder/conv_signature.hpp b/experimental/builder/include/ck_tile/builder/conv_signature.hpp index 165508c123..c9e9857ae7 100644 --- a/experimental/builder/include/ck_tile/builder/conv_signature.hpp +++ b/experimental/builder/include/ck_tile/builder/conv_signature.hpp @@ -10,9 +10,8 @@ namespace ck_tile::builder { // Layouts for grouped convolutions. enum class GroupConvLayout { - NHWGC_GKYXC_NHWGK, // Channels-last - NDHWGC_GKZYXC_NDHWGK, // Channels-last - NGCHW_GKCYX_NGKHW // Channels-first + CHANNELS_LAST, // Channels-last NHWGC_GKYXC_NHWGK + CHANNELS_FIRST // Channels-first NGCHW_GKCYX_NGKHW }; // Spatial dimensionalities of grouped convolutions. diff --git a/experimental/builder/test/test_conv_builder.cpp b/experimental/builder/test/test_conv_builder.cpp index 4eab3cd6d0..38612e53f5 100644 --- a/experimental/builder/test/test_conv_builder.cpp +++ b/experimental/builder/test/test_conv_builder.cpp @@ -13,7 +13,7 @@ struct ConvSignature { int spatial_dim = 2; ckb::ConvDirection direction = ckb::ConvDirection::Forward; - ckb::GroupConvLayout layout = ckb::GroupConvLayout::NHWGC_GKYXC_NHWGK; + ckb::GroupConvLayout layout = ckb::GroupConvLayout::CHANNELS_LAST; ckb::DataType data_type = ckb::DataType::FP16; }; static_assert(ckb::ConvSignatureDescriptor); diff --git a/experimental/builder/test/test_conv_grp_fwd_2d.cpp b/experimental/builder/test/test_conv_grp_fwd_2d.cpp index 45b1741ceb..a27e2c3dc6 100644 --- a/experimental/builder/test/test_conv_grp_fwd_2d.cpp +++ b/experimental/builder/test/test_conv_grp_fwd_2d.cpp @@ -13,7 +13,7 @@ struct ConvSignature { int spatial_dim = 2; ckb::ConvDirection direction = ckb::ConvDirection::Forward; - ckb::GroupConvLayout layout = ckb::GroupConvLayout::NHWGC_GKYXC_NHWGK; + ckb::GroupConvLayout layout = ckb::GroupConvLayout::CHANNELS_LAST; ckb::DataType data_type = ckb::DataType::FP16; }; static_assert(ckb::ConvSignatureDescriptor); diff --git a/experimental/builder/test/test_conv_grp_fwd_3d.cpp b/experimental/builder/test/test_conv_grp_fwd_3d.cpp index 9c7001d4f0..c29b3f46cf 100644 --- a/experimental/builder/test/test_conv_grp_fwd_3d.cpp +++ b/experimental/builder/test/test_conv_grp_fwd_3d.cpp @@ -13,7 +13,7 @@ struct ConvSignature { int spatial_dim = 3; ckb::ConvDirection direction = ckb::ConvDirection::Forward; - ckb::GroupConvLayout layout = ckb::GroupConvLayout::NDHWGC_GKZYXC_NDHWGK; + ckb::GroupConvLayout layout = ckb::GroupConvLayout::CHANNELS_LAST; ckb::DataType data_type = ckb::DataType::BF16; }; static_assert(ckb::ConvSignatureDescriptor); diff --git a/experimental/builder/test/test_conv_instances.cpp b/experimental/builder/test/test_conv_instances.cpp index 3c458ea540..e0aea59116 100644 --- a/experimental/builder/test/test_conv_instances.cpp +++ b/experimental/builder/test/test_conv_instances.cpp @@ -19,7 +19,7 @@ struct ConvSignature { int spatial_dim = 2; ckb::ConvDirection direction = ckb::ConvDirection::Forward; - ckb::GroupConvLayout layout = ckb::GroupConvLayout::NHWGC_GKYXC_NHWGK; + ckb::GroupConvLayout layout = ckb::GroupConvLayout::CHANNELS_LAST; ckb::DataType data_type = ckb::DataType::FP16; }; static_assert(ckb::ConvSignatureDescriptor);