gpu: jit: conv: fix zero padding with sum

oneapi-src · Apr 9, 2022 · 2e4ad3a · 2e4ad3a
1 parent 27668dd
commit 2e4ad3a
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/src/gpu/jit/conv/post_op_support.hpp b/src/gpu/jit/conv/post_op_support.hpp
@@ -396,7 +396,12 @@ class post_op_context_t {
                 if (!eltwise_fwd_pd_t::eltwise_preserves_zero(po.eltwise))
                     return true;
             } else if (po.is_sum(/*require_scale_one=*/false)) {
-                // Preserves zero padding.
+                for (int j = 0; j < cp_ndims(); j++) {
+                    if (!is_cp_dim_zero_padded(j)) continue;
+                    // Size one dimensions are treated as broadcast which does
+                    // not preserve zero padding with block updates.
+                    if (cp_view_.vdims()[j] == 1) return true;
+                }
             } else if (po.is_binary()) {
                 for (int j = 0; j < cp_ndims(); j++) {
                     if (!is_cp_dim_zero_padded(j)) continue;