Skip to content

Commit a0b76be

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 <rehan.qasim@10xengineers.ai>
1 parent 7cbd610 commit a0b76be

3 files changed

Lines changed: 195 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: 101 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,23 @@
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 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. The reused B vector is staged
9+
// into per-shire L2 SCP and the accumulator stays register-resident across K
10+
// via f32_dot_*.
11+
//******************************************************************************
12+
113
#include "block_ops.h"
214
#include "ggml_tensor.h"
15+
#include "math_fp.h"
316
#include "platform.h"
417
#include "quants.h"
18+
#include "tensor.h"
519

6-
#include <etsoc/common/utils.h>
720
#include <stdint.h>
8-
#include <stdio.h>
921

1022
int entry_point(struct ggml_et_binary_params * params, void * env) {
1123
kernel_environment_t * kernel_env = (kernel_environment_t *) env;
@@ -14,18 +26,13 @@ int entry_point(struct ggml_et_binary_params * params, void * env) {
1426
return -1;
1527
}
1628

17-
// Thread coordination
29+
// Thread coordination, use every hart of every active shire.
1830
int thread_id = get_relative_thread_id(kernel_env->shire_mask);
1931
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
32+
if (thread_id < 0) {
33+
return 0;
2334
}
2435

25-
int effective_thread_id = thread_id / 2;
26-
int effective_num_threads = (num_threads + 1) / 2;
27-
28-
// Extract tensor references
2936
struct ggml_tensor * src0 = &params->src0; // Weight matrix A (F32)
3037
struct ggml_tensor * src1 = &params->src1; // Activation matrix B (F16/F32)
3138
struct ggml_tensor * dst = &params->dst; // Output matrix C (F32)
@@ -39,7 +46,7 @@ int entry_point(struct ggml_et_binary_params * params, void * env) {
3946
const float * src0_data = (const float *) src0->data;
4047
float * dst_data = (float *) dst->data;
4148

42-
// Dimensions and Strides
49+
// Dimensions and strides
4350
const int64_t K = src0->ne[0];
4451
const int64_t M = src0->ne[1];
4552
const int64_t N = src1->ne[1];
@@ -52,86 +59,101 @@ int entry_point(struct ggml_et_binary_params * params, void * env) {
5259
const size_t nb11 = src1->nb[1], nb12 = src1->nb[2], nb13 = src1->nb[3];
5360
const size_t nb1 = dst->nb[1], nb2 = dst->nb[2], nb3 = dst->nb[3];
5461

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

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-
6966
// Broadcasting support
7067
const int64_t r2 = ne12 / ne02;
7168
const int64_t r3 = ne13 / ne03;
7269

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;
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_ln = 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_ln;
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_ln += cl;
105+
addr += cl * 64;
106+
remaining -= cl;
78107
}
108+
}
109+
et_barrier(ET_BARRIER_SHIRE); // B is now resident in L2 SCP for all harts
110+
}
79111

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-
}
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 float * f32_row =
130+
(const float *) ((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+
f32_dot_reset();
139+
f32_dot_tile(f32_row, b_col, K_blocks);
140+
sum = f32_dot_reduce();
141+
142+
if (K_remainder > 0) {
143+
const int64_t offset = K_blocks * block_size;
144+
sum += compute_block_dot_product_f32_partial(&f32_row[offset], &b_col[offset], K_remainder);
127145
}
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);
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_f32_f16_partial(f32_row, b_col, K);
133150
}
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);
134155
}
135156

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

‎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)