Add missing check for K padding in XDL GEMM (#1056)

60ecfd73 · Bartlomiej Wroblewski · GitHub · bfecc193 · 60ecfd73
Unverified Commit 60ecfd73 authored Nov 27, 2023 by Bartlomiej Wroblewski Committed by GitHub Nov 27, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 11 additions and 0 deletions

include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdlops_v2r3.hpp ...k/tensor_operation/gpu/grid/gridwise_gemm_xdlops_v2r3.hpp +11 -0

No files found.
--- 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)