Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

hexagon: further improved pipeline of the core bits (L2, DMA, MM, FA) (#26049)

Browse Source 
* hex-l2: use dirty ranges for flushing

* hex-l2: simplify range based flush logic

* hex-l2: optimize dirty range scans

* hex-hvx: support for reduce_max_i32

* hex-mm: optimize fused MUL_MAT+ADD to use vtcm for bias when it fits

* hex-mmid: optimize mmid row-mapping generation

* hex-mmid: optimize mmid row-mapping generation

* hex-mmid: optimize mmid row-mapping generation (round2)

* hmx-mm: optimize output proc by tiling (col-chunking)

* hex-fa: start the next q dmas a bit earlier

* hex-fa: prefetch Q even earlier

* hvx-fa: optimize softmax to keep things in hvx registers

* hex-fa: hoist const register init in softmax loop

* hmx-fa: kick off next-qkv DMAs before o-proc

* hmx-fa: hoist various checks out of the inner loop

* hmx-fa: adjust the cost model to better balance softmax work across hvx threads

* hmx-fa: overlap diag rescale build with last HMX task

* hmx-fa: optimize idx update in output proc

* hmx-fa: unroll the softmax loops for improved perf

* hmx-fa: overlap qk-dot with softmax, double-buffer p and s tiles

* hex-trace: double the default number of trace entries

* hex-trace: add trace events for opbatch and buffer mgmt

* hex-trace: overhaul tracing to simplify runtime event handling and support opbatch stats

* hex-trace: replace ascii timeline diagram with pipeline bubbles detector

* hex-trace: handle missing start/stop events

* hex-dma: always log stop/start trace events even for dummy dmas

* hex-scripts: fix flake warnings
This commit is contained in:
Max Krasnyansky
2026-07-23 19:13:03 -07:00
committed by GitHub
parent c0bc8591e8
commit 0a50d9909a
17 changed files with 1666 additions and 761 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-hexagon/ggml-hexagon.cpp
+92 -153
View File
@@ -143,12 +143,12 @@ static const char * htp_event_name(uint16_t id) {
case HTP_TRACE_EVT_HMX_COMP: return "HMX_COMP"; case HTP_TRACE_EVT_HMX_COMP: return "HMX_COMP";
case HTP_TRACE_EVT_L2FLUSH: return "L2FLUSH"; case HTP_TRACE_EVT_L2FLUSH: return "L2FLUSH";
case HTP_TRACE_EVT_INIT: return "INIT"; case HTP_TRACE_EVT_INIT: return "INIT";
case HTP_TRACE_EVT_BUFF: return "BUFF";
default: return "UNKNOWN"; default: return "UNKNOWN";
} }
} }
static void ggml_hexagon_dump_op_prof(const std::string &sess_name, const htp_opnode & node, static void ggml_hexagon_dump_op_prof(const std::string &sess_name, const htp_opnode & node, const htp_prof_desc & pd) {
const htp_prof_desc & pd) {
if (!opt_profile) return; if (!opt_profile) return;
uint32_t op_usec = pd.usecs; uint32_t op_usec = pd.usecs;
@@ -168,6 +168,43 @@ static void ggml_hexagon_dump_op_prof(const std::string &sess_name, const htp_op
node.op_name().c_str(), fmt.names, fmt.dims, fmt.types, fmt.strides, fmt.kparams, op_usec, op_cycles, pd.cycles_start, mhz, pmu_str); node.op_name().c_str(), fmt.names, fmt.dims, fmt.types, fmt.strides, fmt.kparams, op_usec, op_cycles, pd.cycles_start, mhz, pmu_str);
} }
static void ggml_hexagon_dump_batch_prof(const std::string & sess_name, const htp_opbatch_rsp & rsp) {
uint64_t batch_cycles = rsp.cycles_stop - rsp.cycles_start;
float batch_mhz = rsp.usecs > 0 ? (float) batch_cycles / rsp.usecs : 0.0f;
char evt_str[256] = "----";
if (opt_profile == 3) {
snprintf(evt_str, sizeof(evt_str), "evt-cnt %u,%u,%u,%u,%u,%u,%u,%u,%u,%u,%u",
rsp.n_traces[0], rsp.n_traces[1], rsp.n_traces[2], rsp.n_traces[3],
rsp.n_traces[4], rsp.n_traces[5], rsp.n_traces[6], rsp.n_traces[7],
rsp.n_traces[8], rsp.n_traces[9], rsp.n_traces[10]);
}
GGML_LOG_DEBUG("ggml-hex: %s profile-op OPBATCH|----|n-ops %u|%s|----|----|usec %u cycles %llu start %llu mhz %.1f\n",
sess_name.c_str(), rsp.n_ops, evt_str, rsp.usecs, (unsigned long long) batch_cycles, (unsigned long long) rsp.cycles_start, batch_mhz);
}
static void ggml_hexagon_dump_trace_events(const std::string & sess_name, const htp_opbatch_rsp & rsp,
const htp_trace_desc * trace_events, uint32_t n_traces) {
if (opt_profile == 3 && trace_events) {
uint32_t valid_cnt[HTP_MAX_NTHREADS + 1] = {0};
for (uint32_t t = 0; t <= HTP_MAX_NTHREADS; t++) {
uint32_t count = rsp.n_traces[t];
valid_cnt[t] = count > n_traces ? n_traces : count;
}
for (uint32_t t = 0; t <= HTP_MAX_NTHREADS; t++) {
for (uint32_t idx = 0; idx < valid_cnt[t]; idx++) {
const auto & e = trace_events[t * n_traces + idx];
bool is_stop = (e.info & 0x8000) != 0;
uint16_t info = e.info & 0x7FFF;
GGML_LOG_DEBUG("ggml-hex: %s trace-evt %s: thread %u info %u %s %u\n",
sess_name.c_str(), htp_event_name(e.id), t, info, is_stop ? "stop" : "start", e.cycles);
}
}
}
}
// ** // **
static inline bool ggml_hexagon_is_repack_type(enum ggml_type type) { static inline bool ggml_hexagon_is_repack_type(enum ggml_type type) {
@@ -1128,13 +1165,7 @@ struct ggml_hexagon_opbatch {
std::unordered_map<const ggml_tensor*, int> t_map; // tensor ptr to index std::unordered_map<const ggml_tensor*, int> t_map; // tensor ptr to index
std::unordered_multimap<void*, int> d_map; // tensor data to index std::unordered_multimap<void*, int> d_map; // tensor data to index
struct tensor_range {
uint64_t start;
uint64_t end;
int bi;
std::vector<int> tensors;
};
std::vector<tensor_range> ranges;
unsigned int n_bufs; // num buffers in the batch unsigned int n_bufs; // num buffers in the batch
unsigned int n_tens; // num tensors ... unsigned int n_tens; // num tensors ...
@@ -1155,7 +1186,6 @@ struct ggml_hexagon_opbatch {
b_map.clear(); b_map.clear();
t_map.clear(); t_map.clear();
d_map.clear(); d_map.clear();
ranges.clear();
} }
ggml_hexagon_opbatch(ggml_hexagon_session *sess, size_t batch_size, size_t max_vmem) { ggml_hexagon_opbatch(ggml_hexagon_session *sess, size_t batch_size, size_t max_vmem) {
@@ -1209,70 +1239,7 @@ struct ggml_hexagon_opbatch {
return bi; return bi;
} }
void add_range(const htp_tensor * h, int ti) {
uint64_t t_start = h->data;
uint64_t t_end = t_start + h->size;
int bi = h->bi;
int first_match = -1;
int unused_idx = -1;
for (size_t i = 0; i < ranges.size(); i++) {
if (ranges[i].bi == -1) {
unused_idx = i;
continue;
}
if (ranges[i].bi != bi) {
continue;
}
if (ranges[i].start >= t_end || ranges[i].end <= t_start) {
continue;
}
if (first_match == -1) {
first_match = i;
HEX_VERBOSE("ggml-hex: %s range-grow #%d : bi %d [%p, %p) + #%d [%p, %p) -> [%p, %p)\n",
sess->c_name(), (int) i, ranges[i].bi,
(void *) (h_bufs[ranges[i].bi].base + ranges[i].start),
(void *) (h_bufs[ranges[i].bi].base + ranges[i].end),
ti,
(void *) (h_bufs[bi].base + t_start),
(void *) (h_bufs[bi].base + t_end),
(void *) (h_bufs[ranges[i].bi].base + std::min(ranges[i].start, t_start)),
(void *) (h_bufs[ranges[i].bi].base + std::max(ranges[i].end, t_end)));
ranges[i].start = std::min(ranges[i].start, t_start);
ranges[i].end = std::max(ranges[i].end, t_end);
ranges[i].tensors.push_back(ti);
} else {
HEX_VERBOSE("ggml-hex: %s range-merge #%d [%p, %p) + #%d [%p, %p) -> [%p, %p)\n",
sess->c_name(), first_match,
(void *) (h_bufs[bi].base + ranges[first_match].start),
(void *) (h_bufs[bi].base + ranges[first_match].end),
(int) i,
(void *) (h_bufs[bi].base + ranges[i].start),
(void *) (h_bufs[bi].base + ranges[i].end),
(void *) (h_bufs[bi].base + std::min(ranges[first_match].start, ranges[i].start)),
(void *) (h_bufs[bi].base + std::max(ranges[first_match].end, ranges[i].end)));
ranges[first_match].start = std::min(ranges[first_match].start, ranges[i].start);
ranges[first_match].end = std::max(ranges[first_match].end, ranges[i].end);
ranges[first_match].tensors.insert(
ranges[first_match].tensors.end(),
ranges[i].tensors.begin(),
ranges[i].tensors.end()
);
ranges[i].bi = -1;
}
}
if (first_match == -1) {
if (unused_idx != -1) {
ranges[unused_idx] = {t_start, t_end, bi, {ti}};
} else {
ranges.push_back({t_start, t_end, bi, {ti}});
}
}
}
bool same_shape(const htp_tensor * h, const ggml_tensor * t) const { bool same_shape(const htp_tensor * h, const ggml_tensor * t) const {
int64_t ne0 = t->ne[0]; int64_t ne0 = t->ne[0];
@@ -1341,8 +1308,7 @@ struct ggml_hexagon_opbatch {
h.nb[0] = t->nb[0]; h.nb[1] = t->nb[1]; h.nb[2] = t->nb[2]; h.nb[3] = t->nb[3]; h.nb[0] = t->nb[0]; h.nb[1] = t->nb[1]; h.nb[2] = t->nb[2]; h.nb[3] = t->nb[3];
} }
h.alias = ti;
add_range(&h, ti);
h.flags = 0; h.flags = 0;
if (ggml_backend_buffer_get_usage(t->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS) { if (ggml_backend_buffer_get_usage(t->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS) {
@@ -1424,14 +1390,6 @@ struct ggml_hexagon_opbatch {
} }
void finalize_ranges() { void finalize_ranges() {
for (const auto & r : ranges) {
if (r.bi == -1) {
continue;
}
for (size_t i = 0; i < r.tensors.size(); i++) {
h_tens[r.tensors[i]].alias = r.tensors[(i + 1) % r.tensors.size()];
}
}
} }
}; };
@@ -1582,9 +1540,6 @@ struct ggml_hexagon_opqueue {
if (opt_profile && rsp.n_ops > 0) { if (opt_profile && rsp.n_ops > 0) {
auto & ops = op_cache[rsp.id]; auto & ops = op_cache[rsp.id];
uint64_t batch_usec = ggml_time_us() - start_usec[rsp.id];
uint32_t htp_usec = 0;
GGML_ASSERT(rsp.n_ops <= ops.size()); GGML_ASSERT(rsp.n_ops <= ops.size());
const htp_prof_desc * pd = (const htp_prof_desc *) p_ptr; const htp_prof_desc * pd = (const htp_prof_desc *) p_ptr;
@@ -1595,55 +1550,13 @@ struct ggml_hexagon_opqueue {
trace_events = (const htp_trace_desc *) (p_ptr + p_size); trace_events = (const htp_trace_desc *) (p_ptr + p_size);
} }
uint32_t trace_idx[HTP_MAX_NTHREADS + 1] = {0}; ggml_hexagon_dump_batch_prof(shm_buf->sess->name, rsp);
uint32_t valid_cnt[HTP_MAX_NTHREADS + 1] = {0};
if (opt_profile == 3) {
for (uint32_t t = 0; t <= HTP_MAX_NTHREADS; t++) {
uint32_t count = rsp.n_traces[t];
valid_cnt[t] = count > n_traces ? n_traces : count;
}
}
for (uint32_t i = 0; i < rsp.n_ops; i++) { for (uint32_t i = 0; i < rsp.n_ops; i++) {
htp_usec += pd[i].usecs;
ggml_hexagon_dump_op_prof(shm_buf->sess->name, ops[i], pd[i]); ggml_hexagon_dump_op_prof(shm_buf->sess->name, ops[i], pd[i]);
if (opt_profile == 3) {
uint32_t op_duration = pd[i].cycles_stop - pd[i].cycles_start;
for (uint32_t t = 0; t <= HTP_MAX_NTHREADS; t++) {
while (trace_idx[t] < valid_cnt[t]) {
const auto & e = trace_events[t * n_traces + trace_idx[t]];
uint32_t offset = e.cycles - pd[i].cycles_start;
if (offset >= 0x80000000) {
trace_idx[t]++;
continue;
}
if (offset > op_duration) {
break;
}
bool is_stop = (e.info & 0x8000) != 0;
uint16_t info = e.info & 0x7FFF;
GGML_LOG_DEBUG("ggml-hex: %s trace-op %s: thread %u event %s info %u %s %u\n",
shm_buf->sess->c_name(), ops[i].op_name().c_str(), t, htp_event_name(e.id), info, is_stop ? "stop" : "start", e.cycles);
trace_idx[t]++;
}
}
}
} }
char evt_str[256] = ""; ggml_hexagon_dump_trace_events(shm_buf->sess->name, rsp, trace_events, n_traces);
if (opt_profile == 3) {
snprintf(evt_str, sizeof(evt_str), " evt [%u,%u,%u,%u,%u,%u,%u,%u,%u,%u,%u]",
rsp.n_traces[0], rsp.n_traces[1], rsp.n_traces[2], rsp.n_traces[3],
rsp.n_traces[4], rsp.n_traces[5], rsp.n_traces[6], rsp.n_traces[7],
rsp.n_traces[8], rsp.n_traces[9], rsp.n_traces[10]);
}
GGML_LOG_DEBUG("ggml-hex: %s profile-batch n-ops %u batch-dur-usec %lld htp-ops-usec %u%s\n",
shm_buf->sess->c_name(), rsp.n_ops, (long long) batch_usec, htp_usec, evt_str);
} }
} }
}; };
@@ -2114,7 +2027,7 @@ static bool ggml_hexagon_precompute_flash_attn_params(
const struct ggml_tensor * sinks = op->src[4]; const struct ggml_tensor * sinks = op->src[4];
if (ggml_hexagon_flash_attn_is_hmx_eligible(sess, q, k, v, sinks)) { if (ggml_hexagon_flash_attn_is_hmx_eligible(sess, q, k, v, sinks)) {
size_t Br = 0, Bc = 0; size_t Br = 0, Bc = 0;
int ret = hmx_fa_find_chunk_size(&Br, &Bc, G, DK, DV, neq1, nek1, sess->vtcm_size, sess->n_threads); int ret = hmx_fa_find_chunk_size(&Br, &Bc, G, DK, DV, neq1, nek1, sess->vtcm_size, sess->n_threads, kparams->is_q_fp32 != 0);
if (ret == 0) { if (ret == 0) {
kparams->kernel_type = HTP_FA_KERNEL_HMX; kparams->kernel_type = HTP_FA_KERNEL_HMX;
kparams->Br = Br; kparams->Br = Br;
@@ -2124,7 +2037,7 @@ static bool ggml_hexagon_precompute_flash_attn_params(
kparams->u.hmx.g_br = hex_align_up(G * Br, 32); kparams->u.hmx.g_br = hex_align_up(G * Br, 32);
kparams->u.hmx.pipeline = (kparams->n_kv_blocks >= 3 && sess->n_threads >= 2) ? 1 : 0; kparams->u.hmx.pipeline = (kparams->n_kv_blocks >= 3 && sess->n_threads >= 2) ? 1 : 0;
kparams->vtcm_size = hmx_fa_compute_vtcm_usage(G, DK, DV, Br, Bc, kparams->n_threads, kparams->u.hmx.pipeline != 0); kparams->vtcm_size = hmx_fa_compute_vtcm_usage(G, DK, DV, Br, Bc, kparams->n_threads, kparams->u.hmx.pipeline != 0, kparams->is_q_fp32 != 0);
const size_t row_vec_bytes = hex_align_up(Bc * sizeof(uint16_t), 256); const size_t row_vec_bytes = hex_align_up(Bc * sizeof(uint16_t), 256);
kparams->u.hmx.row_buf_stride = row_vec_bytes / 128; // HVX vector is 128 bytes kparams->u.hmx.row_buf_stride = row_vec_bytes / 128; // HVX vector is 128 bytes
@@ -2413,6 +2326,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
int ne12, int ne12,
int ne13, int ne13,
bool is_matmul_id, bool is_matmul_id,
const size_t src2_row_size,
size_t vtcm_budget, size_t vtcm_budget,
struct htp_mm_kernel_params * kparams struct htp_mm_kernel_params * kparams
) { ) {
@@ -2438,7 +2352,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
for (uint32_t d = max_prefetch; d >= 2; d /= 2) { for (uint32_t d = max_prefetch; d >= 2; d /= 2) {
htp_mm_hvx_vtcm_layout_build( htp_mm_hvx_vtcm_layout_build(
&L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
0, src0->nb[1], 0, d, true, false, false 0, src0->nb[1], 0, src2_row_size, d, true, false, false
); );
if (L.total_bytes <= vtcm_budget) { if (L.total_bytes <= vtcm_budget) {
best_n_prefetch = d; best_n_prefetch = d;
@@ -2448,7 +2362,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
if (best_n_prefetch == 2 && L.total_bytes > vtcm_budget) { if (best_n_prefetch == 2 && L.total_bytes > vtcm_budget) {
htp_mm_hvx_vtcm_layout_build( htp_mm_hvx_vtcm_layout_build(
&L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
0, src0->nb[1], 0, 2, true, false, false 0, src0->nb[1], 0, src2_row_size, 2, true, false, false
); );
} }
kparams->n_prefetch = best_n_prefetch; kparams->n_prefetch = best_n_prefetch;
@@ -2472,7 +2386,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
for (uint32_t d = max_prefetch; d >= 2; d /= 2) { for (uint32_t d = max_prefetch; d >= 2; d /= 2) {
htp_mm_hvx_vtcm_layout_build( htp_mm_hvx_vtcm_layout_build(
&L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
dst->nb[1], src0->nb[1], src1->nb[1], d, false, false, false dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, d, false, false, false
); );
if (L.total_bytes <= vtcm_budget) { if (L.total_bytes <= vtcm_budget) {
best_n_prefetch = d; best_n_prefetch = d;
@@ -2482,7 +2396,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
if (best_n_prefetch == 2 && L.total_bytes > vtcm_budget) { if (best_n_prefetch == 2 && L.total_bytes > vtcm_budget) {
htp_mm_hvx_vtcm_layout_build( htp_mm_hvx_vtcm_layout_build(
&L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
dst->nb[1], src0->nb[1], src1->nb[1], 2, false, false, false dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 2, false, false, false
); );
} }
@@ -2506,7 +2420,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
struct htp_mm_hvx_vtcm_layout L; struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build( htp_mm_hvx_vtcm_layout_build(
&L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false
); );
kparams->n_prefetch = 16; kparams->n_prefetch = 16;
@@ -2526,7 +2440,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
struct htp_mm_hvx_vtcm_layout L; struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build( htp_mm_hvx_vtcm_layout_build(
&L, HTP_MM_KERNEL_HVX_F16_F16_VTCM, wtype, ne10, src1_nrows, sess->n_threads, &L, HTP_MM_KERNEL_HVX_F16_F16_VTCM, wtype, ne10, src1_nrows, sess->n_threads,
dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false
); );
if (!is_batched && !is_permuted && L.total_bytes <= vtcm_budget) { if (!is_batched && !is_permuted && L.total_bytes <= vtcm_budget) {
@@ -2546,7 +2460,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
kparams->src1_row_size = src1->nb[1]; kparams->src1_row_size = src1->nb[1];
htp_mm_hvx_vtcm_layout_build( htp_mm_hvx_vtcm_layout_build(
&L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false
); );
kparams->vtcm_size = L.total_bytes; kparams->vtcm_size = L.total_bytes;
kparams->vtcm_src0_size = L.src0_bytes; kparams->vtcm_src0_size = L.src0_bytes;
@@ -2562,7 +2476,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
struct htp_mm_hvx_vtcm_layout L; struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build( htp_mm_hvx_vtcm_layout_build(
&L, HTP_MM_KERNEL_HVX_F32_F32_VTCM, wtype, ne10, src1_nrows, sess->n_threads, &L, HTP_MM_KERNEL_HVX_F32_F32_VTCM, wtype, ne10, src1_nrows, sess->n_threads,
dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false
); );
if (!is_batched && !is_permuted && L.total_bytes <= vtcm_budget) { if (!is_batched && !is_permuted && L.total_bytes <= vtcm_budget) {
@@ -2578,7 +2492,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
kparams->src1_row_size = src1->nb[1]; kparams->src1_row_size = src1->nb[1];
htp_mm_hvx_vtcm_layout_build( htp_mm_hvx_vtcm_layout_build(
&L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false
); );
kparams->vtcm_size = L.total_bytes; kparams->vtcm_size = L.total_bytes;
kparams->vtcm_src0_size = L.src0_bytes; kparams->vtcm_src0_size = L.src0_bytes;
@@ -2589,11 +2503,12 @@ static void ggml_hexagon_precompute_hvx_mm_params(
} }
} }
static void ggml_hexagon_precompute_matmul_params( static void ggml_hexagon_precompute_matmul_params_impl(
const struct ggml_hexagon_session * sess, const struct ggml_hexagon_session * sess,
const struct ggml_tensor * src0, const struct ggml_tensor * src0,
const struct ggml_tensor * src1, const struct ggml_tensor * src1,
const struct ggml_tensor * dst, const struct ggml_tensor * dst,
const size_t src2_row_size,
struct htp_mm_kernel_params * kparams struct htp_mm_kernel_params * kparams
) { ) {
memset(kparams, 0, sizeof(*kparams)); memset(kparams, 0, sizeof(*kparams));
@@ -2628,7 +2543,7 @@ static void ggml_hexagon_precompute_matmul_params(
} }
// Fallback to HVX parameter computation // Fallback to HVX parameter computation
ggml_hexagon_precompute_hvx_mm_params(sess, src0, src1, dst, wtype, ne02, ne03, ne10, ne11, ne12, ne13, is_matmul_id, vtcm_budget, kparams); ggml_hexagon_precompute_hvx_mm_params(sess, src0, src1, dst, wtype, ne02, ne03, ne10, ne11, ne12, ne13, is_matmul_id, src2_row_size, vtcm_budget, kparams);
finalize: finalize:
kparams->div_ne12_ne1 = init_fastdiv_values(ne12 * ne11); kparams->div_ne12_ne1 = init_fastdiv_values(ne12 * ne11);
@@ -2638,6 +2553,27 @@ finalize:
kparams->div_ne11 = init_fastdiv_values(ne11); kparams->div_ne11 = init_fastdiv_values(ne11);
} }
static void ggml_hexagon_precompute_matmul_params(
const struct ggml_hexagon_session * sess,
const struct ggml_tensor * src0,
const struct ggml_tensor * src1,
const struct ggml_tensor * dst,
struct htp_mm_kernel_params * kparams
) {
ggml_hexagon_precompute_matmul_params_impl(sess, src0, src1, dst, 0, kparams);
}
static void ggml_hexagon_precompute_fused_matmul_add_params(
const struct ggml_hexagon_session * sess,
const struct ggml_tensor * src0,
const struct ggml_tensor * src1,
const struct ggml_tensor * src2,
const struct ggml_tensor * dst,
struct htp_mm_kernel_params * kparams
) {
ggml_hexagon_precompute_matmul_params_impl(sess, src0, src1, dst, src2->nb[1], kparams);
}
static void ggml_hexagon_precompute_unary_params( static void ggml_hexagon_precompute_unary_params(
const struct ggml_hexagon_session * sess, const struct ggml_hexagon_session * sess,
uint32_t op, uint32_t op,
@@ -2731,7 +2667,7 @@ static void ggml_hexagon_precompute_fused_qkv_params(
struct htp_mm_hvx_vtcm_layout L; struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build( htp_mm_hvx_vtcm_layout_build(
&L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads, &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
0, src0_row_size, src1_row_size, d, false, true, false 0, src0_row_size, src1_row_size, 0, d, false, true, false
); );
if (L.total_bytes <= sess->vtcm_size) { if (L.total_bytes <= sess->vtcm_size) {
best_n_prefetch = d; best_n_prefetch = d;
@@ -2746,7 +2682,7 @@ static void ggml_hexagon_precompute_fused_qkv_params(
// Test tiled first // Test tiled first
htp_mm_hvx_vtcm_layout_build( htp_mm_hvx_vtcm_layout_build(
&L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads, &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
0, src0_row_size, src1_row_size, best_n_prefetch, false, true, false 0, src0_row_size, src1_row_size, 0, best_n_prefetch, false, true, false
); );
if (try_tiled && L.total_bytes <= sess->vtcm_size) { if (try_tiled && L.total_bytes <= sess->vtcm_size) {
@@ -2764,7 +2700,7 @@ static void ggml_hexagon_precompute_fused_qkv_params(
htp_mm_hvx_vtcm_layout_build( htp_mm_hvx_vtcm_layout_build(
&L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads, &L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads,
0, src0_row_size, flat_src1_row_size, best_n_prefetch, false, true, false 0, src0_row_size, flat_src1_row_size, 0, best_n_prefetch, false, true, false
); );
kparams->vtcm_src0_size = L.src0_bytes; kparams->vtcm_src0_size = L.src0_bytes;
kparams->vtcm_src1_size = L.src1_bytes; kparams->vtcm_src1_size = L.src1_bytes;
@@ -2801,7 +2737,7 @@ static void ggml_hexagon_precompute_fused_ffn_params(
struct htp_mm_hvx_vtcm_layout L; struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build( htp_mm_hvx_vtcm_layout_build(
&L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads, &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
0, src0_row_size, src1_row_size, d, false, false, true 0, src0_row_size, src1_row_size, 0, d, false, false, true
); );
if (L.total_bytes <= sess->vtcm_size) { if (L.total_bytes <= sess->vtcm_size) {
best_n_prefetch = d; best_n_prefetch = d;
@@ -2816,7 +2752,7 @@ static void ggml_hexagon_precompute_fused_ffn_params(
// Test tiled first // Test tiled first
htp_mm_hvx_vtcm_layout_build( htp_mm_hvx_vtcm_layout_build(
&L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads, &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
0, src0_row_size, src1_row_size, best_n_prefetch, false, false, true 0, src0_row_size, src1_row_size, 0, best_n_prefetch, false, false, true
); );
if (try_tiled && L.total_bytes <= sess->vtcm_size) { if (try_tiled && L.total_bytes <= sess->vtcm_size) {
@@ -2833,7 +2769,7 @@ static void ggml_hexagon_precompute_fused_ffn_params(
htp_mm_hvx_vtcm_layout_build( htp_mm_hvx_vtcm_layout_build(
&L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads, &L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads,
0, src0_row_size, flat_src1_row_size, best_n_prefetch, false, false, true 0, src0_row_size, flat_src1_row_size, 0, best_n_prefetch, false, false, true
); );
kparams->vtcm_src0_size = L.src0_bytes; kparams->vtcm_src0_size = L.src0_bytes;
kparams->vtcm_src1_size = L.src1_bytes; kparams->vtcm_src1_size = L.src1_bytes;
@@ -3656,16 +3592,19 @@ static bool try_fuse_node(const ggml_hexagon_session * sess, const ggml_cgraph *
if (n->op == GGML_OP_MUL_MAT && next_node) { if (n->op == GGML_OP_MUL_MAT && next_node) {
if (next_node->op == GGML_OP_ADD && op_is_compute(next_node) && ggml_can_fuse(graph, i, { GGML_OP_MUL_MAT, GGML_OP_ADD })) { if (next_node->op == GGML_OP_ADD && op_is_compute(next_node) && ggml_can_fuse(graph, i, { GGML_OP_MUL_MAT, GGML_OP_ADD })) {
if (next_node->src[0] == n || next_node->src[1] == n) { if (next_node->src[0] == n || next_node->src[1] == n) {
const struct ggml_tensor * src2 = (next_node->src[0] == n) ? next_node->src[1] : next_node->src[0];
struct htp_mm_kernel_params kparams; struct htp_mm_kernel_params kparams;
ggml_hexagon_precompute_matmul_params(sess, n->src[0], n->src[1], next_node, &kparams); ggml_hexagon_precompute_fused_matmul_add_params(sess, n->src[0], n->src[1], src2, next_node, &kparams);
if ((size_t)kparams.vtcm_size <= sess->vtcm_size) { const int src1_nrows = n->src[1]->ne[1] * n->src[1]->ne[2] * n->src[1]->ne[3];
const bool can_fuse = (kparams.n_hmx > 0) || (src1_nrows == 1);
if (can_fuse && (size_t)kparams.vtcm_size <= sess->vtcm_size) {
htp_opnode node(n, {}, HTP_OP_MUL_MAT_ADD); htp_opnode node(n, {}, HTP_OP_MUL_MAT_ADD);
node.add_fused(next_node); node.add_fused(next_node);
memcpy(node.kernel_params, &kparams, sizeof(kparams)); memcpy(node.kernel_params, &kparams, sizeof(kparams));
nodes.push_back(std::move(node)); nodes.push_back(std::move(node));
i += 1; i += 1;
return true; return true;
} else { } else if (can_fuse) {
HEX_VERBOSE("ggml-hex: skip MUL_MAT_ADD fusion because VTCM needed (%d) > budget (%zu)\n", HEX_VERBOSE("ggml-hex: skip MUL_MAT_ADD fusion because VTCM needed (%d) > budget (%zu)\n",
kparams.vtcm_size, sess->vtcm_size); kparams.vtcm_size, sess->vtcm_size);
} }
@@ -4455,7 +4394,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
opt_opstage = str_opstage ? strtoul(str_opstage, NULL, 0) : opt_opstage; opt_opstage = str_opstage ? strtoul(str_opstage, NULL, 0) : opt_opstage;
opt_opbatch = str_opbatch ? strtoul(str_opbatch, NULL, 0) : opt_opbatch; opt_opbatch = str_opbatch ? strtoul(str_opbatch, NULL, 0) : opt_opbatch;
opt_opqueue = str_opqueue ? strtoul(str_opqueue, NULL, 0) : opt_opqueue; opt_opqueue = str_opqueue ? strtoul(str_opqueue, NULL, 0) : opt_opqueue;
opt_optrace = str_optrace ? strtoul(str_optrace, NULL, 0) : (opt_opbatch * 128); opt_optrace = str_optrace ? strtoul(str_optrace, NULL, 0) : (opt_opbatch * 256);
opt_oppoll = str_oppoll ? strtoul(str_oppoll, NULL, 0) : opt_oppoll; opt_oppoll = str_oppoll ? strtoul(str_oppoll, NULL, 0) : opt_oppoll;
opt_opfusion = str_opfusion ? atoi(str_opfusion) : opt_opfusion; opt_opfusion = str_opfusion ? atoi(str_opfusion) : opt_opfusion;
opt_profile = str_profile ? atoi(str_profile) : 0; opt_profile = str_profile ? atoi(str_profile) : 0;
ggml/src/ggml-hexagon/htp/dma-queue.c
+1 -3
View File
@@ -101,6 +101,4 @@ void dma_queue_alias_free(dma_queue_t q) {
(void) q; (void) q;
} }
void dma_queue_flush(dma_queue_t q) {
while (dma_queue_pop(q).dst != NULL) ;
}
ggml/src/ggml-hexagon/htp/dma-queue.h
+23 -12
View File
@@ -106,7 +106,7 @@ struct dma_queue_s {
bool alias; // When set, dma_queue_delete will not free the ring bool alias; // When set, dma_queue_delete will not free the ring
}; };
void dma_queue_flush(dma_queue_t q);
size_t dma_queue_sizeof(size_t capacity); size_t dma_queue_sizeof(size_t capacity);
size_t dma_queue_alignof(void); size_t dma_queue_alignof(void);
@@ -154,7 +154,6 @@ static inline bool dma_is_vtcm(const dma_queue * q, const void * ptr) {
static inline bool dma_queue_push_single_1d(dma_queue * q, dma_ptr dptr, size_t size) { static inline bool dma_queue_push_single_1d(dma_queue * q, dma_ptr dptr, size_t size) {
dma_ring * r = q->ring; dma_ring * r = q->ring;
if (((r->push_idx + 1) & r->idx_mask) == r->pop_idx) { if (((r->push_idx + 1) & r->idx_mask) == r->pop_idx) {
FARF(HIGH, "dma-push: queue full\n");
return false; return false;
} }
@@ -165,6 +164,8 @@ static inline bool dma_queue_push_single_1d(dma_queue * q, dma_ptr dptr, size_t
r->dptr[r->push_idx] = dptr; r->dptr[r->push_idx] = dptr;
htp_trace_event_start(r->trace, HTP_TRACE_EVT_DMA, r->push_idx);
if (size) { if (size) {
desc->next = NULL; desc->next = NULL;
desc->desc_size = 0; // 1D mode desc->desc_size = 0; // 1D mode
@@ -173,7 +174,6 @@ static inline bool dma_queue_push_single_1d(dma_queue * q, dma_ptr dptr, size_t
desc->order = 0; desc->order = 0;
desc->done = 0; desc->done = 0;
htp_trace_event_start(r->trace, HTP_TRACE_EVT_DMA, r->push_idx);
dmlink(r->tail, desc); dmlink(r->tail, desc);
r->tail = (dma_descriptor_2d *) desc; r->tail = (dma_descriptor_2d *) desc;
} else { } else {
@@ -188,7 +188,6 @@ static inline bool dma_queue_push_single_1d(dma_queue * q, dma_ptr dptr, size_t
static inline bool dma_queue_push_single_2d(dma_queue * q, dma_ptr dptr, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) { static inline bool dma_queue_push_single_2d(dma_queue * q, dma_ptr dptr, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
dma_ring * r = q->ring; dma_ring * r = q->ring;
if (((r->push_idx + 1) & r->idx_mask) == r->pop_idx) { if (((r->push_idx + 1) & r->idx_mask) == r->pop_idx) {
FARF(HIGH, "dma-push: queue full\n");
return false; return false;
} }
@@ -224,8 +223,9 @@ static inline bool dma_queue_push_single_2d(dma_queue * q, dma_ptr dptr, size_t
r->dptr[r->push_idx] = dptr; r->dptr[r->push_idx] = dptr;
htp_trace_event_start(r->trace, HTP_TRACE_EVT_DMA, r->push_idx);
if (nrows) { if (nrows) {
htp_trace_event_start(r->trace, HTP_TRACE_EVT_DMA, r->push_idx);
dmlink(r->tail, desc); dmlink(r->tail, desc);
r->tail = desc; r->tail = desc;
} else { } else {
@@ -252,10 +252,11 @@ static inline dma_ptr dma_queue_pop(dma_queue * q) {
dmpoll(); dmpoll();
} }
} }
htp_trace_event_stop(r->trace, HTP_TRACE_EVT_DMA, r->pop_idx);
dptr = r->dptr[r->pop_idx]; dptr = r->dptr[r->pop_idx];
htp_trace_event_stop(r->trace, HTP_TRACE_EVT_DMA, r->pop_idx);
r->pop_idx = (r->pop_idx + 1) & r->idx_mask; r->pop_idx = (r->pop_idx + 1) & r->idx_mask;
return dptr; return dptr;
} }
@@ -270,6 +271,8 @@ static inline dma_ptr dma_queue_pop_nowait(dma_queue * q) {
dptr = r->dptr[r->pop_idx]; dptr = r->dptr[r->pop_idx];
htp_trace_event_stop(r->trace, HTP_TRACE_EVT_DMA, r->pop_idx);
r->pop_idx = (r->pop_idx + 1) & r->idx_mask; r->pop_idx = (r->pop_idx + 1) & r->idx_mask;
return dptr; return dptr;
} }
@@ -278,6 +281,10 @@ static inline bool dma_queue_empty(dma_queue * q) {
return q->ring->push_idx == q->ring->pop_idx; return q->ring->push_idx == q->ring->pop_idx;
} }
static inline void dma_queue_flush(dma_queue * q) {
while (dma_queue_pop(q).dst != NULL) ;
}
static inline uint32_t dma_queue_depth(dma_queue * q) { static inline uint32_t dma_queue_depth(dma_queue * q) {
return (q->ring->push_idx - q->ring->pop_idx) & q->ring->idx_mask; return (q->ring->push_idx - q->ring->pop_idx) & q->ring->idx_mask;
} }
@@ -314,14 +321,18 @@ static inline bool dma_queue_push(dma_queue *q, dma_ptr dptr, size_t dst_stride,
{ {
const uint8_t *src = (const uint8_t *) dptr.src; const uint8_t *src = (const uint8_t *) dptr.src;
uint8_t *dst = (uint8_t *) dptr.dst; uint8_t *dst = (uint8_t *) dptr.dst;
for (size_t r = 0; r < nrows; ++r) { size_t r = 0;
while (r + 1 < nrows) {
dma_ptr p = dma_make_ptr(dst + r * dst_stride, src + r * src_stride); dma_ptr p = dma_make_ptr(dst + r * dst_stride, src + r * src_stride);
if (!dma_queue_push_single_1d(q, p, row_size)) if (!dma_queue_push_single_1d(q, p, row_size)) {
return false; dma_queue_flush(q);
if (r + 1 < nrows) } else {
dma_queue_pop(q); r++;
}
} }
return true; dma_queue_flush(q);
dma_ptr p = dma_make_ptr(dst + r * dst_stride, src + r * src_stride);
return dma_queue_push_single_1d(q, p, row_size);
} }
} }
ggml/src/ggml-hexagon/htp/flash-attn-ops.c
+515 -179
View File
@@ -123,15 +123,17 @@ struct hmx_fa_context {
uint32_t g_br; // hex_align_up(G * Br, 32) - actual tile row dim uint32_t g_br; // hex_align_up(G * Br, 32) - actual tile row dim
// VTCM buffers (allocated by vtcm_seq_alloc) // VTCM buffers (allocated by vtcm_seq_alloc)
__fp16 * vtcm_q_dma; // Q DMA fetch buffer
__fp16 * vtcm_q_tiles; // Q tile format [g_br, D] __fp16 * vtcm_q_tiles; // Q tile format [g_br, D]
__fp16 * vtcm_o_tiles[2]; // O ping-pong [g_br, D] __fp16 * vtcm_o_tiles[2]; // O ping-pong [g_br, D]
__fp16 * vtcm_k_fp16[2]; // K DMA double-buffer [Bc, D] __fp16 * vtcm_k_fp16[2]; // K DMA double-buffer [Bc, D]
__fp16 * vtcm_v_fp16[2]; // V DMA double-buffer [Bc, D] __fp16 * vtcm_v_fp16[2]; // V DMA double-buffer [Bc, D]
__fp16 * vtcm_k_tiles; // K tiles (transposed) __fp16 * vtcm_k_tiles[2]; // K tiles (transposed, double-buffered)
__fp16 * vtcm_v_tiles[2]; // V tiles (column-major, double-buffered) __fp16 * vtcm_v_tiles[2]; // V tiles (column-major, double-buffered)
__fp16 * vtcm_s_tiles; // S = QK^T [g_br, Bc] __fp16 * vtcm_s_tiles[2]; // S = QK^T [g_br, Bc] (double-buffered)
__fp16 * vtcm_p_tiles; // P = softmax(S) [g_br, Bc] __fp16 * vtcm_p_tiles[2]; // P = softmax(S) [g_br, Bc]
__fp16 * vtcm_d_tiles; // Diagonal rescale [g_br, g_br] __fp16 * vtcm_d_tiles; // Diagonal rescale [g_br, g_br]
__fp16 * vtcm_d_inv_l; // Diagonal rescale (1/l) [g_br, g_br]
HVX_Vector * vtcm_m_vec; // Row max [g_br] HVX_Vector * vtcm_m_vec; // Row max [g_br]
HVX_Vector * vtcm_l_vec; // Row sum [g_br] HVX_Vector * vtcm_l_vec; // Row sum [g_br]
HVX_Vector * vtcm_s_rowmax; // Softmax intermediate [g_br] HVX_Vector * vtcm_s_rowmax; // Softmax intermediate [g_br]
@@ -236,10 +238,6 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
const uint32_t iv3 = fastdiv(iq3, &factx->broadcast_rv3); const uint32_t iv3 = fastdiv(iq3, &factx->broadcast_rv3);
const uint32_t iv2 = fastdiv(iq2, &factx->broadcast_rv2); const uint32_t iv2 = fastdiv(iq2, &factx->broadcast_rv2);
// Fetch Q row
const uint8_t * q_row_ptr = (const uint8_t *) q->data + (iq1*nbq1 + iq2*nbq2 + iq3*nbq3);
dma_queue_push(dma, dma_make_ptr(spad_q, q_row_ptr), factx->size_q_row_padded, nbq1, size_q_row, 1);
const __fp16 * mp_base = NULL; const __fp16 * mp_base = NULL;
if (mask) { if (mask) {
const uint32_t im2 = fastmodulo(iq2, mask->ne[2], &factx->src3_div2); const uint32_t im2 = fastmodulo(iq2, mask->ne[2], &factx->src3_div2);
@@ -247,26 +245,91 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
mp_base = (const __fp16 *) ((const uint8_t *) mask->data + iq1*mask->nb[1] + im2*mask->nb[2] + im3*mask->nb[3]); mp_base = (const __fp16 *) ((const uint8_t *) mask->data + iq1*mask->nb[1] + im2*mask->nb[2] + im3*mask->nb[3]);
} }
// Prefetch first two blocks // Precalculate next row variables if there is a next row
for (uint32_t ib = 0; ib < MIN(factx->n_blocks, 2); ++ib) { bool has_next_ir = (ir + 1 < ir1);
const uint32_t ic_start = ib * FLASH_ATTN_BLOCK_SIZE; uint32_t next_ik2 = 0, next_ik3 = 0, next_iv2 = 0, next_iv3 = 0;
const uint32_t current_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start); const uint8_t * next_q_row_ptr = NULL;
const __fp16 * next_mp_base = NULL;
// K const uint8_t * next_k_src0 = NULL;
const uint8_t * k_src = (const uint8_t *) k->data + (ic_start*nbk1 + ik2*nbk2 + ik3*nbk3); const uint8_t * next_v_src0 = NULL;
uint8_t * k_dst = spad_k + (ib % 2) * factx->size_k_block; const uint8_t * next_m_src0 = NULL;
dma_queue_push(dma, dma_make_ptr(k_dst, k_src), factx->size_k_row_padded, nbk1, size_k_row, current_block_size); uint32_t next_block_size0 = 0;
// V const uint8_t * next_k_src1 = NULL;
const uint8_t * v_src = (const uint8_t *) v->data + (ic_start*nbv1 + iv2*nbv2 + iv3*nbv3); const uint8_t * next_v_src1 = NULL;
uint8_t * v_dst = spad_v + (ib % 2) * factx->size_v_block; const uint8_t * next_m_src1 = NULL;
dma_queue_push(dma, dma_make_ptr(v_dst, v_src), factx->size_v_row_padded, nbv1, size_v_row, current_block_size); uint32_t next_block_size1 = 0;
if (has_next_ir) {
const uint32_t next_ir = ir + 1;
const uint32_t next_iq3 = fastdiv(next_ir, &factx->src0_div21);
const uint32_t next_iq2 = fastdiv(next_ir - next_iq3*neq2*neq1, &factx->src0_div1);
const uint32_t next_iq1 = (next_ir - next_iq3*neq2*neq1 - next_iq2 * neq1);
next_ik3 = fastdiv(next_iq3, &factx->broadcast_rk3);
next_ik2 = fastdiv(next_iq2, &factx->broadcast_rk2);
next_iv3 = fastdiv(next_iq3, &factx->broadcast_rv3);
next_iv2 = fastdiv(next_iq2, &factx->broadcast_rv2);
next_q_row_ptr = (const uint8_t *) q->data + (next_iq1*nbq1 + next_iq2*nbq2 + next_iq3*nbq3);
// Mask
if (mask) { if (mask) {
const uint8_t * m_src = (const uint8_t *) (mp_base + ic_start); const uint32_t next_im2 = fastmodulo(next_iq2, mask->ne[2], &factx->src3_div2);
// Mask is 1D contiguous for this row const uint32_t next_im3 = fastmodulo(next_iq3, mask->ne[3], &factx->src3_div3);
dma_cache_push(dma, &m_cache, m_src, current_block_size * 2, current_block_size * 2, current_block_size * 2, 1); next_mp_base = (const __fp16 *) ((const uint8_t *) mask->data + next_iq1*mask->nb[1] + next_im2*mask->nb[2] + next_im3*mask->nb[3]);
}
// Precalculate next K/V block 0 source pointers
{
const uint32_t ic_start = 0;
next_block_size0 = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start);
next_k_src0 = (const uint8_t *) k->data + (ic_start*nbk1 + next_ik2*nbk2 + next_ik3*nbk3);
next_v_src0 = (const uint8_t *) v->data + (ic_start*nbv1 + next_iv2*nbv2 + next_iv3*nbv3);
if (mask) {
next_m_src0 = (const uint8_t *) (next_mp_base + ic_start);
}
}
// Precalculate next K/V block 1 source pointers (if n_blocks > 1)
if (factx->n_blocks > 1) {
const uint32_t ic_start = 1 * FLASH_ATTN_BLOCK_SIZE;
next_block_size1 = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start);
next_k_src1 = (const uint8_t *) k->data + (ic_start*nbk1 + next_ik2*nbk2 + next_ik3*nbk3);
next_v_src1 = (const uint8_t *) v->data + (ic_start*nbv1 + next_iv2*nbv2 + next_iv3*nbv3);
if (mask) {
next_m_src1 = (const uint8_t *) (next_mp_base + ic_start);
}
}
}
if (ir == ir0) {
// Fetch Q row
const uint8_t * q_row_ptr = (const uint8_t *) q->data + (iq1*nbq1 + iq2*nbq2 + iq3*nbq3);
dma_queue_push(dma, dma_make_ptr(spad_q, q_row_ptr), factx->size_q_row_padded, nbq1, size_q_row, 1);
// Prefetch first two blocks
for (uint32_t ib = 0; ib < MIN(factx->n_blocks, 2); ++ib) {
const uint32_t ic_start = ib * FLASH_ATTN_BLOCK_SIZE;
const uint32_t current_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start);
// K
const uint8_t * k_src = (const uint8_t *) k->data + (ic_start*nbk1 + ik2*nbk2 + ik3*nbk3);
uint8_t * k_dst = spad_k + (ib % 2) * factx->size_k_block;
dma_queue_push(dma, dma_make_ptr(k_dst, k_src), factx->size_k_row_padded, nbk1, size_k_row, current_block_size);
// V
const uint8_t * v_src = (const uint8_t *) v->data + (ic_start*nbv1 + iv2*nbv2 + iv3*nbv3);
uint8_t * v_dst = spad_v + (ib % 2) * factx->size_v_block;
dma_queue_push(dma, dma_make_ptr(v_dst, v_src), factx->size_v_row_padded, nbv1, size_v_row, current_block_size);
// Mask
if (mask) {
const uint8_t * m_src = (const uint8_t *) (mp_base + ic_start);
// Mask is 1D contiguous for this row
dma_cache_push(dma, &m_cache, m_src, current_block_size * 2, current_block_size * 2, current_block_size * 2, 1);
}
} }
} }
@@ -287,6 +350,11 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
const HVX_Vector slope_vec = hvx_vec_splat_f16(slope); const HVX_Vector slope_vec = hvx_vec_splat_f16(slope);
const HVX_Vector v_neg_inf = Q6_Vh_vsplat_R(0xfbff); const HVX_Vector v_neg_inf = Q6_Vh_vsplat_R(0xfbff);
const HVX_Vector v_cap = (factx->logit_softcap != 0.0f) ? hvx_vec_splat_f16(factx->logit_softcap) : Q6_V_vzero();
const HVX_Vector vinf = Q6_Vh_vsplat_R(0xFC00);
const HVX_Vector vmin = Q6_Vh_vsplat_R(0xFBFF);
const HVX_Vector v_log2e = hvx_vec_splat_f16(EXP_LOG2E_F);
const uint32_t stride_v2 = factx->size_v_row_padded * 2;
for (uint32_t ib = 0; ib < factx->n_blocks; ++ib) { for (uint32_t ib = 0; ib < factx->n_blocks; ++ib) {
const uint32_t ic_start = ib * FLASH_ATTN_BLOCK_SIZE; const uint32_t ic_start = ib * FLASH_ATTN_BLOCK_SIZE;
const uint32_t current_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start); const uint32_t current_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start);
@@ -309,7 +377,6 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
// 2. Softcap (in FP16) // 2. Softcap (in FP16)
if (factx->logit_softcap != 0.0f) { if (factx->logit_softcap != 0.0f) {
const HVX_Vector v_cap = hvx_vec_splat_f16(factx->logit_softcap);
scores_f16 = hvx_vec_tanh_f16(scores_f16); scores_f16 = hvx_vec_tanh_f16(scores_f16);
scores_f16 = hvx_vec_mul_f16_f16(scores_f16, v_cap); scores_f16 = hvx_vec_mul_f16_f16(scores_f16, v_cap);
} }
@@ -319,8 +386,6 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
// 3. Mask (in FP16) // 3. Mask (in FP16)
if (mask) { if (mask) {
HVX_Vector m_vals_f16 = *(const HVX_UVector *) m_base; HVX_Vector m_vals_f16 = *(const HVX_UVector *) m_base;
HVX_Vector vinf = Q6_Vh_vsplat_R(0xFC00);
HVX_Vector vmin = Q6_Vh_vsplat_R(0xFBFF);
HVX_VectorPred is_inf = Q6_Q_vcmp_eq_VhVh(m_vals_f16, vinf); HVX_VectorPred is_inf = Q6_Q_vcmp_eq_VhVh(m_vals_f16, vinf);
m_vals_f16 = Q6_V_vmux_QVV(is_inf, vmin, m_vals_f16); m_vals_f16 = Q6_V_vmux_QVV(is_inf, vmin, m_vals_f16);
@@ -335,10 +400,30 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
HVX_Vector v_max = Q6_V_lo_W(hvx_vec_f16_to_f32(v_max_f16)); // splat block max in FP32 HVX_Vector v_max = Q6_V_lo_W(hvx_vec_f16_to_f32(v_max_f16)); // splat block max in FP32
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_QK, ir); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_QK, ir);
if (ib + 1 == factx->n_blocks && has_next_ir) {
// Queue next row's Q row!
dma_queue_push(dma, dma_make_ptr(spad_q, next_q_row_ptr), factx->size_q_row_padded, nbq1, size_q_row, 1);
if (factx->n_blocks % 2 == 0) {
// Queue next row's block 0 (into buffer slot 0)
uint8_t * k_dst = spad_k + 0 * factx->size_k_block;
uint8_t * v_dst = spad_v + 0 * factx->size_v_block;
// K (block 0 of next row)
dma_queue_push(dma, dma_make_ptr(k_dst, next_k_src0), factx->size_k_row_padded, nbk1, size_k_row, next_block_size0);
// V (block 0 of next row)
dma_queue_push(dma, dma_make_ptr(v_dst, next_v_src0), factx->size_v_row_padded, nbv1, size_v_row, next_block_size0);
// Mask (block 0 of next row)
if (mask) {
dma_cache_push(dma, &m_cache, next_m_src0, next_block_size0 * 2, next_block_size0 * 2, next_block_size0 * 2, 1);
}
}
}
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_SFM, ir); htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_SFM, ir);
{ {
const HVX_Vector v_log2e = hvx_vec_splat_f16(EXP_LOG2E_F);
// 4. Online Softmax Update // 4. Online Softmax Update
HVX_Vector M_new_vec = Q6_Vsf_vmax_VsfVsf(v_max, M_vec); HVX_Vector M_new_vec = Q6_Vsf_vmax_VsfVsf(v_max, M_vec);
HVX_Vector diff_vec = HVX_OP_SUB_F32(M_vec, M_new_vec); HVX_Vector diff_vec = HVX_OP_SUB_F32(M_vec, M_new_vec);
@@ -370,24 +455,20 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
S_vec = HVX_OP_ADD_F32(HVX_OP_MUL_F32(S_vec, ms_vec), p_sum_vec); S_vec = HVX_OP_ADD_F32(HVX_OP_MUL_F32(S_vec, ms_vec), p_sum_vec);
// 5. Accumulate V (F16 * F16 -> F32 accumulator) // 5. Accumulate V (F16 * F16 -> F32 accumulator)
__fp16 __attribute__((aligned(128))) p_arr[VLEN_FP16]; const uint8_t * v_ptr = v_base;
hvx_vec_store_a(p_arr, 128, P);
for (uint32_t j = 0; j < current_block_size; j += 2) { for (uint32_t j = 0; j < current_block_size; j += 2) {
if (j + 1 == current_block_size) { if (j + 1 == current_block_size) {
if (p_arr[j] != 0.0f) { HVX_Vector S0 = hvx_vec_repl_f16(Q6_V_vror_VR(P, j * 2));
const uint8_t * v_ptr = v_base + j * factx->size_v_row_padded; hvx_mad_f32_f16_aa_vec(VKQ32, v_ptr, S0, DV);
hvx_mad_f32_f16_aa(VKQ32, v_ptr, (p_arr + j), DV);
}
break; break;
} }
if (p_arr[j] == 0.0f && p_arr[j + 1] == 0.0f) { HVX_Vector S0 = hvx_vec_repl_f16(Q6_V_vror_VR(P, j * 2));
continue; HVX_Vector S1 = hvx_vec_repl_f16(Q6_V_vror_VR(P, (j + 1) * 2));
}
const uint8_t * v_ptr = v_base + j * factx->size_v_row_padded; hvx_mad_f32_f16_aa_rx2_vec(VKQ32, v_ptr, v_ptr + factx->size_v_row_padded, S0, S1, DV);
hvx_mad_f32_f16_aa_rx2(VKQ32, v_ptr, v_ptr + factx->size_v_row_padded, (p_arr + j), (p_arr + j + 1), DV); v_ptr += stride_v2;
} }
} }
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_SFM, ir); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_SFM, ir);
@@ -414,6 +495,61 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
} }
} }
if (has_next_ir) {
if (factx->n_blocks % 2 == 0) {
// Queue next row's block 1 (into buffer slot 1, if n_blocks > 1)
if (factx->n_blocks > 1) {
uint8_t * k_dst = spad_k + 1 * factx->size_k_block;
uint8_t * v_dst = spad_v + 1 * factx->size_v_block;
// K (block 1 of next row)
dma_queue_push(dma, dma_make_ptr(k_dst, next_k_src1), factx->size_k_row_padded, nbk1, size_k_row, next_block_size1);
// V (block 1 of next row)
dma_queue_push(dma, dma_make_ptr(v_dst, next_v_src1), factx->size_v_row_padded, nbv1, size_v_row, next_block_size1);
// Mask (block 1 of next row)
if (mask) {
dma_cache_push(dma, &m_cache, next_m_src1, next_block_size1 * 2, next_block_size1 * 2, next_block_size1 * 2, 1);
}
}
} else {
// Queue next row's block 0 (into buffer slot 0)
{
uint8_t * k_dst = spad_k + 0 * factx->size_k_block;
uint8_t * v_dst = spad_v + 0 * factx->size_v_block;
// K (block 0 of next row)
dma_queue_push(dma, dma_make_ptr(k_dst, next_k_src0), factx->size_k_row_padded, nbk1, size_k_row, next_block_size0);
// V (block 0 of next row)
dma_queue_push(dma, dma_make_ptr(v_dst, next_v_src0), factx->size_v_row_padded, nbv1, size_v_row, next_block_size0);
// Mask (block 0 of next row)
if (mask) {
dma_cache_push(dma, &m_cache, next_m_src0, next_block_size0 * 2, next_block_size0 * 2, next_block_size0 * 2, 1);
}
}
// Queue next row's block 1 (into buffer slot 1, if n_blocks > 1)
if (factx->n_blocks > 1) {
uint8_t * k_dst = spad_k + 1 * factx->size_k_block;
uint8_t * v_dst = spad_v + 1 * factx->size_v_block;
// K (block 1 of next row)
dma_queue_push(dma, dma_make_ptr(k_dst, next_k_src1), factx->size_k_row_padded, nbk1, size_k_row, next_block_size1);
// V (block 1 of next row)
dma_queue_push(dma, dma_make_ptr(v_dst, next_v_src1), factx->size_v_row_padded, nbv1, size_v_row, next_block_size1);
// Mask (block 1 of next row)
if (mask) {
dma_cache_push(dma, &m_cache, next_m_src1, next_block_size1 * 2, next_block_size1 * 2, next_block_size1 * 2, 1);
}
}
}
}
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, ir); htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, ir);
// sinks // sinks
float M = hvx_vec_get_f32(M_vec); float M = hvx_vec_get_f32(M_vec);
@@ -471,6 +607,7 @@ typedef struct {
void * curr_k; void * curr_k;
uint32_t kv_start; uint32_t kv_start;
uint32_t rows_per_t; uint32_t rows_per_t;
size_t buf_idx;
} fa_k_int_args_t; } fa_k_int_args_t;
static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data) { static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data) {
@@ -488,19 +625,19 @@ static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data)
struct htp_thread_trace * tr = &factx->octx->ctx->trace[i]; struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start)); htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start));
hmx_interleave_rows_to_tiles(factx->vtcm_k_tiles, (const __fp16 *) args->curr_k, total_rows, factx->DK, hmx_interleave_rows_to_tiles(factx->vtcm_k_tiles[args->buf_idx], (const __fp16 *) args->curr_k, total_rows, factx->DK,
args->src_stride, start, end); args->src_stride, start, end);
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start)); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start));
} }
static void fa_phase_k_interleave(struct hmx_fa_context * factx, uint32_t kv_rows, size_t src_stride, void * curr_k, uint32_t kv_start) { static void fa_phase_k_interleave(struct hmx_fa_context * factx, uint32_t kv_rows, size_t src_stride, void * curr_k, uint32_t kv_start, size_t buf_idx) {
work_queue_t wp = factx->octx->ctx->work_queue; work_queue_t wp = factx->octx->ctx->work_queue;
uint32_t n = 1; uint32_t n = 1;
if (factx->n_threads > 1 && kv_rows >= factx->n_threads * 2) { if (factx->n_threads > 1 && kv_rows >= factx->n_threads * 2) {
n = factx->n_threads; n = factx->n_threads;
} }
uint32_t rows_per_t = hex_align_up(hmx_ceil_div(kv_rows, n), 2); uint32_t rows_per_t = hex_align_up(hmx_ceil_div(kv_rows, n), 2);
fa_k_int_args_t args = { factx, kv_rows, src_stride, curr_k, kv_start, rows_per_t }; fa_k_int_args_t args = { factx, kv_rows, src_stride, curr_k, kv_start, rows_per_t, buf_idx };
if (n > 1) { if (n > 1) {
work_queue_run(wp, fa_k_interleave_thread, &args, n); work_queue_run(wp, fa_k_interleave_thread, &args, n);
} else { } else {
@@ -645,12 +782,13 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) {
} }
} }
// Initialize vtcm_d_tiles to 0 // Initialize vtcm_d_tiles and vtcm_d_inv_l to 0
const size_t d_bytes_per_t = hex_align_up(d_tile_bytes / n, 128); const size_t d_bytes_per_t = hex_align_up(d_tile_bytes / n, 128);
const size_t d_start = i * d_bytes_per_t; const size_t d_start = i * d_bytes_per_t;
const size_t d_end = hex_smin(d_start + d_bytes_per_t, d_tile_bytes); const size_t d_end = hex_smin(d_start + d_bytes_per_t, d_tile_bytes);
if (d_start < d_tile_bytes) { if (d_start < d_tile_bytes) {
hvx_splat_u8_a((char *) factx->vtcm_d_tiles + d_start, 0, d_end - d_start); hvx_splat_u8_a((char *) factx->vtcm_d_tiles + d_start, 0, d_end - d_start);
hvx_splat_u8_a((char *) factx->vtcm_d_inv_l + d_start, 0, d_end - d_start);
} }
} }
@@ -662,15 +800,14 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) {
assert(factx->DK == factx->DV); assert(factx->DK == factx->DV);
const size_t o_tile_bytes = factx->o_tile_bytes; const bool use_q_dma = (factx->vtcm_q_dma != NULL);
const bool use_q_dma = (2 * o_tile_bytes >= factx->g_br * DK * (factx->is_q_fp32 ? 4 : 2));
__fp16 * q_tiles = factx->vtcm_q_tiles; __fp16 * q_tiles = factx->vtcm_q_tiles;
if (use_q_dma) { if (use_q_dma) {
const size_t g_rows_end = hex_smin(end, n_rows_g); const size_t g_rows_end = hex_smin(end, n_rows_g);
const uint32_t d_limit = factx->is_q_fp32 ? DK / 32 : DK / 64; const uint32_t d_limit = factx->is_q_fp32 ? DK / 32 : DK / 64;
uint8_t * q_flat = (uint8_t *) factx->vtcm_o_tiles[0]; uint8_t * q_flat = (uint8_t *) factx->vtcm_q_dma;
if (factx->is_q_fp32) { if (factx->is_q_fp32) {
switch (d_limit) { switch (d_limit) {
case 2: hmx_fa_q_prep_fp32_d2(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); break; case 2: hmx_fa_q_prep_fp32_d2(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); break;
@@ -781,10 +918,10 @@ static void fa_o_store_thread_f32(unsigned int n, unsigned int i, void * data) {
const uint32_t kv_head = args->kv_head; const uint32_t kv_head = args->kv_head;
const uint32_t ib3 = args->ib3; const uint32_t ib3 = args->ib3;
for (size_t r = start; r < end; ++r) { size_t q_idx = fastdiv(start, &factx->div_G);
const size_t q_idx = fastdiv(r, &factx->div_G); size_t h_idx = fastmodulo(start, G, &factx->div_G);
const size_t h_idx = fastmodulo(r, G, &factx->div_G);
for (size_t r = start; r < end; ++r) {
float * out = (float *) ((uint8_t *) dst->data + (kv_head * G + h_idx) * dst->nb[1] + float * out = (float *) ((uint8_t *) dst->data + (kv_head * G + h_idx) * dst->nb[1] +
(q_start + q_idx) * dst->nb[2] + ib3 * dst->nb[3]); (q_start + q_idx) * dst->nb[2] + ib3 * dst->nb[3]);
@@ -801,6 +938,12 @@ static void fa_o_store_thread_f32(unsigned int n, unsigned int i, void * data) {
*(HVX_UVector *) (out + d * 32) = Q6_V_hi_W(vp); *(HVX_UVector *) (out + d * 32) = Q6_V_hi_W(vp);
} }
} }
h_idx++;
if (h_idx == G) {
h_idx = 0;
q_idx++;
}
} }
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) (args->q_start * G + start)); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) (args->q_start * G + start));
} }
@@ -829,10 +972,10 @@ static void fa_o_store_thread_f16(unsigned int n, unsigned int i, void * data) {
const uint32_t kv_head = args->kv_head; const uint32_t kv_head = args->kv_head;
const uint32_t ib3 = args->ib3; const uint32_t ib3 = args->ib3;
for (size_t r = start; r < end; ++r) { size_t q_idx = fastdiv(start, &factx->div_G);
const size_t q_idx = fastdiv(r, &factx->div_G); size_t h_idx = fastmodulo(start, G, &factx->div_G);
const size_t h_idx = fastmodulo(r, G, &factx->div_G);
for (size_t r = start; r < end; ++r) {
__fp16 * out = (__fp16 *) ((uint8_t *) dst->data + (kv_head * G + h_idx) * dst->nb[1] + __fp16 * out = (__fp16 *) ((uint8_t *) dst->data + (kv_head * G + h_idx) * dst->nb[1] +
(q_start + q_idx) * dst->nb[2] + ib3 * dst->nb[3]); (q_start + q_idx) * dst->nb[2] + ib3 * dst->nb[3]);
@@ -851,6 +994,12 @@ static void fa_o_store_thread_f16(unsigned int n, unsigned int i, void * data) {
*(HVX_UVector *) (out + d * 64) = Q6_V_hi_W(vp); *(HVX_UVector *) (out + d * 64) = Q6_V_hi_W(vp);
} }
} }
h_idx++;
if (h_idx == G) {
h_idx = 0;
q_idx++;
}
} }
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) (args->q_start * G + start)); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) (args->q_start * G + start));
} }
@@ -879,6 +1028,7 @@ static void fa_phase_o_store(struct hmx_fa_context * factx,
typedef struct { typedef struct {
struct hmx_fa_context * factx; struct hmx_fa_context * factx;
size_t buf_idx;
size_t kv_rows; size_t kv_rows;
size_t n_rows_g; size_t n_rows_g;
size_t n_col_tiles; size_t n_col_tiles;
@@ -960,8 +1110,8 @@ static inline void fa_softmax_impl(
uint32_t r0 = r / HMX_FP16_TILE_N_ROWS; uint32_t r0 = r / HMX_FP16_TILE_N_ROWS;
uint32_t r1 = r % HMX_FP16_TILE_N_ROWS; uint32_t r1 = r % HMX_FP16_TILE_N_ROWS;
const __fp16 * s_ld_base = factx->vtcm_s_tiles + r0 * HMX_FP16_TILE_N_ROWS * Bc; const __fp16 * s_ld_base = factx->vtcm_s_tiles[args->buf_idx] + r0 * HMX_FP16_TILE_N_ROWS * Bc;
__fp16 * p_st_base = factx->vtcm_p_tiles + r0 * HMX_FP16_TILE_N_ROWS * Bc; __fp16 * p_st_base = factx->vtcm_p_tiles[args->buf_idx] + r0 * HMX_FP16_TILE_N_ROWS * Bc;
// Decode 2 rows from S tiles into per-thread row buffers // Decode 2 rows from S tiles into per-thread row buffers
if (has_softcap) { if (has_softcap) {
@@ -983,7 +1133,26 @@ static inline void fa_softmax_impl(
my_row_buf1[ci] = hvx_vec_mul_f16_f16(t1, v_cap); my_row_buf1[ci] = hvx_vec_mul_f16_f16(t1, v_cap);
} }
} else { } else {
for (size_t c = 0; c < kv_rows; c += 64) { size_t c = 0;
for (; c + 64 < kv_rows; c += 128) {
size_t ci0 = c / 64;
size_t ci1 = ci0 + 1;
const __fp16 * in_dtile0 = s_ld_base + ci0 * HMX_FP16_TILE_N_ELMS * 2;
const __fp16 * in_dtile1 = s_ld_base + ci1 * HMX_FP16_TILE_N_ELMS * 2;
const HVX_Vector * pv_s_in0_0 = ((const HVX_Vector *) in_dtile0) + r1 / 2;
const HVX_Vector * pv_s_in1_0 = pv_s_in0_0 + 16;
const HVX_Vector * pv_s_in0_1 = ((const HVX_Vector *) in_dtile1) + r1 / 2;
const HVX_Vector * pv_s_in1_1 = pv_s_in0_1 + 16;
HVX_VectorPair vp_s_drow0 = Q6_W_vdeal_VVR(*pv_s_in1_0, *pv_s_in0_0, -2);
my_row_buf0[ci0] = Q6_V_lo_W(vp_s_drow0);
my_row_buf1[ci0] = Q6_V_hi_W(vp_s_drow0);
HVX_VectorPair vp_s_drow1 = Q6_W_vdeal_VVR(*pv_s_in1_1, *pv_s_in0_1, -2);
my_row_buf0[ci1] = Q6_V_lo_W(vp_s_drow1);
my_row_buf1[ci1] = Q6_V_hi_W(vp_s_drow1);
}
for (; c < kv_rows; c += 64) {
size_t ci = c / 64; size_t ci = c / 64;
const __fp16 * in_dtile = s_ld_base + ci * HMX_FP16_TILE_N_ELMS * 2; const __fp16 * in_dtile = s_ld_base + ci * HMX_FP16_TILE_N_ELMS * 2;
const HVX_Vector * pv_s_in0 = ((const HVX_Vector *) in_dtile) + r1 / 2; const HVX_Vector * pv_s_in0 = ((const HVX_Vector *) in_dtile) + r1 / 2;
@@ -1007,12 +1176,12 @@ static inline void fa_softmax_impl(
HVX_Vector v_s_rowmax0 = v_neg_inf; HVX_Vector v_s_rowmax0 = v_neg_inf;
HVX_Vector v_s_rowmax1 = v_neg_inf; HVX_Vector v_s_rowmax1 = v_neg_inf;
for (size_t c = 0; c < kv_rows; c += 64) { if (has_mask) {
size_t ci = c / 64; for (size_t c = 0; c < kv_rows; c += 64) {
const size_t ne = hex_smin(kv_rows - c, 64); size_t ci = c / 64;
HVX_VectorPred q_tail_keep = Q6_Q_vsetq2_R(ne * sizeof(__fp16)); const size_t ne = hex_smin(kv_rows - c, 64);
HVX_VectorPred q_tail_keep = Q6_Q_vsetq2_R(ne * sizeof(__fp16));
if (has_mask) {
HVX_Vector v_mask0, v_mask1; HVX_Vector v_mask0, v_mask1;
if (mask_broadcast) { if (mask_broadcast) {
@@ -1066,15 +1235,31 @@ static inline void fa_softmax_impl(
my_row_buf0[ci] = Q6_V_vmux_QVV(q_keep0, hvx_vec_add_f16_f16(my_row_buf0[ci], v_mask0_scaled), v_neg_inf); my_row_buf0[ci] = Q6_V_vmux_QVV(q_keep0, hvx_vec_add_f16_f16(my_row_buf0[ci], v_mask0_scaled), v_neg_inf);
my_row_buf1[ci] = Q6_V_vmux_QVV(q_keep1, hvx_vec_add_f16_f16(my_row_buf1[ci], v_mask1_scaled), v_neg_inf); my_row_buf1[ci] = Q6_V_vmux_QVV(q_keep1, hvx_vec_add_f16_f16(my_row_buf1[ci], v_mask1_scaled), v_neg_inf);
} }
} else {
v_s_rowmax0 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax0, my_row_buf0[ci]);
v_s_rowmax1 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax1, my_row_buf1[ci]);
}
} else {
size_t c = 0;
for (; c + 64 < kv_rows; c += 128) {
size_t ci0 = c / 64;
size_t ci1 = ci0 + 1;
v_s_rowmax0 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax0, my_row_buf0[ci0]);
v_s_rowmax1 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax1, my_row_buf1[ci0]);
v_s_rowmax0 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax0, my_row_buf0[ci1]);
v_s_rowmax1 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax1, my_row_buf1[ci1]);
}
for (; c < kv_rows; c += 64) {
size_t ci = c / 64;
const size_t ne = hex_smin(kv_rows - c, 64);
HVX_VectorPred q_tail_keep = Q6_Q_vsetq2_R(ne * sizeof(__fp16));
if (ne < 64) { if (ne < 64) {
my_row_buf0[ci] = Q6_V_vmux_QVV(q_tail_keep, my_row_buf0[ci], v_neg_inf); my_row_buf0[ci] = Q6_V_vmux_QVV(q_tail_keep, my_row_buf0[ci], v_neg_inf);
my_row_buf1[ci] = Q6_V_vmux_QVV(q_tail_keep, my_row_buf1[ci], v_neg_inf); my_row_buf1[ci] = Q6_V_vmux_QVV(q_tail_keep, my_row_buf1[ci], v_neg_inf);
} }
v_s_rowmax0 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax0, my_row_buf0[ci]);
v_s_rowmax1 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax1, my_row_buf1[ci]);
} }
v_s_rowmax0 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax0, my_row_buf0[ci]);
v_s_rowmax1 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax1, my_row_buf1[ci]);
} }
v_s_rowmax0 = hvx_vec_reduce_max_f16(v_s_rowmax0); v_s_rowmax0 = hvx_vec_reduce_max_f16(v_s_rowmax0);
@@ -1121,8 +1306,48 @@ static inline void fa_softmax_impl(
HVX_Vector v_p_rowsum0 = v_zero; HVX_Vector v_p_rowsum0 = v_zero;
HVX_Vector v_p_rowsum1 = v_zero; HVX_Vector v_p_rowsum1 = v_zero;
for (size_t c = 0; c < kv_rows; c += 64) { size_t c = 0;
size_t ci = c / 64; for (; c + 64 < kv_rows; c += 128) {
size_t ci0 = c / 64;
size_t ci1 = ci0 + 1;
HVX_Vector v_s_minus_m0_0 = Q6_Vqf16_vsub_VhfVhf(my_row_buf0[ci0], v_dup_m0);
HVX_Vector v_s_minus_m1_0 = Q6_Vqf16_vsub_VhfVhf(my_row_buf1[ci0], v_dup_m1);
HVX_Vector v_s_minus_m0_1 = Q6_Vqf16_vsub_VhfVhf(my_row_buf0[ci1], v_dup_m0);
HVX_Vector v_s_minus_m1_1 = Q6_Vqf16_vsub_VhfVhf(my_row_buf1[ci1], v_dup_m1);
HVX_Vector v_p_row0_hf_0 = hvx_vec_exp2_f16(Q6_Vhf_equals_Vqf16(v_s_minus_m0_0));
HVX_Vector v_p_row1_hf_0 = hvx_vec_exp2_f16(Q6_Vhf_equals_Vqf16(v_s_minus_m1_0));
HVX_Vector v_p_row0_hf_1 = hvx_vec_exp2_f16(Q6_Vhf_equals_Vqf16(v_s_minus_m0_1));
HVX_Vector v_p_row1_hf_1 = hvx_vec_exp2_f16(Q6_Vhf_equals_Vqf16(v_s_minus_m1_1));
__fp16 * out_dtile0 = p_st_base + ci0 * HMX_FP16_TILE_N_ELMS * 2;
__fp16 * out_dtile1 = p_st_base + ci1 * HMX_FP16_TILE_N_ELMS * 2;
HVX_Vector * pv_p_out0_0 = ((HVX_Vector *) out_dtile0) + r1 / 2;
HVX_Vector * pv_p_out1_0 = pv_p_out0_0 + 16;
HVX_Vector * pv_p_out0_1 = ((HVX_Vector *) out_dtile1) + r1 / 2;
HVX_Vector * pv_p_out1_1 = pv_p_out0_1 + 16;
HVX_VectorPair vp_p_dual0 = Q6_W_vshuff_VVR(v_p_row1_hf_0, v_p_row0_hf_0, -2);
*pv_p_out0_0 = Q6_V_lo_W(vp_p_dual0);
*pv_p_out1_0 = Q6_V_hi_W(vp_p_dual0);
HVX_VectorPair vp_p_dual1 = Q6_W_vshuff_VVR(v_p_row1_hf_1, v_p_row0_hf_1, -2);
*pv_p_out0_1 = Q6_V_lo_W(vp_p_dual1);
*pv_p_out1_1 = Q6_V_hi_W(vp_p_dual1);
HVX_VectorPair vp_p0_0 = hvx_vec_f16_to_f32_shuff(v_p_row0_hf_0);
HVX_VectorPair vp_p1_0 = hvx_vec_f16_to_f32_shuff(v_p_row1_hf_0);
HVX_VectorPair vp_p0_1 = hvx_vec_f16_to_f32_shuff(v_p_row0_hf_1);
HVX_VectorPair vp_p1_1 = hvx_vec_f16_to_f32_shuff(v_p_row1_hf_1);
v_p_rowsum0 = Q6_Vqf32_vadd_Vqf32Vqf32(v_p_rowsum0, Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(vp_p0_0), Q6_V_hi_W(vp_p0_0)));
v_p_rowsum0 = Q6_Vqf32_vadd_Vqf32Vqf32(v_p_rowsum0, Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(vp_p0_1), Q6_V_hi_W(vp_p0_1)));
v_p_rowsum1 = Q6_Vqf32_vadd_Vqf32Vqf32(v_p_rowsum1, Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(vp_p1_0), Q6_V_hi_W(vp_p1_0)));
v_p_rowsum1 = Q6_Vqf32_vadd_Vqf32Vqf32(v_p_rowsum1, Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(vp_p1_1), Q6_V_hi_W(vp_p1_1)));
}
for (size_t c_rem = c; c_rem < kv_rows; c_rem += 64) {
size_t ci = c_rem / 64;
HVX_Vector v_s_minus_m0 = Q6_Vqf16_vsub_VhfVhf(my_row_buf0[ci], v_dup_m0); HVX_Vector v_s_minus_m0 = Q6_Vqf16_vsub_VhfVhf(my_row_buf0[ci], v_dup_m0);
HVX_Vector v_s_minus_m1 = Q6_Vqf16_vsub_VhfVhf(my_row_buf1[ci], v_dup_m1); HVX_Vector v_s_minus_m1 = Q6_Vqf16_vsub_VhfVhf(my_row_buf1[ci], v_dup_m1);
@@ -1281,7 +1506,7 @@ static __attribute__((noinline)) void fa_build_d_diag_inv_l(struct hmx_fa_contex
v_content = Q6_V_vror_VR(v_content, 64); v_content = Q6_V_vror_VR(v_content, 64);
} }
__fp16 * out_base = factx->vtcm_d_tiles + i * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS; __fp16 * out_base = factx->vtcm_d_inv_l + i * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
Q6_vscatter_QRMVhV(q_32_mask, (size_t) out_base, HMX_FP16_TILE_SIZE - 1, v_offsets, v_content); Q6_vscatter_QRMVhV(q_32_mask, (size_t) out_base, HMX_FP16_TILE_SIZE - 1, v_offsets, v_content);
} }
} }
@@ -1514,6 +1739,27 @@ static void fa_pop_mask_dma_gqa(dma_queue * dma, uint32_t G) {
} }
} }
static inline void fa_prefetch_block(dma_queue * dma, const struct htp_tensor * k, const struct htp_tensor * v, const struct htp_tensor * mask,
uint32_t b, size_t Bc, size_t size_k_row_padded, size_t size_k_row, size_t size_v_row_padded, size_t size_v_row,
uint32_t ik2, uint32_t ik3, uint32_t iv2, uint32_t iv3, uint32_t q_start, uint32_t im3, uint32_t kv_head, uint32_t G,
size_t m_line_bytes, size_t n_rows_q, size_t nek1, size_t prefetch_buf, struct hmx_fa_context * factx) {
const uint32_t prefetch_start = b * Bc;
const uint32_t prefetch_rows = hex_smin(Bc, nek1 - prefetch_start);
const uint8_t * k_prefetch_src = (const uint8_t *) k->data + prefetch_start * k->nb[1] + ik2 * k->nb[2] + ik3 * k->nb[3];
dma_queue_push(dma, dma_make_ptr(factx->vtcm_k_fp16[prefetch_buf], k_prefetch_src), size_k_row_padded, k->nb[1], size_k_row, prefetch_rows);
const uint8_t * v_prefetch_src = (const uint8_t *) v->data + prefetch_start * v->nb[1] + iv2 * v->nb[2] + iv3 * v->nb[3];
dma_queue_push(dma, dma_make_ptr(factx->vtcm_v_fp16[prefetch_buf], v_prefetch_src), size_v_row_padded, v->nb[1], size_v_row, prefetch_rows);
if (mask) {
if (__builtin_expect(factx->mask_broadcast, true)) {
const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + prefetch_start * sizeof(__fp16);
dma_cache_push(dma, &factx->m_cache, ms_src, m_line_bytes, mask->nb[1], prefetch_rows * sizeof(__fp16), n_rows_q);
} else {
fa_push_mask_dma_gqa(dma, mask, q_start, im3, prefetch_start, kv_head, G, m_line_bytes, prefetch_rows, n_rows_q, factx);
}
}
}
// ============================================================================ // ============================================================================
// Core HMX flash attention algorithm (GQA-merged) // Core HMX flash attention algorithm (GQA-merged)
// ============================================================================ // ============================================================================
@@ -1612,7 +1858,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
// Build the VTCM layout once (shared with the host estimator) and place every // Build the VTCM layout once (shared with the host estimator) and place every
// scratch buffer at its computed offset. // scratch buffer at its computed offset.
struct hmx_fa_vtcm_layout L; struct hmx_fa_vtcm_layout L;
hmx_fa_vtcm_layout_build(&L, G, DK, DV, Br, Bc, n_threads, pipeline); hmx_fa_vtcm_layout_build(&L, G, DK, DV, Br, Bc, n_threads, pipeline, factx.is_q_fp32);
if (L.total_bytes > ctx->vtcm_size) { if (L.total_bytes > ctx->vtcm_size) {
return HTP_STATUS_VTCM_TOO_SMALL; return HTP_STATUS_VTCM_TOO_SMALL;
@@ -1620,6 +1866,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
uint8_t * const base = ctx->vtcm_base; uint8_t * const base = ctx->vtcm_base;
factx.vtcm_q_dma = VTCM_LAYOUT_PTR(__fp16, base, L.off_q_dma);
factx.vtcm_q_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_q_tiles); factx.vtcm_q_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_q_tiles);
factx.vtcm_o_tiles[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_o_tiles[0]); factx.vtcm_o_tiles[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_o_tiles[0]);
factx.vtcm_o_tiles[1] = VTCM_LAYOUT_PTR(__fp16, base, L.off_o_tiles[1]); factx.vtcm_o_tiles[1] = VTCM_LAYOUT_PTR(__fp16, base, L.off_o_tiles[1]);
@@ -1627,12 +1874,16 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
factx.vtcm_k_fp16[1] = VTCM_LAYOUT_PTR(__fp16, base, L.off_k_fp16[1]); factx.vtcm_k_fp16[1] = VTCM_LAYOUT_PTR(__fp16, base, L.off_k_fp16[1]);
factx.vtcm_v_fp16[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_v_fp16[0]); factx.vtcm_v_fp16[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_v_fp16[0]);
factx.vtcm_v_fp16[1] = VTCM_LAYOUT_PTR(__fp16, base, L.off_v_fp16[1]); factx.vtcm_v_fp16[1] = VTCM_LAYOUT_PTR(__fp16, base, L.off_v_fp16[1]);
factx.vtcm_k_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_k_tiles); factx.vtcm_k_tiles[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_k_tiles[0]);
factx.vtcm_k_tiles[1] = VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_k_tiles[1], pipeline);
factx.vtcm_v_tiles[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_v_tiles[0]); factx.vtcm_v_tiles[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_v_tiles[0]);
factx.vtcm_v_tiles[1] = VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_v_tiles[1], pipeline); factx.vtcm_v_tiles[1] = VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_v_tiles[1], pipeline);
factx.vtcm_s_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_s_tiles); factx.vtcm_s_tiles[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_s_tiles[0]);
factx.vtcm_p_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_p_tiles); factx.vtcm_s_tiles[1] = VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_s_tiles[1], pipeline);
factx.vtcm_p_tiles[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_p_tiles[0]);
factx.vtcm_p_tiles[1] = VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_p_tiles[1], pipeline);
factx.vtcm_d_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_d_tiles); factx.vtcm_d_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_d_tiles);
factx.vtcm_d_inv_l = VTCM_LAYOUT_PTR(__fp16, base, L.off_d_inv_l);
factx.vtcm_m_vec = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_m_vec); factx.vtcm_m_vec = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_m_vec);
factx.vtcm_l_vec = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_l_vec); factx.vtcm_l_vec = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_l_vec);
factx.vtcm_s_rowmax = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_s_rowmax); factx.vtcm_s_rowmax = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_s_rowmax);
@@ -1670,6 +1921,12 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
const size_t qo_element_size = factx.is_q_fp32 ? sizeof(float) : sizeof(__fp16); const size_t qo_element_size = factx.is_q_fp32 ? sizeof(float) : sizeof(__fp16);
const bool q_transposed = q->nb[1] < q->nb[2];
const size_t q_src_stride = q_transposed ? q->nb[2] : q->nb[1];
const size_t q_row_bytes_untransposed = factx.G * factx.DK * qo_element_size;
const size_t q_row_bytes_trans_factor = factx.DK * qo_element_size;
const uint32_t kv_rows0 = hex_smin(Bc, nek1);
// ======== Reusable job descriptors for pipeline ======== // ======== Reusable job descriptors for pipeline ========
hmx_fa_qk_job_t qk_job; hmx_fa_qk_job_t qk_job;
hmx_fa_o_update_job_t ou_job; hmx_fa_o_update_job_t ou_job;
@@ -1690,34 +1947,34 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
const uint32_t iv2 = kv_head; const uint32_t iv2 = kv_head;
const uint32_t iv3 = fastdiv(ib3, &kparams->broadcast_rv3); const uint32_t iv3 = fastdiv(ib3, &kparams->broadcast_rv3);
// 1. Push Q DMA (if Q DMA is used) // 1. Push Q and KV DMAs for the very first iteration.
const size_t o_tile_bytes = factx.o_tile_bytes; // Subsequent iterations are enqueued early at the end of the previous iteration.
const bool use_q_dma = (2 * o_tile_bytes >= factx.g_br * factx.DK * (factx.is_q_fp32 ? 4 : 2)); if (ib3 == 0 && q_start == 0 && kv_head == 0) {
if (use_q_dma) { const uint8_t * q_ptr = (const uint8_t *) q->data;
const bool q_transposed = q->nb[1] < q->nb[2]; const size_t q_row_bytes = q_transposed ? n_rows_q * q_row_bytes_trans_factor : q_row_bytes_untransposed;
const uint8_t * q_ptr = (const uint8_t *) q->data + q_start * q->nb[1] + (kv_head * factx.G) * q->nb[2] + ib3 * q->nb[3];
const size_t el_size = factx.is_q_fp32 ? sizeof(float) : sizeof(__fp16);
const size_t q_row_bytes = q_transposed ? n_rows_q * factx.DK * el_size : factx.G * factx.DK * el_size;
const size_t src_stride = q_transposed ? q->nb[2] : q->nb[1];
const size_t n_rows = q_transposed ? factx.G : n_rows_q; const size_t n_rows = q_transposed ? factx.G : n_rows_q;
dma_queue_push(dma, dma_make_ptr(factx.vtcm_o_tiles[0], q_ptr), q_row_bytes, hex_smax(src_stride, q_row_bytes), q_row_bytes, n_rows); dma_queue_push(dma, dma_make_ptr(factx.vtcm_q_dma, q_ptr), q_row_bytes, hex_smax(q_src_stride, q_row_bytes), q_row_bytes, n_rows);
if (factx.n_kv_blocks > 0) {
const uint8_t * k_src = (const uint8_t *) k->data + ik2 * k->nb[2] + ik3 * k->nb[3];
dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[0], k_src), size_k_row_padded, k->nb[1], size_k_row, kv_rows0);
const uint8_t * v_src = (const uint8_t *) v->data + iv2 * v->nb[2] + iv3 * v->nb[3];
dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[0], v_src), size_v_row_padded, v->nb[1], size_v_row, kv_rows0);
if (factx.pipeline && mask) {
if (__builtin_expect(factx.mask_broadcast, true)) {
const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + 0;
dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows0 * sizeof(__fp16), n_rows_q);
} else {
fa_push_mask_dma_gqa(dma, mask, q_start, im3, 0, kv_head, G, m_line_bytes, kv_rows0, n_rows_q, &factx);
}
}
}
} }
// 2. Prefetch first KV block // 2. Pop Q DMA (blocks until Q is loaded)
if (factx.n_kv_blocks > 0) { dma_queue_pop(dma);
const uint32_t kv_rows0 = hex_smin(Bc, nek1);
const uint8_t * k_src = (const uint8_t *) k->data + ik2 * k->nb[2] + ik3 * k->nb[3];
dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[0], k_src), size_k_row_padded, k->nb[1], size_k_row, kv_rows0);
const uint8_t * v_src = (const uint8_t *) v->data + iv2 * v->nb[2] + iv3 * v->nb[3];
dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[0], v_src), size_v_row_padded, v->nb[1], size_v_row, kv_rows0);
}
// 3. Pop Q DMA (blocks until Q is loaded)
if (use_q_dma) {
dma_queue_pop(dma);
}
// ---- Load Q block & Initialize per-block state ---- // ---- Load Q block & Initialize per-block state ----
fa_phase_q_load(&factx, q, q_start, kv_head, ib3, n_rows_g); fa_phase_q_load(&factx, q, q_start, kv_head, ib3, n_rows_g);
@@ -1738,76 +1995,40 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
hmx_queue_t hmx_q = ctx->hmx_queue; hmx_queue_t hmx_q = ctx->hmx_queue;
if (factx.pipeline) { if (factx.pipeline) {
// Pipeline path // Double-buffered job structs because HMX queue runs asynchronously
hmx_fa_qk_job_t qk_job[2];
hmx_fa_o_update_job_t ou_job[2];
// Prefetch block 1 early if there are multiple blocks
if (factx.n_kv_blocks > 1) {
fa_prefetch_block(dma, k, v, mask, 1, Bc, size_k_row_padded, size_k_row, size_v_row_padded, size_v_row,
ik2, ik3, iv2, iv3, q_start, im3, kv_head, G, m_line_bytes, n_rows_q, nek1, 1, &factx);
}
// Prep and start QK-dot(0)
void * curr_k0 = dma_queue_pop(dma).dst;
fa_phase_k_interleave(&factx, kv_rows0, k_src_stride, curr_k0, 0, 0);
qk_job[0].q_tiles = factx.vtcm_q_tiles;
qk_job[0].k_tiles = factx.vtcm_k_tiles[0];
qk_job[0].s_tiles = factx.vtcm_s_tiles[0];
qk_job[0].n_row_tiles = n_row_tiles;
qk_job[0].n_col_tiles = hmx_ceil_div(kv_rows0, HMX_FP16_TILE_N_COLS);
qk_job[0].n_dot_tiles = DK / 32;
qk_job[0].n_tiles_per_bc = n_tiles_per_bc;
qk_job[0].hmx_scales = factx.vtcm_hmx_scales_qk;
hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_qk_dot_worker, &qk_job[0]));
for (uint32_t kv_blk = 0; kv_blk < factx.n_kv_blocks; ++kv_blk) { for (uint32_t kv_blk = 0; kv_blk < factx.n_kv_blocks; ++kv_blk) {
const uint32_t kv_start = kv_blk * Bc; const uint32_t kv_start = kv_blk * Bc;
const uint32_t kv_rows = hex_smin(Bc, nek1 - kv_start); const uint32_t kv_rows = hex_smin(Bc, nek1 - kv_start);
const size_t n_col_tiles = hmx_ceil_div(kv_rows, HMX_FP16_TILE_N_COLS); const size_t n_col_tiles = hmx_ceil_div(kv_rows, HMX_FP16_TILE_N_COLS);
// Push mask DMA // ---- 1. Pop and run V-prep for current block ----
if (mask) {
if (__builtin_expect(factx.mask_broadcast, true)) {
const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + kv_start * sizeof(__fp16);
dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q);
} else {
fa_push_mask_dma_gqa(dma, mask, q_start, im3, kv_start, kv_head, G, m_line_bytes, kv_rows, n_rows_q, &factx);
}
}
// Prefetch next KV block early
if (kv_blk + 1 < factx.n_kv_blocks) {
const uint32_t prefetch_start = (kv_blk + 1) * Bc;
const uint32_t prefetch_rows = hex_smin(Bc, nek1 - prefetch_start);
const size_t prefetch_buf = 1 - buf_idx;
const uint8_t * k_prefetch_src = (const uint8_t *) k->data + prefetch_start * k->nb[1] + ik2 * k->nb[2] + ik3 * k->nb[3];
dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[prefetch_buf], k_prefetch_src), size_k_row_padded, k->nb[1], size_k_row, prefetch_rows);
const uint8_t * v_prefetch_src = (const uint8_t *) v->data + prefetch_start * v->nb[1] + iv2 * v->nb[2] + iv3 * v->nb[3];
dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[prefetch_buf], v_prefetch_src), size_v_row_padded, v->nb[1], size_v_row, prefetch_rows);
}
// ---- Phase 1: K_int ----
if (kv_blk > 0) {
ou_job.o_curr = o_tile_curr;
ou_job.o_prev = o_tile_prev;
ou_job.p_tiles = factx.vtcm_p_tiles;
ou_job.v_tiles = factx.vtcm_v_tiles[1 - buf_idx];
ou_job.d_tiles = factx.vtcm_d_tiles;
ou_job.hmx_scales = factx.vtcm_hmx_scales_id;
ou_job.n_row_tiles = n_row_tiles;
ou_job.n_col_tiles = hmx_ceil_div(hex_smin(Bc, nek1 - (kv_blk - 1) * Bc), HMX_FP16_TILE_N_COLS);
ou_job.n_row_tiles_g_br = n_row_tiles_g_br;
ou_job.n_tiles_per_bc = n_tiles_per_bc;
ou_job.DV = DV;
hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job));
}
// Wait for current K DMA and interleave
void * curr_k = dma_queue_pop(dma).dst;
fa_phase_k_interleave(&factx, kv_rows, k_src_stride, curr_k, kv_start);
// ---- Phase 2: qk_dot ----
qk_job.q_tiles = factx.vtcm_q_tiles;
qk_job.k_tiles = factx.vtcm_k_tiles;
qk_job.s_tiles = factx.vtcm_s_tiles;
qk_job.n_row_tiles = n_row_tiles;
qk_job.n_col_tiles = n_col_tiles;
qk_job.n_dot_tiles = DK / 32;
qk_job.n_tiles_per_bc = n_tiles_per_bc;
qk_job.hmx_scales = factx.vtcm_hmx_scales_qk;
hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_qk_dot_worker, &qk_job));
// Wait for current V DMA and interleave
void * curr_v = dma_queue_pop(dma).dst; void * curr_v = dma_queue_pop(dma).dst;
fa_phase_v_interleave(&factx, kv_rows, v_src_stride, curr_v, factx.vtcm_v_tiles[buf_idx], n_tiles_per_bc, kv_start); fa_phase_v_interleave(&factx, kv_rows, v_src_stride, curr_v, factx.vtcm_v_tiles[buf_idx], n_tiles_per_bc, kv_start);
if (kv_blk > 0) { // ---- 2. Pop and run mask-prep for current block ----
hmx_queue_pop(hmx_q);
hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev);
}
hmx_queue_pop(hmx_q);
// ---- Phase 3: softmax + build_D ----
__fp16 * current_mask_vtcm = NULL; __fp16 * current_mask_vtcm = NULL;
if (mask) { if (mask) {
if (__builtin_expect(factx.mask_broadcast, true)) { if (__builtin_expect(factx.mask_broadcast, true)) {
@@ -1818,9 +2039,34 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
} }
} }
// ---- 3. Pop and run K-prep for next block & push next QK-dot ----
if (kv_blk + 1 < factx.n_kv_blocks) {
const uint32_t next_start = (kv_blk + 1) * Bc;
const uint32_t next_rows = hex_smin(Bc, nek1 - next_start);
const size_t next_buf = 1 - buf_idx;
void * next_k = dma_queue_pop(dma).dst;
fa_phase_k_interleave(&factx, next_rows, k_src_stride, next_k, next_start, next_buf);
qk_job[next_buf].q_tiles = factx.vtcm_q_tiles;
qk_job[next_buf].k_tiles = factx.vtcm_k_tiles[next_buf];
qk_job[next_buf].s_tiles = factx.vtcm_s_tiles[next_buf];
qk_job[next_buf].n_row_tiles = n_row_tiles;
qk_job[next_buf].n_col_tiles = hmx_ceil_div(next_rows, HMX_FP16_TILE_N_COLS);
qk_job[next_buf].n_dot_tiles = DK / 32;
qk_job[next_buf].n_tiles_per_bc = n_tiles_per_bc;
qk_job[next_buf].hmx_scales = factx.vtcm_hmx_scales_qk;
hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_qk_dot_worker, &qk_job[next_buf]));
}
// ---- 4. Wait for current block's QK-dot to finish ----
hmx_queue_pop(hmx_q);
// ---- 5. Phase 2: softmax + build_D ----
fa_softmax_args_t sargs; fa_softmax_args_t sargs;
memset(&sargs, 0, sizeof(sargs)); memset(&sargs, 0, sizeof(sargs));
sargs.factx = &factx; sargs.factx = &factx;
sargs.buf_idx = buf_idx;
sargs.kv_rows = kv_rows; sargs.kv_rows = kv_rows;
sargs.n_rows_g = n_rows_g; sargs.n_rows_g = n_rows_g;
sargs.n_col_tiles = n_col_tiles; sargs.n_col_tiles = n_col_tiles;
@@ -1838,8 +2084,39 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
sargs.mask_vtcm = current_mask_vtcm; sargs.mask_vtcm = current_mask_vtcm;
sargs.mask_vtcm_row_stride = factx.mask_buf_row_stride; sargs.mask_vtcm_row_stride = factx.mask_buf_row_stride;
sargs.slopes = factx.vtcm_slopes; sargs.slopes = factx.vtcm_slopes;
// Start HMX O update for block kv_blk - 1 (reads P[1 - buf_idx], V[1 - buf_idx])
if (kv_blk > 0) {
const size_t prev_buf = 1 - buf_idx;
ou_job[prev_buf].o_curr = o_tile_curr;
ou_job[prev_buf].o_prev = o_tile_prev;
ou_job[prev_buf].p_tiles = factx.vtcm_p_tiles[prev_buf];
ou_job[prev_buf].v_tiles = factx.vtcm_v_tiles[prev_buf];
ou_job[prev_buf].d_tiles = factx.vtcm_d_tiles;
ou_job[prev_buf].hmx_scales = factx.vtcm_hmx_scales_id;
ou_job[prev_buf].n_row_tiles = n_row_tiles;
ou_job[prev_buf].n_col_tiles = hmx_ceil_div(hex_smin(Bc, nek1 - (kv_blk - 1) * Bc), HMX_FP16_TILE_N_COLS);
ou_job[prev_buf].n_row_tiles_g_br = n_row_tiles_g_br;
ou_job[prev_buf].n_tiles_per_bc = n_tiles_per_bc;
ou_job[prev_buf].DV = DV;
hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job[prev_buf]));
}
// Run Softmax on HVX (blocking call)
fa_phase_softmax_and_build_d(&factx, &sargs, n_row_tiles, n_row_tiles_g_br); fa_phase_softmax_and_build_d(&factx, &sargs, n_row_tiles, n_row_tiles_g_br);
// Wait for HMX O update for block kv_blk - 1 to finish
if (kv_blk > 0) {
hmx_queue_pop(hmx_q);
hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev);
}
// Prefetch block kv_blk + 2
if (kv_blk + 2 < factx.n_kv_blocks) {
fa_prefetch_block(dma, k, v, mask, kv_blk + 2, Bc, size_k_row_padded, size_k_row, size_v_row_padded, size_v_row,
ik2, ik3, iv2, iv3, q_start, im3, kv_head, G, m_line_bytes, n_rows_q, nek1, buf_idx, &factx);
}
buf_idx = 1 - buf_idx; buf_idx = 1 - buf_idx;
} }
@@ -1847,18 +2124,23 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
if (factx.n_kv_blocks > 0) { if (factx.n_kv_blocks > 0) {
const uint32_t last_blk = factx.n_kv_blocks - 1; const uint32_t last_blk = factx.n_kv_blocks - 1;
const size_t last_cols = hmx_ceil_div(hex_smin(Bc, nek1 - last_blk * Bc), HMX_FP16_TILE_N_COLS); const size_t last_cols = hmx_ceil_div(hex_smin(Bc, nek1 - last_blk * Bc), HMX_FP16_TILE_N_COLS);
ou_job.o_curr = o_tile_curr; ou_job[0].o_curr = o_tile_curr;
ou_job.o_prev = o_tile_prev; ou_job[0].o_prev = o_tile_prev;
ou_job.p_tiles = factx.vtcm_p_tiles; ou_job[0].p_tiles = factx.vtcm_p_tiles[1 - buf_idx];
ou_job.v_tiles = factx.vtcm_v_tiles[1 - buf_idx]; ou_job[0].v_tiles = factx.vtcm_v_tiles[1 - buf_idx];
ou_job.d_tiles = factx.vtcm_d_tiles; ou_job[0].d_tiles = factx.vtcm_d_tiles;
ou_job.hmx_scales = factx.vtcm_hmx_scales_id; ou_job[0].hmx_scales = factx.vtcm_hmx_scales_id;
ou_job.n_row_tiles = n_row_tiles; ou_job[0].n_row_tiles = n_row_tiles;
ou_job.n_col_tiles = last_cols; ou_job[0].n_col_tiles = last_cols;
ou_job.n_row_tiles_g_br = n_row_tiles_g_br; ou_job[0].n_row_tiles_g_br = n_row_tiles_g_br;
ou_job.n_tiles_per_bc = n_tiles_per_bc; ou_job[0].n_tiles_per_bc = n_tiles_per_bc;
ou_job.DV = DV; ou_job[0].DV = DV;
hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job)); hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job[0]));
// Overlapped: run HVX build diag inv L while HMX is busy executing the update
htp_trace_event_start(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br);
htp_trace_event_stop(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
hmx_queue_pop(hmx_q); hmx_queue_pop(hmx_q);
hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev); hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev);
@@ -1892,12 +2174,12 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
// Wait for current K DMA and interleave // Wait for current K DMA and interleave
void * curr_k = dma_queue_pop(dma).dst; void * curr_k = dma_queue_pop(dma).dst;
fa_phase_k_interleave(&factx, kv_rows, k_src_stride, curr_k, kv_start); fa_phase_k_interleave(&factx, kv_rows, k_src_stride, curr_k, kv_start, 0);
{ {
qk_job.q_tiles = factx.vtcm_q_tiles; qk_job.q_tiles = factx.vtcm_q_tiles;
qk_job.k_tiles = factx.vtcm_k_tiles; qk_job.k_tiles = factx.vtcm_k_tiles[0];
qk_job.s_tiles = factx.vtcm_s_tiles; qk_job.s_tiles = factx.vtcm_s_tiles[0];
qk_job.n_row_tiles = n_row_tiles; qk_job.n_row_tiles = n_row_tiles;
qk_job.n_col_tiles = n_col_tiles; qk_job.n_col_tiles = n_col_tiles;
qk_job.n_dot_tiles = (size_t) (DK / 32); qk_job.n_dot_tiles = (size_t) (DK / 32);
@@ -1948,7 +2230,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
{ {
ou_job.o_curr = o_tile_curr; ou_job.o_curr = o_tile_curr;
ou_job.o_prev = o_tile_prev; ou_job.o_prev = o_tile_prev;
ou_job.p_tiles = factx.vtcm_p_tiles; ou_job.p_tiles = factx.vtcm_p_tiles[0];
ou_job.v_tiles = factx.vtcm_v_tiles[0]; ou_job.v_tiles = factx.vtcm_v_tiles[0];
ou_job.d_tiles = factx.vtcm_d_tiles; ou_job.d_tiles = factx.vtcm_d_tiles;
ou_job.hmx_scales = factx.vtcm_hmx_scales_id; ou_job.hmx_scales = factx.vtcm_hmx_scales_id;
@@ -1959,6 +2241,12 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
ou_job.DV = DV; ou_job.DV = DV;
hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job)); hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job));
if (kv_blk + 1 == factx.n_kv_blocks) {
// Overlapped: run HVX build diag inv L while HMX is busy executing the update
htp_trace_event_start(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br);
htp_trace_event_stop(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
}
hmx_queue_pop(ctx->hmx_queue); hmx_queue_pop(ctx->hmx_queue);
hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev); hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev);
@@ -1968,15 +2256,63 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
} }
} }
// Enqueue DMAs for the next iteration early so they overlap with O-PROC
uint32_t next_kv_head = kv_head + 1;
uint32_t next_q_start = q_start;
uint32_t next_ib3 = ib3;
if (next_kv_head >= n_kv_heads) {
next_kv_head = 0;
next_q_start = q_start + Br;
if (next_q_start >= neq1) {
next_q_start = 0;
next_ib3 = ib3 + 1;
}
}
bool has_next = (next_ib3 < neq3);
if (has_next) {
const uint32_t next_n_rows_q = hex_smin(Br, neq1 - next_q_start);
const uint8_t * next_q_ptr = (const uint8_t *) q->data + next_q_start * q->nb[1] + (next_kv_head * factx.G) * q->nb[2] + next_ib3 * q->nb[3];
const size_t next_q_row_bytes = q_transposed ? next_n_rows_q * q_row_bytes_trans_factor : q_row_bytes_untransposed;
const size_t next_n_rows = q_transposed ? factx.G : next_n_rows_q;
dma_queue_push(dma, dma_make_ptr(factx.vtcm_q_dma, next_q_ptr), next_q_row_bytes, hex_smax(q_src_stride, next_q_row_bytes), next_q_row_bytes, next_n_rows);
if (factx.n_kv_blocks > 0) {
const uint32_t next_ik2 = next_kv_head;
const uint32_t next_iv2 = next_kv_head;
uint32_t next_ik3 = ik3;
uint32_t next_iv3 = iv3;
if (next_ib3 != ib3) {
next_ik3 = fastdiv(next_ib3, &kparams->broadcast_rk3);
next_iv3 = fastdiv(next_ib3, &kparams->broadcast_rv3);
}
const uint8_t * next_k_src = (const uint8_t *) k->data + next_ik2 * k->nb[2] + next_ik3 * k->nb[3];
dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[0], next_k_src), size_k_row_padded, k->nb[1], size_k_row, kv_rows0);
const uint8_t * next_v_src = (const uint8_t *) v->data + next_iv2 * v->nb[2] + next_iv3 * v->nb[3];
dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[0], next_v_src), size_v_row_padded, v->nb[1], size_v_row, kv_rows0);
if (factx.pipeline && mask) {
uint32_t next_im3 = im3;
if (next_ib3 != ib3) {
next_im3 = fastmodulo(next_ib3, mask->ne[3], &factx.src3_div3);
}
if (__builtin_expect(factx.mask_broadcast, true)) {
const uint8_t * ms_src = (const uint8_t *) mask->data + next_q_start * mask->nb[1] + next_im3 * mask->nb[3] + 0;
dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows0 * sizeof(__fp16), next_n_rows_q);
} else {
fa_push_mask_dma_gqa(dma, mask, next_q_start, next_im3, 0, next_kv_head, G, m_line_bytes, kv_rows0, next_n_rows_q, &factx);
}
}
}
}
// ---- Final normalization ---- // ---- Final normalization ----
{ {
htp_trace_event_start(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br);
htp_trace_event_stop(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
on_job.o_curr = o_tile_curr; on_job.o_curr = o_tile_curr;
on_job.o_prev = o_tile_prev; on_job.o_prev = o_tile_prev;
on_job.d_tiles = factx.vtcm_d_tiles; on_job.d_tiles = factx.vtcm_d_inv_l;
on_job.hmx_scales = factx.vtcm_hmx_scales_id; on_job.hmx_scales = factx.vtcm_hmx_scales_id;
on_job.n_row_tiles = n_row_tiles; on_job.n_row_tiles = n_row_tiles;
on_job.n_row_tiles_g_br = n_row_tiles_g_br; on_job.n_row_tiles_g_br = n_row_tiles_g_br;
ggml/src/ggml-hexagon/htp/flash-attn-ops.h
+57 -25
View File
@@ -101,14 +101,16 @@ static_assert(sizeof(struct htp_fa_kernel_params) <= 128, "htp_fa_kernel_params
struct hmx_fa_vtcm_layout { struct hmx_fa_vtcm_layout {
// Byte offsets from vtcm_base for each region. // Byte offsets from vtcm_base for each region.
size_t off_q_tiles; size_t off_q_tiles;
size_t off_q_dma;
size_t off_o_tiles[2]; size_t off_o_tiles[2];
size_t off_k_fp16[2]; size_t off_k_fp16[2];
size_t off_v_fp16[2]; size_t off_v_fp16[2];
size_t off_k_tiles; size_t off_k_tiles[2];
size_t off_v_tiles[2]; // [1] allocated only when pipeline, else 0 size_t off_v_tiles[2];
size_t off_s_tiles; size_t off_s_tiles[2];
size_t off_p_tiles; size_t off_p_tiles[2];
size_t off_d_tiles; size_t off_d_tiles;
size_t off_d_inv_l;
size_t off_m_vec; size_t off_m_vec;
size_t off_l_vec; size_t off_l_vec;
size_t off_s_rowmax; size_t off_s_rowmax;
@@ -140,7 +142,7 @@ struct hmx_fa_vtcm_layout {
static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L, static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,
size_t gqa_factor, size_t DK, size_t DV, size_t gqa_factor, size_t DK, size_t DV,
size_t Br, size_t Bc, size_t n_threads, bool pipeline) { size_t Br, size_t Bc, size_t n_threads, bool pipeline, bool is_q_fp32) {
const size_t g_br = hex_align_up(gqa_factor * Br, HMX_FP16_TILE_N_ROWS); const size_t g_br = hex_align_up(gqa_factor * Br, HMX_FP16_TILE_N_ROWS);
const size_t q_tile_size = hex_align_up(g_br * DK * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE); const size_t q_tile_size = hex_align_up(g_br * DK * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
const size_t o_tile_size = hex_align_up(g_br * DV * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE); const size_t o_tile_size = hex_align_up(g_br * DV * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
@@ -149,6 +151,7 @@ static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,
const size_t s_tile_size = hex_align_up(g_br * Bc * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE); const size_t s_tile_size = hex_align_up(g_br * Bc * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
const size_t d_tile_size = hex_align_up(g_br * g_br * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE); const size_t d_tile_size = hex_align_up(g_br * g_br * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
const size_t q_dma_size = hex_align_up(g_br * DK * (is_q_fp32 ? sizeof(float) : sizeof(__fp16)), 128);
const size_t k_dma_size = hex_align_up(Bc * hex_round_up(DK * sizeof(__fp16), 128), 128); const size_t k_dma_size = hex_align_up(Bc * hex_round_up(DK * sizeof(__fp16), 128), 128);
const size_t v_dma_size = hex_align_up(Bc * hex_round_up(DV * sizeof(__fp16), 128), 128); const size_t v_dma_size = hex_align_up(Bc * hex_round_up(DV * sizeof(__fp16), 128), 128);
const size_t col_vec_size = hex_align_up(g_br * sizeof(float), 256); const size_t col_vec_size = hex_align_up(g_br * sizeof(float), 256);
@@ -160,27 +163,47 @@ static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,
size_t off = 0; size_t off = 0;
// Section 1: HMX Tiled Buffers (FA_HMX_TILE_SIZE = 2KB Aligned) // Group A (Part 1 - HMX Tiled buffers)
VTCM_LAYOUT_ALLOC(off, off_q_tiles, q_tile_size); VTCM_LAYOUT_ALLOC(off, off_q_tiles, q_tile_size);
VTCM_LAYOUT_ALLOC(off, off_o_tiles[0], o_tile_size); VTCM_LAYOUT_ALLOC(off, off_o_tiles[0], o_tile_size);
VTCM_LAYOUT_ALLOC(off, off_o_tiles[1], o_tile_size); VTCM_LAYOUT_ALLOC(off, off_o_tiles[1], o_tile_size);
VTCM_LAYOUT_ALLOC(off, off_k_tiles, k_tile_size);
VTCM_LAYOUT_ALLOC(off, off_v_tiles[0], v_tile_size);
VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_v_tiles[1], v_tile_size, pipeline);
VTCM_LAYOUT_ALLOC(off, off_s_tiles, s_tile_size);
VTCM_LAYOUT_ALLOC(off, off_p_tiles, s_tile_size);
VTCM_LAYOUT_ALLOC(off, off_d_tiles, d_tile_size); VTCM_LAYOUT_ALLOC(off, off_d_tiles, d_tile_size);
VTCM_LAYOUT_ALLOC(off, off_d_inv_l, d_tile_size);
// Section 2: HVX/DMA flat and vector buffers (128B / 256B Aligned) // Group B & C share start offset (Group B tiles must be 2KB aligned)
size_t off_group_b_c = hex_align_up(off, HTP_FA_HMX_TILE_SIZE);
// Group B: Compute-only buffers
size_t off_group_b = off_group_b_c;
VTCM_LAYOUT_ALLOC(off_group_b, off_k_tiles[0], k_tile_size);
VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_k_tiles[1], k_tile_size, pipeline);
VTCM_LAYOUT_ALLOC(off_group_b, off_v_tiles[0], v_tile_size);
VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_v_tiles[1], v_tile_size, pipeline);
VTCM_LAYOUT_ALLOC(off_group_b, off_s_tiles[0], s_tile_size);
VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_s_tiles[1], s_tile_size, pipeline);
VTCM_LAYOUT_ALLOC(off_group_b, off_p_tiles[0], s_tile_size);
VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_p_tiles[1], s_tile_size, pipeline);
VTCM_LAYOUT_ALLOC(off_group_b, off_s_rowmax, col_vec_size);
VTCM_LAYOUT_ALLOC(off_group_b, off_p_rowsum, col_vec_size);
VTCM_LAYOUT_ALLOC(off_group_b, off_row_bufs, row_vec_size * 2 * n_threads);
const size_t group_b_size = off_group_b - off_group_b_c;
// Group C: Q fetch DMA buffer
size_t off_group_c = off_group_b_c;
VTCM_LAYOUT_ALLOC(off_group_c, off_q_dma, q_dma_size);
const size_t group_c_size = off_group_c - off_group_b_c;
off = off_group_b_c + hex_smax(group_b_size, group_c_size);
// Group A (Part 2 - remaining non-HMX buffers)
VTCM_LAYOUT_ALLOC(off, off_k_fp16[0], k_dma_size); VTCM_LAYOUT_ALLOC(off, off_k_fp16[0], k_dma_size);
VTCM_LAYOUT_ALLOC(off, off_k_fp16[1], k_dma_size); VTCM_LAYOUT_ALLOC(off, off_k_fp16[1], k_dma_size);
VTCM_LAYOUT_ALLOC(off, off_v_fp16[0], v_dma_size); VTCM_LAYOUT_ALLOC(off, off_v_fp16[0], v_dma_size);
VTCM_LAYOUT_ALLOC(off, off_v_fp16[1], v_dma_size); VTCM_LAYOUT_ALLOC(off, off_v_fp16[1], v_dma_size);
VTCM_LAYOUT_ALLOC(off, off_m_vec, col_vec_size); VTCM_LAYOUT_ALLOC(off, off_m_vec, col_vec_size);
VTCM_LAYOUT_ALLOC(off, off_l_vec, col_vec_size); VTCM_LAYOUT_ALLOC(off, off_l_vec, col_vec_size);
VTCM_LAYOUT_ALLOC(off, off_s_rowmax, col_vec_size);
VTCM_LAYOUT_ALLOC(off, off_p_rowsum, col_vec_size);
VTCM_LAYOUT_ALLOC(off, off_row_bufs, row_vec_size * 2 * n_threads);
VTCM_LAYOUT_ALLOC(off, off_hmx_scales_id, 256); VTCM_LAYOUT_ALLOC(off, off_hmx_scales_id, 256);
VTCM_LAYOUT_ALLOC(off, off_hmx_scales_qk, 256); VTCM_LAYOUT_ALLOC(off, off_hmx_scales_qk, 256);
VTCM_LAYOUT_ALLOC(off, off_mask_buf, m_buf_size); VTCM_LAYOUT_ALLOC(off, off_mask_buf, m_buf_size);
@@ -200,9 +223,9 @@ static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,
} }
// Exact VTCM usage for a given (gqa_factor, DK, DV, Br, Bc) configuration. // Exact VTCM usage for a given (gqa_factor, DK, DV, Br, Bc) configuration.
static inline size_t hmx_fa_compute_vtcm_usage(size_t gqa_factor, size_t DK, size_t DV, size_t Br, size_t Bc, size_t n_threads, bool pipeline) { static inline size_t hmx_fa_compute_vtcm_usage(size_t gqa_factor, size_t DK, size_t DV, size_t Br, size_t Bc, size_t n_threads, bool pipeline, bool is_q_fp32) {
struct hmx_fa_vtcm_layout L; struct hmx_fa_vtcm_layout L;
hmx_fa_vtcm_layout_build(&L, gqa_factor, DK, DV, Br, Bc, n_threads, pipeline); hmx_fa_vtcm_layout_build(&L, gqa_factor, DK, DV, Br, Bc, n_threads, pipeline, is_q_fp32);
return L.total_bytes; return L.total_bytes;
} }
@@ -239,7 +262,8 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
size_t qo_len, size_t qo_len,
size_t kv_len, size_t kv_len,
size_t vtcm_budget, size_t vtcm_budget,
size_t n_threads) { size_t n_threads,
bool is_q_fp32) {
const size_t T = HMX_FP16_TILE_N_ROWS; // 32 const size_t T = HMX_FP16_TILE_N_ROWS; // 32
const size_t br_unit = hmx_ceil_div(T, gqa_factor); const size_t br_unit = hmx_ceil_div(T, gqa_factor);
const size_t bc_unit = HMX_FP16_TILE_N_COLS * 2; // 64 const size_t bc_unit = HMX_FP16_TILE_N_COLS * 2; // 64
@@ -253,8 +277,9 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
const size_t Bc_limit = can_pipeline ? hex_align_down(kv_len / FA_MIN_KV_BLOCKS, bc_unit) : const size_t Bc_limit = can_pipeline ? hex_align_down(kv_len / FA_MIN_KV_BLOCKS, bc_unit) :
(kv_len >= bc_unit ? hex_align_down(kv_len, bc_unit) : bc_unit); (kv_len >= bc_unit ? hex_align_down(kv_len, bc_unit) : bc_unit);
// Cost coefficients calibrated from profiling // Cost coefficients calibrated from profiling
const size_t c_q_fixed = 1400; // per-Q-block: q_load + epilogue o_update + o_norm + o_store const size_t c_q_fixed = 800; // per-Q-block: q_load + epilogue o_update + o_norm + o_store
const size_t c_iter_fixed = 200; // per-KV-iter: HMX queue push/pop + DMA pop + barriers const size_t c_iter_base = 200; // per-KV-iter base (HMX dot/update + DMA)
const size_t c_softmax = 600; // per 64-row vector chunk on HVX
size_t best_cost = SIZE_MAX, best_mn = 0; size_t best_cost = SIZE_MAX, best_mn = 0;
size_t best_Br = 0, best_Bc = 0; size_t best_Br = 0, best_Bc = 0;
@@ -262,13 +287,20 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
for (size_t Br = Br_max; Br >= br_unit; Br -= br_unit) { for (size_t Br = Br_max; Br >= br_unit; Br -= br_unit) {
// Try all Bc candidates from Bc_limit down to bc_unit // Try all Bc candidates from Bc_limit down to bc_unit
for (size_t Bc = Bc_limit; Bc >= bc_unit; Bc -= bc_unit) { for (size_t Bc = Bc_limit; Bc >= bc_unit; Bc -= bc_unit) {
size_t vtcm_needed = hmx_fa_compute_vtcm_usage(gqa_factor, DK, DV, Br, Bc, n_threads, can_pipeline); size_t vtcm_needed = hmx_fa_compute_vtcm_usage(gqa_factor, DK, DV, Br, Bc, n_threads, can_pipeline, is_q_fp32);
if (vtcm_needed <= vtcm_budget) { if (vtcm_needed <= vtcm_budget) {
// This Bc fits for this Br! // This Bc fits for this Br!
const size_t q_blocks = (qo_len + Br - 1) / Br; const size_t q_blocks = (qo_len + Br - 1) / Br;
const size_t kv_blocks = (kv_len + Bc - 1) / Bc; const size_t kv_blocks = (kv_len + Bc - 1) / Bc;
const size_t cost = q_blocks * (c_q_fixed + kv_blocks * c_iter_fixed); const size_t actual_threads = (kv_blocks >= 3 && n_threads >= 2) ? n_threads : 1;
const size_t mn = Br * Bc; const size_t n_rows_g = Br * gqa_factor;
const size_t n_row_vec_cnt = (n_rows_g + 63) / 64;
const size_t n_use = n_row_vec_cnt < actual_threads ? n_row_vec_cnt : actual_threads;
const size_t vecs_per_t = n_use > 0 ? (n_row_vec_cnt + n_use - 1) / n_use : 1;
const size_t c_iter_actual = c_iter_base + c_softmax * vecs_per_t;
const size_t cost = q_blocks * (c_q_fixed + kv_blocks * c_iter_actual);
const size_t mn = Br * Bc;
if (cost < best_cost || (cost == best_cost && mn > best_mn)) { if (cost < best_cost || (cost == best_cost && mn > best_mn)) {
best_cost = cost; best_cost = cost;
ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h
+23 -5
View File
@@ -767,23 +767,25 @@ static void core_mma_chunk_fp16(__fp16 *restrict c, const __fp16 *restrict a, co
// output : fp16 -> f32p // output : fp16 -> f32p
static void transfer_output_chunk_fp16_to_fp32( static void transfer_output_chunk_fp16_to_fp32_col_chunk(
float *restrict dst, float *restrict dst,
const float *restrict src2, const float *restrict src2,
const __fp16 *restrict vtcm_src, const __fp16 *restrict vtcm_src,
uint32_t start_row, uint32_t start_row,
uint32_t n_rows, uint32_t n_rows,
uint32_t n_cols, uint32_t c_len,
uint32_t total_n_cols,
uint32_t dst_stride, uint32_t dst_stride,
uint32_t src2_stride, uint32_t src2_stride,
uint32_t dst_cols uint32_t dst_cols
) { ) {
assert(n_cols % HTP_MM_HMX_TILE_N_COLS == 0); assert(c_len % HTP_MM_HMX_TILE_N_COLS == 0);
const size_t tile_row_stride = (n_cols / HTP_MM_HMX_TILE_N_COLS) * HTP_MM_HMX_TILE_N_ELMS; assert(total_n_cols % HTP_MM_HMX_TILE_N_COLS == 0);
const size_t tile_row_stride = (total_n_cols / HTP_MM_HMX_TILE_N_COLS) * HTP_MM_HMX_TILE_N_ELMS;
const HVX_Vector one = hvx_vec_splat_f16(1.0); const HVX_Vector one = hvx_vec_splat_f16(1.0);
const size_t limit_c = hex_smin(n_cols, dst_cols); const size_t limit_c = hex_smin(c_len, dst_cols);
const size_t limit_c_aligned = (limit_c & ~31); const size_t limit_c_aligned = (limit_c & ~31);
for (size_t r = 0; r < n_rows; r += 2) { for (size_t r = 0; r < n_rows; r += 2) {
@@ -848,6 +850,22 @@ static void transfer_output_chunk_fp16_to_fp32(
} }
} }
static inline void transfer_output_chunk_fp16_to_fp32(
float *restrict dst,
const float *restrict src2,
const __fp16 *restrict vtcm_src,
uint32_t start_row,
uint32_t n_rows,
uint32_t n_cols,
uint32_t dst_stride,
uint32_t src2_stride,
uint32_t dst_cols
) {
transfer_output_chunk_fp16_to_fp32_col_chunk(
dst, src2, vtcm_src, start_row, n_rows, n_cols, n_cols, dst_stride, src2_stride, dst_cols
);
}
typedef struct { typedef struct {
const __fp16 *vtcm_src; const __fp16 *vtcm_src;
float *dst; float *dst;
ggml/src/ggml-hexagon/htp/htp-ctx.h
+7 -1
View File
@@ -19,6 +19,8 @@
#endif #endif
#define HTP_MAX_MMAPS 16 #define HTP_MAX_MMAPS 16
#define HTP_MAX_DIRTY_RANGES 16
// Memory mapping // Memory mapping
struct htp_mmap { struct htp_mmap {
uint64_t size; uint64_t size;
@@ -95,7 +97,11 @@ struct htp_context {
atomic_bool vtcm_needs_release; atomic_bool vtcm_needs_release;
uint64_t max_vmem; uint64_t max_vmem;
uint32_t dirty_map[HTP_OP_MAX_TENSORS / 32]; struct htp_dirty_range {
uint32_t start;
uint32_t end;
uint32_t bi;
} dirty_ranges[HTP_MAX_DIRTY_RANGES];
// Persistent DDR scratchpad for MUL_MAT_ID mappings // Persistent DDR scratchpad for MUL_MAT_ID mappings
void * ddr_spad_base; void * ddr_spad_base;
ggml/src/ggml-hexagon/htp/htp-ops.h
+6 -2
View File
@@ -123,7 +123,7 @@ enum htp_tensor_flags {
// Tensor descriptor // Tensor descriptor
struct htp_tensor { struct htp_tensor {
uint32_t data; // Buffer offset in the messages, and data pointer on the NPU uint32_t data; // Buffer offset in the messages, and data pointer on the NPU
uint32_t alias; // Index of the canonical tensor for this memory buffer uint32_t reserved; // Reserved for alignment padding (must be multiple of 8)
uint32_t size; // Data size in bytes uint32_t size; // Data size in bytes
uint32_t flags; // Buffer / tensor flags uint32_t flags; // Buffer / tensor flags
uint32_t type; // Data type uint32_t type; // Data type
@@ -173,6 +173,7 @@ enum htp_trace_event_id {
HTP_TRACE_EVT_DMA = 0, HTP_TRACE_EVT_DMA = 0,
HTP_TRACE_EVT_L2FLUSH = 1, HTP_TRACE_EVT_L2FLUSH = 1,
HTP_TRACE_EVT_INIT = 2, HTP_TRACE_EVT_INIT = 2,
HTP_TRACE_EVT_BUFF = 3,
HTP_TRACE_EVT_HVX_COMP = 20, HTP_TRACE_EVT_HVX_COMP = 20,
HTP_TRACE_EVT_HVX_A_QUANT = 21, HTP_TRACE_EVT_HVX_A_QUANT = 21,
@@ -225,7 +226,10 @@ struct htp_opbatch_rsp {
uint32_t n_tensors; // Number of tensors uint32_t n_tensors; // Number of tensors
uint32_t n_ops; // Number of op profile descriptors uint32_t n_ops; // Number of op profile descriptors
uint32_t n_traces[HTP_MAX_NTHREADS + 1]; uint32_t n_traces[HTP_MAX_NTHREADS + 1];
uint8_t pad[8]; // align to 8 bytes uint32_t usecs; // Number of usec
uint32_t pad; // align to 8 bytes
uint64_t cycles_start; // Start cycle counter
uint64_t cycles_stop; // Stop cycle counter
// struct htp_prof_desc profs[]; -- dspqueue buf 0 // struct htp_prof_desc profs[]; -- dspqueue buf 0
}; };
ggml/src/ggml-hexagon/htp/htp-tensor.c
+188 -99
View File
@@ -2,6 +2,7 @@
#include <qurt.h> #include <qurt.h>
#include <qurt_memory.h> #include <qurt_memory.h>
#include <HAP_farf.h>
#include "hex-common.h" #include "hex-common.h"
#include "hex-utils.h" #include "hex-utils.h"
@@ -10,84 +11,6 @@
#include "htp-ctx.h" #include "htp-ctx.h"
#include "work-queue.h" #include "work-queue.h"
struct l2flush_task {
struct htp_thread_trace * trace;
uint32_t start;
uint32_t end;
uint32_t chunk_size;
uint32_t ti;
};
static void l2flush_thread_worker(unsigned int n, unsigned int i, void * data) {
struct l2flush_task * task = (struct l2flush_task *) data;
const uint32_t start = task->start;
const uint32_t end = task->end;
const uint32_t ti = task->ti;
const uint32_t chunk_size = task->chunk_size;
const uint32_t thread_s = start + i * chunk_size;
if (thread_s >= end) {
return;
}
uint32_t thread_e = thread_s + chunk_size;
if (thread_e > end) {
thread_e = end;
}
struct htp_thread_trace * tr = &task->trace[i];
htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, ti);
hex_l2flush((void *) (uintptr_t) thread_s, thread_e - thread_s);
htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, ti);
}
static void flush_all_dcache(struct htp_context * ctx) {
struct htp_thread_trace * tr = &ctx->trace[0];
htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, 0);
qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
hex_l2fetch_block(ctx, ctx->footprint);
htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, 0);
bitmap_reset(ctx->dirty_map, HTP_OP_MAX_TENSORS);
}
static void flush_tensor_range(struct htp_context * ctx, const struct htp_tensor * t) {
struct htp_thread_trace * tr = &ctx->trace[0];
if (t->size > HEX_L2_FLUSH_WQ_THRESHOLD && ctx->n_threads > 1) {
struct l2flush_task task;
task.start = hex_align_down((size_t) t->data, HEX_L2_LINE_SIZE);
task.end = hex_align_up((size_t) t->data + t->size, HEX_L2_LINE_SIZE);
task.ti = t->ti;
task.trace = ctx->trace;
const uint32_t total_size = task.end - task.start;
const uint32_t n_blocks = (total_size + HEX_L2_BLOCK_SIZE - 1) / HEX_L2_BLOCK_SIZE;
const uint32_t blocks_per_thread = fastdiv(n_blocks + ctx->n_threads - 1, &ctx->n_threads_div);
task.chunk_size = blocks_per_thread * HEX_L2_BLOCK_SIZE;
work_queue_run(ctx->work_queue, l2flush_thread_worker, &task, ctx->n_threads);
} else {
htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, t->ti);
hex_l2flush((void *) t->data, t->size);
htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, t->ti);
}
htp_tensor_make_clean(t, ctx->dirty_map);
}
void htp_tensor_flush(struct htp_context * ctx, const struct htp_tensor * t) {
if (!bitmap_test(ctx->dirty_map, t->ti)) {
return;
}
if (t->size > HEX_L2_FLUSH_ALL_THRESHOLD) {
flush_all_dcache(ctx);
return;
}
flush_tensor_range(ctx, t);
}
// One dirty tensor's line-aligned range, placed in the flattened global block space.
struct l2flush_range { struct l2flush_range {
uint32_t start; // line-aligned start address uint32_t start; // line-aligned start address
uint32_t end; // line-aligned end address uint32_t end; // line-aligned end address
@@ -103,9 +26,18 @@ struct l2flush_multi_task {
uint32_t blocks_per_thread; uint32_t blocks_per_thread;
}; };
static void flush_all_dcache(struct htp_context * ctx) {
struct htp_thread_trace * tr = &ctx->trace[0];
htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, 0);
qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
hex_l2fetch_block(ctx, ctx->footprint);
htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, 0);
memset(ctx->dirty_ranges, 0, sizeof(ctx->dirty_ranges));
}
static void l2flush_multi_worker(unsigned int n, unsigned int i, void * data) { static void l2flush_multi_worker(unsigned int n, unsigned int i, void * data) {
(void) n;
struct l2flush_multi_task * task = (struct l2flush_multi_task *) data; struct l2flush_multi_task * task = (struct l2flush_multi_task *) data;
(void) n;
const uint32_t gb_first = i * task->blocks_per_thread; const uint32_t gb_first = i * task->blocks_per_thread;
uint32_t gb_last = gb_first + task->blocks_per_thread; uint32_t gb_last = gb_first + task->blocks_per_thread;
@@ -141,11 +73,177 @@ static void l2flush_multi_worker(unsigned int n, unsigned int i, void * data) {
htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, gb_first); htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, gb_first);
} }
void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n) { void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n) {
uint64_t total_dirty = 0; const struct htp_tensor * pending[HTP_OP_MAX_OUTPUTS];
uint32_t n_pending = 0;
for (uint32_t i = 0; i < n; i++) { for (uint32_t i = 0; i < n; i++) {
const struct htp_tensor * t = tensors[i]; const struct htp_tensor * t = tensors[i];
if (t && bitmap_test(ctx->dirty_map, t->ti)) { if (!t) continue;
uint32_t t_start = t->data;
uint32_t t_end = t_start + t->size;
bool merged = false;
for (uint32_t j = 0; j < HTP_MAX_DIRTY_RANGES; j++) {
struct htp_dirty_range * r = &ctx->dirty_ranges[j];
if (!r->start) continue;
if (r->start <= t_end && t_start <= r->end) {
uint32_t new_start = (t_start < r->start) ? t_start : r->start;
uint32_t new_end = (t_end > r->end) ? t_end : r->end;
r->start = new_start;
r->end = new_end;
merged = true;
}
}
if (!merged) {
pending[n_pending++] = t;
}
}
if (n_pending == 0) {
return;
}
uint32_t empty_indices[HTP_MAX_DIRTY_RANGES];
uint32_t active_indices[HTP_MAX_DIRTY_RANGES];
uint32_t n_active = 0;
uint32_t n_empty = 0;
for (uint32_t j = 0; j < HTP_MAX_DIRTY_RANGES; j++) {
if (ctx->dirty_ranges[j].start) {
active_indices[n_active++] = j;
} else {
empty_indices[n_empty++] = j;
}
}
if (n_pending <= n_empty) {
for (uint32_t i = 0; i < n_pending; i++) {
uint32_t idx = empty_indices[i];
struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
r->start = pending[i]->data;
r->end = pending[i]->data + pending[i]->size;
r->bi = pending[i]->bi;
}
return;
}
uint32_t n_evict = n_pending - n_empty;
uint32_t total_evict_size = 0;
for (uint32_t i = 0; i < n_evict; i++) {
uint32_t idx = active_indices[i];
struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
total_evict_size += r->end - r->start;
}
if (total_evict_size > HEX_L2_FLUSH_ALL_THRESHOLD) {
flush_all_dcache(ctx);
for (uint32_t i = 0; i < n_pending; i++) {
struct htp_dirty_range * r = &ctx->dirty_ranges[i];
r->start = pending[i]->data;
r->end = pending[i]->data + pending[i]->size;
r->bi = pending[i]->bi;
}
return;
}
if (total_evict_size > HEX_L2_FLUSH_WQ_THRESHOLD && ctx->n_threads > 1 && n_evict <= HTP_OP_MAX_INPUTS) {
struct l2flush_multi_task task;
task.trace = ctx->trace;
task.n_ranges = n_evict;
uint32_t block_acc = 0;
for (uint32_t i = 0; i < n_evict; i++) {
uint32_t idx = active_indices[i];
struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
struct l2flush_range * rg = &task.ranges[i];
rg->start = hex_align_down((size_t) r->start, HEX_L2_LINE_SIZE);
rg->end = hex_align_up((size_t) r->end, HEX_L2_LINE_SIZE);
rg->block_first = block_acc;
rg->n_blocks = (rg->end - rg->start + HEX_L2_BLOCK_SIZE - 1) / HEX_L2_BLOCK_SIZE;
block_acc += rg->n_blocks;
}
task.total_blocks = block_acc;
task.blocks_per_thread = fastdiv(block_acc + ctx->n_threads - 1, &ctx->n_threads_div);
work_queue_run(ctx->work_queue, l2flush_multi_worker, &task, ctx->n_threads);
} else {
struct htp_thread_trace * tr = &ctx->trace[0];
htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, 0);
for (uint32_t i = 0; i < n_evict; i++) {
uint32_t idx = active_indices[i];
struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
uint32_t size = r->end - r->start;
hex_l2flush((void *) (uintptr_t) r->start, size);
}
htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, 0);
}
for (uint32_t i = 0; i < n_evict; i++) {
uint32_t idx = active_indices[i];
struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
r->start = pending[i]->data;
r->end = pending[i]->data + pending[i]->size;
r->bi = pending[i]->bi;
}
for (uint32_t i = 0; i < n_empty; i++) {
uint32_t idx = empty_indices[i];
struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
r->start = pending[n_evict + i]->data;
r->end = pending[n_evict + i]->data + pending[n_evict + i]->size;
r->bi = pending[n_evict + i]->bi;
}
}
static void make_tensor_clean(struct htp_context * ctx, const struct htp_tensor * t) {
uint32_t t_start = t->data;
uint32_t t_end = t_start + t->size;
for (uint32_t i = 0; i < HTP_MAX_DIRTY_RANGES; i++) {
struct htp_dirty_range * r = &ctx->dirty_ranges[i];
if (!r->start) continue;
if (r->start < t_end && t_start < r->end) {
if (t_start <= r->start && r->end <= t_end) {
r->start = 0;
} else if (t_start <= r->start) {
r->start = t_end;
} else if (r->end <= t_end) {
r->end = t_start;
}
}
}
}
static inline bool is_tensor_dirty(struct htp_context * ctx, const struct htp_tensor * t) {
uint32_t t_start = t->data;
uint32_t t_end = t_start + t->size;
for (uint32_t i = 0; i < HTP_MAX_DIRTY_RANGES; i++) {
struct htp_dirty_range * r = &ctx->dirty_ranges[i];
if (!r->start) continue;
if (r->start < t_end && t_start < r->end) {
return true;
}
}
return false;
}
void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n) {
const struct htp_tensor * dirty_tensors[HTP_OP_MAX_INPUTS];
uint32_t n_dirty = 0;
uint64_t total_dirty = 0;
for (uint32_t i = 0; i < n; i++) {
const struct htp_tensor * t = tensors[i];
if (t && (t->flags & HTP_TENSOR_COMPUTE) && is_tensor_dirty(ctx, t)) {
dirty_tensors[n_dirty++] = t;
total_dirty += t->size; total_dirty += t->size;
} }
} }
@@ -159,21 +257,15 @@ void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * co
return; return;
} }
// Aggregate is small enough to walk. Thread it across all dirty ranges at once if (total_dirty >= HEX_L2_FLUSH_WQ_THRESHOLD && ctx->n_threads > 1) {
// when it is worth the dispatch, otherwise flush sequentially.
if (total_dirty > HEX_L2_FLUSH_WQ_THRESHOLD && ctx->n_threads > 1) {
struct l2flush_multi_task task; struct l2flush_multi_task task;
task.trace = ctx->trace; task.trace = ctx->trace;
task.n_ranges = 0; task.n_ranges = 0;
uint32_t block_acc = 0; uint32_t block_acc = 0;
for (uint32_t i = 0; i < n; i++) { for (uint32_t i = 0; i < n_dirty; i++) {
const struct htp_tensor * t = tensors[i]; const struct htp_tensor * t = dirty_tensors[i];
if (!t || !bitmap_test(ctx->dirty_map, t->ti)) { make_tensor_clean(ctx, t);
continue;
}
// Clear as we go: dedups a tensor passed as multiple srcs (e.g. mul(x,x)).
htp_tensor_make_clean(t, ctx->dirty_map);
struct l2flush_range * rg = &task.ranges[task.n_ranges++]; struct l2flush_range * rg = &task.ranges[task.n_ranges++];
rg->start = hex_align_down((size_t) t->data, HEX_L2_LINE_SIZE); rg->start = hex_align_down((size_t) t->data, HEX_L2_LINE_SIZE);
@@ -191,14 +283,11 @@ void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * co
} }
struct htp_thread_trace * tr = &ctx->trace[0]; struct htp_thread_trace * tr = &ctx->trace[0];
for (uint32_t i = 0; i < n; i++) { for (uint32_t i = 0; i < n_dirty; i++) {
const struct htp_tensor * t = tensors[i]; const struct htp_tensor * t = dirty_tensors[i];
if (!t || !bitmap_test(ctx->dirty_map, t->ti)) {
continue;
}
htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, t->ti); htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, t->ti);
hex_l2flush((void *) t->data, t->size); hex_l2flush((void *) (uintptr_t) t->data, t->size);
htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, t->ti); htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, t->ti);
htp_tensor_make_clean(t, ctx->dirty_map); make_tensor_clean(ctx, t);
} }
} }
ggml/src/ggml-hexagon/htp/htp-tensor.h
+1 -17
View File
@@ -5,10 +5,6 @@
#include "htp-ops.h" #include "htp-ops.h"
#include "hex-bitmap.h" #include "hex-bitmap.h"
static inline struct htp_tensor * htp_tensor_alias(const struct htp_tensor * t) {
return (struct htp_tensor *) (uintptr_t) t->alias;
}
static inline void * htp_tensor_data(const struct htp_tensor * t) { static inline void * htp_tensor_data(const struct htp_tensor * t) {
return (void *) (uintptr_t) t->data; return (void *) (uintptr_t) t->data;
} }
@@ -17,20 +13,8 @@ static inline uint32_t * htp_tensor_flags(const struct htp_tensor * t) {
return (uint32_t *) &t->flags; return (uint32_t *) &t->flags;
} }
static inline void htp_tensor_make_dirty(const struct htp_tensor * t, uint32_t * dirty_map) {
struct htp_tensor * curr = (struct htp_tensor *) t;
do {
bitmap_set(dirty_map, curr->ti);
curr = htp_tensor_alias(curr);
} while (curr != t);
}
static inline void htp_tensor_make_clean(const struct htp_tensor * t, uint32_t * dirty_map) {
bitmap_clear(dirty_map, t->ti);
}
struct htp_context; struct htp_context;
void htp_tensor_flush(struct htp_context * ctx, const struct htp_tensor * t);
void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n); void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n);
void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n);
#endif // HTP_TENSOR_H #endif // HTP_TENSOR_H
ggml/src/ggml-hexagon/htp/hvx-fa-kernels.h
+71
View File
@@ -208,6 +208,77 @@ static inline void hvx_mad_f32_f16_aa_rx2(float * restrict y, const void * restr
} }
} }
} }
static inline void hvx_mad_f32_f16_aa_vec(float * restrict y, const void * restrict x, HVX_Vector S0, uint32_t n) {
const HVX_Vector * restrict vx0 = (const HVX_Vector *) x;
HVX_VectorPair * restrict vy_p = (HVX_VectorPair *) y;
HVX_Vector * restrict vy = (HVX_Vector *) y;
uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
uint32_t nloe = n % VLEN_FP16; // leftover elements
uint32_t i = 0;
#pragma unroll(2)
for (i = 0; i < nvec; ++i) {
vy_p[i] = hvx_vec_mpyacc_f32_f16(vy_p[i], Q6_Vh_vshuff_Vh(vx0[i]), S0);
}
if (nloe) {
HVX_VectorPair xy_p = vy_p[i];
xy_p = hvx_vec_mpyacc_f32_f16(xy_p, Q6_Vh_vshuff_Vh(vx0[i]), S0);
HVX_Vector xy = Q6_V_lo_W(xy_p);
i = 2 * i; // index for vy
if (nloe >= VLEN_FP32) {
vy[i] = xy;
nloe -= VLEN_FP32; ++i; xy = Q6_V_hi_W(xy_p);
}
if (nloe) {
hvx_vec_store_a(&vy[i], nloe * 4, xy);
}
}
}
static inline void hvx_mad_f32_f16_aa_rx2_vec(float * restrict y, const void * restrict x0, const void * restrict x1,
HVX_Vector S0, HVX_Vector S1, uint32_t n) {
const HVX_Vector * restrict vx0 = (const HVX_Vector *) x0;
const HVX_Vector * restrict vx1 = (const HVX_Vector *) x1;
HVX_VectorPair * restrict vy_p = (HVX_VectorPair *) y;
HVX_Vector * restrict vy = (HVX_Vector *) y;
uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
uint32_t nloe = n % VLEN_FP16; // leftover elements
uint32_t i = 0;
#pragma unroll(2)
for (i = 0; i < nvec; ++i) {
vy_p[i] = hvx_vec_mpyacc_f32_f16(vy_p[i], Q6_Vh_vshuff_Vh(vx0[i]), S0);
vy_p[i] = hvx_vec_mpyacc_f32_f16(vy_p[i], Q6_Vh_vshuff_Vh(vx1[i]), S1);
}
if (nloe) {
HVX_VectorPair xy_p = vy_p[i];
xy_p = hvx_vec_mpyacc_f32_f16(xy_p, Q6_Vh_vshuff_Vh(vx0[i]), S0);
xy_p = hvx_vec_mpyacc_f32_f16(xy_p, Q6_Vh_vshuff_Vh(vx1[i]), S1);
HVX_Vector xy = Q6_V_lo_W(xy_p);
i = 2 * i; // index for vy
if (nloe >= VLEN_FP32) {
vy[i] = xy;
nloe -= VLEN_FP32; ++i; xy = Q6_V_hi_W(xy_p);
}
if (nloe) {
hvx_vec_store_a(&vy[i], nloe * 4, xy);
}
}
}
static inline void hvx_scale_vec_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, const uint32_t n, HVX_Vector vs) { static inline void hvx_scale_vec_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, const uint32_t n, HVX_Vector vs) {
assert((size_t) dst % 128 == 0); assert((size_t) dst % 128 == 0);
ggml/src/ggml-hexagon/htp/hvx-reduce.h
+40
View File
@@ -286,6 +286,46 @@ static inline float hvx_sum_of_squares_f32(const uint8_t * restrict src, const i
} }
} }
// Signed 32-bit Integer Max variants
static inline HVX_Vector hvx_vec_reduce_max_n_i32(HVX_Vector in, unsigned int n) {
unsigned int total = n * 4; // total vec nbytes
unsigned int width = 4; // int32 nbytes
HVX_Vector max_val = in, max_t;
while (width < total) {
max_t = Q6_V_vror_VR(max_val, width); // rotate right
max_val = Q6_Vw_vmax_VwVw(max_t, max_val); // elementwise signed max
width = width << 1;
}
return max_val;
}
static inline HVX_Vector hvx_vec_reduce_max_i32(HVX_Vector in) {
return hvx_vec_reduce_max_n_i32(in, 32);
}
static inline int32_t hvx_reduce_max_i32_a(const uint8_t * restrict src, const int num_elems) {
HVX_Vector init_vec = Q6_V_vsplat_R(((const int32_t *) src)[0]);
HVX_Vector pad_vec = Q6_V_vsplat_R(0x80000000);
assert((uintptr_t) src % 128 == 0);
hvx_reduce_loop_body(HVX_Vector, init_vec, pad_vec, Q6_Vw_vmax_VwVw, hvx_vec_reduce_max_i32, hvx_vec_get_i32);
}
static inline int32_t hvx_reduce_max_i32_u(const uint8_t * restrict src, const int num_elems) {
HVX_Vector init_vec = Q6_V_vsplat_R(((const int32_t *) src)[0]);
HVX_Vector pad_vec = Q6_V_vsplat_R(0x80000000);
hvx_reduce_loop_body(HVX_UVector, init_vec, pad_vec, Q6_Vw_vmax_VwVw, hvx_vec_reduce_max_i32, hvx_vec_get_i32);
}
static inline int32_t hvx_reduce_max_i32(const uint8_t * restrict src, const int num_elems) {
if (hex_is_aligned((void *) src, 128)) {
return hvx_reduce_max_i32_a(src, num_elems);
} else {
return hvx_reduce_max_i32_u(src, num_elems);
}
}
#undef hvx_reduce_loop_body #undef hvx_reduce_loop_body
#undef HVX_REDUCE_MAX_OP #undef HVX_REDUCE_MAX_OP
#undef HVX_REDUCE_SUM_OP #undef HVX_REDUCE_SUM_OP
ggml/src/ggml-hexagon/htp/main.c
+37 -27
View File
@@ -901,10 +901,8 @@ static void prep_tensor(struct htp_context *ctx, struct htp_buf_desc *bufs, stru
uint32_t offset = t->data; uint32_t offset = t->data;
uint32_t size = t->size; uint32_t size = t->size;
uint32_t bi = t->bi; uint32_t bi = t->bi;
uint32_t alias = t->alias;
t->data = (uint32_t) (bufs[bi].base + offset); // update data to the actual pointer t->data = (uint32_t) (bufs[bi].base + offset); // update data to the actual pointer
t->alias = (uint32_t) (tens + alias); // update alias to the actual pointer
FARF(HIGH, "prep-tensor #%u: bi %u offset %u size %u data %p : %u:%u:%u:%u", idx, t->bi, offset, t->size, (void*) t->data, FARF(HIGH, "prep-tensor #%u: bi %u offset %u size %u data %p : %u:%u:%u:%u", idx, t->bi, offset, t->size, (void*) t->data,
t->ne[0], t->ne[1], t->ne[3], t->ne[3]); t->ne[0], t->ne[1], t->ne[3], t->ne[3]);
@@ -955,14 +953,14 @@ static int proc_op_req(struct htp_ops_context * octx, struct htp_tensor *tens, u
octx->dsts[i] = dst; octx->dsts[i] = dst;
octx->dst_dma[i] = octx->ctx->dma; // FIXME: ? octx->ctx->dma_cached : octx->ctx->dma; octx->dst_dma[i] = octx->ctx->dma; // FIXME: ? octx->ctx->dma_cached : octx->ctx->dma;
htp_tensor_make_dirty(dst, octx->ctx->dirty_map);
FARF(HIGH, "prep-dst[%u] #%u: data %p size %u : %u:%u:%u:%u", i, dst_idx, (void*) dst->data, dst->size, FARF(HIGH, "prep-dst[%u] #%u: data %p size %u : %u:%u:%u:%u", i, dst_idx, (void*) dst->data, dst->size,
dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]); dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]);
} }
int status = execute_op(octx); int status = execute_op(octx);
htp_tensor_dirty_all(octx->ctx, octx->dsts, HTP_OP_MAX_OUTPUTS);
octx->src0_spad.src = NULL; octx->src0_spad.src = NULL;
octx->src1_spad.src = NULL; octx->src1_spad.src = NULL;
octx->src2_spad.src = NULL; octx->src2_spad.src = NULL;
@@ -994,12 +992,6 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r
FARF(HIGH, "processing opbatch #%u: n-bufs %u n-tensors %u n-ops %u n-traces %u : m-size %u b-size %u t-size %u o-size %u", req->id, FARF(HIGH, "processing opbatch #%u: n-bufs %u n-tensors %u n-ops %u n-traces %u : m-size %u b-size %u t-size %u o-size %u", req->id,
n_bufs, n_tens, n_ops, req->n_traces, dbuf->size, b_size, t_size, o_size); n_bufs, n_tens, n_ops, req->n_traces, dbuf->size, b_size, t_size, o_size);
// Clean cache at the start of the batch
// We cant trace this part because the trace buffer is setup later
qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
hex_l2fetch_block(ctx, ctx->footprint);
bitmap_reset(ctx->dirty_map, HTP_OP_MAX_TENSORS);
// Setup descriptor pointers // Setup descriptor pointers
uint8_t * m_ptr = dbuf->ptr; uint8_t * m_ptr = dbuf->ptr;
struct htp_buf_desc* bufs = (struct htp_buf_desc*) m_ptr; m_ptr += b_size; struct htp_buf_desc* bufs = (struct htp_buf_desc*) m_ptr; m_ptr += b_size;
@@ -1007,13 +999,8 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r
struct htp_op_desc* ops = (struct htp_op_desc*) m_ptr; m_ptr += o_size; struct htp_op_desc* ops = (struct htp_op_desc*) m_ptr; m_ptr += o_size;
struct htp_prof_desc* pds = (struct htp_prof_desc*) m_ptr; struct htp_prof_desc* pds = (struct htp_prof_desc*) m_ptr;
prep_op_bufs(ctx, bufs, n_bufs); struct profile_data batch_prof;
prep_tensors(ctx, bufs, tens, n_tens); profile_start(HTP_PROF_BASIC, &batch_prof);
struct htp_ops_context *octx = &ctx->octx;
memset(octx, 0, sizeof(*octx));
octx->n_threads = ctx->n_threads;
octx->ctx = ctx;
memset(ctx->trace, 0, sizeof(ctx->trace)); memset(ctx->trace, 0, sizeof(ctx->trace));
if (ctx->profiler == HTP_PROF_TRACE) { if (ctx->profiler == HTP_PROF_TRACE) {
@@ -1024,6 +1011,24 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r
} }
} }
// Clean cache at the start of the batch
htp_trace_event_start(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
hex_l2fetch_block(ctx, ctx->footprint);
memset(ctx->dirty_ranges, 0, sizeof(ctx->dirty_ranges));
htp_trace_event_stop(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
htp_trace_event_start(&ctx->trace[0], HTP_TRACE_EVT_BUFF, 0);
prep_op_bufs(ctx, bufs, n_bufs);
htp_trace_event_stop(&ctx->trace[0], HTP_TRACE_EVT_BUFF, 0);
prep_tensors(ctx, bufs, tens, n_tens);
struct htp_ops_context *octx = &ctx->octx;
memset(octx, 0, sizeof(*octx));
octx->n_threads = ctx->n_threads;
octx->ctx = ctx;
work_queue_wakeup(ctx->work_queue); work_queue_wakeup(ctx->work_queue);
if (ctx->hmx_queue) { if (ctx->hmx_queue) {
hmx_queue_wakeup(ctx->hmx_queue); hmx_queue_wakeup(ctx->hmx_queue);
@@ -1056,13 +1061,23 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r
} }
work_queue_suspend(ctx->work_queue); work_queue_suspend(ctx->work_queue);
// Flush remaining dirty tensors at the end of the batch
htp_trace_event_start(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
htp_trace_event_stop(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
profile_stop(HTP_PROF_BASIC, &batch_prof);
struct htp_opbatch_rsp rsp; struct htp_opbatch_rsp rsp;
memset(&rsp, 0, sizeof(rsp)); memset(&rsp, 0, sizeof(rsp));
rsp.id = req->id; rsp.id = req->id;
rsp.status = op_status; rsp.status = op_status;
rsp.n_bufs = n_bufs; rsp.n_bufs = n_bufs;
rsp.n_tensors = n_tens; rsp.n_tensors = n_tens;
rsp.n_ops = n_ops; rsp.n_ops = n_ops;
rsp.usecs = batch_prof.usecs;
rsp.cycles_start = batch_prof.cycles_start;
rsp.cycles_stop = batch_prof.cycles_stop;
if (ctx->profiler == HTP_PROF_TRACE) { if (ctx->profiler == HTP_PROF_TRACE) {
for (int t = 0; t <= HTP_MAX_NTHREADS; t++) { for (int t = 0; t <= HTP_MAX_NTHREADS; t++) {
@@ -1073,11 +1088,6 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r
struct dspqueue_buffer write_dbuf = *dbuf; struct dspqueue_buffer write_dbuf = *dbuf;
write_dbuf.flags = DSPQUEUE_BUFFER_FLAG_FLUSH_SENDER | DSPQUEUE_BUFFER_FLAG_INVALIDATE_RECIPIENT; write_dbuf.flags = DSPQUEUE_BUFFER_FLAG_FLUSH_SENDER | DSPQUEUE_BUFFER_FLAG_INVALIDATE_RECIPIENT;
// Flush remaining dirty tensors at the end of the batch
htp_trace_event_start(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
htp_trace_event_stop(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
err = dspqueue_write(queue, 0, 1, &write_dbuf, sizeof(rsp), (const uint8_t *) &rsp, DSPQUEUE_TIMEOUT_NONE); err = dspqueue_write(queue, 0, 1, &write_dbuf, sizeof(rsp), (const uint8_t *) &rsp, DSPQUEUE_TIMEOUT_NONE);
if (err != 0) { if (err != 0) {
FARF(ERROR, "dspqueue_write failed: 0x%08x", (unsigned) err); FARF(ERROR, "dspqueue_write failed: 0x%08x", (unsigned) err);
ggml/src/ggml-hexagon/htp/matmul-ops.c
+236 -104
View File
@@ -14,6 +14,8 @@
#include "hex-dma.h" #include "hex-dma.h"
#include "hvx-utils.h" #include "hvx-utils.h"
#include "hvx-dump.h" #include "hvx-dump.h"
#include "hvx-arith.h"
#include "hvx-reduce.h"
#define GGML_COMMON_DECL_C #define GGML_COMMON_DECL_C
#include "ggml-common.h" #include "ggml-common.h"
@@ -82,6 +84,8 @@ struct htp_mm_context {
// Precomputed values // Precomputed values
uint32_t src0_nrows_per_thread; uint32_t src0_nrows_per_thread;
uint32_t src0_row_size_padded;
uint32_t src1_nrows;
struct fastdiv_values mm_div_ne12_ne1; struct fastdiv_values mm_div_ne12_ne1;
struct fastdiv_values mm_div_ne1; struct fastdiv_values mm_div_ne1;
@@ -103,6 +107,7 @@ struct htp_mm_context {
// Fields for scattered mapping & HMX support in MUL_MAT_ID // Fields for scattered mapping & HMX support in MUL_MAT_ID
const uint32_t * matrix_row_counts; const uint32_t * matrix_row_counts;
const struct mmid_row_mapping * matrix_rows; const struct mmid_row_mapping * matrix_rows;
uint32_t mapping_stride;
// Dynamic VTCM pointers allocated sequentially // Dynamic VTCM pointers allocated sequentially
uint8_t * vtcm_src0; uint8_t * vtcm_src0;
@@ -154,8 +159,6 @@ static const uint8_t __attribute__((aligned(VLEN))) kvalues_mxfp4_lut[] = {
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
}; };
#define htp_matmul_tensors_preamble \ #define htp_matmul_tensors_preamble \
const struct htp_tensor * restrict src0 = octx->src[0]; \ const struct htp_tensor * restrict src0 = octx->src[0]; \
const struct htp_tensor * restrict src1 = octx->src[1]; \ const struct htp_tensor * restrict src1 = octx->src[1]; \
@@ -444,6 +447,16 @@ static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void
\ \
uint32_t push_ct = ct_start; \ uint32_t push_ct = ct_start; \
if (src0_start_row < src0_end_row) { \ if (src0_start_row < src0_end_row) { \
if (src2) { \
float * vtcm_src2_ptr = (float *) mmctx->vtcm_src2 + src0_start_row; \
const float * src2_ptr = (const float *) src2->data + src0_start_row; \
int slice_size = (int)MIN(src0_end_row, ne0) - (int)src0_start_row; \
if (slice_size > 0) { \
dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr, src2_ptr), \
slice_size * sizeof(float), slice_size * sizeof(float), slice_size * sizeof(float), 1); \
dma_queue_pop_nowait(dma_queue); \
} \
} \
for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { \ for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { \
dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, \ dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, \
src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \
@@ -465,7 +478,7 @@ static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void
\ \
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); \
DOT_2X1(ne10, dst_ptr, w_tile, src1_col, valid_rows, NULL); \ DOT_2X1(ne10, dst_ptr, w_tile, src1_col, valid_rows, NULL); \
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \
\ \
if (push_ct < ct_end) { \ if (push_ct < ct_end) { \
dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile, src0_row + push_ct * tile_row_stride), \ dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile, src0_row + push_ct * tile_row_stride), \
@@ -476,24 +489,16 @@ static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void
\ \
int copy_cnt = (int)MIN(src0_end_row, ne0) - (int)src0_start_row; \ int copy_cnt = (int)MIN(src0_end_row, ne0) - (int)src0_start_row; \
if (copy_cnt > 0) { \ if (copy_cnt > 0) { \
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct_end); \
if (src2) { \ if (src2) { \
float * dst_ptr = &dst_col[src0_start_row]; \ hvx_add_f32_uaa((uint8_t *) &dst_col[src0_start_row], \
const float * src2_ptr = (const float *) src2->data + src0_start_row; \ (const uint8_t *) tmp, \
float * tmp_ptr = tmp; \ (const uint8_t *) ((const float *) mmctx->vtcm_src2 + src0_start_row), \
int remaining = copy_cnt; \ copy_cnt); \
while (remaining > 0) { \
int n = MIN(remaining, 32); \
HVX_Vector v_out = hvx_vmemu(tmp_ptr); \
HVX_Vector v_z = hvx_vmemu(src2_ptr); \
hvx_vec_store_u(dst_ptr, n * sizeof(float), hvx_vec_add_f32_f32(v_out, v_z)); \
dst_ptr += n; \
src2_ptr += n; \
tmp_ptr += n; \
remaining -= n; \
} \
} else { \ } else { \
hvx_copy_f32_ua((uint8_t *) &dst_col[src0_start_row], (uint8_t *) tmp, copy_cnt); \ hvx_copy_f32_ua((uint8_t *) &dst_col[src0_start_row], (uint8_t *) tmp, copy_cnt); \
} \ } \
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct_end); \
} \ } \
} }
@@ -1069,6 +1074,16 @@ static void hvx_mv_2d(unsigned int nth, unsigned int ith, void * data) {
// Prefill vtcm with 2x src0 rows // Prefill vtcm with 2x src0 rows
if (src0_start_row < src0_end_row) { if (src0_start_row < src0_end_row) {
if (src2) {
float * vtcm_src2_ptr = (float *) mmctx->vtcm_src2 + src0_start_row;
const float * src2_ptr = (const float *) src2->data + src0_start_row;
int slice_size = (int)src0_end_row - (int)src0_start_row;
if (slice_size > 0) {
dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr, src2_ptr),
slice_size * sizeof(float), slice_size * sizeof(float), slice_size * sizeof(float), 1);
dma_queue_pop_nowait(dma_queue);
}
}
for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) {
const uint32_t is0 = (ir0 - src0_start_row); const uint32_t is0 = (ir0 - src0_start_row);
if (is0 >= n_prefetch) { if (is0 >= n_prefetch) {
@@ -1114,27 +1129,21 @@ static void hvx_mv_2d(unsigned int nth, unsigned int ith, void * data) {
} }
int copy_cnt = src0_end_row - src0_start_row; int copy_cnt = src0_end_row - src0_start_row;
if (src2) { if (copy_cnt > 0) {
float * dst_ptr = &dst_col[src0_start_row]; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, src0_end_row);
const float * src2_ptr = (const float *) src2->data + src0_start_row; if (src2) {
float * tmp_ptr = tmp; hvx_add_f32_uaa((uint8_t *) &dst_col[src0_start_row],
int remaining = copy_cnt; (const uint8_t *) tmp,
while (remaining > 0) { (const uint8_t *) ((const float *) mmctx->vtcm_src2 + src0_start_row),
int n = MIN(remaining, 32); copy_cnt);
HVX_Vector v_out = hvx_vmemu(tmp_ptr); } else {
HVX_Vector v_z = hvx_vmemu(src2_ptr); hvx_copy_f32_ua((uint8_t *) &dst_col[src0_start_row], (uint8_t *) tmp, copy_cnt);
hvx_vec_store_u(dst_ptr, n * sizeof(float), hvx_vec_add_f32_f32(v_out, v_z));
dst_ptr += n;
src2_ptr += n;
tmp_ptr += n;
remaining -= n;
} }
} else { htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, src0_end_row);
hvx_copy_f32_ua((uint8_t *) &dst_col[src0_start_row], (uint8_t *) tmp, copy_cnt);
} }
} }
#define MMID_MATRIX_ROW(row_id, i1) matrix_rows[(row_id) * ids->ne[0] * ids->ne[1] + (i1)] #define MMID_MATRIX_ROW(row_id, i1) matrix_rows[(row_id) * mmctx->mapping_stride + (i1)]
static void hvx_mm_id(unsigned int nth, unsigned int ith, void * data) { static void hvx_mm_id(unsigned int nth, unsigned int ith, void * data) {
htp_matmul_preamble; htp_matmul_preamble;
@@ -1519,7 +1528,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) {
struct htp_mm_hvx_vtcm_layout L; struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads, htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads,
dst_row_size, src0_row_size, src1_row_size, kparams->n_prefetch, false, false, false); dst_row_size, src0_row_size, src1_row_size, src2 ? src2->nb[1] : 0, kparams->n_prefetch, false, false, false);
if (kparams->kernel_type == HTP_MM_KERNEL_HVX_F16_F16_VTCM || if (kparams->kernel_type == HTP_MM_KERNEL_HVX_F16_F16_VTCM ||
kparams->kernel_type == HTP_MM_KERNEL_HVX_F32_F32_VTCM || kparams->kernel_type == HTP_MM_KERNEL_HVX_F32_F32_VTCM ||
@@ -1551,6 +1560,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) {
uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base; uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1); mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0); mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
mmctx->vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src2);
mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst); mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
octx->src1_spad.src = NULL; octx->src1_spad.src = NULL;
@@ -2346,12 +2356,77 @@ static void dequantize_tiled_weight_chunk_to_fp16_tiles(
} }
} }
typedef struct {
float *dst;
const float *src2;
const __fp16 *vtcm_src;
uint32_t n_rows;
uint32_t n_cols;
uint32_t dst_stride;
uint32_t src2_stride;
uint32_t dst_cols;
struct fastdiv_values n_threads_div;
struct htp_thread_trace *traces;
struct htp_context *ctx;
} output_transfer_col_chunk_state_t;
static void transfer_output_chunk_col_chunk_worker_fn(unsigned int n, unsigned int i, void *data) {
(void) n;
output_transfer_col_chunk_state_t *st = (output_transfer_col_chunk_state_t *) data;
struct htp_thread_trace * tr = &st->traces[i];
uint32_t n_blocks = st->n_cols / 32;
uint32_t b_first = fastdiv(n_blocks * i, &st->n_threads_div);
uint32_t b_last = fastdiv(n_blocks * (i + 1), &st->n_threads_div);
uint32_t c_first = b_first * 32;
uint32_t c_last = b_last * 32;
uint32_t c_len = c_last - c_first;
if (c_len == 0) return;
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, c_first);
float *dst = st->dst + c_first;
const float *src2 = st->src2 ? (st->src2 + c_first) : NULL;
const __fp16 *vtcm_src = st->vtcm_src + b_first * HTP_MM_HMX_TILE_N_ELMS;
int chunk_dst_cols = (int)st->dst_cols - (int)c_first;
if (chunk_dst_cols > 0) {
transfer_output_chunk_fp16_to_fp32_col_chunk(
dst, src2, vtcm_src, 0, st->n_rows, c_len, st->n_cols,
st->dst_stride, st->src2_stride, (uint32_t)chunk_dst_cols
);
}
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, c_first);
}
static void transfer_output_chunk_threaded(struct htp_context *ctx, float *dst, const float *src2, const __fp16 *vtcm_src, static void transfer_output_chunk_threaded(struct htp_context *ctx, float *dst, const float *src2, const __fp16 *vtcm_src,
int n_rows, int n_cols, int dst_stride, uint32_t src2_stride, int dst_cols, int n_threads) { int n_rows, int n_cols, int dst_stride, uint32_t src2_stride, int dst_cols, int n_threads) {
assert(n_cols % HTP_MM_HMX_TILE_N_COLS == 0); assert(n_cols % HTP_MM_HMX_TILE_N_COLS == 0);
if (n_rows <= 0) return; if (n_rows <= 0) return;
uint32_t n_blocks = (uint32_t)n_cols / 32;
if (n_threads > 1 && n_blocks >= (uint32_t)n_threads) {
struct fastdiv_values n_threads_div = init_fastdiv_values(n_threads);
output_transfer_col_chunk_state_t col_state;
col_state.dst = dst;
col_state.src2 = src2;
col_state.vtcm_src = vtcm_src;
col_state.n_rows = (uint32_t)n_rows;
col_state.n_cols = (uint32_t)n_cols;
col_state.dst_stride = (uint32_t)dst_stride;
col_state.src2_stride = src2_stride;
col_state.dst_cols = (uint32_t)dst_cols;
col_state.n_threads_div = n_threads_div;
col_state.traces = ctx->trace;
col_state.ctx = ctx;
worker_pool_run_func(ctx->worker_pool, transfer_output_chunk_col_chunk_worker_fn, &col_state, n_threads);
return;
}
size_t n_tot_chunks = n_rows; size_t n_tot_chunks = n_rows;
size_t n_chunks_per_task = (n_threads == 1) ? n_tot_chunks : hmx_ceil_div(n_rows, n_threads); size_t n_chunks_per_task = (n_threads == 1) ? n_tot_chunks : hmx_ceil_div(n_rows, n_threads);
n_chunks_per_task = hex_align_up(n_chunks_per_task, 2); n_chunks_per_task = hex_align_up(n_chunks_per_task, 2);
@@ -3338,12 +3413,10 @@ int op_matmul(struct htp_ops_context * octx) {
static int hmx_mm_op_matmul_id( static int hmx_mm_op_matmul_id(
struct htp_ops_context * octx, struct htp_ops_context * octx,
struct htp_mm_context * mmctx, struct htp_mm_context * mmctx
const uint32_t * matrix_row_counts,
const struct mmid_row_mapping * matrix_rows,
void * mapping_buf,
bool must_free_mapping
) { ) {
const uint32_t * matrix_row_counts = mmctx->matrix_row_counts;
const struct mmid_row_mapping * matrix_rows = mmctx->matrix_rows;
htp_matmul_tensors_preamble; htp_matmul_tensors_preamble;
const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
const int n_ids = octx->src[2]->ne[0]; const int n_ids = octx->src[2]->ne[0];
@@ -3361,28 +3434,24 @@ static int hmx_mm_op_matmul_id(
nb11, nb12, nb11, nb12,
nb1, nb2, nb1, nb2,
(int) src0->nb[1], (int) src0->type, (int) src0->nb[1], (int) src0->type,
matrix_rows, cur_a, n_ids * octx->src[2]->ne[1]); matrix_rows, cur_a, mmctx->mapping_stride);
if (ret != 0) { if (ret != 0) {
FARF(ERROR, "HMX matmul failed for expert %u, error %d\n", cur_a, ret); FARF(ERROR, "HMX matmul failed for expert %u, error %d\n", cur_a, ret);
if (must_free_mapping) free(mapping_buf);
return HTP_STATUS_NO_SUPPORT; return HTP_STATUS_NO_SUPPORT;
} }
} }
if (must_free_mapping) free(mapping_buf);
return HTP_STATUS_OK; return HTP_STATUS_OK;
} }
static int hvx_mm_matmul_id( static int hvx_mm_matmul_id(
struct htp_ops_context * octx, struct htp_ops_context * octx,
struct htp_mm_context * mmctx, struct htp_mm_context * mmctx,
size_t src0_row_size_padded, work_queue_func_t hvx_mmid_task_func
uint32_t src1_nrows,
worker_callback_t matmul_id_job_func,
void * mapping_buf,
bool must_free_mapping
) { ) {
htp_matmul_tensors_preamble; htp_matmul_tensors_preamble;
const uint32_t src0_row_size_padded = mmctx->src0_row_size_padded;
const uint32_t src1_nrows = mmctx->src1_nrows;
struct htp_thread_trace * tr = &octx->ctx->trace[0]; struct htp_thread_trace * tr = &octx->ctx->trace[0];
htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0); htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);
@@ -3395,7 +3464,7 @@ static int hvx_mm_matmul_id(
const uint32_t nb = (ne10 + qk - 1) / qk; const uint32_t nb = (ne10 + qk - 1) / qk;
const uint32_t total_nb = src1_nrows * nb; const uint32_t total_nb = src1_nrows * nb;
worker_callback_t quant_task_func; work_queue_func_t quant_task_func;
uint32_t n_quant_tasks = 1; uint32_t n_quant_tasks = 1;
if (src1_nrows < octx->n_threads) { if (src1_nrows < octx->n_threads) {
n_quant_tasks = MIN(total_nb, octx->n_threads); n_quant_tasks = MIN(total_nb, octx->n_threads);
@@ -3416,7 +3485,7 @@ static int hvx_mm_matmul_id(
struct htp_mm_hvx_vtcm_layout L; struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads, htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads,
0, src0_row_size, src1_row_size, kparams->n_prefetch, true, false, false); 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, true, false, false);
size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes; size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;
@@ -3431,7 +3500,6 @@ static int hvx_mm_matmul_id(
// Make sure the reserved vtcm size is sufficient // Make sure the reserved vtcm size is sufficient
if (octx->ctx->vtcm_size < vtcm_size) { if (octx->ctx->vtcm_size < vtcm_size) {
FARF(ERROR, "matmul-id-%s : current VTCM reservation %zu is too small, needed %zu\n", mmctx->type, octx->ctx->vtcm_size, vtcm_size); FARF(ERROR, "matmul-id-%s : current VTCM reservation %zu is too small, needed %zu\n", mmctx->type, octx->ctx->vtcm_size, vtcm_size);
if (must_free_mapping) free(mapping_buf);
return HTP_STATUS_VTCM_TOO_SMALL; return HTP_STATUS_VTCM_TOO_SMALL;
} }
@@ -3461,12 +3529,78 @@ static int hvx_mm_matmul_id(
htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0);
worker_pool_run_func(octx->ctx->worker_pool, matmul_id_job_func, mmctx, octx->n_threads); worker_pool_run_func(octx->ctx->worker_pool, hvx_mmid_task_func, mmctx, octx->n_threads);
if (must_free_mapping) free(mapping_buf);
return HTP_STATUS_OK; return HTP_STATUS_OK;
} }
static inline void scan_expert_ids_n(
const struct htp_tensor * ids,
const uint32_t n_ids,
uint32_t n_as,
uint32_t * counts,
struct mmid_row_mapping * matrix_rows,
uint32_t mapping_stride
) {
const size_t ids_nb1 = ids->nb[1];
const uint8_t * ids_data = (const uint8_t *) ids->data;
for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
const int32_t * row_ptr = (const int32_t *) (ids_data + iid1 * ids_nb1);
for (uint32_t id = 0; id < n_ids; ++id) {
const int32_t i02 = row_ptr[id];
if (i02 < 0) {
continue;
}
assert(i02 < n_as);
if (matrix_rows) {
matrix_rows[i02 * mapping_stride + counts[i02]] = (struct mmid_row_mapping) { id, iid1 };
}
counts[i02] += 1;
}
}
}
static inline void scan_expert_ids(
const struct htp_tensor * ids,
uint32_t n_ids,
uint32_t n_as,
uint32_t * counts,
struct mmid_row_mapping * matrix_rows,
uint32_t mapping_stride
) {
const size_t ids_nb0 = ids->nb[0];
if (ids_nb0 == 4) {
switch (n_ids) {
case 8: scan_expert_ids_n(ids, 8, n_as, counts, matrix_rows, mapping_stride); break;
case 4: scan_expert_ids_n(ids, 4, n_as, counts, matrix_rows, mapping_stride); break;
case 2: scan_expert_ids_n(ids, 2, n_as, counts, matrix_rows, mapping_stride); break;
default: scan_expert_ids_n(ids, n_ids, n_as, counts, matrix_rows, mapping_stride); break;
}
} else {
// Strided fallback
const size_t ids_nb1 = ids->nb[1];
const uint8_t * ids_data = (const uint8_t *) ids->data;
for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
const int32_t * row_ptr = (const int32_t *) (ids_data + iid1 * ids_nb1);
for (uint32_t id = 0; id < n_ids; ++id) {
const int32_t i02 = *(const int32_t *) ((const uint8_t *) row_ptr + id * ids_nb0);
if (i02 < 0) {
continue;
}
assert(i02 < n_as);
if (matrix_rows) {
matrix_rows[i02 * mapping_stride + counts[i02]] = (struct mmid_row_mapping) { id, iid1 };
}
counts[i02] += 1;
}
}
}
}
int op_matmul_id(struct htp_ops_context * octx) { int op_matmul_id(struct htp_ops_context * octx) {
htp_matmul_tensors_preamble; htp_matmul_tensors_preamble;
@@ -3489,74 +3623,72 @@ int op_matmul_id(struct htp_ops_context * octx) {
const uint32_t src0_nrows = ne01; // per expert const uint32_t src0_nrows = ne01; // per expert
const uint32_t src1_nrows = ne11 * ne12 * ne13; const uint32_t src1_nrows = ne11 * ne12 * ne13;
worker_callback_t quant_task_func; mmctx->src0_nrows_per_thread = (src0_nrows + octx->n_threads - 1) / octx->n_threads;
worker_callback_t matmul_id_job_func = src1_nrows > 1 ? hvx_mm_id : hvx_mv_id; mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32);
// Compute src0_nrows_per_thread
mmctx->src0_nrows_per_thread = (src0_nrows + octx->n_threads - 1) / octx->n_threads;
mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32);
// row groups // row groups
const int n_ids = ids->ne[0]; // n_expert_used const int n_ids = ids->ne[0]; // n_expert_used
const int n_as = ne02; // n_expert const int n_as = ne02; // n_expert
size_t matrix_row_counts_size = n_as * sizeof(uint32_t); uint8_t * mapping_buf = octx->ctx->ddr_spad_base;
size_t matrix_row_map_size = n_as * ids->ne[0] * ids->ne[1] * sizeof(struct mmid_row_mapping); uint32_t mapping_stride = 1;
const size_t total_map_size = matrix_row_counts_size + matrix_row_map_size; uint32_t * matrix_row_counts = (uint32_t *) mapping_buf;
struct mmid_row_mapping * matrix_rows = NULL;
void * mapping_buf = NULL;
bool must_free_mapping = false;
if (octx->ctx->ddr_spad_base && total_map_size <= octx->ctx->ddr_spad_size) {
mapping_buf = octx->ctx->ddr_spad_base;
} else {
mapping_buf = memalign(128, total_map_size);
if (mapping_buf) {
must_free_mapping = true;
} else {
return HTP_STATUS_INTERNAL_ERR;
}
}
uint32_t * matrix_row_counts = (uint32_t *) mapping_buf;
struct mmid_row_mapping * matrix_rows = (struct mmid_row_mapping *) ((uint8_t *) mapping_buf + matrix_row_counts_size);
mmctx->matrix_row_counts = matrix_row_counts;
mmctx->matrix_rows = matrix_rows;
mmctx->mm_div_ne11 = kparams->div_ne11;
if (hvx_mm_init_vec_dot(mmctx, src0->type) != 0) {
if (must_free_mapping) free(mapping_buf);
return HTP_STATUS_NO_SUPPORT;
}
if (src1_nrows > 1) { if (src1_nrows > 1) {
// initialize matrix_row_counts and map const size_t matrix_row_counts_size = n_as * sizeof(uint32_t);
memset(matrix_row_counts, 0, n_as * sizeof(uint32_t)); assert(octx->ctx->ddr_spad_size >= matrix_row_counts_size);
// group rows by src0 matrix hex_l2fetch_block((const void *) ids->data, ids->ne[1] * ids->nb[1]);
for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) { // token idx
for (uint32_t id = 0; id < n_ids; ++id) { // expert idx
const int32_t i02 = *(const int32_t *) ((const uint8_t *) ids->data + iid1 * ids->nb[1] + id * ids->nb[0]);
if (i02 < 0) { memset(matrix_row_counts, 0, matrix_row_counts_size);
continue; scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, NULL, 0);
}
assert(i02 < n_as);
matrix_rows[i02 * n_ids * ids->ne[1] + matrix_row_counts[i02]] = (struct mmid_row_mapping) { id, iid1 }; uint32_t max_count = hvx_reduce_max_i32((const uint8_t *) matrix_row_counts, n_as);
matrix_row_counts[i02] += 1; mapping_stride = max_count > 0 ? max_count : 1;
size_t matrix_row_map_size = n_as * mapping_stride * sizeof(struct mmid_row_mapping);
const size_t total_map_size = matrix_row_counts_size + matrix_row_map_size;
if (total_map_size > octx->ctx->ddr_spad_size) {
mapping_buf = memalign(128, total_map_size);
if (!mapping_buf) {
return HTP_STATUS_INTERNAL_ERR;
} }
} }
matrix_row_counts = (uint32_t *) mapping_buf;
matrix_rows = (struct mmid_row_mapping *) (mapping_buf + matrix_row_counts_size);
memset(matrix_row_counts, 0, n_as * sizeof(uint32_t));
scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, matrix_rows, mapping_stride);
} }
mmctx->matrix_row_counts = matrix_row_counts;
mmctx->matrix_rows = matrix_rows;
mmctx->mapping_stride = mapping_stride;
mmctx->mm_div_ne11 = kparams->div_ne11;
mmctx->src0_row_size_padded = src0_row_size_padded;
mmctx->src1_nrows = src1_nrows;
htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0);
int s;
if (kparams->n_hmx) { if (kparams->n_hmx) {
return hmx_mm_op_matmul_id(octx, mmctx, matrix_row_counts, matrix_rows, mapping_buf, must_free_mapping); s = hmx_mm_op_matmul_id(octx, mmctx);
} else {
if (hvx_mm_init_vec_dot(mmctx, src0->type) == 0) {
s = hvx_mm_matmul_id(octx, mmctx, src1_nrows > 1 ? hvx_mm_id : hvx_mv_id);
} else {
s = HTP_STATUS_NO_SUPPORT;
}
} }
return hvx_mm_matmul_id(octx, mmctx, src0_row_size_padded, src1_nrows, matmul_id_job_func, mapping_buf, must_free_mapping); if (mapping_buf != octx->ctx->ddr_spad_base) {
free(mapping_buf);
}
return s;
} }
int op_matmul_qkv(struct htp_ops_context * octx) { int op_matmul_qkv(struct htp_ops_context * octx) {
@@ -3633,7 +3765,7 @@ int op_matmul_qkv(struct htp_ops_context * octx) {
struct htp_mm_hvx_vtcm_layout L; struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, src1->ne[0], src1_nrows, octx->n_threads, htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, src1->ne[0], src1_nrows, octx->n_threads,
0, src0_row_size, src1_row_size, kparams->n_prefetch, false, true, false); 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, false, true, false);
size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes; size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;
@@ -3778,7 +3910,7 @@ int op_matmul_ffn(struct htp_ops_context * octx) {
struct htp_mm_hvx_vtcm_layout L; struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, src1->ne[0], src1_nrows, octx->n_threads, htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, src1->ne[0], src1_nrows, octx->n_threads,
0, src0_row_size, src1_row_size, kparams->n_prefetch, false, false, true); 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, false, false, true);
size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes; size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;
ggml/src/ggml-hexagon/htp/matmul-ops.h
+2 -1
View File
@@ -460,6 +460,7 @@ static inline void htp_mm_hvx_vtcm_layout_build(
size_t dst_row_size, size_t dst_row_size,
size_t src0_row_size, size_t src0_row_size,
size_t src1_row_size, size_t src1_row_size,
size_t src2_row_size,
uint32_t n_prefetch, uint32_t n_prefetch,
bool is_matmul_id, bool is_matmul_id,
bool is_fused_qkv, bool is_fused_qkv,
@@ -467,7 +468,7 @@ static inline void htp_mm_hvx_vtcm_layout_build(
) { ) {
size_t src0_sz = 0; size_t src0_sz = 0;
size_t src1_sz = 0; size_t src1_sz = 0;
size_t src2_sz = 0; size_t src2_sz = src2_row_size > 0 ? htp_mm_round_up(src2_row_size, 128) : 0;
size_t src3_sz = 0; size_t src3_sz = 0;
size_t dst_sz = 0; size_t dst_sz = 0;
scripts/snapdragon/ggml-hexagon-profile.py
+237 -93
View File
@@ -6,6 +6,7 @@ import re
import argparse import argparse
import statistics import statistics
import logging import logging
import bisect
from typing import Any, Dict, List, Optional from typing import Any, Dict, List, Optional
from collections import defaultdict from collections import defaultdict
@@ -30,7 +31,7 @@ op_pattern = re.compile(
) )
trace_pattern = re.compile( trace_pattern = re.compile(
r"trace-op\s+(?P<op_name>[A-Z_0-9+]+):\s+thread\s+(?P<thread>\d+)\s+event\s+(?P<event>[A-Z_0-9\-]+)\s+info\s+(?P<info>\d+)\s+(?P<state>start|stop)\s+(?P<cycles>\d+)" r"trace-evt\s+(?P<event>[A-Z_0-9\-]+):\s+thread\s+(?P<thread>\d+)\s+info\s+(?P<info>\d+)\s+(?P<state>start|stop)\s+(?P<cycles>\d+)"
) )
logger = logging.getLogger("ggml-hexagon-profile") logger = logging.getLogger("ggml-hexagon-profile")
@@ -50,9 +51,13 @@ def normalize_event_name(evt_type):
class CycleUnwrapper: class CycleUnwrapper:
def __init__(self): def __init__(self, initial_val=None):
self.last_raw = None if initial_val is not None:
self.high_part = 0 self.last_raw = initial_val & 0xFFFFFFFF
self.high_part = initial_val & 0xFFFFFFFF00000000
else:
self.last_raw = None
self.high_part = 0
def unwrap(self, raw): def unwrap(self, raw):
if self.last_raw is None: if self.last_raw is None:
@@ -78,10 +83,12 @@ def parse_log(file_path, pmu_index=None):
sys.exit(1) sys.exit(1)
all_ops: List[Dict[str, Any]] = [] all_ops: List[Dict[str, Any]] = []
all_traces: List[Dict[str, Any]] = []
current_op: Optional[Dict[str, Any]] = None current_op: Optional[Dict[str, Any]] = None
timestamp_pattern = re.compile(r"^(?P<min>\d+)\.(?P<sec>\d+)\.(?P<ms>\d+)\.(?P<us>\d+)\s+[A-Z]\s+") timestamp_pattern = re.compile(r"^(?P<min>\d+)\.(?P<sec>\d+)\.(?P<ms>\d+)\.(?P<us>\d+)\s+[A-Z]\s+")
unwrapper = CycleUnwrapper() unwrapper = None
trace_unwrapper = None
for line in f: for line in f:
ts_match = timestamp_pattern.match(line) ts_match = timestamp_pattern.match(line)
@@ -100,6 +107,7 @@ def parse_log(file_path, pmu_index=None):
if not prefix_match: if not prefix_match:
continue continue
names = parts[1]
if len(parts) == 7: if len(parts) == 7:
dims, types, timings = parts[2], parts[3], parts[6] dims, types, timings = parts[2], parts[3], parts[6]
elif len(parts) == 6: elif len(parts) == 6:
@@ -120,6 +128,7 @@ def parse_log(file_path, pmu_index=None):
op_match = op_pattern.search(line) op_match = op_pattern.search(line)
if op_match: if op_match:
op_name = op_match.group('op_name') op_name = op_match.group('op_name')
names = ""
dims = op_match.group('dims').strip() dims = op_match.group('dims').strip()
types = op_match.group('types').strip() types = op_match.group('types').strip()
else: else:
@@ -136,24 +145,31 @@ def parse_log(file_path, pmu_index=None):
except (ValueError, IndexError): except (ValueError, IndexError):
pmu_val = None pmu_val = None
evt_raw = op_match.group('evt') if 'evt' in op_match.groupdict() else None
evt_val = None evt_val = None
if evt_raw: evt_val = None
if types.startswith("evt-cnt "):
try: try:
evt_val = [int(x.strip()) for x in evt_raw.split(',')] evt_val = [int(x.strip()) for x in types[8:].split(',')]
except ValueError: except ValueError:
evt_val = None evt_val = None
cycles_start_raw = op_match.group('start') cycles_start_raw = op_match.group('start')
unwrapped_cycles_start = None unwrapped_cycles_start = None
if cycles_start_raw: if op_name == "OPBATCH":
unwrapped_cycles_start = unwrapper.unwrap(int(cycles_start_raw)) if cycles_start_raw:
unwrapped_cycles_start = int(cycles_start_raw)
unwrapper = CycleUnwrapper(unwrapped_cycles_start)
trace_unwrapper = CycleUnwrapper(unwrapped_cycles_start)
else:
if cycles_start_raw and unwrapper is not None:
unwrapped_cycles_start = unwrapper.unwrap(int(cycles_start_raw))
idx = line.find("profile-op ") idx = line.find("profile-op ")
op_text = line[idx + 11:].strip() if idx != -1 else line.strip() op_text = line[idx + 11:].strip() if idx != -1 else line.strip()
current_op = { current_op = {
'name': op_name, 'name': op_name,
'names': names,
'dims': dims, 'dims': dims,
'types': types, 'types': types,
'op_text': op_text, 'op_text': op_text,
@@ -170,110 +186,239 @@ def parse_log(file_path, pmu_index=None):
continue continue
trace_match = trace_pattern.search(line) trace_match = trace_pattern.search(line)
if trace_match and current_op: if trace_match:
if trace_match.group('op_name') == current_op['name']: raw_cyc = int(trace_match.group('cycles'))
raw_cyc = int(trace_match.group('cycles')) unwrapped_cyc = None
current_op['trace_events'].append({ if trace_unwrapper is not None:
'thread': int(trace_match.group('thread')), unwrapped_cyc = trace_unwrapper.unwrap(raw_cyc)
'event': trace_match.group('event'), all_traces.append({
'info': int(trace_match.group('info')), 'thread': int(trace_match.group('thread')),
'cycles': raw_cyc, 'event': trace_match.group('event'),
'unwrapped_cycles': unwrapper.unwrap(raw_cyc), 'info': int(trace_match.group('info')),
'state': trace_match.group('state') 'cycles': raw_cyc,
}) 'unwrapped_cycles': unwrapped_cyc,
'state': trace_match.group('state')
})
f.close() f.close()
# Assign start/end cycles to all ops
for op in all_ops:
op['start_cycles'] = op['unwrapped_cycles_start']
op['end_cycles'] = op['start_cycles'] + op['cycles'] if op['start_cycles'] is not None else None
# Filter ops with valid start_cycles
valid_ops = [op for op in all_ops if op['start_cycles'] is not None and op['end_cycles'] is not None]
# Separate OPBATCH ops from other ops
opbatch_ops = [op for op in valid_ops if op['name'] == "OPBATCH"]
other_ops = [op for op in valid_ops if op['name'] != "OPBATCH"]
# Sort them by start_cycles to enable binary search
opbatch_ops.sort(key=lambda op: op['start_cycles'])
other_ops.sort(key=lambda op: op['start_cycles'])
opbatch_starts = [op['start_cycles'] for op in opbatch_ops]
other_starts = [op['start_cycles'] for op in other_ops]
# Map trace events to any operator whose cycles contain them
for e in all_traces:
cyc = e['unwrapped_cycles']
if cyc is None:
continue
# Map to OPBATCH
idx = bisect.bisect_right(opbatch_starts, cyc) - 1
if idx >= 0:
op = opbatch_ops[idx]
if op['start_cycles'] <= cyc <= op['end_cycles']:
op['trace_events'].append(e)
# Map to other ops
idx = bisect.bisect_right(other_starts, cyc) - 1
if idx >= 0:
op = other_ops[idx]
if op['start_cycles'] <= cyc <= op['end_cycles']:
op['trace_events'].append(e)
return all_ops return all_ops
def print_ascii_timeline(op_name, dims, types, usec, cycles, events, evt_val=None): def print_bubbles_timeline(op):
evt_str = "" op_name = op['name']
if evt_val: dims = op['dims']
evt_str = " - evt [" + ",".join(str(x) for x in evt_val) + "]" types = op['types']
usec = op['usec']
cycles = op['cycles']
events = op['trace_events']
logger.info("=" * 100) logger.info("=" * 100)
logger.info(f"{op_name} ({dims} : {types}) - {usec} usec {cycles} cycles{evt_str}") logger.info(f"{op_name} ({dims} : {types}) - {usec} usec {cycles} cycles")
logger.info("=" * 100) logger.info("=" * 100)
events = sorted(events, key=lambda e: e['cycles'])
if not events: if not events:
logger.info(" No trace events recorded.") logger.info(" No trace events recorded.")
return return
min_cycles = events[0]['cycles'] # Identify start and end cycles for this operator
op_start = op['start_cycles']
op_end = op['end_cycles']
if op_start is None or op_end is None:
logger.info(" Cannot analyze bubbles: missing start/end cycle counts.")
return
logger.info("Cycles %-30s" % "EventDetails" + " ".join(f"T{i:<2}" for i in range(10)) + " HMX") batch_duration = op_end - op_start
logger.info("-" * 100) if batch_duration <= 0:
logger.info(" Cannot analyze bubbles: batch duration is 0.")
thread_stacks = [[] for _ in range(11)] return
# Group events by (thread, track_type)
tracks = defaultdict(list)
for e in events: for e in events:
t = e['thread'] t = e['thread']
if t < 0 or t > 10: is_dma = (normalize_event_name(e['event']) == 'DMA')
continue track_type = 'dma' if is_dma else 'compute'
tracks[(t, track_type)].append(e)
if e['cycles'] >= min_cycles: active_threads = sorted(list(set(t for (t, track_type) in tracks.keys())))
rel_cycles = e['cycles'] - min_cycles if not active_threads:
else: logger.info(" No active threads in trace.")
rel_cycles = (e['cycles'] + 0x100000000) - min_cycles return
state = e['state'] bubble_threshold = 10000 # 10k cycles
evt_type = e['event']
# Determine char representing the event thread_stats = {}
norm_evt = normalize_event_name(evt_type) for t in active_threads:
char = '?' thread_stats[t] = {
if norm_evt == 'V-COMP': 'compute_idle_cycles': batch_duration,
char = 'V' 'compute_idle_pct': 100.0,
elif norm_evt == 'M-COMP': 'compute_bubbles': [],
char = 'H'
elif norm_evt == 'A-QUANT':
char = 'Q'
elif norm_evt == 'A-PREP':
char = 'A'
elif norm_evt == 'Q-PREP':
char = 'q'
elif norm_evt == 'K-PREP':
char = 'k'
elif norm_evt == 'V-PREP':
char = 'v'
elif norm_evt == 'W-DEQUANT':
char = 'D'
elif norm_evt == 'O-PROC':
char = 'O'
elif norm_evt == 'W-PREP':
char = 'P'
elif norm_evt == 'DMA':
char = 'M'
if state == 'start': 'dma_idle_cycles': batch_duration,
thread_stacks[t].append(char) 'dma_idle_pct': 100.0,
elif state == 'stop': 'dma_bubbles': []
if thread_stacks[t]: }
if thread_stacks[t][-1] == char:
thread_stacks[t].pop()
elif char in thread_stacks[t]:
thread_stacks[t].remove(char)
else:
thread_stacks[t].pop()
cols = [] total_compute_idle_pct = 0.0
for i in range(11): total_dma_idle_pct = 0.0
if thread_stacks[i]:
cols.append(f"[{thread_stacks[i][-1]}]") for t in active_threads:
for track_type in ['compute', 'dma']:
key = (t, track_type)
track_events = tracks.get(key, [])
if not track_events:
gaps = [(op_start, op_end)]
idle_cycles = batch_duration
else: else:
cols.append(" | ") track_events = sorted(track_events, key=lambda e: e.get('unwrapped_cycles') or e['cycles'])
evt_desc = f"T{t}: {evt_type} {state} ({e['info']})" active_intervals = []
logger.info(f"{rel_cycles:10d} %-30s" % evt_desc + " ".join(cols[:10]) + " " + cols[10]) active_count = 0
curr_start = None
for e in track_events:
cyc = e.get('unwrapped_cycles') or e['cycles']
cyc = max(op_start, min(op_end, cyc))
state = e['state']
if state == 'start':
if active_count == 0:
curr_start = cyc
active_count += 1
elif state == 'stop':
if active_count > 0:
active_count -= 1
if active_count == 0:
active_intervals.append((curr_start, cyc))
else:
active_intervals.append((op_start, cyc))
if active_count > 0 and curr_start is not None:
active_intervals.append((curr_start, op_end))
# Merge intervals
active_intervals.sort(key=lambda x: x[0])
merged_intervals = []
for start, end in active_intervals:
if not merged_intervals:
merged_intervals.append([start, end])
else:
last_start, last_end = merged_intervals[-1]
if start <= last_end:
merged_intervals[-1][1] = max(last_end, end)
else:
merged_intervals.append([start, end])
# Calculate gaps
gaps = []
curr_time = op_start
for start, end in merged_intervals:
if start > curr_time:
gaps.append((curr_time, start))
curr_time = max(curr_time, end)
if curr_time < op_end:
gaps.append((curr_time, op_end))
idle_cycles = sum(end - start for start, end in gaps)
idle_pct = (idle_cycles / batch_duration) * 100.0
bubbles = []
for start, end in gaps:
dur = end - start
if dur >= bubble_threshold:
bubbles.append((start, end, dur))
if track_type == 'compute':
thread_stats[t]['compute_idle_cycles'] = idle_cycles
thread_stats[t]['compute_idle_pct'] = idle_pct
thread_stats[t]['compute_bubbles'] = bubbles
total_compute_idle_pct += idle_pct
else:
thread_stats[t]['dma_idle_cycles'] = idle_cycles
thread_stats[t]['dma_idle_pct'] = idle_pct
thread_stats[t]['dma_bubbles'] = bubbles
total_dma_idle_pct += idle_pct
avg_compute_idle = total_compute_idle_pct / len(active_threads)
avg_dma_idle = total_dma_idle_pct / len(active_threads)
logger.info(" Combined Idle Statistics:")
logger.info(f" Active Threads : {', '.join(str(t) for t in active_threads)}")
logger.info(f" Avg Thread Compute IDLE : {avg_compute_idle:.1f}%")
logger.info(f" Avg Thread DMA IDLE : {avg_dma_idle:.1f}%")
logger.info("-" * 100) logger.info("-" * 100)
logger.info(" Per-Thread Idle Analysis:")
for t in active_threads:
stats = thread_stats[t]
thread_name = f"Thread {t:<2} (HVX)" if t != 10 else "Thread 10 (HMX)"
logger.info(f" {thread_name} -> Compute Idle: {stats['compute_idle_pct']:.1f}% | DMA Idle: {stats['dma_idle_pct']:.1f}%")
def print_ascii_summary(op_name, dims, types, usec, cycles, events, evt_val=None): all_bubbles = []
evt_str = "" for t in active_threads:
if evt_val: stats = thread_stats[t]
evt_str = " - evt [" + ",".join(str(x) for x in evt_val) + "]" for start, end, dur in stats['compute_bubbles']:
pct = (dur / batch_duration) * 100.0
all_bubbles.append((dur, f"Thread {t} Compute: bubble of {dur} cycles ({pct:.1f}%) at {start - op_start} to {end - op_start}"))
for start, end, dur in stats['dma_bubbles']:
pct = (dur / batch_duration) * 100.0
all_bubbles.append((dur, f"Thread {t} DMA : bubble of {dur} cycles ({pct:.1f}%) at {start - op_start} to {end - op_start}"))
if all_bubbles:
logger.info("-" * 100)
logger.info(f" Significant Bubbles (>= {bubble_threshold} cycles):")
all_bubbles.sort(key=lambda x: x[0], reverse=True)
for dur, desc in all_bubbles[:15]:
logger.info(f" {desc}")
else:
logger.info("-" * 100)
logger.info(f" No significant bubbles detected (all idle gaps < {bubble_threshold} cycles).")
def print_ascii_summary(op_name, dims, types, usec, cycles, events):
logger.info("=" * 100) logger.info("=" * 100)
logger.info(f"{op_name} ({dims} : {types}) - {usec} usec {cycles} cycles{evt_str}") logger.info(f"{op_name} ({dims} : {types}) - {usec} usec {cycles} cycles")
logger.info("=" * 100) logger.info("=" * 100)
events = sorted(events, key=lambda e: e['cycles']) events = sorted(events, key=lambda e: e['cycles'])
@@ -415,8 +560,8 @@ def main():
parser.add_argument("--pmu-index", type=int) parser.add_argument("--pmu-index", type=int)
parser.add_argument("--pmu-name", type=str) parser.add_argument("--pmu-name", type=str)
parser.add_argument("--width", action='append', default=['dims:40'], help="Override column width, e.g. --width dims:50") parser.add_argument("--width", action='append', default=['dims:40'], help="Override column width, e.g. --width dims:50")
parser.add_argument("--timeline", type=str, nargs='?', const='summary', choices=["summary", "diagram"], parser.add_argument("--timeline", type=str, nargs='?', const='summary', choices=["summary", "bubbles"],
help="Output ASCII art event summary or timing diagram (default: summary)") help="Output ASCII art event summary or thread idle bubble analysis (default: summary)")
parser.add_argument("--filter", type=str, help="Regex filter matching against the original profile-op line") parser.add_argument("--filter", type=str, help="Regex filter matching against the original profile-op line")
group = parser.add_mutually_exclusive_group() group = parser.add_mutually_exclusive_group()
@@ -457,12 +602,11 @@ def main():
ops = ops[-args.tail:] ops = ops[-args.tail:]
if args.timeline: if args.timeline:
logger.info(f"\n# ASCII Timing {args.timeline.capitalize()}\n")
for op in ops: for op in ops:
if args.timeline == "summary": if args.timeline == "summary":
print_ascii_summary(op['name'], op['dims'], op['types'], op['usec'], op['cycles'], op['trace_events'], op.get('evt_val')) print_ascii_summary(op['name'], op['dims'], op['types'], op['usec'], op['cycles'], op['trace_events'])
elif args.timeline == "diagram": elif args.timeline == "bubbles":
print_ascii_timeline(op['name'], op['dims'], op['types'], op['usec'], op['cycles'], op['trace_events'], op.get('evt_val')) print_bubbles_timeline(op)
else: else:
generate_report(ops, args.top, overrides, args.sort, pmu_name=final_pmu_name) generate_report(ops, args.top, overrides, args.sort, pmu_name=final_pmu_name)
scripts/snapdragon/ggml-hexagon-trace.py
+130 -40
View File
@@ -6,6 +6,7 @@ import re
import argparse import argparse
import statistics import statistics
import logging import logging
import bisect
from typing import Any, Dict, List, Optional from typing import Any, Dict, List, Optional
from collections import defaultdict from collections import defaultdict
@@ -16,11 +17,11 @@ op_pattern = re.compile(
) )
trace_pattern = re.compile( trace_pattern = re.compile(
r"trace-op\s+(?P<op_name>[A-Z_0-9+]+):\s+thread\s+(?P<thread>\d+)\s+event\s+(?P<event>[A-Z_0-9\-]+)\s+info\s+(?P<info>\d+)\s+(?P<state>start|stop)\s+(?P<cycles>\d+)" r"trace-evt\s+(?P<event>[A-Z_0-9\-]+):\s+thread\s+(?P<thread>\d+)\s+info\s+(?P<info>\d+)\s+(?P<state>start|stop)\s+(?P<cycles>\d+)"
) )
def normalize_event_name(evt_type): def normalize_event_name(evt_type, info=0):
if evt_type == "HVX_COMP": if evt_type == "HVX_COMP":
return "V-COMP" return "V-COMP"
if evt_type == "HMX_COMP": if evt_type == "HMX_COMP":
@@ -32,9 +33,13 @@ def normalize_event_name(evt_type):
class CycleUnwrapper: class CycleUnwrapper:
def __init__(self): def __init__(self, initial_val=None):
self.last_raw = None if initial_val is not None:
self.high_part = 0 self.last_raw = initial_val & 0xFFFFFFFF
self.high_part = initial_val & 0xFFFFFFFF00000000
else:
self.last_raw = None
self.high_part = 0
def unwrap(self, raw): def unwrap(self, raw):
if self.last_raw is None: if self.last_raw is None:
@@ -60,8 +65,10 @@ def parse_log(file_path):
sys.exit(1) sys.exit(1)
all_ops: List[Dict[str, Any]] = [] all_ops: List[Dict[str, Any]] = []
all_traces: List[Dict[str, Any]] = []
current_op: Optional[Dict[str, Any]] = None current_op: Optional[Dict[str, Any]] = None
unwrapper = CycleUnwrapper() unwrapper = None
trace_unwrapper = None
line_idx = 0 line_idx = 0
for line in f: for line in f:
@@ -73,6 +80,7 @@ def parse_log(file_path):
if not prefix_match: if not prefix_match:
continue continue
names = parts[1]
if len(parts) == 7: if len(parts) == 7:
dims, types, strides, params, timings = parts[2], parts[3], parts[4], parts[5], parts[6] dims, types, strides, params, timings = parts[2], parts[3], parts[4], parts[5], parts[6]
elif len(parts) == 6: elif len(parts) == 6:
@@ -93,6 +101,7 @@ def parse_log(file_path):
op_match = op_pattern.search(line) op_match = op_pattern.search(line)
if op_match: if op_match:
op_name = op_match.group('op_name') op_name = op_match.group('op_name')
names = ""
dims = op_match.group('dims').strip() if op_match.group('dims') else '' dims = op_match.group('dims').strip() if op_match.group('dims') else ''
types = op_match.group('types').strip() if op_match.group('types') else '' types = op_match.group('types').strip() if op_match.group('types') else ''
strides = op_match.group('strides').strip() if op_match.group('strides') else '' strides = op_match.group('strides').strip() if op_match.group('strides') else ''
@@ -103,18 +112,30 @@ def parse_log(file_path):
if op_match: if op_match:
cycles_start_raw = op_match.group('start') cycles_start_raw = op_match.group('start')
unwrapped_cycles_start = None unwrapped_cycles_start = None
if cycles_start_raw: if op_name == "OPBATCH":
unwrapped_cycles_start = unwrapper.unwrap(int(cycles_start_raw)) if cycles_start_raw:
unwrapped_cycles_start = int(cycles_start_raw)
unwrapper = CycleUnwrapper(unwrapped_cycles_start)
trace_unwrapper = CycleUnwrapper(unwrapped_cycles_start)
else:
if cycles_start_raw and unwrapper is not None:
unwrapped_cycles_start = unwrapper.unwrap(int(cycles_start_raw))
idx = line.find("profile-op ") idx = line.find("profile-op ")
op_text = line[idx + 11:].strip() if idx != -1 else line.strip() op_text = line[idx + 11:].strip() if idx != -1 else line.strip()
evt_str = None
if types.startswith("evt-cnt "):
evt_str = types[8:].strip()
current_op = { current_op = {
'name': op_name, 'name': op_name,
'names': names,
'dims': dims, 'dims': dims,
'types': types, 'types': types,
'strides': strides, 'strides': strides,
'params': params, 'params': params,
'evt': evt_str,
'op_text': op_text, 'op_text': op_text,
'usec': int(op_match.group('usec')), 'usec': int(op_match.group('usec')),
'cycles': int(op_match.group('cycles')), 'cycles': int(op_match.group('cycles')),
@@ -127,20 +148,22 @@ def parse_log(file_path):
continue continue
trace_match = trace_pattern.search(line) trace_match = trace_pattern.search(line)
if trace_match and current_op: if trace_match:
if trace_match.group('op_name') == current_op['name']: raw_cyc = int(trace_match.group('cycles'))
raw_cyc = int(trace_match.group('cycles')) unwrapped_cyc = None
current_op['trace_events'].append({ if trace_unwrapper is not None:
'thread': int(trace_match.group('thread')), unwrapped_cyc = trace_unwrapper.unwrap(raw_cyc)
'event': trace_match.group('event'), all_traces.append({
'info': int(trace_match.group('info')), 'thread': int(trace_match.group('thread')),
'cycles': raw_cyc, 'event': trace_match.group('event'),
'unwrapped_cycles': unwrapper.unwrap(raw_cyc), 'info': int(trace_match.group('info')),
'state': trace_match.group('state') 'cycles': raw_cyc,
}) 'unwrapped_cycles': unwrapped_cyc,
'state': trace_match.group('state')
})
f.close() f.close()
return all_ops return all_ops, all_traces
# --- Simple protobuf encoder --- # --- Simple protobuf encoder ---
@@ -246,7 +269,7 @@ def write_trace_packet_to_file(f, packet_bytes):
# --- End Protobuf Encoder --- # --- End Protobuf Encoder ---
def generate_perfetto_trace(filtered_ops, output_path): def generate_perfetto_trace(filtered_ops, trace_events, output_path):
if not filtered_ops: if not filtered_ops:
logger.warning("No operators found after filtering.") logger.warning("No operators found after filtering.")
return return
@@ -269,14 +292,12 @@ def generate_perfetto_trace(filtered_ops, output_path):
# Process events # Process events
completed_events = [] completed_events = []
for op in filtered_ops: if trace_events:
events = op['trace_events'] trace_events = sorted(trace_events, key=lambda e: e['unwrapped_cycles'])
if not events: one_usec_cycles = max(avg_freq_mhz, 1.0)
continue
events = sorted(events, key=lambda e: e['unwrapped_cycles'])
active_starts = {} active_starts = {}
for e in events: for e in trace_events:
t = e['thread'] t = e['thread']
evt = e['event'] evt = e['event']
info = e['info'] info = e['info']
@@ -285,6 +306,17 @@ def generate_perfetto_trace(filtered_ops, output_path):
key = (t, evt, info) key = (t, evt, info)
if state == 'start': if state == 'start':
# Handle missing stop (start followed by another start)
if key in active_starts:
prev_start = active_starts[key]
completed_events.append({
'thread': t,
'event': evt,
'info': info,
'start_cyc': prev_start,
'end_cyc': prev_start + one_usec_cycles,
'missing_stop': True,
})
active_starts[key] = cyc active_starts[key] = cyc
elif state == 'stop': elif state == 'stop':
if key in active_starts: if key in active_starts:
@@ -296,8 +328,29 @@ def generate_perfetto_trace(filtered_ops, output_path):
'info': info, 'info': info,
'start_cyc': start_cyc, 'start_cyc': start_cyc,
'end_cyc': cyc, 'end_cyc': cyc,
'op_name': op['name']
}) })
else:
# Handle missing start (stop without start)
completed_events.append({
'thread': t,
'event': evt,
'info': info,
'start_cyc': cyc - one_usec_cycles,
'end_cyc': cyc,
'missing_start': True,
})
# Clear remaining unmatched starts
for key, start_cyc in active_starts.items():
t, evt, info = key
completed_events.append({
'thread': t,
'event': evt,
'info': info,
'start_cyc': start_cyc,
'end_cyc': start_cyc + one_usec_cycles,
'missing_stop': True,
})
completed_events.sort(key=lambda e: e['start_cyc']) completed_events.sort(key=lambda e: e['start_cyc'])
@@ -316,7 +369,7 @@ def generate_perfetto_trace(filtered_ops, output_path):
ts = e['ts_ns'] ts = e['ts_ns']
dur = e['dur_ns'] dur = e['dur_ns']
norm_evt = normalize_event_name(evt) norm_evt = normalize_event_name(evt, e['info'])
if norm_evt == "DMA": if norm_evt == "DMA":
track_key = (t, "DMA") track_key = (t, "DMA")
elif t == 10: elif t == 10:
@@ -343,7 +396,7 @@ def generate_perfetto_trace(filtered_ops, output_path):
evt = e['event'] evt = e['event']
slot = e['slot'] slot = e['slot']
norm_evt = normalize_event_name(evt) norm_evt = normalize_event_name(evt, e['info'])
if norm_evt == "DMA": if norm_evt == "DMA":
track_evt = "DMA" track_evt = "DMA"
evt_id = 1 evt_id = 1
@@ -421,18 +474,26 @@ def generate_perfetto_trace(filtered_ops, output_path):
for op in filtered_ops: for op in filtered_ops:
op_start_ns = int(round(((op['start_cycles'] - global_min_cyc) / avg_freq_mhz) * 1000)) op_start_ns = int(round(((op['start_cycles'] - global_min_cyc) / avg_freq_mhz) * 1000))
op_dur_ns = int(round((op['cycles'] / avg_freq_mhz) * 1000)) op_dur_ns = int(round((op['cycles'] / avg_freq_mhz) * 1000))
if op_start_ns < last_op_end_ns: if op['name'] != "OPBATCH":
op_start_ns = last_op_end_ns if op_start_ns < last_op_end_ns:
clamped_dur = max(op_dur_ns, 100) # Clamp to 100ns (0.1us) op_start_ns = last_op_end_ns
clamped_dur = max(op_dur_ns, 100) # Clamp to 100ns (0.1us)
last_op_end_ns = op_start_ns + clamped_dur
else:
clamped_dur = max(op_dur_ns, 100)
# Debug annotations for Ops # Debug annotations for Ops
debug_annots = [] debug_annots = []
if 'line_num' in op: if 'line_num' in op:
debug_annots.append(make_debug_annotation("line", int_val=op['line_num'])) debug_annots.append(make_debug_annotation("line", int_val=op['line_num']))
if 'strides' in op and op['strides']: if 'names' in op and op['names'] and op['names'] != '----':
debug_annots.append(make_debug_annotation("names", string_val=op['names']))
if 'strides' in op and op['strides'] and op['strides'] != '----':
debug_annots.append(make_debug_annotation("strides", string_val=op['strides'])) debug_annots.append(make_debug_annotation("strides", string_val=op['strides']))
if 'params' in op and op['params'] and op['params'] != '----': if 'params' in op and op['params'] and op['params'] != '----':
debug_annots.append(make_debug_annotation("params", string_val=op['params'])) debug_annots.append(make_debug_annotation("params", string_val=op['params']))
if 'evt' in op and op['evt']:
debug_annots.append(make_debug_annotation("evt", string_val=op['evt']))
# Slice Begin # Slice Begin
evt_begin = make_track_event(1, 2, name=f"{op['name']} ({op['dims']})", category="operator", debug_annotations=debug_annots) evt_begin = make_track_event(1, 2, name=f"{op['name']} ({op['dims']})", category="operator", debug_annotations=debug_annots)
@@ -444,15 +505,21 @@ def generate_perfetto_trace(filtered_ops, output_path):
packet_end = make_trace_packet(op_start_ns + clamped_dur, track_event=evt_end) packet_end = make_trace_packet(op_start_ns + clamped_dur, track_event=evt_end)
write_trace_packet_to_file(f, packet_end) write_trace_packet_to_file(f, packet_end)
last_op_end_ns = op_start_ns + clamped_dur
# Emit Thread Trace Events # Emit Thread Trace Events
for e in completed_events: for e in completed_events:
norm_name = normalize_event_name(e['event']) norm_name = normalize_event_name(e['event'], e['info'])
name = f"DMA {e['info']}" if norm_name == "DMA" else norm_name name = f"DMA {e['info']}" if norm_name == "DMA" else norm_name
if e.get('missing_start') or e.get('missing_stop'):
name += "!"
debug_annots = []
if e.get('missing_start'):
debug_annots.append(make_debug_annotation("missing_start", string_val="true"))
if e.get('missing_stop'):
debug_annots.append(make_debug_annotation("missing_stop", string_val="true"))
# Slice Begin # Slice Begin
evt_begin = make_track_event(1, e['uuid'], name=name, category="trace") evt_begin = make_track_event(1, e['uuid'], name=name, category="trace", debug_annotations=debug_annots if debug_annots else None)
packet_begin = make_trace_packet(e['ts_ns'], track_event=evt_begin) packet_begin = make_trace_packet(e['ts_ns'], track_event=evt_begin)
write_trace_packet_to_file(f, packet_begin) write_trace_packet_to_file(f, packet_begin)
@@ -477,7 +544,7 @@ def main():
args = parser.parse_args() args = parser.parse_args()
logging.basicConfig(level=logging.INFO, format='%(message)s') logging.basicConfig(level=logging.INFO, format='%(message)s')
ops = parse_log(args.logfile) ops, traces = parse_log(args.logfile)
if args.filter: if args.filter:
try: try:
@@ -492,7 +559,30 @@ def main():
elif args.tail is not None: elif args.tail is not None:
ops = ops[-args.tail:] ops = ops[-args.tail:]
generate_perfetto_trace(ops, args.output) if args.filter or args.head is not None or args.tail is not None:
valid_ranges = []
for op in ops:
start_cyc = op['unwrapped_cycles_start']
end_cyc = start_cyc + op['cycles'] if start_cyc is not None else None
if start_cyc is not None and end_cyc is not None:
valid_ranges.append((start_cyc, end_cyc))
valid_ranges.sort(key=lambda r: r[0])
range_starts = [r[0] for r in valid_ranges]
filtered_traces = []
for e in traces:
cyc = e['unwrapped_cycles']
if cyc is None:
continue
idx = bisect.bisect_right(range_starts, cyc) - 1
if idx >= 0:
start, end = valid_ranges[idx]
if start <= cyc <= end:
filtered_traces.append(e)
traces = filtered_traces
generate_perfetto_trace(ops, traces, args.output)
if __name__ == "__main__": if __name__ == "__main__":