Skip to content

Commit cbba22a

Browse files
RehanQasim-devRehanQasim-dev
committed
et-backend: F16 decode GEMV rewrite for register-resident vecdot
Stripe output elements across every hart of all 32 shires instead of blocking work into 16-element chunks, which only filled 8 shires for a typical decode GEMV. Adds a register-resident f16 row-dot helper and stages the reused B activation vector into per-shire L2 SCP so it survives weight-matrix streaming. Co-authored-by: Rehan Qasim <rehan.qasim@10xengineers.ai>
1 parent 7cbd610 commit cbba22a

2 files changed

Lines changed: 202 additions & 79 deletions

File tree

‎ggml/src/ggml-et/et-kernels/src/block_ops.h‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -670,6 +670,111 @@ static inline float compute_block_dot_product_f16(const uint16_t * a_block, cons
670670
return compute_block_dot_product_f16_partial(a_block, b_col_start, QK_F16);
671671
}
672672

673+
//******************************************************************************
674+
// Hoisted F16 row-dot API (register-resident accumulator)
675+
//
676+
// Mirrors the q8_dot_* API above: the running sum lives in vector registers
677+
// for the whole row instead of being spilled to memory every 8 elements, so
678+
// the inner loop is load + convert + fmadd with no accumulator round-trip.
679+
//
680+
// Register contract (shared convention with q8_dot_*):
681+
// f20-f23 row accumulators (persistent across tiles, reset per row)
682+
// f31 gather pattern (byte offsets of 8 consecutive f16)
683+
// f11-f18 scratch within tile
684+
// t0, f1-f5 scratch within reduce
685+
// Caller sets the vector mask to 0xFF once around the surrounding loop.
686+
//******************************************************************************
687+
688+
static inline void __attribute__((always_inline))
689+
f16_dot_reset(void) {
690+
// Four independent lane accumulators break the fmadd dependency chain so
691+
// multiple A/B loads stay in flight, hiding load latency.
692+
__asm__ volatile(
693+
"fbci.pi f20, 0\n"
694+
"fbci.pi f21, 0\n"
695+
"fbci.pi f22, 0\n"
696+
"fbci.pi f23, 0\n"
697+
::: "f20", "f21", "f22", "f23");
698+
}
699+
700+
// Accumulate n_blocks blocks of QK_F16 (=32) f16 A values times f32 B into
701+
// f20..f23 (one accumulator per 8-lane chunk within the block).
702+
static inline void __attribute__((always_inline))
703+
f16_dot_tile(const uint16_t * a_row, const float * b_col, int64_t n_blocks) {
704+
static const int32_t gather_pattern[8] = {0, 2, 4, 6, 8, 10, 12, 14};
705+
__asm__ volatile("flw.ps f31, %[g]\n"
706+
: : [g] "m"(*(const int32_t(*)[8])gather_pattern)
707+
: "f31");
708+
709+
// fgh.ps/flw.ps are blocking loads, so a non-blocking prefetch hint issued
710+
// a few blocks ahead turns a DRAM miss into an L1 hit by the time the
711+
// gather runs.
712+
const int64_t PF = 4;
713+
for (int64_t kb = 0; kb < n_blocks; kb++) {
714+
const uint16_t * a_ptr = a_row + (kb << 5); // 32 f16 per block
715+
const float * b_ptr = b_col + (kb << 5);
716+
if (kb + PF < n_blocks) {
717+
const uint16_t * pfa = a_ptr + (PF << 5); // A block = 64B = 1 line
718+
const float * pfb = b_ptr + (PF << 5); // B block = 128B = 2 lines
719+
__asm__ volatile(
720+
"lb x0, 0(%[pa])\n"
721+
"lb x0, 0(%[pb])\n"
722+
"lb x0, 64(%[pb])\n"
723+
: : [pa] "r"(pfa), [pb] "r"(pfb));
724+
}
725+
__asm__ volatile(
726+
// Issue all 8 independent memory ops first (4 A gathers + 4 B loads)
727+
// so several misses are outstanding before any consumer stalls.
728+
"fgh.ps f11, f31(%[a0])\n"
729+
"fgh.ps f13, f31(%[a1])\n"
730+
"fgh.ps f15, f31(%[a2])\n"
731+
"fgh.ps f17, f31(%[a3])\n"
732+
"flw.ps f12, %[b0]\n"
733+
"flw.ps f14, %[b1]\n"
734+
"flw.ps f16, %[b2]\n"
735+
"flw.ps f18, %[b3]\n"
736+
"fcvt.ps.f16 f11, f11\n"
737+
"fcvt.ps.f16 f13, f13\n"
738+
"fcvt.ps.f16 f15, f15\n"
739+
"fcvt.ps.f16 f17, f17\n"
740+
"fmadd.ps f20, f11, f12, f20\n"
741+
"fmadd.ps f21, f13, f14, f21\n"
742+
"fmadd.ps f22, f15, f16, f22\n"
743+
"fmadd.ps f23, f17, f18, f23\n"
744+
:
745+
: [a0] "r"(a_ptr), [a1] "r"(a_ptr + 8),
746+
[a2] "r"(a_ptr + 16), [a3] "r"(a_ptr + 24),
747+
[b0] "m"(*(const float(*)[8])&b_ptr[0]),
748+
[b1] "m"(*(const float(*)[8])&b_ptr[8]),
749+
[b2] "m"(*(const float(*)[8])&b_ptr[16]),
750+
[b3] "m"(*(const float(*)[8])&b_ptr[24])
751+
: "f11", "f12", "f13", "f14", "f15", "f16", "f17", "f18",
752+
"f20", "f21", "f22", "f23"
753+
);
754+
}
755+
}
756+
757+
static inline float __attribute__((always_inline))
758+
f16_dot_reduce(void) {
759+
float result;
760+
__asm__ __volatile__ (
761+
// Combine the 4 lane accumulators, then horizontal-sum the 8 lanes.
762+
"fadd.ps f20, f20, f21, rne \n\t"
763+
"fadd.ps f22, f22, f23, rne \n\t"
764+
"fadd.ps f20, f20, f22, rne \n\t"
765+
"fswizz.ps f1, f20, 0xB1 \n\t"
766+
"fadd.ps f2, f20, f1, rne \n\t"
767+
"fswizz.ps f3, f2, 0x4E \n\t"
768+
"fadd.ps f4, f2, f3, rne \n\t"
769+
"fmvz.x.ps t0, f4, 4 \n\t"
770+
"fbcx.ps f5, t0 \n\t"
771+
"fadd.ps %[vout], f4, f5, rne \n\t"
772+
: [vout] "=f" (result)
773+
:: "t0", "f1", "f2", "f3", "f4", "f5", "f20", "f21", "f22", "f23"
774+
);
775+
return result;
776+
}
777+
673778
// Compute dot product between f32 block and f32 column vector
674779
// Vectorized: processes 8 elements at a time using ET vector instructions
675780
// Block size: up to 16 f32 values (can handle partial blocks for misaligned K)

‎ggml/src/ggml-et/et-kernels/src/mul_mat_f16.c‎

Lines changed: 97 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,22 @@
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

Comments
 (0)