Commit dca216ad6 for llama.cpp
commit dca216ad63c586f0db465324f785c5a418faef58
Author: Pranav Dhinakar <pdhinaka@qti.qualcomm.com>
Date: Sat Nov 14 15:03:31 2026 -0800
hexagon: add TIMESTEP_EMBEDDING HVX kernel (#32228)
* hexagon: add TIMESTEP_EMBEDDING HVX kernel
* hex-timestep: add dma, improve kparams, and work-splits, and tracing
* timestep: add testcases that cover WAN2.2
* hex-timestep: relax the gates to cont_rows
---------
Co-authored-by: Max Krasnyansky <maxk@qti.qualcomm.com>
diff --git a/ggml/src/ggml-hexagon/htp-kparams.cpp b/ggml/src/ggml-hexagon/htp-kparams.cpp
index 4c771946a..6d365b811 100644
--- a/ggml/src/ggml-hexagon/htp-kparams.cpp
+++ b/ggml/src/ggml-hexagon/htp-kparams.cpp
@@ -65,6 +65,7 @@ htp_op_code htp_op_remap(const struct ggml_tensor * t) {
case GGML_OP_CONV_2D_DW: return HTP_OP_CONV_2D_DW;
case GGML_OP_CONV_3D: return HTP_OP_CONV_3D;
case GGML_OP_ROLL: return HTP_OP_ROLL;
+ case GGML_OP_TIMESTEP_EMBEDDING: return HTP_OP_TIMESTEP_EMBEDDING;
case GGML_OP_UPSCALE: return HTP_OP_UPSCALE;
case GGML_OP_LIGHTNING_INDEXER: return HTP_OP_LIGHTNING_INDEXER;
@@ -3808,6 +3809,94 @@ static bool htp_kparams_solve_fill(
return true;
}
+static bool htp_kparams_solve_timestep_embedding(
+ const struct ggml_hexagon_session * sess,
+ const struct ggml_tensor * op,
+ struct htp_timestep_kernel_params * kparams
+) {
+ memset(kparams, 0, sizeof(*kparams));
+
+ const struct ggml_tensor * src0 = op->src[0];
+ const struct ggml_tensor * dst = op;
+
+ if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+ return false;
+ }
+
+ const int32_t dim = ggml_get_op_params_i32(op, 0);
+ const int32_t max_period = ggml_get_op_params_i32(op, 1);
+
+ if (dim <= 0 || max_period <= 0) {
+ return false;
+ }
+ if (src0->ne[1] != 1 || src0->ne[2] != 1 || src0->ne[3] != 1 ||
+ dst->ne[0] != dim || dst->ne[1] != src0->ne[0] || dst->ne[2] != 1 || dst->ne[3] != 1) {
+ return false;
+ }
+ if (!ggml_is_contiguous_rows(src0) || !ggml_is_contiguous_rows(dst)) {
+ return false;
+ }
+
+ if (src0->ne[0] <= 0 || src0->ne[0] > INT32_MAX || dim < 2) {
+ return false;
+ }
+
+ // below this many output elements threading costs more than it saves
+ constexpr uint32_t min_split_elems = 1024;
+ // target size of one VTCM output buffer per thread
+ constexpr size_t out_buf_budget = 16 * 1024;
+
+ const uint32_t n_rows = (uint32_t) src0->ne[0];
+ const uint32_t half = (uint32_t) dim / 2;
+ const uint32_t n_blocks = (half + HTP_TIMESTEP_BLOCK - 1) / HTP_TIMESTEP_BLOCK;
+
+ // split rows or freq blocks, whichever keeps more threads busy
+ const uint32_t max_hvx = (uint64_t) n_rows * (uint32_t) dim < min_split_elems ? 1 : htp_session_n_hvx(sess);
+ const uint32_t hvx_rows = (std::min)(max_hvx, n_rows);
+ const uint32_t hvx_cols = (std::min)(max_hvx, n_blocks);
+ const bool split_cols = hvx_cols > hvx_rows;
+
+ uint32_t rows_per_hvx = n_rows;
+ uint32_t blocks_per_hvx = n_blocks;
+ uint32_t n_hvx = 1;
+ if (split_cols) {
+ blocks_per_hvx = (n_blocks + hvx_cols - 1) / hvx_cols;
+ n_hvx = (n_blocks + blocks_per_hvx - 1) / blocks_per_hvx;
+ } else {
+ rows_per_hvx = (n_rows + hvx_rows - 1) / hvx_rows;
+ n_hvx = (n_rows + rows_per_hvx - 1) / rows_per_hvx;
+ }
+
+ struct htp_timestep_vtcm_layout layout;
+ htp_timestep_vtcm_layout_build(&layout, (uint32_t) dim, blocks_per_hvx, 1, n_hvx);
+
+ const size_t row_bytes = layout.out_buf_size; // one [cos | sin] row
+ uint32_t chunk_rows = (uint32_t) (std::max)((size_t) 1, (std::min)((size_t) rows_per_hvx, out_buf_budget / row_bytes));
+
+ htp_timestep_vtcm_layout_build(&layout, (uint32_t) dim, blocks_per_hvx, chunk_rows, n_hvx);
+ while (layout.total_bytes > htp_session_vtcm_size(sess) && chunk_rows > 1) {
+ chunk_rows /= 2;
+ htp_timestep_vtcm_layout_build(&layout, (uint32_t) dim, blocks_per_hvx, chunk_rows, n_hvx);
+ }
+ if (layout.total_bytes > htp_session_vtcm_size(sess)) {
+ return false;
+ }
+
+ kparams->n_hvx = n_hvx;
+ kparams->n_rows = n_rows;
+ kparams->dim = (uint32_t) dim;
+ kparams->half = half;
+ kparams->n_blocks = n_blocks;
+ kparams->split_cols = split_cols ? 1 : 0;
+ kparams->rows_per_hvx = rows_per_hvx;
+ kparams->blocks_per_hvx = blocks_per_hvx;
+ kparams->chunk_rows = chunk_rows;
+ kparams->vtcm_size = (uint32_t) layout.total_bytes;
+ kparams->freq_scale = -logf((float) max_period) / (float) half;
+
+ return true;
+}
+
static bool htp_kparams_solve_roll(
const struct ggml_hexagon_session * sess,
const struct ggml_tensor * op,
@@ -5663,6 +5752,9 @@ bool htp_kparams_solve(
case HTP_OP_ROLL:
return htp_kparams_invoke<htp_kparams_solve_roll>(sess, op, kparams);
+ case HTP_OP_TIMESTEP_EMBEDDING:
+ return htp_kparams_invoke<htp_kparams_solve_timestep_embedding>(sess, op, kparams);
+
case HTP_OP_LIGHTNING_INDEXER:
return htp_kparams_invoke<htp_kparams_solve_lightning_indexer>(sess, op, kparams);
diff --git a/ggml/src/ggml-hexagon/htp-kparams.h b/ggml/src/ggml-hexagon/htp-kparams.h
index 0ef485509..bae41ef1b 100644
--- a/ggml/src/ggml-hexagon/htp-kparams.h
+++ b/ggml/src/ggml-hexagon/htp-kparams.h
@@ -28,6 +28,7 @@
#include "htp/set-rows-ops.h"
#include "htp/softmax-ops.h"
#include "htp/solve-tri-ops.h"
+#include "htp/timestep-ops.h"
#include "htp/ssm-conv.h"
#include "htp/ssm-scan-ops.h"
#include "htp/sum-rows-ops.h"
diff --git a/ggml/src/ggml-hexagon/htp-opnode.cpp b/ggml/src/ggml-hexagon/htp-opnode.cpp
index 025ff7ddb..b7be0cea2 100644
--- a/ggml/src/ggml-hexagon/htp-opnode.cpp
+++ b/ggml/src/ggml-hexagon/htp-opnode.cpp
@@ -21,6 +21,7 @@
#include "htp/ssm-conv.h"
#include "htp/ssm-scan-ops.h"
#include "htp/sum-rows-ops.h"
+#include "htp/timestep-ops.h"
#include "htp/top-k-mask-ops.h"
#include "htp/unary-ops.h"
#include "htp/upscale-ops.h"
@@ -276,6 +277,9 @@ static std::string format_kernel_params(const htp_opnode & node) {
case HTP_UPSCALE_KERNEL_BICUBIC: ktype_str = "bicubic"; break;
}
snprintf(str, max_size, "%s nhvx %u chk %u vtcm %u", ktype_str, (unsigned int) kparams->n_hvx, (unsigned int) kparams->dst_nrows_per_chunk, (unsigned int) kparams->vtcm_size);
+ } else if (node.opcode == HTP_OP_TIMESTEP_EMBEDDING) {
+ const auto * kparams = node.kparams<htp_timestep_kernel_params>();
+ snprintf(str, max_size, "%s nhvx %u chk %u vtcm %u", kparams->split_cols ? "cols" : "rows", (unsigned int) kparams->n_hvx, (unsigned int) kparams->chunk_rows, (unsigned int) kparams->vtcm_size);
} else {
snprintf(str, max_size, "----");
}
diff --git a/ggml/src/ggml-hexagon/htp/CMakeLists.txt b/ggml/src/ggml-hexagon/htp/CMakeLists.txt
index 90980a0ee..029f9df99 100644
--- a/ggml/src/ggml-hexagon/htp/CMakeLists.txt
+++ b/ggml/src/ggml-hexagon/htp/CMakeLists.txt
@@ -40,6 +40,7 @@ add_library(${HTP_LIB} SHARED
cumsum-ops.c
top-k-mask-ops.c
fill-ops.c
+ timestep-ops.c
concat-ops.c
diag-ops.c
solve-tri-ops.c
diff --git a/ggml/src/ggml-hexagon/htp/htp-ctx.h b/ggml/src/ggml-hexagon/htp/htp-ctx.h
index ccbf6afee..27b56ff0b 100644
--- a/ggml/src/ggml-hexagon/htp/htp-ctx.h
+++ b/ggml/src/ggml-hexagon/htp/htp-ctx.h
@@ -201,6 +201,7 @@ int op_pool_2d(struct htp_ops_context * octx);
int op_pool_1d(struct htp_ops_context * octx);
int op_allreduce(struct htp_ops_context * octx);
int op_roll(struct htp_ops_context * octx);
+int op_timestep_embedding(struct htp_ops_context * octx);
int op_upscale(struct htp_ops_context * octx);
int op_lightning_indexer(struct htp_ops_context * octx);
diff --git a/ggml/src/ggml-hexagon/htp/htp-ops.h b/ggml/src/ggml-hexagon/htp/htp-ops.h
index 4f24eda1e..2cb49c0a8 100644
--- a/ggml/src/ggml-hexagon/htp/htp-ops.h
+++ b/ggml/src/ggml-hexagon/htp/htp-ops.h
@@ -128,7 +128,7 @@ enum htp_op_code {
HTP_OP_TOP_K_MASK,
HTP_OP_LIGHTNING_INDEXER,
HTP_OP_SSM_SCAN,
-
+ HTP_OP_TIMESTEP_EMBEDDING,
HTP_OP_INVALID
};
diff --git a/ggml/src/ggml-hexagon/htp/main.c b/ggml/src/ggml-hexagon/htp/main.c
index 9d97f0186..712e898d2 100644
--- a/ggml/src/ggml-hexagon/htp/main.c
+++ b/ggml/src/ggml-hexagon/htp/main.c
@@ -1056,6 +1056,9 @@ static int execute_op(struct htp_ops_context * octx) {
case HTP_OP_ROLL:
return op_roll(octx);
+ case HTP_OP_TIMESTEP_EMBEDDING:
+ return op_timestep_embedding(octx);
+
case HTP_OP_POOL_2D:
return op_pool_2d(octx);
diff --git a/ggml/src/ggml-hexagon/htp/timestep-ops.c b/ggml/src/ggml-hexagon/htp/timestep-ops.c
new file mode 100644
index 000000000..28451e38b
--- /dev/null
+++ b/ggml/src/ggml-hexagon/htp/timestep-ops.c
@@ -0,0 +1,217 @@
+#pragma clang diagnostic ignored "-Wunused-variable"
+#pragma clang diagnostic ignored "-Wunused-function"
+#pragma clang diagnostic ignored "-Wunused-but-set-variable"
+
+#include <HAP_farf.h>
+#include <stdint.h>
+
+#include "hex-common.h"
+#include "htp-trace.h"
+
+#include "hvx-utils.h"
+
+#define GGML_COMMON_DECL_C
+#include "ggml-common.h"
+#include "htp-ctx.h"
+#include "htp-ops.h"
+#include "htp-tensor.h"
+#include "htp-vtcm.h"
+#include "dma-queue.h"
+#include "timestep-ops.h"
+
+// ggml op_params layout for TIMESTEP_EMBEDDING:
+// op_params[0] - dim
+// op_params[1] - max_period
+
+// f32 lanes 0..31 built in registers
+static inline HVX_Vector timestep_iota_f32(void) {
+ const HVX_Vector zero = Q6_V_vzero();
+ const HVX_Vector one = Q6_V_vsplat_R(1);
+ HVX_Vector v = one;
+ v = Q6_Vw_vadd_VwVw(v, Q6_V_vlalign_VVR(v, zero, 4));
+ v = Q6_Vw_vadd_VwVw(v, Q6_V_vlalign_VVR(v, zero, 8));
+ v = Q6_Vw_vadd_VwVw(v, Q6_V_vlalign_VVR(v, zero, 16));
+ v = Q6_Vw_vadd_VwVw(v, Q6_V_vlalign_VVR(v, zero, 32));
+ v = Q6_Vw_vadd_VwVw(v, Q6_V_vlalign_VVR(v, zero, 64));
+ return Q6_Vsf_equals_Vw(Q6_Vw_vsub_VwVw(v, one));
+}
+
+struct htp_timestep_context {
+ struct htp_ops_context * octx;
+ const struct htp_timestep_kernel_params * kparams;
+ struct htp_timestep_vtcm_layout layout;
+};
+
+static __attribute__((noinline)) void timestep_compute_freq(float * restrict freq, uint32_t jb0, uint32_t nb, float freq_scale) {
+ const HVX_Vector v_iota = timestep_iota_f32();
+ const HVX_Vector v_scale = hvx_vec_splat_f32(freq_scale);
+ const HVX_Vector v_step = hvx_vec_splat_f32((float) HTP_TIMESTEP_BLOCK);
+
+ HVX_Vector v_j = hvx_vec_add_f32_f32(hvx_vec_splat_f32((float) (jb0 * HTP_TIMESTEP_BLOCK)), v_iota);
+
+ // exp constants fill the register file, pipelining this loop spills them
+ #pragma clang loop pipeline(disable) unroll(disable)
+ for (uint32_t b = 0; b < nb; ++b) {
+ hvx_vmem(freq + b * HTP_TIMESTEP_BLOCK) = hvx_vec_exp_f32(hvx_vec_mul_f32_f32(v_j, v_scale));
+ v_j = hvx_vec_add_f32_f32(v_j, v_step);
+ }
+}
+
+// out row r: [cos (col_stride) | sin (col_stride)], all VTCM
+static __attribute__((noinline)) void timestep_compute_rows(float * restrict out, const float * restrict ts, const float * restrict freq,
+ uint32_t nr, uint32_t nb, uint32_t col_stride) {
+ for (uint32_t r = 0; r < nr; ++r) {
+ const HVX_Vector v_t = hvx_vec_repl4(Q6_V_vror_VR(hvx_vmem(ts + (r & ~31u)), (r & 31) * sizeof(float)));
+ float * out_cos = out + (size_t) r * 2 * col_stride;
+ float * out_sin = out_cos + col_stride;
+
+ // sincos constants fill the register file, pipelining this loop spills them
+ #pragma clang loop pipeline(disable) unroll(disable)
+ for (uint32_t b = 0; b < nb; ++b) {
+ HVX_Vector v_cos;
+ HVX_Vector v_sin;
+ hvx_vec_sincos_f32(hvx_vec_mul_f32_f32(v_t, hvx_vmem(freq + b * HTP_TIMESTEP_BLOCK)), &v_cos, &v_sin);
+ hvx_vmem(out_cos + b * HTP_TIMESTEP_BLOCK) = v_cos;
+ hvx_vmem(out_sin + b * HTP_TIMESTEP_BLOCK) = v_sin;
+ }
+ }
+}
+
+// odd dim: zero sin lanes from n_cols on, this writes the pad element
+static __attribute__((noinline)) void timestep_zero_pad(float * restrict out, uint32_t nr, uint32_t col_stride, uint32_t n_cols) {
+ float * pad = out + col_stride + (n_cols & ~(HTP_TIMESTEP_BLOCK - 1));
+ const HVX_VectorPred q_keep = Q6_Q_vsetq_R((n_cols & (HTP_TIMESTEP_BLOCK - 1)) * sizeof(float));
+ const HVX_Vector v_zero = Q6_V_vzero();
+ for (uint32_t r = 0; r < nr; ++r) {
+ hvx_vmem(pad) = Q6_V_vmux_QVV(q_keep, hvx_vmem(pad), v_zero);
+ pad += 2 * col_stride;
+ }
+}
+
+// Each thread owns a tile: rows [ir0, ir1) x freq blocks [jb0, jb1).
+// Rows are processed in chunks: DMA timesteps in, compute [cos | sin] rows in VTCM, DMA out.
+// Two buffers for both, so compute of chunk c overlaps the writeback of chunk c-1.
+static void timestep_thread(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
+ const struct htp_timestep_context * tctx = (const struct htp_timestep_context *) data;
+ const struct htp_timestep_kernel_params * kparams = tctx->kparams;
+ const struct htp_timestep_vtcm_layout * L = &tctx->layout;
+ struct htp_ops_context * octx = tctx->octx;
+
+ const struct htp_tensor * src0 = octx->src[0];
+ const struct htp_tensor * dst = octx->dst;
+
+ uint32_t ir0 = 0, ir1 = kparams->n_rows;
+ uint32_t jb0 = 0, jb1 = kparams->n_blocks;
+ if (kparams->split_cols) {
+ jb0 = kparams->blocks_per_hvx * ith;
+ jb1 = MIN(jb0 + kparams->blocks_per_hvx, kparams->n_blocks);
+ } else {
+ ir0 = kparams->rows_per_hvx * ith;
+ ir1 = MIN(ir0 + kparams->rows_per_hvx, kparams->n_rows);
+ }
+ if (ir0 >= ir1 || jb0 >= jb1) {
+ return;
+ }
+
+ const uint32_t half = kparams->half;
+ const uint32_t chunk_rows = kparams->chunk_rows;
+ const uint32_t n_chunks = (ir1 - ir0 + chunk_rows - 1) / chunk_rows;
+ const uint32_t j0 = jb0 * HTP_TIMESTEP_BLOCK;
+ const uint32_t n_cols = MIN(jb1 * HTP_TIMESTEP_BLOCK, half) - j0;
+ const uint32_t n_tail = (jb1 == kparams->n_blocks) ? (kparams->dim & 1) : 0;
+ const size_t col_stride = L->col_stride;
+ const size_t dst_nb1 = dst->nb[1];
+
+ uint8_t * vtcm = (uint8_t *) octx->ctx->vtcm_base + ith * L->bytes_per_hvx;
+ uint8_t * ts_buf[2] = { vtcm + L->off_ts, vtcm + L->off_ts + L->ts_buf_size };
+ float * freq = (float *) (vtcm + L->off_freq);
+ uint8_t * out_buf[2] = { vtcm + L->off_out, vtcm + L->off_out + L->out_buf_size };
+
+ const dma_addr_t src_data = src0->data;
+ const dma_addr_t dst_data = dst->data;
+
+ dma_queue * dma_q = th->dma_q;
+ struct htp_thread_trace * tr = th->trace;
+
+ {
+ const uint32_t nr = MIN(chunk_rows, ir1 - ir0);
+ dma_queue_push(dma_q, dma_make_data(ts_buf[0], src_data + (size_t) ir0 * sizeof(float)),
+ nr * sizeof(float), nr * sizeof(float), nr * sizeof(float), 1);
+ }
+
+ // freq table for this tile, overlaps the first timestep fetch
+ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) jb0);
+ timestep_compute_freq(freq, jb0, jb1 - jb0, kparams->freq_scale);
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) jb0);
+
+ for (uint32_t c = 0; c < n_chunks; ++c) {
+ const uint32_t r0 = ir0 + c * chunk_rows;
+ const uint32_t nr = MIN(chunk_rows, ir1 - r0);
+
+ // queue order is: wb(c-2) cos, wb(c-2) sin, ts(c), wb(c-1) cos, wb(c-1) sin
+ if (c >= 2) {
+ dma_queue_pop(dma_q);
+ dma_queue_pop(dma_q);
+ }
+ dma_queue_pop(dma_q);
+
+ if (c + 1 < n_chunks) {
+ const uint32_t r0n = r0 + chunk_rows;
+ const uint32_t nrn = MIN(chunk_rows, ir1 - r0n);
+ dma_queue_push(dma_q, dma_make_data(ts_buf[(c + 1) & 1], src_data + (size_t) r0n * sizeof(float)),
+ nrn * sizeof(float), nrn * sizeof(float), nrn * sizeof(float), 1);
+ }
+
+ float * out = (float *) out_buf[c & 1];
+
+ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) c);
+ timestep_compute_rows(out, (const float *) ts_buf[c & 1], freq, nr, jb1 - jb0, col_stride);
+ if (n_tail) {
+ timestep_zero_pad(out, nr, col_stride, n_cols);
+ }
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) c);
+
+ const dma_addr_t dst_row = dst_data + (size_t) r0 * dst_nb1;
+ const size_t src_nb = 2 * col_stride * sizeof(float);
+ dma_queue_push(dma_q, dma_make_data(dst_row + (size_t) j0 * sizeof(float), out),
+ dst_nb1, src_nb, n_cols * sizeof(float), nr);
+ dma_queue_push(dma_q, dma_make_data(dst_row + (size_t) (half + j0) * sizeof(float), out + col_stride),
+ dst_nb1, src_nb, (n_cols + n_tail) * sizeof(float), nr);
+ }
+
+ dma_queue_flush(dma_q);
+}
+
+int op_timestep_embedding(struct htp_ops_context * octx) {
+ const struct htp_timestep_kernel_params * kparams =
+ (const struct htp_timestep_kernel_params *) octx->kernel_params;
+
+ const int32_t max_period = (int32_t) octx->op_params[1];
+
+ // Only device 0 writes the output when running multi-device
+ if (octx->ctx->mdev.count > 1 && octx->ctx->mdev.idx != 0) {
+ return HTP_STATUS_OK;
+ }
+
+ if (!htp_ops_context_set_n_hvx(octx, kparams->n_hvx)) {
+ return HTP_STATUS_INVAL_PARAMS;
+ }
+
+ struct htp_timestep_context tctx = {
+ .octx = octx,
+ .kparams = kparams,
+ };
+
+ htp_timestep_vtcm_layout_build(&tctx.layout, kparams->dim, kparams->blocks_per_hvx, kparams->chunk_rows, kparams->n_hvx);
+
+ if (octx->ctx->vtcm_size < tctx.layout.total_bytes) {
+ return HTP_STATUS_VTCM_TOO_SMALL;
+ }
+
+ FARF(HIGH, "timestep-embedding: %u timesteps dim %u max-period %d split-%s n-hvx %u chunk-rows %u\n",
+ kparams->n_rows, kparams->dim, max_period, kparams->split_cols ? "cols" : "rows", kparams->n_hvx, kparams->chunk_rows);
+
+ work_queue_run(octx->ctx->hvx_queue, timestep_thread, &tctx, kparams->n_hvx);
+
+ return HTP_STATUS_OK;
+}
diff --git a/ggml/src/ggml-hexagon/htp/timestep-ops.h b/ggml/src/ggml-hexagon/htp/timestep-ops.h
new file mode 100644
index 000000000..9af96df6e
--- /dev/null
+++ b/ggml/src/ggml-hexagon/htp/timestep-ops.h
@@ -0,0 +1,67 @@
+#ifndef HTP_TIMESTEP_OPS_H
+#define HTP_TIMESTEP_OPS_H
+
+#include "hex-common.h"
+#include "htp-vtcm.h"
+#include <stddef.h>
+#include <stdint.h>
+
+#define HTP_TIMESTEP_BLOCK 32 // f32 lanes per HVX vector
+
+// Per-hvx-thread VTCM layout (all buffers are 128B aligned):
+// ts - 2 x chunk_rows timesteps (ping-pong)
+// freq - blocks_per_hvx frequency vectors
+// out - 2 x chunk_rows x [cos | sin] rows (ping-pong)
+struct htp_timestep_vtcm_layout {
+ size_t total_bytes;
+ size_t bytes_per_hvx;
+ size_t off_ts;
+ size_t off_freq;
+ size_t off_out;
+ size_t ts_buf_size;
+ size_t out_buf_size;
+ size_t col_stride; // floats per cos (or sin) row in VTCM
+};
+
+static inline void htp_timestep_vtcm_layout_build(
+ struct htp_timestep_vtcm_layout * L,
+ uint32_t dim,
+ uint32_t blocks_per_hvx,
+ uint32_t chunk_rows,
+ uint32_t n_hvx
+) {
+ // one extra slot holds the zero pad for odd dim
+ L->col_stride = hex_round_up((size_t) blocks_per_hvx * HTP_TIMESTEP_BLOCK + (dim & 1), HTP_TIMESTEP_BLOCK);
+ L->ts_buf_size = hex_round_up((size_t) chunk_rows * sizeof(float), 128);
+ L->out_buf_size = (size_t) chunk_rows * 2 * L->col_stride * sizeof(float);
+
+ size_t off = 0;
+ VTCM_LAYOUT_ALLOC(off, off_ts, 2 * L->ts_buf_size);
+ VTCM_LAYOUT_ALLOC(off, off_freq, (size_t) blocks_per_hvx * HTP_TIMESTEP_BLOCK * sizeof(float));
+ VTCM_LAYOUT_ALLOC(off, off_out, 2 * L->out_buf_size);
+
+ L->bytes_per_hvx = off;
+ L->total_bytes = off * n_hvx;
+}
+
+struct htp_timestep_kernel_params {
+ uint32_t n_hvx;
+ uint32_t n_rows;
+ uint32_t dim;
+ uint32_t half;
+ uint32_t n_blocks;
+ uint32_t split_cols; // 0: split rows across threads, 1: split freq blocks
+ uint32_t rows_per_hvx;
+ uint32_t blocks_per_hvx;
+ uint32_t chunk_rows;
+ uint32_t vtcm_size;
+ float freq_scale; // -ln(max_period) / half
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_timestep_kernel_params) <= 128, "htp_timestep_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_timestep_kernel_params) <= 128, "htp_timestep_kernel_params is too large for kernel_params blob");
+#endif
+
+#endif // HTP_TIMESTEP_OPS_H
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 3ff5ced0b..1ad7cc2ad 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -12859,6 +12859,9 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
}
test_cases.emplace_back(new test_arange(GGML_TYPE_F32, 0.0f, 1048576.0f, 1.0f));
test_cases.emplace_back(new test_timestep_embedding());
+ test_cases.emplace_back(new test_timestep_embedding(GGML_TYPE_F32, { 1, 1, 1, 1 }, 256, 10000));
+ test_cases.emplace_back(new test_timestep_embedding(GGML_TYPE_F32, { 21, 1, 1, 1 }, 256, 10000));
+ test_cases.emplace_back(new test_timestep_embedding(GGML_TYPE_F32, { 31, 1, 1, 1 }, 256, 10000));
test_cases.emplace_back(new test_leaky_relu());
test_cases.emplace_back(new test_cumsum(GGML_TYPE_F32, { 10, 5, 4, 3 }));