Fix Nan issue in convolution asm Kernel with NCHW layout (#3403)

* fix clear buffer for NCHW layout * Clean up commented code --------- Co-authored-by: xinlipn <[email protected]>
ROCm · Dec 2, 2024 · 531f2c8 · 531f2c8
1 parent 377eaeb
commit 531f2c8
Showing 1 changed file with 8 additions and 15 deletions.
diff --git a/src/conv/invokers/impl_gemm_dynamic.cpp b/src/conv/invokers/impl_gemm_dynamic.cpp
@@ -178,9 +178,7 @@ InvokerFactory MakeImplGemmDynamicBackwardDataInvokerFactory(const ProblemDescri
         const int gemm_k_gid = k * y_dot_slice_gid[gemm_id] * x_dot_slice_gid[gemm_id];
         is_gemm_not_empty.emplace_back(gemm_k_gid > 0);
     }
-    bool need_set_zero = false;
-    if(y < stride_h || x < stride_w || dilation_h != 1 || dilation_w != 1)
-        need_set_zero = true;
+    bool need_set_zero = true;
 
     return [=](const std::vector<Kernel>& kernels) {
         const auto kernel = kernels[0];
@@ -315,13 +313,10 @@ MakeImplGemmDynamicBackwardDataInvokerFactory(const ProblemDescription& problem,
         const int gemm_k_gid = k * y_dot_slice_gid[gemm_id] * x_dot_slice_gid[gemm_id];
         is_gemm_not_empty.emplace_back(gemm_k_gid > 0);
     }
-    bool need_set_zero = false;
-    if(y < stride_h || x < stride_w || dilation_h != 1 || dilation_w != 1)
-        need_set_zero = true;
-
-    int nxb = cfg.nxb;
-    int b   = h_tilda_slice * w_tilda_slice;
-    b       = (cfg.nxe == 0) ? (b) : ((b + nxb - 1) / nxb) * nxb; // pad to nxb modulo when nxe != 0
+    bool need_set_zero = true;
+    int nxb            = cfg.nxb;
+    int b              = h_tilda_slice * w_tilda_slice;
+    b = (cfg.nxe == 0) ? (b) : ((b + nxb - 1) / nxb) * nxb; // pad to nxb modulo when nxe != 0
 
     uint32_t nb_n0          = cfg.tensor_b_cluster_lengths[2] * cfg.tensor_b_thread_lengths[2];
     uint32_t nb_n1b         = cfg.tensor_b_cluster_lengths[3] * cfg.tensor_b_thread_lengths[3];
@@ -496,7 +491,8 @@ InvokerFactory MakeImplGemmDynamicForwardXdlopsNHWCInvokerFactory(
         shift_pack_1 = 0;
     }
 
-    bool need_set_zero                 = config.gemm_k_global_split > 0;
+    // Clear buffer for all condition to resolve the NaN issue.
+    bool need_set_zero                 = true;
     bool use_fp32_global_split_on_fp16 = config.vector_store == 1 && config.gemm_k_global_split > 0;
 
     std::vector<OpKernelArg> opArgs;
@@ -797,11 +793,8 @@ InvokerFactory MakeImplGemmDynamicBackwardDataXdlopsNHWCInvokerFactory(
     int dtile_h  = num_of_gemms > 1 ? static_cast<int>(mdiv_group_mn.magic) : h_tilda;
     int dtile_w  = num_of_gemms > 1 ? static_cast<int>(mdiv_group_mn.shift) : w_tilda;
 
-    bool need_set_zero                 = false;
+    bool need_set_zero                 = true;
     bool use_fp32_global_split_on_fp16 = config.vector_store == 1 && config.gemm_k_global_split > 0;
-    if(y < stride_h || x < stride_w || dilation_h != 1 || dilation_w != 1)
-        need_set_zero = true;
-    need_set_zero |= config.gemm_k_global_split > 0;
 
     std::vector<OpKernelArg> opArgs;
     opArgs.emplace_back(0); // placeholder