Hugging Face's logo

Kernels:
phanerozoic
/
dpx-decode

Kernel card Files
xet
 Community
dpx-decode / dpx_decode_cuda /dpx_decode.cu
phanerozoic's picture
phanerozoic
dpx-decode v1: kernel-builder source tree
1a42548 verified
Raw
History Blame
16.8 kB
// Dynamic-programming decoders on DPX instructions.
//
// Three wavefront DP kernels for speech and alignment pipelines:
// viterbi max-plus HMM decoding with backtrace
// dtw monotonic dynamic time warping (min-plus, 3-way)
// ctc_forced_align CTC trellis Viterbi aligning a transcript to emissions
//
// Each op has two numeric paths. The float32 path matches float reference
// implementations. The int32 path quantizes log-domain values to fixed point
// (scale quanta per unit) and runs on the DPX fused min/max instructions
// (__vimin3_s32 / __vimax3_s32), single hardware instructions on sm_90+ and
// bit-identically emulated by the compiler on sm_80-89. Viterbi packs
// (score << 16 | argument) into one int32 so the 3-way max resolves the
// argmax in the same instruction; per-timestep renormalization keeps packed
// scores inside 16 bits. Both paths are deterministic: fixed iteration
// order, no atomics, bitwise-reproducible outputs.
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/all.h>
#include <cfloat>
#include <climits>
#include <cstdint>
namespace {
constexpr int kThreads = 256;
constexpr int kNegBig = -(1 << 24); // safely below any renormalized score
__device__ __forceinline__ int imax3(int a, int b, int c) {
#if __CUDA_ARCH__ >= 800
 return __vimax3_s32(a, b, c);
#else
 return max(a, max(b, c));
#endif
}
__device__ __forceinline__ int imin3(int a, int b, int c) {
#if __CUDA_ARCH__ >= 800
 return __vimin3_s32(a, b, c);
#else
 return min(a, min(b, c));
#endif
}
__device__ int block_imax(int v) {
 __shared__ int smem[kThreads];
 smem[threadIdx.x] = v;
 __syncthreads();
 for (int s = kThreads / 2; s > 0; s >>= 1) {
 if (threadIdx.x < s) smem[threadIdx.x] = max(smem[threadIdx.x], smem[threadIdx.x + s]);
 __syncthreads();
 }
 int r = smem[0];
 __syncthreads();
 return r;
}
// ---------------------------------------------------------------------------
// Viterbi, int32 DPX path. One block per sequence. States parallel over
// threads; time sequential. delta values are renormalized each step so the
// packed (score << 16 | state) representation stays exact.
// ---------------------------------------------------------------------------
__global__ void viterbi_i32_kernel(const int32_t* __restrict__ emit /* [B,T,S] */,
 const int32_t* __restrict__ trans /* [S,S] */,
 const int32_t* __restrict__ prior /* [S] */, int64_t B,
 int64_t T, int64_t S, int32_t* __restrict__ bp /* [B,T,S] */,
 int32_t* __restrict__ path /* [B,T] */,
 int64_t* __restrict__ score /* [B] */,
 int32_t* __restrict__ work /* [B, 2S] */) {
 const int64_t b = blockIdx.x;
 int32_t* prev = work + b * 2 * S;
 int32_t* curr = prev + S;
 __shared__ int64_t norm_acc;
 if (threadIdx.x == 0) norm_acc = 0;
for (int64_t j = threadIdx.x; j < S; j += blockDim.x) {
 prev[j] = prior[j] + emit[(b * T + 0) * S + j];
 }
 __syncthreads();
for (int64_t t = 1; t < T; t++) {
 // renormalize prev to <= 0 so packed scores fit 16 bits
 int mx = INT_MIN;
 for (int64_t j = threadIdx.x; j < S; j += blockDim.x) mx = max(mx, prev[j]);
 mx = block_imax(mx);
 if (threadIdx.x == 0) norm_acc += mx;
 for (int64_t j = threadIdx.x; j < S; j += blockDim.x) prev[j] -= mx;
 __syncthreads();
for (int64_t j = threadIdx.x; j < S; j += blockDim.x) {
 // scores clamp to the int16 floor before packing; only candidates far
 // below the winner clamp, so the argmax is unaffected
 int best = INT_MIN;
 int64_t i = 0;
 for (; i + 2 <= S; i += 2) {
 const int s1 = max(prev[i] + trans[i * S + j], -32768);
 const int s2 = max(prev[i + 1] + trans[(i + 1) * S + j], -32768);
 const int p1 = (s1 << 16) | (int)i;
 const int p2 = (s2 << 16) | (int)(i + 1);
 best = imax3(best, p1, p2);
 }
 if (i < S) {
 const int s1 = max(prev[i] + trans[i * S + j], -32768);
 best = max(best, (s1 << 16) | (int)i);
 }
 curr[j] = (best >> 16) + emit[(b * T + t) * S + j];
 bp[(b * T + t) * S + j] = best & 0xFFFF;
 }
 __syncthreads();
 for (int64_t j = threadIdx.x; j < S; j += blockDim.x) prev[j] = curr[j];
 __syncthreads();
 }
if (threadIdx.x == 0) {
 int best = INT_MIN;
 int bj = 0;
 for (int64_t j = 0; j < S; j++) {
 if (prev[j] > best) { best = prev[j]; bj = (int)j; }
 }
 score[b] = (int64_t)best + norm_acc;
 path[b * T + T - 1] = bj;
 for (int64_t t = T - 1; t > 0; t--) {
 bj = bp[(b * T + t) * S + bj];
 path[b * T + t - 1] = bj;
 }
 }
}
// Viterbi, float32 path (explicit argmax compares; same structure).
__global__ void viterbi_f32_kernel(const float* __restrict__ emit, const float* __restrict__ trans,
 const float* __restrict__ prior, int64_t B, int64_t T,
 int64_t S, int32_t* __restrict__ bp,
 int32_t* __restrict__ path, float* __restrict__ score,
 float* __restrict__ work) {
 const int64_t b = blockIdx.x;
 float* prev = work + b * 2 * S;
 float* curr = prev + S;
 for (int64_t j = threadIdx.x; j < S; j += blockDim.x) {
 prev[j] = prior[j] + emit[(b * T + 0) * S + j];
 }
 __syncthreads();
 for (int64_t t = 1; t < T; t++) {
 for (int64_t j = threadIdx.x; j < S; j += blockDim.x) {
 float best = -FLT_MAX;
 int bi = 0;
 for (int64_t i = 0; i < S; i++) {
 const float c = prev[i] + trans[i * S + j];
 if (c > best) { best = c; bi = (int)i; }
 }
 curr[j] = best + emit[(b * T + t) * S + j];
 bp[(b * T + t) * S + j] = bi;
 }
 __syncthreads();
 for (int64_t j = threadIdx.x; j < S; j += blockDim.x) prev[j] = curr[j];
 __syncthreads();
 }
 if (threadIdx.x == 0) {
 float best = -FLT_MAX;
 int bj = 0;
 for (int64_t j = 0; j < S; j++) {
 if (prev[j] > best) { best = prev[j]; bj = (int)j; }
 }
 score[b] = best;
 path[b * T + T - 1] = bj;
 for (int64_t t = T - 1; t > 0; t--) {
 bj = bp[(b * T + t) * S + bj];
 path[b * T + t - 1] = bj;
 }
 }
}
// ---------------------------------------------------------------------------
// DTW (min-plus, steps: match/insert/delete). Anti-diagonal wavefront, one
// block per sequence, full D matrix stored for the backtrace.
// ---------------------------------------------------------------------------
template <typename T_, bool kIsInt>
__global__ void dtw_kernel(const T_* __restrict__ cost /* [B,N,M] */, int64_t B, int64_t N,
 int64_t M, T_* __restrict__ D /* [B,N,M] */,
 int32_t* __restrict__ path /* [B,N+M,2], -1 padded */,
 int32_t* __restrict__ path_len /* [B] */) {
 const int64_t b = blockIdx.x;
 const T_* c = cost + b * N * M;
 T_* d = D + b * N * M;
 const T_ BIG = kIsInt ? (T_)(1 << 30) : (T_)FLT_MAX;
for (int64_t diag = 0; diag < N + M - 1; diag++) {
 const int64_t i_lo = diag < M ? 0 : diag - M + 1;
 const int64_t i_hi = diag < N ? diag : N - 1;
 for (int64_t i = i_lo + threadIdx.x; i <= i_hi; i += blockDim.x) {
 const int64_t j = diag - i;
 T_ up = (i > 0) ? d[(i - 1) * M + j] : BIG;
 T_ left = (j > 0) ? d[i * M + j - 1] : BIG;
 T_ ul = (i > 0 && j > 0) ? d[(i - 1) * M + j - 1] : ((i == 0 && j == 0) ? (T_)0 : BIG);
 T_ m;
 if (kIsInt) {
 m = (T_)imin3((int)up, (int)left, (int)ul);
 } else {
 m = fminf((float)up, fminf((float)left, (float)ul));
 }
 d[i * M + j] = c[i * M + j] + m;
 }
 __syncthreads();
 }
if (threadIdx.x == 0) {
 int64_t i = N - 1, j = M - 1, n = 0;
 int32_t* p = path + b * (N + M) * 2;
 while (true) {
 p[n * 2] = (int32_t)i;
 p[n * 2 + 1] = (int32_t)j;
 n++;
 if (i == 0 && j == 0) break;
 const T_ up = (i > 0) ? d[(i - 1) * M + j] : BIG;
 const T_ left = (j > 0) ? d[i * M + j - 1] : BIG;
 const T_ ul = (i > 0 && j > 0) ? d[(i - 1) * M + j - 1] : BIG;
 if (ul <= up && ul <= left) { i--; j--; }
 else if (up <= left) { i--; }
 else { j--; }
 }
 path_len[b] = (int32_t)n;
 for (int64_t k = n; k < N + M; k++) { p[k * 2] = -1; p[k * 2 + 1] = -1; }
 }
}
// ---------------------------------------------------------------------------
// CTC forced alignment: Viterbi over the 2L+1 blank-interleaved trellis.
// Allowed predecessors of state s at time t: s, s-1, and s-2 when s is a
// non-blank whose label differs from the label at s-2. Backpointers store
// the step (0/1/2). One block per sequence.
// ---------------------------------------------------------------------------
template <typename T_, bool kIsInt>
__global__ void ctc_align_kernel(const T_* __restrict__ logp /* [B,T,C] */,
 const int64_t* __restrict__ tgt /* [B,L], -1 pad */, int64_t B,
 int64_t T, int64_t C, int64_t L, int64_t blank,
 uint8_t* __restrict__ bp /* [B,T,2L+1] */,
 int32_t* __restrict__ frames /* [B,T] label per frame */,
 T_* __restrict__ score /* [B] */,
 T_* __restrict__ work /* [B, 2*(2L+1)] */) {
 const int64_t b = blockIdx.x;
 int64_t Lb = 0;
 for (int64_t l = 0; l < L; l++) {
 if (tgt[b * L + l] >= 0) Lb++;
 }
 const int64_t Sn = 2 * Lb + 1;
 T_* prev = work + b * 2 * (2 * L + 1);
 T_* curr = prev + (2 * L + 1);
 const T_ NEG = kIsInt ? (T_)kNegBig : (T_)(-FLT_MAX / 2);
auto state_label = [&](int64_t s) -> int64_t {
 return (s % 2 == 0) ? blank : tgt[b * L + s / 2];
 };
for (int64_t s = threadIdx.x; s < Sn; s += blockDim.x) {
 prev[s] = (s <= 1) ? logp[(b * T) * C + state_label(s)] : NEG;
 }
 __syncthreads();
for (int64_t t = 1; t < T; t++) {
 for (int64_t s = threadIdx.x; s < Sn; s += blockDim.x) {
 const T_ stay = prev[s];
 const T_ from1 = (s >= 1) ? prev[s - 1] : NEG;
 T_ from2 = NEG;
 if (s >= 2 && s % 2 == 1 && state_label(s) != state_label(s - 2)) from2 = prev[s - 2];
 T_ best;
 int choice;
 if (kIsInt) {
 best = (T_)imax3((int)stay, (int)from1, (int)from2);
 } else {
 best = fmaxf((float)stay, fmaxf((float)from1, (float)from2));
 }
 choice = (best == stay) ? 0 : (best == from1 ? 1 : 2);
 curr[s] = best + logp[(b * T + t) * C + state_label(s)];
 bp[(b * T + t) * (2 * L + 1) + s] = (uint8_t)choice;
 }
 __syncthreads();
 for (int64_t s = threadIdx.x; s < Sn; s += blockDim.x) prev[s] = curr[s];
 __syncthreads();
 }
if (threadIdx.x == 0) {
 int64_t s = (Sn >= 2 && prev[Sn - 1] < prev[Sn - 2]) ? Sn - 2 : Sn - 1;
 score[b] = prev[s];
 for (int64_t t = T - 1; t >= 0; t--) {
 frames[b * T + t] = (int32_t)state_label(s);
 if (t > 0) s -= bp[(b * T + t) * (2 * L + 1) + s];
 }
 }
}
} // namespace
void viterbi(torch::Tensor& path, torch::Tensor& score, torch::Tensor& bp,
 torch::Tensor const& emissions, torch::Tensor const& transitions,
 torch::Tensor const& priors) {
 TORCH_CHECK(emissions.is_cuda() && emissions.dim() == 3 && emissions.is_contiguous(),
 "emissions must be contiguous [B, T, S] CUDA");
 const int64_t B = emissions.size(0), T = emissions.size(1), S = emissions.size(2);
 TORCH_CHECK(S <= 32768, "S <= 32768 (packed argmax uses 16 bits)");
 TORCH_CHECK(transitions.sizes() == torch::IntArrayRef({S, S}) && transitions.is_contiguous(),
 "transitions must be [S, S]");
 TORCH_CHECK(priors.numel() == S && priors.is_contiguous(), "priors must be [S]");
 TORCH_CHECK(path.dtype() == torch::kInt32 && path.sizes() == torch::IntArrayRef({B, T}));
 TORCH_CHECK(bp.dtype() == torch::kInt32 && bp.sizes() == torch::IntArrayRef({B, T, S}));
 const at::cuda::CUDAGuard guard(emissions.device());
 cudaStream_t stream = at::cuda::getCurrentCUDAStream();
if (emissions.dtype() == torch::kInt32) {
 TORCH_CHECK(transitions.dtype() == torch::kInt32 && priors.dtype() == torch::kInt32,
 "dtype mismatch");
 TORCH_CHECK(score.dtype() == torch::kInt64 && score.numel() == B);
 auto work = torch::empty({B, 2 * S}, emissions.options());
 viterbi_i32_kernel<<<B, kThreads, 0, stream>>>(
 emissions.const_data_ptr<int32_t>(), transitions.const_data_ptr<int32_t>(),
 priors.const_data_ptr<int32_t>(), B, T, S, bp.data_ptr<int32_t>(),
 path.data_ptr<int32_t>(), score.data_ptr<int64_t>(), work.data_ptr<int32_t>());
 } else if (emissions.dtype() == torch::kFloat32) {
 TORCH_CHECK(transitions.dtype() == torch::kFloat32 && priors.dtype() == torch::kFloat32,
 "dtype mismatch");
 TORCH_CHECK(score.dtype() == torch::kFloat32 && score.numel() == B);
 auto work = torch::empty({B, 2 * S}, emissions.options());
 viterbi_f32_kernel<<<B, kThreads, 0, stream>>>(
 emissions.const_data_ptr<float>(), transitions.const_data_ptr<float>(),
 priors.const_data_ptr<float>(), B, T, S, bp.data_ptr<int32_t>(),
 path.data_ptr<int32_t>(), score.data_ptr<float>(), work.data_ptr<float>());
 } else {
 TORCH_CHECK(false, "emissions must be int32 (DPX path) or float32");
 }
}
void dtw(torch::Tensor& path, torch::Tensor& path_len, torch::Tensor& D,
 torch::Tensor const& cost) {
 TORCH_CHECK(cost.is_cuda() && cost.dim() == 3 && cost.is_contiguous(),
 "cost must be contiguous [B, N, M] CUDA");
 const int64_t B = cost.size(0), N = cost.size(1), M = cost.size(2);
 TORCH_CHECK(D.sizes() == cost.sizes() && D.dtype() == cost.dtype());
 TORCH_CHECK(path.dtype() == torch::kInt32 &&
 path.sizes() == torch::IntArrayRef({B, N + M, 2}));
 TORCH_CHECK(path_len.dtype() == torch::kInt32 && path_len.numel() == B);
 const at::cuda::CUDAGuard guard(cost.device());
 cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 if (cost.dtype() == torch::kInt32) {
 dtw_kernel<int32_t, true><<<B, kThreads, 0, stream>>>(
 cost.const_data_ptr<int32_t>(), B, N, M, D.data_ptr<int32_t>(), path.data_ptr<int32_t>(),
 path_len.data_ptr<int32_t>());
 } else if (cost.dtype() == torch::kFloat32) {
 dtw_kernel<float, false><<<B, kThreads, 0, stream>>>(
 cost.const_data_ptr<float>(), B, N, M, D.data_ptr<float>(), path.data_ptr<int32_t>(),
 path_len.data_ptr<int32_t>());
 } else {
 TORCH_CHECK(false, "cost must be int32 (DPX path) or float32");
 }
}
void ctc_forced_align(torch::Tensor& frames, torch::Tensor& score, torch::Tensor const& log_probs,
 torch::Tensor const& targets, int64_t blank) {
 TORCH_CHECK(log_probs.is_cuda() && log_probs.dim() == 3 && log_probs.is_contiguous(),
 "log_probs must be contiguous [B, T, C] CUDA");
 const int64_t B = log_probs.size(0), T = log_probs.size(1), C = log_probs.size(2);
 TORCH_CHECK(targets.is_cuda() && targets.dtype() == torch::kInt64 && targets.dim() == 2 &&
 targets.size(0) == B && targets.is_contiguous(),
 "targets must be contiguous int64 [B, L] (-1 padded)");
 const int64_t L = targets.size(1);
 TORCH_CHECK(frames.dtype() == torch::kInt32 && frames.sizes() == torch::IntArrayRef({B, T}));
 TORCH_CHECK(blank >= 0 && blank < C, "blank out of range");
 const at::cuda::CUDAGuard guard(log_probs.device());
 cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 auto bp = torch::empty({B, T, 2 * L + 1}, log_probs.options().dtype(torch::kUInt8));
 if (log_probs.dtype() == torch::kInt32) {
 TORCH_CHECK(score.dtype() == torch::kInt32 && score.numel() == B);
 auto work = torch::empty({B, 2 * (2 * L + 1)}, log_probs.options());
 ctc_align_kernel<int32_t, true><<<B, kThreads, 0, stream>>>(
 log_probs.const_data_ptr<int32_t>(), targets.const_data_ptr<int64_t>(), B, T, C, L, blank,
 bp.data_ptr<uint8_t>(), frames.data_ptr<int32_t>(), score.data_ptr<int32_t>(),
 work.data_ptr<int32_t>());
 } else if (log_probs.dtype() == torch::kFloat32) {
 TORCH_CHECK(score.dtype() == torch::kFloat32 && score.numel() == B);
 auto work = torch::empty({B, 2 * (2 * L + 1)}, log_probs.options());
 ctc_align_kernel<float, false><<<B, kThreads, 0, stream>>>(
 log_probs.const_data_ptr<float>(), targets.const_data_ptr<int64_t>(), B, T, C, L, blank,
 bp.data_ptr<uint8_t>(), frames.data_ptr<int32_t>(), score.data_ptr<float>(),
 work.data_ptr<float>());
 } else {
 TORCH_CHECK(false, "log_probs must be int32 (DPX path) or float32");
 }
}