From d466096e25edd0ac0159f4fad2d3f0104697bba7 Mon Sep 17 00:00:00 2001
From: Rostyslav Geyyer <rosty.geyyer@amd.com>
Date: Mon, 10 Mar 2025 21:09:38 +0000
Subject: [PATCH] Update chunk size for f4x2

---
 test/mx_mfma_op/mx_mfma_op.hpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/test/mx_mfma_op/mx_mfma_op.hpp b/test/mx_mfma_op/mx_mfma_op.hpp
index d46f234f98..b68488f5ce 100644
--- a/test/mx_mfma_op/mx_mfma_op.hpp
+++ b/test/mx_mfma_op/mx_mfma_op.hpp
@@ -226,7 +226,8 @@ __device__ AFragT load_A_row_major(AType const* input_ptr)
     static constexpr int32_t WAVE_SIZE = 64;
 
     // Here we want to load from rows of A in chunks of 16 elements each.
-    static constexpr uint32_t chunk_size = 16;
+    static constexpr uint32_t chunk_size =
+        16 / (ck::is_same_v<ck::remove_cvref_t<AType>, ck::f4x2_pk_t> ? 2 : 1);
 
     // each chunk is separated by offset
     static constexpr uint32_t chunk_offset = chunk_size * WAVE_SIZE / BLOCK_M;