| """dpx-decode: dynamic-programming decoders on DPX instructions. |
| |
| Wavefront DP kernels for speech and alignment pipelines: Viterbi decoding, |
| monotonic DTW, and CTC forced alignment. Each op has a float32 path and an |
| int32 fixed-point path running on the DPX fused min/max instructions, |
| hardware on Hopper and newer (sm_90+), compiler-emulated bit-identically on |
| Ampere and Ada. Both paths are deterministic and bitwise reproducible: |
| fixed iteration order, documented tie-breaking, no atomics. |
| |
| The int32 path quantizes log-domain inputs to fixed point (`quantize`, |
| default 256 quanta per unit). Viterbi packs (score << 16 | state) so the |
| 3-way max resolves the argmax in the same DPX instruction; scores clamp to |
| the int16 floor after per-step renormalization, which cannot affect the |
| argmax. |
| """ |
| from typing import Optional, Tuple |
|
|
| import torch |
|
|
| from ._ops import ops |
|
|
| DEFAULT_SCALE = 256.0 |
|
|
|
|
| def quantize(x: torch.Tensor, scale: float = DEFAULT_SCALE) -> torch.Tensor: |
| """Fixed-point quantization for the int32 DPX paths.""" |
| return torch.round(x * scale).to(torch.int32) |
|
|
|
|
| def viterbi(emissions: torch.Tensor, transitions: torch.Tensor, |
| priors: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]: |
| """Max-plus Viterbi decoding. |
| |
| Args: |
| emissions: [B, T, S] float32 or int32 (quantized) CUDA log-scores. |
| transitions: [S, S] log transition scores, transitions[i, j] = i -> j. |
| priors: [S] initial log-scores (zeros when omitted). |
| Returns: |
| (path [B, T] int32, score [B]); score dtype is float32 on the float |
| path and int64 (in quanta) on the int path. Ties select the larger |
| predecessor index on the int path and the smaller on the float path; |
| both are deterministic. |
| """ |
| B, T, S = emissions.shape |
| if priors is None: |
| priors = torch.zeros(S, dtype=emissions.dtype, device=emissions.device) |
| path = torch.empty(B, T, dtype=torch.int32, device=emissions.device) |
| bp = torch.empty(B, T, S, dtype=torch.int32, device=emissions.device) |
| sdt = torch.int64 if emissions.dtype == torch.int32 else torch.float32 |
| score = torch.empty(B, dtype=sdt, device=emissions.device) |
| ops.viterbi(path, score, bp, emissions.contiguous(), transitions.contiguous(), |
| priors.contiguous()) |
| return path, score |
|
|
|
|
| def dtw(cost: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| """Monotonic DTW over a cost matrix (steps: diagonal, down, right). |
| |
| Args: |
| cost: [B, N, M] float32 or int32 CUDA costs (lower is better). |
| Returns: |
| (path [B, N+M, 2] int32, start-to-end, -1 padded; path_len [B] int32; |
| D [B, N, M] accumulated-cost matrix). Backtrace prefers diagonal, |
| then vertical, deterministically. |
| """ |
| B, N, M = cost.shape |
| D = torch.empty_like(cost.contiguous()) |
| path = torch.empty(B, N + M, 2, dtype=torch.int32, device=cost.device) |
| path_len = torch.empty(B, dtype=torch.int32, device=cost.device) |
| ops.dtw(path, path_len, D, cost.contiguous()) |
| return path.flip(1), path_len, D |
|
|
|
|
| def ctc_forced_align(log_probs: torch.Tensor, targets: torch.Tensor, |
| blank: int = 0) -> Tuple[torch.Tensor, torch.Tensor]: |
| """CTC forced alignment: per-frame labels for a known transcript. |
| |
| Args: |
| log_probs: [B, T, C] float32 or int32 (quantized) CUDA log-probs. |
| targets: [B, L] int64 label ids, -1 padded on the right. |
| blank: blank label id. |
| Returns: |
| (frames [B, T] int32 per-frame label ids including blanks, |
| score [B] best-path log-score; int32 in quanta on the int path). |
| """ |
| B, T, C = log_probs.shape |
| frames = torch.empty(B, T, dtype=torch.int32, device=log_probs.device) |
| sdt = torch.int32 if log_probs.dtype == torch.int32 else torch.float32 |
| score = torch.empty(B, dtype=sdt, device=log_probs.device) |
| ops.ctc_forced_align(frames, score, log_probs.contiguous(), targets.contiguous(), blank) |
| return frames, score |
|
|
|
|
| __all__ = ["viterbi", "dtw", "ctc_forced_align", "quantize", "DEFAULT_SCALE"] |
|
|
