reasoning
We need to write a custom CUDA kernel that replaces the forward computation of the given PyTorch model. The model implements a reparameterization trick for a mixture of Gaussians (when num_components > 1 and mu.dim() == 3) and a simple diagonal Gaussian otherwise. We need to write a ModelNew class that accepts the same constructor arguments and implements forward with the same signature, producing numerically equivalent outputs (fp32 tolerance atol=rtol=1e-4). We need to use torch.utils.cpp_extension.load_inline to compile and bind CUDA source at module load time.
We need to optimize the kernel to run faster than the PyTorch reference. So we need to write an efficient CUDA kernel.
First, let's understand the reference implementation.
Case 1: num_components > 1 and mu.dim() == 3:
- std = exp(0.5 * log_var) (shape: (B, K, D))
- weights = softmax(logits, dim=0) (shape: (K,))
- Sample component: comp = torch.multinomial(weights.expand(B, -1), 1).squeeze(-1) (shape: (B,))
- batch_idx = arange(B)
- mu_k = mu[batch_idx, comp] (shape: (B, D))
- std_k = std[batch_idx, comp] (shape: (B, D))
- z = mu_k + std_k * randn_like(mu_k) (shape: (B, D))
- log_probs = -0.5 * ((z.unsqueeze(1) - mu) / std)^2 - log_var * 0.5 - 0.5 * log(2*pi) (shape: (B, K, D))
- log_probs = log_probs.sum(-1) + weights.log().unsqueeze(0) (shape: (B, K))
- log_q_z = torch.logsumexp(log_probs, dim=1) (shape: (B,))
Note: The PyTorch reference uses torch.logsumexp which is numerically stable.
Case 2: else. If num_components > 1 but mu.dim() != 3? Actually the else handles mu.dim() == 3 (but num_components == 1? Wait: if num_components > 1 and mu.dim() == 3, it's case 1. Else: if mu.dim() == 3, then mu = mu[:, 0], std = std[:, 0], log_var = log_var[:, 0]. Then z = mu + std * randn_like(mu). log_q_z = sum over D of -0.5 * ((z - mu)/std)^2 - 0.5 * log_var - 0.5 * log(2*pi). Since z - mu = std * eps, ((z - mu)/std)^2 = eps^2. But we compute it directly.
If mu.dim() == 2 (B, D), then the else directly computes z and log_q_z.
So the main heavy work is in case 1 (mixture). We need to implement a CUDA kernel that can handle both cases efficiently.
Observations:
- The reference uses torch.multinomial which calls a CPU-based sampler? Actually torch.multinomial on GPU uses a GPU kernel, but it's sequential per batch element (maybe not super efficient). We can try to implement a parallel sampling of components from the categorical distribution. However, for B=256 and K=4, the sampling is trivial overhead. The heavy part is computing log_probs for all K components and D dimensions.
- The log_probs computation in case 1 involves: for each batch element b, each component k, each dimension d:
diff = z[b, d] - mu[b, k, d]
term = -0.5 * (diff / std[b, k, d])^2 - 0.5 * log_var[b, k, d] - 0.5 * log(2*pi)
Then sum over d, add log(weights[k]), then logsumexp over k.
- For K=4, D=64, B=256, total elements in log_probs before sum: B*K*D = 256*4*64 = 65536. That's not huge. But we need to be faster than PyTorch. PyTorch reference does several launches: exp, softmax, multinomial, indexing, randn, unsqueeze, subtraction, division, square, sum, logsumexp. Fusing these into a single kernel or a few kernels can reduce kernel launch overhead and memory bandwidth.
- The main bottleneck is memory bandwidth if the tensors are small. Fusing the computation of z and log_q_z in a single kernel can be beneficial.
- Also note that the model uses self.logits as a parameter. In ModelNew forward, we need to compute weights = softmax(logits) and log_weights = log(weights). This is tiny (K elements), can be done on CPU or GPU.
- Since we need to use a custom CUDA kernel, we need to handle the random number generation. PyTorch's torch.randn_like uses Philox random number generator. In CUDA, we can use curand or custom Philox. However, PyTorch provides ATen/cuda/CUDAGeneratorImpl.h and at::cuda::philox::unpack to generate random numbers, but this is complicated. Since the kernel is small, we can generate the random numbers using PyTorch's torch.randn_like before calling the custom kernel, or we can generate them on the fly using curand. But since we need speed, generating random numbers in the kernel avoids an extra memory read/write. However, using curand may be slower due to state setup.
Alternatively, we can call torch.randn_like outside the kernel and pass it as an argument. The kernel then just does mu_k + std_k * eps. But we still need to sample the component indices comp. For the mixture case, we need comp to select mu_k and std_k. We can generate the component samples in the kernel or outside.
Simplify: Write two CUDA kernels: one for the simple case (num_components==1 or mu.dim()==2) that fuses everything, and one for the mixture case (num_components>1 and mu.dim()==3). Actually we can fuse everything for the mixture case into a single kernel that computes z and log_q_z.
For the mixture case:
Inputs: mu (B, K, D), log_var (B, K, D), weights (K,), log_weights (K,).
Outputs: z (B, D), log_q_z (B,).
We also need random numbers: eps (B, D) and comp_rand (B,) to sample components.
Wait: The component sampling uses multinomial from weights. For K=4, we can sample using the inverse CDF method: generate u ~ Uniform(0,1), then find the smallest k such that cumulative weight >= u. This is sequential over K but K is only 4. We can do it per thread. Alternatively, we can generate comp on CPU and pass it in, or generate it using PyTorch's multinomial. Since B=256, generating on CPU and copying to GPU is fine, but we need to be careful about device synchronization. Actually, we can just use torch.multinomial on GPU (it returns a tensor on GPU). But if we want to fuse, we can generate uniform random numbers in the kernel and do the selection. However, generating uniform random numbers from a Philox state is nontrivial without PyTorch's generator. We can generate uniform using CUDA's curandState but need to initialize states. For B=256, it's fine.
Actually, a simpler approach: In ModelNew.forward, we can generate the random numbers using PyTorch's standard functions (randn, rand) before the kernel launch, then pass them to the custom kernel. The kernel then does the deterministic fused computation. This avoids dealing with random number generation in CUDA and is likely sufficient for speedup because the deterministic part dominates.
But we can also generate random numbers inside the kernel using a simple LCG or Philox if we want to eliminate the extra kernel launch. However, the evaluation may check determinism? Actually, the forward pass is stochastic, but correctness is checked with tolerance, so any valid random numbers are fine as long as the distribution is correct. But the reference uses torch.randn_like which uses a specific RNG. We don't need to match the exact random draws; we just need the output to be within tolerance for the same input? Wait, the correctness check calls ModelNew.forward with the same inputs, but the reference and new model may use different random draws. The output depends on random draws, so if we use different random numbers, the outputs will be different and the check may fail if tolerance is tight. Actually, the correctness check probably runs the model on the same inputs and compares the outputs. Since the forward pass is stochastic, the reference and ModelNew will be called once each, and their outputs compared. If we use independent random numbers, z will differ, and log_q_z will be computed based on z, so it will differ. The tolerance is fp32 atol=rtol=1e-4. The difference in z could be large (order 1). That would fail.
Therefore, we must use the same random draws as the reference? Or we must seed the RNG deterministically? Wait, how does the correctness check handle stochasticity? It might set the random seed before calling forward. Let's assume that the correctness check sets a fixed random seed, and both Model and ModelNew are called with the same seed, so if we use the same RNG calls (e.g., torch.randn_like), we get the same random numbers. If we generate random numbers inside the kernel using a different method, we might get different numbers and fail.
However, we can call the same PyTorch RNG functions before launching the kernel, and then use those values. For example, in the mixture case, we need both component samples and Gaussian noise. The reference calls torch.multinomial and torch.randn_like. We can call torch.multinomial and torch.randn_like in ModelNew.forward as well, then pass the results to the kernel. This ensures the same random draws (assuming same seed and same order of RNG calls). But if we want speed, we should avoid the overhead of the separate PyTorch ops. But correctness is paramount.
Actually, we can replace torch.multinomial with a custom parallel sampler that uses the same uniform random numbers? No, we can't guarantee the same draws. So to be safe, we should rely on PyTorch's RNG for the random bits.
But note: The reference calls torch.multinomial, which consumes random numbers from the CUDA generator. Then it calls torch.randn_like, which consumes more random numbers. If we replicate these calls in ModelNew, the same random numbers will be drawn in the same order, producing identical comp and eps. Then the rest of the computation is deterministic, so the outputs will match exactly (within floating point error).
Therefore, I propose:
- In ModelNew.forward, compute std = torch.exp(0.5 * log_var) (or compute it in kernel, but doing it in PyTorch is fine).
- If mixture case:
- Compute weights = torch.softmax(self.logits, dim=0) and log_weights = weights.log(). (This is deterministic and tiny.)
- Sample comp using torch.multinomial(weights.expand(B, -1), 1).squeeze(-1). (Same as reference.)
- Generate eps = torch.randn_like(mu_k? Actually we need eps of shape (B, D). The reference does torch.randn_like(mu_k) which is (B, D). So we can do eps = torch.randn(B, D, device=mu.device, dtype=mu.dtype). Or we can do torch.randn_like(mu[:,0]) or something.
- Then launch a custom CUDA kernel that takes mu, log_var, std, comp, eps, weights, log_weights and computes z and log_q_z.
Wait, we need mu_k and std_k. Instead of indexing in PyTorch (which causes memory copy), we can pass the full mu and std and the comp indices to the kernel, and have each thread read the appropriate component. This avoids a separate indexing kernel.
- Similarly, for the simple case:
- If mu.dim() == 3: we need to slice mu[:,0], std[:,0], log_var[:,0]. We can pass the full tensors and have the kernel read the first component. Or slice in PyTorch (it's a view, not a copy). Slicing is cheap. Then generate eps = torch.randn_like(mu_sliced).
- Launch a kernel that computes z = mu + std * eps and log_q_z.
The heavy computation is log_q_z. Let's write a fused kernel for both cases.
Given the small sizes (B=256, K=4, D=64), a simple one-thread-per-element or thread-per-batch with reduction should be efficient.
For the mixture case:
- For each batch element b:
- Determine component k = comp[b].
- For each dimension d (thread within block):
- Read mu_bkd = mu[b][k][d], std_bkd = std[b][k][d], eps_bd = eps[b][d].
- z[b][d] = mu_bkd + std_bkd * eps_bd.
- Then compute log_q_z[b] via logsumexp over components.
- For each component k (0..K-1):
- Compute sum over d of -0.5 * ((z[b][d] - mu[b][k][d]) / std[b][k][d])^2 - 0.5 * log_var[b][k][d] - 0.5 * log(2*pi).
- Add log_weights[k].
- Then logsumexp over k.
Since D=64 and K=4, we can compute the per-component sum using warp reduction or shared memory. Each block can handle one batch element. With D=64, we can have 64 threads per block (or 128). K=4, so after computing z, each block computes 4 component log-likelihoods. This is very small.
For the simple case:
- z[b][d] = mu[b][0][d] + std[b][0][d] * eps[b][d] (if mu.dim()==3) or directly if mu.dim()==2.
- log_q_z[b] = sum_d [ -0.5 * ((z[b][d] - mu[b][d]) / std[b][d])^2 - 0.5 * log_var[b][d] - 0.5 * log(2*pi) ].
This can also be fused: one thread per d, warp sum.
But note: The reference computes log_q_z using the sliced mu, std, log_var. If we pass the full (B,K,D) tensors and comp to the kernel, we don't need slicing or indexing in PyTorch. However, we do need to generate eps of shape (B,D). That's fine.
Alternatively, we can generate eps inside the kernel using curand, but as argued, it's safer to use PyTorch's RNG. But we can also generate eps using PyTorch's randn and then fuse the deterministic part.
Wait: For the simple case, if mu.dim()==3, the reference does:
mu = mu[:, 0]
std = std[:, 0]
log_var = log_var[:, 0]
z = mu + std * torch.randn_like(mu)
log_q_z = (-0.5 * ((z - mu)/std)^2 - 0.5*log_var - 0.5*log(2*pi)).sum(-1)
So the reference slices mu, std, log_var. We can replicate that in ModelNew by slicing them (which is a view) and then passing those views to the kernel. The kernel then computes z and log_q_z for mu shape (B,D).
Let's design the kernel API.
We can define two kernels:
1. `mixture_forward_kernel`: inputs: `const float* mu`, `const float* log_var`, `const float* std`, `const long* comp`, `const float* eps`, `const float* log_weights`, `int B`, `int K`, `int D`, `float* z`, `float* log_q_z`.
- Computes z using comp and eps.
- Computes log_q_z via logsumexp over components.
2. `simple_forward_kernel`: inputs: `const float* mu`, `const float* log_var`, `const float* std`, `const float* eps`, `int B`, `int D`, `float* z`, `float* log_q_z`.
- Computes z = mu + std * eps.
- Computes log_q_z = sum_d (...).
But wait, for the simple case, the reference has a subtlety: if mu.dim()==3, it slices to mu[:,0]. So the kernel sees (B,D). If mu.dim()==2, it's already (B,D). So the kernel is the same.
Actually, we can write a single kernel that handles both cases if we pass flags, but it's easier to write two.
Memory layout: We assume contiguous tensors. In PyTorch, mu and log_var passed to forward are contiguous? We can't guarantee, but we can call .contiguous() in forward before kernel launch. That adds overhead but ensures correctness. For speed, we can assume they are contiguous or use strides.
Better: Use TensorAccessor or pointer with strides. In CUDA, using strided access is fine.
Let's implement the kernels.
### Kernel design for mixture case
```
__global__ void mixture_kernel(
const float* mu,
const float* log_var,
const float* std,
const int64_t* comp,
const float* eps,
const float* log_weights,
int B, int K, int D,
float* z,
float* log_q_z)
{
int b = blockIdx.x;
if (b >= B) return;
int tid = threadIdx.x;
int num_threads = blockDim.x;
// compute z[b, d] = mu[b, comp[b], d] + std[b, comp[b], d] * eps[b, d]
int k_sel = comp[b];
// mu stride: assume mu is (B, K, D) contiguous => stride K*D, D, 1.
for (int d = tid; d < D; d += num_threads) {
int idx = ((b * K) + k_sel) * D + d;
float m = mu[idx];
float s = std[idx];
float e = eps[b * D + d];
float z_val = m + s * e;
z[b * D + d] = z_val;
}
__syncthreads();
// compute per-component log-likelihood sum over D
// We'll compute for each k, sum over d.
// Use shared memory for partial sums. Since K is small (<=4 maybe) and D=64,
// we can have each thread compute partial sums for all k in a loop.
// Then reduce.
__shared__ float s_log_probs[4]; // max components
// initialize s_log_probs[k] = log_weights[k]
if (tid < K) {
s_log_probs[tid] = log_weights[tid];
}
__syncthreads();
// Each thread processes a subset of d and accumulates into registers for each k.
float local_sum[4]; // K max 4
#pragma unroll
for (int k = 0; k < K; ++k) local_sum[k] = 0.0f;
const float log_2pi = 1.8378770664093453f; // log(2*pi)
for (int d = tid; d < D; d += num_threads) {
float z_val = z[b * D + d];
for (int k = 0; k < K; ++k) {
int idx = ((b * K) + k) * D + d;
float m = mu[idx];
float s = std[idx];
float lv = log_var[idx];
float diff = z_val - m;
float term = -0.5f * (diff / s) * (diff / s) - 0.5f * lv - 0.5f * log_2pi;
local_sum[k] += term;
}
}
// reduce local_sum across threads using shared memory or warp shuffle
// Since K is small, we can reduce each k separately.
__shared__ float s_reduce[4];
// Wait, we need to reduce per k. We can use warp shuffle for each k.
for (int k = 0; k < K; ++k) {
float val = local_sum[k];
// warp reduce
for (int offset = warpSize/2; offset > 0; offset /= 2) {
val += __shfl_down_sync(0xffffffff, val, offset);
}
if (tid % warpSize == 0) {
atomicAdd(&s_reduce[k], val);
}
}
__syncthreads();
if (tid == 0) {
float max_val = -1e20f;
for (int k = 0; k < K; ++k) {
float val = s_reduce[k] + s_log_probs[k];
s_reduce[k] = val; // store log prob
if (val > max_val) max_val = val;
}
float lse = 0.0f;
for (int k = 0; k < K; ++k) {
lse += expf(s_reduce[k] - max_val);
}
log_q_z[b] = max_val + logf(lse);
}
}
```
Wait, the atomicAdd may have issues because s_reduce isn't initialized to zero. We need to zero s_reduce before atomicAdd. Also, using atomicAdd in shared memory is safe but we can avoid it by having the first warp do the reduction or by using shared memory in a more structured way.
Alternative: Since D=64 and K=4, we can have 64 threads per block (one per d). Each thread computes the term for each k? That would be 64 threads * 4 = 256 FMAs, which is fine. Then we can use warp-level reduction.
But for generality (D could be larger though latent_dim is 64 in given config), we can handle any D.
Better approach: Have blockDim.x = 128 (or next power of 2 >= D). Each thread computes a partial sum for all k over d in stride. Then we reduce per k.
We can allocate shared memory: `__shared__ float shmem[128 * 4];` where each thread stores its partial sum for each k. Then after the loop, we do a block-level reduction per k. But that's a bit more complex.
Simpler: Use warp shuffle reduction per k per thread. Since K is small, we can just loop over k and do warp shuffle.
However, atomicAdd to shared memory is okay if we initialize to zero. But there is a subtlety: multiple warps may write to the same shared memory location. atomicAdd handles it, but it's slower. However, with only 4 components and maybe 2 warps, it's negligible.
Let's refine the reduction:
```
__shared__ float s_reduce[4];
if (tid < K) s_reduce[tid] = 0.0f;
__syncthreads();
for (int k = 0; k < K; ++k) {
float val = local_sum[k];
// warp reduce
for (int offset = 16; offset > 0; offset /= 2) {
val += __shfl_down_sync(0xffffffff, val, offset);
}
if (tid % 32 == 0) {
atomicAdd(&s_reduce[k], val);
}
}
__syncthreads();
if (tid == 0) {
// compute logsumexp
}
```
This is fine. Note that if blockDim.x < 32, we need to be careful, but we can set blockDim.x to 32 or more.
Actually, we can also avoid atomicAdd by having a single warp (32 threads) handle the reduction if D <= 32? But D=64, so we need two warps. We can have the first thread of each warp add to shared memory after warp reduction:
```
if (tid % 32 == 0) {
s_reduce[k] += val;
}
```
But then we need to initialize s_reduce to zero. And we need to avoid race conditions between warps writing to the same location. Since they are adding, we could do:
```
if (tid % 32 == 0) {
atomicAdd(&s_reduce[k], val);
}
```
Or we can have each warp write to a different location and then reduce:
```
__shared__ float swarp[32][4]; // max warps
int warp_id = tid / 32;
if (tid % 32 == 0) {
for (int k=0; k<K; ++k) swarp[warp_id][k] = local_sum[k]; // after warp reduce
}
__syncthreads();
if (tid == 0) {
for (int k=0; k<K; ++k) {
float sum = 0;
for (int w=0; w<(blockDim.x+31)/32; ++w) sum += swarp[w][k];
// ...
}
}
```
That's also fine.
Let's keep it simple: Since D=64 and K is small, we can compute log_q_z using a loop over d in each thread and then use a simple block reduction with shared memory.
I'll write a helper for block reduction:
```
__inline__ __device__ float warp_reduce_sum(float val) {
for (int offset = warpSize/2; offset > 0; offset /= 2)
val += __shfl_down_sync(0xffffffff, val, offset);
return val;
}
__inline__ __device__ float block_reduce_sum(float val, float* shared) {
int lane = threadIdx.x % warpSize;
int wid = threadIdx.x / warpSize;
val = warp_reduce_sum(val);
if (lane == 0) shared[wid] = val;
__syncthreads();
val = (threadIdx.x < blockDim.x / warpSize) ? shared[lane] : 0.0f;
if (wid == 0) val = warp_reduce_sum(val);
return val;
}
```
Wait, the standard block reduce: After warp reduce, store per-warp sum to shared mem. Then the first warp loads these and reduces.
```
template <int BLOCK_SIZE>
__inline__ __device__ float block_reduce_sum(float val) {
__shared__ float s[BLOCK_SIZE / 32];
int lane = threadIdx.x % 32;
int wid = threadIdx.x / 32;
val = warp_reduce_sum(val);
if (lane == 0) s[wid] = val;
__syncthreads();
val = (threadIdx.x < blockDim.x / 32) ? s[threadIdx.x] : 0.0f;
if (wid == 0) val = warp_reduce_sum(val);
return val;
}
```
But the block size may not be known at compile time. We can use a fixed maximum block size, say 128.
However, since K is small, we can just have a loop over K and compute block_reduce_sum for each k.
Pseudo:
```
__shared__ float smem[128/32]; // for block reduce
float local[4];
for k: local[k]=0;
for d=tid; d<D; d+=blockDim.x:
float zv = z[b*D+d];
for k=0..K-1:
float m = mu[((b*K)+k)*D+d];
float s = std[...];
float lv = log_var[...];
float diff = zv - m;
local[k] += -0.5f*(diff/s)*(diff/s) - 0.5f*lv - 0.5f*log_2pi;
for k=0..K-1:
float sum_k = block_reduce_sum(local[k]);
if (tid==0) s reduce[k] = sum_k + log_weights[k];
__syncthreads();
if (tid==0) {
log_q_z[b] = logsumexp(s_reduce, K);
}
```
This is clean.
But note: The block_reduce_sum as written above requires `smem` to be sized with `blockDim.x / 32`. Since blockDim.x is a runtime variable, we can size smem based on the maximum block size we choose, e.g., 128. We can pass it as a template parameter.
Alternatively, we can just use `blockDim.x` to size shared memory dynamically: `extern __shared__ float smem[];` passed as the third kernel launch argument. That's more elegant.
Let's write:
```
__inline__ __device__ float block_reduce_sum(float val, float* shared) {
int lane = threadIdx.x % 32;
int wid = threadIdx.x / 32;
// warp reduce
for (int offset = 16; offset > 0; offset /= 2) {
val += __shfl_down_sync(0xffffffff, val, offset);
}
if (lane == 0) shared[wid] = val;
__syncthreads();
// read from shared memory
val = (threadIdx.x < (blockDim.x + 31) / 32) ? shared[threadIdx.x] : 0.0f;
if (wid == 0) {
for (int offset = 16; offset > 0; offset /= 2) {
val += __shfl_down_sync(0xffffffff, val, offset);
}
}
return val;
}
```
Wait, the second warp reduce: Only the first warp should do it. So:
```
if (wid == 0) {
for (int offset = 16; offset > 0; offset /= 2) {
val += __shfl_down_sync(0xffffffff, val, offset);
}
shared[0] = val;
}
__syncthreads();
return shared[0];
```
But we have to be careful: the read of shared[threadIdx.x] by all threads may be okay because threads beyond number of warps read 0. However, we only need the result in thread 0? Actually, we can have the result in all threads if we broadcast. Let's just have thread 0 compute logsumexp. So we can return the sum to all threads or just to thread 0.
Simplify: After warp reduce, store to shared. Then read and warp reduce in warp 0, store to shared[0]. Then `__syncthreads()`. Return `shared[0]` to all threads. That's fine.
Implementation:
```
__inline__ __device__ float block_reduce_sum(float val, float* shared) {
int lane = threadIdx.x % 32;
int wid = threadIdx.x / 32;
for (int offset = 16; offset > 0; offset /= 2)
val += __shfl_down_sync(0xffffffff, val, offset);
if (lane == 0) shared[wid] = val;
__syncthreads();
if (wid == 0) {
val = (threadIdx.x < (blockDim.x + 31) / 32) ? shared[lane] : 0.0f;
for (int offset = 16; offset > 0; offset /= 2)
val += __shfl_down_sync(0xffffffff, val, offset);
if (lane == 0) shared[0] = val;
}
__syncthreads();
return shared[0];
}
```
This is a standard pattern. Good.
Now, for the mixture kernel:
```
__global__ void mixture_forward_kernel(
const float* __restrict__ mu,
const float* __restrict__ log_var,
const float* __restrict__ std,
const int64_t* __restrict__ comp,
const float* __restrict__ eps,
const float* __restrict__ log_weights,
int B, int K, int D,
float* __restrict__ z,
float* __restrict__ log_q_z)
{
int b = blockIdx.x;
if (b >= B) return;
int tid = threadIdx.x;
int k_sel = static_cast<int>(comp[b]);
// compute z
for (int d = tid; d < D; d += blockDim.x) {
int idx = ((b * K) + k_sel) * D + d;
float m = mu[idx];
float s = std[idx];
float e = eps[b * D + d];
z[b * D + d] = m + s * e;
}
__syncthreads();
// compute per-component log-likelihoods
float local[4]; // max K
#pragma unroll
for (int k = 0; k < K; ++k) local[k] = 0.0f;
const float log_2pi = 1.8378770664093453f;
for (int d = tid; d < D; d += blockDim.x) {
float zv = z[b * D + d];
for (int k = 0; k < K; ++k) {
int idx = ((b * K) + k) * D + d;
float m = mu[idx];
float s = std[idx];
float lv = log_var[idx];
float diff = zv - m;
float term = -0.5f * (diff / s) * (diff / s) - 0.5f * lv - 0.5f * log_2pi;
local[k] += term;
}
}
extern __shared__ float smem[];
float comp_log_prob[4];
for (int k = 0; k < K; ++k) {
float sum_k = block_reduce_sum(local[k], smem);
if (tid == 0) {
comp_log_prob[k] = sum_k + log_weights[k];
}
}
__syncthreads();
if (tid == 0) {
// logsumexp
float max_val = comp_log_prob[0];
#pragma unroll
for (int k = 1; k < K; ++k) {
if (comp_log_prob[k] > max_val) max_val = comp_log_prob[k];
}
float lse = 0.0f;
#pragma unroll
for (int k = 0; k < K; ++k) {
lse += expf(comp_log_prob[k] - max_val);
}
log_q_z[b] = max_val + logf(lse);
}
}
```
But there is a subtlety: `block_reduce_sum` uses `__syncthreads()` internally, and we call it in a loop over k. Each call to `block_reduce_sum` has `__syncthreads()`. That's okay if the loop is unrolled and all threads participate. However, `__syncthreads()` inside a loop is fine.
But the `extern __shared__ float smem[];` is fine.
Wait: The block_reduce_sum uses `(blockDim.x + 31) / 32`. This is a runtime computation. Good.
Also, we need to make sure that blockDim.x is at most 1024. For D=64, we can set blockDim.x to e.g., 64. But we should choose a block size that is a multiple of 32. Let's choose min(128, nearest power of 2 >= D). Actually, since D=64, block size 64 is perfect. But to handle arbitrary D, we can compute block size in host code.
For the simple case:
```
__global__ void simple_forward_kernel(
const float* __restrict__ mu,
const float* __restrict__ log_var,
const float* __restrict__ std,
const float* __restrict__ eps,
int B, int D,
float* __restrict__ z,
float* __restrict__ log_q_z)
{
int b = blockIdx.x;
if (b >= B) return;
int tid = threadIdx.x;
float local = 0.0f;
const float log_2pi = 1.8378770664093453f;
for (int d = tid; d < D; d += blockDim.x) {
int idx = b * D + d;
float m = mu[idx];
float s = std[idx];
float lv = log_var[idx];
float e = eps[idx];
float zv = m + s * e;
z[idx] = zv;
float diff = zv - m; // = s*e
float term = -0.5f * (diff / s) * (diff / s) - 0.5f * lv - 0.5f * log_2pi;
local += term;
}
extern __shared__ float smem[];
float sum_log_prob = block_reduce_sum(local, smem);
if (tid == 0) {
log_q_z[b] = sum_log_prob;
}
}
```
Note: For the simple case, if mu.dim()==3, we need to pass mu[:,0], std[:,0], log_var[:,0] to the kernel. We can do that in Python by slicing, or we can pass the full tensor with strides. Slicing is easiest.
Now, about the log(2*pi) constant: The reference uses `torch.log(torch.tensor(2 * 3.14159265))`. That's a scalar. We should compute it in Python and pass as a float argument to the kernel to ensure exact match? Actually, computing it as a float constant `logf(2.0f * 3.14159265f)` in CUDA is fine. But there might be tiny difference in precision. Better to compute `float log_2pi = 1.8378770664093453f;` which is exactly `log(2*pi)`. Wait, `torch.log(torch.tensor(2 * 3.14159265))` uses double precision on CPU? Actually, `torch.tensor(2 * 3.14159265)` creates a float64 tensor by default, then `torch.log` returns float64. When subtracted from float32 tensors, it gets cast to float32. The float32 value of `log(2 * 3.14159265)` might be slightly different from `logf(2.0f * 3.14159265f)`. Let's compute:
2 * 3.14159265 = 6.2831853
log(6.2831853) in double = 1.8378770664093453
In float32: it will be cast to float32, which is exactly representable? Not exactly, but it's a single float32 value.
If we compute in CUDA as `logf(2.0f * 3.14159265f)`, `2.0f * 3.14159265f` is a float32 multiplication, giving a slightly different value than the double multiplication, and then `logf` of that float32 value might differ slightly. It could cause small differences beyond tolerance? Unlikely, but to be safe, we can compute the constant in Python and pass it as a scalar argument to the kernel. Or we can compute `log(2*pi)` on CPU in float64 and cast to float32, then pass as a kernel argument.
In the reference, the constant is computed inside `forward`, so if we compute it the same way in Python and pass to the kernel, we match exactly.
Wait, we can also compute it in the kernel using the same literal: `-0.5f * logf(2.0f * 3.14159265f)`. The literal `3.14159265` in C is a double constant. If we write `3.14159265f`, it's float. The reference uses Python float `3.14159265`, which is double precision. In C, `2.0 * 3.14159265` is double. But if we use `2.0f * 3.14159265f`, it's float. The difference is minuscule.
Better: In Python, compute `log_2pi = torch.log(torch.tensor(2 * 3.14159265)).item()`. Then pass it as a float argument to the kernel. In CUDA, use that argument. This ensures exact match.
Similarly for the `softmax` of logits: In reference, `weights = torch.softmax(self.logits, dim=0)`. In our code, we can compute `weights = torch.softmax(self.logits, dim=0)` and `log_weights = weights.log()`. That's fine.
Now, the issue with random numbers: We need to replicate the exact random draws. For the mixture case:
- Reference: `comp = torch.multinomial(weights.expand(mu.shape[0], -1), 1).squeeze(-1)`
- Reference: `z = mu_k + std_k * torch.randn_like(mu_k)`
So we need to call these in the same order. In ModelNew.forward, we can call:
```
comp = torch.multinomial(weights.expand(B, -1), 1).squeeze(-1)
eps = torch.randn(B, D, device=mu.device, dtype=mu.dtype)
```
But `torch.randn_like(mu_k)` is equivalent to `torch.randn_like(mu[:,0])`. Wait, `mu_k` is (B,D). So `torch.randn_like(mu_k)` draws B*D normal floats. In our code, we need to draw exactly `B*D` floats. So `eps = torch.randn(B, D, device=mu.device, dtype=mu.dtype)` is fine.
But note: In the reference, `mu_k` is computed before `torch.randn_like(mu_k)`. So the RNG call for `randn` happens after `multinomial`. If we call `multinomial` then `randn`, the order is same. Good.
For the simple case (mu.dim()==3 and num_components>1? Actually that combination doesn't happen in reference because if num_components>1 and mu.dim()==3, it goes to mixture case. So simple case includes:
- mu.dim()==3 and num_components==1 (though in get_init_inputs num_components could be anything? Actually simple case includes `else`, which handles `mu.dim()==3` by slicing to component 0. This happens if num_components==1? Wait, if num_components==1, the init doesn't create logits. But if num_components>1 and mu.dim()==3, it's mixture. So in the else branch, if mu.dim()==3, it must be that either num_components==1 or the user passed a 3D mu even though num_components>1? Actually, the reference says if num_components > 1 and mu.dim() == 3: mixture. Else: if mu.dim()==3: slice. This means if num_components > 1 but mu.dim() != 3, it goes to else. But in the task, inputs are (B, K, D) where K=num_components. So if num_components>1, mu.dim()==3 and mixture is used.
If num_components==1, mu.dim() might be 3 or 2. In get_init_inputs, num_components=4. So the mixture case is the one used in evaluation. But the model should handle all cases.
Given the input generation uses `num_components=4` and `mu` shape `(batch_size, num_components, latent_dim)`, the forward will always take the mixture branch. So we should optimize that path primarily.
However, for correctness, we need to handle the else case as well, but it's less critical for speed.
Given that the mixture case includes:
- `std` computation via `exp(0.5 * log_var)`.
- `weights`, `comp`, `eps`.
- Then the fused kernel.
We need to be careful about numerical equivalence. The reference computes `std = torch.exp(0.5 * log_var)`. In our code, we compute `std = torch.exp(0.5 * log_var)` and pass it to the kernel. That's identical.
But wait: in the reference, `log_q_z` is computed using `std` in the denominator: `(z.unsqueeze(1) - mu) / std`. We compute `(diff / s) * (diff / s)`. That is mathematically equivalent. But we need to ensure no numerical issues. Since std is positive (exp of real), it's fine.
One potential issue: In the reference, `log_probs.sum(-1)` sums over D. The term is `-0.5 * ((z.unsqueeze(1) - mu) / std) ** 2 - log_var * 0.5 - 0.5 * torch.log(torch.tensor(2 * 3.14159265))`. Note that `log_var * 0.5` is computed as `0.5 * log_var`. Since `std = exp(0.5 * log_var)`, then `log(std) = 0.5 * log_var`. The Gaussian log-density is `-0.5 * (diff^2 / std^2) - log(std) - 0.5*log(2*pi)`. The reference uses `log_var * 0.5` which is exactly `log(std)`. So our kernel uses `std` (which came from `exp(0.5*log_var)`) and `log_var`. The term `-0.5 * (diff/s)^2 - 0.5*lv - 0.5*log_2pi` matches.
Potential tiny difference: In reference, `(z.unsqueeze(1) - mu) / std` is a broadcast division. PyTorch may compute this as `z - mu` then divide by `std`, same as ours. The order of operations in the fused kernel is the same.
But we need to be careful: The reference computes `log_probs = log_probs.sum(-1) + weights.log().unsqueeze(0)`. `weights.log()` is natural log. We pass `log_weights` to kernel.
Then `log_q_z = torch.logsumexp(log_probs, dim=1)`. Our kernel computes logsumexp manually. Need to ensure it's stable and accurate. The manual logsumexp with max shift is standard and matches PyTorch's implementation.
One more issue: In the reference, `comp` is sampled from `weights`. Then `mu_k` and `std_k` are selected. The random `eps` is multiplied by `std_k`. Our kernel does the same.
But wait: In the reference, `mu_k + std_k * torch.randn_like(mu_k)`. If we call `eps = torch.randn_like(mu_k)` in Python, `mu_k` is not a tensor variable in our code if we avoid indexing. But we can call `eps = torch.randn(B, D, device=mu.device, dtype=mu.dtype)`. Since `mu_k` is (B,D), `randn_like(mu_k)` is equivalent to `randn(B, D, ...)`. But is the RNG state advanced the same way? Yes, because `randn_like` dispatches to `randn` with the same shape, and the number of elements is the same. However, `randn_like` may use the strides of the input, but for a contiguous (B,D) tensor, it's exactly the same. Since we draw a new tensor of shape (B,D), it's fine. The key is that we draw the same number of random numbers in the same order. So as long as we call `torch.multinomial` first (which consumes some random numbers), then `torch.randn` second, the state will be identical.
Wait: In the reference, after `multinomial`, the next RNG call is `torch.randn_like(mu_k)`. In our code, we will call `torch.multinomial` then `torch.randn(B, D, ...)`. That should produce the exact same random numbers as `randn_like(mu_k)` because both dispatch to the same underlying `normal_` kernel with total elements B*D.
However, we need to verify that `torch.randn_like(mu_k)` and `torch.randn(B, D, device=mu.device, dtype=mu.dtype)` use the same generator state advancement. They should, but there might be differences in how the kernel is launched. Let's check: In PyTorch, `randn_like` calls `at::native::randn_like`, which creates a tensor of the same size and calls `normal_` on it. `torch.randn` creates a tensor of the given size and calls `normal_`. The `normal_` kernel uses the Philox generator with a counter incremented by the number of elements. Since both use the same size, the generator state after the call should be the same. So yes, they are equivalent.
To be extra safe, we can even call `eps = torch.randn_like(mu[:,0])` if mu.dim()==3, or `eps = torch.randn_like(mu)` if mu.dim()==2. This directly mirrors the reference pattern. For the mixture case, the reference draws `torch.randn_like(mu_k)`, where `mu_k` is (B,D). We can compute `mu_k = mu[torch.arange(B), comp]` in Python (as a separate op) and then `eps = torch.randn_like(mu_k)`. But that introduces an indexing op that we wanted to avoid. Actually, we can create a dummy tensor of shape (B,D) via `mu.new_empty(B, D)` and call `torch.randn_like` on it? `torch.randn_like` requires an existing tensor with values to mirror shape and dtype. But we can do:
```
mu_slice = mu[:,0] # shape (B,D), view
eps = torch.randn_like(mu_slice)
```
This will produce a (B,D) random tensor. Does it consume the same RNG state as `torch.randn_like(mu_k)`? The number of elements is the same, so yes. But we need to make sure that `mu_slice` is contiguous? Actually, `mu[:,0]` is a view that may not be contiguous, but `randn_like` on a non-contiguous tensor still creates a contiguous output of the same shape and calls `normal_`. The underlying `normal_` kernel uses the total number of elements. So the generator state advancement is the same.
But if we create `mu_slice = mu[:,0]` just to get a reference for `randn_like`, that is cheap. We don't need to use `mu_slice`'s data. Wait, but `randn_like` uses the shape, not the data. So yes.
Alternatively, we can keep it simple: `eps = torch.empty(B, D, device=mu.device, dtype=mu.dtype).normal_()`. This explicitly calls `normal_` on an empty tensor. That's guaranteed to match the generator behavior? Actually, `torch.randn` and `randn_like` both call `normal_`. So `torch.empty(...).normal_()` is also equivalent.
I think the safest is to call the exact same PyTorch functions as the reference in the same order, then pass their results to the kernel. For performance, the overhead of `multinomial` and `randn` is tiny compared to the fused computation, especially since we avoid multiple intermediate kernels for log_q_z. But we could also fuse the random generation inside the kernel to avoid the overhead. However, for correctness of RNG state and simplicity, I'll call them externally.
Wait, another thought: The correctness check might compare outputs for a fixed seed. But if our kernel runs in a different order (e.g., we call `randn` before `multinomial`), the random draws will differ and outputs will not match. So we must preserve the exact order of RNG calls: first `multinomial` (if mixture), then `randn`. This matches the reference.
Let's plan ModelNew.forward:
```python
class ModelNew(nn.Module):
def __init__(self, latent_dim, num_components=1):
super().__init__()
self.latent_dim = latent_dim
self.num_components = num_components
if num_components > 1:
self.logits = nn.Parameter(torch.zeros(num_components))
def forward(self, mu: torch.Tensor, log_var: torch.Tensor) -> tuple:
std = torch.exp(0.5 * log_var)
B = mu.shape[0]
D = self.latent_dim
K = self.num_components
if K > 1 and mu.dim() == 3:
weights = torch.softmax(self.logits, dim=0)
log_weights = weights.log()
# sample component
comp = torch.multinomial(weights.expand(B, -1), 1).squeeze(-1)
# generate noise: same as reference order
# We need a dummy tensor of shape (B, D) for randn_like
# Actually we can just call torch.randn(B, D, device=mu.device, dtype=mu.dtype)
# But to be safe with RNG state, let's use randn_like on a view.
mu_k_dummy = mu[:, 0] # shape (B, D)
eps = torch.randn_like(mu_k_dummy)
# launch fused kernel
z = torch.empty(B, D, device=mu.device, dtype=mu.dtype)
log_q_z = torch.empty(B, device=mu.device, dtype=mu.dtype)
# call cuda extension
cuda_ext.mixture_forward(
mu, log_var, std, comp, eps, log_weights,
B, K, D,
z, log_q_z,
self._log_2pi # constant
)
else:
if mu.dim() == 3:
mu_s = mu[:, 0]
std_s = std[:, 0]
log_var_s = log_var[:, 0]
else:
mu_s = mu
std_s = std
log_var_s = log_var
eps = torch.randn_like(mu_s)
z = torch.empty_like(mu_s)
log_q_z = torch.empty(B, device=mu.device, dtype=mu.dtype)
cuda_ext.simple_forward(
mu_s, log_var_s, std_s, eps,
B, D,
z, log_q_z,
self._log_2pi
)
return z, log_q_z
```
We need to precompute `_log_2pi` in `__init__` or as a constant.
Actually, in the reference, the constant is computed inside forward: `-0.5 * torch.log(torch.tensor(2 * 3.14159265))`. Wait, the constant used inside `log_q_z` is `-0.5 * log(2*pi)`. But the reference computes `-0.5 * torch.log(torch.tensor(2 * 3.14159265))`. This is a scalar tensor. In our kernel, we compute `-0.5f * log_2pi` where `log_2pi` is passed as `torch.log(torch.tensor(2 * 3.14159265)).item()`. Actually, the reference uses `log_var * 0.5` (wait the reference has `- log_var * 0.5 - 0.5 * torch.log(...)`). So each dimension's term is `-0.5 * ((diff)/std)^2 - 0.5 * log_var - 0.5 * log(2*pi)`. So the constant per dimension is `-0.5 * log(2*pi)`. We can pass `log_2pi` as `log(2*pi)` and multiply by -0.5 in the kernel, or we can precompute `-0.5 * log(2*pi)` and pass it as `neg_half_log_2pi`. This avoids a multiplication per element. But we need to match the reference exactly. Since multiplication by 0.5 is exact in float32 (0.5 is exactly representable), the result of `-0.5 * log(2*pi)` computed in float32 should match exactly whether done on CPU or GPU. But to be safe, we can compute it in `__init__` using the same PyTorch expression and convert to float:
```python
self._const = (-0.5 * torch.log(torch.tensor(2 * 3.14159265))).item()
```
Then pass it to the kernel. The kernel adds `self._const` per dimension. Wait, but there are two terms: `-0.5 * lv` and `self._const`. So per dimension we add `-0.5f * lv + const`. Actually the term is `-0.5f * lv + const` where `const = -0.5 * log(2*pi)`. So we can compute `float term = -0.5f * lv + const;` then add `-0.5f * (diff/s)^2`. This saves one multiplication per element but doesn't affect correctness.
Alternatively, we can keep the kernel formula exactly as the mathematical expression and compute `log_2pi` on the device using the same literal. I think it's fine.
But wait: The reference uses `torch.log(torch.tensor(2 * 3.14159265))`. If we compute the constant in Python as above, we get a Python float (double precision) which is the result of evaluating `torch.log` on a float64 tensor. If we pass it to CUDA `float` argument, it gets cast to float32. This should produce exactly the same float32 value as if we had computed it in CUDA using `logf(2.0f * 3.14159265f)`? Not necessarily, because `2 * 3.14159265` in Python is a double, giving a more accurate `2*pi` representation, then `log` in double is accurate, then cast to float. In CUDA, `2.0f * 3.14159265f` is a float multiplication, which gives a less accurate representation of `2*pi`, and `logf` operates on that less accurate input. So the final float32 constant might differ by a few ULPs. Could that cause tolerance issues? Probably not for 1e-4 tolerance, but let's be safe and compute the constant in Python exactly as the reference does: use PyTorch's double computation and cast to float32.
Actually, in PyTorch reference, `torch.log(torch.tensor(2 * 3.14159265))` returns a float64 scalar. When subtracted from a float32 tensor, it is cast to float32 at the moment of the binary op. The cast from double to float is well-defined. So we can replicate this by computing:
```python
log_2pi = torch.log(torch.tensor(2 * 3.14159265, dtype=torch.float32)).item()
```
Wait, the reference doesn't specify dtype for the tensor, so it's float64. But the cast to float32 happens during the op. The float32 value of the double-precision log(2*pi) is the correctly rounded float32 value of the true log(2*pi). If we compute `logf(2.0f * 3.14159265f)`, we get the float32 log of the float32 number `2.0f * 3.14159265f`. The float32 number `6.2831853` is already a rounded version of `2*pi`. The log of that might differ from the true log(2*pi) by more than 1 ULP. Let's estimate: `2*pi` is about 6.283185307179586. The float32 representation is about 6.2831854820251465 (actually let's compute: 3.14159265 is a float? 3.14159265 in float32 is 3.14159274? Wait, 3.14159265 as a double literal cast to float is ~3.14159274. Then 2.0f * 3.14159274 = 6.28318548. The true 2*pi is 6.283185307... So the float32 product is off by ~1.7e-7. The log derivative is 1/x, so error in log is ~2.7e-8. That's tiny, well within 1e-4 tolerance.
But still, to be absolutely safe, compute the constant in Python exactly as the reference does:
```python
with torch.no_grad():
const_val = (-0.5 * torch.log(torch.tensor(2 * 3.14159265))).item()
```
Then pass as float32.
Actually, maybe simpler: just compute `std` and `log_var` in Python, and compute the entire log_q_z in the kernel using the passed constant. No need to overthink.
Let's now think about the correctness of the fused kernel vs reference.
In reference:
```
log_probs = -0.5 * ((z.unsqueeze(1) - mu) / std) ** 2 - log_var * 0.5 - 0.5 * torch.log(torch.tensor(2 * 3.14159265))
```
Note: We need to replicate `log_var * 0.5` exactly. Since `0.5` is exactly representable, `log_var * 0.5` is the same as `0.5 * log_var`. In our kernel, we use `-0.5f * lv`. That's the same.
Potential issue: The reference computes `(z.unsqueeze(1) - mu) / std` as `z.unsqueeze(1) - mu` first, then divides by `std`. In our kernel, we compute `diff = zv - m` then `diff / s`. The order is identical.
Another issue: The reference computes `log_probs = log_probs.sum(-1) + weights.log().unsqueeze(0)`. `weights.log()` is natural log of weights computed by `softmax`. In our kernel, we pass `log_weights` computed by `weights.log()`. To ensure same `weights`, we compute `softmax` in Python. That's identical.
Potential issue: The reference uses `torch.logsumexp(log_probs, dim=1)`. Our kernel does manual logsumexp. Could there be a difference in precision? `torch.logsumexp` uses the same max-shift trick. So as long as we do max shift, it should be fine.
Another potential issue: The reference `comp` sampling uses `torch.multinomial`. We call `torch.multinomial` too. Good.
Potential issue: For the simple case, if `mu.dim() == 3`, reference slices `mu = mu[:, 0]`, `std = std[:, 0]`, `log_var = log_var[:, 0]`. This yields views (non-contiguous if the original tensor is (B,K,D) with K>1? Actually, slicing the second dimension of a contiguous (B,K,D) tensor yields a tensor with shape (B,D) and stride (K*D, 1), which is not contiguous. Then `z = mu + std * torch.randn_like(mu)` operates on non-contiguous tensors. In our code, if we pass these non-contiguous tensors to the kernel, we need to handle strides. Our kernel uses pointer indexing assuming contiguous (B,D) layout. If we pass a non-contiguous (B,D) view of a (B,K,D) tensor, the indexing `b*D + d` will be wrong because the stride between batch elements is K*D, not D.
Therefore, we have two options:
1. Call `.contiguous()` on the sliced views before passing to the kernel. This adds a copy overhead, but it's small.
2. Pass strides to the kernel and compute indexing accordingly.
For simplicity and correctness, we can call `.contiguous()` on the sliced views. But it does add overhead. Another option: We can pass the original (B,K,D) tensors to the kernel along with a `component_offset=0` flag, and read from index `((b*K)+offset)*D + d`. This avoids copies and indexing issues. This is better.
Similarly, for the mixture case, we pass the original (B,K,D) tensors and the `comp` array, and read `k_sel` per batch.
So we can unify the kernel signatures to accept the original tensors with shape (B,K,D) and a component index per batch.
For the simple case, we can pass a `comp` tensor of all zeros (or just a scalar 0) and use the same mixture kernel? No, because the simple case does not do logsumexp over components. But we can write a single kernel that handles both by passing K and a flag for mixture vs simple. Or we can write a simple kernel that reads the first component.
Let's write two separate kernels but both accept the original (B,K,D) tensors.
For simple case: if mu.dim() == 3, we pass mu (B,K,D), std (B,K,D), log_var (B,K,D) and a scalar comp=0 (or just pass them as pointers and the kernel uses component 0). Since we know we want component 0, we can have a simple kernel that does:
```
int idx = ((b * K) + 0) * D + d = (b * K * D) + d;
```
But wait, the stride between batch elements in the original contiguous tensor is K*D. So `idx = b * K * D + d` is correct only if we want component 0. In general, if we want component `c`, the index is `b * K * D + c * D + d`. So we can have the simple kernel accept a `component` integer (default 0) and compute `idx = ((b * K) + component) * D + d`.
But for the simple case where mu.dim() == 2 (K effectively 1), we can treat K=1 and idx = b*D + d. So we can have the kernel accept a stride between components (`comp_stride = K*D` if K>1, else `D`). Actually, the simplest is to pass a 2D tensor view (via `mu[:,0].contiguous()` for K>1). But that copies. To avoid copies, pass stride explicitly.
Given that the main path is the mixture case, let's not overcomplicate the simple case. We can just pass the sliced `.contiguous()` copy for the simple case. The overhead is negligible and it's rarely used in evaluation. Or we can pass the original 3D tensor and a `component=0` to the kernel.
Let's write the simple kernel to accept `mu`, `log_var`, `std` as 2D tensors. In Python, if `mu.dim() == 3`, we pass `mu[:,0].contiguous(), log_var[:,0].contiguous(), std[:,0].contiguous()`. This is safe. Since the simple case is not the main path, the copy cost is fine.
Alternatively, we can avoid copies by computing in Python:
```
if mu.dim() == 3:
mu_ptr = mu
std_ptr = std
lv_ptr = log_var
K_ptr = mu.shape[1]
comp_val = 0
# call special simple kernel with 3D tensors and comp_val
```
Let's write a simple kernel that works with 3D tensors and a fixed component index:
```
__global__ void simple_forward_kernel_3d(
const float* __restrict__ mu,
const float* __restrict__ log_var,
const float* __restrict__ std,
const float* __restrict__ eps,
int B, int K, int D, int comp,
float* __restrict__ z,
float* __restrict__ log_q_z,
float const_val)
```
But we can just write one kernel that handles both 2D and 3D by passing K (which is 1 for 2D). If K=1, the index is `b*D + d`. If K>1 and comp is fixed, we can pass `comp` and index as `((b*K)+comp)*D + d`.
Actually, we can unify by passing a `batch_stride` which is the stride between batch elements. For a 2D contiguous tensor, stride = D. For a 3D tensor sliced at comp=0, stride = K*D. So passing `batch_stride = mu.stride(0)` would work. In CUDA, using strides is straightforward.
Let's use strides for both kernels. This avoids copies and is general.
For the mixture kernel:
- `mu` is a 3D tensor. We need `mu_stride_batch`, `mu_stride_comp`, `mu_stride_dim`. Since it's contiguous, `mu_stride_batch = K*D`, `mu_stride_comp = D`, `mu_stride_dim = 1`. But we can just compute index as `b*K*D + k*D + d` assuming contiguous. However, we should not assume contiguous. Better to pass strides, or call `.contiguous()` in Python. For speed, calling `.contiguous()` on inputs ensures they are contiguous but adds overhead. Given that `get_inputs()` returns contiguous tensors (from `torch.randn`), we can probably assume contiguous. But just in case, we can call `.contiguous()` on inputs at the start of forward. This adds a small overhead but ensures correctness.
Actually, the simplest and safest approach for the kernel is to call `.contiguous()` on mu, log_var, std, comp, eps, etc., before passing to the kernel. Since they are likely already contiguous, `.contiguous()` is cheap (just returns the same tensor). For non-contiguous cases (like sliced views in simple case), it causes a copy. But that's okay.
Let's design the host code accordingly:
In `forward`:
```
mu = mu.contiguous()
log_var = log_var.contiguous()
std = std.contiguous()
```
Then for mixture case, `comp` is contiguous (from multinomial). `eps` is contiguous (from randn). `weights` and `log_weights` are contiguous.
Now, the kernel can assume contiguous layout and compute indices directly.
But wait, `mu[:,0].contiguous()` in simple case is a copy. If we want to avoid that, we can just treat the simple case by passing the original 3D tensor and a component index of 0 to a kernel that reads from it. Let's do that: write a simple kernel that takes 3D tensors and reads component `c` (a scalar). For 2D tensors, we can treat them as 3D with K=1, but `mu.dim()` is 2. Actually, we can just create a `component` tensor or pass an int.
Let's write a generic kernel:
```
__global__ void simple_forward_kernel(
const float* __restrict__ mu,
const float* __restrict__ log_var,
const float* __restrict__ std,
const float* __restrict__ eps,
int B, int K, int D, int comp,
float* __restrict__ z,
float* __restrict__ log_q_z,
float const_val)
{
int b = blockIdx.x;
if (b >= B) return;
int tid = threadIdx.x;
float local = 0.0f;
for (int d = tid; d < D; d += blockDim.x) {
int idx = ((b * K) + comp) * D + d;
float m = mu[idx];
float s = std[idx];
float lv = log_var[idx];
float e = eps[b * D + d];
float zv = m + s * e;
z[b * D + d] = zv;
float diff = zv - m;
float term = -0.5f * (diff / s) * (diff / s) - 0.5f * lv + const_val;
local += term;
}
extern __shared__ float smem[];
float sum_log_prob = block_reduce_sum(local, smem);
if (tid == 0) {
log_q_z[b] = sum_log_prob;
}
}
```
This works for both 2D and 3D by setting K=1, comp=0 for 2D.
But for 2D, `eps[b * D + d]` assumes eps is (B,D) contiguous. That's fine.
For the mixture case:
```
__global__ void mixture_forward_kernel(
const float* __restrict__ mu,
const float* __restrict__ log_var,
const float* __restrict__ std,
const int64_t* __restrict__ comp,
const float* __restrict__ eps,
const float* __restrict__ log_weights,
int B, int K, int D,
float* __restrict__ z,
float* __restrict__ log_q_z,
float const_val)
```
Same as before, but `const_val = -0.5 * log(2*pi)`.
Now, the host wrapper functions. We need to write C++ wrapper functions that take `at::Tensor` arguments and launch the kernels.
We need to be careful about tensor types: `mu` is float32. `comp` is int64 (LongTensor). We need to get pointers.
Let's write the C++ host wrappers:
```cpp
#include <torch/extension.h>
#include <cuda_runtime.h>
#define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
#define CHECK_FLOAT(x) AT_ASSERTM(x.scalar_type() == at::kFloat, #x " must be float32")
#define CHECK_LONG(x) AT_ASSERTM(x.scalar_type() == at::kLong, #x " must be int64")
__inline__ __device__ float warp_reduce_sum(float val) {
for (int offset = 16; offset > 0; offset /= 2)
val += __shfl_down_sync(0xffffffff, val, offset);
return val;
}
__inline__ __device__ float block_reduce_sum(float val, float* shared) {
int lane = threadIdx.x % 32;
int wid = threadIdx.x / 32;
val = warp_reduce_sum(val);
if (lane == 0) shared[wid] = val;
__syncthreads();
if (wid == 0) {
val = (threadIdx.x < (blockDim.x + 31) / 32) ? shared[lane] : 0.0f;
val = warp_reduce_sum(val);
if (lane == 0) shared[0] = val;
}
__syncthreads();
return shared[0];
}
__global__ void mixture_forward_kernel(
const float* __restrict__ mu,
const float* __restrict__ log_var,
const float* __restrict__ std,
const int64_t* __restrict__ comp,
const float* __restrict__ eps,
const float* __restrict__ log_weights,
int B, int K, int D,
float* __restrict__ z,
float* __restrict__ log_q_z,
float const_val)
{
int b = blockIdx.x;
if (b >= B) return;
int tid = threadIdx.x;
int k_sel = static_cast<int>(comp[b]);
// compute z
for (int d = tid; d < D; d += blockDim.x) {
int idx = ((b * K) + k_sel) * D + d;
float m = mu[idx];
float s = std[idx];
float e = eps[b * D + d];
z[b * D + d] = m + s * e;
}
__syncthreads();
float local[4];
#pragma unroll
for (int k = 0; k < K; ++k) local[k] = 0.0f;
for (int d = tid; d < D; d += blockDim.x) {
float zv = z[b * D + d];
for (int k = 0; k < K; ++k) {
int idx = ((b * K) + k) * D + d;
float m = mu[idx];
float s = std[idx];
float lv = log_var[idx];
float diff = zv - m;
float term = -0.5f * (diff / s) * (diff / s) - 0.5f * lv + const_val;
local[k] += term;
}
}
extern __shared__ float smem[];
float comp_log_prob[4];
for (int k = 0; k < K; ++k) {
float sum_k = block_reduce_sum(local[k], smem);
if (tid == 0) {
comp_log_prob[k] = sum_k + log_weights[k];
}
}
__syncthreads();
if (tid == 0) {
float max_val = comp_log_prob[0];
#pragma unroll
for (int k = 1; k < K; ++k) {
if (comp_log_prob[k] > max_val) max_val = comp_log_prob[k];
}
float lse = 0.0f;
#pragma unroll
for (int k = 0; k < K; ++k) {
lse += expf(comp_log_prob[k] - max_val);
}
log_q_z[b] = max_val + logf(lse);
}
}
__global__ void simple_forward_kernel(
const float* __restrict__ mu,
const float* __restrict__ log_var,
const float* __restrict__ std,
const float* __restrict__ eps,
int B, int K, int D, int comp,
float* __restrict__ z,
float* __restrict__ log_q_z,
float const_val)
{
int b = blockIdx.x;
if (b >= B) return;
int tid = threadIdx.x;
float local = 0.0f;
for (int d = tid; d < D; d += blockDim.x) {
int idx = ((b * K) + comp) * D + d;
float m = mu[idx];
float s = std[idx];
float lv = log_var[idx];
float e = eps[b * D + d];
float zv = m + s * e;
z[b * D + d] = zv;
float diff = zv - m;
float term = -0.5f * (diff / s) * (diff / s) - 0.5f * lv + const_val;
local += term;
}
extern __shared__ float smem[];
float sum_log_prob = block_reduce_sum(local, smem);
if (tid == 0) {
log_q_z[b] = sum_log_prob;
}
}
std::pair<at::Tensor, at::Tensor> mixture_forward_cuda(
at::Tensor mu,
at::Tensor log_var,
at::Tensor std,
at::Tensor comp,
at::Tensor eps,
at::Tensor log_weights,
float const_val)
{
CHECK_CUDA(mu); CHECK_CONTIGUOUS(mu); CHECK_FLOAT(mu);
CHECK_CUDA(log_var); CHECK_CONTIGUOUS(log_var); CHECK_FLOAT(log_var);
CHECK_CUDA(std); CHECK_CONTIGUOUS(std); CHECK_FLOAT(std);
CHECK_CUDA(comp); CHECK_CONTIGUOUS(comp); CHECK_LONG(comp);
CHECK_CUDA(eps); CHECK_CONTIGUOUS(eps); CHECK_FLOAT(eps);
CHECK_CUDA(log_weights); CHECK_CONTIGUOUS(log_weights); CHECK_FLOAT(log_weights);
int B = mu.size(0);
int K = mu.size(1);
int D = mu.size(2);
int threads = 64;
if (D > 64) threads = 128;
else if (D > 128) threads = 256;
else threads = ((D + 31) / 32) * 32;
if (threads < 32) threads = 32;
dim3 block(threads);
dim3 grid(B);
auto z = torch::empty({B, D}, mu.options());
auto log_q_z = torch::empty({B}, mu.options());
int shared_mem_size = (threads / 32 + 1) * sizeof(float);
mixture_forward_kernel<<<grid, block, shared_mem_size>>>(
mu.data_ptr<float>(),
log_var.data_ptr<float>(),
std.data_ptr<float>(),
comp.data_ptr<int64_t>(),
eps.data_ptr<float>(),
log_weights.data_ptr<float>(),
B, K, D,
z.data_ptr<float>(),
log_q_z.data_ptr<float>(),
const_val
);
return {z, log_q_z};
}
std::pair<at::Tensor, at::Tensor> simple_forward_cuda(
at::Tensor mu,
at::Tensor log_var,
at::Tensor std,
at::Tensor eps,
int comp,
float const_val)
{
CHECK_CUDA(mu); CHECK_CONTIGUOUS(mu); CHECK_FLOAT(mu);
CHECK_CUDA(log_var); CHECK_CONTIGUOUS(log_var); CHECK_FLOAT(log_var);
CHECK_CUDA(std); CHECK_CONTIGUOUS(std); CHECK_FLOAT(std);
CHECK_CUDA(eps); CHECK_CONTIGUOUS(eps); CHECK_FLOAT(eps);
int B = mu.size(0);
int K = mu.dim() > 2 ? mu.size(1) : 1;
int D = mu.dim() > 2 ? mu.size(2) : mu.size(1);
int threads = 64;
if (D > 64) threads = 128;
else if (D > 128) threads = 256;
else threads = ((D + 31) / 32) * 32;
if (threads < 32) threads = 32;
dim3 block(threads);
dim3 grid(B);
auto z = torch::empty({B, D}, mu.options());
auto log_q_z = torch::empty({B}, mu.options());
int shared_mem_size = (threads / 32 + 1) * sizeof(float);
simple_forward_kernel<<<grid, block, shared_mem_size>>>(
mu.data_ptr<float>(),
log_var.data_ptr<float>(),
std.data_ptr<float>(),
eps.data_ptr<float>(),
B, K, D, comp,
z.data_ptr<float>(),
log_q_z.data_ptr<float>(),
const_val
);
return {z, log_q_z};
}
```
Note: In `simple_forward_kernel`, we need to handle the case where `mu` is 2D. The kernel expects 3D indexing with K. If `mu` is 2D, we pass K=1, comp=0. The indexing `((b * K) + comp) * D + d` becomes `b*D + d`. Good.
But we need to check that `mu.dim() > 2` in the host code. For 2D case, K=1, D=mu.size(1). The `eps` shape is (B,D). The kernel reads `eps[b*D+d]`. That's fine.
In `simple_forward_cuda`, we need to handle `comp` as an integer argument. C++ wrapper can take `int64_t comp`.
Wait, the indexing in `simple_forward_kernel` for `mu`, `log_var`, `std` uses `int idx = ((b * K) + comp) * D + d;`. If `mu` is 2D and contiguous, `mu.size(1) = D`, `K=1`. The pointer arithmetic is correct.
If `mu` is 3D and contiguous, `K = mu.size(1)`, `comp` is the component index. The pointer arithmetic is `b * K * D + comp * D + d`. That's correct.
One issue: `CHECK_LONG(comp)` is needed only for mixture, not for simple. We'll write separate wrappers.
Another thing: The host wrapper returns `std::pair<at::Tensor, at::Tensor>`. In Python, `load_inline` with `functions=['mixture_forward_cuda', 'simple_forward_cuda']` will generate bindings that return a tuple of two tensors.
Now, let's think about the `block_reduce_sum` function. It uses `shared` memory. In the mixture kernel, we call `block_reduce_sum` K times. The shared memory is declared externally once per kernel. Since `block_reduce_sum` uses `__syncthreads()`, and we call it in a loop, we must ensure that all threads in the block participate in each `__syncthreads()`. They do, because the loop runs for all threads. However, the `extern __shared__ float smem[];` is allocated per kernel launch. In the mixture kernel, we call `block_reduce_sum(local[k], smem)`. The `smem` is reused for each k. Since `block_reduce_sum` ends with `__syncthreads()`, the next iteration will start with all threads synced. Good.
But there is a subtlety: `block_reduce_sum` has a conditional `if (wid == 0) { ... }`. Inside that, it uses `__syncthreads()`. Wait, `block_reduce_sum` does:
```
val = warp_reduce_sum(val);
if (lane == 0) shared[wid] = val;
__syncthreads();
if (wid == 0) {
val = ...;
val = warp_reduce_sum(val);
if (lane == 0) shared[0] = val;
}
__syncthreads();
return shared[0];
```
This is a standard pattern. The first `__syncthreads()` ensures all warps have written. Then only warp 0 reads and reduces. Then second `__syncthreads()` ensures warp 0 is done before any thread reads `shared[0]`. This is okay.
But the `warp_reduce_sum` uses `__shfl_down_sync`. We need to ensure all threads in the warp participate. They do.
One more thing: `mixture_forward_kernel` uses `float local[4];` with a pragma unroll on K. K could be up to maybe 8 or 16? In the problem statement, `num_components=4`. But the model accepts arbitrary `num_components`. We should allocate a local array of a reasonable maximum, or use dynamic shared memory. Using a fixed-size local array of, say, 16 is fine, but if K is larger, it will break. However, the instantiation uses `num_components=4`. The constructor accepts any `num_components`. To be safe, we should not hardcode a small maximum. But if `num_components` is large, registers might be an issue. We can allocate a shared-memory array for component sums instead of registers.
Given the constraints (latent_dim=64, batch_size=256, num_components=4), we can hardcode a maximum like 8 or 16 and check in host code. But it's better to use a dynamic shared memory array for `comp_log_prob` and perhaps compute per-component sums sequentially.
Actually, we can avoid storing per-component local sums in registers by computing them on the fly: For each component k, we compute the block reduction of local[k] and store it in shared memory. To do that, we need to loop over k, but each thread has a `local` value for that k. So we can do:
```
extern __shared__ float smem[];
__shared__ float s_log_probs[8]; // or dynamic
for (int k=0; k<K; ++k) {
float local = 0.0f;
for (int d=tid; d<D; d+=blockDim.x) {
// compute term for component k
local += term_k;
}
float sum_k = block_reduce_sum(local, smem);
if (tid == 0) s_log_probs[k] = sum_k + log_weights[k];
}
__syncthreads();
if (tid == 0) {
// logsumexp over s_log_probs[0..K-1]
}
```
This only needs one float `local` per thread (instead of K floats), and a small shared memory array `s_log_probs[maxK]`. Since K is at most `num_components`, which is small, `s_log_probs[16]` is safe. But if num_components could be larger (e.g., 100), shared memory size is limited but still okay (400 bytes). However, the kernel would be slower because it recomputes `z` and reads `z` from memory inside the inner loop. Wait, `zv` is the same for all k. If we compute `zv` inside the `d` loop, we compute it K times. That is a recompute but avoids storing `z` in global memory. Actually, we can store `z` in global memory (as output) and then read it back. That's memory bandwidth, but we have to write `z` anyway. Alternatively, we can store `z` in registers per thread, but each thread handles multiple d's in a strided fashion, so it doesn't have all z values simultaneously. So reading `z` from global memory is fine.
In the previous design, we wrote `z` to global memory, synced, then read it back. That's an extra global memory read/write. For such a small problem, it might be okay. But we can avoid it by recomputing `z` inside the inner loop: we know `k_sel` and `eps`. So:
```
float zv = mu_k_sel + std_k_sel * eps;
```
Then for each k, compute term using `zv`. This avoids writing and reading `z` entirely for the log_q_z computation, but we still need to write `z` for output. So we can write `z` once and then reuse `zv` computed on the fly.
Let's redesign:
```
// compute z and write to output
for (int d = tid; d < D; d += blockDim.x) {
int idx = ((b * K) + k_sel) * D + d;
float m = mu[idx];
float s = std[idx];
float e = eps[b * D + d];
float zv = m + s * e;
z[b * D + d] = zv;
}
__syncthreads();
// compute log_q_z
extern __shared__ float smem[];
__shared__ float s_log_probs[16]; // enough for typical K
for (int k = 0; k < K; ++k) {
float local = 0.0f;
for (int d = tid; d < D; d += blockDim.x) {
int idx = ((b * K) + k) * D + d;
float m = mu[idx];
float s = std[idx];
float lv = log_var[idx];
float zv = z[b * D + d]; // read from global memory
float diff = zv - m;
local += -0.5f * (diff / s) * (diff / s) - 0.5f * lv + const_val;
}
float sum_k = block_reduce_sum(local, smem);
if (tid == 0) {
s_log_probs[k] = sum_k + log_weights[k];
}
}
__syncthreads();
if (tid == 0) {
// logsumexp
}
```
This writes and reads `z`. The read is coalesced. It's okay.
Alternatively, to avoid the extra global read, we can recompute `zv` from `mu`, `std`, `eps` for `k_sel` inside the inner loop. That would require reading `mu` and `std` again for `k_sel`. We already read them for each k. So we can do:
```
for (int k=0; k<K; ++k) {
float local = 0.0f;
for (int d=tid; d<D; d+=blockDim.x) {
// get zv: if k == k_sel, it's just the computed zv, else we need it anyway
// We can compute zv once per d before the k loop and store it in shared memory or registers.
// Since each thread handles a subset of d, we can store them in local registers if we unroll.
// But D is unknown.
}
}
```
Better: Each thread can load `zv` before the k loop into a local register because each thread iterates over its subset of d. But since the subset is strided, we can just load it inside the inner loop. Actually, the inner loop is:
```
for (int d=tid; d<D; d+=blockDim.x) {
float zv = ???;
for (int k=0; k<K; ++k) {
...
}
}
```
If we compute `zv` before the k loop:
```
for (int d=tid; d<D; d+=blockDim.x) {
int idx_sel = ((b*K)+k_sel)*D + d;
float zv = mu[idx_sel] + std[idx_sel] * eps[b*D+d];
z[b*D+d] = zv;
for (int k=0; k<K; ++k) {
int idx = ((b*K)+k)*D + d;
...
float diff = zv - mu[idx];
local[k] += term;
}
}
```
This computes `zv` once per d, stores it to global memory (for output), and reuses it in the inner k loop. It also reads `mu` and `std` for `k_sel` again in the zv computation, and then for each k in the inner loop. This is more arithmetic but fewer memory reads than reading z back from global memory. Since memory bandwidth is the bottleneck, this is likely better.
Wait, we still need to write `z` to global memory. But we don't need to read it back. So total memory traffic:
- Read mu[k_sel], std[k_sel] for zv: 2 reads per d.
- Write z: 1 write per d.
- Read mu[k], std[k], log_var[k] for each k: 3*K reads per d.
- Read eps: 1 read per d.
Total reads per d: 2 + 3*K + 1 = 3*K + 3. For K=4, that's 15 reads + 1 write.
If we read z back, total reads: 2 (for zv) + 1 (eps) + 3*K (for terms) + 1 (read z) = 3*K + 4. For K=4, 16 reads. So recomputing zv saves one read. Actually, we also wrote z, so the write is still there. The difference is negligible. But we can keep the design where we read z back to simplify the code. For such small sizes, it doesn't matter much.
However, we can do the fused computation without reading z back:
```
// 1. compute z and write to global memory
for (int d=tid; d<D; d+=blockDim.x) {
int idx = ((b*K)+k_sel)*D + d;
float zv = mu[idx] + std[idx] * eps[b*D+d];
z[b*D+d] = zv;
}
__syncthreads();
// 2. compute log_q_z using z from global memory
...
```
This is cleaner. Let's keep it.
Now, we need to define the maximum K for the local array. The problem says num_components=4. But the constructor accepts any. We can use dynamic allocation in shared memory or just a fixed upper bound.
Option: Use `constexpr int MAX_K = 8;` or `16`. If K > MAX_K, we can fallback to PyTorch? But the instructions say no fallback to reference. So we must handle arbitrary K.
If we don't know K at compile time, we can't use `float local[K]` or `s_log_probs[K]` in shared memory without dynamic sizing. But we can allocate dynamic shared memory for both the reduction workspace and the component log-prob array.
Example:
```
extern __shared__ float smem[];
float* reduce_buf = smem; // size: (blockDim.x/32 + 1) maybe
float* comp_buf = smem + ((blockDim.x + 31)/32); // size K
```
But actually, `block_reduce_sum` uses `smem` as a buffer of size `(blockDim.x/32 + 1)`. It returns a float value, but we need to store K floats. We can declare `__shared__ float comp_buf[64];` with a fixed max (say 64). For typical use, K is small. To be safe for larger K, we can allocate dynamic shared memory for `comp_buf`.
Let's allocate total dynamic shared memory as:
```
int shared_mem_size = ((threads + 31) / 32) * sizeof(float) + K * sizeof(float);
```
In the kernel, we split it:
```
extern __shared__ float smem[];
float* reduce_smem = smem;
float* comp_smem = smem + ((blockDim.x + 31) / 32);
```
Then in the loop:
```
for (int k = 0; k < K; ++k) {
float local = 0.0f;
for (int d = tid; d < D; d += blockDim.x) {
...
local += term;
}
float sum_k = block_reduce_sum(local, reduce_smem);
if (tid == 0) comp_smem[k] = sum_k + log_weights[k];
}
__syncthreads();
if (tid == 0) {
float max_val = comp_smem[0];
for (int k = 1; k < K; ++k) if (comp_smem[k] > max_val) max_val = comp_smem[k];
float lse = 0.0f;
for (int k = 0; k < K; ++k) lse += expf(comp_smem[k] - max_val);
log_q_z[b] = max_val + logf(lse);
}
```
This avoids fixed-size arrays.
Now, rewrite `block_reduce_sum` to take a pointer to the shared memory buffer.
```
__inline__ __device__ float block_reduce_sum(float val, float* shared) {
int lane = threadIdx.x % 32;
int wid = threadIdx.x / 32;
for (int offset = 16; offset > 0; offset /= 2)
val += __shfl_down_sync(0xffffffff, val, offset);
if (lane == 0) shared[wid] = val;
__syncthreads();
if (wid == 0) {
val = (threadIdx.x < (blockDim.x + 31) / 32) ? shared[lane] : 0.0f;
for (int offset = 16; offset > 0; offset /= 2)
val += __shfl_down_sync(0xffffffff, val, offset);
if (lane == 0) shared[0] = val;
}
__syncthreads();
return shared[0];
}
```
This is fine.
Now, the simple kernel can also use dynamic shared memory.
Let's write the final C++ code.
One more thing: The constant `const_val`. In the reference, the term is `-0.5 * torch.log(torch.tensor(2 * 3.14159265))`. We can compute this in C++ as `const_val` argument. In `__init__`, we compute:
```python
self._const_val = (-0.5 * torch.log(torch.tensor(2 * 3.14159265))).item()
```
But wait: This is a Python float. When passed to C++, it becomes a C++ float. Good.
Alternatively, we can compute it inside the C++ wrapper:
```cpp
float const_val = -0.5f * std::log(2.0 * 3.14159265);
```
But again, double vs float precision differences. Safer to compute in Python and pass as float.
Actually, the Python computation `(-0.5 * torch.log(torch.tensor(2 * 3.14159265))).item()` returns a Python float (double). When passed to C++ as `float`, it's cast to float32. This cast matches the reference's cast from float64 to float32. So it's good.
Now, the host code: We need to write the C++ wrappers. Let's craft them carefully.
```cpp
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cmath>
#define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
#define CHECK_FLOAT(x) AT_ASSERTM(x.scalar_type() == at::kFloat, #x " must be float32")
#define CHECK_LONG(x) AT_ASSERTM(x.scalar_type() == at::kLong, #x " must be int64")
__inline__ __device__ float warp_reduce_sum(float val) {
for (int offset = 16; offset > 0; offset /= 2)
val += __shfl_down_sync(0xffffffff, val, offset);
return val;
}
__inline__ __device__ float block_reduce_sum(float val, float* shared) {
int lane = threadIdx.x % 32;
int wid = threadIdx.x / 32;
val = warp_reduce_sum(val);
if (lane == 0) shared[wid] = val;
__syncthreads();
if (wid == 0) {
val = (threadIdx.x < (blockDim.x + 31) / 32) ? shared[lane] : 0.0f;
val = warp_reduce_sum(val);
if (lane == 0) shared[0] = val;
}
__syncthreads();
return shared[0];
}
__global__ void mixture_forward_kernel(
const float* __restrict__ mu,
const float* __restrict__ log_var,
const float* __restrict__ std,
const int64_t* __restrict__ comp,
const float* __restrict__ eps,
const float* __restrict__ log_weights,
int B, int K, int D,
float* __restrict__ z,
float* __restrict__ log_q_z,
float const_val)
{
int b = blockIdx.x;
if (b >= B) return;
int tid = threadIdx.x;
int k_sel = static_cast<int>(comp[b]);
// Compute z
for (int d = tid; d < D; d += blockDim.x) {
int idx = ((b * K) + k_sel) * D + d;
float m = mu[idx];
float s = std[idx];
float e = eps[b * D + d];
z[b * D + d] = m + s * e;
}
__syncthreads();
extern __shared__ float smem[];
float* reduce_smem = smem;
float* comp_smem = smem + ((blockDim.x + 31) / 32);
for (int k = 0; k < K; ++k) {
float local = 0.0f;
for (int d = tid; d < D; d += blockDim.x) {
int idx = ((b * K) + k) * D + d;
float m = mu[idx];
float s = std[idx];
float lv = log_var[idx];
float zv = z[b * D + d];
float diff = zv - m;
float term = -0.5f * (diff / s) * (diff / s) - 0.5f * lv + const_val;
local += term;
}
float sum_k = block_reduce_sum(local, reduce_smem);
if (tid == 0) {
comp_smem[k] = sum_k + log_weights[k];
}
}
__syncthreads();
if (tid == 0) {
float max_val = comp_smem[0];
for (int k = 1; k < K; ++k) {
if (comp_smem[k] > max_val) max_val = comp_smem[k];
}
float lse = 0.0f;
for (int k = 0; k < K; ++k) {
lse += expf(comp_smem[k] - max_val);
}
log_q_z[b] = max_val + logf(lse);
}
}
__global__ void simple_forward_kernel(
const float* __restrict__ mu,
const float* __restrict__ log_var,
const float* __restrict__ std,
const float* __restrict__ eps,
int B, int K, int D, int comp,
float* __restrict__ z,
float* __restrict__ log_q_z,
float const_val)
{
int b = blockIdx.x;
if (b >= B) return;
int tid = threadIdx.x;
float local = 0.0f;
for (int d = tid; d < D; d += blockDim.x) {
int idx = ((b * K) + comp) * D + d;
float m = mu[idx];
float s = std[idx];
float lv = log_var[idx];
float e = eps[b * D + d];
float zv = m + s * e;
z[b * D + d] = zv;
float diff = zv - m;
float term = -0.5f * (diff / s) * (diff / s) - 0.5f * lv + const_val;
local += term;
}
extern __shared__ float smem[];
float sum_log_prob = block_reduce_sum(local, smem);
if (tid == 0) {
log_q_z[b] = sum_log_prob;
}
}
std::pair<at::Tensor, at::Tensor> mixture_forward_cuda(
at::Tensor mu,
at::Tensor log_var,
at::Tensor std,
at::Tensor comp,
at::Tensor eps,
at::Tensor log_weights,
float const_val)
{
CHECK_CUDA(mu); CHECK_CONTIGUOUS(mu); CHECK_FLOAT(mu);
CHECK_CUDA(log_var); CHECK_CONTIGUOUS(log_var); CHECK_FLOAT(log_var);
CHECK_CUDA(std); CHECK_CONTIGUOUS(std); CHECK_FLOAT(std);
CHECK_CUDA(comp); CHECK_CONTIGUOUS(comp); CHECK_LONG(comp);
CHECK_CUDA(eps); CHECK_CONTIGUOUS(eps); CHECK_FLOAT(eps);
CHECK_CUDA(log_weights); CHECK_CONTIGUOUS(log_weights); CHECK_FLOAT(log_weights);
int B = mu.size(0);
int K = mu.size(1);
int D = mu.size(2);
int threads = ((D + 31) / 32) * 32;
if (threads < 32) threads = 32;
if (threads > 1024) threads = 1024;
dim3 block(threads);
dim3 grid(B);
auto z = torch::empty({B, D}, mu.options());
auto log_q_z = torch::empty({B}, mu.options());
int shared_mem_size = ((threads + 31) / 32) * sizeof(float) + K * sizeof(float);
mixture_forward_kernel<<<grid, block, shared_mem_size>>>(
mu.data_ptr<float>(),
log_var.data_ptr<float>(),
std.data_ptr<float>(),
comp.data_ptr<int64_t>(),
eps.data_ptr<float>(),
log_weights.data_ptr<float>(),
B, K, D,
z.data_ptr<float>(),
log_q_z.data_ptr<float>(),
const_val
);
return {z, log_q_z};
}
std::pair<at::Tensor, at::Tensor> simple_forward_cuda(
at::Tensor mu,
at::Tensor log_var,
at::Tensor std,
at::Tensor eps,
int64_t comp,
float const_val)
{
CHECK_CUDA(mu); CHECK_CONTIGUOUS(mu); CHECK_FLOAT(mu);
CHECK_CUDA(log_var); CHECK_CONTIGUOUS(log_var); CHECK_FLOAT(log_var);
CHECK_CUDA(std); CHECK_CONTIGUOUS(std); CHECK_FLOAT(std);
CHECK_CUDA(eps); CHECK_CONTIGUOUS(eps); CHECK_FLOAT(eps);
int B = mu.size(0);
int K = mu.dim() > 2 ? mu.size(1) : 1;
int D = mu.dim() > 2 ? mu.size(2) : mu.size(1);
int threads = ((D + 31) / 32) * 32;
if (threads < 32) threads = 32;
if (threads > 1024) threads = 1024;
dim3 block(threads);
dim3 grid(B);
auto z = torch::empty({B, D}, mu.options());
auto log_q_z = torch::empty({B}, mu.options());
int shared_mem_size = ((threads + 31) / 32) * sizeof(float);
simple_forward_kernel<<<grid, block, shared_mem_size>>>(
mu.data_ptr<float>(),
log_var.data_ptr<float>(),
std.data_ptr<float>(),
eps.data_ptr<float>(),
B, K, D, static_cast<int>(comp),
z.data_ptr<float>(),
log_q_z.data_ptr<float>(),
const_val
);
return {z, log_q_z};
}
```
A few issues:
1. `block_reduce_sum` in the simple kernel uses `extern __shared__ float smem[];` just like in the mixture kernel. That's fine.
2. In `mixture_forward_kernel`, `comp_smem` is written by thread 0 only after `block_reduce_sum`. There is a `__syncthreads()` at the end of `block_reduce_sum`. Then after the loop, we have `__syncthreads()` before reading `comp_smem`. But we also need a `__syncthreads()` after writing `comp_smem` in the loop? Actually, `block_reduce_sum` ends with `__syncthreads()`. Then thread 0 writes `comp_smem[k]`. Then the next iteration of the loop starts. At the start of the next iteration, we compute `local`. But we need a `__syncthreads()` after thread 0 writes `comp_smem[k]` before any other thread reads `comp_smem`? That's only needed after the loop. However, `block_reduce_sum` inside the next iteration will call `__syncthreads()`. But before that, threads might read `comp_smem`? No, they don't. So it's okay.
3. But the final `__syncthreads()` after the loop is necessary because thread 0 reads `comp_smem` and other threads might still be writing? Actually, only thread 0 writes. The other threads are done. But to be safe, we have `__syncthreads()` before the `if (tid == 0)` block. Wait, in the code above, after the loop there is `__syncthreads();` then `if (tid == 0)`. That is good.
4. However, inside the `for (int k = 0; k < K; ++k)` loop, after `if (tid == 0) comp_smem[k] = ...;`, the next iteration of the loop starts immediately. But `block_reduce_sum` is called again. The first `__syncthreads()` inside `block_reduce_sum` will sync all threads. But there is a potential race: thread 0 writes `comp_smem[k]`, then proceeds to the next iteration and may start executing `for (int d = tid; ...)` while other threads are still behind? No, because all threads are at the same loop iteration. They all call `block_reduce_sum` which syncs them. After `block_reduce_sum`, thread 0 writes to `comp_smem[k]`, but other threads exit the `if` and continue. The loop increments `k`. There is no barrier between iterations, but that's okay because the next iteration's computation (`local = 0.0f; for d ...)`) does not depend on `comp_smem` from previous iterations. However, `block_reduce_sum` uses `reduce_smem`, not `comp_smem`. So no race.
5. The `reduce_smem` pointer is `smem`. `block_reduce_sum` writes to `shared[wid]`. The maximum `wid` is `(threads+31)/32 - 1`. The size of `reduce_smem` is `(threads+31)/32` floats. That's correct.
6. `comp_smem` starts at `smem + ((blockDim.x + 31) / 32)`. This is also correct.
Now, about the Python side. We need to define `ModelNew`. We'll use `torch.utils.cpp_extension.load_inline`.
Important: Do not write own `PYBIND11_MODULE`. `load_inline` auto-generates bindings from the `functions` list.
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
CUDA_SOURCE = r'''
#include <torch/extension.h>
#include <cuda_runtime.h>
#define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
#define CHECK_FLOAT(x) AT_ASSERTM(x.scalar_type() == at::kFloat, #x " must be float32")
#define CHECK_LONG(x) AT_ASSERTM(x.scalar_type() == at::kLong, #x " must be int64")
__inline__ __device__ float warp_reduce_sum(float val) {
for (int offset = 16; offset > 0; offset /= 2)
val += __shfl_down_sync(0xffffffff, val, offset);
return val;
}
__inline__ __device__ float block_reduce_sum(float val, float* shared) {
int lane = threadIdx.x % 32;
int wid = threadIdx.x / 32;
val = warp_reduce_sum(val);
if (lane == 0) shared[wid] = val;
__syncthreads();
if (wid == 0) {
val = (threadIdx.x < (blockDim.x + 31) / 32) ? shared[lane] : 0.0f;
val = warp_reduce_sum(val);
if (lane == 0) shared[0] = val;
}
__syncthreads();
return shared[0];
}
__global__ void mixture_forward_kernel(
const float* __restrict__ mu,
const float* __restrict__ log_var,
const float* __restrict__ std,
const int64_t* __restrict__ comp,
const float* __restrict__ eps,
const float* __restrict__ log_weights,
int B, int K, int D,
float* __restrict__ z,
float* __restrict__ log_q_z,
float const_val)
{
int b = blockIdx.x;
if (b >= B) return;
int tid = threadIdx.x;
int k_sel = static_cast<int>(comp[b]);
for (int d = tid; d < D; d += blockDim.x) {
int idx = ((b * K) + k_sel) * D + d;
float m = mu[idx];
float s = std[idx];
float e = eps[b * D + d];
z[b * D + d] = m + s * e;
}
__syncthreads();
extern __shared__ float smem[];
float* reduce_smem = smem;
float* comp_smem = smem + ((blockDim.x + 31) / 32);
for (int k = 0; k < K; ++k) {
float local = 0.0f;
for (int d = tid; d < D; d += blockDim.x) {
int idx = ((b * K) + k) * D + d;
float m = mu[idx];
float s = std[idx];
float lv = log_var[idx];
float zv = z[b * D + d];
float diff = zv - m;
float term = -0.5f * (diff / s) * (diff / s) - 0.5f * lv + const_val;
local += term;
}
float sum_k = block_reduce_sum(local, reduce_smem);
if (tid == 0) {
comp_smem[k] = sum_k + log_weights[k];
}
}
__syncthreads();
if (tid == 0) {
float max_val = comp_smem[0];
for (int k = 1; k < K; ++k) {
if (comp_smem[k] > max_val) max_val = comp_smem[k];
}
float lse = 0.0f;
for (int k = 0; k < K; ++k) {
lse += expf(comp_smem[k] - max_val);
}
log_q_z[b] = max_val + logf(lse);
}
}
__global__ void simple_forward_kernel(
const float* __restrict__ mu,
const float* __restrict__ log_var,
const float* __restrict__ std,
const float* __restrict__ eps,
int B, int K, int D, int comp,
float* __restrict__ z,
float* __restrict__ log_q_z,
float const_val)
{
int b = blockIdx.x;
if (b >= B) return;
int tid = threadIdx.x;
float local = 0.0f;
for (int d = tid; d < D; d += blockDim.x) {
int idx = ((b * K) + comp) * D + d;
float m = mu[idx];
float s = std[idx];
float lv = log_var[idx];
float e = eps[b * D + d];
float zv = m + s * e;
z[b * D + d] = zv;
float diff = zv - m;
float term = -0.5f * (diff / s) * (diff / s) - 0.5f * lv + const_val;
local += term;
}
extern __shared__ float smem[];
float sum_log_prob = block_reduce_sum(local, smem);
if (tid == 0) {
log_q_z[b] = sum_log_prob;
}
}
std::pair<at::Tensor, at::Tensor> mixture_forward_cuda(
at::Tensor mu,
at::Tensor log_var,
at::Tensor std,
at::Tensor comp,
at::Tensor eps,
at::Tensor log_weights,
float const_val)
{
CHECK_CUDA(mu); CHECK_CONTIGUOUS(mu); CHECK_FLOAT(mu);
CHECK_CUDA(log_var); CHECK_CONTIGUOUS(log_var); CHECK_FLOAT(log_var);
CHECK_CUDA(std); CHECK_CONTIGUOUS(std); CHECK_FLOAT(std);
CHECK_CUDA(comp); CHECK_CONTIGUOUS(comp); CHECK_LONG(comp);
CHECK_CUDA(eps); CHECK_CONTIGUOUS(eps); CHECK_FLOAT(eps);
CHECK_CUDA(log_weights); CHECK_CONTIGUOUS(log_weights); CHECK_FLOAT(log_weights);
int B = mu.size(0);
int K = mu.size(1);
int D = mu.size(2);
int threads = ((D + 31) / 32) * 32;
if (threads < 32) threads = 32;
if (threads > 1024) threads = 1024;
dim3 block(threads);
dim3 grid(B);
auto z = torch::empty({B, D}, mu.options());
auto log_q_z = torch::empty({B}, mu.options());
int shared_mem_size = ((threads + 31) / 32) * sizeof(float) + K * sizeof(float);
mixture_forward_kernel<<<grid, block, shared_mem_size>>>(
mu.data_ptr<float>(),
log_var.data_ptr<float>(),
std.data_ptr<float>(),
comp.data_ptr<int64_t>(),
eps.data_ptr<float>(),
log_weights.data_ptr<float>(),
B, K, D,
z.data_ptr<float>(),
log_q_z.data_ptr<float>(),
const_val
);
return {z, log_q_z};
}
std::pair<at::Tensor, at::Tensor> simple_forward_cuda(
at::Tensor mu,
at::Tensor log_var,
at::Tensor std,
at::Tensor eps,
int64_t comp,
float const_val)
{
CHECK_CUDA(mu); CHECK_CONTIGUOUS(mu); CHECK_FLOAT(mu);
CHECK_CUDA(log_var); CHECK_CONTIGUOUS(log_var); CHECK_FLOAT(log_var);
CHECK_CUDA(std); CHECK_CONTIGUOUS(std); CHECK_FLOAT(std);
CHECK_CUDA(eps); CHECK_CONTIGUOUS(eps); CHECK_FLOAT(eps);
int B = mu.size(0);
int K = mu.dim() > 2 ? mu.size(1) : 1;
int D = mu.dim() > 2 ? mu.size(2) : mu.size(1);
int threads = ((D + 31) / 32) * 32;
if (threads < 32) threads = 32;
if (threads > 1024) threads = 1024;
dim3 block(threads);
dim3 grid(B);
auto z = torch::empty({B, D}, mu.options());
auto log_q_z = torch::empty({B}, mu.options());
int shared_mem_size = ((threads + 31) / 32) * sizeof(float);
simple_forward_kernel<<<grid, block, shared_mem_size>>>(
mu.data_ptr<float>(),
log_var.data_ptr<float>(),
std.data_ptr<float>(),
eps.data_ptr<float>(),
B, K, D, static_cast<int>(comp),
z.data_ptr<float>(),
log_q_z.data_ptr<float>(),
const_val
);
return {z, log_q_z};
}
'''
_cuda_ext = load_inline(
name='reparam_cuda',
cpp_sources='',
cuda_sources=CUDA_SOURCE,
functions=['mixture_forward_cuda', 'simple_forward_cuda'],
extra_cuda_cflags=['-O3', '--use_fast_math'],
verbose=False
)
class ModelNew(nn.Module):
def __init__(self, latent_dim, num_components=1):
super().__init__()
self.latent_dim = latent_dim
self.num_components = num_components
if num_components > 1:
self.logits = nn.Parameter(torch.zeros(num_components))
self._const_val = (-0.5 * torch.log(torch.tensor(2 * 3.14159265))).item()
def forward(self, mu, log_var):
std = torch.exp(0.5 * log_var)
B = mu.shape[0]
D = self.latent_dim
K = self.num_components
if K > 1 and mu.dim() == 3:
weights = torch.softmax(self.logits, dim=0)
log_weights = weights.log()
comp = torch.multinomial(weights.expand(B, -1), 1).squeeze(-1)
eps = torch.randn(B, D, device=mu.device, dtype=mu.dtype)
mu = mu.contiguous()
log_var = log_var.contiguous()
std = std.contiguous()
log_weights = log_weights.contiguous()
z, log_q_z = _cuda_ext.mixture_forward_cuda(
mu, log_var, std, comp, eps, log_weights, self._const_val
)
else:
comp_idx = 0
if mu.dim() == 3:
comp_idx = 0
eps = torch.randn_like(mu[:, 0] if mu.dim() == 3 else mu)
mu = mu.contiguous()
log_var = log_var.contiguous()
std = std.contiguous()
z, log_q_z = _cuda_ext.simple_forward_cuda(
mu, log_var, std, eps, comp_idx, self._const_val
)
return z, log_q_z
```
Wait, there is an issue in the simple case: `eps = torch.randn_like(mu[:, 0] if mu.dim() == 3 else mu)`. For `mu.dim() == 2`, `mu[:,0]` would be shape (B,), which is wrong. Actually, the ternary handles that: if `mu.dim() == 3`, use `mu[:,0]`, else use `mu`. But we need `eps` of shape (B, D). `mu[:,0]` is shape (B,D) if mu is (B,K,D). So `torch.randn_like(mu[:,0])` returns (B,D). For `mu.dim() == 2`, `mu` is (B,D), so `torch.randn_like(mu)` returns (B,D). Good.
But we need to ensure that `torch.randn_like` is called after any earlier RNG calls. In the mixture branch, we call `multinomial` then `randn`. In the simple branch, there is no `multinomial`, just `randn`. The reference simple branch also has just `randn`. So order matches.
However, what about `std = torch.exp(0.5 * log_var)`? The reference computes `std` before `multinomial` or `randn`. So that's fine.
One subtlety: In the simple case of the reference where `mu.dim() == 3`, it does:
```
mu = mu[:, 0]
std = std[:, 0]
log_var = log_var[:, 0]
z = mu + std * torch.randn_like(mu)
```
Note: `mu` is reassigned to `mu[:,0]`. This means `torch.randn_like(mu)` is called on the already-sliced `mu`. In our code, we do `eps = torch.randn_like(mu[:, 0] if mu.dim() == 3 else mu)`. This matches.
Another subtlety: In the simple case, `log_q_z` computation uses the sliced `mu`, `std`, `log_var`, and `z`. Our kernel passes the original 3D tensors and `comp_idx=0`. This is mathematically identical. But we need to verify pointer indexing.
If `mu` is contiguous (B, K, D), then for `b` batch and `d` dim, the address of `mu[b, 0, d]` is `(b * K * D + 0 * D + d)`. The kernel does `((b * K) + comp) * D + d`. With `comp=0`, that's `b*K*D + d`. That's correct. So no copy needed for simple case. Great.
But we call `mu = mu.contiguous()`. If it's already contiguous, no copy. Good.
Wait: the reference simple case slices before calling `randn_like`. The sliced tensors are views but the `randn_like` consumes RNG state based on shape. Our `torch.randn_like(mu[:,0])` also consumes the same amount. So it's fine.
One more thing: In the simple case, `std` is computed from the original `log_var` (which is 3D). Then `std = std[:,0]`. Our kernel receives the original 3D `std` and reads component 0. That's correct.
Now, check the correctness of the log_q_z computation in the simple case.
Reference:
```
log_q_z = (-0.5 * ((z - mu) / std) ** 2 - 0.5 * log_var - 0.5 * torch.log(torch.tensor(2 * 3.14159265))).sum(-1)
```
Our kernel:
```
float diff = zv - m;
float term = -0.5f * (diff / s) * (diff / s) - 0.5f * lv + const_val;
```
where `const_val = -0.5 * log(2 * pi)`.
So `term = -0.5 * (diff/s)^2 - 0.5*lv - 0.5*log(2*pi)`. Then sum over d. This matches.
Now, mixture case correctness.
Reference:
```
log_probs = -0.5 * ((z.unsqueeze(1) - mu) / std) ** 2 - log_var * 0.5 - 0.5 * torch.log(torch.tensor(2 * 3.14159265))
log_probs = log_probs.sum(-1) + weights.log().unsqueeze(0)
log_q_z = torch.logsumexp(log_probs, dim=1)
```
Our kernel:
- For each k, compute sum over d of `-0.5 * (diff/std)^2 - 0.5*log_var + const_val`. This is `log_probs.sum(-1)`.
- Add `log_weights[k]` (which is `weights.log().unsqueeze(0)` per k).
- Store in `comp_smem[k]`.
- Then compute `logsumexp` over k.
This matches.
Potential floating point difference: The reference computes `weights.log().unsqueeze(0)`, which is a 2D tensor of shape (1, K). Then it is broadcast-added to `log_probs` of shape (B, K). In our kernel, we add `log_weights[k]` after the block reduction. Since `log_weights[k]` is a scalar, the addition order is per-component sum then add scalar. In PyTorch, the addition is done elementwise after the sum, which is the same mathematically.
Potential difference: The reference computes `-0.5 * ((z.unsqueeze(1) - mu) / std) ** 2`. The exponentiation is `**2`, same as multiplying by self. PyTorch may use `square` or `pow(2)`. Our kernel does `(diff / s) * (diff / s)`. This is equivalent.
Another potential issue: `std` is computed as `exp(0.5 * log_var)`. If `log_var` is large positive, `std` can be large. Division by `std` is fine. If `std` is zero (log_var -> -inf), division by zero occurs. But `log_var` comes from `torch.randn(...) * 0.5`, so it's finite. Not a concern.
Now, let's consider performance relative to reference.
Reference operations for mixture case:
- `exp`: launch kernel for 0.5*log_var -> B*K*D = 65536 elements.
- `softmax`: launch kernel for logits -> K=4 elements, trivial.
- `multinomial`: launch kernel for sampling -> B*K = 1024 elements.
- indexing `mu_k`, `std_k`: launch indexing kernels -> B*D = 16384 elements each.
- `randn_like`: launch normal kernel -> B*D = 16384 elements.
- `z = mu_k + std_k * eps`: launch kernel -> B*D.
- `z.unsqueeze(1) - mu`: broadcast subtraction -> B*K*D.
- division by `std`: broadcast division -> B*K*D.
- `**2`: elementwise pow -> B*K*D.
- multiply by -0.5: elementwise -> B*K*D.
- subtract log_var*0.5: broadcast -> B*K*D.
- subtract constant: broadcast -> B*K*D.
- `sum(-1)`: reduction over D -> B*K.
- add log weights: broadcast -> B*K.
- `logsumexp`: reduction over K -> B.
Total elementwise launches: many. Our fused kernel does it in a single launch. This should be much faster due to fewer memory round trips and kernel launch overhead.
Shared memory usage: `((threads+31)/32 + K) * sizeof(float)`. For threads=64, K=4, this is `(2 + 4) * 4 = 24` bytes per block. Negligible.
Occupancy: blocks = B = 256. threads per block = 64. Total threads = 16384. On a modern GPU, occupancy is high.
One possible issue: The mixture kernel writes `z` to global memory, then reads it back. This is necessary for output anyway. Could we avoid reading `z` back by keeping `z` in registers? Since each thread handles a strided subset of D, it cannot keep all D in registers. But we can recompute `z` on the fly for `k_sel` inside the inner loop. That saves a global memory read. Let's modify the kernel to recompute `z` instead of reading it.
Actually, reading `z` from global memory is coalesced and L2 cache should capture it (since we just wrote it). But we can avoid it to be slightly faster.
Modified mixture kernel:
```
__global__ void mixture_forward_kernel(...)
{
int b = blockIdx.x;
if (b >= B) return;
int tid = threadIdx.x;
int k_sel = static_cast<int>(comp[b]);
// write z to global memory
for (int d = tid; d < D; d += blockDim.x) {
int idx = ((b * K) + k_sel) * D + d;
float m = mu[idx];
float s = std[idx];
float e = eps[b * D + d];
z[b * D + d] = m + s * e;
}
__syncthreads();
extern __shared__ float smem[];
float* reduce_smem = smem;
float* comp_smem = smem + ((blockDim.x + 31) / 32);
for (int k = 0; k < K; ++k) {
float local = 0.0f;
for (int d = tid; d < D; d += blockDim.x) {
int idx = ((b * K) + k) * D + d;
float m = mu[idx];
float s = std[idx];
float lv = log_var[idx];
float zv;
if (k == k_sel) {
// recompute z to avoid reading from global memory
int idx_sel = ((b * K) + k_sel) * D + d;
zv = mu[idx_sel] + std[idx_sel] * eps[b * D + d];
} else {
// read z from global memory (only actually needed if k != k_sel)
zv = z[b * D + d];
}
float diff = zv - m;
float term = -0.5f * (diff / s) * (diff / s) - 0.5f * lv + const_val;
local += term;
}
float sum_k = block_reduce_sum(local, reduce_smem);
if (tid == 0) {
comp_smem[k] = sum_k + log_weights[k];
}
}
__syncthreads();
if (tid == 0) {
...
}
}
```
But this adds a branch inside the inner loop. The branch is uniform across all threads (`k == k_sel` is the same for all threads in a block since `b` is fixed). So it's a uniform branch, which is fine. However, it adds extra reads of `mu[idx_sel]` and `std[idx_sel]` for `k_sel`. Actually, for `k == k_sel`, we recompute `zv` using `mu` and `std`. For `k != k_sel`, we read `z` from global memory. This means for the `k_sel` iteration, we avoid the global read of `z`, but for other iterations we still read `z`. Alternatively, we can read `z` from global memory for all iterations, which is simpler.
Given the small size, reading `z` is fine. Let's keep it simple and read `z` to avoid extra instructions.
One more thing: The mixture kernel uses `eps[b * D + d]`. In the reference, `eps` is computed as `torch.randn_like(mu_k)`, which is (B,D). We generate `eps` as `torch.randn(B, D, device=mu.device, dtype=mu.dtype)`. This is contiguous. Good.
Now, let's think about the `block_reduce_sum` function again. It uses `shared[wid]`. If `blockDim.x` is not a multiple of 32, some warps may be incomplete. But `warp_reduce_sum` uses `__shfl_down_sync` with `0xffffffff`. In an incomplete warp (e.g., only 16 active threads), the `__shfl_down_sync` with `0xffffffff` will still work but the extra threads are inactive? Actually, `__shfl_down_sync` requires the mask to include all participating threads. If we use `0xffffffff`, it includes all 32 threads in the warp, but some may be inactive. CUDA semantics: `__shfl_down_sync(mask, val, offset)` is only safe if all threads in the mask are convergent and active. If some threads are inactive because `tid >= blockDim.x`, they won't execute the instruction. However, in our kernel, we don't guard against `tid >= blockDim.x` when calling `block_reduce_sum`. But the loops `for (int d = tid; d < D; d += blockDim.x)` may have some threads with no iterations. For those threads, `local` will be 0.0f. They still participate in `block_reduce_sum`, including `__shfl_down_sync(0xffffffff, ...)`. Since they are active (they are in the block), the mask is fine. `blockDim.x` is a multiple of 32? We set it to `((D + 31) / 32) * 32`, so yes, it's a multiple of 32. Good.
Now, compile it and test.
But first, there is a potential issue with `load_inline`: The `cpp_sources` argument must contain the C++ wrapper code, and `cuda_sources` must contain the CUDA code. Usually, you put the host code in `cpp_sources` and the device code in `cuda_sources`. Wait, the typical usage:
```python
load_inline(
name='...',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=...,
)
```
Where `cpp_source` contains the host-side wrappers (including pybind if you were to do it manually, but `load_inline` auto-generates it). Actually, `load_inline` splits: `cpp_sources` are compiled with the C++ compiler, and `cuda_sources` are compiled with nvcc. You can put everything in `cuda_sources` and leave `cpp_sources` empty? I think so, because nvcc can compile host code too. But usually, the C++ wrappers that call CUDA kernels need to be visible to the C++ compiler. If we put them in `cuda_sources`, nvcc will compile them into the CUDA object, and `load_inline` will link them. Since `load_inline` supports having wrappers in `cuda_sources`, that is fine.
But to be safe, we should put the host code (the wrapper functions) in `cpp_sources` and the device code (kernels) in `cuda_sources`. However, the kernels need to be declared before use. In separate compilation, we need to declare the kernels in the C++ file or include a header. But `load_inline` compiles both sources and links them. If we put wrappers in `cpp_sources`, they need to see the kernel declarations. We can forward-declare the kernels in `cpp_sources`.
Alternatively, we can put everything in `cuda_sources` and leave `cpp_sources` empty. NVCC will handle both device and host code in the same compilation unit. That is simpler and avoids linkage issues.
Let's try putting everything in `cuda_sources`.
Wait, `load_inline` documentation: The `cuda_sources` are passed to nvcc, which compiles them into a .cu object. The `cpp_sources` are passed to the C++ compiler. Both are linked together. If we put host code in `cuda_sources`, NVCC will compile it. That's fine.
But there is a subtlety with `AT_ASSERTM` and torch includes in `cuda_sources`. It should work because NVCC includes the same headers. So let's put everything into `cuda_sources` and leave `cpp_sources` empty.
Actually, I recall that in some PyTorch versions, `load_inline` with `cpp_sources=''` and `cuda_sources=...` works fine. Let's do that.
One more thing: We need to avoid name collisions. The module name `reparam_cuda` should be unique.
Potential compilation issue: `std::pair<at::Tensor, at::Tensor>` requires `<utility>` or it is included by `<torch/extension.h>`. Should be fine.
Also, `torch::empty` might require using `at::Tensor` or `torch::Tensor`. `torch::empty` is in the `torch` namespace and returns `at::Tensor`. That's fine.
Now, test compilation.
But wait, `CHECK_CONTIGUOUS` macro: `x.is_contiguous()` checks all dimensions. For the simple case with 2D tensors, it's fine. For the mixture case with 3D tensors, it's fine.
Another issue: `mu` and `log_var` are passed to `forward`. The reference computes `std = torch.exp(0.5 * log_var)`. In our code, we compute `std = torch.exp(0.5 * log_var)`. We then call `std = std.contiguous()`. Good.
In the mixture branch, we call `_cuda_ext.mixture_forward_cuda(...)`. The function expects `log_weights` as a 1D tensor. We compute `log_weights = weights.log().contiguous()`. Good.
Now, let's think about the `const_val` argument. It's a Python float. When passed to C++, it becomes a C++ `float`. In the wrapper signature, we have `float const_val`. `load_inline` generated pybind11 code should handle `float` arguments. Yes, pybind11 casts Python float to C++ float.
One more thing: In the mixture branch, we call `eps = torch.randn(B, D, device=mu.device, dtype=mu.dtype)`. But `mu.dtype` could be float32 or float64. The reference uses default float32. `get_inputs()` uses `torch.randn` which defaults to float32. So `mu.dtype` is float32. Our `eps` will be float32. Good.
But there's an important detail: In the mixture branch, the reference creates `mu_k` and `std_k` by indexing, then calls `torch.randn_like(mu_k)`. Since `mu_k` is float32, `eps` is float32. In our code, we call `torch.randn(B, D, device=mu.device, dtype=mu.dtype)`. This is equivalent. However, the reference's `mu_k` is computed on the same device as `mu`, which is CPU or GPU. Since we will run on GPU, it's GPU.
Now, in the simple case, we call `eps = torch.randn_like(mu[:, 0] if mu.dim() == 3 else mu)`. This is also equivalent.
One more subtlety: The reference code for the simple case where `mu.dim() == 3` reassigns `mu = mu[:, 0]`. Then it uses this new `mu` to compute `z` and `log_q_z`. The original `mu` is not used after that. In our code, we keep the original 3D `mu` and pass it to the kernel. The kernel reads `comp=0`. This is correct.
But the reference also reassigns `log_var = log_var[:, 0]`. Our kernel uses the original 3D `log_var`. Correct.
Now, compilation check. Let's generate the code and compile it.
Wait, the `cpp_sources` argument in `load_inline` might need to be non-empty? Let's check typical usage. In many examples, `load_inline` is used with `cpp_sources` containing the host wrappers and `cuda_sources` containing the kernels. But if we put everything in `cuda_sources`, it should be okay because NVCC compiles host code too.
Let's see the `torch.utils.cpp_extension.load_inline` signature:
```python
def load_inline(name,
cpp_sources,
cuda_sources=None,
functions=None,
extra_cflags=None,
extra_cuda_cflags=None,
extra_ldflags=None,
extra_include_paths=None,
build_directory=None,
verbose=False,
with_cuda=None,
is_python_module=True,
with_pytorch_error_handling=True,
keep_intermediates=True):
```
If `cuda_sources` is provided, `with_cuda` is set to True. The `cpp_sources` are compiled as C++ files, and `cuda_sources` as `.cu` files. Then linked. If we put the C++ host code in `cuda_sources`, it will be compiled by NVCC as a `.cu` file. That should work.
Let's do that.
We also need to make sure to include `<torch/extension.h>`.
Let's also consider using `torch::checkAllSameGPU` or not. We have our own checks.
Now, compile.
But there might be an issue: `load_inline` auto-generates a pybind11 module that calls the C++ functions. The functions must be declared in the `cpp_sources` or `cuda_sources` such that the generated bindings can find them. Since we put them in `cuda_sources`, NVCC will compile them with external linkage. The generated bindings will include declarations of these functions and then call them. This should work.
Potential issue: `std::pair<at::Tensor, at::Tensor>` might not be recognized by the auto-generated pybind11 code as a tuple? Actually, `torch.utils.cpp_extension` uses a custom binding generator that specifically supports `std::pair<at::Tensor, at::Tensor>` and returns a tuple of tensors. This is standard.
Now, let's write the complete code.
Before finalizing, let's consider an extra optimization: In the mixture case, `z` is computed and then each component's log-prob uses `z`. For `k_sel`, `diff = z - mu_k_sel = std_k_sel * eps`. So the term is `-0.5 * eps^2 - 0.5 * log_var_k_sel + const_val`. We could compute this specially to avoid division by `std_k_sel`, but it's a minor optimization. Not needed.
Also, note that in the reference, `log_var` is used directly, while `std` is used for division. We could avoid passing `std` and compute `std` as `exp(0.5*lv)` inside the kernel, but that would be slower. Passing `std` is fine.
Actually, we compute `std` in Python and pass it. That matches the reference exactly.
One more potential issue: The reference's `std = torch.exp(0.5 * log_var)` may be computed as an intermediate. In PyTorch, this is fused into a single kernel if using TorchInductor or JIT, but not in eager mode. Our code also computes it eagerly. That's fine.
Let's also verify the shape of `log_weights`. In reference: `weights = torch.softmax(self.logits, dim=0)`. Then `weights.log().unsqueeze(0)` is added to `log_probs`. So `log_weights` is (K,). We pass it as a 1D tensor.
Now, let's write the code.
One more thing: The `mixture_forward_cuda` wrapper uses `comp.data_ptr<int64_t>()`. This is fine.
Now, test compilation. I'll write the code into the tool.