cpu: x64: brgemm: fix zp_comp_b when bd_block >= 14 on avx2+

oneapi-src · Jun 30, 2023 · ec0b2ee · ec0b2ee
1 parent 6d2e567
commit ec0b2ee
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/src/cpu/x64/brgemm/jit_brgemm_kernel.cpp b/src/cpu/x64/brgemm/jit_brgemm_kernel.cpp
@@ -242,7 +242,12 @@ struct jit_brgemm_kernel_t : public jit_generator {
         }
     }
 
-    Vmm vmm_tmp(int i) { return Vmm(i); }
+    Vmm vmm_tmp(int i) {
+        assert(i >= 0
+                && i < max_effective_vregs - brg.bd_block * brg.ld_block2);
+        return Vmm(i);
+    }
+
     Vmm vmm_tail_mask() { return vmm_tmp(1); }
     Vmm vmm_one_bytes() const noexcept { return Vmm(3); }
     Vmm vmm_zp_a_shift() const noexcept { return Vmm(2); }
@@ -1222,7 +1227,7 @@ void jit_brgemm_kernel_t<isa, Wmm>::apply_compensation(
                             reg_aux_zp_comp_b, zp_comp_b_off, true);
                     uni_vpaddd(vmm, vmm, zp_comp_b_addr);
                 } else {
-                    const auto vmm_zp_comp_b = vmm_tmp(2);
+                    const auto vmm_zp_comp_b = vmm_tmp(0);
                     uni_vpbroadcastd(vmm_zp_comp_b,
                             ptr[reg_aux_zp_comp_b + zp_comp_b_off]);
                     uni_vpaddd(vmm, vmm, vmm_zp_comp_b);