11//******************************************************************************
2- // MUL_MAT Kernel
2+ // MUL_MAT Kernel (F16 weights)
33// Matrix multiplication: C[M,N] = A[M,K] * B[K,N]
4+ //
5+ // Decode-optimized distribution: one output element (row m of A dotted with
6+ // column n of B) is the unit of work, striped across all harts of every active
7+ // shire. The previous version blocked work into chunks of 16 output elements,
8+ // so a GEMV (M=4096, N=1) only ever filled the first 8 shires, leaving 24 idle
9+ // and off-chip bandwidth stuck near 1/4 of peak. Striping by thread_id spreads
10+ // rows across all 32 shires and the register-resident f16_dot row helper
11+ // keeps the accumulator in registers for the whole K dimension.
412//******************************************************************************
513
614#include "block_ops.h"
715#include "ggml_tensor.h"
816#include "math_fp.h"
917#include "platform.h"
1018#include "quants.h"
19+ #include "tensor.h"
1120
1221#include <stdint.h>
1322
@@ -18,18 +27,13 @@ int entry_point(struct ggml_et_binary_params * params, void * env) {
1827 return -1 ;
1928 }
2029
21- // Thread coordination
30+ // Thread coordination, use every hart of every active shire.
2231 int thread_id = get_relative_thread_id (kernel_env -> shire_mask );
2332 int num_threads = get_num_threads (kernel_env -> shire_mask );
24-
25- if (thread_id < 0 || (thread_id & 1 )) {
26- return 0 ; // Skip odd threads to avoid resource contention
33+ if (thread_id < 0 ) {
34+ return 0 ;
2735 }
2836
29- int effective_thread_id = thread_id / 2 ;
30- int effective_num_threads = (num_threads + 1 ) / 2 ;
31-
32- // Extract tensor references
3337 struct ggml_tensor * src0 = & params -> src0 ; // Weight matrix A (F16)
3438 struct ggml_tensor * src1 = & params -> src1 ; // Activation matrix B (F16/F32)
3539 struct ggml_tensor * dst = & params -> dst ; // Output matrix C (F32)
@@ -43,7 +47,7 @@ int entry_point(struct ggml_et_binary_params * params, void * env) {
4347 const uint16_t * src0_data = (const uint16_t * ) src0 -> data ;
4448 float * dst_data = (float * ) dst -> data ;
4549
46- // Dimensions and Strides
50+ // Dimensions and strides
4751 const int64_t K = src0 -> ne [0 ];
4852 const int64_t M = src0 -> ne [1 ];
4953 const int64_t N = src1 -> ne [1 ];
@@ -56,87 +60,101 @@ int entry_point(struct ggml_et_binary_params * params, void * env) {
5660 const size_t nb11 = src1 -> nb [1 ], nb12 = src1 -> nb [2 ], nb13 = src1 -> nb [3 ];
5761 const size_t nb1 = dst -> nb [1 ], nb2 = dst -> nb [2 ], nb3 = dst -> nb [3 ];
5862
59- // F16 specific block size (Usually QK_F16)
60- const int block_size = QK_F16 ;
63+ const int block_size = QK_F16 ; // 32
6164 const int64_t K_blocks = K / block_size ;
6265 const int64_t K_remainder = K % block_size ;
6366
64- // Threading distribution
65- const uint64_t total_elements = M * N * ne2 * ne3 ;
66- const uint64_t per_thread = 16 ;
67- const uint64_t threads_stride = per_thread * effective_num_threads ;
68-
69- if (effective_thread_id * per_thread >= total_elements ) {
70- return 0 ;
71- }
72-
7367 // Broadcasting support
7468 const int64_t r2 = ne12 / ne02 ;
7569 const int64_t r3 = ne13 / ne03 ;
7670
77- for (uint64_t base_idx = effective_thread_id * per_thread ; base_idx < total_elements ; base_idx += threads_stride ) {
78- for (uint64_t j = 0 ; j < per_thread ; j ++ ) {
79- const uint64_t idx = base_idx + j ;
80- if (idx >= total_elements ) {
81- break ;
71+ const int is_f32_b = (src1 -> type == GGML_TYPE_F32 );
72+ const uint64_t total_elements = (uint64_t ) M * N * ne2 * ne3 ;
73+
74+ // Stage the reused B activation vector into per-shire L2 SCP. Streaming the
75+ // weight matrix thrashes L2 and evicts B, so B is re-read from DRAM
76+ // repeatedly during the decode GEMV. Stage it once per shire into L2 SCP (a
77+ // separate SRAM partition cache streaming cannot evict) via
78+ // et_tensor_load_l2scp; one hart per shire issues the DMA loop, then a
79+ // shire barrier lets all harts read B on-chip. Only applies to the common
80+ // decode case: F32 B, a single contiguous non-broadcast B that fits the SCP
81+ // budget.
82+ const int b_contig = (nb11 == (size_t ) K * sizeof (float ));
83+ const uint64_t b_lines = ((uint64_t ) N * K * sizeof (float ) + 63 ) / 64 ;
84+ const int stage_b = is_f32_b && b_contig &&
85+ ne12 == 1 && ne13 == 1 && ne02 == 1 && ne03 == 1 &&
86+ b_lines <= 8192 ; // <= 512 KB, within the SCP budget
87+ const float * b_scp = (const float * ) et_shire_l2scp_local (0 );
88+
89+ if (stage_b ) {
90+ if ((get_hart_id () & 63 ) == 0 ) {
91+ et_tensor_load_l2scp_conf_t conf ;
92+ conf .use_tmask = false;
93+ conf .stride = 64 ; // advance one 64B cache line per line loaded
94+ uint64_t remaining = b_lines ;
95+ uint64_t dst = 0 ;
96+ uint64_t addr = (uint64_t ) src1 -> data ;
97+ while (remaining > 0 ) {
98+ uint64_t cl = (remaining >= 16 ) ? 16 : remaining ;
99+ conf .dst_start = dst ;
100+ conf .addr = addr ;
101+ conf .num_lines = cl - 1 ; // 4-bit field encodes (lines - 1)
102+ conf .id = 0 ;
103+ et_tensor_load_l2scp (& conf );
104+ WAIT_TENSOR_LOAD_L2_0 ;
105+ dst += cl ;
106+ addr += cl * 64 ;
107+ remaining -= cl ;
82108 }
109+ }
110+ et_barrier (ET_BARRIER_SHIRE ); // B is now resident in L2 SCP for all harts
111+ }
83112
84- // Index decoding
85- const int64_t i3 = idx / (M * N * ne2 );
86- const int64_t rem3 = idx % (M * N * ne2 );
87- const int64_t i2 = rem3 / (M * N );
88- const int64_t rem2 = rem3 % (M * N );
89- const int64_t n = rem2 / M ;
90- const int64_t m = rem2 % M ;
91-
92- const int64_t i03 = i3 / r3 , i02 = i2 / r2 ;
93- const int64_t i13 = (ne13 > 1 ) ? i3 : 0 , i12 = (ne12 > 1 ) ? i2 : 0 ;
94-
95- float sum = 0.0f ;
96- const uint16_t * f16_row =
97- (const uint16_t * ) ((const char * ) src0_data + m * nb01 + i02 * nb02 + i03 * nb03 );
98-
99- if (src1 -> type == GGML_TYPE_F32 ) {
100- const float * src1_data = (const float * ) src1 -> data ;
101-
102- for (int64_t kb = 0 ; kb < K_blocks ; kb ++ ) {
103- const float * b_col_ptr =
104- (const float * ) ((const char * ) src1_data + (kb * block_size ) * sizeof (float ) + n * nb11 +
105- i12 * nb12 + i13 * nb13 );
106- sum += compute_block_dot_product_f16_naive (& f16_row [kb * block_size ], b_col_ptr );
107- }
108-
109- if (K_remainder > 0 ) {
110- const int64_t offset = K_blocks * block_size ;
111- const float * b_col_ptr = (const float * ) ((const char * ) src1_data + offset * sizeof (float ) +
112- n * nb11 + i12 * nb12 + i13 * nb13 );
113- sum += compute_block_dot_product_f16_partial (& f16_row [offset ], b_col_ptr , K_remainder );
114- }
115- } else {
116- const uint16_t * src1_data = (const uint16_t * ) src1 -> data ;
117-
118- for (int64_t kb = 0 ; kb < K_blocks ; kb ++ ) {
119- const uint16_t * b_col_ptr =
120- (const uint16_t * ) ((const char * ) src1_data + (kb * block_size ) * sizeof (uint16_t ) + n * nb11 +
121- i12 * nb12 + i13 * nb13 );
122- sum += compute_block_dot_product_f16_f16_partial (& f16_row [kb * block_size ], b_col_ptr , block_size );
123- }
124-
125- if (K_remainder > 0 ) {
126- const int64_t offset = K_blocks * block_size ;
127- const uint16_t * b_col_ptr =
128- (const uint16_t * ) ((const char * ) src1_data + offset * sizeof (uint16_t ) + n * nb11 +
129- i12 * nb12 + i13 * nb13 );
130- sum += compute_block_dot_product_f16_f16_partial (& f16_row [offset ], b_col_ptr , K_remainder );
131- }
113+ // Set the vector mask (all 8 lanes) once for the whole row loop.
114+ unsigned long saved_mask ;
115+ __asm__ volatile ("mova.x.m %0" : "=r" (saved_mask ));
116+ __asm__ volatile ("mov.m.x m0, x0, 0xFF" );
117+
118+ for (uint64_t idx = (uint64_t ) thread_id ; idx < total_elements ; idx += (uint64_t ) num_threads ) {
119+ // Index decoding
120+ const int64_t i3 = idx / (M * N * ne2 );
121+ const int64_t rem3 = idx % (M * N * ne2 );
122+ const int64_t i2 = rem3 / (M * N );
123+ const int64_t rem2 = rem3 % (M * N );
124+ const int64_t n = rem2 / M ;
125+ const int64_t m = rem2 % M ;
126+
127+ const int64_t i03 = i3 / r3 , i02 = i2 / r2 ;
128+ const int64_t i13 = (ne13 > 1 ) ? i3 : 0 , i12 = (ne12 > 1 ) ? i2 : 0 ;
129+
130+ const uint16_t * f16_row =
131+ (const uint16_t * ) ((const char * ) src0_data + m * nb01 + i02 * nb02 + i03 * nb03 );
132+
133+ float sum ;
134+ if (is_f32_b ) {
135+ const float * b_col = stage_b
136+ ? (b_scp + n * K )
137+ : (const float * ) ((const char * ) src1 -> data + n * nb11 + i12 * nb12 + i13 * nb13 );
138+
139+ f16_dot_reset ();
140+ f16_dot_tile (f16_row , b_col , K_blocks );
141+ sum = f16_dot_reduce ();
142+
143+ if (K_remainder > 0 ) {
144+ const int64_t offset = K_blocks * block_size ;
145+ sum += compute_block_dot_product_f16_partial (& f16_row [offset ], & b_col [offset ], K_remainder );
132146 }
133-
134- // Atomic store for output
135- volatile float * c_element =
136- (volatile float * ) ((char * ) dst_data + m * dst -> nb [0 ] + n * nb1 + i2 * nb2 + i3 * nb3 );
137- atomic_store_f32 (c_element , sum );
147+ } else {
148+ const uint16_t * b_col =
149+ (const uint16_t * ) ((const char * ) src1 -> data + n * nb11 + i12 * nb12 + i13 * nb13 );
150+ sum = compute_block_dot_product_f16_f16_partial (f16_row , b_col , K );
138151 }
152+
153+ volatile float * c_element =
154+ (volatile float * ) ((char * ) dst_data + m * dst -> nb [0 ] + n * nb1 + i2 * nb2 + i3 * nb3 );
155+ atomic_store_f32 (c_element , sum );
139156 }
140157
158+ __asm__ volatile ("mova.m.x %0" ::"r" (saved_mask ));
141159 return 0 ;
142160}
0 commit comments