def compute_alphas_kernel(
acts: torch.Tensor,
denom: torch.Tensor,
alphas: torch.Tensor,
llForward: torch.Tensor,
xlen: torch.Tensor,
ylen: torch.Tensor,
mlabels: torch.Tensor, # [B]
minibatch: int,
maxT: int,
maxU: int,
alphabet_size: int,
blank_: int,
):
"""
Compute alpha (forward variable) probabilities over the transduction step.
Args:
acts: Tensor of shape [B, T, U, V+1] flattened. Represents the logprobs activation tensor.
denom: Tensor of shape [B, T, U] flattened. Represents the denominator of the logprobs activation tensor
across entire vocabulary.
alphas: Zero tensor of shape [B, T, U]. Will be updated inside the kernel with the forward variable
probabilities.
llForward: Zero tensor of shape [B]. Represents the log-likelihood of the forward pass.
Returned as the forward pass loss that is reduced by the optimizer.
xlen: Vector of length B which contains the actual acoustic sequence lengths in the padded
activation tensor.
ylen: Vector of length B which contains the actual target sequence lengths in the padded
activation tensor.
mlabels: Matrix of shape [B, U+1] (+1 here is due to <SOS> token - usually the RNNT blank).
The matrix contains the padded target transcription that must be predicted.
minibatch: Int representing the batch size.
maxT: The maximum possible acoustic sequence length. Represents T in the logprobs tensor.
maxU: The maximum possible target sequence length. Represents U in the logprobs tensor.
alphabet_size: The vocabulary dimension V+1 (inclusive of RNNT blank).
blank_: Index of the RNNT blank token in the vocabulary. Generally the first or last token in the vocab.
Updates:
Kernel inplace updates the following inputs:
- alphas: forward variable scores.
- llForward: log-likelihood of forward variable.
"""
# // launch B blocks, each block has U threads
b = cuda.blockIdx.x # // batch id
u = cuda.threadIdx.x # label id, u
T = xlen[b] # select AM length of current sample
U = ylen[b] + 1 # select target length of current sample, +1 for the blank token
labels: torch.Tensor = mlabels[b] # mb label start point, equivalent to mlabels + b * (maxU - 1)
offset = b * maxT * maxU # pointer indexing offset
# alphas += offset # pointer offset, ignored since we explicitly add offset
# Initilize alpha[b, t=0, u=0] for all b in B
if u == 0:
alphas[offset] = 0
# sync until all alphas are initialized
cuda.syncthreads()
# Ordinary alpha calculations, broadcast across B=b and U=u
# Look up forward variable calculation from rnnt_numpy.forward_pass()
for n in range(1, T + U - 1):
t = n - u
if u == 0:
# for t in range(1, T) step to initialize alphas[b, t, 0]
if t > 0 and t < T:
alphas[offset + t * maxU + u] = alphas[offset + (t - 1) * maxU + u] + logp(
denom, acts, maxT, maxU, alphabet_size, b, t - 1, 0, blank_
)
elif u < U:
# for u in range(1, U) step to initialize alphas[b, 0, u]
if t == 0:
alphas[offset + u] = alphas[offset + u - 1] + logp(
denom, acts, maxT, maxU, alphabet_size, b, 0, u - 1, labels[u - 1]
)
# for t in range(1, T) for u in range(1, U) step to compute alphas[b, t, u]
elif t > 0 and t < T:
no_emit = alphas[offset + (t - 1) * maxU + u] + logp(
denom, acts, maxT, maxU, alphabet_size, b, t - 1, u, blank_
)
emit = alphas[offset + t * maxU + u - 1] + logp(
denom, acts, maxT, maxU, alphabet_size, b, t, u - 1, labels[u - 1]
)
alphas[offset + t * maxU + u] = rnnt_helper.log_sum_exp(emit, no_emit)
# sync across all B=b and U=u
cuda.syncthreads()
# After final sync, alphas[b, T-1, U - 1] + logprobs[b, T-1, U-1, blank] + denom[b, T-1, U-1] gives
# log-likelihood of forward pass.
if u == 0:
loglike = alphas[offset + (T - 1) * maxU + U - 1] + logp(
denom, acts, maxT, maxU, alphabet_size, b, T - 1, U - 1, blank_
)
llForward[b] = loglike
def _kernel(x, y):
x_pos = cuda.grid(1)
if x_pos < x.shape[0] and x_pos < y.shape[0]:
x[x_pos] = rnnt_helper.log_sum_exp(x[x_pos], y[x_pos])