Skip to content

Commit 97416e8

Browse files
RehanQasim-devRehanQasim-dev
committed
et-backend: F32 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 f32 row-dot helper with L2 prefetch for the streamed weight row, and stages the reused B activation vector into per-shire L2 SCP. Also fixes the matrix-engine GEMV dispatch check, which compared src1->ne[0] instead of src1->ne[1] and so never actually caught the n=1 decode case, sending it to the matrix engine kernel where it stalls on a single-column matmul. Co-authored-by: Rehan Qasim <rehanbhatti0317@gmail.com>
1 parent 7cbd610 commit 97416e8

3 files changed

Lines changed: 196 additions & 80 deletions

File tree

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

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
# define BLOCK_OPS_H
88

99
# include "math_fp.h"
10+
# include "platform.h"
1011
# include "quants.h"
1112

1213
# include <stdint.h>
@@ -735,6 +736,98 @@ static inline float compute_block_dot_product_f32_f16_partial(const float * a
735736
return sum;
736737
}
737738

739+
//******************************************************************************
740+
// F32 x F32 register-resident row dot (decode GEMV). The running sum lives in
741+
// vector registers f20..f23 for the whole row instead of being spilled to
742+
// memory every few elements, so the inner loop is just paired flw.ps loads and
743+
// fmadd, no gather or convert needed since both operands are already f32.
744+
// Four independent lane accumulators break the fmadd dependency chain so
745+
// several A/B loads stay in flight, hiding load latency.
746+
//******************************************************************************
747+
static inline void __attribute__((always_inline))
748+
f32_dot_reset(void) {
749+
__asm__ volatile(
750+
"fbci.pi f20, 0\n"
751+
"fbci.pi f21, 0\n"
752+
"fbci.pi f22, 0\n"
753+
"fbci.pi f23, 0\n"
754+
::: "f20", "f21", "f22", "f23");
755+
}
756+
757+
// Accumulate n_blocks blocks of 32 f32 A values times 32 f32 B into f20..f23.
758+
static inline void __attribute__((always_inline))
759+
f32_dot_tile(const float * a_row, const float * b_col, int64_t n_blocks) {
760+
// Prefetch A ahead into L2, not L1. The weight row is streamed once and
761+
// never reused, but L1 is shared by both harts of a minion, so a simple
762+
// L1 prefetch gets evicted before the load consumes it and the line is
763+
// re-fetched from DRAM. l2_prefetch stages the line into L2 instead, which
764+
// is large enough to hold the in-flight window, so a re-touch hits L2
765+
// rather than causing a second DRAM read. B lives in on-chip L2 SCP
766+
// already, so it needs no prefetch.
767+
//
768+
// Each block is 128B = 2 lines. Prefetch 16 lines (8 blocks) at a time,
769+
// staying PF_AHEAD blocks in front of the consuming load.
770+
const int64_t PF_AHEAD = 16;
771+
for (int64_t kb = 0; kb < n_blocks; kb++) {
772+
const float * a_ptr = a_row + (kb << 5); // 32 f32 per block = 128B
773+
const float * b_ptr = b_col + (kb << 5);
774+
if ((kb & 7) == 0) {
775+
const int64_t pf_block = kb + PF_AHEAD;
776+
if (pf_block + 8 <= n_blocks) {
777+
l2_prefetch(a_row + (pf_block << 5), 16, 64); // 16 lines = 8 blocks
778+
}
779+
}
780+
__asm__ volatile(
781+
// Issue all 8 independent loads first (4 A + 4 B) so several misses
782+
// are outstanding before any consumer stalls in-order issue.
783+
"flw.ps f11, %[a0]\n"
784+
"flw.ps f13, %[a1]\n"
785+
"flw.ps f15, %[a2]\n"
786+
"flw.ps f17, %[a3]\n"
787+
"flw.ps f12, %[b0]\n"
788+
"flw.ps f14, %[b1]\n"
789+
"flw.ps f16, %[b2]\n"
790+
"flw.ps f18, %[b3]\n"
791+
"fmadd.ps f20, f11, f12, f20\n"
792+
"fmadd.ps f21, f13, f14, f21\n"
793+
"fmadd.ps f22, f15, f16, f22\n"
794+
"fmadd.ps f23, f17, f18, f23\n"
795+
:
796+
: [a0] "m"(*(const float(*)[8])&a_ptr[0]),
797+
[a1] "m"(*(const float(*)[8])&a_ptr[8]),
798+
[a2] "m"(*(const float(*)[8])&a_ptr[16]),
799+
[a3] "m"(*(const float(*)[8])&a_ptr[24]),
800+
[b0] "m"(*(const float(*)[8])&b_ptr[0]),
801+
[b1] "m"(*(const float(*)[8])&b_ptr[8]),
802+
[b2] "m"(*(const float(*)[8])&b_ptr[16]),
803+
[b3] "m"(*(const float(*)[8])&b_ptr[24])
804+
: "f11", "f12", "f13", "f14", "f15", "f16", "f17", "f18",
805+
"f20", "f21", "f22", "f23"
806+
);
807+
}
808+
}
809+
810+
static inline float __attribute__((always_inline))
811+
f32_dot_reduce(void) {
812+
float result;
813+
__asm__ __volatile__ (
814+
// Combine the 4 lane accumulators, then horizontal-sum the 8 lanes.
815+
"fadd.ps f20, f20, f21, rne \n\t"
816+
"fadd.ps f22, f22, f23, rne \n\t"
817+
"fadd.ps f20, f20, f22, rne \n\t"
818+
"fswizz.ps f1, f20, 0xB1 \n\t"
819+
"fadd.ps f2, f20, f1, rne \n\t"
820+
"fswizz.ps f3, f2, 0x4E \n\t"
821+
"fadd.ps f4, f2, f3, rne \n\t"
822+
"fmvz.x.ps t0, f4, 4 \n\t"
823+
"fbcx.ps f5, t0 \n\t"
824+
"fadd.ps %[vout], f4, f5, rne \n\t"
825+
: [vout] "=f" (result)
826+
:: "t0", "f1", "f2", "f3", "f4", "f5", "f20", "f21", "f22", "f23"
827+
);
828+
return result;
829+
}
830+
738831
// Compute dot product between f32 block and f32 column vector
739832
// Vectorized: processes 8 elements at a time using ET vector instructions
740833
// Block size: 16 f32 values (64 bytes = 1 cache line)

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

