[CK_BUILDER] Refactor convolution signature to provide data type/layout/elementwise op per tensor #3331

vpietila-amd · 2025-12-01T14:34:02Z

Proposed changes

Refactored the CK Builder convolution signature such that

data type
data type for computation
layout
elementwise operation

are defined per tensor. This refactoring allows us build instances for complex fused operations such as scale-add-scale-add-relu. I added a fwd conv builder test that demonstrates how the complex ops benefit from the new signature design pattern. At the high level, the signature is composed as follows

┌─────────────────────────────────────────────────────────┐
│                    ConvSignature                        │
├─────────────────────────────────────────────────────────┤
│ Properties:                                             │
│   • spatial_dim: int           (1D, 2D, or 3D)          │
│   • direction: ConvDirection   (Fwd/BwdData/BwdWeight)  │
│   • data_type: DataType        (default data type)      │
│   • accumulation_data_type: DataType                    │
│   • input: ConvTensor          ──┐                      │
│   • weight: ConvTensor         ──│                      │
│   • output: ConvTensor         ──│                      │
└──────────────────────────────────┼──────────────────────┘
                                   │
                                   ▼
              ┌─────────────────────────────────────────┐
              │           ConvTensor                    │
              ├─────────────────────────────────────────┤
              │ ╔═════════════════════════════════════╗ │
              │ ║ TensorConfig (required)             ║ │
              │ ╠═════════════════════════════════════╣ │
              │ ║  • layout: ConvLayout               ║ │
              │ ║  • data_type: DataType (optional)   ║ │
              │ ║  • compute_type: DataType (optional)║ │
              │ ╚═════════════════════════════════════╝ │
              │                                         │
              │ ┌─────────────────────────────────────┐ │
              │ │ TensorOperation (optional)          │ │
              │ ├─────────────────────────────────────┤ │
              │ │  • elementwise_operation            │ │
              │ │  • auxiliary_operand_configs[]      │ │
              │ │    (each is also ConvTensor)  ◄───────┼─┐
              │ └─────────────────────────────────────┘ │ │
              └─────────────────────────────────────────┘ │
                                                          │
                                 Recursive ───────────────┘

Design Points:

ConvSignature contains three ConvTensor instances (input, weight, output)
All tensors share the same ConvTensor structure
Each ConvTensor has:
- TensorConfig (required): Defines layout as well as optional data and compute type overrides
- TensorOperation (optional): Defines fused elementwise operations
Auxiliary operands (e.g., bias) in TensorOperation also use the ConvTensor type

The convolution traits and descriptors do not fully utilize the new structure, but all tests are passing. They could be refactored separately.

… tensors.

…ve-elementwise-ops

shumway · 2025-12-04T07:22:15Z

experimental/builder/include/ck_tile/builder/README.md

@@ -0,0 +1,244 @@
+# Composable Kernel Builder Design Documentation


This is really nice documentation! Thanks for taking time to create this!

shumway · 2025-12-04T07:28:02Z

experimental/builder/include/ck_tile/builder/types.hpp


 enum class DataType
 {
+    UNDEFINDED = 0,


Nit: UNDEFINED

Only fix if you need to make other changes.

shumway · 2025-12-04T07:32:03Z

experimental/builder/include/ck_tile/builder/types.hpp

-// G: Group, N: Batch, K: Output Channel, C: Input Channel, W: Width
-// Enum defines Input, Weight, and Output tensor layouts respectively.
-enum class GroupConvLayout1D
+enum class TensorLayout


This is pretty clean. Another strategy I sometimes see with flat lists of enums like this is to add extra enum values for start and end of ranges of grouped elements to enable classification checks with math. We can add that later if it helps logic, but concepts and lists with constexpr std::array<TensorLayout> can also be used if we need more structure.

shumway · 2025-12-04T07:34:27Z

experimental/builder/include/ck_tile/builder/types.hpp

-inline std::ostream&
-operator<<(std::ostream& os,
-           const std::variant<GroupConvLayout1D, GroupConvLayout2D, GroupConvLayout3D>& layout)
+inline std::ostream& operator<<(std::ostream& os, TensorLayout layout)


It might be better to have a "toString" function with the switch case, and then use that to define ostream printing.

shumway · 2025-12-04T07:38:10Z

experimental/builder/include/ck_tile/builder/README.md

+              │           ConvTensor                    │
+              ├─────────────────────────────────────────┤
+              │ ╔═════════════════════════════════════╗ │
+              │ ║ TensorConfig (required)             ║ │


It's not clear to me why we need the TensorConfig wrapper instead of just directly having a layout, datatype, and compute_type for the tensor.

shumway · 2025-12-04T07:41:10Z

experimental/builder/test/conv/test_conv_traits.cpp

    EXPECT_EQ(Traits::direction, ck_tile::builder::ConvDirection::FORWARD);
-    EXPECT_EQ(Traits::layout, ck_tile::builder::GroupConvLayout2D::GNHWC_GKYXC_GNHWK);
+    EXPECT_THAT(Traits::layout,
+                ::testing::ElementsAre(ck_tile::builder::TensorLayout::GNHWC,


Note: generally add a using declaration at the top of the test file:

using ::testing::ElementsAre

There are a lot of advantages to that, including shorter code and documenting what helpers you are importing.

(The same goes for test utils, prefer using declarations instead of importing entire namespaces. Once you import an entire namespace, it can quickly become difficult to deduce which functions you are using or even which function is being called.)

shumway · 2025-12-04T07:45:54Z

experimental/builder/test/unit_conv_tensor_layout.cpp

-using ::ck_tile::builder::factory::internal::GetTensorLayout;
+using ::ck_tile::builder::factory::internal::LayoutToCK;
+
+using namespace ::ck_tile::builder::test;


It's much better to have specific using delarations for symbols in a namespace instead of importing everything from a namespace.

shumway · 2025-12-04T07:46:59Z

experimental/builder/test/unit_conv_tensor_layout.cpp

+using ::ck_tile::builder::factory::internal::LayoutToCK;
+
+using namespace ::ck_tile::builder::test;
 using enum ::ck_tile::builder::ConvDirection;


It's much cleaner to only have using enum in a function scope (even if it's repeated many times) rather than have it at file scope.

shumway · 2025-12-04T07:51:16Z

experimental/builder/test/unit_conv_tensor_layout.cpp

+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>));


