From 6a7d07fcdcf83dce2a8335cc14440717c9e7d469 Mon Sep 17 00:00:00 2001
From: Bartlomiej Wroblewski <bwroblewski10@gmail.com>
Date: Mon, 27 Nov 2023 11:31:39 +0100
Subject: [PATCH] Add missing check for K padding in XDL GEMM (#1056)

[ROCm/composable_kernel commit: 60ecfd73f9d006e02ef2a33820115222f971bd98]
---
 .../gpu/grid/gridwise_gemm_xdlops_v2r3.hpp            | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdlops_v2r3.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdlops_v2r3.hpp
index 2c42ac9345..0b50601648 100644
--- a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdlops_v2r3.hpp
+++ b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdlops_v2r3.hpp
@@ -996,6 +996,17 @@ struct GridwiseGemm_k0mk1_k0nk1_mn_xdlops_v2r3_ext
             }
         }
 
+        if constexpr(!(GemmSpec == tensor_operation::device::GemmSpecialization::KPadding ||
+                       GemmSpec == tensor_operation::device::GemmSpecialization::MKPadding ||
+                       GemmSpec == tensor_operation::device::GemmSpecialization::NKPadding ||
+                       GemmSpec == tensor_operation::device::GemmSpecialization::MNKPadding))
+        {
+            if(!(problem.K0 % K0PerBlock == 0))
+            {
+                return false;
+            }
+        }
+
         if constexpr(is_same<tensor_layout::gemm::RowMajor, ALayout>::value)
         {
             if(problem.K % ABlockTransferSrcScalarPerVector != 0)