Cleanup and refactoring related to tile loading by SamiAario-AMD · Pull Request #3505 · ROCm/composable_kernel

SamiAario-AMD · 2026-01-02T09:46:20Z

Proposed changes

Cleanup and refactoring done while implementing mixed precision for fp16/bf16 x fp8

Key changes:

Renamed load_interleaved_pk_type.hpp to load_and_convert_tile.hpp and refactored the API to use consistent naming conventions
Updated load_tile_transpose functions to use output parameters instead of return values for consistency
Removed unused variable declarations and simplified type deduction logic
Define load_tile_with_elementwise to use tuple types explicitly for clarity

Checklist

Please put an x into the boxes that apply. You can also fill these out after creating the PR. If you're not sure, please don't hesitate to ask.

I have added tests relevant to the introduced functionality, and the unit tests are passing locally
I have added the test to REGRESSION_TESTS list defined at the top of CMakeLists.txt in tests/CMakeLists.txt, IF the test takes more than 30 seconds to run.
I have added inline documentation which enables the maintainers with understanding the motivation
I have removed the stale documentation which is no longer relevant after this pull request
(If this change is user-facing) I have added release notes which provide the end users with a brief summary of the improvement from this pull request
I have run clang-format on all changed files
Any dependent changes have been merged

Discussion

If this is a relatively large or complex change, feel free to start a discussion by explaining why you chose the solution you did and what alternatives you considered

SamiAario-AMD · 2026-01-02T15:55:44Z