Google Test has a static assert types eq matcher.
https://google.github.io/googletest/gmock_cook_book.html#restricting-the-type-of-an-argument-or-parameter-in-an-action

shumway

This looks great! I added a lot of comments, but you don't have to implement them unless you end up needing to do more edits (no reason to run this through testing again). The feedback is more for incremental improvement and to help us converge on common best practices.

Ville Pietilä added 28 commits November 26, 2025 13:00

Separate layouts into separate entities for input, weight, and output…

bbe5426

… tensors.

Add test for handling bias tensor layouts.

6d7fdf6

Use instance string in builder tests.

626610f

Add handling of output bias data types and layouts.

edd180a

Generalize handling of the elementwise ops.

0bd50a5

Test fix.

c7c814f

Create builder for layouts.

809c8b4

Layout builder improvements.

f1cff54

Improve layout builder.

3173d94

Simplify bias layout handling.

d358895

Code clean-up.

69bfe64

Move layout utils into separate file.

aba3eea

Remove hard-coded layout combinations.

f00ac4e

Small code clean-up.

74bc17a

Move data type utils into a separate file.

54c58f1

Add data types, layouts, and elementwise ops per conv tensor.

d6fc6c8

Builder bug fixes after refactoring.

05a4067

Working baseline.

c25eb65

Make signature definition look nice in the test code.

a4252c4

Move TensorConfig into test implementations.

7ed339f

Fix all fwd conv builder tests.

beac0e8

Fix conv traits and descriptors tests.

05c46ba

Merge remote-tracking branch 'origin/develop' into vpietila/ckb-impro…

67a607d

…ve-elementwise-ops

More factory assets under a separate directory.

bab26ee

Fix building conv traits.

8a9e22c

Fix clang-format.

979a851

Add Readme doc to describe the design.

e995fcf

Add link to main Readme. Fix links in the builder design doc.

a468fce

vpietila-amd changed the title ~~[CK_TILE] Refactor convolution signature to provide data type/layout/elementwis op per tensor~~ [CK_TILE] Refactor convolution signature to provide data type/layout/elementwise op per tensor Dec 1, 2025

Clean-up data type/layout/elementwise op conversions.

95fbe2d

vpietila-amd marked this pull request as ready for review December 3, 2025 15:09

vpietila-amd requested review from a team, ThomasNing, afagaj, andriy-ca, aosewski, asleepzzz, bartekxk, carlushuang, cgmillette, coderfeli, ddembeckAMD, geyyer, illsilin, poyenc, qianfengz, shumway, tenpercent and vidyasagar-amd as code owners December 3, 2025 15:09

Improve elementwise op handling.

cd37f4e

shumway reviewed Dec 4, 2025

View reviewed changes

shumway approved these changes Dec 4, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[CK_BUILDER] Refactor convolution signature to provide data type/layout/elementwise op per tensor #3331

[CK_BUILDER] Refactor convolution signature to provide data type/layout/elementwise op per tensor #3331

vpietila-amd commented Dec 1, 2025

Uh oh!

shumway Dec 4, 2025

Uh oh!

shumway Dec 4, 2025

Uh oh!

shumway Dec 4, 2025

Uh oh!

shumway Dec 4, 2025

Uh oh!

shumway Dec 4, 2025

Uh oh!

shumway Dec 4, 2025 •

edited

Loading

Uh oh!

shumway Dec 4, 2025

Uh oh!

shumway Dec 4, 2025

Uh oh!

shumway Dec 4, 2025

Uh oh!

shumway left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		@@ -0,0 +1,244 @@
		# Composable Kernel Builder Design Documentation

[CK_BUILDER] Refactor convolution signature to provide data type/layout/elementwise op per tensor #3331

Are you sure you want to change the base?

[CK_BUILDER] Refactor convolution signature to provide data type/layout/elementwise op per tensor #3331

Conversation

vpietila-amd commented Dec 1, 2025

Proposed changes

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shumway Dec 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shumway left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

shumway Dec 4, 2025 •

edited

Loading