Lines changed: 102 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,24 @@
1+
//******************************************************************************
2+
// MUL_MAT Kernel (F32 weights)
3+
// 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. This mirrors the F16 GEMV
10+
// fix: stripe rows by thread_id, stage the reused B vector into per-shire L2
11+
// SCP and keep the accumulator register-resident across K via f32_dot_*.
12+
//******************************************************************************
13+
114
#include "block_ops.h"
215
#include "ggml_tensor.h"
16+
#include "math_fp.h"
317
#include "platform.h"
418
#include "quants.h"
19+
#include "tensor.h"
520

6-
#include <etsoc/common/utils.h>
721
#include <stdint.h>
8-
#include <stdio.h>
922

1023
int entry_point(struct ggml_et_binary_params * params, void * env) {
1124
kernel_environment_t * kernel_env = (kernel_environment_t *) env;
@@ -14,18 +27,13 @@ int entry_point(struct ggml_et_binary_params * params, void * env) {
1427
return -1;
1528
}
1629

17-
// Thread coordination
30+
// Thread coordination, use every hart of every active shire.
1831
int thread_id = get_relative_thread_id(kernel_env->shire_mask);
1932
int num_threads = get_num_threads(kernel_env->shire_mask);
20-
21-
if (thread_id < 0 || (thread_id & 1)) {
22-
return 0; // Skip odd threads to avoid resource contention
33+
if (thread_id < 0) {
34+
return 0;
2335
}
2436

25-
int effective_thread_id = thread_id / 2;
26-
int effective_num_threads = (num_threads + 1) / 2;
27-
28-
// Extract tensor references
2937
struct ggml_tensor * src0 = &params->src0; // Weight matrix A (F32)
3038
struct ggml_tensor * src1 = &params->src1; // Activation matrix B (F16/F32)
3139
struct ggml_tensor * dst = &params->dst; // Output matrix C (F32)
@@ -39,7 +47,7 @@ int entry_point(struct ggml_et_binary_params * params, void * env) {
3947
const float * src0_data = (const float *) src0->data;
4048
float * dst_data = (float *) dst->data;
4149

42-
// Dimensions and Strides
50+
// Dimensions and strides
4351
const int64_t K = src0->ne[0];
4452
const int64_t M = src0->ne[1];
4553
const int64_t N = src1->ne[1];
@@ -52,86 +60,101 @@ int entry_point(struct ggml_et_binary_params * params, void * env) {
5260
const size_t nb11 = src1->nb[1], nb12 = src1->nb[2], nb13 = src1->nb[3];
5361
const size_t nb1 = dst->nb[1], nb2 = dst->nb[2], nb3 = dst->nb[3];
5462

55-
// F32 specific block size and counts
56-
const int block_size = QK_F32;
63+
const int block_size = 32; // 32 f32 per row-dot tile (128B)
5764
const int64_t K_blocks = K / block_size;
5865
const int64_t K_remainder = K % block_size;
5966

60-
// Threading distribution
61-
const uint64_t total_elements = M * N * ne2 * ne3;
62-
const uint64_t per_thread = 16;
63-
const uint64_t threads_stride = per_thread * effective_num_threads;
64-
65-
if (effective_thread_id * per_thread >= total_elements) {
66-
return 0;
67-
}
68-
6967
// Broadcasting support
7068
const int64_t r2 = ne12 / ne02;
7169
const int64_t r3 = ne13 / ne03;
7270

73-
for (uint64_t base_idx = effective_thread_id * per_thread; base_idx < total_elements; base_idx += threads_stride) {
74-
for (uint64_t j = 0; j < per_thread; j++) {
75-
const uint64_t idx = base_idx + j;
76-
if (idx >= total_elements) {
77-
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_ln = 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_ln;
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_ln += cl;
106+
addr += cl * 64;
107+
remaining -= cl;
78108
}
109+
}
110+
et_barrier(ET_BARRIER_SHIRE); // B is now resident in L2 SCP for all harts
111+
}
79112

80-
// Index decoding
81-
const int64_t i3 = idx / (M * N * ne2);
82-
const int64_t rem3 = idx % (M * N * ne2);
83-
const int64_t i2 = rem3 / (M * N);
84-
const int64_t rem2 = rem3 % (M * N);
85-
const int64_t n = rem2 / M;
86-
const int64_t m = rem2 % M;
87-
88-
const int64_t i03 = i3 / r3, i02 = i2 / r2;
89-
const int64_t i13 = (ne13 > 1) ? i3 : 0, i12 = (ne12 > 1) ? i2 : 0;
90-
91-
float sum = 0.0f;
92-
const float * f32_row = (const float *) ((const char *) src0_data + m * nb01 + i02 * nb02 + i03 * nb03);
93-
94-
if (src1->type == GGML_TYPE_F32) {
95-
const float * src1_data = (const float *) src1->data;
96-
97-
for (int64_t kb = 0; kb < K_blocks; kb++) {
98-
const float * b_col_ptr =
99-
(const float *) ((const char *) src1_data + (kb * block_size) * sizeof(float) + n * nb11 +
100-
i12 * nb12 + i13 * nb13);
101-
sum += compute_block_dot_product_f32(&f32_row[kb * block_size], b_col_ptr);
102-
}
103-
104-
if (K_remainder > 0) {
105-
const int64_t offset = K_blocks * block_size;
106-
const float * b_col_ptr = (const float *) ((const char *) src1_data + offset * sizeof(float) +
107-
n * nb11 + i12 * nb12 + i13 * nb13);
108-
sum += compute_block_dot_product_f32_partial(&f32_row[offset], b_col_ptr, K_remainder);
109-
}
110-
} else {
111-
const uint16_t * src1_data = (const uint16_t *) src1->data;
112-
113-
for (int64_t kb = 0; kb < K_blocks; kb++) {
114-
const uint16_t * b_col_ptr =
115-
(const uint16_t *) ((const char *) src1_data + (kb * block_size) * sizeof(uint16_t) + n * nb11 +
116-
i12 * nb12 + i13 * nb13);
117-
sum += compute_block_dot_product_f32_f16_partial(&f32_row[kb * block_size], b_col_ptr, block_size);
118-
}
119-
120-
if (K_remainder > 0) {
121-
const int64_t offset = K_blocks * block_size;
122-
const uint16_t * b_col_ptr =
123-
(const uint16_t *) ((const char *) src1_data + offset * sizeof(uint16_t) + n * nb11 +
124-
i12 * nb12 + i13 * nb13);
125-
sum += compute_block_dot_product_f32_f16_partial(&f32_row[offset], b_col_ptr, K_remainder);
126-
}
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 float * f32_row =
131+
(const float *) ((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+
f32_dot_reset();
140+
f32_dot_tile(f32_row, b_col, K_blocks);
141+
sum = f32_dot_reduce();
142+
143+
if (K_remainder > 0) {
144+
const int64_t offset = K_blocks * block_size;
145+
sum += compute_block_dot_product_f32_partial(&f32_row[offset], &b_col[offset], K_remainder);
127146
}
128-
129-
// Atomic store for output
130-
volatile float * c_element =
131-
(volatile float *) ((char *) dst_data + m * dst->nb[0] + n * nb1 + i2 * nb2 + i3 * nb3);
132-
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_f32_f16_partial(f32_row, b_col, K);
133151
}
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);
134156
}
135157

158+
__asm__ volatile("mova.m.x %0" ::"r"(saved_mask));
136159
return 0;
137160
}

‎ggml/src/ggml-et/ggml-et-ops.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -749,7 +749,7 @@ bool ggml_et_op_mul_mat(ggml_backend_et_device_context * dev_ctx,
749749

750750
} else if (node->type == GGML_TYPE_F32 && node->src[0]->type == GGML_TYPE_F32 &&
751751
node->src[1]->type == GGML_TYPE_F32 && node->ne[0] % 16 == 0 && node->src[0]->ne[0] % 16 == 0 &&
752-
node->src[0]->ne[1] % 16 == 0 && node->src[1]->ne[0] != 1) { // GEMV is faster with the generic path
752+
node->src[0]->ne[1] % 16 == 0 && node->src[1]->ne[1] != 1) { // GEMV (n=1) is faster with the generic path
753753

754754
kernel_name = "mul_mat_f32_matrix_engine";
755755
src0_type_name = "F32";

0 commit comments

Comments
 (0)