xe: jit: gemm: expand decomp cases, enforce fpmath

Author: kealan-barbieri

src/gpu/intel/jit/gemm/gen_gemm.hpp

@@ -74,8 +74,8 @@ struct gen_gemm_t : public gpu_gemm_t {
             wei_decomp_ = (utils::one_of(d->c_type(), f32, f16, bf16, f8_e5m2,
                                    f8_e4m3)
                                   && utils::one_of(d->a_type(), u8, s8, s4, u4)
-                                  && utils::one_of(d->b_type(), f16, f32, bf16,
-                                          f8_e5m2, f8_e4m3))
+                                  && utils::one_of(d->b_type(), u8, s8, s4, u4,
+                                          f16, f32, bf16, f8_e5m2, f8_e4m3))
                     && attr()->mayiconvert(d->a_type(), f32);
             dy_quant_enabled_
                     = (utils::one_of(d->c_type(), f32, f16, bf16)
@@ -224,6 +224,9 @@ struct gen_gemm_t : public gpu_gemm_t {
 
             if (!attr()->zero_points_.has_default_values()) {
                 if (!attr_zps.has_default_values(DNNL_ARG_A)) {
+                    // Only apply to integers inputs.
+                    VDISPATCH_GEMM(utils::one_of(d->a_type(), s4, u4, s8, u8),
+                            VERBOSE_UNSUPPORTED_ZP_CFG);
                     const int cmask_a = attr_zps.get_mask(DNNL_ARG_A);
                     ao_dims_ = cmask_a > 0;
 
@@ -253,10 +256,17 @@ struct gen_gemm_t : public gpu_gemm_t {
                         VDISPATCH_GEMM(utils::one_of(cmask_a, 0, mask_per_oc,
                                                mask_per_ic),
                                 VERBOSE_UNSUPPORTED_ZP_CFG);
+                        // Weights zp can only be performantly enabled during upconversion.
+                        VDISPATCH_GEMM(wei_decomp_
+                                        || utils::one_of(d->b_type(), s4, u4),
+                                VERBOSE_UNSUPPORTED_ZP_CFG);
                     }
                 }
 
                 if (!attr_zps.has_default_values(DNNL_ARG_B)) {
+                    // Only apply to integers inputs.
+                    VDISPATCH_GEMM(utils::one_of(d->b_type(), s4, u4, s8, u8),
+                            VERBOSE_UNSUPPORTED_ZP_CFG);
                     const int cmask_b = attr_zps.get_mask(DNNL_ARG_B);
                     bo_dims_ = cmask_b > 0;
 
@@ -390,6 +400,7 @@ struct gen_gemm_t : public gpu_gemm_t {
                     : data_type::s32;
             if (swap_ab_) std::swap(ao_type, bo_type);
             bool int_acc = utils::one_of(eff_a_type(), s8, u8);
+            int_acc &= !wei_scales_2d_;
             auto co_type = with_bias() ? d->bias_type()
                     : with_sum_ab()    ? d->sum_ab_type
                     : int_acc          ? s32

src/gpu/intel/jit/gemm/gen_gemm_kernel.cpp

@@ -535,19 +535,13 @@ status_t gen_gemm_nocopy_kernel_desc_t::select_kernel(compute::gpu_arch_t arch,
 
     auto add_mode_matches = [&](bool has_mode, const char *(*match)(Type)) {
         if (!has_mode) return;
-        auto &def = base.selector.precisions;
         if (match(problem_.Ta)) {
-            match_params.push_back(base);
             match_params.back().selector.precisions[0] = match(problem_.Ta);
-            match_params.back().selector.precisions[1] = def[1];
         }
         if (match(problem_.Tb)) {
-            match_params.push_back(base);
-            match_params.back().selector.precisions[0] = def[0];
             match_params.back().selector.precisions[1] = match(problem_.Tb);
         }
         if (match(problem_.Ta) && match(problem_.Tb)) {
-            match_params.push_back(base);
             match_params.back().selector.precisions[0] = match(problem_.Ta);
             match_params.back().selector.precisions[1] = match(problem_.Tb);
         }