cpu: x64: matmul: restrict K-collapse to int8 and bound load concurrency

The K-collapse in set_blocking_parameters (32909e9) keeps C resident in
the AMX tiles across the whole K reduction, saving the per-chunk C read-
modify-write. Cold-weight A/B on GNR (int8, M=64) shows this saving is a net
win for most shapes but can regress at large reduction depth: collapsing
serializes the cold B-tile loads into one call, so the out-of-order engine
can no longer overlap that DRAM latency the way the chunked baseline does
with its separate brgemm calls.

The exposed cost scales with the aggregate concurrent memory demand, which
is the reduction depth times the number of threads streaming B at once
(nthr_ * k_tiles_full). A cold sweep over K and thread count shows the
win/regress boundary tracks this product: at 32 threads the crossover is
k_tiles_full == 96, and it moves proportionally with fewer threads (deep
reductions that regress at 32 threads win at 8-24). VTune confirms the
mechanism on the regressing case -- collapse eliminates the C read-modify-
write (Store-Bound 2.0%->0.2%) but raises Memory-Bound via lost memory-level
parallelism (DRAM 27%->36%, CPI 6.8->10.3); Front-End Bound stays <1%, so
the deeper unroll itself is not the cost.

Replace the fixed reduction-depth cap with a load-concurrency cap
(nthr_ * k_tiles_full <= 32 * 96) alongside the existing tmul-count cap,
which alone does not catch the regressors at M=64 (c_tiles == 8 keeps
tmuls == 896 < 1024 even at deep K; M>=128 has c_tiles >= 16 and is already
excluded). Also restrict the collapse to int8, where it was calibrated. The
new cap reduces exactly to the prior k_tiles_full <= 96 at 32 threads, so
behavior in that regime is unchanged, while recovering the deep-K wins at
lower thread counts. Validated cold across N and thread counts (wins hold at
all tested N for nthr <= 24, behavior identical to the prior guard at 32
threads); all affected shapes pass benchdnn mode=C.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: asimonov1

src/cpu/x64/matmul/amx_blocking_heuristics.cpp

@@ -1191,27 +1191,38 @@ bool matmul_amx_blocking_params_macro_t::set_blocking_parameters(
         use_fused_copy_a_ = false;
 
         // Collapse to a single K-chunk only when one thread owns the whole K
-        // reduction (nthr_k_ == 1) and the heuristic split K (best_k_v > 1).
-        if (nthr_k_ == 1 && best_k_v > 1) {
+        // reduction (nthr_k_ == 1) and the heuristic split K into more than one
+        // chunk (best_k_v > 1). Restricted to int8 (GNR-calibrated). The
+        // postops-bound branch above leaves best_k_v == 1, so best_k_v > 1 also
+        // keeps these two paths mutually exclusive.
+        if (nthr_k_ == 1 && best_k_v > 1 && one_of(src_dt, s8, u8)) {
             // Keep C resident in the AMX tiles across the whole K reduction
             // instead of chunking K, which forces a C read-modify-write
             // (tilestored/tileloadd of every C tile) at each chunk boundary
             // that contends with the tmul and B-load ops on the AMX port. Here
             // C fits in tiles and B is streamed once, so chunking only bounds
             // the unrolled kernel size -- collapse to one K-chunk while the
-            // full-K working set fits L2 and the unrolled tmul count stays
-            // capped.
+            // caps below hold.
             const dim_t k_tiles_full = k_per_thread / wei_k_blk;
+            // c_tiles = 16x16 output sub-tiles in the m_blk_ x n_blk_ block (the
+            // output work grid), not AMX register occupancy.
             const dim_t c_tiles = div_up(m_blk_, 16) * div_up(n_blk_, 16);
             const size_t full_k_working_set
-                    = (size_t)(m_blk_ + n_blk_) * k_per_thread * gemm_dt_sz;
-            // k_tiles_full * c_tiles = unrolled tmul count, a proxy for code
-            // size (~50 B/op). The collapse saves one C RMW regardless of K,
-            // but this footprint grows with K; past the cap I-cache pressure
-            // outweighs the win. Calibrated on GNR (int8 M=64/K=7168): 896 ops
-            // wins ~2x, >= 1344 neutral. Re-tune per-uarch vs L1i.
+                    = (m_blk_ + n_blk_) * k_per_thread * gemm_dt_sz;
+            // k_tiles_full * c_tiles is the total tmul count, a proxy for the
+            // unrolled kernel size; past this cap the collapse stops paying off.
             constexpr dim_t max_unrolled_tmuls = 1024;
+            // Collapsing removes the per-chunk C read-modify-write (fixed
+            // saving) but serializes the cold B-tile loads in one call, so the
+            // out-of-order engine can no longer overlap that DRAM latency. The
+            // exposed cost scales with aggregate concurrent demand ~
+            // nthr_ * k_tiles_full (reduction depth x threads streaming B at
+            // once); collapse wins while that product stays under the value
+            // below (measured cold on GNR, int8). At 32 threads this is
+            // k_tiles_full <= 96; fewer threads tolerate deeper reductions.
+            constexpr dim_t max_load_concurrency = 32 * 96;
             if (k_tiles_full * c_tiles <= max_unrolled_tmuls
+                    && nthr_ * k_tiles_full <= max_load_concurrency
                     && full_k_working_set <= L2_threshold()) {
                 k_blk_v = k_per_thread;
                 best_k_v = 1;