andodeki/makepad
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
makepad/libs/ai/cuda/kernels/train.cu

823 lines
42 KiB
Text
Raw Permalink Blame History

// Training kernels: forward and backward of the small op set the singing
// model (libs/ai/models/sing, nn.rs) is built from, plus the optimiser,
// the STFT family and monotonic alignment. f32 row-major [rows, cols]
// tensors; "seg" is the number of rows per batch item (convolutions, RoPE
// positions and attention never cross an item). Every *_bwd kernel
// ACCUMULATES into its gradient buffer (+=). Launch wrappers are
// `extern "C" cudaError_t mkt_*(..., cudaStream_t)`.
#include <cuda_runtime.h>
#include <stdint.h>
#include <math.h>
#define MKT_THREADS 256
static inline unsigned mkt_blocks(size_t n) {
size_t b = (n + MKT_THREADS - 1) / MKT_THREADS;
if (b > 65535u * 8u) b = 65535u * 8u;
return (unsigned)(b ? b : 1);
}
#define GRID_LOOP(i, n) for (size_t i = blockIdx.x * (size_t)blockDim.x + threadIdx.x; i < (n); i += (size_t)blockDim.x * gridDim.x)
// ---------------------------------------------------------------------------
// Activations: 0 gelu(tanh), 1 silu, 2 relu, 3 leaky relu 0.1, 4 tanh,
// 5 sigmoid, 6 exp(min(x, 10))
// ---------------------------------------------------------------------------
__device__ __forceinline__ float act_f(int op, float x) {
switch (op) {
case 0: return 0.5f * x * (1.0f + tanhf(0.7978846f * (x + 0.044715f * x * x * x)));
case 1: return x / (1.0f + __expf(-x));
case 2: return fmaxf(x, 0.0f);
case 3: return x > 0.0f ? x : 0.1f * x;
case 4: return tanhf(x);
case 5: return 1.0f / (1.0f + __expf(-x));
default: return __expf(fminf(x, 10.0f));
}
}
__device__ __forceinline__ float act_g(int op, float x) {
switch (op) {
case 0: {
float u = 0.7978846f * (x + 0.044715f * x * x * x);
float t = tanhf(u);
return 0.5f * (1.0f + t) + 0.5f * x * (1.0f - t * t) * 0.7978846f * (1.0f + 3.0f * 0.044715f * x * x);
}
case 1: { float s = 1.0f / (1.0f + __expf(-x)); return s * (1.0f + x * (1.0f - s)); }
case 2: return x > 0.0f ? 1.0f : 0.0f;
case 3: return x > 0.0f ? 1.0f : 0.1f;
case 4: { float t = tanhf(x); return 1.0f - t * t; }
case 5: { float s = 1.0f / (1.0f + __expf(-x)); return s * (1.0f - s); }
default: return x < 10.0f ? __expf(x) : 0.0f;
}
}
__global__ void k_act_fwd(const float* x, float* y, size_t n, int op) {
GRID_LOOP(i, n) y[i] = act_f(op, x[i]);
}
__global__ void k_act_bwd(const float* x, const float* d, float* dx, size_t n, int op) {
GRID_LOOP(i, n) dx[i] += d[i] * act_g(op, x[i]);
}
// ---------------------------------------------------------------------------
// Elementwise
// ---------------------------------------------------------------------------
__global__ void k_add(const float* a, const float* b, float* y, size_t n) { GRID_LOOP(i, n) y[i] = a[i] + b[i]; }
__global__ void k_mul(const float* a, const float* b, float* y, size_t n) { GRID_LOOP(i, n) y[i] = a[i] * b[i]; }
__global__ void k_scale(const float* x, float* y, float s, size_t n) { GRID_LOOP(i, n) y[i] = x[i] * s; }
__global__ void k_axpy(const float* x, float* y, float s, size_t n) { GRID_LOOP(i, n) y[i] += s * x[i]; }
// y += x * other (the product rule's half)
__global__ void k_mul_acc(const float* x, const float* other, float* y, size_t n) { GRID_LOOP(i, n) y[i] += x[i] * other[i]; }
__global__ void k_fill(float* y, float v, size_t n) { GRID_LOOP(i, n) y[i] = v; }
// y += s * x[0] (scalar broadcast; used by scalar loss sums)
__global__ void k_axpy_scalar(const float* x, float* y, float s) { if (threadIdx.x == 0 && blockIdx.x == 0) y[0] += s * x[0]; }
__global__ void k_log_eps(const float* x, float* y, float eps, size_t n) { GRID_LOOP(i, n) y[i] = logf(x[i] + eps); }
__global__ void k_log_eps_bwd(const float* x, const float* d, float* dx, float eps, size_t n) { GRID_LOOP(i, n) dx[i] += d[i] / (x[i] + eps); }
// Row-broadcast: y[r,c] = x[r,c] (+|*) v[c]
__global__ void k_add_row(const float* x, const float* v, float* y, size_t rows, size_t cols) {
GRID_LOOP(i, rows * cols) y[i] = x[i] + v[i % cols];
}
__global__ void k_mul_row(const float* x, const float* v, float* y, size_t rows, size_t cols) {
GRID_LOOP(i, rows * cols) y[i] = x[i] * v[i % cols];
}
__global__ void k_mul_row_bwd_x(const float* d, const float* v, float* dx, size_t rows, size_t cols) {
GRID_LOOP(i, rows * cols) dx[i] += d[i] * v[i % cols];
}
__global__ void k_mul_col(const float* x, const float* s, float* y, size_t rows, size_t cols) {
GRID_LOOP(i, rows * cols) y[i] = x[i] * s[i / cols];
}
__global__ void k_mul_col_bwd_x(const float* d, const float* s, float* dx, size_t rows, size_t cols) {
GRID_LOOP(i, rows * cols) dx[i] += d[i] * s[i / cols];
}
// dv[c] += sum_r d[r,c] * (x ? x[r,c] : 1). grid.x covers columns in 32s,
// block (32, 8): each thread strides rows, a shared reduction per column.
__global__ void k_col_dot_acc(const float* d, const float* x, float* dv, size_t rows, size_t cols) {
__shared__ float sh[8][33];
size_t c = blockIdx.x * 32 + threadIdx.x;
float acc = 0.0f;
if (c < cols) {
for (size_t r = blockIdx.y * 8 + threadIdx.y; r < rows; r += 8 * gridDim.y) {
float v = d[r * cols + c];
acc += x ? v * x[r * cols + c] : v;
}
}
sh[threadIdx.y][threadIdx.x] = acc;
__syncthreads();
if (threadIdx.y == 0 && c < cols) {
float s = 0.0f;
for (int k = 0; k < 8; k++) s += sh[k][threadIdx.x];
atomicAdd(&dv[c], s);
}
}
// ds[r] += sum_c d[r,c] * x[r,c]; one warp per row.
__global__ void k_row_dot_acc(const float* d, const float* x, float* ds, size_t rows, size_t cols) {
size_t r = blockIdx.x * (blockDim.x / 32) + threadIdx.x / 32;
int lane = threadIdx.x & 31;
if (r >= rows) return;
float acc = 0.0f;
for (size_t c = lane; c < cols; c += 32) acc += d[r * cols + c] * x[r * cols + c];
for (int o = 16; o > 0; o >>= 1) acc += __shfl_down_sync(0xffffffff, acc, o);
if (lane == 0) ds[r] += acc;
}
// Copy a block of columns: dst[r, dst_off + c] (=|+=) src[r, src_off + c].
__global__ void k_copy_cols(const float* src, size_t src_cols, size_t src_off, float* dst, size_t dst_cols, size_t dst_off, size_t rows, size_t width, int acc) {
GRID_LOOP(i, rows * width) {
size_t r = i / width, c = i % width;
float v = src[r * src_cols + src_off + c];
float* o = &dst[r * dst_cols + dst_off + c];
if (acc) *o += v; else *o = v;
}
}
__global__ void k_gather_rows(const float* x, const uint32_t* idx, float* y, size_t rows_out, size_t cols) {
GRID_LOOP(i, rows_out * cols) {
size_t r = i / cols, c = i % cols;
y[i] = x[(size_t)idx[r] * cols + c];
}
}
__global__ void k_scatter_rows_acc(const float* d, const uint32_t* idx, float* dx, size_t rows_out, size_t cols) {
GRID_LOOP(i, rows_out * cols) {
size_t r = i / cols, c = i % cols;
atomicAdd(&dx[(size_t)idx[r] * cols + c], d[i]);
}
}
// ---------------------------------------------------------------------------
// Convolutions over time within segments of `seg` rows (same padding)
// ---------------------------------------------------------------------------
// col[r, j*C + c] = x[r + j*dil - pad, c] inside r's segment, else 0.
__global__ void k_im2col(const float* x, float* col, size_t rows, size_t C, int k, int dil, size_t seg) {
int pad = (k - 1) * dil / 2;
GRID_LOOP(i, rows * (size_t)k * C) {
size_t r = i / ((size_t)k * C);
size_t rem = i % ((size_t)k * C);
int j = (int)(rem / C);
size_t c = rem % C;
size_t base = (r / seg) * seg;
long src = (long)(r - base) + (long)j * dil - pad;
col[i] = (src >= 0 && (size_t)src < seg && base + src < rows) ? x[(base + src) * C + c] : 0.0f;
}
}
// dx[s, c] += sum_j dcol[s - j*dil + pad, j*C + c] (inside the segment)
__global__ void k_col2im_acc(const float* dcol, float* dx, size_t rows, size_t C, int k, int dil, size_t seg) {
int pad = (k - 1) * dil / 2;
GRID_LOOP(i, rows * C) {
size_t s = i / C, c = i % C;
size_t base = (s / seg) * seg;
float acc = 0.0f;
for (int j = 0; j < k; j++) {
long r = (long)(s - base) - (long)j * dil + pad;
if (r >= 0 && (size_t)r < seg && base + r < rows) acc += dcol[(base + r) * (size_t)k * C + (size_t)j * C + c];
}
dx[i] += acc;
}
}
// Depthwise: y[r,c] = sum_j x[r + j - pad, c] * w[j, c]
__global__ void k_dwconv(const float* x, const float* w, float* y, size_t rows, size_t C, int k, size_t seg) {
int pad = (k - 1) / 2;
GRID_LOOP(i, rows * C) {
size_t r = i / C, c = i % C;
size_t base = (r / seg) * seg;
float acc = 0.0f;
for (int j = 0; j < k; j++) {
long s = (long)(r - base) + j - pad;
if (s >= 0 && (size_t)s < seg && base + s < rows) acc += x[(base + s) * C + c] * w[(size_t)j * C + c];
}
y[i] = acc;
}
}
__global__ void k_dwconv_bwd_x(const float* d, const float* w, float* dx, size_t rows, size_t C, int k, size_t seg) {
int pad = (k - 1) / 2;
GRID_LOOP(i, rows * C) {
size_t s = i / C, c = i % C;
size_t base = (s / seg) * seg;
float acc = 0.0f;
for (int j = 0; j < k; j++) {
long r = (long)(s - base) - j + pad;
if (r >= 0 && (size_t)r < seg && base + r < rows) acc += d[(base + r) * C + c] * w[(size_t)j * C + c];
}
dx[i] += acc;
}
}
// dw[j, c] += sum_r d[r, c] * x[r + j - pad, c]; grid (ceil(C/32), k, ychunks), block (32, 8)
__global__ void k_dwconv_bwd_w(const float* d, const float* x, float* dw, size_t rows, size_t C, int k, size_t seg) {
__shared__ float sh[8][33];
int pad = (k - 1) / 2;
size_t c = blockIdx.x * 32 + threadIdx.x;
int j = blockIdx.y;
float acc = 0.0f;
if (c < C) {
for (size_t r = blockIdx.z * 8 + threadIdx.y; r < rows; r += 8 * gridDim.z) {
size_t base = (r / seg) * seg;
long s = (long)(r - base) + j - pad;
if (s >= 0 && (size_t)s < seg && base + s < rows) acc += d[r * C + c] * x[(base + s) * C + c];
}
}
sh[threadIdx.y][threadIdx.x] = acc;
__syncthreads();
if (threadIdx.y == 0 && c < C) {
float s = 0.0f;
for (int q = 0; q < 8; q++) s += sh[q][threadIdx.x];
atomicAdd(&dw[(size_t)j * C + c], s);
}
}
// ---------------------------------------------------------------------------
// LayerNorm (no affine; the gain and bias are mul_row / add_row)
// ---------------------------------------------------------------------------
__device__ float block_sum(float v, float* sh) {
for (int o = 16; o > 0; o >>= 1) v += __shfl_down_sync(0xffffffff, v, o);
int lane = threadIdx.x & 31, w = threadIdx.x >> 5;
if (lane == 0) sh[w] = v;
__syncthreads();
int nw = blockDim.x >> 5;
v = threadIdx.x < nw ? sh[threadIdx.x] : 0.0f;
if (w == 0) for (int o = 16; o > 0; o >>= 1) v += __shfl_down_sync(0xffffffff, v, o);
if (threadIdx.x == 0) sh[0] = v;
__syncthreads();
float r = sh[0];
__syncthreads();
return r;
}
// One block per row.
__global__ void k_layer_norm(const float* x, float* xhat, float* rstd, size_t cols) {
__shared__ float sh[32];
size_t r = blockIdx.x;
const float* xr = x + r * cols;
float s = 0.0f;
for (size_t c = threadIdx.x; c < cols; c += blockDim.x) s += xr[c];
float mean = block_sum(s, sh) / cols;
float v = 0.0f;
for (size_t c = threadIdx.x; c < cols; c += blockDim.x) { float e = xr[c] - mean; v += e * e; }
float var = block_sum(v, sh) / cols;
float rs = rsqrtf(var + 1e-5f);
for (size_t c = threadIdx.x; c < cols; c += blockDim.x) xhat[r * cols + c] = (xr[c] - mean) * rs;
if (threadIdx.x == 0) rstd[r] = rs;
}
__global__ void k_layer_norm_bwd(const float* d, const float* xhat, const float* rstd, float* dx, size_t cols) {
__shared__ float sh[32];
size_t r = blockIdx.x;
const float* dr = d + r * cols;
const float* xr = xhat + r * cols;
float a = 0.0f, b = 0.0f;
for (size_t c = threadIdx.x; c < cols; c += blockDim.x) { a += dr[c]; b += dr[c] * xr[c]; }
float m1 = block_sum(a, sh) / cols;
float m2 = block_sum(b, sh) / cols;
float rs = rstd[r];
for (size_t c = threadIdx.x; c < cols; c += blockDim.x) dx[r * cols + c] += rs * (dr[c] - m1 - xr[c] * m2);
}
// ---------------------------------------------------------------------------
// RoPE: per head, pairs (i, i + dh/2); position = row index in its segment.
// sign = +1 forward, -1 the inverse rotation (the backward).
// ---------------------------------------------------------------------------
__global__ void k_rope(const float* x, float* y, size_t rows, size_t C, int heads, size_t seg, float sign, int acc) {
int dh = (int)(C / heads), half = dh / 2;
GRID_LOOP(i, rows * (size_t)heads * half) {
size_t r = i / ((size_t)heads * half);
size_t rem = i % ((size_t)heads * half);
int h = (int)(rem / half), q = (int)(rem % half);
float pos = (float)(r % seg);
float f = powf(10000.0f, -(float)q / half);
float sn, cs;
sincosf(pos * f, &sn, &cs);
size_t a = r * C + (size_t)h * dh + q, b = a + half;
float xa = x[a], xb = x[b];
float ya = xa * cs - sign * xb * sn, yb = sign * xa * sn + xb * cs;
if (acc) { y[a] += ya; y[b] += yb; } else { y[a] = ya; y[b] = yb; }
}
}
// ---------------------------------------------------------------------------
// Attention helpers: [B*T, H*dh] <-> [B, H, T, dh], masked softmax
// ---------------------------------------------------------------------------
__global__ void k_split_heads(const float* x, float* y, size_t B, size_t T, int H, int dh) {
GRID_LOOP(i, B * T * (size_t)H * dh) {
size_t e = i % dh, t = (i / dh) % T, h = (i / ((size_t)dh * T)) % H, b = i / ((size_t)dh * T * H);
y[i] = x[(b * T + t) * (size_t)H * dh + h * dh + e];
}
}
__global__ void k_merge_heads(const float* x, float* y, size_t B, size_t T, int H, int dh, int acc) {
GRID_LOOP(i, B * T * (size_t)H * dh) {
size_t e = i % dh, t = (i / dh) % T, h = (i / ((size_t)dh * T)) % H, b = i / ((size_t)dh * T * H);
float v = x[i];
float* o = &y[(b * T + t) * (size_t)H * dh + h * dh + e];
if (acc) *o += v; else *o = v;
}
}
// In place: rows of s [B*H*T, S] scaled, keys >= key_len[b] masked, softmax.
// One block per row.
__global__ void k_softmax_rows(float* s, size_t T, size_t S, int H, const uint32_t* key_len, float scale) {
__shared__ float sh[32];
size_t row = blockIdx.x;
size_t b = row / ((size_t)H * T);
size_t valid = key_len ? key_len[b] : S;
float* p = s + row * S;
float mx = -1e30f;
for (size_t j = threadIdx.x; j < S; j += blockDim.x) if (j < valid) mx = fmaxf(mx, p[j] * scale);
// block max
for (int o = 16; o > 0; o >>= 1) mx = fmaxf(mx, __shfl_down_sync(0xffffffff, mx, o));
int lane = threadIdx.x & 31, w = threadIdx.x >> 5;
if (lane == 0) sh[w] = mx;
__syncthreads();
if (threadIdx.x == 0) { float m = -1e30f; for (int k = 0; k < (int)(blockDim.x >> 5); k++) m = fmaxf(m, sh[k]); sh[0] = m; }
__syncthreads();
mx = sh[0];
__syncthreads();
float sum = 0.0f;
for (size_t j = threadIdx.x; j < S; j += blockDim.x) {
float e = j < valid ? __expf(p[j] * scale - mx) : 0.0f;
p[j] = e;
sum += e;
}
sum = block_sum(sum, sh);
float inv = 1.0f / fmaxf(sum, 1e-30f);
for (size_t j = threadIdx.x; j < S; j += blockDim.x) p[j] *= inv;
}
// dp -> ds in place: ds = p * (dp - sum(dp * p)) * scale
__global__ void k_softmax_bwd_rows(const float* p, float* dp, size_t S, float scale) {
__shared__ float sh[32];
size_t row = blockIdx.x;
const float* pr = p + row * S;
float* dr = dp + row * S;
float dot = 0.0f;
for (size_t j = threadIdx.x; j < S; j += blockDim.x) dot += pr[j] * dr[j];
dot = block_sum(dot, sh);
for (size_t j = threadIdx.x; j < S; j += blockDim.x) dr[j] = pr[j] * (dr[j] - dot) * scale;
}
// ---------------------------------------------------------------------------
// Losses. Row weights w[r] (the mask), columns >= col_lim[item] ignored
// (col_lim per item = rows/seg items; null = all). Writes the per-element
// gradient coefficient into g (later scaled by the upstream scalar) and
// accumulates the loss sum into out[0].
// ---------------------------------------------------------------------------
__global__ void k_loss(const float* x, const float* t, const float* w, const uint32_t* col_lim, size_t seg, float* g, float* out, size_t rows, size_t cols, float inv_denom, int sq) {
__shared__ float sh[32];
float acc = 0.0f;
for (size_t i = blockIdx.x * (size_t)blockDim.x + threadIdx.x; i < rows * cols; i += (size_t)blockDim.x * gridDim.x) {
size_t r = i / cols, c = i % cols;
float wr = w ? w[r] : 1.0f;
if (col_lim && c >= col_lim[r / seg]) wr = 0.0f;
float e = x[i] - t[i];
if (sq) { acc += wr * e * e; g[i] = 2.0f * wr * e * inv_denom; }
else { acc += wr * fabsf(e); g[i] = wr * (e > 0.0f ? 1.0f : (e < 0.0f ? -1.0f : 0.0f)) * inv_denom; }
}
acc = block_sum(acc, sh);
if (threadIdx.x == 0) atomicAdd(out, acc * inv_denom);
}
// dx += g * d[0]
__global__ void k_scaled_acc(const float* g, const float* d, float* dx, size_t n) {
float s = d[0];
GRID_LOOP(i, n) dx[i] += g[i] * s;
}
__global__ void k_bce(const float* x, const float* t, float* g, float* out, size_t n) {
__shared__ float sh[32];
float acc = 0.0f;
float inv = 1.0f / n;
for (size_t i = blockIdx.x * (size_t)blockDim.x + threadIdx.x; i < n; i += (size_t)blockDim.x * gridDim.x) {
float z = x[i], y = t[i];
acc += fmaxf(z, 0.0f) - z * y + log1pf(__expf(-fabsf(z)));
g[i] = (1.0f / (1.0f + __expf(-z)) - y) * inv;
}
acc = block_sum(acc, sh);
if (threadIdx.x == 0) atomicAdd(out, acc * inv);
}
// ---------------------------------------------------------------------------
// The vocoder's source filter (see nn.rs hn_filter)
// ---------------------------------------------------------------------------
__global__ void k_hn_filter(const float* gh, const float* gn, const float* pa, const float* pb,
const float* hr, const float* hi, const float* nr, const float* ni,
float* y, size_t rows, size_t F) {
GRID_LOOP(i, rows * F) {
size_t r = i / F, q = i % F;
float ah = __expf(fminf(gh[i], 10.0f)), an = __expf(fminf(gn[i], 10.0f));
float a = pa[i], b = pb[i];
float m = sqrtf(a * a + b * b + 1e-6f);
float cs = a / m, sn = b / m;
float xr = hr[i] * cs - hi[i] * sn, xi = hr[i] * sn + hi[i] * cs;
y[r * 2 * F + q] = ah * xr + an * nr[i];
y[r * 2 * F + F + q] = ah * xi + an * ni[i];
}
}
__global__ void k_hn_filter_bwd(const float* d, const float* gh, const float* gn, const float* pa, const float* pb,
const float* hr, const float* hi, const float* nr, const float* ni,
float* dgh, float* dgn, float* dpa, float* dpb, size_t rows, size_t F) {
GRID_LOOP(i, rows * F) {
size_t r = i / F, q = i % F;
float dyr = d[r * 2 * F + q], dyi = d[r * 2 * F + F + q];
float ah = __expf(fminf(gh[i], 10.0f)), an = __expf(fminf(gn[i], 10.0f));
float a = pa[i], b = pb[i];
float m2 = a * a + b * b + 1e-6f, m = sqrtf(m2);
float cs = a / m, sn = b / m;
float xr = hr[i] * cs - hi[i] * sn, xi = hr[i] * sn + hi[i] * cs;
if (dgh) dgh[i] += (dyr * xr + dyi * xi) * ah * (gh[i] < 10.0f ? 1.0f : 0.0f);
if (dgn) dgn[i] += (dyr * nr[i] + dyi * ni[i]) * an * (gn[i] < 10.0f ? 1.0f : 0.0f);
float dcs = ah * (dyr * hr[i] + dyi * hi[i]);
float dsn = ah * (-dyr * hi[i] + dyi * hr[i]);
float m3 = m2 * m;
if (dpa) dpa[i] += dcs * (b * b + 1e-6f) / m3 - dsn * a * b / m3;
if (dpb) dpb[i] += dsn * (a * a + 1e-6f) / m3 - dcs * a * b / m3;
}
}
// ---------------------------------------------------------------------------
// FFT: one block per transform, in shared memory, n a power of two <= 4096.
// Forward e^{-i}, inverse e^{+i} unscaled.
// ---------------------------------------------------------------------------
__global__ void k_fft(float* re, float* im, int n, int log2n, int inverse) {
extern __shared__ float s[];
float* sr = s;
float* si = s + n;
size_t base = (size_t)blockIdx.x * n;
for (int i = threadIdx.x; i < n; i += blockDim.x) {
int j = __brev(i) >> (32 - log2n);
sr[j] = re[base + i];
si[j] = im[base + i];
}
__syncthreads();
float sign = inverse ? 1.0f : -1.0f;
for (int len = 2; len <= n; len <<= 1) {
int half = len >> 1;
for (int t = threadIdx.x; t < n / 2; t += blockDim.x) {
int grp = t / half, k = t % half;
int a = grp * len + k, b = a + half;
float ang = sign * 2.0f * 3.14159265358979f * k / len;
float wr, wi;
sincosf(ang, &wi, &wr);
float tr = sr[b] * wr - si[b] * wi;
float ti = sr[b] * wi + si[b] * wr;
sr[b] = sr[a] - tr; si[b] = si[a] - ti;
sr[a] += tr; si[a] += ti;
}
__syncthreads();
}
for (int i = threadIdx.x; i < n; i += blockDim.x) {
re[base + i] = sr[i];
im[base + i] = si[i];
}
}
// Frames of B waveforms (each `len` samples) into [B*frames, n] with the
// window (length n, zero-padded) and an optional gain per sample of an item
// (`len` values, shared by all items), im = 0.
__global__ void k_frame(const float* x, const float* gain, const float* win, float* re, float* im, size_t B, size_t len, size_t frames, int n, int hop) {
GRID_LOOP(i, B * frames * (size_t)n) {
size_t k = i % n, f = (i / n) % frames, b = i / ((size_t)n * frames);
long src = (long)(f * hop) - n / 2 + (long)k;
float v = 0.0f;
if (src >= 0 && (size_t)src < len) {
size_t si = b * len + src;
v = x[si] * win[k] * (gain ? gain[src] : 1.0f);
}
re[i] = v;
im[i] = 0.0f;
}
}
// Full-spectrum [rows, n] -> half [rows, bins] (re | separate im)
__global__ void k_half(const float* fr, const float* fi, float* hr, float* hi, size_t rows, int n, float scale) {
int bins = n / 2 + 1;
GRID_LOOP(i, rows * (size_t)bins) {
size_t r = i / bins, q = i % bins;
hr[i] = fr[r * n + q] * scale;
hi[i] = fi[r * n + q] * scale;
}
}
// Full -> half with per-bin weights (c_edge at DC/Nyquist, whose imaginary
// part is zeroed; c_mid elsewhere): the iSTFT's adjoint.
__global__ void k_half_c(const float* fr, const float* fi, float* hr, float* hi, size_t rows, int n, float c_edge, float c_mid) {
int bins = n / 2 + 1;
GRID_LOOP(i, rows * (size_t)bins) {
size_t r = i / bins; int q = (int)(i % bins);
bool edge = q == 0 || q == n / 2;
float c = edge ? c_edge : c_mid;
hr[i] = fr[r * n + q] * c;
hi[i] = edge ? 0.0f : fi[r * n + q] * c;
}
}
// Half spectrum (with per-bin weights c_edge at DC/Nyquist and c_mid elsewhere,
// imaginary part dropped at the edges) -> Hermitian full spectrum.
__global__ void k_hermitian(const float* hr, const float* hi, float* fr, float* fi, size_t rows, int n, float c_edge, float c_mid) {
int bins = n / 2 + 1;
GRID_LOOP(i, rows * (size_t)n) {
size_t r = i / n;
int k = (int)(i % n);
int q = k < bins ? k : n - k;
float c = (q == 0 || q == n / 2) ? c_edge : c_mid;
float vr = hr[r * bins + q] * c;
float vi = (q == 0 || q == n / 2) ? 0.0f : hi[r * bins + q] * c;
fr[i] = vr;
fi[i] = k < bins ? vi : -vi;
}
}
// Overlap-add of windowed frames [B*frames, n] (real part) into B waveforms:
// y[b, s] (=|+=) norm[s] * sum_f win[k] * fr[(b, f), k] * scale
__global__ void k_overlap_add(const float* fr, const float* win, const float* norm, float* y, size_t B, size_t len, size_t frames, int n, int hop, float scale, int acc) {
GRID_LOOP(i, B * len) {
size_t b = i / len, s = i % len;
long c = (long)s + n / 2;
long f_hi = c / hop;
long f_lo = (c - n + 1 + hop - 1) / hop;
if (f_lo < 0) f_lo = 0;
if (f_hi > (long)frames - 1) f_hi = (long)frames - 1;
float v = 0.0f;
for (long f = f_lo; f <= f_hi; f++) {
long k = c - f * hop;
if (k >= 0 && k < n) v += win[k] * fr[(b * frames + f) * n + k];
}
v *= scale * (norm ? norm[s] : 1.0f);
if (acc) y[i] += v; else y[i] = v;
}
}
__global__ void k_mag(const float* re, const float* im, float* m, size_t n) {
GRID_LOOP(i, n) m[i] = sqrtf(re[i] * re[i] + im[i] * im[i] + 1e-9f);
}
__global__ void k_mag_bwd(const float* d, const float* re, const float* im, const float* m, float* dre, float* dim, size_t n) {
GRID_LOOP(i, n) { float s = d[i] / m[i]; dre[i] = s * re[i]; dim[i] = s * im[i]; }
}
// Deinterleave [rows, 2F] (re | im) <-> separate [rows, F] planes.
__global__ void k_split_ri(const float* y, float* re, float* im, size_t rows, size_t F) {
GRID_LOOP(i, rows * F) { size_t r = i / F, q = i % F; re[i] = y[r * 2 * F + q]; im[i] = y[r * 2 * F + F + q]; }
}
__global__ void k_join_ri_acc(const float* re, const float* im, float* y, size_t rows, size_t F) {
GRID_LOOP(i, rows * F) { size_t r = i / F, q = i % F; y[r * 2 * F + q] += re[i]; y[r * 2 * F + F + q] += im[i]; }
}
// ---------------------------------------------------------------------------
// Optimiser
// ---------------------------------------------------------------------------
__global__ void k_sumsq(const float* g, size_t n, float* out) {
__shared__ float sh[32];
float acc = 0.0f;
for (size_t i = blockIdx.x * (size_t)blockDim.x + threadIdx.x; i < n; i += (size_t)blockDim.x * gridDim.x) acc += g[i] * g[i];
acc = block_sum(acc, sh);
if (threadIdx.x == 0) atomicAdd(out, acc);
}
// AdamW with the gradient scaled by gscale_ptr[0] (clipping), decoupled decay.
__global__ void k_adamw(float* p, const float* g, float* m, float* v, size_t n, float lr, float b1, float b2, float eps, float wd, float bc1, float bc2, const float* gscale_ptr) {
float gs = gscale_ptr ? gscale_ptr[0] : 1.0f;
GRID_LOOP(i, n) {
float gi = g[i] * gs;
float mi = b1 * m[i] + (1.0f - b1) * gi;
float vi = b2 * v[i] + (1.0f - b2) * gi * gi;
m[i] = mi; v[i] = vi;
float upd = (mi / bc1) / (sqrtf(vi / bc2) + eps);
p[i] = p[i] * (1.0f - lr * wd) - lr * upd;
}
}
// clip: scale = min(1, max_norm / sqrt(sumsq))
__global__ void k_clip_scale(const float* sumsq, float max_norm, float* scale) {
if (threadIdx.x == 0 && blockIdx.x == 0) { float nrm = sqrtf(sumsq[0]); scale[0] = nrm > max_norm ? max_norm / nrm : 1.0f; }
}
__global__ void k_ema(float* e, const float* p, float decay, size_t n) { GRID_LOOP(i, n) e[i] = decay * e[i] + (1.0f - decay) * p[i]; }
// Counter-based normal noise (for the flow refiner's training targets).
__device__ __forceinline__ uint64_t mix64(uint64_t z) {
z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ull;
z = (z ^ (z >> 27)) * 0x94D049BB133111EBull;
return z ^ (z >> 31);
}
__global__ void k_randn(float* y, size_t n, uint64_t seed) {
GRID_LOOP(i, n) {
uint64_t a = mix64(seed + 2 * i * 0x9E3779B97F4A7C15ull), b = mix64(seed + (2 * i + 1) * 0x9E3779B97F4A7C15ull);
float u1 = ((a >> 40) + 1) * (1.0f / 16777217.0f);
float u2 = (b >> 40) * (1.0f / 16777216.0f);
y[i] = sqrtf(-2.0f * logf(u1)) * cospif(2.0f * u2);
}
}
// The harmonic source from per-sample controls (see dsp.rs harmonic_sample).
__global__ void k_harmonic(const float* ph, const float* hz, const float* amp, float* y, size_t n, float top, float level) {
GRID_LOOP(i, n) {
float f = hz[i], a = amp[i];
float acc = 0.0f;
if (a > 0.0f && f > 0.0f) {
float taper_lo = top * 0.8f;
int count = (int)(top / f);
float c2 = 2.0f * cosf(ph[i]);
float sp = 0.0f, s = sinf(ph[i]);
for (int k = 1; k <= count; k++) {
float fk = k * f;
float g = fk > taper_lo ? (top - fk) / (top - taper_lo) : 1.0f;
acc += s * g;
float nx = c2 * s - sp;
sp = s;
s = nx;
}
}
y[i] = acc * level * a;
}
}
// ---------------------------------------------------------------------------
// Monotonic alignment search: logp [B, N, T] (token x frame log-likelihood),
// n_tok[b], n_frame[b]; writes each token's frame count dur[b, n] (u32).
// One block per item; threads over tokens; the DP table is written over logp.
// ---------------------------------------------------------------------------
__global__ void k_mas(float* logp, const uint32_t* n_tok, const uint32_t* n_frame, uint32_t* dur, int N, int T) {
int b = blockIdx.x;
int nt = n_tok[b], nf = n_frame[b];
float* q = logp + (size_t)b * N * T;
const float NEG = -1e30f;
// Column 0.
for (int i = threadIdx.x; i < nt; i += blockDim.x) q[(size_t)i * T] = i == 0 ? q[0] : NEG;
__syncthreads();
for (int j = 1; j < nf; j++) {
// Update rows from high to low is unnecessary: read column j-1 only.
for (int i = threadIdx.x; i < nt; i += blockDim.x) {
float stay = q[(size_t)i * T + j - 1];
float move = i > 0 ? q[(size_t)(i - 1) * T + j - 1] : NEG;
float best = fmaxf(stay, move);
// A token cannot start before frame i, nor leave too few frames.
if (i > j || nt - 1 - i > nf - 1 - j) best = NEG;
q[(size_t)i * T + j] = best + q[(size_t)i * T + j];
}
__syncthreads();
}
if (threadIdx.x == 0) {
for (int i = 0; i < N; i++) dur[(size_t)b * N + i] = 0;
int i = nt - 1;
for (int j = nf - 1; j >= 0; j--) {
dur[(size_t)b * N + i] += 1;
if (i > 0 && j > 0 && (i == j || q[(size_t)(i - 1) * T + j - 1] > q[(size_t)i * T + j - 1])) i--;
}
}
}
// logp[b, i, j] = dot[b, i, j] - 0.5 |mu_bi|^2 - 0.5 |mel_bj|^2 (in place over dot)
__global__ void k_mas_logp(float* dot, const float* mu, const float* mel, size_t B, int N, int T, int D) {
GRID_LOOP(x, B * (size_t)N * T) {
size_t b = x / ((size_t)N * T);
int i = (int)((x / T) % N), j = (int)(x % T);
const float* m = mu + (b * N + i) * D;
const float* e = mel + (b * T + j) * D;
float a = 0.0f, c = 0.0f;
for (int k = 0; k < D; k++) { a += m[k] * m[k]; c += e[k] * e[k]; }
dot[x] = dot[x] - 0.5f * a - 0.5f * c;
}
}
// Durations -> the encoder row of every frame (padded frames: the item's last row),
// and ln(1 + frames) per token as f32.
__global__ void k_dur_to_idx(const uint32_t* dur, uint32_t* idx, float* logdur, size_t B, int N, int T) {
size_t b = blockIdx.x * (size_t)blockDim.x + threadIdx.x;
if (b >= B) return;
int f = 0;
for (int i = 0; i < N; i++) {
uint32_t d = dur[b * N + i];
logdur[b * N + i] = log1pf((float)d);
for (uint32_t k = 0; k < d && f < T; k++) idx[b * T + f++] = (uint32_t)(b * N + i);
}
for (; f < T; f++) idx[b * T + f] = (uint32_t)(b * N + N - 1);
}
// ---------------------------------------------------------------------------
// Launch wrappers
// ---------------------------------------------------------------------------
#define L1(k, n, ...) do { k<<<mkt_blocks(n), MKT_THREADS, 0, st>>>(__VA_ARGS__); return cudaGetLastError(); } while (0)
extern "C" {
cudaError_t mkt_act_fwd(const float* x, float* y, size_t n, int op, cudaStream_t st) { L1(k_act_fwd, n, x, y, n, op); }
cudaError_t mkt_act_bwd(const float* x, const float* d, float* dx, size_t n, int op, cudaStream_t st) { L1(k_act_bwd, n, x, d, dx, n, op); }
cudaError_t mkt_add(const float* a, const float* b, float* y, size_t n, cudaStream_t st) { L1(k_add, n, a, b, y, n); }
cudaError_t mkt_mul(const float* a, const float* b, float* y, size_t n, cudaStream_t st) { L1(k_mul, n, a, b, y, n); }
cudaError_t mkt_scale(const float* x, float* y, float s, size_t n, cudaStream_t st) { L1(k_scale, n, x, y, s, n); }
cudaError_t mkt_axpy(const float* x, float* y, float s, size_t n, cudaStream_t st) { L1(k_axpy, n, x, y, s, n); }
cudaError_t mkt_mul_acc(const float* x, const float* o, float* y, size_t n, cudaStream_t st) { L1(k_mul_acc, n, x, o, y, n); }
cudaError_t mkt_fill(float* y, float v, size_t n, cudaStream_t st) { L1(k_fill, n, y, v, n); }
cudaError_t mkt_axpy_scalar(const float* x, float* y, float s, cudaStream_t st) { k_axpy_scalar<<<1, 32, 0, st>>>(x, y, s); return cudaGetLastError(); }
cudaError_t mkt_log_eps(const float* x, float* y, float eps, size_t n, cudaStream_t st) { L1(k_log_eps, n, x, y, eps, n); }
cudaError_t mkt_log_eps_bwd(const float* x, const float* d, float* dx, float eps, size_t n, cudaStream_t st) { L1(k_log_eps_bwd, n, x, d, dx, eps, n); }
cudaError_t mkt_add_row(const float* x, const float* v, float* y, size_t rows, size_t cols, cudaStream_t st) { L1(k_add_row, rows * cols, x, v, y, rows, cols); }
cudaError_t mkt_mul_row(const float* x, const float* v, float* y, size_t rows, size_t cols, cudaStream_t st) { L1(k_mul_row, rows * cols, x, v, y, rows, cols); }
cudaError_t mkt_mul_row_bwd_x(const float* d, const float* v, float* dx, size_t rows, size_t cols, cudaStream_t st) { L1(k_mul_row_bwd_x, rows * cols, d, v, dx, rows, cols); }
cudaError_t mkt_mul_col(const float* x, const float* s, float* y, size_t rows, size_t cols, cudaStream_t st) { L1(k_mul_col, rows * cols, x, s, y, rows, cols); }
cudaError_t mkt_mul_col_bwd_x(const float* d, const float* s, float* dx, size_t rows, size_t cols, cudaStream_t st) { L1(k_mul_col_bwd_x, rows * cols, d, s, dx, rows, cols); }
cudaError_t mkt_col_dot_acc(const float* d, const float* x, float* dv, size_t rows, size_t cols, cudaStream_t st) {
unsigned gy = (unsigned)((rows + 255) / 256); if (gy > 128) gy = 128; if (gy == 0) gy = 1;
dim3 grid((unsigned)((cols + 31) / 32), gy);
k_col_dot_acc<<<grid, dim3(32, 8), 0, st>>>(d, x, dv, rows, cols);
return cudaGetLastError();
}
cudaError_t mkt_row_dot_acc(const float* d, const float* x, float* ds, size_t rows, size_t cols, cudaStream_t st) {
k_row_dot_acc<<<(unsigned)((rows + 7) / 8), 256, 0, st>>>(d, x, ds, rows, cols);
return cudaGetLastError();
}
cudaError_t mkt_copy_cols(const float* src, size_t src_cols, size_t src_off, float* dst, size_t dst_cols, size_t dst_off, size_t rows, size_t width, int acc, cudaStream_t st) {
L1(k_copy_cols, rows * width, src, src_cols, src_off, dst, dst_cols, dst_off, rows, width, acc);
}
cudaError_t mkt_gather_rows(const float* x, const uint32_t* idx, float* y, size_t rows_out, size_t cols, cudaStream_t st) { L1(k_gather_rows, rows_out * cols, x, idx, y, rows_out, cols); }
cudaError_t mkt_scatter_rows_acc(const float* d, const uint32_t* idx, float* dx, size_t rows_out, size_t cols, cudaStream_t st) { L1(k_scatter_rows_acc, rows_out * cols, d, idx, dx, rows_out, cols); }
cudaError_t mkt_im2col(const float* x, float* col, size_t rows, size_t C, int k, int dil, size_t seg, cudaStream_t st) { L1(k_im2col, rows * (size_t)k * C, x, col, rows, C, k, dil, seg); }
cudaError_t mkt_col2im_acc(const float* dcol, float* dx, size_t rows, size_t C, int k, int dil, size_t seg, cudaStream_t st) { L1(k_col2im_acc, rows * C, dcol, dx, rows, C, k, dil, seg); }
cudaError_t mkt_dwconv(const float* x, const float* w, float* y, size_t rows, size_t C, int k, size_t seg, cudaStream_t st) { L1(k_dwconv, rows * C, x, w, y, rows, C, k, seg); }
cudaError_t mkt_dwconv_bwd_x(const float* d, const float* w, float* dx, size_t rows, size_t C, int k, size_t seg, cudaStream_t st) { L1(k_dwconv_bwd_x, rows * C, d, w, dx, rows, C, k, seg); }
cudaError_t mkt_dwconv_bwd_w(const float* d, const float* x, float* dw, size_t rows, size_t C, int k, size_t seg, cudaStream_t st) {
unsigned gz = (unsigned)((rows + 255) / 256); if (gz > 64) gz = 64; if (gz == 0) gz = 1;
dim3 grid((unsigned)((C + 31) / 32), (unsigned)k, gz);
k_dwconv_bwd_w<<<grid, dim3(32, 8), 0, st>>>(d, x, dw, rows, C, k, seg);
return cudaGetLastError();
}
cudaError_t mkt_layer_norm(const float* x, float* xhat, float* rstd, size_t rows, size_t cols, cudaStream_t st) {
k_layer_norm<<<(unsigned)rows, cols >= 512 ? 256 : 128, 0, st>>>(x, xhat, rstd, cols);
return cudaGetLastError();
}
cudaError_t mkt_layer_norm_bwd(const float* d, const float* xhat, const float* rstd, float* dx, size_t rows, size_t cols, cudaStream_t st) {
k_layer_norm_bwd<<<(unsigned)rows, cols >= 512 ? 256 : 128, 0, st>>>(d, xhat, rstd, dx, cols);
return cudaGetLastError();
}
cudaError_t mkt_rope(const float* x, float* y, size_t rows, size_t C, int heads, size_t seg, float sign, int acc, cudaStream_t st) {
L1(k_rope, rows * (size_t)heads * (C / heads / 2), x, y, rows, C, heads, seg, sign, acc);
}
cudaError_t mkt_split_heads(const float* x, float* y, size_t B, size_t T, int H, int dh, cudaStream_t st) { L1(k_split_heads, B * T * (size_t)H * dh, x, y, B, T, H, dh); }
cudaError_t mkt_merge_heads(const float* x, float* y, size_t B, size_t T, int H, int dh, int acc, cudaStream_t st) { L1(k_merge_heads, B * T * (size_t)H * dh, x, y, B, T, H, dh, acc); }
cudaError_t mkt_softmax_rows(float* s, size_t rows, size_t T, size_t S, int H, const uint32_t* key_len, float scale, cudaStream_t st) {
k_softmax_rows<<<(unsigned)rows, 128, 0, st>>>(s, T, S, H, key_len, scale);
return cudaGetLastError();
}
cudaError_t mkt_softmax_bwd_rows(const float* p, float* dp, size_t rows, size_t S, float scale, cudaStream_t st) {
k_softmax_bwd_rows<<<(unsigned)rows, 128, 0, st>>>(p, dp, S, scale);
return cudaGetLastError();
}
cudaError_t mkt_loss(const float* x, const float* t, const float* w, const uint32_t* col_lim, size_t seg, float* g, float* out, size_t rows, size_t cols, float inv_denom, int sq, cudaStream_t st) {
unsigned b = mkt_blocks(rows * cols); if (b > 1024) b = 1024;
k_loss<<<b, MKT_THREADS, 0, st>>>(x, t, w, col_lim, seg, g, out, rows, cols, inv_denom, sq);
return cudaGetLastError();
}
cudaError_t mkt_scaled_acc(const float* g, const float* d, float* dx, size_t n, cudaStream_t st) { L1(k_scaled_acc, n, g, d, dx, n); }
cudaError_t mkt_bce(const float* x, const float* t, float* g, float* out, size_t n, cudaStream_t st) {
unsigned b = mkt_blocks(n); if (b > 1024) b = 1024;
k_bce<<<b, MKT_THREADS, 0, st>>>(x, t, g, out, n);
return cudaGetLastError();
}
cudaError_t mkt_hn_filter(const float* gh, const float* gn, const float* pa, const float* pb, const float* hr, const float* hi, const float* nr, const float* ni, float* y, size_t rows, size_t F, cudaStream_t st) {
L1(k_hn_filter, rows * F, gh, gn, pa, pb, hr, hi, nr, ni, y, rows, F);
}
cudaError_t mkt_hn_filter_bwd(const float* d, const float* gh, const float* gn, const float* pa, const float* pb, const float* hr, const float* hi, const float* nr, const float* ni, float* dgh, float* dgn, float* dpa, float* dpb, size_t rows, size_t F, cudaStream_t st) {
L1(k_hn_filter_bwd, rows * F, d, gh, gn, pa, pb, hr, hi, nr, ni, dgh, dgn, dpa, dpb, rows, F);
}
cudaError_t mkt_fft(float* re, float* im, size_t batch, int n, int inverse, cudaStream_t st) {
int log2n = 0; while ((1 << log2n) < n) log2n++;
int threads = n / 2 < 512 ? n / 2 : 512;
size_t done = 0;
while (done < batch) {
size_t chunk = batch - done; if (chunk > 65535) chunk = 65535;
k_fft<<<(unsigned)chunk, threads, 2 * n * sizeof(float), st>>>(re + done * n, im + done * n, n, log2n, inverse);
done += chunk;
}
return cudaGetLastError();
}
cudaError_t mkt_frame(const float* x, const float* gain, const float* win, float* re, float* im, size_t B, size_t len, size_t frames, int n, int hop, cudaStream_t st) {
L1(k_frame, B * frames * (size_t)n, x, gain, win, re, im, B, len, frames, n, hop);
}
cudaError_t mkt_half(const float* fr, const float* fi, float* hr, float* hi, size_t rows, int n, float scale, cudaStream_t st) { L1(k_half, rows * (size_t)(n / 2 + 1), fr, fi, hr, hi, rows, n, scale); }
cudaError_t mkt_half_c(const float* fr, const float* fi, float* hr, float* hi, size_t rows, int n, float c_edge, float c_mid, cudaStream_t st) { L1(k_half_c, rows * (size_t)(n / 2 + 1), fr, fi, hr, hi, rows, n, c_edge, c_mid); }
cudaError_t mkt_hermitian(const float* hr, const float* hi, float* fr, float* fi, size_t rows, int n, float c_edge, float c_mid, cudaStream_t st) { L1(k_hermitian, rows * (size_t)n, hr, hi, fr, fi, rows, n, c_edge, c_mid); }
cudaError_t mkt_overlap_add(const float* fr, const float* win, const float* norm, float* y, size_t B, size_t len, size_t frames, int n, int hop, float scale, int acc, cudaStream_t st) {
L1(k_overlap_add, B * len, fr, win, norm, y, B, len, frames, n, hop, scale, acc);
}
cudaError_t mkt_mag(const float* re, const float* im, float* m, size_t n, cudaStream_t st) { L1(k_mag, n, re, im, m, n); }
cudaError_t mkt_mag_bwd(const float* d, const float* re, const float* im, const float* m, float* dre, float* dim, size_t n, cudaStream_t st) { L1(k_mag_bwd, n, d, re, im, m, dre, dim, n); }
cudaError_t mkt_split_ri(const float* y, float* re, float* im, size_t rows, size_t F, cudaStream_t st) { L1(k_split_ri, rows * F, y, re, im, rows, F); }
cudaError_t mkt_join_ri_acc(const float* re, const float* im, float* y, size_t rows, size_t F, cudaStream_t st) { L1(k_join_ri_acc, rows * F, re, im, y, rows, F); }
cudaError_t mkt_sumsq(const float* g, size_t n, float* out, cudaStream_t st) {
unsigned b = mkt_blocks(n); if (b > 1024) b = 1024;
k_sumsq<<<b, MKT_THREADS, 0, st>>>(g, n, out);
return cudaGetLastError();
}
cudaError_t mkt_adamw(float* p, const float* g, float* m, float* v, size_t n, float lr, float b1, float b2, float eps, float wd, float bc1, float bc2, const float* gscale, cudaStream_t st) {
L1(k_adamw, n, p, g, m, v, n, lr, b1, b2, eps, wd, bc1, bc2, gscale);
}
cudaError_t mkt_clip_scale(const float* sumsq, float max_norm, float* scale, cudaStream_t st) { k_clip_scale<<<1, 32, 0, st>>>(sumsq, max_norm, scale); return cudaGetLastError(); }
cudaError_t mkt_ema(float* e, const float* p, float decay, size_t n, cudaStream_t st) { L1(k_ema, n, e, p, decay, n); }
cudaError_t mkt_randn(float* y, size_t n, uint64_t seed, cudaStream_t st) { L1(k_randn, n, y, n, seed); }
cudaError_t mkt_mas_logp(float* dot, const float* mu, const float* mel, size_t B, int N, int T, int D, cudaStream_t st) { L1(k_mas_logp, B * (size_t)N * T, dot, mu, mel, B, N, T, D); }
cudaError_t mkt_dur_to_idx(const uint32_t* dur, uint32_t* idx, float* logdur, size_t B, int N, int T, cudaStream_t st) {
k_dur_to_idx<<<(unsigned)((B + 63) / 64), 64, 0, st>>>(dur, idx, logdur, B, N, T);
return cudaGetLastError();
}
cudaError_t mkt_harmonic(const float* ph, const float* hz, const float* amp, float* y, size_t n, float top, float level, cudaStream_t st) { L1(k_harmonic, n, ph, hz, amp, y, n, top, level); }
cudaError_t mkt_mas(float* logp, const uint32_t* n_tok, const uint32_t* n_frame, uint32_t* dur, size_t B, int N, int T, cudaStream_t st) {
k_mas<<<(unsigned)B, 256, 0, st>>>(logp, n_tok, n_frame, dur, N, T);
return cudaGetLastError();
}
}
Reference in a new issue View git blame Copy permalink