reasoning
Let me analyze the problem:
1. We have a PyTorch model that computes importance sampling weights
2. Input: `samples` of shape (N, dim) = (4096, 16)
3. The model computes:
- `log_p`: log density of a Gaussian mixture with K=3 components, D=16
- `log_q`: log density of a single Gaussian proposal
- `log_w = log_p - log_q`
- Normalize weights: `log_w = log_w - logsumexp(log_w)`
- `w = exp(log_w)`
- `weighted_mean = (w.unsqueeze(-1) * samples).sum(dim=0)`
- `ess = 1.0 / (w^2).sum()`
4. We need to write a CUDA kernel that's faster than the reference.
Let me think about optimization opportunities:
- The reference does multiple passes over the data and allocates intermediate tensors
- We can fuse all operations into a single kernel (or a few kernels)
- N=4096, D=16, K=3 - this is relatively small but we can still optimize
Memory bandwidth bound analysis:
- Input: 4096 * 16 * 4 = 256 KB
- Buffers (mix_mus, mix_stds, etc.) are tiny
- Output is tiny
- Total memory traffic if we fuse: ~256 KB read + a few bytes write
- At 3.35 TB/s, this is ~0.000076 ms = 0.076 us theoretical
- Real kernels take ~10-100 us for this size due to kernel launch overhead and inefficiency
A fused kernel approach:
1. First kernel: compute log_w for each sample and store in shared memory
2. Compute logsumexp using parallel reduction
3. Normalize, compute exp, accumulate weighted_mean and w^2 sum
Actually, let me break this into a couple of kernels:
1. Compute log_p and log_q for each sample, output log_w
2. Compute logsumexp of log_w
3. Compute weighted mean and ess
But we can do even better by fusing step 2 and 3, or even everything.
Let me think about the math more carefully:
log_p computation for each sample i:
- diff[i,k,d] = samples[i,d] - mix_mus[k,d]
- log_comp[i,k] = sum_d[-0.5 * (diff[i,k,d] / mix_stds[k,d])^2 - log(mix_stds[k,d])]
- log_p[i] = logsumexp_k[log_comp[i,k] + log_mix_weights[k]]
log_q computation for each sample i:
- log_q[i] = sum_d[-0.5 * ((samples[i,d] - q_mu[d]) / q_std[d])^2 - log(q_std[d])]
Then log_w[i] = log_p[i] - log_q[i]
We need a two-pass algorithm:
- Pass 1: compute all log_w[i]
- Pass 2: compute logsumexp(log_w)
- Pass 3: compute exp(log_w - lse), weighted sum, and ESS
Since N=4096, we could potentially have each block handle a subset of samples. Let me design a kernel where:
- Each thread block computes log_w for its samples
- We do a parallel reduction for logsumexp
- Then compute final results
Actually, doing this in two kernels is straightforward and efficient:
1. `compute_logw`: compute log_w per sample
2. `compute_results`: compute lse (via atomic or another approach), weighted mean, ess
Wait, for N=4096 we can actually use a single kernel with shared memory reduction. Here's my plan:
Kernel design:
- Grid: enough blocks to cover N samples (e.g., 32 blocks of 128 threads = 4096)
- Each thread handles one sample
- Phase 1: compute log_w[i] and store in shared memory
- Phase 2: parallel reduction for logsumexp(log_w) across all threads in the block
- Phase 3: if using multiple blocks, we need to reduce across blocks using atomics or a second kernel
For N=4096, we could use just one block with 4096 threads, but that's hardware-limited. Better: use multiple blocks, each writes partial sum to global memory, then reduce.
Actually, simpler approach:
1. `compute_logw_kernel`: launch with N threads, each computes log_w[i] and stores in a temporary buffer
2. `final_kernel`: use a single block with N threads (or multiple blocks with final atomic reduction):
- Each thread reads log_w[i], computes normalized w[i]
- Does parallel reduction for weighted_mean (D=16 elements per thread contributes)
- Does parallel reduction for sum_w2
Wait, final reduction of weighted_mean across D=16 dimensions... we need a sum of w[i] * samples[i,d] for each d. And sum of w[i]^2.
Let me think more carefully:
For kernel 2, we can have:
- BlockDim = 128 threads
- GridDim = ceil(N / 128) = 32 blocks
- Each thread in a block handles one or more samples
- Uses shared memory for:
- partial weighted_mean[16] per block
- partial sum_w2 per block
- At the end, atomically accumulate to global output
Actually, let me reconsider. N=4096 with D=16 and K=3. The compute is:
- For each of 4096 samples: compute log_p (3 components * 16 dims each = 48 operations) and log_q (16 dims). Total ~64 operations per sample.
- Total compute: ~262k ops - very small
- Memory: read samples (256KB), read model parameters (tiny), write outputs (tiny)
This is going to be heavily memory-bandwidth-bound and kernel-launch-overhead bound. Fusing into 1-2 kernels is critical.
Let me design a 2-kernel approach:
Kernel 1: `compute_logw`
- N threads
- Each thread computes log_w for one sample
- Writes to `logw` buffer of size N
Kernel 2: `compute_results`
- Single block with N threads (or multiple blocks with atomic reduce)
- Actually, for N=4096, we can use 1 block with 1024 threads, each handling 4 samples
- Loads logw[4] and samples[4*16]
- Computes lse using shared memory reduction (logsumexp trick)
- Then computes w = exp(logw - lse)
- Accumulates weighted_mean[16] and sum_w2
- Final block writes to output
Hmm, but doing logsumexp and then the weighted sums in the same kernel requires either:
a) Two passes within the kernel (compute all logw -> reduce -> compute weighted sums), or
b) Precompute logsumexp separately
For approach (a):
- Use 1 block, so we can do:
- Phase 1: each thread computes logw for its samples, stores in smem
- Phase 2: parallel reduction for logsumexp in smem
- Phase 3: broadcast lse, compute w and accumulate weighted_mean + sum_w2
This is clean for N=4096! We can use 1 block with 1024 threads (each thread handles 4 samples).
Wait, but K=3 and D=16... computing log_p requires reading samples[16] and model params. If we use 1024 threads but only 4096 samples, we'd have 32 warps. We should be careful about divergence and memory access patterns.
Actually, let me use a simpler approach: just a single block with N threads (4096 threads). CUDA supports up to 1024 threads per block. So we need multiple blocks or fewer threads per sample.
Let me go with: 1 block, 1024 threads, each thread handles 4 samples sequentially.
Shared memory layout:
- `logw[4096]` - actually 4096 floats = 16KB, well within shared memory
Actually, let me use a different approach to minimize shared memory:
- We can compute logsumexp using a warp shuffle approach without storing all logw in smem if we're clever, but it's more complex.
Let me use 2 kernels for simplicity and robustness:
Kernel 1: compute_logw (N threads)
Kernel 2: reduction_kernel (1 block, N threads, but N=4096 so actually 1024 threads with 4 samples each)
For kernel 2:
```cuda
__global__ void reduction_kernel(const float* __restrict__ logw,
const float* __restrict__ samples,
float* __restrict__ out_mean,
float* __restrict__ out_ess,
int N, int D) {
// Block handles all N samples
// 1024 threads, each handles 4 samples (for N=4096)
const int tid = threadIdx.x;
const int num_threads = blockDim.x;
// Phase 1: compute logsumexp
// Each thread computes partial logsumexp for its samples
float local_max = -INFINITY;
for (int i = tid; i < N; i += num_threads) {
local_max = fmaxf(local_max, logw[i]);
}
// Reduce max across warp
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
local_max = fmaxf(local_max, __shfl_down_sync(0xffffffff, local_max, offset));
}
// Store warp max in shared memory
__shared__ float warp_max[32];
if (tid % 32 == 0) warp_max[tid / 32] = local_max;
__syncthreads();
// Reduce across warps
if (tid < 32) {
local_max = (tid < blockDim.x / 32) ? warp_max[tid] : -INFINITY;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
local_max = fmaxf(local_max, __shfl_down_sync(0xffffffff, local_max, offset));
}
if (tid == 0) warp_max[0] = local_max;
}
__syncthreads();
float global_max = warp_max[0];
// Now compute sum of exp(logw - global_max)
float local_sum = 0.0f;
for (int i = tid; i < N; i += num_threads) {
local_sum += expf(logw[i] - global_max);
}
// Reduce sum
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
local_sum += __shfl_down_sync(0xffffffff, local_sum, offset);
}
if (tid % 32 == 0) warp_max[tid / 32] = local_sum;
__syncthreads();
if (tid < 32) {
local_sum = (tid < blockDim.x / 32) ? warp_max[tid] : 0.0f;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
local_sum += __shfl_down_sync(0xffffffff, local_sum, offset);
}
if (tid == 0) warp_max[0] = local_sum;
}
__syncthreads();
float lse = global_max + logf(warp_max[0]);
// Phase 2: compute weighted_mean and sum_w2
__shared__ float smem_mean[16]; // D=16
__shared__ float smem_w2;
if (tid < D) smem_mean[tid] = 0.0f;
if (tid == 0) smem_w2 = 0.0f;
__syncthreads();
for (int i = tid; i < N; i += num_threads) {
float w = expf(logw[i] - lse);
float w2 = w * w;
atomicAdd(&smem_w2, w2);
const float* sample = samples + i * D;
#pragma unroll
for (int d = 0; d < D; d++) {
atomicAdd(&smem_mean[d], w * sample[d]);
}
}
__syncthreads();
if (tid < D) out_mean[tid] = smem_mean[tid];
if (tid == 0) *out_ess = 1.0f / smem_w2;
}
```
Wait, using atomics in a loop might be slow. Let me think of a better approach.
For weighted_mean, we can have each thread accumulate its own local_mean[16], then reduce across threads. But reducing 16 floats across 1024 threads is... possible but a lot of shared memory (16 * 1024 * 4 = 64KB, which is too much).
Alternative: use 16 warps, each warp computes weighted_mean for one dimension.
- Warp 0 computes sum_d0, warp 1 computes sum_d1, etc.
- For D=16, we need 16 warps = 512 threads. We can use 512 threads.
- Each warp handles N samples (seq4 for N=4096).
- Each warp also handles sum_w2? Actually sum_w2 can be computed by all warps redundantly, or by warp 0.
Let me redesign:
Use 512 threads (16 warps).
Each thread handles `N / 512 = 8` samples? Wait, N=4096 / 512 = 8 samples per thread.
Actually, let's use 256 threads (8 warps), each handling 16 samples.
Or 128 threads, each handling 32 samples.
Given the small N, let me use 128 threads (4 warps), each handling 32 samples sequentially.
But actually, sequential handling might be fine since compute is higher than memory bandwidth for this kernel.
Let me try a different approach:
- Kernel 1: compute logw (N threads)
- Kernel 2: compute partial sums per block, then final reduce
- Use 32 blocks of 128 threads each
- Each block handles 128 samples
- Block computes local weighted_mean[16] and local sum_w2
- Writes to global partial buffers
- Another tiny kernel or atomic reduce for final result
Actually, since N=4096 (fixed? Actually it seems fixed at 4096 from the code), but we shouldn't hardcode.
Let me use a cleaner architecture:
- Kernel 1: compute_logw
- Kernel 2: compute_logsumexp (using CUB or custom tree reduction, writing one float)
- Kernel 3: compute_weighted_mean_ess (N threads, using atomicAdd for mean and w2)
But 3 kernels = 3 launches = overhead. For N=4096 this might still be okay.
Actually the simplest fast approach: fuse everything into 2 kernels.
Kernel 1: compute log_w for each sample.
Kernel 2: single block, compute lse via reduction, then compute weighted_mean and ess.
For kernel 2, I'll use a different reduction strategy. I'll have each thread process samples in a strided loop and accumulate into thread-local registers, then reduce.
For logsumexp:
```cuda
float local_max = -INFINITY;
for (int i = tid; i < N; i += blockDim.x) {
local_max = fmaxf(local_max, logw[i]);
}
// blockReduceMax
float global_max = blockReduceMax(local_max);
// block broadcast
float local_sum = 0;
for (int i = tid; i < N; i += blockDim.x) {
local_sum += expf(logw[i] - global_max);
}
float total_sum = blockReduceSum(local_sum);
float lse = logf(total_sum) + global_max;
```
For weighted_mean and ess, I can do something similar. But for 16 dimensions, I need a blockReduce for each dimension.
I could do:
```cuda
float local_mean[16] = {0};
float local_w2 = 0;
for (int i = tid; i < N; i += blockDim.x) {
float w = expf(logw[i] - lse);
local_w2 += w * w;
#pragma unroll
for (int d = 0; d < D; d++) {
local_mean[d] += w * samples[i * D + d];
}
}
// Now reduce local_mean[16] and local_w2 across block
float total_w2 = blockReduceSum(local_w2);
__shared__ float shared_mean[16];
#pragma unroll
for (int d = 0; d < D; d++) {
shared_mean[d] = blockReduceSum(local_mean[d]);
}
```
Wait, `blockReduceSum` for arrays. I can use shared memory to store per-thread partial sums, but that requires blockDim.x * sizeof(float) * D which is too much for D=16 and blockDim=1024 (64KB).
Alternative: reduce each dimension separately using warp shuffles but sequentially. This is what I mentioned above.
Actually, for D=16, we can simply do:
```cuda
for (int d = 0; d < D; d++) {
float local_m = local_mean[d];
// warp reduce
for (int offset = 16; offset > 0; offset /= 2) {
local_m += __shfl_down_sync(0xffffffff, local_m, offset);
}
if (tid % 32 == 0) s_warp_sums[tid / 32 + d * 32] = local_m;
__syncthreads();
if (tid < 32) {
float val = (tid < num_warps) ? s_warp_sums[tid + d * 32] : 0;
for (int offset = 16; offset > 0; offset /= 2) {
val += __shfl_down_sync(0xffffffff, val, offset);
}
if (tid == 0) out_mean[d] = val;
}
__syncthreads();
}
```
This does 16 separate warp reductions. With 32 warps and 16 dimensions, the shared memory needed is 16 * 32 * 4 = 2KB.
For `sum_w2`, we do one more reduction. And for `max` and `sum` in logsumexp, two more reductions.
Total: 16 + 1 + 2 = 19 reductions. At ~10-20 cycles each, this is ~200-400 cycles. For a kernel that processes 4096 samples, this is negligible.
Actually wait, we can optimize further. For the two logsumexp reductions (max and sum), we need them before computing w. So we must do:
1. Find global max of logw
2. Compute sum of exp(logw - max)
3. Compute lse = log(sum) + max
4. For each sample, compute w = exp(logw - lse), accumulate local_mean[d] and local_w2
5. Reduce local_mean and local_w2
This means we need two passes over the data in the kernel (one for max, one for accumulating). But since it's in shared memory / L1 cache, the second pass is cheap.
Let me write this out properly:
```cuda
extern "C" __global__ void importance_sampling_kernel(
const float* __restrict__ samples,
const float* __restrict__ logw,
float* __restrict__ out_mean,
float* __restrict__ out_ess,
int N, int D) {
const int tid = threadIdx.x;
const int num_threads = blockDim.x;
const int num_warps = (num_threads + 31) / 32;
__shared__ float s_warp_buf[32 * 17]; // enough for 16 dims + w2, per warp
// Phase 1: find max(logw)
float local_max = -INFINITY;
for (int i = tid; i < N; i += num_threads) {
local_max = fmaxf(local_max, logw[i]);
}
// Warp reduce max
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
local_max = fmaxf(local_max, __shfl_down_sync(0xFFFFFFFF, local_max, offset));
}
if (tid % 32 == 0) {
s_warp_buf[tid / 32] = local_max;
}
__syncthreads();
// Block reduce max
if (tid < 32) {
local_max = (tid < num_warps) ? s_warp_buf[tid] : -INFINITY;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
local_max = fmaxf(local_max, __shfl_down_sync(0xFFFFFFFF, local_max, offset));
}
if (tid == 0) s_warp_buf[0] = local_max;
}
__syncthreads();
float global_max = s_warp_buf[0];
// Phase 2: compute sum of exp(logw - global_max)
float local_sum = 0.0f;
for (int i = tid; i < N; i += num_threads) {
local_sum += expf(logw[i] - global_max);
}
// Warp reduce sum
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
local_sum += __shfl_down_sync(0xFFFFFFFF, local_sum, offset);
}
if (tid % 32 == 0) {
s_warp_buf[tid / 32] = local_sum;
}
__syncthreads();
// Block reduce sum
if (tid < 32) {
local_sum = (tid < num_warps) ? s_warp_buf[tid] : 0.0f;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
local_sum += __shfl_down_sync(0xFFFFFFFF, local_sum, offset);
}
if (tid == 0) s_warp_buf[0] = local_sum;
}
__syncthreads();
float total_sum = s_warp_buf[0];
float lse = global_max + logf(total_sum);
// Phase 3: compute weighted_mean and sum_w2
float local_mean[16];
#pragma unroll
for (int d = 0; d < D; d++) {
local_mean[d] = 0.0f;
}
float local_w2 = 0.0f;
for (int i = tid; i < N; i += num_threads) {
float w = expf(logw[i] - lse);
float w2 = w * w;
local_w2 += w2;
const float* sample = samples + i * D;
#pragma unroll
for (int d = 0; d < D; d++) {
local_mean[d] += w * sample[d];
}
}
// Reduce each dimension
#pragma unroll
for (int d = 0; d < D; d++) {
float val = local_mean[d];
// Warp reduce
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
val += __shfl_down_sync(0xFFFFFFFF, val, offset);
}
if (tid % 32 == 0) {
s_warp_buf[tid / 32 + d * num_warps] = val;
}
}
// w2 reduce
float w2_val = local_w2;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
w2_val += __shfl_down_sync(0xFFFFFFFF, w2_val, offset);
}
if (tid % 32 == 0) {
s_warp_buf[tid / 32 + D * num_warps] = w2_val;
}
__syncthreads();
// Block reduce for each dimension
if (tid < 32) {
// For each dimension d
#pragma unroll
for (int d = 0; d < D; d++) {
float val = (tid < num_warps) ? s_warp_buf[tid + d * num_warps] : 0.0f;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
val += __shfl_down_sync(0xFFFFFFFF, val, offset);
}
if (tid == 0) {
out_mean[d] = val;
}
}
// w2
float w2_sum = (tid < num_warps) ? s_warp_buf[tid + D * num_warps] : 0.0f;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
w2_sum += __shfl_down_sync(0xFFFFFFFF, w2_sum, offset);
}
if (tid == 0) {
*out_ess = 1.0f / w2_sum;
}
}
}
```
Wait, there's a problem. If D is not known at compile time, we can't unroll the loop. But D=16 is fixed in the model's `dim`. We should pass D as a template parameter or make the kernel handle it dynamically.
Since `dim` is a constructor argument, we need to handle it at runtime. But in practice, the evaluation will use `dim=16`. Let me write it to handle D generically but optimize for D=16.
Actually for simplicity, I can pass D as a regular argument and loop dynamically. The compiler might not unroll, but for small D it might still be okay. Or I can template it.
Let me just use a dynamic loop but with `#pragma unroll` - actually, the compiler won't unroll a loop with variable bound. For maximum performance, I should template. But the evaluation only tests dim=16, so let me use a compile-time constant or pass it as a parameter.
Actually, `cpp_extension.load_inline` doesn't easily support templates in the function signature. Let me just use a runtime parameter and not worry too much about unrolling for now. If it's not fast enough, I can optimize later.
Actually wait, for the inner loop `for (int d = 0; d < D; d++)`, since D=16 is small, even without unrolling it should be fine. Let me use the runtime parameter.
Now, for kernel 1 (compute logw):
```cuda
extern "C" __global__ void compute_logw_kernel(
const float* __restrict__ samples,
const float* __restrict__ mix_mus,
const float* __restrict__ mix_stds,
const float* __restrict__ mix_weights,
const float* __restrict__ q_mu,
const float* __restrict__ q_std,
float* __restrict__ logw,
int N, int D, int K) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= N) return;
const float* sample = samples + idx * D;
// Compute log_p
// log_comp[k] = sum_d[-0.5 * (diff / std)^2 - log(std)]
float log_comp[3]; // K=3
#pragma unroll
for (int k = 0; k < K; k++) {
const float* mu = mix_mus + k * D;
const float* std = mix_stds + k * D;
float sum = 0.0f;
#pragma unroll
for (int d = 0; d < D; d++) {
float diff = sample[d] - mu[d];
float scaled = diff / std[d];
sum += -0.5f * scaled * scaled - logf(std[d]);
}
log_comp[k] = sum + logf(mix_weights[k]);
}
// logsumexp over K components
float max_val = log_comp[0];
#pragma unroll
for (int k = 1; k < K; k++) {
max_val = fmaxf(max_val, log_comp[k]);
}
float sum_exp = 0.0f;
#pragma unroll
for (int k = 0; k < K; k++) {
sum_exp += expf(log_comp[k] - max_val);
}
float log_p = max_val + logf(sum_exp);
// Compute log_q
float log_q = 0.0f;
#pragma unroll
for (int d = 0; d < D; d++) {
float diff = sample[d] - q_mu[d];
float scaled = diff / q_std[d];
log_q += -0.5f * scaled * scaled - logf(q_std[d]);
}
logw[idx] = log_p - log_q;
}
```
Since K=3, we can use a fixed-size array `log_comp[3]`. But to make it general, let me use a loop. Actually, K is also fixed at 3 by default. Let me keep it general but hopefully the compiler handles K=3 well.
Wait, I realize K is a constructor arg too, defaulting to 3. I can pass it as a runtime parameter.
Actually, looking more carefully: `mix_weights` is a buffer of size `num_components`. So K = num_components.
Hmm, for better performance, let me try to simplify. Since this is bandwidth/compute bound, let me see if I can just write a clean 2-kernel solution and see if it beats PyTorch.
PyTorch does:
1. `diff = x.unsqueeze(1) - self.mix_mus.unsqueeze(0)` - allocates (B, K, D) = (4096, 3, 16) * 4 = 768KB
2. `(diff / stds)^2` - another allocation
3. `log_comp` computation - more allocations
4. `logsumexp` - more allocations
5. `log_q` - similar ops but fewer allocations
6. `log_w = log_p - log_q`
7. `log_w - logsumexp(log_w)` - another allocation + reduction
8. `w = exp(log_w)` - another allocation
9. `w.unsqueeze(-1) * samples` - (4096, 16) allocation
10. `.sum(dim=0)` - reduction
11. `(w**2).sum()` - another allocation + reduction
This is a LOT of kernel launches and memory allocations. My fused approach should be significantly faster.
Let me also consider fusing into a single kernel. For a single kernel:
- We'd need to compute logw, then do block reduction for lse, then compute weighted sums.
- If we use a single block, we can do this easily. But N=4096 with a single block of 1024 threads means each thread does 4 samples.
- With 1024 threads, we have 32 warps. The model parameters (mix_mus, mix_stds, etc.) can stay in registers / L1 cache.
- The samples need to be read from global memory.
- This would be ONE kernel launch instead of two.
Single kernel approach:
```cuda
__global__ void full_kernel(
const float* __restrict__ samples,
const float* __restrict__ mix_mus,
const float* __restrict__ mix_stds,
const float* __restrict__ mix_weights,
const float* __restrict__ q_mu,
const float* __restrict__ q_std,
float* __restrict__ out_mean,
float* __restrict__ out_ess,
int N, int D, int K) {
// One block handles all N samples
const int tid = threadIdx.x;
const int num_threads = blockDim.x;
// Compute logw for each sample assigned to this thread
// We need to keep logw around for the second pass, but we can recompute...
// Actually, recomputing is expensive. Let's store in shared memory.
// N=4096, float logw[4096] = 16KB shared memory. That's fine!
// Wait, we only need to keep logw if we want to avoid recomputing.
// For N=4096, 16KB is small. But for larger N, it might exceed shared memory.
// The problem says num_samples=4096, dim=16. But should we handle arbitrary N?
// Let's support arbitrary N by processing in chunks or just using a two-kernel approach.
```
Actually, a two-kernel approach is cleaner and handles arbitrary N. Let me stick with it.
But I just realized: for kernel 2, using a single block limits scalability to large N. However, since N=4096 is small, a single block is actually fine and avoids cross-block synchronization overhead.
Wait, what if the evaluation uses larger N? The code says `num_samples = 4096`, but maybe they test with different sizes. Let me make it handle arbitrary N by using multiple blocks with a final atomic reduce or by using a separate reduction kernel.
Actually, for safety, let me use a two-block approach for kernel 2:
- Multiple blocks process chunks of samples
- Each block computes local_mean[16] and local_w2 and writes to global partial buffers
- First thread of each block also participates in computing lse... but lse requires global knowledge.
- Actually, we need global lse before computing w. So we'd need:
- Kernel 1: compute logw per sample
- Kernel 2: compute lse (device-wide reduction)
- Kernel 3: compute weighted_mean and ess (device-wide reduction)
With 3 kernels, we can scale to arbitrary N easily.
For device-wide reduction of lse:
- We can use `cub::DeviceReduce::Max` followed by a custom kernel for sum of exp.
- Or we can write a simple tree reduction kernel.
- Or we can just use a single block for N <= 4096.
Given the constraints of `load_inline`, using CUB might be tricky unless we include the headers.
Actually, CUB is included with CUDA toolkit, so we can use it. But compiling with CUB through `load_inline` might add compile time.
Let me just use a simple approach: for N <= 8192, use a single block. For larger N, use multiple blocks with atomics.
Actually, atomicAdd for logsumexp is tricky. For max, we can use atomicMax (for float bits). For sum, after we have the global max, we can atomicAdd the exp terms.
Or even simpler: just use a single block. CUDA blocks can have up to 1024 threads. If N > 1024 * work_per_thread, we can just increase work_per_thread. For N = 4096 and 1024 threads, each thread handles 4 samples. For N = 65536 and 1024 threads, each thread handles 64 samples. The loop is still efficient. A single block can handle quite large N before becoming a bottleneck.
But there is a limit: with 1024 threads, we want each thread to do enough work to amortize the reduction overhead. If N = 1M, each thread does 1000 iterations. That might be okay but a single block would underutilize the GPU.
For safety and generality, let me use a multi-block reduction approach.
Actually, let me compromise:
- I'll write a 2-kernel solution
- Kernel 1: compute logw (multi-block, N threads)
- Kernel 2: single block reduction (with arbitrary N supported by strided loops)
This handles arbitrary N (up to whatever fits in reasonable time per block). For really large N, performance won't be optimal due to single block, but the reference says N=4096.
Wait, I should double check: `num_samples = 4096` and `dim = 16` are module-level globals. The evaluation probably uses these exact values. Let me optimize for these.
For N=4096, single block approach is completely fine. Let me optimize for that.
Let me now write the complete Python file.
Actually, I can do even better by fusing into ONE kernel with a single block:
```cuda
__global__ void fused_kernel(...) {
// Single block, T threads
// Phase 1: each thread computes logw for its assigned samples
// Store logw in shared memory (N floats = 16KB)
// Phase 2: block reduction for logsumexp
// Phase 3: each thread computes w and accumulates weighted_mean + w2
// Phase 4: block reduction for mean and w2
}
```
Advantage: only one kernel launch, no intermediate global memory write/read for logw. Everything stays in shared memory / registers.
Let me try this! For N=4096 with 1024 threads:
- Each thread computes 4 logw values and stores them in shared memory
- Block reduction (1024 elements in shared memory)
- Each thread processes its 4 samples again, computes w, accumulates mean and w2
- Block reductions for mean[16] and w2
Wait, storing 4096 floats in shared memory is 16KB. Plus some scratch space for reductions. We need:
- `logw[4096]` = 16KB
- `warp_sums` for reduction = maybe a few hundred bytes
Total well under 48KB shared memory limit. This should work great!
But what if N is larger than what fits in shared memory? Well, for N up to ~10000 it still fits. For larger N, we'd need to spill to global memory or process in tiles.
For the evaluation, N=4096, so this is perfect.
Let me design the single-kernel version:
```cuda
__global__ void importance_sampling_fused(
const float* __restrict__ samples,
const float* __restrict__ mix_mus,
const float* __restrict__ mix_stds,
const float* __restrict__ mix_weights,
const float* __restrict__ q_mu,
const float* __restrict__ q_std,
float* __restrict__ out_mean,
float* __restrict__ out_ess,
int N, int D, int K) {
const int tid = threadIdx.x;
const int num_threads = blockDim.x;
// Shared memory: logw for all samples
extern __shared__ float shared[];
float* s_logw = shared; // size N
// Each thread computes logw for its samples
for (int i = tid; i < N; i += num_threads) {
const float* sample = samples + i * D;
// Compute log_p
float log_comp[8]; // Support up to K=8
for (int k = 0; k < K; k++) {
const float* mu = mix_mus + k * D;
const float* std = mix_stds + k * D;
float sum = 0.0f;
for (int d = 0; d < D; d++) {
float diff = sample[d] - mu[d];
float scaled = diff / std[d];
sum += -0.5f * scaled * scaled - logf(std[d]);
}
log_comp[k] = sum + logf(mix_weights[k]);
}
float max_val = log_comp[0];
for (int k = 1; k < K; k++) {
max_val = fmaxf(max_val, log_comp[k]);
}
float sum_exp = 0.0f;
for (int k = 0; k < K; k++) {
sum_exp += expf(log_comp[k] - max_val);
}
float log_p = max_val + logf(sum_exp);
// Compute log_q
float log_q = 0.0f;
for (int d = 0; d < D; d++) {
float diff = sample[d] - q_mu[d];
float scaled = diff / q_std[d];
log_q += -0.5f * scaled * scaled - logf(q_std[d]);
}
s_logw[i] = log_p - log_q;
}
__syncthreads();
// Phase 2: Find global max of logw
float local_max = -INFINITY;
for (int i = tid; i < N; i += num_threads) {
local_max = fmaxf(local_max, s_logw[i]);
}
// Block reduce max
#pragma unroll
for (int offset = warpSize/2; offset > 0; offset /= 2) {
local_max = fmaxf(local_max, __shfl_down_sync(0xFFFFFFFF, local_max, offset));
}
if (tid % warpSize == 0) {
s_logw[N + tid / warpSize] = local_max; // reuse shared memory past logw
}
__syncthreads();
int num_warps = num_threads / warpSize;
if (tid < warpSize) {
local_max = (tid < num_warps) ? s_logw[N + tid] : -INFINITY;
#pragma unroll
for (int offset = warpSize/2; offset > 0; offset /= 2) {
local_max = fmaxf(local_max, __shfl_down_sync(0xFFFFFFFF, local_max, offset));
}
if (tid == 0) s_logw[N] = local_max;
}
__syncthreads();
float global_max = s_logw[N];
// Phase 3: Sum of exp(logw - global_max)
float local_sum = 0.0f;
for (int i = tid; i < N; i += num_threads) {
local_sum += expf(s_logw[i] - global_max);
}
#pragma unroll
for (int offset = warpSize/2; offset > 0; offset /= 2) {
local_sum += __shfl_down_sync(0xFFFFFFFF, local_sum, offset);
}
if (tid % warpSize == 0) {
s_logw[N + tid / warpSize] = local_sum;
}
__syncthreads();
if (tid < warpSize) {
local_sum = (tid < num_warps) ? s_logw[N + tid] : 0.0f;
#pragma unroll
for (int offset = warpSize/2; offset > 0; offset /= 2) {
local_sum += __shfl_down_sync(0xFFFFFFFF, local_sum, offset);
}
if (tid == 0) s_logw[N + 1] = local_sum;
}
__syncthreads();
float total_sum = s_logw[N + 1];
float lse = global_max + logf(total_sum);
// Phase 4: Compute weighted_mean and sum_w2
float local_mean[16]; // D=16
for (int d = 0; d < D; d++) local_mean[d] = 0.0f;
float local_w2 = 0.0f;
for (int i = tid; i < N; i += num_threads) {
float w = expf(s_logw[i] - lse);
local_w2 += w * w;
const float* sample = samples + i * D;
for (int d = 0; d < D; d++) {
local_mean[d] += w * sample[d];
}
}
// Reduce each dimension
float* s_buf = s_logw + N + 2; // scratch space after logw and two floats
for (int d = 0; d < D; d++) {
float val = local_mean[d];
#pragma unroll
for (int offset = warpSize/2; offset > 0; offset /= 2) {
val += __shfl_down_sync(0xFFFFFFFF, val, offset);
}
if (tid % warpSize == 0) {
s_buf[(tid / warpSize) + d * num_warps] = val;
}
}
float w2_val = local_w2;
#pragma unroll
for (int offset = warpSize/2; offset > 0; offset /= 2) {
w2_val += __shfl_down_sync(0xFFFFFFFF, w2_val, offset);
}
if (tid % warpSize == 0) {
s_buf[num_warps * D + tid / warpSize] = w2_val;
}
__syncthreads();
if (tid < warpSize) {
for (int d = 0; d < D; d++) {
float val = (tid < num_warps) ? s_buf[tid + d * num_warps] : 0.0f;
#pragma unroll
for (int offset = warpSize/2; offset > 0; offset /= 2) {
val += __shfl_down_sync(0xFFFFFFFF, val, offset);
}
if (tid == 0) out_mean[d] = val;
}
float w2_sum = (tid < num_warps) ? s_buf[num_warps * D + tid] : 0.0f;
#pragma unroll
for (int offset = warpSize/2; offset > 0; offset /= 2) {
w2_sum += __shfl_down_sync(0xFFFFFFFF, w2_sum, offset);
}
if (tid == 0) *out_ess = 1.0f / w2_sum;
}
}
```
Wait, there are issues with shared memory sizing. Let me be more careful.
Shared memory needed:
- `s_logw[N]` for storing logw values
- `s_warp_max[num_warps]` for max reduction
- `s_warp_sum[num_warps]` for sum reduction
- `s_warp_mean[D * num_warps]` for mean reduction
- `s_warp_w2[num_warps]` for w2 reduction
Total: N + num_warps + num_warps + D*num_warps + num_warps floats
For N=4096, D=16, num_threads=1024, num_warps=32:
= 4096 + 32 + 32 + 16*32 + 32 = 4096 + 32 + 32 + 512 + 32 = 4704 floats = ~18.4KB
That's well within limits.
Actually, let me be smarter about shared memory reuse. The areas for max and sum reduction can overlap with mean reduction since they're used at different times.
Let me define:
- `s_logw[N]`
- `s_temp[max(num_warps, D * num_warps + num_warps)]`
Actually, to keep it simple and safe:
```cuda
extern __shared__ float shared[];
float* s_logw = shared;
float* s_scratch = shared + N; // used for various reductions
```
Size of s_scratch needs to be at least `max(num_warps * 2, D * num_warps + num_warps)`.
For D=16, num_warps=32: max(64, 512+32) = 544.
Total shared memory = (N + 544) * 4 bytes.
For N=4096: ~18.4KB.
For N=8192: ~34KB.
For N=16384: ~66KB (still under 48KB? No, 66KB > 48KB).
So this fused single-block kernel works well for N up to about 10000. For larger N, we'd need a different approach. But for the problem, N=4096.
Actually, I see that the maximum shared memory per block is 48KB by default (or up to 227KB on H100). Let me check: on H100, maximum shared memory per thread block is 227KB. So we can handle quite large N.
But actually, limiting to 48KB might allow more concurrent blocks. But with a single block, it doesn't matter.
OK, let me write the actual code now. I'll use a 2-kernel approach to be safer for arbitrary N, and because the shared memory gymnastics might not be worth it for marginal gains.
Actually wait, I should benchmark both approaches. A single block with shared memory for logw avoids writing logw to global memory. But with N=4096, writing logw to global memory is only 16KB. Memory bandwidth is 3.35 TB/s, so reading/writing 16KB takes about 5 nanoseconds. The kernel launch overhead is much larger (~5 microseconds).
So the advantage of a single kernel is: fewer kernel launches.
PyTorch reference does ~10-15 kernel launches.
My 2-kernel approach does 2 launches.
My 1-kernel approach does 1 launch.
2 vs 1 launch: difference is ~5 us. But the PyTorch reference might take ~100-200 us, so even 2 kernels should be 10x faster.
Let me just go with the 2-kernel approach since it's simpler and more robust.
Wait, but I want to maximize speedup. Let me try the 1-kernel fused approach.
Hold on, I need to be careful about how `load_inline` works. The `cuda_sources` should contain the CUDA kernels, and `functions` should list the C++ wrapper functions that call them.
Actually, let me re-read the instructions: "List the host-side wrapper function names in `functions=[...]` instead. Do not write your own `PYBIND11_MODULE(...)` block."
So I need C++ host wrapper functions that call the CUDA kernels. These wrapper functions are what pybind11 will bind to.
Here's the structure:
```cpp
// CUDA kernel
__global__ void my_kernel(...) { ... }
// Host wrapper
torch::Tensor my_wrapper(torch::Tensor input, ...) {
auto output = torch::empty(...);
my_kernel<<<grid, block>>>(...);
return output;
}
```
Then in Python: `functions=['my_wrapper']`
OK let me write it.
Actually, for maximum performance, let me reconsider fusing everything into one kernel. The inputs are:
- `samples` (N, D) = 256KB read
- Model parameters: negligible (a few hundred bytes)
- Output: mean (D) = 64 bytes, ess = 4 bytes
With one fused kernel:
- Read samples once
- Compute everything
- Write mean and ess
Total memory traffic: ~256KB read + ~68 bytes write = essentially nothing.
With two kernels:
- Kernel 1: read samples (256KB), read model params, write logw (16KB)
- Kernel 2: read samples (256KB), read logw (16KB), write mean + ess (68 bytes)
- Total: ~544KB read + ~16KB write
Still small, but double the memory traffic. For bandwidth-bound problems, this matters. But we're also doing compute (log_p, log_q for each sample).
Let me estimate:
- For each sample: ~3*16 sub/mul/div/mul/add/log = ~100 FLOPs for log_p, ~16*5 = ~80 FLOPs for log_q, plus exp and some reductions.
- Total per sample: ~200 FLOPs
- For 4096 samples: ~800k FLOPs
- At 67 TFLOPS, pure compute would take ~12 nanoseconds
- But memory bandwidth: reading 256KB at 3.35 TB/s = 76 nanoseconds
So this is memory bandwidth bound (but not by much; the compute and memory are roughly balanced when memory is read once).
With two kernels, memory traffic is ~544KB = 162 nanoseconds. Still very fast.
Kernel launch overhead: ~2 * 5 us = 10 us.
PyTorch probably takes ~100-200 us due to many allocations and kernel launches.
So both 1-kernel and 2-kernel should be ~10x faster than PyTorch. Let me go with 1-kernel for maximum performance.
Actually, there might be another issue: compiler optimization for the single kernel. The kernel might be large and complex, which could lead to register spilling. Let me think...
For the single kernel, each thread:
1. Reads D=16 floats from samples
2. Computes log_p (K=3 nested loops over D=16)
3. Computes log_q (loop over D=16)
4. Writes logw to shared memory
5. Participates in 4 reductions
6. Reads samples again (or keeps in registers? Can't keep 16*4=64 floats)
7. Computes weighted sums
8. Participates in D+1 reductions
Register usage might be high. Let me see... the compiler can keep `sample[16]` in registers during log_p/log_q computation, then reuse those registers later. Actually, sample values aren't needed after computing logw. So registers can be reused.
For the second pass, we need to re-read samples. That's fine, they'll be in L1 cache.
Let me try both approaches. First, the simpler 2-kernel approach, compile and test. If it's fast enough, great. If not, I'll try the 1-kernel approach.
Let me write the code.
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>
__global__ void compute_logw_kernel(
const float* __restrict__ samples,
const float* __restrict__ mix_mus,
const float* __restrict__ mix_stds,
const float* __restrict__ mix_weights,
const float* __restrict__ q_mu,
const float* __restrict__ q_std,
float* __restrict__ logw,
int N, int D, int K) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= N) return;
const float* sample = samples + idx * D;
// Compute log_p: logsumexp over mixture components
float max_log_comp = -1e30f;
for (int k = 0; k < K; k++) {
const float* mu = mix_mus + k * D;
const float* std = mix_stds + k * D;
float sum = 0.0f;
for (int d = 0; d < D; d++) {
float diff = sample[d] - mu[d];
float inv_std = 1.0f / std[d];
float scaled = diff * inv_std;
sum += -0.5f * scaled * scaled - logf(std[d]);
}
float log_comp = sum + logf(mix_weights[k]);
if (k == 0 || log_comp > max_log_comp) {
max_log_comp = log_comp;
}
}
float sum_exp = 0.0f;
for (int k = 0; k < K; k++) {
const float* mu = mix_mus + k * D;
const float* std = mix_stds + k * D;
float sum = 0.0f;
for (int d = 0; d < D; d++) {
float diff = sample[d] - mu[d];
float inv_std = 1.0f / std[d];
float scaled = diff * inv_std;
sum += -0.5f * scaled * scaled - logf(std[d]);
}
float log_comp = sum + logf(mix_weights[k]);
sum_exp += expf(log_comp - max_log_comp);
}
float log_p = max_log_comp + logf(sum_exp);
// Compute log_q
float log_q = 0.0f;
for (int d = 0; d < D; d++) {
float diff = sample[d] - q_mu[d];
float inv_std = 1.0f / q_std[d];
float scaled = diff * inv_std;
log_q += -0.5f * scaled * scaled - logf(q_std[d]);
}
logw[idx] = log_p - log_q;
}
```
Wait, I'm computing log_comp twice - that's wasteful! In the first loop I find the max, then in the second loop I compute the sum. Each loop does the same math. Let me restructure to avoid this.
Better approach:
```cuda
float log_comp[8];
for (int k = 0; k < K; k++) {
const float* mu = mix_mus + k * D;
const float* std = mix_stds + k * D;
float sum = 0.0f;
for (int d = 0; d < D; d++) {
float diff = sample[d] - mu[d];
float inv_std = 1.0f / std[d];
float scaled = diff * inv_std;
sum += -0.5f * scaled * scaled - logf(std[d]);
}
log_comp[k] = sum + logf(mix_weights[k]);
}
float max_val = log_comp[0];
for (int k = 1; k < K; k++) {
if (log_comp[k] > max_val) max_val = log_comp[k];
}
float sum_exp = 0.0f;
for (int k = 0; k < K; k++) {
sum_exp += expf(log_comp[k] - max_val);
}
float log_p = max_val + logf(sum_exp);
```
This is much better. `log_comp[8]` stays in registers. K=3, so only 3 registers.
Now for the second kernel:
```cuda
__global__ void compute_results_kernel(
const float* __restrict__ samples,
const float* __restrict__ logw,
float* __restrict__ out_mean,
float* __restrict__ out_ess,
int N, int D) {
const int tid = threadIdx.x;
const int num_threads = blockDim.x;
const int num_warps = (num_threads + 31) / 32;
extern __shared__ float shared[];
float* s_buf = shared;
// Phase 1: Find max(logw)
float local_max = -1e30f;
for (int i = tid; i < N; i += num_threads) {
if (logw[i] > local_max) local_max = logw[i];
}
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
local_max = fmaxf(local_max, __shfl_down_sync(0xFFFFFFFF, local_max, offset));
}
if (tid % 32 == 0) {
s_buf[tid / 32] = local_max;
}
__syncthreads();
if (tid < 32) {
local_max = (tid < num_warps) ? s_buf[tid] : -1e30f;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
local_max = fmaxf(local_max, __shfl_down_sync(0xFFFFFFFF, local_max, offset));
}
if (tid == 0) s_buf[0] = local_max;
}
__syncthreads();
float global_max = s_buf[0];
// Phase 2: Sum of exp(logw - global_max)
float local_sum = 0.0f;
for (int i = tid; i < N; i += num_threads) {
local_sum += expf(logw[i] - global_max);
}
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
local_sum += __shfl_down_sync(0xFFFFFFFF, local_sum, offset);
}
if (tid % 32 == 0) {
s_buf[tid / 32] = local_sum;
}
__syncthreads();
if (tid < 32) {
local_sum = (tid < num_warps) ? s_buf[tid] : 0.0f;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
local_sum += __shfl_down_sync(0xFFFFFFFF, local_sum, offset);
}
if (tid == 0) s_buf[1] = local_sum;
}
__syncthreads();
float total_sum = s_buf[1];
float lse = global_max + logf(total_sum);
// Phase 3: Compute local weighted_mean and local_w2
float local_mean[16];
#pragma unroll
for (int d = 0; d < 16; d++) {
local_mean[d] = 0.0f;
}
float local_w2 = 0.0f;
for (int i = tid; i < N; i += num_threads) {
float w = expf(logw[i] - lse);
local_w2 += w * w;
const float* sample = samples + i * D;
#pragma unroll
for (int d = 0; d < 16; d++) {
if (d < D) {
local_mean[d] += w * sample[d];
}
}
}
// Reduce each dimension
// s_buf layout: [num_warps entries for dim0, num_warps for dim1, ..., num_warps for w2]
for (int d = 0; d < D; d++) {
float val = local_mean[d];
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
val += __shfl_down_sync(0xFFFFFFFF, val, offset);
}
if (tid % 32 == 0) {
s_buf[(tid / 32) + d * num_warps] = val;
}
}
float w2_val = local_w2;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
w2_val += __shfl_down_sync(0xFFFFFFFF, w2_val, offset);
}
if (tid % 32 == 0) {
s_buf[D * num_warps + tid / 32] = w2_val;
}
__syncthreads();
if (tid < 32) {
for (int d = 0; d < D; d++) {
float val = (tid < num_warps) ? s_buf[tid + d * num_warps] : 0.0f;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
val += __shfl_down_sync(0xFFFFFFFF, val, offset);
}
if (tid == 0) {
out_mean[d] = val;
}
}
float w2_sum = (tid < num_warps) ? s_buf[D * num_warps + tid] : 0.0f;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
w2_sum += __shfl_down_sync(0xFFFFFFFF, w2_sum, offset);
}
if (tid == 0) {
*out_ess = 1.0f / w2_sum;
}
}
}
```
Wait, I used a fixed-size `local_mean[16]` but the loop `for (int d = 0; d < D; d++)` might access out of bounds if D > 16. To be safe, I should make it dynamic or ensure D <= 16. But for this problem D=16.
Actually, the issue is that `local_mean[16]` is declared but if D < 16, only some elements are used. And the `#pragma unroll for (int d = 0; d < 16; d++)` loop initializes all 16. This is fine.
Wait, but what if D > 16? Then we'd have buffer overflow. But the problem uses D=16, and we can assert or handle it. For safety, let me just use a dynamic approach or a large enough fixed size.
Actually, using a template or dynamic allocation would be better. But for simplicity, let me use a fixed size of 64 or something and add a check. Or I can just use a dynamic approach:
```cuda
float local_mean[64]; // Support up to D=64
for (int d = 0; d < D; d++) {
local_mean[d] = 0.0f;
}
```
But this might use a lot of registers if not optimized. For D=16, it's fine.
Actually, I can do better. Let me not overthink this and just write a clean version.
One more thing: I should mark pointers with `__restrict__` for better optimization.
Let me also think about using `float` vs `double`. The reference uses `torch` defaults which for this problem is fp32 (atol=1e-4). So float is fine.
Now for the host wrappers:
```cpp
std::vector<torch::Tensor> importance_sampling_forward(
torch::Tensor samples,
torch::Tensor mix_mus,
torch::Tensor mix_stds,
torch::Tensor mix_weights,
torch::Tensor q_mu,
torch::Tensor q_std) {
int N = samples.size(0);
int D = samples.size(1);
int K = mix_mus.size(0);
auto logw = torch::empty({N}, samples.options());
auto out_mean = torch::empty({D}, samples.options());
auto out_ess = torch::empty({}, samples.options());
const int threads = 256;
const int blocks = (N + threads - 1) / threads;
compute_logw_kernel<<<blocks, threads>>>(
samples.data_ptr<float>(),
mix_mus.data_ptr<float>(),
mix_stds.data_ptr<float>(),
mix_weights.data_ptr<float>(),
q_mu.data_ptr<float>(),
q_std.data_ptr<float>(),
logw.data_ptr<float>(),
N, D, K);
// Second kernel: single block
const int reduction_threads = 1024;
const int smem_size = (reduction_threads / 32) * (D + 1) * sizeof(float) + 2 * sizeof(float);
compute_results_kernel<<<1, reduction_threads, smem_size>>>(
samples.data_ptr<float>(),
logw.data_ptr<float>(),
out_mean.data_ptr<float>(),
out_ess.data_ptr<float>(),
N, D);
return {out_mean, out_ess};
}
```
Wait, I need to make sure the shared memory calculation is correct. In the kernel, I use `s_buf` for:
1. Warp max values: `num_warps` floats
2. Then warp sum values: `num_warps` floats
3. Then per-dimension warp sums: `D * num_warps` floats
4. Then w2 warp sum: `num_warps` floats
The maximum need at any time is `max(num_warps, D * num_warps + num_warps)`.
For D=16, num_warps=32: max(32, 16*32 + 32) = max(32, 544) = 544.
So `smem_size = 544 * sizeof(float)` should suffice. Actually, during phase 1 and 2, we only need `num_warps` floats. During phase 4, we need `D * num_warps + num_warps` floats. So the max is `D * num_warps + num_warps`.
But wait, in phase 4, after writing per-dimension and w2 sums, we do the final reduction. This reuses `s_buf`. Let me trace through exactly:
Phase 1:
- Write `num_warps` max values to `s_buf[0..num_warps-1]`
- Reduce to `s_buf[0]`
Phase 2:
- Write `num_warps` sum values to `s_buf[0..num_warps-1]`
- Reduce to `s_buf[1]` (I wrote `s_buf[1]` but that's overlapping with the first element!)
Wait, that's a bug! Let me re-read:
```cuda
if (tid == 0) s_buf[1] = local_sum;
```
But during the block reduction in phase 2:
```cuda
if (tid < 32) {
local_sum = (tid < num_warps) ? s_buf[tid] : 0.0f;
// reduce
if (tid == 0) s_buf[1] = local_sum;
}
```
This is fine because after `local_sum = (tid < num_warps) ? s_buf[tid]`, all threads in the warp load from `s_buf`. Then during the shuffle reduction, they don't access shared memory. Only at the end, tid==0 writes to `s_buf[1]`. No other thread accesses `s_buf` at that moment.
But wait, `s_buf[0]` holds `global_max` from phase 1. Writing `s_buf[1]` doesn't overwrite it. Good.
Phase 3:
- Read `s_buf[1]` as `total_sum`
- Compute `lse`
Phase 4:
- Each warp writes `D` reduced values starting at `s_buf[warp_id + d * num_warps]`
- Wait, the layout is `(tid / 32) + d * num_warps` per dimension
- For warp 0, dim 0: writes to `s_buf[0]`
- For warp 0, dim 1: writes to `s_buf[32]`
- etc.
- For w2, warp 0 writes to `s_buf[D * num_warps + 0]`
Then in the block reduction:
```cuda
float val = (tid < num_warps) ? s_buf[tid + d * num_warps] : 0.0f;
```
This reads from `s_buf` which was just written by the warp leaders. There's a `__syncthreads()` before this, so it's fine.
After the block reduction writes `out_mean[d]`, it doesn't need `s_buf` anymore.
So the shared memory footprint is: `max(num_warps, D * num_warps + num_warps)`.
For D=16, num_warps=32: max(32, 544) = 544 floats.
OK, let me fix the code to be cleaner.
Let me also consider: for the `log_comp` arrays, I used `float log_comp[8]`. For K=3 this is fine. But what if larger K is used? The constructor says `num_components=3` by default. Let me make it general by using a fixed-size array with a maximum. Or I can just dynamically compute without storing `log_comp`. Actually, I showed a way to store log_comp, which requires K registers. Let me use a fixed max size.
Wait, on second thought, let me keep it simple and just restructure the log_p computation to compute everything in one pass without storing per-component values, then a second pass for the sum. Since K is small (K=3), doing two passes over the component computations is not a big deal, and it avoids stack arrays.
Or, even simpler, using the stored array approach is cleaner:
```cuda
float log_comp[8]; // max 8 components
#pragma unroll
for (int k = 0; k < K; k++) {
// compute log_comp[k]
}
float max_val = log_comp[0];
#pragma unroll
for (int k = 1; k < K; k++) {
max_val = fmaxf(max_val, log_comp[k]);
}
float sum_exp = 0.0f;
#pragma unroll
for (int k = 0; k < K; k++) {
sum_exp += expf(log_comp[k] - max_val);
}
```
For K=3, `#pragma unroll` might not work since K is a runtime variable. But the compiler will unroll if K is a compile-time constant... which it's not. So the loops won't be unrolled. That's fine for K=3.
Actually, let me check if I can pass K as a template parameter. No, because `load_inline` with pybind11 auto-generated wrappers has fixed signatures. The wrapper function takes tensors, and the kernel is launched with runtime parameters.
OK, runtime K it is.
One more optimization: I can precompute `log(mix_weights[k])`, `log(std[d])`, and `1/std[d]` on the CPU and pass them as tensors. But the reference model already has `mix_weights`, `mix_stds`, `q_std` as buffers. Computing `log()` and `1/x` once per kernel vs once on CPU... well, in the PyTorch reference, `mix_stds.log()` is computed in `_log_p` every forward call. So precomputing these on CPU in `ModelNew.__init__` would be a win.
Let me do that! In `ModelNew.__init__`, I'll create:
- `log_mix_weights = log(mix_weights)`
- `log_mix_stds = log(mix_stds)`
- `inv_mix_stds = 1.0 / mix_stds`
- `log_q_std = log(q_std)`
- `inv_q_std = 1.0 / q_std`
Then pass these to the kernel. This saves log() and division operations in the kernel.
Wait, but `mix_stds` is `torch.ones(num_components, dim)` in the reference. So `log(mix_stds) = 0` and `1/mix_stds = 1`. But to be numerically equivalent and general, let me precompute.
Actually, looking at the reference more carefully:
```python
log_comp = -0.5 * (diff / self.mix_stds.unsqueeze(0)) ** 2 - self.mix_stds.log().unsqueeze(0)
```
`mix_stds.log().unsqueeze(0)` is computed every forward pass! In PyTorch, `self.mix_stds.log()` should be computed once and cached as a buffer attribute, but in this code it's computed inline. That's fine.
For the CUDA kernel, precomputing `log_stds` and `inv_stds` on the CPU makes sense.
Let me design the final code:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
CUDA_SOURCE = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>
__global__ void compute_logw_kernel(
const float* __restrict__ samples,
const float* __restrict__ mix_mus,
const float* __restrict__ mix_inv_stds,
const float* __restrict__ mix_log_stds,
const float* __restrict__ mix_log_weights,
const float* __restrict__ q_mu,
const float* __restrict__ q_inv_std,
const float* __restrict__ q_log_std,
float* __restrict__ logw,
int N, int D, int K) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= N) return;
const float* sample = samples + idx * D;
float log_comp[8];
#pragma unroll
for (int k = 0; k < K; k++) {
const float* mu = mix_mus + k * D;
const float* inv_std = mix_inv_stds + k * D;
const float* log_std = mix_log_stds + k * D;
float sum = 0.0f;
for (int d = 0; d < D; d++) {
float diff = sample[d] - mu[d];
float scaled = diff * inv_std[d];
sum += -0.5f * scaled * scaled - log_std[d];
}
log_comp[k] = sum + mix_log_weights[k];
}
float max_val = log_comp[0];
#pragma unroll
for (int k = 1; k < K; k++) {
max_val = fmaxf(max_val, log_comp[k]);
}
float sum_exp = 0.0f;
#pragma unroll
for (int k = 0; k < K; k++) {
sum_exp += expf(log_comp[k] - max_val);
}
float log_p = max_val + logf(sum_exp);
float log_q = 0.0f;
for (int d = 0; d < D; d++) {
float diff = sample[d] - q_mu[d];
float scaled = diff * q_inv_std[d];
log_q += -0.5f * scaled * scaled - q_log_std[d];
}
logw[idx] = log_p - log_q;
}
"""
```
OK wait, I need to actually write the complete code. Let me focus and write the final version.
Actually, I just realized that I can potentially avoid the temporary `logw` buffer entirely by doing everything in one kernel. Let me reconsider.
For a single kernel with one block:
```cuda
__global__ void fused_kernel(...) {
// Single block
int tid = threadIdx.x;
int num_threads = blockDim.x;
extern __shared__ float smem[];
float* s_logw = smem;
// Compute logw for each sample
for (int i = tid; i < N; i += num_threads) {
// ... compute logw[i] ...
s_logw[i] = log_p - log_q;
}
__syncthreads();
// Block reduce for max
// ...
// Block reduce for sum
// ...
// Compute weighted mean and ess
// ...
}
```
The shared memory for logw is N floats.
For N=4096, that's 16KB. Plus scratch space for reductions (~1KB). Total ~17KB. Well within limits.
This avoids the global memory write for logw entirely. For N=4096, logw buffer is 16KB. Reading and writing it costs ~32KB of memory traffic, which at 3.35 TB/s is ~10 nanoseconds. The kernel launch saved is ~5 microseconds. So the 1-kernel approach is ~5us faster.
Given that PyTorch probably takes 100-200us, both approaches will be 10-50x faster. But for maximum performance, let me do the 1-kernel approach.
However, the 1-kernel approach has a limitation: it only supports N that fits in shared memory. For N=4096, it's fine. For larger N, we'd need tiling or a different approach. But the problem parameters are fixed.
Let me write the single-kernel version.
Actually, there's another consideration: I need to re-read `samples` for the weighted mean computation. With a single block, the L1 cache might still have the samples from the first pass. But with 4096 samples * 16 floats = 262KB, and L1 cache is 128KB per SM on H100, some of it might have spilled to L2. Still, L2 is fast.
The 2-kernel approach re-reads samples from global memory into the second kernel. L2 cache helps here.
I'm overthinking this. Let me just write a clean 2-kernel version first, compile, test, and see the timing. If it's not fast enough, I'll optimize further.
Let me write the code now.
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
CUDA_SOURCE = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>
__global__ void compute_logw_kernel(
const float* __restrict__ samples,
const float* __restrict__ mix_mus,
const float* __restrict__ mix_inv_stds,
const float* __restrict__ mix_log_stds,
const float* __restrict__ mix_log_weights,
const float* __restrict__ q_mu,
const float* __restrict__ q_inv_std,
const float* __restrict__ q_log_std,
float* __restrict__ logw,
int N, int D, int K) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= N) return;
const float* sample = samples + idx * D;
float log_comp[8];
for (int k = 0; k < K; k++) {
const float* mu = mix_mus + k * D;
const float* inv_std = mix_inv_stds + k * D;
const float* log_std = mix_log_stds + k * D;
float sum = 0.0f;
for (int d = 0; d < D; d++) {
float diff = sample[d] - mu[d];
float scaled = diff * inv_std[d];
sum += -0.5f * scaled * scaled - log_std[d];
}
log_comp[k] = sum + mix_log_weights[k];
}
float max_val = log_comp[0];
for (int k = 1; k < K; k++) {
max_val = fmaxf(max_val, log_comp[k]);
}
float sum_exp = 0.0f;
for (int k = 0; k < K; k++) {
sum_exp += expf(log_comp[k] - max_val);
}
float log_p = max_val + logf(sum_exp);
float log_q = 0.0f;
for (int d = 0; d < D; d++) {
float diff = sample[d] - q_mu[d];
float scaled = diff * q_inv_std[d];
log_q += -0.5f * scaled * scaled - q_log_std[d];
}
logw[idx] = log_p - log_q;
}
__global__ void compute_results_kernel(
const float* __restrict__ samples,
const float* __restrict__ logw,
float* __restrict__ out_mean,
float* __restrict__ out_ess,
int N, int D) {
const int tid = threadIdx.x;
const int num_threads = blockDim.x;
const int num_warps = (num_threads + 31) / 32;
extern __shared__ float shared[];
float* s_buf = shared;
// Phase 1: Find global max of logw
float local_max = -1e30f;
for (int i = tid; i < N; i += num_threads) {
if (logw[i] > local_max) local_max = logw[i];
}
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
local_max = fmaxf(local_max, __shfl_down_sync(0xFFFFFFFF, local_max, offset));
}
if (tid % 32 == 0) {
s_buf[tid / 32] = local_max;
}
__syncthreads();
if (tid < 32) {
local_max = (tid < num_warps) ? s_buf[tid] : -1e30f;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
local_max = fmaxf(local_max, __shfl_down_sync(0xFFFFFFFF, local_max, offset));
}
if (tid == 0) s_buf[0] = local_max;
}
__syncthreads();
float global_max = s_buf[0];
// Phase 2: Sum of exp(logw - global_max)
float local_sum = 0.0f;
for (int i = tid; i < N; i += num_threads) {
local_sum += expf(logw[i] - global_max);
}
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
local_sum += __shfl_down_sync(0xFFFFFFFF, local_sum, offset);
}
if (tid % 32 == 0) {
s_buf[tid / 32] = local_sum;
}
__syncthreads();
if (tid < 32) {
local_sum = (tid < num_warps) ? s_buf[tid] : 0.0f;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
local_sum += __shfl_down_sync(0xFFFFFFFF, local_sum, offset);
}
if (tid == 0) s_buf[1] = local_sum;
}
__syncthreads();
float total_sum = s_buf[1];
float lse = global_max + logf(total_sum);
// Phase 3: Compute local weighted_mean and local_w2
const int MAX_DIM = 64;
float local_mean[MAX_DIM];
for (int d = 0; d < D; d++) {
local_mean[d] = 0.0f;
}
float local_w2 = 0.0f;
for (int i = tid; i < N; i += num_threads) {
float w = expf(logw[i] - lse);
local_w2 += w * w;
const float* sample = samples + i * D;
for (int d = 0; d < D; d++) {
local_mean[d] += w * sample[d];
}
}
// Phase 4: Block reduction for each dimension
for (int d = 0; d < D; d++) {
float val = local_mean[d];
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
val += __shfl_down_sync(0xFFFFFFFF, val, offset);
}
if (tid % 32 == 0) {
s_buf[(tid / 32) + d * num_warps] = val;
}
}
float w2_val = local_w2;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
w2_val += __shfl_down_sync(0xFFFFFFFF, w2_val, offset);
}
if (tid % 32 == 0) {
s_buf[D * num_warps + tid / 32] = w2_val;
}
__syncthreads();
if (tid < 32) {
for (int d = 0; d < D; d++) {
float val = (tid < num_warps) ? s_buf[tid + d * num_warps] : 0.0f;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
val += __shfl_down_sync(0xFFFFFFFF, val, offset);
}
if (tid == 0) {
out_mean[d] = val;
}
}
float w2_sum = (tid < num_warps) ? s_buf[D * num_warps + tid] : 0.0f;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
w2_sum += __shfl_down_sync(0xFFFFFFFF, w2_sum, offset);
}
if (tid == 0) {
*out_ess = 1.0f / w2_sum;
}
}
}
std::vector<torch::Tensor> importance_sampling_forward(
torch::Tensor samples,
torch::Tensor mix_mus,
torch::Tensor mix_inv_stds,
torch::Tensor mix_log_stds,
torch::Tensor mix_log_weights,
torch::Tensor q_mu,
torch::Tensor q_inv_std,
torch::Tensor q_log_std) {
int N = samples.size(0);
int D = samples.size(1);
int K = mix_mus.size(0);
auto options = samples.options();
auto logw = torch::empty({N}, options);
auto out_mean = torch::empty({D}, options);
auto out_ess = torch::empty({}, options);
const int threads = 256;
const int blocks = (N + threads - 1) / threads;
compute_logw_kernel<<<blocks, threads>>>(
samples.data_ptr<float>(),
mix_mus.data_ptr<float>(),
mix_inv_stds.data_ptr<float>(),
mix_log_stds.data_ptr<float>(),
mix_log_weights.data_ptr<float>(),
q_mu.data_ptr<float>(),
q_inv_std.data_ptr<float>(),
q_log_std.data_ptr<float>(),
logw.data_ptr<float>(),
N, D, K);
const int reduction_threads = 1024;
int num_warps = (reduction_threads + 31) / 32;
int smem_size = (D * num_warps + num_warps + 2) * sizeof(float);
compute_results_kernel<<<1, reduction_threads, smem_size>>>(
samples.data_ptr<float>(),
logw.data_ptr<float>(),
out_mean.data_ptr<float>(),
out_ess.data_ptr<float>(),
N, D);
return {out_mean, out_ess};
}
"""
class ModelNew(nn.Module):
def __init__(self, dim, num_components=3):
super().__init__()
self.dim = dim
self.num_components = num_components
self.register_buffer("mix_weights", torch.ones(num_components) / num_components)
mus = torch.randn(num_components, dim) * 2.0
self.register_buffer("mix_mus", mus)
self.register_buffer("mix_stds", torch.ones(num_components, dim))
self.register_buffer("q_mu", torch.zeros(dim))
self.register_buffer("q_std", torch.ones(dim) * 3.0)
# Precompute constants for the kernel
self.register_buffer("mix_inv_stds", 1.0 / self.mix_stds)
self.register_buffer("mix_log_stds", torch.log(self.mix_stds))
self.register_buffer("mix_log_weights", torch.log(self.mix_weights))
self.register_buffer("q_inv_std", 1.0 / self.q_std)
self.register_buffer("q_log_std", torch.log(self.q_std))
self.cuda_module = load_inline(
name="importance_sampling_cuda",
cpp_sources="",
cuda_sources=CUDA_SOURCE,
functions=["importance_sampling_forward"],
extra_cuda_cflags=["-O3"],
verbose=False
)
def forward(self, samples):
return self.cuda_module.importance_sampling_forward(
samples,
self.mix_mus,
self.mix_inv_stds,
self.mix_log_stds,
self.mix_log_weights,
self.q_mu,
self.q_inv_std,
self.q_log_std
)
```
Wait, I need to make sure `load_inline` is called at module load time, not in `__init__`. The instructions say: "Use `torch.utils.cpp_extension.load_inline` to compile and bind CUDA source at module load time, then call the compiled extension from `ModelNew.forward`."
So I should compile at module level. But the compiled extension doesn't depend on any instance variables. So I can compile it once at module level.
However, there's a catch: `load_inline` compiles the CUDA code. If I do it at module level, it's fine. But I need to make sure it only runs once. Let me use a module-level variable.
Actually, looking again at the reference code, `get_inputs()` and `get_init_inputs()` are used by the evaluator to create the model and inputs. The model is instantiated once per test. But CUDA compilation is expensive. Doing it at module level is better.
Let me restructure:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
CUDA_SOURCE = """
...
"""
cuda_module = load_inline(
name="importance_sampling_cuda",
cpp_sources="",
cuda_sources=CUDA_SOURCE,
functions=["importance_sampling_forward"],
extra_cuda_cflags=["-O3"],
verbose=False
)
class ModelNew(nn.Module):
def __init__(self, dim, num_components=3):
...
def forward(self, samples):
return cuda_module.importance_sampling_forward(...)
```
This compiles once at import time.
But wait, `load_inline` might fail if CUDA is not available at import time. But the problem states this runs on GPU, so CUDA should be available.
One concern: if there are multiple `ModelNew` instances, they all share the same `cuda_module`. That's fine since the kernels are stateless.
Actually, I need to be careful about `name="importance_sampling_cuda"`. If the module is already loaded, `load_inline` might use the cached version. That's good for repeated runs.
Let me also think about numerical precision. The reference uses `torch.logsumexp` which is numerically stable. My kernel uses the standard max + log(sum(exp(x - max))) approach, which is equally stable.
One more thing: I should make sure `out_ess` is a scalar (0-d tensor). In the reference, `ess = 1.0 / (w ** 2).sum()` returns a Python float. But PyTorch models typically return tensors, and `1.0 / tensor.sum()` gives a scalar tensor.
Wait, the reference returns `ess` which is a Python float because of `1.0 / ...`. Let me check:
```python
ess = 1.0 / (w ** 2).sum()
```
`w` is a tensor, `(w ** 2).sum()` is a tensor. `1.0 / tensor` gives a tensor. Wait, no:
Actually, in PyTorch:
```python
>>> import torch
>>> w = torch.randn(4)
>>> type(1.0 / (w**2).sum())
<class 'torch.Tensor'>
```
Yes, `1.0 / tensor` produces a tensor. So `ess` is a scalar tensor. My kernel outputs `out_ess` as a 0-d tensor, which is correct.
But wait, in the reference, `ess` could potentially be on CPU if `w` operations happen that way... no, everything should be on the same device as the inputs. Since `samples` is passed and model buffers are on the same device, everything stays on that device.
OK, let me finalize the code.
One issue: `load_inline` with an empty `cpp_sources` might not work properly. Actually, looking at PyTorch docs, `load_inline` takes `cuda_sources` and optionally `cpp_sources`. Let me check...
Actually, `load_inline` parameter order might differ in different PyTorch versions. Let me use keyword arguments to be safe.
In some PyTorch versions it's:
```python
load_inline(name, cpp_sources, cuda_sources, functions, extra_cuda_cflags, verbose)
```
Let me use this signature.
Also, I should add `--use_fast_math` to the CUDA flags for extra speed. But this might reduce precision. Let me try without it first, and add it if needed.
Wait, using `--use_fast_math` trades precision for speed. For fp32 with atol=1e-4, it might still be fine, but `__expf` and `__logf` can have errors of a few ULPs. Let's keep `-O3` without `--use_fast_math` initially.
Another optimization: I can specialize for the known problem sizes. The problem says `dim=16`, `num_samples=4096`, `num_components=3`. I can pass these as template parameters to avoid dynamic loops, but that requires compiling multiple versions or hardcoding.
A simpler approach: I noticed that `mix_stds` and `q_std` are all ones (or fixed values). Let me check:
- `self.mix_stds` = `torch.ones(num_components, dim)`
- `self.q_std` = `torch.ones(dim) * 3.0`
Wait, `mix_stds` is ones! And `q_std` is `3.0 * ones`. So:
- `mix_inv_stds = 1.0 / 1.0 = 1.0`
- `mix_log_stds = log(1.0) = 0.0`
- `q_inv_std = 1.0 / 3.0`
- `q_log_std = log(3.0)`
This means in the kernel, for `mix_stds`, we can simplify:
```cuda
sum += -0.5f * diff * diff; // since inv_std = 1, log_std = 0
```
And for `q_std`:
```cuda
log_q += -0.5f * scaled * scaled - log_3; // where scaled = diff / 3.0
```
But this is specific to the problem setup. To be general (since the constructor could theoretically be called with different values, though the evaluation probably uses defaults), I should keep the general form. The compiler will optimize away the multiplications by 1.0 and additions of 0.0 anyway.
Actually, looking at the constructor: the buffers are fixed during `__init__`. So if we want maximum speed, we could detect that `mix_stds` is all ones and optimize. But that's overkill. The compiler should handle it.
Let me now write the complete, clean code and submit for compilation.
I realize I should also check: the `functions` argument to `load_inline` should be a list of function names to expose. The C++ wrapper function `importance_sampling_forward` returns `std::vector<torch::Tensor>`. PyTorch's pybind11 binding should handle this correctly.
One more thing: I need to make sure the CUDA code compiles with the C++ compiler. Let me double-check for any issues.
Potential issues:
1. `__shfl_down_sync` requires CUDA >= 9.0. H100 has CUDA 12+, so it's fine.
2. `extern __shared__ float shared[];` is correct syntax.
3. Using `s_buf[1]` after `s_buf[0]` is fine as long as num_warps >= 2, which it is for blockDim >= 64. I'll use 1024 threads, so 32 warps. Fine.
4. I need to make sure D * num_warps doesn't overflow shared memory. For D=64 and num_warps=32: 2048 + 32 + 2 = 2082 floats ≈ 8KB. Fine.
Wait, I want to change how s_buf is used for clarity. Let me trace through again to make sure there are no races:
Phase 1:
- Threads write to `s_buf[warp_id]` after warp-level reduction
- `__syncthreads()`
- Threads in first warp read from `s_buf[tid]` and reduce
- `tid == 0` writes to `s_buf[0]`
- `__syncthreads()`
- All threads read `s_buf[0]` as `global_max`
Phase 2:
- Threads write to `s_buf[warp_id]` after warp-level reduction
- `__syncthreads()`
- Threads in first warp read from `s_buf[tid]` and reduce
- `tid == 0` writes to `s_buf[1]`
- `__syncthreads()`
- All threads read `s_buf[1]`
Phase 3:
- Each thread computes `local_mean` and `local_w2`
- Warp-level reduction for each dimension
- Warp leader writes to `s_buf[warp_id + d * num_warps]`
- Warp leader writes w2 to `s_buf[D * num_warps + warp_id]`
- `__syncthreads()`
- First warp reads and reduces each dimension
- `tid == 0` writes to `out_mean[d]`
- `tid == 0` writes to `out_ess`
Potential issue: after phase 2, `__syncthreads()`, then phase 3 starts. In phase 3, warp leaders from different warps write to `s_buf` simultaneously but to different locations. This is fine.
Then another `__syncthreads()`, and the first warp reads. Fine.
Wait, in phase 3, the first warp also participates in the local computation and warp-level reduction. When `tid % 32 == 0`, those threads (warp leaders) write to `s_buf`. Specifically, warp 0's leader (tid=0) writes to:
- `s_buf[0 + d * num_warps]` for each d
- `s_buf[D * num_warps + 0]` for w2
Then all warps hit `__syncthreads()`.
Then `if (tid < 32)`, meaning the first warp reads:
- `s_buf[tid + d * num_warps]` for each d
- `s_buf[D * num_warps + tid]` for w2
For warp 0, `tid` ranges from 0 to 31, but only `tid < num_warps` (32) are valid. So it reads all the values written by the warp leaders. Good.
One subtle issue: during phase 1, `s_buf[0]` is written. During phase 2, `s_buf[warp_id]` is written again. For warp 0, `warp_id=0`, so `s_buf[0]` is overwritten. But this is after phase 1 is completely done and all threads have read `global_max`. Good.
Another subtle issue: during phase 4, when `tid == 0` writes `out_mean[d]` and `out_ess`, this happens inside the `if (tid < 32)` block. But `tid` only goes 0..31, and only `tid == 0` writes. No conflicts.
I think it's correct.
Let me also think about whether I should use `.contiguous()` on the inputs. The reference model's buffers are contiguous by default (created with `torch.ones`, `torch.zeros`, `torch.randn`). The samples input from `get_inputs()` is `torch.randn(num_samples, dim) * q_std`, which is contiguous. So I don't need to explicitly make things contiguous, but it doesn't hurt.
Actually, I should check if `samples` might not be contiguous. In general PyTorch usage, inputs might not be contiguous. But the `get_inputs()` returns contiguous tensors. For safety, I could ensure it.
Now, for the shared memory size in the host wrapper:
```cpp
int smem_size = (D * num_warps + num_warps + 2) * sizeof(float);
```
But wait, in the kernel during phase 1 and 2, I only use `s_buf[0]` to `s_buf[num_warps]`. But in phase 4, I use up to `s_buf[D * num_warps + num_warps]`. So the total needed is `D * num_warps + num_warps` plus a couple of extra.
Actually, the maximum index used is:
- Phase 1: `s_buf[num_warps - 1]`
- Phase 2: `s_buf[1]` (but also using `s_buf[tid]` where tid < 32, so up to `s_buf[31]`)
- Phase 4: `s_buf[D * num_warps + num_warps - 1]`
So the maximum index is `D * num_warps + num_warps - 1`.
With D=16, num_warps=32: max index = 16*32 + 32 - 1 = 543.
We need 544 floats.
Size = 544 * sizeof(float) = 2176 bytes.
That seems small. Let me double check: `D * num_warps + num_warps` = 16 * 32 + 32 = 544. Yes.
Actually, I want to be more precise. In phase 4:
- For each dimension d, warp leader writes to `s_buf[warp_id + d * num_warps]`
- The max index for dimension data is when d = D-1 and warp_id = num_warps-1: `(D-1) * num_warps + (num_warps - 1) = D * num_warps - 1`
- For w2, warp leader writes to `s_buf[D * num_warps + warp_id]`
- Max index: `D * num_warps + num_warps - 1`
Total elements needed: `D * num_warps + num_warps`.
In my host code I had `+ 2` which is fine (over-allocating).
Wait, but during phase 1 and 2, I also write `s_buf[0]` and `s_buf[1]` which are within the first `num_warps` elements. Since `num_warps <= D * num_warps + num_warps` for D >= 1, the total allocation is sufficient.
Actually, one more thing: I had `s_buf[1] = local_sum` in phase 2. But `local_sum` is only written by `tid == 0`. And `s_buf[1]` might also be used as a warp buffer during phase 1 (`s_buf[warp_id]` where warp_id can be 1 for the second warp). But phase 1 is completely done before phase 2, and all threads have synchronized. So overwriting `s_buf[1]` is fine.
OK, I think the code is correct. Let me write the final version.
Actually, let me think about whether I should avoid the temporary `logw` allocation. I'll write both versions: first the 2-kernel version, and if it's not fast enough, I'll try the 1-kernel fused version.
Let me write the 2-kernel version now.
Wait, one more optimization: for the first kernel, instead of `int idx = blockIdx.x * blockDim.x + threadIdx.x`, I could use a grid-stride loop to handle larger N and improve occupancy. But N=4096 with 256-thread blocks gives 16 blocks. On an H100 with many SMs, 16 blocks might not fully utilize the GPU. However, this kernel is memory-bandwidth-bound, and each block processes 256 samples. H100 has 132 SMs, each can run multiple blocks. With only 16 blocks, some SMs will be idle, but the memory bandwidth should still be saturated.
For better occupancy, I could use smaller blocks (like 128 threads) giving 32 blocks, or even 64 threads giving 64 blocks. 64 blocks would spread better across 132 SMs. But the overhead of launching more blocks might not matter.
Actually, let me think about this more carefully. Each thread processes one sample. A sample has D=16 floats = 64 bytes. Each thread reads 64 bytes and does ~200 FLOPs. The arithmetic intensity is ~3 FLOPs/byte. At 3.35 TB/s bandwidth, we need ~10 TFLOPs to be compute-bound. H100 does 67 TFLOPS in FP32. So we're memory bandwidth bound.
For bandwidth-bound kernels, we want to maximize memory throughput. Using enough threads to saturate memory bandwidth is key. 4096 threads reading 64 bytes each = 256KB total. At 3.35 TB/s, this takes ~76 ns. A block of 256 threads reading coalesced memory can achieve near-peak bandwidth. 16 blocks of 256 threads = 4096 threads total.
Actually, 16 blocks might be a bit low for H100. H100 has 132 SMs, and each SM can run up to 32 blocks. But the real limit is often threads per SM (2048). So each SM can run up to 8 blocks of 256 threads, or 2048 threads total.
With 16 blocks, they can be scheduled on 16 SMs (1 block each) or fewer SMs (multiple blocks each). Either way, we're only using a fraction of the SMs. But for a problem that takes ~76ns of memory time, the kernel is likely to be dominated by launch overhead anyway.
Let me not overthink this and just use 256 threads.
Actually, I want to point out: when I compile the kernel, the compiler will tell me if there are issues. Let me write the code and compile it.
One last check: I'm using `s_buf[1]` in the kernel, but `s_buf` is a `float*`. I write `if (tid == 0) s_buf[1] = local_sum;`. In phase 1, I wrote `if (tid == 0) s_buf[0] = local_max;`. So after phase 1, `s_buf[0] = global_max`. Then in phase 2, I write `s_buf[1] = total_sum`. But in between, during the block-level reduction of phase 2:
```cuda
if (tid < 32) {
local_sum = (tid < num_warps) ? s_buf[tid] : 0.0f;
// ...
if (tid == 0) s_buf[1] = local_sum;
}
```
Wait, `s_buf[tid]` for tid=0 reads `s_buf[0]`. But `s_buf[0]` was set to `global_max` in phase 1! Then during the warp-level shuffle, all threads get `local_sum`. Then tid==0 writes to `s_buf[1]`. That's fine.
But wait, I read `s_buf[tid]` where `tid` can be 0 in the first warp. `s_buf[0]` still contains `global_max` from phase 1. But that's OK because all threads in the first warp load their own `s_buf[warp_id]` value. Warp 0 loads `s_buf[0]` which was written by warp 0's leader in the warp reduction step of phase 2... WAIT!
Let me re-read phase 2 more carefully:
```cuda
// Phase 2: Sum of exp(logw - global_max)
float local_sum = 0.0f;
for (int i = tid; i < N; i += num_threads) {
local_sum += expf(logw[i] - global_max);
}
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
local_sum += __shfl_down_sync(0xFFFFFFFF, local_sum, offset);
}
if (tid % 32 == 0) {
s_buf[tid / 32] = local_sum;
}
__syncthreads();
if (tid < 32) {
local_sum = (tid < num_warps) ? s_buf[tid] : 0.0f;
// ...
if (tid == 0) s_buf[1] = local_sum;
}
```
Ah! The warp leader writes `s_buf[warp_id] = local_sum` in phase 2. For warp 0, this writes to `s_buf[0]`, overwriting the `global_max` that was there from phase 1! Then the first warp reads `s_buf[tid]` for tid < num_warps. For tid=0, it reads `s_buf[0]`, now containing warp 0's sum. Good!
Then after the block-level reduction, `tid == 0` writes to `s_buf[1]`.
Then `__syncthreads()`.
Then `total_sum = s_buf[1]`.
This is correct. The overwrite of `s_buf[0]` happens before the read of `s_buf[0]` by the first warp, and there's a `__syncthreads()` in between. Good.
Phase 4:
```cuda
for (int d = 0; d < D; d++) {
float val = local_mean[d];
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
val += __shfl_down_sync(0xFFFFFFFF, val, offset);
}
if (tid % 32 == 0) {
s_buf[(tid / 32) + d * num_warps] = val;
}
}
```
For d=0, the first warp leader (tid=0) writes to `s_buf[0 + 0 * 32] = s_buf[0]`. This overwrites whatever was there. But between phase 2 and phase 4, we've already read `total_sum = s_buf[1]`, and `s_buf[0]` is no longer needed. Good.
OK, I think it's all correct.
Let me write the final code now. I'll make sure to handle potential edge cases and be careful about the shared memory layout.
One more detail: I should use `#include <cuda_fp16.h>`? No, we're using float.
I should also make sure to add `__restrict__` to all pointer parameters for better compiler optimization.
And I'll precompute `log(mix_weights)` etc. in `__init__`.
Let me write the Python code now.
Actually, I want to make the first kernel more efficient. For K=3 and D=16, each thread computes 3 * 16 = 48 differences. Let me see if I can vectorize or otherwise optimize.
One optimization: I can transpose `mix_mus` so that memory access is coalesced within a warp. Currently, each thread reads `mix_mus[k*D + d]` for d=0..D-1. Since all threads in a warp process different samples, they read the same `mix_mus` address. This is actually broadcast from global memory (via L1/L2 cache), which is efficient.
Similarly for `mix_inv_stds`, `mix_log_stds`. All threads read the same model parameters, so it's a broadcast pattern - very efficient.
For `samples`, each thread reads a different sample. If threads in a warp have consecutive `idx`, then `samples + idx * D` for consecutive idx means the memory accesses are strided by D*4=64 bytes. With 32 threads, that's 32*64 = 2048 bytes per warp load. This is not fully coalesced; each thread accesses a separate cache line (assuming 128B cache lines). So each thread causes a 64B load, but the hardware will coalesce adjacent accesses.
Actually, for samples with shape (N, D), thread i accesses `samples[i, 0..D-1]` and thread i+1 accesses `samples[i+1, 0..D-1]`. These are separated by D floats. For D=16, that's 64 bytes. On H100 with 128B L1 cache lines, threads i and i+1 access different cache lines. But the hardware can still service up to 32 independent memory requests per warp per cycle. However, it might take multiple cycles.
Wait, global memory loads are done at the warp level. All threads in a warp issue their load requests, and the memory subsystem coalesces them into as few transactions as possible. For `D=16` and `float` (4 bytes), `threadIdx.x` is the same for all threads in the warp if they have the same `idx % blockDim.x`. Well, in the same warp, threads have consecutive `threadIdx.x` values (0-31), so they have consecutive `idx` values (assuming blockIdx.x=0), so they read consecutive samples.
`sample[d]` for a fixed `d`: thread i reads `samples[i*D + d]`, thread i+1 reads `samples[(i+1)*D + d]`. The distance is D=16 floats = 64 bytes. With 32 threads, the accesses span 32*64 = 2048 bytes. Each 128B cache line holds 2 consecutive `sample[d]` values. So we need 16 cache lines per warp per `d`.
This is not great coalescing. Better coalescing would happen if threads in a warp accessed adjacent elements. But here, each thread processes one complete sample, so threads access different samples.
An alternative would be to have each warp process one sample together, with each thread in the warp computing one or a few dimensions. But D=16 < 32, so some threads would be idle. And each warp would need to do 3 * 16 = 48 operations, which is not enough to hide latency.
Alternatively, we could transpose the samples or process the data differently. But for N=4096 and D=16, the total memory traffic is only 256KB. Even with poor coalescing, this should be fast because the data will fit in L2 cache and the working set is small.
PyTorch also has the same access pattern, so my kernel should be competitive.
One more thought: I could reduce memory traffic by loading `sample` into registers once and reusing them for both log_p and log_q computations. My kernel already does this!
Let me check my planned kernel:
```cuda
const float* sample = samples + idx * D;
// Compute log_p using sample[d]
// Compute log_q using sample[d]
```
Yes, `sample` is loaded fresh for each d, but since D is small and the compiler might keep values in registers... actually, no, the compiler loads from global memory for each access unless it can keep them in registers. But for D=16, there's not enough register pressure to keep everything. However, L1/L2 cache will help.
Wait, I could explicitly load the sample into a local array:
```cuda
float s[16];
for (int d = 0; d < D; d++) s[d] = sample[d];
```
This forces the compiler to load from global memory once and keep in registers. For D=16, that's 16 registers. Since each thread has up to 255 registers, this is fine. For K=3 log_p computation plus log_q, the total register usage might be 16 (sample) + 3 (log_comp) + some temporaries ≈ 30-40 registers. Well within limits.
Let me add explicit local array loading. This might improve performance by ensuring the sample is loaded once.
Actually, the compiler might already do this optimization if it can prove the memory isn't aliased. With `__restrict__`, the compiler knows there are no aliases. But explicitly loading into registers is safer.
Let me update the first kernel:
```cuda
float s[16];
for (int d = 0; d < D; d++) {
s[d] = sample[d];
}
float log_comp[8];
for (int k = 0; k < K; k++) {
const float* mu = mix_mus + k * D;
const float* inv_std = mix_inv_stds + k * D;
const float* log_std = mix_log_stds + k * D;
float sum = 0.0f;
for (int d = 0; d < D; d++) {
float diff = s[d] - mu[d];
float scaled = diff * inv_std[d];
sum += -0.5f * scaled * scaled - log_std[d];
}
log_comp[k] = sum + mix_log_weights[k];
}
// ...
for (int d = 0; d < D; d++) {
float diff = s[d] - q_mu[d];
// ...
}
```
This is much better. The sample is loaded once into registers.
But wait, with `float s[16]` as a local array, if D is a runtime parameter, the compiler might not be able to fully unroll the loop or optimize register allocation. On some GPUs, local arrays spill to local memory if the compiler can't keep them in registers. But with small D=16 and plenty of registers, it should be fine.
To be safe, I can use `#pragma unroll` on the load loop if D is a template parameter. But it's not. I'll just hope the compiler optimizes it well.
Actually, I can use a fixed-size array and only use the first D elements:
```cuda
float s[64];
#pragma unroll
for (int d = 0; d < D; d++) {
s[d] = sample[d];
}
```
But `#pragma unroll` on a loop with runtime bound doesn't do anything. Without unrolling, the array might not stay in registers.
Hmm, let me just use the direct memory access pattern. For D=16 and K=3, the compiler should be smart enough to cache the loads. And since sample is accessed multiple times, L1 cache will help.
I'll keep it simple and directly access `sample[d]`.
Now, one more thing: I noticed that the reference uses `self.mix_stds` and `self.mix_stds.log().unsqueeze(0)` in `_log_p`. The `unsqueeze(0)` adds a dimension for broadcasting. In my kernel, I don't need this because I handle the indexing explicitly.
Let me finalize the complete code. Here's my plan:
1. Define CUDA source with two kernels and a wrapper function
2. Compile with `load_inline` at module level
3. Define `ModelNew` class
4. In `__init__`, create precomputed buffers
5. In `forward`, call the compiled function
I'm ready to write it.
Wait, I should also consider whether to pass the precomputed buffers as contiguous tensors. The reference uses `.register_buffer()` which makes them contiguous by default. But to be safe, I'll ensure they are contiguous in `__init__`.
Also, I need to make sure `q_mu` and other buffers are on the same device as `samples`. The reference model handles this automatically. In my case, the `load_inline` function just calls CUDA kernels, so all tensors need to be on CUDA. If `samples` is on CPU, the kernel won't work. But `get_inputs()` likely returns CUDA tensors. Let me check: `get_inputs()` returns `[torch.randn(num_samples, dim) * q_std]`. In the evaluator, these are probably moved to the model's device. But to be safe, I can handle device transfer in `forward` if needed.
Actually, looking at how the evaluation typically works: the model is instantiated and then `model(inputs)` is called, where inputs are on the same device as the model. The reference model inherits from `nn.Module`, and its `register_buffer` calls put buffers on CPU by default. But in the evaluator, the model is likely moved to CUDA before being called.
For `ModelNew`, I need to ensure that when the model is moved to CUDA (via `.to('cuda')` or `.cuda()`), the precomputed buffers also move. Since I use `register_buffer` for them too, they will move automatically.
Good.
Now, one remaining question: should I call `.cuda()` on inputs inside `forward` if they're not already on CUDA? The reference doesn't do this, so I shouldn't either. The evaluator handles device placement.
Let me write the code!
Also, I want to note: the reference model has this in `_log_q`:
```python
return (-0.5 * ((x - self.q_mu) / self.q_std) ** 2 - self.q_std.log()).sum(-1)
```
Wait, there's no `.unsqueeze(0)` for q_mu or q_std, but x is (B, D), q_mu is (D,), and broadcasting works fine in PyTorch. My kernel does the equivalent.
And `_log_p`:
```python
diff = x.unsqueeze(1) - self.mix_mus.unsqueeze(0) # (B, K, D)
log_comp = -0.5 * (diff / self.mix_stds.unsqueeze(0)) ** 2 - self.mix_stds.log().unsqueeze(0)
log_comp = log_comp.sum(-1) # (B, K)
return torch.logsumexp(log_comp + self.mix_weights.log().unsqueeze(0), dim=-1)
```
My kernel computes the same thing.
One numerical detail: `self.mix_stds.log().unsqueeze(0)` has shape (1, K, D). For `diff / stds`, PyTorch broadcasts. In my kernel, I compute `diff * inv_std[d]`, which is equivalent.
Actually wait, PyTorch broadcasts `self.mix_stds.unsqueeze(0)` from (K, D) to (B, K, D). My kernel uses `mix_mus + k * D`, so it's accessing the k-th component directly. This is correct.
Now for `mix_weights.log().unsqueeze(0)`: shape (1, K). In my kernel, `mix_log_weights[k]` is the log weight for the k-th component. Correct.
OK, let me write the final code.
One more thought: I can potentially do a single-kernel version to save one kernel launch. For N=4096, the kernel launch overhead is non-trivial compared to the actual computation. Let me try the 2-kernel version first, and if timing is not good enough, I'll try the 1-kernel version.
But actually, writing the 1-kernel version now might save time. Let me think if it's significantly more complex...
The main complexity of the 1-kernel version is the shared memory for `logw`. For a single block of 1024 threads processing N=4096 samples:
- Each thread computes logw for 4 samples
- Store all 4096 logw values in shared memory (16KB)
- Reduce for global max and sum
- Each thread computes weighted sums for its 4 samples
- Reduce for mean and ess
This is actually not much more complex, and it avoids the global memory read/write for `logw`. Let me try it!
The single-kernel version:
```cuda
__global__ void fused_kernel(
const float* __restrict__ samples,
const float* __restrict__ mix_mus,
const float* __restrict__ mix_inv_stds,
const float* __restrict__ mix_log_stds,
const float* __restrict__ mix_log_weights,
const float* __restrict__ q_mu,
const float* __restrict__ q_inv_std,
const float* __restrict__ q_log_std,
float* __restrict__ out_mean,
float* __restrict__ out_ess,
int N, int D, int K) {
const int tid = threadIdx.x;
const int num_threads = blockDim.x;
const int num_warps = (num_threads + 31) / 32;
extern __shared__ float smem[];
float* s_logw = smem;
float* s_buf = smem + N;
// Phase 1: Compute logw for each sample
for (int i = tid; i < N; i += num_threads) {
const float* sample = samples + i * D;
float log_comp[8];
for (int k = 0; k < K; k++) {
const float* mu = mix_mus + k * D;
const float* inv_std = mix_inv_stds + k * D;
const float* log_std = mix_log_stds + k * D;
float sum = 0.0f;
for (int d = 0; d < D; d++) {
float diff = sample[d] - mu[d];
float scaled = diff * inv_std[d];
sum += -0.5f * scaled * scaled - log_std[d];
}
log_comp[k] = sum + mix_log_weights[k];
}
float max_val = log_comp[0];
for (int k = 1; k < K; k++) {
max_val = fmaxf(max_val, log_comp[k]);
}
float sum_exp = 0.0f;
for (int k = 0; k < K; k++) {
sum_exp += expf(log_comp[k] - max_val);
}
float log_p = max_val + logf(sum_exp);
float log_q = 0.0f;
for (int d = 0; d < D; d++) {
float diff = sample[d] - q_mu[d];
float scaled = diff * q_inv_std[d];
log_q += -0.5f * scaled * scaled - q_log_std[d];
}
s_logw[i] = log_p - log_q;
}
__syncthreads();
// Phase 2: Find global max of logw
float local_max = -1e30f;
for (int i = tid; i < N; i += num_threads) {
if (s_logw[i] > local_max) local_max = s_logw[i];
}
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
local_max = fmaxf(local_max, __shfl_down_sync(0xFFFFFFFF, local_max, offset));
}
if (tid % 32 == 0) {
s_buf[tid / 32] = local_max;
}
__syncthreads();
if (tid < 32) {
local_max = (tid < num_warps) ? s_buf[tid] : -1e30f;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
local_max = fmaxf(local_max, __shfl_down_sync(0xFFFFFFFF, local_max, offset));
}
if (tid == 0) s_buf[0] = local_max;
}
__syncthreads();
float global_max = s_buf[0];
// Phase 3: Sum of exp(logw - global_max)
float local_sum = 0.0f;
for (int i = tid; i < N; i += num_threads) {
local_sum += expf(s_logw[i] - global_max);
}
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
local_sum += __shfl_down_sync(0xFFFFFFFF, local_sum, offset);
}
if (tid % 32 == 0) {
s_buf[tid / 32] = local_sum;
}
__syncthreads();
if (tid < 32) {
local_sum = (tid < num_warps) ? s_buf[tid] : 0.0f;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
local_sum += __shfl_down_sync(0xFFFFFFFF, local_sum, offset);
}
if (tid == 0) s_buf[1] = local_sum;
}
__syncthreads();
float total_sum = s_buf[1];
float lse = global_max + logf(total_sum);
// Phase 4: Compute local weighted_mean and local_w2
float local_mean[16];
for (int d = 0; d < D; d++) {
local_mean[d] = 0.0f;
}
float local_w2 = 0.0f;
for (int i = tid; i < N; i += num_threads) {
float w = expf(s_logw[i] - lse);
local_w2 += w * w;
const float* sample = samples + i * D;
for (int d = 0; d < D; d++) {
local_mean[d] += w * sample[d];
}
}
// Phase 5: Block reduction for each dimension
for (int d = 0; d < D; d++) {
float val = local_mean[d];
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
val += __shfl_down_sync(0xFFFFFFFF, val, offset);
}
if (tid % 32 == 0) {
s_buf[(tid / 32) + d * num_warps] = val;
}
}
float w2_val = local_w2;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
w2_val += __shfl_down_sync(0xFFFFFFFF, w2_val, offset);
}
if (tid % 32 == 0) {
s_buf[D * num_warps + tid / 32] = w2_val;
}
__syncthreads();
if (tid < 32) {
for (int d = 0; d < D; d++) {
float val = (tid < num_warps) ? s_buf[tid + d * num_warps] : 0.0f;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
val += __shfl_down_sync(0xFFFFFFFF, val, offset);
}
if (tid == 0) {
out_mean[d] = val;
}
}
float w2_sum = (tid < num_warps) ? s_buf[D * num_warps + tid] : 0.0f;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
w2_sum += __shfl_down_sync(0xFFFFFFFF, w2_sum, offset);
}
if (tid == 0) {
*out_ess = 1.0f / w2_sum;
}
}
}
```
Shared memory needed: `N + D * num_warps + num_warps + 2` floats. For N=4096, D=16, num_warps=32: 4096 + 512 + 32 + 2 = 4642 floats = 18.6KB. Well within limits!
And we launch with a single block: `<<<1, 1024, smem_size>>>`.
This is elegant! One kernel launch, no intermediate global memory for logw. Let me use this.
One concern: the kernel is quite large. Register usage might be high, potentially causing spilling. Let me estimate:
- `sample[d]` accesses: if the compiler doesn't keep them in registers, it loads from global memory each time. The L1 cache should help.
- `log_comp[8]`: 8 registers.
- `local_mean[16]`: 16 registers.
- Various scalars: ~10 registers.
- Total: ~40 registers, well under 255. No spilling expected.
Actually, `local_mean[16]` is declared at the end of the kernel. The compiler should allocate registers for it. For D=16, that's 16 registers. Plus the earlier temporaries might have been freed. Total should be under 64 registers, which allows 32 threads/warp * 64 = 2048 registers/warp, and each SM has 64K registers, so 32 warps per SM. Wait, 64K / 64 = 1024 threads per SM at 64 registers each. This is fine.
Actually, on H100, there are 64K registers per SM and max 2048 threads per SM. With 1024 threads in a block and ~64 registers each, we get 64K registers used, which is exactly the limit. This might not allow two blocks per SM, but we only have 1 block anyway.
If register usage is higher, say 80 registers, then 1024 * 80 = 81.6K > 64K, which would cause spilling. I need to be careful. Let me try to keep register usage low.
I can reduce register pressure by not using `local_mean[16]` and instead doing the accumulation into shared memory with atomics or a different reduction strategy. But with 1024 threads and 16 dimensions, each dimension only has 64 accumulators if we use shared memory per-thread, which would be 16 * 1024 * 4 = 64KB. That's more than shared memory.
Alternative: process the dimensions sequentially with a single float accumulator, and do multiple warp reductions:
```cuda
for (int d = 0; d < D; d++) {
float local_m = 0.0f;
for (int i = tid; i < N; i += num_threads) {
float w = expf(s_logw[i] - lse);
local_m += w * samples[i * D + d];
}
// warp and block reduce local_m
// write out_mean[d]
}
```
This way, we only need 1 float for local accumulation instead of 16. The tradeoff is we compute `w = expf(s_logw[i] - lse)` D times instead of once. But we have to re-read `s_logw[i]` each time (from shared memory, which is fast) and re-read `samples[i*D + d]` (from global memory or L1 cache).
For D=16, we'd compute `expf` 16 times per sample instead of 1 time. That's 15 extra expf per sample, or 15 * 4096 = 61k extra expf calls. expf is relatively expensive (~20 cycles). Total extra cost: ~1.2M cycles. At 2GHz, that's ~600us. That seems too high.
Wait, that's wrong. If each thread handles 4 samples and D=16, and we do D sequential passes:
- For each of the 4 samples, compute w once, store in register
- Then for each d, accumulate w * sample[d]
Actually, we can do:
```cuda
for (int i = tid; i < N; i += num_threads) {
float w = expf(s_logw[i] - lse);
s_logw[i] = w; // Reuse shared memory to store w
}
__syncthreads();
for (int d = 0; d < D; d++) {
float local_m = 0.0f;
for (int i = tid; i < N; i += num_threads) {
local_m += s_logw[i] * samples[i * D + d];
}
// reduce local_m and write out_mean[d]
}
// For ESS, compute sum of w^2
float local_w2 = 0.0f;
for (int i = tid; i < N; i += num_threads) {
float w = s_logw[i];
local_w2 += w * w;
}
// reduce local_w2
```
This reuses `s_logw` to store `w`! That saves us from recomputing expf. And it only requires 1 float local accumulator per thread per dimension. No need for `local_mean[16]`. This significantly reduces register pressure.
The tradeoff:
- We write to shared memory once (N floats)
- We read from shared memory D+1 times per sample
- Shared memory bandwidth is ~20 TB/s, much faster than global memory
- Global memory reads: samples are read D+1 times (D times for mean, 0 times for w2 since it's in smem)
Wait, samples are read D times for computing mean, but if samples don't fit in L1, that's D * 256KB = 4MB of global memory reads. Hmm.
Actually, samples are 256KB. L2 cache is ~50MB on H100. L1 is 128KB per SM. Since the kernel runs on one SM (one block), L1 is 128KB. Samples (256KB) don't fit in L1. But they do fit in L2, and reading from L2 is fast (~10 TB/s).
Actually, let me reconsider. In the first approach with `local_mean[16]`, each thread reads each sample once. Total global memory reads for samples in phase 4: 256KB. In the sequential approach, each thread reads samples D times: 4MB. That's 16x more global memory traffic!
Wait, but the sample values might stay in L2 cache. The first pass (phase 1) reads all samples. In phase 4, with sequential per-dimension processing, we'd read them again. L2 cache hit rate would be high if the cache is large enough and we're not evicting.
But it's still more traffic. And 4MB at L2 bandwidth (~10 TB/s) is ~400ns. The `local_mean[16]` approach does 256KB in ~25ns (if L1 hit) or ~100ns (if L2 hit). Actually, with 1024 threads each reading 4 samples of 16 floats = 64 floats = 256 bytes, total per-thread memory traffic in phase 4 is 256 bytes. Over 1024 threads, all reading from L2, it's 256KB. The L2 can handle this. But doing it 16 times is 4MB.
Actually wait, in the `local_mean[16]` approach, each thread also reads the same 256 bytes of sample data. It just does it once and accumulates all 16 dimensions. In the sequential approach, it reads dimension 0 for all samples, then dimension 1 for all samples, etc. The total bytes read are the same! (16 floats per sample * 4 bytes * 4096 samples = 256KB for one pass, * 16 passes = 4MB).
Wait no, that can't be right. In the `local_mean[16]` approach, for 4096 samples * 16 dims * 4 bytes = 256KB of sample data. Each thread processes 4 samples, so it reads 4 * 16 * 4 = 256 bytes. Across 1024 threads, the total is 1024 * 256 bytes = 256KB. But each sample is read by exactly one thread, so the total global memory traffic is 256KB.
In the sequential approach:
- For d=0: all threads read sample[i, 0] for their assigned samples. Total traffic: 4096 * 4 = 16KB.
- For d=1: all threads read sample[i, 1]. Total: 16KB.
- ...
- For d=15: total: 16KB.
- Total: 256KB... wait, that's the same!
I was mistakenly thinking each thread reads the full sample each time. No, in the sequential approach, for each dimension d, each thread only reads `samples[i * D + d]`, not the full sample. So total traffic is still 4096 * 16 * 4 = 256KB.
The difference is memory access pattern:
- `local_mean[16]`: each thread reads a contiguous block of 16 floats (one sample), then another block, etc. This is strided by D=16 between threads in a warp.
- Sequential: each thread reads a single float, then the next single float for the next sample. This is strided by D=16 between threads.
Actually, both have the same strided pattern. In `local_mean[16]`, thread 0 reads `samples[0..15]`, thread 1 reads `samples[16..31]`. In sequential for d=0, thread 0 reads `samples[0]`, thread 1 reads `samples[16]`. The pattern for a single dimension is even more strided, but the coalescing hardware can handle it.
Actually, looking more carefully:
- In `local_mean[16]`, for a fixed d, threads in a warp read `samples[i*D + d]` where i differs by 1 between consecutive threads. The addresses are spaced by D=16 floats = 64 bytes. This is the same as the sequential approach!
So the total memory traffic is the same (256KB), and the access pattern is the same. The only difference is:
- `local_mean[16]`: reads sample once, accumulates 16 dimensions. Needs 16 registers for local_mean.
- Sequential: reads each dimension separately, accumulates into 1 register. Reuses the register 16 times. Needs 1 register for local accumulation.
The sequential approach uses fewer registers but requires either:
1. Recomputing `w = expf(s_logw[i] - lse)` D times, or
2. Storing `w` in shared memory or registers
Option 2 is good: store `w` in shared memory by overwriting `s_logw[i]`. Then in each dimension pass, read `w` from shared memory.
So the approach would be:
1. Phase 1: compute `s_logw[i] = log_p - log_q`
2. Phase 2: reduce for lse
3. Overwrite `s_logw[i] = expf(s_logw[i] - lse)`
4. Sequential dimension reduction: for each d, accumulate `s_logw[i] * samples[i*D + d]`
5. Reduce w2: `s_logw[i] * s_logw[i]`
This uses `s_logw` for two purposes, which is fine.
The advantages:
- Lower register pressure
- Same memory traffic
- More elegant?
The disadvantage:
- One extra pass over `s_logw` to compute `w`
- D passes over `samples` instead of 1 pass (but same total traffic)
Actually, D passes over samples vs 1 pass: with L2 cache, each sample read for dimension d might hit L2 if the cache can hold all samples. But L2 on H100 is 50MB, and samples are 256KB. So all of samples can stay in L2! Reading them 16 times is fine.
Wait, but in the `local_mean[16]` approach, we also need to store `local_mean[16]` in registers. For 16 registers, the compiler might decide to spill if other temporaries push us over the limit.
Let me just go with the `local_mean[16]` approach. Register pressure should be manageable, and it has fewer passes over the data.
Actually, let me count potential register usage more carefully for the `local_mean[16]` approach:
Phase 1 (per sample, per 4 samples total):
- `s_logw[i]` write
- `log_comp[8]`: up to 8 registers, but can be reused
- `diff`, `scaled`, `sum`, `max_val`, `sum_exp`, `log_p`, `log_q`, `w`: ~10 registers
- Temporary loop variables: a few
Max registers at any point: maybe 20-30.
Phase 4 (per sample):
- `w`, `w2`, `sample[d]`, `local_mean[d]`: 16 + a few = ~20 registers
Total peak: maybe ~30-40 registers. Safe.
I'll use the `local_mean[16]` approach.
One more thing: since the kernel is quite large, the compiler might not be able to keep `local_mean[16]` in registers if the function is too complex. A common trick is to split the computation into smaller device functions, but that can add call overhead.
Actually, thinking about it more: `local_mean[16]` is used after the `lse` computation. The compiler knows its live range. The earlier temporaries (`log_comp`, etc.) are no longer needed. So registers can be reused.
OK, I'll write the single-kernel version. Let me be very careful about shared memory.
For the shared memory layout:
- `smem[0..N-1]` = `s_logw`
- `smem[N..N+num_warps-1]` = scratch for warp max/sum
- But wait, in the block reduce step, I use `s_buf` which starts at `smem + N`. I need enough space for:
- Phase 1/2: `num_warps` floats
- Phase 4: `D * num_warps + num_warps` floats
- The max of these is `D * num_warps + num_warps`
- So total shared memory = `N + D * num_warps + num_warps` floats
Actually, in Phase 4, I write to:
- `s_buf[warp_id + d * num_warps]` for each dimension
- `s_buf[D * num_warps + warp_id]` for w2
But `s_buf` starts at `smem + N`. So the total shared memory indices used are:
- `N` to `N + D * num_warps + num_warps - 1`
- Total elements: `N + D * num_warps + num_warps`
For N=4096, D=16, num_warps=32: 4096 + 512 + 32 = 4640 floats.
Let me allocate: `extern __shared__ float smem[];` and in host: `smem_size = (N + D * num_warps + num_warps) * sizeof(float)`.
Wait, N is a runtime variable. In the host wrapper, N is known. But the kernel launches with `<<<1, 1024, smem_size>>>` where `smem_size` is dynamic.
Actually, the maximum shared memory per block might limit N. H100 supports up to 227KB. For N=4096, we need ~18.5KB. For N=8192, ~37KB. For N=16384, ~73KB. For N=32768, ~146KB. For N=65536, ~292KB > 227KB limit.
So the single-block kernel works for N up to about 50K. For the problem, N=4096, so it's fine.
But what if the evaluator tests with different N? The problem defines `num_samples = 4096`, but sometimes evaluators vary the inputs. To be safe, I can check N in `forward` and fall back to a multi-block approach for large N. Or I can just use the 2-kernel approach which scales to arbitrary N.
You know what, let me just use the 2-kernel approach. It's simpler, more robust, and avoids dealing with shared memory size limits. The overhead of an extra kernel launch and 16KB of global memory traffic is negligible compared to PyTorch's many kernel launches.
Let me write the 2-kernel version and see how it performs.
Full code:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
CUDA_SOURCE = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>
__global__ void compute_logw_kernel(
const float* __restrict__ samples,
const float* __restrict__ mix_mus,
const float* __restrict__ mix_inv_stds,
const float* __restrict__ mix_log_stds,
const float* __restrict__ mix_log_weights,
const float* __restrict__ q_mu,
const float* __restrict__ q_inv_std,
const float* __restrict__ q_log_std,
float* __restrict__ logw,
int N, int D, int K) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= N) return;
const float* sample = samples + idx * D;
float log_comp[8];
for (int k = 0; k < K; k++) {
const float* mu = mix_mus + k * D;
const float* inv_std = mix_inv_stds + k * D;
const float* log_std = mix_log_stds + k * D;
float sum = 0.0f;
for (int d = 0; d < D; d++) {
float diff = sample[d] - mu[d];
float scaled = diff * inv_std[d];
sum += -0.5f * scaled * scaled - log_std[d];
}
log_comp[k] = sum + mix_log_weights[k];
}
float max_val = log_comp[0];
for (int k = 1; k < K; k++) {
max_val = fmaxf(max_val, log_comp[k]);
}
float sum_exp = 0.0f;
for (int k = 0; k < K; k++) {
sum_exp += expf(log_comp[k] - max_val);
}
float log_p = max_val + logf(sum_exp);
float log_q = 0.0f;
for (int d = 0; d < D; d++) {
float diff = sample[d] - q_mu[d];
float scaled = diff * q_inv_std[d];
log_q += -0.5f * scaled * scaled - q_log_std[d];
}
logw[idx] = log_p - log_q;
}
__global__ void compute_results_kernel(
const float* __restrict__ samples,
const float* __restrict__ logw,
float* __restrict__ out_mean,
float* __restrict__ out_ess,
int N, int D) {
const int tid = threadIdx.x;
const int num_threads = blockDim.x;
const int num_warps = (num_threads + 31) / 32;
extern __shared__ float shared[];
float* s_buf = shared;
// Phase 1: Find global max of logw
float local_max = -1e30f;
for (int i = tid; i < N; i += num_threads) {
if (logw[i] > local_max) local_max = logw[i];
}
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
local_max = fmaxf(local_max, __shfl_down_sync(0xFFFFFFFF, local_max, offset));
}
if (tid % 32 == 0) {
s_buf[tid / 32] = local_max;
}
__syncthreads();
if (tid < 32) {
local_max = (tid < num_warps) ? s_buf[tid] : -1e30f;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
local_max = fmaxf(local_max, __shfl_down_sync(0xFFFFFFFF, local_max, offset));
}
if (tid == 0) s_buf[0] = local_max;
}
__syncthreads();
float global_max = s_buf[0];
// Phase 2: Sum of exp(logw - global_max)
float local_sum = 0.0f;
for (int i = tid; i < N; i += num_threads) {
local_sum += expf(logw[i] - global_max);
}
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
local_sum += __shfl_down_sync(0xFFFFFFFF, local_sum, offset);
}
if (tid % 32 == 0) {
s_buf[tid / 32] = local_sum;
}
__syncthreads();
if (tid < 32) {
local_sum = (tid < num_warps) ? s_buf[tid] : 0.0f;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
local_sum += __shfl_down_sync(0xFFFFFFFF, local_sum, offset);
}
if (tid == 0) s_buf[1] = local_sum;
}
__syncthreads();
float total_sum = s_buf[1];
float lse = global_max + logf(total_sum);
// Phase 3: Compute local weighted_mean and local_w2
float local_mean[16];
for (int d = 0; d < D; d++) {
local_mean[d] = 0.0f;
}
float local_w2 = 0.0f;
for (int i = tid; i < N; i += num_threads) {
float w = expf(logw[i] - lse);
local_w2 += w * w;
const float* sample = samples + i * D;
for (int d = 0; d < D; d++) {
local_mean[d] += w * sample[d];
}
}
// Phase 4: Block reduction for each dimension
for (int d = 0; d < D; d++) {
float val = local_mean[d];
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
val += __shfl_down_sync(0xFFFFFFFF, val, offset);
}
if (tid % 32 == 0) {
s_buf[(tid / 32) + d * num_warps] = val;
}
}
float w2_val = local_w2;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
w2_val += __shfl_down_sync(0xFFFFFFFF, w2_val, offset);
}
if (tid % 32 == 0) {
s_buf[D * num_warps + tid / 32] = w2_val;
}
__syncthreads();
if (tid < 32) {
for (int d = 0; d < D; d++) {
float val = (tid < num_warps) ? s_buf[tid + d * num_warps] : 0.0f;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
val += __shfl_down_sync(0xFFFFFFFF, val, offset);
}
if (tid == 0) {
out_mean[d] = val;
}
}
float w2_sum = (tid < num_warps) ? s_buf[D * num_warps + tid] : 0.0f;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
w2_sum += __shfl_down_sync(0xFFFFFFFF, w2_sum, offset);
}
if (tid == 0) {
*out_ess = 1.0f / w2_sum;
}
}
}
std::vector<torch::Tensor> importance_sampling_forward(
torch::Tensor samples,
torch::Tensor mix_mus,
torch::Tensor mix_inv_stds,
torch::Tensor mix_log_stds,
torch::Tensor mix_log_weights,
torch::Tensor q_mu,
torch::Tensor q_inv_std,
torch::Tensor q_log_std) {
int N = samples.size(0);
int D = samples.size(1);
int K = mix_mus.size(0);
auto options = samples.options();
auto logw = torch::empty({N}, options);
auto out_mean = torch::empty({D}, options);
auto out_ess = torch::empty({}, options);
const int threads = 256;
const int blocks = (N + threads - 1) / threads;
compute_logw_kernel<<<blocks, threads>>>(
samples.data_ptr<float>(),
mix_mus.data_ptr<float>(),
mix_inv_stds.data_ptr<float>(),
mix_log_stds.data_ptr<float>(),
mix_log_weights.data_ptr<float>(),
q_mu.data_ptr<float>(),
q_inv_std.data_ptr<float>(),
q_log_std.data_ptr<float>(),
logw.data_ptr<float>(),
N, D, K);
const int reduction_threads = 1024;
int num_warps = (reduction_threads + 31) / 32;
int smem_size = (D * num_warps + num_warps + 2) * sizeof(float);
compute_results_kernel<<<1, reduction_threads, smem_size>>>(
samples.data_ptr<float>(),
logw.data_ptr<float>(),
out_mean.data_ptr<float>(),
out_ess.data_ptr<float>(),
N, D);
return {out_mean, out_ess};
}
"""
```
Wait, there's a subtle issue. In `compute_results_kernel`, I use `s_buf[1]` in phase 2. But what if `num_warps == 1`? Then `s_buf[1]` would be outside the phase 1/2 buffer. But for 1024 threads, `num_warps = 32`, so `s_buf[1]` is within the first 32 elements. Fine.
For safety, I could put:
```cuda
if (tid == 0) s_buf[num_warps] = local_sum;
```
instead of `s_buf[1]`. This way, the sum is stored at a fixed location after the warp-level buffers. Let me change this.
Actually, looking at the shared memory layout:
- `s_buf[0..num_warps-1]`: used in phase 1/2 for warp-level max/sum
- `s_buf[num_warps..num_warps+1]`: could store global_max and total_sum
Currently I use `s_buf[0]` for global_max and `s_buf[1]` for total_sum. Since `num_warps >= 1`, `s_buf[1]` might be overwritten in phase 4 if `num_warps >= 2` and `d=0`, `warp_id=1` writes to `s_buf[1 + 0*num_warps] = s_buf[1]`!
Oh no! That's a bug!
In phase 4:
```cuda
for (int d = 0; d < D; d++) {
// ...
if (tid % 32 == 0) {
s_buf[(tid / 32) + d * num_warps] = val;
}
}
```
For `d=0`, `warp_id=1` (tid=32), this writes to `s_buf[1 + 0*32] = s_buf[1]`. But `s_buf[1]` was used to store `total_sum` in phase 2!
After phase 2, all threads read `total_sum = s_buf[1]`. Then `__syncthreads()`. Then in phase 4, threads overwrite `s_buf`. Since `total_sum` was already read into the register `total_sum`, overwriting `s_buf[1]` is fine!
Wait, let me check:
1. Phase 2 ends with `if (tid == 0) s_buf[1] = local_sum;`
2. `__syncthreads()`
3. `float total_sum = s_buf[1];` - all threads read this
4. Then phase 3 and 4 happen, which write to `s_buf[0..D*num_warps+num_warps-1]`
Since `total_sum` is already loaded into a register by the time phase 4 writes to `s_buf`, there's no issue. But I need to make sure `__syncthreads()` is placed correctly.
Phase 2 code:
```cuda
if (tid < 32) {
// ...
if (tid == 0) s_buf[1] = local_sum;
}
__syncthreads();
float total_sum = s_buf[1];
```
Yes! `__syncthreads()` is after the write to `s_buf[1]` and before all threads read `s_buf[1]`. Good.
Then phase 4 writes to `s_buf` starting after all threads have read `total_sum`. Fine.
BUT WAIT: in phase 2, during the block-level reduction, threads in the first warp read from `s_buf[tid]`. What if `tid=1` reads `s_buf[1]` before `tid=0` writes to it? Let me check:
```cuda
if (tid % 32 == 0) {
s_buf[tid / 32] = local_sum; // All warp leaders write
}
__syncthreads();
if (tid < 32) {
local_sum = (tid < num_warps) ? s_buf[tid] : 0.0f; // First warp reads
// ...
if (tid == 0) s_buf[1] = local_sum; // tid=0 writes
}
__syncthreads();
```
Yes! The first `__syncthreads()` ensures all warp leaders have written before the first warp reads. Then within the first warp, only tid=0 writes `s_buf[1]` at the end. Then `__syncthreads()` ensures the write is visible before all threads read `s_buf[1]`.
OK, so `s_buf[1]` is fine. But I still feel uneasy using `s_buf[1]` as a special slot when `s_buf[0]` is also a special slot. For clarity, let me separate the scratch space:
```cuda
// s_buf[0] = global_max (after phase 1)
// s_buf[1] = total_sum (after phase 2)
// Phase 4 writes to s_buf[0..D*num_warps+num_warps-1]
```
Actually, since phase 4 starts at `s_buf[0]` for `d=0, warp_id=0`, it will overwrite `s_buf[0]` (global_max) and `s_buf[1]` (total_sum). But that's OK because both values have already been loaded into registers (`global_max` was loaded as `s_buf[0]` by all threads after phase 1's `__syncthreads()`, and `total_sum` was loaded as `s_buf[1]` after phase 2's `__syncthreads()`).
OK, it's correct. But let me make it more explicit by using named offsets. Or I can just keep the code as is, since it's correct.
Actually, to be extra safe and avoid confusion, I'll use `s_buf[num_warps]` for total_sum and `s_buf[num_warps + 1]` for any other globals. And phase 4 will write starting at `s_buf[0]`. No wait, if phase 4 writes to `s_buf[0]`, it might overwrite `s_buf[0]` which was used in phase 2...
You know what, I'll just resize the shared memory and use clearly separated regions. Shared memory layout:
- Region A: `s_buf[0..num_warps-1]` for phase 1/2 warp buffers
- Region B: `s_buf[num_warps]` for global_max
- Region C: `s_buf[num_warps + 1]` for total_sum
- Region D: `s_buf[num_warps + 2 .. num_warps + 2 + D*num_warps + num_warps - 1]` for phase 4
Wait, I don't need to do this. The current code is correct. Let me just keep it simple and move on.
One more issue to check: does `num_warps = (num_threads + 31) / 32` calculate correctly for 1024? `1024 + 31 = 1055; 1055 / 32 = 32` (integer division). Yes.
OK, let me finalize the code and compile it.
Actually, I want to address one more thing: `log_comp[8]` in the kernel. If K > 8, we'd have a buffer overflow. Should I increase it or use dynamic allocation?
The reference defaults to `num_components=3`. The evaluation probably uses 3. To be safe, I'll increase it to a reasonable maximum, say 64 or 128. But `log_comp` is stored in registers, so a large array could cause spilling. Let me use 64.
Actually, for safety without excessive register pressure, I can just compute log_p using a two-pass approach without storing log_comp:
```cuda
// First pass: find max
float max_val = -1e30f;
for (int k = 0; k < K; k++) {
const float* mu = mix_mus + k * D;
const float* inv_std = mix_inv_stds + k * D;
const float* log_std = mix_log_stds + k * D;
float sum = 0.0f;
for (int d = 0; d < D; d++) {
float diff = sample[d] - mu[d];
float scaled = diff * inv_std[d];
sum += -0.5f * scaled * scaled - log_std[d];
}
float log_comp = sum + mix_log_weights[k];
if (log_comp > max_val) max_val = log_comp;
}
// Second pass: compute sum_exp
float sum_exp = 0.0f;
for (int k = 0; k < K; k++) {
const float* mu = mix_mus + k * D;
const float* inv_std = mix_inv_stds + k * D;
const float* log_std = mix_log_stds + k * D;
float sum = 0.0f;
for (int d = 0; d < D; d++) {
float diff = sample[d] - mu[d];
float scaled = diff * inv_std[d];
sum += -0.5f * scaled * scaled - log_std[d];
}
float log_comp = sum + mix_log_weights[k];
sum_exp += expf(log_comp - max_val);
}
float log_p = max_val + logf(sum_exp);
```
This avoids the `log_comp` array entirely but computes the component sums twice. For K=3, that's fine. And it supports arbitrary K without register array limits.
However, there's a subtle numerical issue: computing the sum twice might give slightly different results due to floating point differences. But wait, the same inputs (sample, mu, std) are used, so the results should be identical. The only issue is if the compiler optimizes differently, but with `-O3` and deterministic floating point, it should be the same.
Actually, for maximum numerical stability and to match PyTorch's behavior, I should compute exactly the same thing. The two-pass approach computes the exact same intermediate values, so it's fine.
Wait, but PyTorch computes `diff = x.unsqueeze(1) - mix_mus`, then `log_comp = ...`, then `logsumexp`. In my kernel, I compute each component's contribution directly. As long as the math is the same, the result should match within fp32 tolerance.
For performance, computing each component twice means 2 * K * D operations instead of K * D + K. For K=3, D=16: two-pass is 96 ops per sample, one-pass with storage is 48 + 3 = 51 ops. The difference is 45 ops, which is negligible compared to memory bandwidth.
I'll use the two-pass approach to avoid the fixed-size array and support arbitrary K.
Let me update the first kernel:
```cuda
__global__ void compute_logw_kernel(...) {
// ...
float max_val = -1e30f;
for (int k = 0; k < K; k++) {
const float* mu = mix_mus + k * D;
const float* inv_std = mix_inv_stds + k * D;
const float* log_std = mix_log_stds + k * D;
float sum = 0.0f;
for (int d = 0; d < D; d++) {
float diff = sample[d] - mu[d];
float scaled = diff * inv_std[d];
sum += -0.5f * scaled * scaled - log_std[d];
}
float log_comp = sum + mix_log_weights[k];
if (log_comp > max_val) max_val = log_comp;
}
float sum_exp = 0.0f;
for (int k = 0; k < K; k++) {
const float* mu = mix_mus + k * D;
const float* inv_std = mix_inv_stds + k * D;
const float* log_std = mix_log_stds + k * D;
float sum = 0.0f;
for (int d = 0; d < D; d++) {
float diff = sample[d] - mu[d];
float scaled = diff * inv_std[d];
sum += -0.5f * scaled * scaled - log_std[d];
}
float log_comp = sum + mix_log_weights[k];
sum_exp += expf(log_comp - max_val);
}
float log_p = max_val + logf(sum_exp);
// ...
}
```
Actually, I realize I might be able to use `powf` or something instead of manual computation for the Gaussian terms. But manual is simpler and probably faster.
One more thing: `logf(sum_exp)` where `sum_exp` could potentially be 0 if all `log_comp` are -inf. But with finite inputs and K >= 1, `sum_exp` should be positive. I don't need to handle the 0 case.
Wait, actually, `log_comp - max_val` could underflow to 0 in expf for very negative values, but `sum_exp` would still be at least 1.0 (from the max term). So `log_p` is well-defined.
OK, let me finalize the code and compile.
Let me also note: I should check if `out_ess` needs to be a 0-dimensional tensor or can be a 1-D tensor of size 1. The reference returns a scalar tensor. Let me make sure.
In the wrapper:
```cpp
auto out_ess = torch::empty({}, options);
```
This creates a 0-dimensional tensor. Good.
And in the kernel:
```cuda
*out_ess = 1.0f / w2_sum;
```
This writes to `out_ess.data_ptr<float>()` which points to the single float in the 0-D tensor. Good.
Alright, I'm going to write the final code now.
Let me double-check the `ModelNew` signature. It should accept `(dim, num_components=3)` and behave like `Model`.
```python
class ModelNew(nn.Module):
def __init__(self, dim, num_components=3):
super().__init__()
self.dim = dim
self.num_components = num_components
self.register_buffer("mix_weights", torch.ones(num_components) / num_components)
mus = torch.randn(num_components, dim) * 2.0
self.register_buffer("mix_mus", mus)
self.register_buffer("mix_stds", torch.ones(num_components, dim))
self.register_buffer("q_mu", torch.zeros(dim))
self.register_buffer("q_std", torch.ones(dim) * 3.0)
# Precompute
self.register_buffer("mix_inv_stds", 1.0 / self.mix_stds)
self.register_buffer("mix_log_stds", torch.log(self.mix_stds))
self.register_buffer("mix_log_weights", torch.log(self.mix_weights))
self.register_buffer("q_inv_std", 1.0 / self.q_std)
self.register_buffer("q_log_std", torch.log(self.q_std))
def forward(self, samples):
return cuda_module.importance_sampling_forward(
samples,
self.mix_mus,
self.mix_inv_stds,
self.mix_log_stds,
self.mix_log_weights,
self.q_mu,
self.q_inv_std,
self.q_log_std
)
```
One issue: if `q_std` contains zeros, `1.0 / q_std` would be inf. But `q_std = torch.ones(dim) * 3.0`, so it's fine.
Another issue: in the reference, the buffers are created in `__init__`, and then `forward` is called. In `ModelNew.__init__`, I create the same buffers and additional precomputed buffers. This should be fine.
Let me also consider: should I handle fp16 or other dtypes? The problem says "produce numerically equivalent outputs at fp32 precision". The inputs are likely fp32. I'll check `samples.dtype` and ensure everything matches. The reference model uses default fp32 for `torch.randn`. So my kernel can assume float.
One potential issue: `load_inline` might fail due to compilation errors. Let me make sure the CUDA code is syntactically correct.
Checklist:
- All `__global__` kernels have proper syntax
- C++ wrapper function uses correct ATen types
- No `PYBIND11_MODULE` block
- `functions=['importance_sampling_forward']` lists the correct wrapper
Let me make sure `#include <torch/extension.h>` is included. Yes.
Also, I should make sure to use `AT_CUDA_CHECK` or similar for error checking? Not strictly necessary for correctness, but could help debugging. However, `load_inline` examples often don't use it. I'll skip it for simplicity.
One more thing: I might want to use `torch::Tensor` const references in the C++ wrapper. The pybind11 bindings generated by `load_inline` expect specific signatures. For `std::vector<torch::Tensor> func(torch::Tensor, ...)` with by-value tensors, it should work.
Wait, let me check: pybind11 typically binds functions with by-value `torch::Tensor`. Passing by const reference might also work since pybind11 handles it. I'll use by-value to be safe.
OK, writing the code now.
Let me also think about whether I should handle the case where the model or inputs are on CPU. The evaluator likely puts everything on GPU. I'll add `.contiguous()` just in case, but actually, since the model buffers are contiguous and inputs from `get_inputs()` are contiguous, I don't need it.
Actually, let me check: the evaluator might not call `.cuda()` before calling `forward`. Looking at typical evaluation setups:
```python
model = Model(*get_init_inputs())
inputs = get_inputs()
# Might move to GPU
if torch.cuda.is_available():
model = model.cuda()
inputs = [x.cuda() for x in inputs]
output = model(*inputs)
```
So the model and inputs are on GPU when `forward` is called.
In my `ModelNew`, the `forward` just passes tensors to the CUDA kernel. If they're on CPU, the kernel launch will fail. But since the evaluator moves them to GPU, it's fine.
One final consideration: I should ensure all tensors passed to the kernel are contiguous. In the reference, the buffers are contiguous. The input `samples` from `get_inputs()` is contiguous. But if someone passes a non-contiguous tensor, my kernel would access it incorrectly. For robustness:
```python
def forward(self, samples):
if not samples.is_contiguous():
samples = samples.contiguous()
return cuda_module.importance_sampling_forward(...)
```
This adds a tiny overhead but improves robustness.
Also, I should ensure `samples` is float32. The problem guarantees fp32 precision, but just in case:
```python
if samples.dtype != torch.float32:
samples = samples.float()
```
Hmm, but this would change the type and might not match the reference if the reference also handles dtype... No, the reference doesn't handle dtype either. And the problem says fp32 precision. I'll assume inputs are fp32.
OK, let me write and compile.