From 1998fcda26051dc63c81ffa05fcd561ffc165591 Mon Sep 17 00:00:00 2001
From: MHYang-gh <meng-hsuan.yang@amd.com>
Date: Sun, 23 Mar 2025 14:13:50 +0800
Subject: [PATCH] Fix A/B lds transform (#2007)

[ROCm/composable_kernel commit: c027637a8fc39cd3aefe9ce5314bb8636f801d77]
---
 .../pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp  | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp
index f5b3523f60..c504a51ad0 100644
--- a/include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp
+++ b/include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp
@@ -362,7 +362,7 @@ struct UniversalGemmPipelineAgBgCrPolicy
         constexpr auto a_lds_block_desc_xk0_mnldslayer_mn_xk1 = transform_tensor_descriptor(
             a_lds_block_desc_permuted,
             make_tuple(make_unmerge_transform(
-                           make_tuple(number<KPerBlock / KPack>{}, number<MLdsLayer>{})),
+                           make_tuple(number<MLdsLayer>{}, number<KPerBlock / KPack>{})),
                        make_pass_through_transform(number<MPerBlock / MLdsLayer>{}),
                        make_pass_through_transform(number<KPack>{})),
             make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}),
@@ -374,7 +374,7 @@ struct UniversalGemmPipelineAgBgCrPolicy
                            make_tuple(number<MPerBlock / MLdsLayer>{}, number<MLdsLayer>{})),
                        make_merge_transform_v3_division_mod(
                            make_tuple(number<KPerBlock / KPack>{}, number<KPack>{}))),
-            make_tuple(sequence<1, 2>{}, sequence<0, 3>{}),
+            make_tuple(sequence<1, 0>{}, sequence<2, 3>{}),
             make_tuple(sequence<0>{}, sequence<1>{}));
 
         return a_lds_block_desc;
@@ -421,7 +421,7 @@ struct UniversalGemmPipelineAgBgCrPolicy
 
             constexpr auto b_lds_block_desc_bk0_nldslayer_n_bk1 = transform_tensor_descriptor(
                 b_lds_block_desc_permuted,
-                make_tuple(make_unmerge_transform(make_tuple(BK0, number<NLdsLayer>{})),
+                make_tuple(make_unmerge_transform(make_tuple(number<NLdsLayer>{}, BK0)),
                            make_pass_through_transform(number<NPerBlock / NLdsLayer>{}),
                            make_pass_through_transform(number<KPack>{})),
                 make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}),
@@ -432,7 +432,7 @@ struct UniversalGemmPipelineAgBgCrPolicy
                 make_tuple(make_merge_transform_v3_division_mod(
                                make_tuple(number<NPerBlock / NLdsLayer>{}, number<NLdsLayer>{})),
                            make_merge_transform_v3_division_mod(make_tuple(BK0, number<KPack>{}))),
-                make_tuple(sequence<1, 2>{}, sequence<0, 3>{}),
+                make_tuple(sequence<1, 0>{}, sequence<2, 3>{}),
                 make_tuple(sequence<0>{}, sequence<1>{}));
             return b_lds_block_desc;
         }