Skip to content

Commit 67937d7

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 67937d7

2 files changed

Lines changed: 201 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: 96 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,21 @@
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 by thread_id across every hart
7+
// of every active shire. This keeps all 32 shires busy for a GEMV
8+
// (M=4096, N=1) instead of leaving most idle, and the register-resident
9+
// f16_dot row helper keeps the accumulator in registers for the whole K
10+
// dimension.
411
//******************************************************************************
512

613
#include "block_ops.h"
714
#include "ggml_tensor.h"
815
#include "math_fp.h"
916
#include "platform.h"
1017
#include "quants.h"
18+
#include "tensor.h"
1119

1220
#include <stdint.h>
1321

@@ -18,18 +26,13 @@ int entry_point(struct ggml_et_binary_params * params, void * env) {
1826
return -1;
1927
}
2028

21-
// Thread coordination
29+
// Thread coordination, use every hart of every active shire.
2230
int thread_id = get_relative_thread_id(kernel_env->shire_mask);
2331
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
32+
if (thread_id < 0) {
33+
return 0;
2734
}
2835

29-
int effective_thread_id = thread_id / 2;
30-
int effective_num_threads = (num_threads + 1) / 2;
31-
32-
// Extract tensor references
3336
struct ggml_tensor * src0 = &params->src0; // Weight matrix A (F16)
3437
struct ggml_tensor * src1 = &params->src1; // Activation matrix B (F16/F32)
3538
struct ggml_tensor * dst = &params->dst; // Output matrix C (F32)
@@ -43,7 +46,7 @@ int entry_point(struct ggml_et_binary_params * params, void * env) {
4346
const uint16_t * src0_data = (const uint16_t *) src0->data;
4447
float * dst_data = (float *) dst->data;
4548

46-
// Dimensions and Strides
49+
// Dimensions and strides
4750
const int64_t K = src0->ne[0];
4851
const int64_t M = src0->ne[1];
4952
const int64_t N = src1->ne[1];
@@ -56,87 +59,101 @@ int entry_point(struct ggml_et_binary_params * params, void * env) {
5659
const size_t nb11 = src1->nb[1], nb12 = src1->nb[2], nb13 = src1->nb[3];
5760
const size_t nb1 = dst->nb[1], nb2 = dst->nb[2], nb3 = dst->nb[3];
5861

59-
// F16 specific block size (Usually QK_F16)
60-
const int block_size = QK_F16;
62+
const int block_size = QK_F16; // 32
6163
const int64_t K_blocks = K / block_size;
6264
const int64_t K_remainder = K % block_size;
6365

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-
7366
// Broadcasting support
7467
const int64_t r2 = ne12 / ne02;
7568
const int64_t r3 = ne13 / ne03;
7669

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;
70+
const int is_f32_b = (src1->type == GGML_TYPE_F32);
71+
const uint64_t total_elements = (uint64_t) M * N * ne2 * ne3;
72+
73+
// Stage the reused B activation vector into per-shire L2 SCP. Streaming the
74+
// weight matrix thrashes L2 and evicts B, so B is re-read from DRAM
75+
// repeatedly during the decode GEMV. Stage it once per shire into L2 SCP (a
76+
// separate SRAM partition cache streaming cannot evict) via
77+
// et_tensor_load_l2scp; one hart per shire issues the DMA loop, then a
78+
// shire barrier lets all harts read B on-chip. Only applies to the common
79+
// decode case: F32 B, a single contiguous non-broadcast B that fits the SCP
80+
// budget.
81+
const int b_contig = (nb11 == (size_t) K * sizeof(float));
82+
const uint64_t b_lines = ((uint64_t) N * K * sizeof(float) + 63) / 64;
83+
const int stage_b = is_f32_b && b_contig &&
84+
ne12 == 1 && ne13 == 1 && ne02 == 1 && ne03 == 1 &&
85+
b_lines <= 8192; // <= 512 KB, within the SCP budget
86+
const float * b_scp = (const float *) et_shire_l2scp_local(0);
87+
88+
if (stage_b) {
89+
if ((get_hart_id() & 63) == 0) {
90+
et_tensor_load_l2scp_conf_t conf;
91+
conf.use_tmask = false;
92+
conf.stride = 64; // advance one 64B cache line per line loaded
93+
uint64_t remaining = b_lines;
94+
uint64_t dst = 0;
95+
uint64_t addr = (uint64_t) src1->data;
96+
while (remaining > 0) {
97+
uint64_t cl = (remaining >= 16) ? 16 : remaining;
98+
conf.dst_start = dst;
99+
conf.addr = addr;
100+
conf.num_lines = cl - 1; // 4-bit field encodes (lines - 1)
101+
conf.id = 0;
102+
et_tensor_load_l2scp(&conf);
103+
WAIT_TENSOR_LOAD_L2_0;
104+
dst += cl;
105+
addr += cl * 64;
106+
remaining -= cl;
82107
}
108+
}
109+
et_barrier(ET_BARRIER_SHIRE); // B is now resident in L2 SCP for all harts
110+
}
83111

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-
}
112+
// Set the vector mask (all 8 lanes) once for the whole row loop.
113+
unsigned long saved_mask;
114+
__asm__ volatile("mova.x.m %0" : "=r"(saved_mask));
115+
__asm__ volatile("mov.m.x m0, x0, 0xFF");
116+
117+
for (uint64_t idx = (uint64_t) thread_id; idx < total_elements; idx += (uint64_t) num_threads) {
118+
// Index decoding
119+
const int64_t i3 = idx / (M * N * ne2);
120+
const int64_t rem3 = idx % (M * N * ne2);
121+
const int64_t i2 = rem3 / (M * N);
122+
const int64_t rem2 = rem3 % (M * N);
123+
const int64_t n = rem2 / M;
124+
const int64_t m = rem2 % M;
125+
126+
const int64_t i03 = i3 / r3, i02 = i2 / r2;
127+
const int64_t i13 = (ne13 > 1) ? i3 : 0, i12 = (ne12 > 1) ? i2 : 0;
128+
129+
const uint16_t * f16_row =
130+
(const uint16_t *) ((const char *) src0_data + m * nb01 + i02 * nb02 + i03 * nb03);
131+
132+
float sum;
133+
if (is_f32_b) {
134+
const float * b_col = stage_b
135+
? (b_scp + n * K)
136+
: (const float *) ((const char *) src1->data + n * nb11 + i12 * nb12 + i13 * nb13);
137+
138+
f16_dot_reset();
139+
f16_dot_tile(f16_row, b_col, K_blocks);
140+
sum = f16_dot_reduce();
141+
142+
if (K_remainder > 0) {
143+
const int64_t offset = K_blocks * block_size;
144+
sum += compute_block_dot_product_f16_partial(&f16_row[offset], &b_col[offset], K_remainder);
132145
}
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);
146+
} else {
147+
const uint16_t * b_col =
148+
(const uint16_t *) ((const char *) src1->data + n * nb11 + i12 * nb12 + i13 * nb13);
149+
sum = compute_block_dot_product_f16_f16_partial(f16_row, b_col, K);
138150
}
151+
152+
volatile float * c_element =
153+
(volatile float *) ((char *) dst_data + m * dst->nb[0] + n * nb1 + i2 * nb2 + i3 * nb3);
154+
atomic_store_f32(c_element, sum);
139155
}
140156

157+
__asm__ volatile("mova.m.x %0" ::"r"(saved_mask));
141158
return 0;
142159
}

0 commit comments

Comments
 (0)