include/ck_tile/ops/common/load_and_convert_tile.hpp

    else if constexpr(LoadTranspose)
    {
-        dst = load_tile_transpose(src);
+        load_tile_transpose(dst, src);


Copilot seems to agree with me that assignment cannot be optimized by the compiler in such a way that the creation of temporaries is avoided, when assigning to a complex type. It then ought to make sense to avoid assignment when possible, and pass the object to be assigned to as a reference instead.

@SamiAario-AMD Have you checked the assembly for this change? Does the register usage is different?

I checked the assembly, and unfortunately there is no difference between using dst = load_tile_transpose(src) and load_tile_transpose(dst, src): the generated assembly is identical for my toy example. Adding the latter does have the benefit of making the API more similar to that of load_tile however. This PR does not remove the existing load_tile_transpose but just adds the new overload.

Copilot

Pull request overview

This PR performs cleanup and refactoring in preparation for mixed precision fp16/bf16 x fp8 implementation. The changes focus on standardizing the tile loading infrastructure and improving code organization.

Key changes:

Renamed load_interleaved_pk_type.hpp to load_and_convert_tile.hpp and refactored the API to use consistent naming conventions
Updated load_tile_transpose functions to use output parameters instead of return values for consistency
Removed unused variable declarations and simplified type deduction logic

Reviewed changes

Copilot reviewed 41 out of 41 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
include/ck_tile/ops/topk_softmax.hpp	Updated include from `load_interleaved_pk_type.hpp` to `load_and_convert_tile.hpp`
include/ck_tile/ops/topk.hpp	Updated include from `load_interleaved_pk_type.hpp` to `load_and_convert_tile.hpp`
include/ck_tile/ops/softmax.hpp	Updated include from `load_interleaved_pk_type.hpp` to `load_and_convert_tile.hpp`
include/ck_tile/ops/smoothquant.hpp	Updated include from `load_interleaved_pk_type.hpp` to `load_and_convert_tile.hpp`
include/ck_tile/ops/rmsnorm2d.hpp	Updated include from `load_interleaved_pk_type.hpp` to `load_and_convert_tile.hpp`
include/ck_tile/ops/reduce.hpp	Updated include from `load_interleaved_pk_type.hpp` to `load_and_convert_tile.hpp`
include/ck_tile/ops/pooling.hpp	Updated include from `load_interleaved_pk_type.hpp` to `load_and_convert_tile.hpp`
include/ck_tile/ops/permute.hpp	Updated include from `load_interleaved_pk_type.hpp` to `load_and_convert_tile.hpp`
include/ck_tile/ops/norm_reduce.hpp	Updated include from `load_interleaved_pk_type.hpp` to `load_and_convert_tile.hpp`
include/ck_tile/ops/layernorm2d.hpp	Updated include from `load_interleaved_pk_type.hpp` to `load_and_convert_tile.hpp`
include/ck_tile/ops/image_to_column.hpp	Updated include from `load_interleaved_pk_type.hpp` to `load_and_convert_tile.hpp`
include/ck_tile/ops/grouped_convolution.hpp	Updated include from `load_interleaved_pk_type.hpp` to `load_and_convert_tile.hpp`
include/ck_tile/ops/gemm_quant/pipeline/gemm_wp_bquant_pipeline_ag_bg_cr_v2.hpp	Replaced `load_int4_tile` calls with simplified `load_and_convert_tile` API
include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp	Simplified tile loading by removing explicit type parameters
include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp	Simplified tile loading by removing explicit type parameters
include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_mem.hpp	Simplified tile loading by removing explicit type parameters
include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp	Updated to use `load_and_convert_tile` API
include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_aquant_bs_cr.hpp	Updated to use `load_and_convert_tile` API
include/ck_tile/ops/gemm_quant.hpp	Updated include from `load_interleaved_pk_type.hpp` to `load_and_convert_tile.hpp`
include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_v2.hpp	Replaced `load_int4_tile` calls with `load_and_convert_tile` maintaining explicit template parameters
include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp	Moved `KPack` variable declaration to narrower scope where it's actually used
include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_base.hpp	Updated to use output parameter for `load_tile_transpose` and simplified type deduction
include/ck_tile/ops/gemm/block/block_universal_gemm_as_bs_cr.hpp	Updated to use `load_and_convert_tile` API and changed tile distribution to use `decltype`
include/ck_tile/ops/gemm.hpp	Updated include from `load_interleaved_pk_type.hpp` to `load_and_convert_tile.hpp`
include/ck_tile/ops/fused_moe.hpp	Updated include from `load_interleaved_pk_type.hpp` to `load_and_convert_tile.hpp`
include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async_trload.hpp	Changed `load_tile_transpose` to use output parameter instead of assignment
include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_v3_pipeline.hpp	Changed `load_tile_transpose` to use output parameter instead of assignment
include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_trload_qr_qtr_dor.hpp	Changed all `load_tile_transpose` calls to use output parameter instead of assignment
include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_trload_kr_ktr_vr.hpp	Changed all `load_tile_transpose` calls to use output parameter instead of assignment
include/ck_tile/ops/fmha.hpp	Updated include from `load_interleaved_pk_type.hpp` to `load_and_convert_tile.hpp`
include/ck_tile/ops/flatmm.hpp	Updated include from `load_interleaved_pk_type.hpp` to `load_and_convert_tile.hpp`
include/ck_tile/ops/epilogue.hpp	Updated include from `load_interleaved_pk_type.hpp` to `load_and_convert_tile.hpp`
include/ck_tile/ops/elementwise.hpp	Updated include from `load_interleaved_pk_type.hpp` to `load_and_convert_tile.hpp`
include/ck_tile/ops/common/load_and_convert_tile.hpp	Refactored and renamed from `load_interleaved_pk_type.hpp` with simplified API
include/ck_tile/ops/common.hpp	Updated include from `load_interleaved_pk_type.hpp` to `load_and_convert_tile.hpp`
include/ck_tile/ops/batched_transpose.hpp	Updated include from `load_interleaved_pk_type.hpp` to `load_and_convert_tile.hpp`
include/ck_tile/ops/batched_contraction.hpp	Updated include from `load_interleaved_pk_type.hpp` to `load_and_convert_tile.hpp`
include/ck_tile/ops/add_rmsnorm2d_rdquant.hpp	Updated include from `load_interleaved_pk_type.hpp` to `load_and_convert_tile.hpp`
include/ck_tile/core/tensor/tile_window.hpp	Refactored `load` function signature to accept tuple of tile windows
include/ck_tile/core/tensor/load_tile_transpose.hpp	Changed functions to use output parameters and updated documentation
include/ck_tile/core/tensor/load_tile.hpp	Changed return types from `auto` to `void` for functions that populate output parameters

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_v2.hpp

…DE integration

aosewski · 2026-01-09T13:54:50Z

include/ck_tile/ops/common/load_and_convert_tile.hpp

        static_assert(WarpTile::get_thread_buffer_size() % UnaryOpSize == 0);
        constexpr index_t thread_buffer_size = WarpTile::get_thread_buffer_size() / UnaryOpSize;
-        const auto in_dstr_tensors           = load_tile(warp_window);
+        const auto tmp                       = load_tile(src);


tmp is not the best name. What about naming it src and the window src_window ?

I will do this in a separate PR as discussed.

aosewski

It's not alwyas possible to just use data type from the tensor views. Please see comments.

aosewski · 2026-01-09T14:18:27Z

include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_base.hpp

-        using BLdsDataType =
-            std::conditional_t<std::is_same_v<typename Problem::BDataType, pk_fp4_raw_t>,
-                               typename Problem::ADataType,
-                               typename Problem::BDataType>;
-
        auto b_lds_load_tile_distr = []() {
            if constexpr(is_b_load_tr)
+            {
                return make_static_tile_distribution(
-                    typename InputTileDistributionTraits<typename BLdsLoadTileDistr::DstrEncode,
-                                                         BLdsDataType>::TransposedDstrEncode{});
-
+                    typename InputTileDistributionTraits<
+                        typename BLdsLoadTileDistr::DstrEncode,
+                        typename BLdsTensorView::DataType>::TransposedDstrEncode{});
+            }


You can't just use BLdsTensorView::DataType here, because of above conditional type selection. See here:

composable_kernel/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_base.hpp

Lines 116 to 135 in e62c96f

CK_TILE_DEVICE auto GetABLdsTensorViews(void* p_smem) const

{

// A tile in LDS

OverrideADataType* __restrict__ p_a_lds = static_cast<OverrideADataType*>(p_smem);

constexpr auto a_lds_block_desc =

Policy::template MakeALdsBlockDescriptor<Problem, OverrideADataType>();

auto a_lds_block = make_tensor_view<address_space_enum::lds>(p_a_lds, a_lds_block_desc);

// TODO: LDS alignment should come from Policy!

constexpr index_t a_lds_block_space_size_aligned = integer_least_multiple(

sizeof(OverrideADataType) * a_lds_block_desc.get_element_space_size(), 16);

// B tile in LDS

OverrideBDataType* __restrict__ p_b_lds = static_cast<OverrideBDataType*>(

static_cast<void*>(static_cast<char*>(p_smem) + a_lds_block_space_size_aligned));

constexpr auto b_lds_block_desc = Policy::template MakeBLdsBlockDescriptor<Problem>();

auto b_lds_block = make_tensor_view<address_space_enum::lds>(p_b_lds, b_lds_block_desc);

return make_tuple(std::move(a_lds_block), std::move(b_lds_block));

}

that the B Lds tensor view normally would get Problem::BDataType then you get different logic here.

OverrideBDataType in GetABLdsTensorViews should be the same as BLdsDataType in MakeBLdsWindows. The idea is that once we have determined OverrideBDataType and used it to define the B LDS tensor view, then we can reuse the B LDS tensor view's data type later in MakeBLdsWindows instead of re-determining it with the conditional type selection that I removed.

Ok - in general I like your idea to use data type from tensor views. So if you'd make sure that in all places this change is properly handled it's fine.
I think this should be mentioned in changelog, along to other load tile API changes.

Two changelog entries added.

include/ck_tile/ops/common/load_and_convert_tile.hpp

ammallya · 2026-02-03T22:02:09Z

Imported to ROCm/rocm-libraries

SamiAario-AMD requested review from ThomasNing, afagaj, andriy-ca, aosewski, asleepzzz, bartekxk, carlushuang, cgmillette, coderfeli, geyyer, illsilin, poyenc, qianfengz, shumway, tenpercent and vidyasagar-amd as code owners January 2, 2026 09:46

SamiAario-AMD force-pushed the LWPCK-3549-cleanups branch from 47bbab6 to 0308965 Compare January 2, 2026 15:50

SamiAario-AMD commented Jan 2, 2026

View reviewed changes

tenpercent requested a review from Copilot January 2, 2026 18:55

Copilot AI reviewed Jan 2, 2026

View reviewed changes

include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_v2.hpp Outdated Show resolved Hide resolved

SamiAario-AMD force-pushed the LWPCK-3549-cleanups branch 2 times, most recently from c811e2a to 211b784 Compare January 7, 2026 16:20

SamiAario-AMD added 8 commits January 7, 2026 16:21

Make some functions return void explicitly instead of auto

4d77856

Use decltype for consistency in Interwave variant of BlockGemmImpl

9691569

Add braces

bda5a7a

Fix a comment

825d17c

Reduce the scope of KPack in MakeALdsBlockDescriptor

ca71cd7

Minor refactoring of load_interleaved_pk_type

994b8f4

Rename load_interleaved_pk_type to load_and_convert_tile

74533b4

Include ck_tile/core.hpp in load_interleaved_pk_type.hpp for better I…

3a094e2

…DE integration

aosewski reviewed Jan 9, 2026

View reviewed changes

aosewski requested changes Jan 9, 2026

View reviewed changes

SamiAario-AMD changed the title ~~Lwpck 3549 cleanups~~ Cleanup and refactoring related to tile loading Jan 13, 2026

joyeamd reviewed Jan 13, 2026

View reviewed changes

include/ck_tile/ops/common/load_and_convert_tile.hpp Show resolved Hide resolved

SamiAario-AMD added 4 commits January 14, 2026 10:59

Merge branch 'develop' into LWPCK-3549-cleanups

ea4e543

Fix a build break introduced when merging

c020a42

Merge branch 'develop' into LWPCK-3549-cleanups

35c620e

Merge branch 'develop' into LWPCK-3549-cleanups

2ab79eb

SamiAario-AMD requested review from Snektron and vpietila-amd as code owners January 16, 2026 07:47

SamiAario-AMD added 3 commits January 19, 2026 17:16

Merge branch 'develop' into LWPCK-3549-cleanups

d71fe5b

Merge branch 'develop' into LWPCK-3549-cleanups

4b26eac

Merge branch 'develop' into LWPCK-3549-cleanups

72fa29b

illsilin assigned aosewski Jan 27, 2026

SamiAario-AMD mentioned this pull request Jan 28, 2026

Implement fp16 x fp8, bf16 x fp8, fp8 x fp16, and fp8 x bf16 for the V3 pipeline #3669

Closed

7 tasks

SamiAario-AMD added 2 commits January 28, 2026 17:14

Merge branch 'develop' into LWPCK-3549-cleanups

d0e9dc5

Fix a build break

fc1b683

SamiAario-AMD force-pushed the LWPCK-3549-cleanups branch from 87bb39b to fc1b683 Compare January 28, 2026 15:44

SamiAario-AMD added 6 commits February 2, 2026 10:00

Merge remote-tracking branch 'origin/develop' into LWPCK-3549-cleanups

348b555

Fix a build break

70be645

Fix a build break

aa247e2

Fix a build break

4eceb2f

Merge branch 'develop' into LWPCK-3549-cleanups

8c87159

Added changelog entries

7a62c52

SamiAario-AMD requested review from a team and ddembeckAMD as code owners February 3, 2026 17:34

This was referenced Feb 3, 2026

Implement fp16 x fp8, bf16 x fp8, fp8 x fp16, and fp8 x bf16 for the V3 pipeline ROCm/rocm-libraries#4276

Draft

Cleanup and refactoring related to tile loading ROCm/rocm-libraries#4294

Open

ammallya closed this Feb 3, 2026

	CK_TILE_DEVICE auto GetABLdsTensorViews(void* p_smem) const
	{
	// A tile in LDS
	OverrideADataType* __restrict__ p_a_lds = static_cast<OverrideADataType*>(p_smem);
	constexpr auto a_lds_block_desc =
	Policy::template MakeALdsBlockDescriptor<Problem, OverrideADataType>();
	auto a_lds_block = make_tensor_view<address_space_enum::lds>(p_a_lds, a_lds_block_desc);

	// TODO: LDS alignment should come from Policy!
	constexpr index_t a_lds_block_space_size_aligned = integer_least_multiple(
	sizeof(OverrideADataType) * a_lds_block_desc.get_element_space_size(), 16);

	// B tile in LDS
	OverrideBDataType* __restrict__ p_b_lds = static_cast<OverrideBDataType*>(
	static_cast<void>(static_cast<char>(p_smem) + a_lds_block_space_size_aligned));
	constexpr auto b_lds_block_desc = Policy::template MakeBLdsBlockDescriptor<Problem>();
	auto b_lds_block = make_tensor_view<address_space_enum::lds>(p_b_lds, b_lds_block_desc);

	return make_tuple(std::move(a_lds_block), std::move(b_lds_block));
	}

Conversation

SamiAario-AMD commented Jan 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Proposed changes

Checklist

Discussion

Uh oh!

SamiAario-AMD Jan 2, 2026

Choose a reason for hiding this comment

Uh oh!

aosewski Jan 9, 2026

Choose a reason for hiding this comment

Uh oh!

SamiAario-AMD Jan 19, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

aosewski Jan 9, 2026

Choose a reason for hiding this comment

Uh oh!

SamiAario-AMD Jan 12, 2026

Choose a reason for hiding this comment

Uh oh!

aosewski left a comment • 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!

aosewski Jan 9, 2026

Choose a reason for hiding this comment

Uh oh!

SamiAario-AMD Jan 12, 2026 • 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!

aosewski Jan 12, 2026 • 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!

SamiAario-AMD Feb 3, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ammallya commented Feb 3, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

SamiAario-AMD commented Jan 2, 2026 •

edited

Loading

aosewski left a comment •

edited

Loading

SamiAario-AMD Jan 12, 2026 •

edited

Loading

aosewski Jan 12, 2026 •

edited

Loading