xe: jit: gemm: loosen k-parallel restriction for attrs

kealan-barbieri · kealan-barbieri · commit 1a5a454270d9 · 2025-08-29T17:10:11.000-07:00
diff --git a/src/gpu/intel/jit/gemm/gen_gemm.hpp b/src/gpu/intel/jit/gemm/gen_gemm.hpp
@@ -325,8 +325,8 @@ struct gen_gemm_t : public gpu_gemm_t {
             //   accumulation unless fusion is enabled.
             if (kernel_desc_.driver_info()->kParallel()
                     && !kernel_desc_.driver_info()->fusedPostOps()) {
-                VDISPATCH_GEMM(!with_eltwise && !with_binary
-                                && utils::one_of(d->c_type(), f32, s32),
+                VDISPATCH_GEMM(
+                        !non_scale_po_ && utils::one_of(d->c_type(), f32, s32),
                         VERBOSE_UNSUPPORTED_POSTOP);
             }
 
diff --git a/src/gpu/intel/jit/gemm/jit_gemm_pd.cpp b/src/gpu/intel/jit/gemm/jit_gemm_pd.cpp
@@ -114,6 +114,7 @@ status_t jit_gemm_pd_t::init_post_ops() {
                 break;
             default: return status::unimplemented;
         }
+        non_scale_po_ = true;
     }
 
     if (!ok) return status::unimplemented;
@@ -134,6 +135,7 @@ status_t jit_gemm_pd_t::init_post_ops() {
         binary_srcs_.insert(
                 binary_srcs_.begin(), binary_src_t {binary_src_t::bias, 0});
     }
+    non_scale_po_ |= bias_via_binary_;
 
     auto maybe_convert_scales_to_postop
             = [this](const dims_t &scales_dims, int arg, data_type_t dt,
diff --git a/src/gpu/intel/jit/gemm/jit_gemm_pd.hpp b/src/gpu/intel/jit/gemm/jit_gemm_pd.hpp
@@ -82,6 +82,7 @@ struct jit_gemm_pd_t : public gpu_gemm_pd_t {
     int wei_q2d_group_k_ = 0;
     int src_q2d_group_k_ = 0;
     bool src_po_sc_ = false;
+    bool non_scale_po_ = false;
     data_type_t wei_scales_type_ = data_type::undef;
     data_type_t src_scales_type_ = data_type::undef;
 
