From df929f094ab0ec7551321e8f5b3d6ced7b4ad361 Mon Sep 17 00:00:00 2001
From: Anca Hamuraru <anca@streamhpc.com>
Date: Thu, 1 May 2025 12:51:21 +0000
Subject: [PATCH] Updates to support mixed precision

(cherry picked from commit e65d71180393e7b66169c56565a6bac740427de6)

Co-authored-by: Anca Hamuraru <anca@streamhpc.com>
---
 .../gpu/block/blockwise_gemm_pipeline_wmmaops_base.hpp        | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_base.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_base.hpp
index a63d32802e..0adf9e7241 100644
--- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_base.hpp
+++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_base.hpp
@@ -282,7 +282,7 @@ struct BlockwiseGemmWmmaops_pipeline_base
 
     using AThreadCopy =
         ThreadwiseTensorSliceTransfer_v4<ADataType,
-                                         ADataType,
+                                         ComputeTypeA,
                                          decltype(a_block_desc_k0_m0_m1_m2_k1),
                                          decltype(a_thread_desc_),
                                          Sequence<KPack / A_K1 / A_KRow, 1, 1, 1, 1, A_K1>,
@@ -293,7 +293,7 @@ struct BlockwiseGemmWmmaops_pipeline_base
 
     using BThreadCopy =
         ThreadwiseTensorSliceTransfer_v4<BDataType,
-                                         BDataType,
+                                         ComputeTypeB,
                                          decltype(b_block_desc_k0_n0_n1_n2_k1),
                                          decltype(b_thread_desc_),
                                          Sequence<KPack / B_K1 / B_KRow, 1, 1, 1, 1, B_K1>,