reasoning
Let me analyze this problem carefully.
The model does:
1. Conv2d: in_channels=8, out_channels=32, kernel_size=3
2. Add bias: shape (out_channels, 1, 1)
3. Multiply by scale: shape (out_channels, 1, 1)
4. Sigmoid
5. GroupNorm: num_groups=8, out_channels=32
Input shape: [128, 8, 256, 256]
I need to fuse these operations into a custom CUDA kernel for performance.
Let me think about the dimensions:
- Batch: 128
- In channels: 8
- Out channels: 32
- Height: 256
- Width: 256
- Kernel: 3x3
The conv output is [128, 32, 256, 256].
For group norm with num_groups=8 and out_channels=32, each group has 4 channels.
This is a classic case for kernel fusion. I should:
1. Compute conv
2. Apply bias, scale, sigmoid
3. Compute group norm statistics
4. Normalize
But doing all in one kernel might be complex. Let me think about a good strategy.
Strategy 1: Use cuDNN for conv (highly optimized), then custom kernel for the rest.
Strategy 2: Custom conv + everything else fused.
Given the small conv parameters (8->32, 3x3), a custom implicit GEMM or direct conv might work, but cuDNN is likely very optimized. However, fusing the rest can save memory bandwidth.
Actually, for H100, the best approach might be:
1. Use cuDNN for conv (create a wrapper or just use nn.Conv2d internally, but that defeats fusion)
2. Or implement a custom fused conv+bias+scale+sigmoid+groupnorm kernel
Wait, I can use torch's conv and then a custom kernel for element-wise + group norm. But that costs memory bandwidth.
Let me think about a single fused kernel:
- For each output pixel, compute conv(3x3x8 = 72 mult-adds)
- Add bias, multiply scale, sigmoid
- Then for group norm, need mean/var across H*W for each group of 4 channels
The issue is that group norm needs reduction across spatial dimensions per group. This is tricky to fuse with conv unless we tile properly.
A better approach might be:
1. Custom conv kernel that outputs to shared memory / registers
2. But conv writes are scattered, while group norm needs spatial reduction.
Actually, let me think about a tile-based approach:
- Process a tile of output (e.g., multiple channels, block of H*W)
- For each group of 4 channels, accumulate mean and variance across the spatial tile
- Then normalize
But the mean/var need to be across ALL spatial locations for each group, so we'd need a two-pass approach or tree reduction across blocks.
Two-pass approach:
- Kernel 1: Conv + bias + scale + sigmoid, write output to intermediate buffer
- Kernel 2: Group norm
But that's what PyTorch does essentially. Can we do better?
Actually, for group norm we can use warp-level reductions. If we assign each thread block to handle one (batch, group) pair and a subset of spatial locations, we can do:
- First pass: each block computes local mean/var for its spatial subset
- Second pass: global reduction of mean/var (could use atomic operations or shared buffer)
- Third pass: normalize
But this is getting complex. Let me think if there's a simpler win.
Wait - maybe I can just use TF32 tensor cores for the conv in a custom kernel, but that's complex too.
Alternative: Keep the conv as a standard convolution (PyTorch will call cuDNN), but fuse bias+scale+sigmoid+groupnorm into a single kernel. The conv output is [128, 32, 256, 256] ≈ 1GB of data at fp32. Reading and writing this intermediate is expensive.
Actually, 128 * 32 * 256 * 256 * 4 bytes = ~1GB. At 3.35 TB/s, that's ~0.3ms for a read+write. But there are multiple passes in PyTorch.
Let me compute roughly what PyTorch does:
1. Conv: reads input + weights, writes output
2. Add bias: reads output, writes output
3. Mul scale: reads output, writes output
4. Sigmoid: reads output, writes output
5. GroupNorm: reads output, computes mean/var, writes normalized output
That's a lot of memory traffic.
If I can fuse steps 2-5, that's a big win. But ideally I fuse conv too.
For a custom fused kernel, I need to think about convolution implementation. With 8 input channels and 32 output channels, 3x3 kernel:
- Weight size: 32 * 8 * 3 * 3 = 2304 floats
- This easily fits in shared memory or registers.
A direct convolution approach where each thread computes one output pixel for multiple output channels could work well.
Let me design a kernel where:
- Each block processes one (batch, group) and a tile of spatial locations
- Each group has 4 channels
- Block size: e.g., 256 threads
- Each thread handles some spatial pixels for all 4 channels in the group
For conv, each output channel needs 8*3*3=72 input mults. For 4 channels, that's 288 mults per pixel.
Actually, let me think about this differently. If each warp handles a few output channels and a few spatial locations, we can store the 3x3x8=72 input values in registers/shared memory and compute.
Hmm, but the tiling for conv is tricky with group norm.
Let me consider another approach: Use CUTLASS or implicit GEMM style for conv to get good performance, but fuse the elementwise ops.
Actually, for a 3x3 conv with padding=1 (default in PyTorch), the simplest approach might be to use the standard conv but fuse everything after it.
Wait, let me check: nn.Conv2d default padding is 0! So output is [128, 32, 254, 254].
Actually wait, let me re-read. `nn.Conv2d(in_channels, out_channels, kernel_size)` - default stride=1, padding=0. So output is 254x254.
128 * 32 * 254 * 254 ≈ 263M elements ≈ 1.05 GB at fp32.
Hmm, that's still significant memory traffic.
OK let me think about a practical kernel design:
**Approach 1: Fused conv + post-processing + group norm**
Kernel design:
- Grid: (batch_size, num_groups, tiles_over_spatial)
- Each block handles one (batch, group) and a tile of spatial locations
- For group norm, each block computes local mean and variance, then we need a global reduction
Actually, a cleaner design for group norm:
- Each block handles one (batch, group) and ALL spatial locations for that group
- Block size = 256 or 512 threads
- Spatial size per group = 254 * 254 = 64516
- Each thread handles ~64516/256 ≈ 252 pixels
- For each pixel, we need to compute conv for 4 channels
Wait, that's 252 * 4 * 72 = 72576 multiply-adds per thread. That's a lot.
Better: use more blocks. Grid = (batch_size * num_groups, spatial_tiles)
Actually for Hopper, we can use up to 128 SMs. Batch * groups = 128 * 8 = 1024, which is plenty for parallelism.
But the issue is that group norm mean/var need to be computed across all spatial locations for each group. If multiple blocks handle the same (batch, group), they need to coordinate.
Let me try a different approach: each block handles one (batch, out_channel) or pair of out_channels, not groups. Then within each block, we have fewer channels but still need group-level stats.
Actually num_groups=8, out_channels=32, so each group is 4 channels. If a block handles 1 group = 4 channels across all spatial locations, that's:
- 4 * 254 * 254 = 258,064 output values
- With 256 threads, each thread does ~1008 output values
- Each output value needs 72 multiply-adds for conv
- Total ~72,576 mult-adds per thread
That's a lot of compute but maybe OK if we can keep weights in registers.
Actually wait - H100 has insane compute (67 TFLOPS fp32, 989 TFLOPS TF32). But memory bandwidth is 3.35 TB/s. The conv is compute-bound if done well.
Let me think of using shared memory for the input tile. For a 3x3 conv, each output pixel needs a 3x3 neighborhood of all 8 input channels.
For a block of threads working on a tile of output:
- Load a tile of input into shared memory (including halo)
- Compute conv
This is the standard direct convolution approach.
Actually, maybe I should keep it simple and just do:
1. Use standard conv (or maybe a custom simple conv)
2. Fuse the rest
Let me start with a simpler approach that has high chance of working:
- Custom conv kernel that stores output to intermediate buffer
- Then a separate group norm kernel
But wait, that doesn't fuse enough. Let me think about fusing conv with the bias/scale/sigmoid and writing to intermediate, then a group norm kernel.
Even better: can I fuse conv + bias + scale + sigmoid into one kernel, then do group norm in another? That saves 3 memory round-trips.
Let me design a kernel that does:
1. Load input weights for conv (can be in constant memory / shared memory)
2. For each output position and channel, compute conv
3. Apply bias, scale, sigmoid
4. Write to output buffer
5. Then a group norm kernel reads this and normalizes
That's 1 write + 1 read for the intermediate. vs PyTorch's multiple reads/writes.
Actually, I can do even better with a single kernel if I'm clever about the group norm reduction.
For the group norm reduction across 254*254 = 64516 elements per group:
- Launch 1024 blocks (128 batches * 8 groups)
- Wait, only 1024 blocks, each doing spatial reduction. With 256 threads per block, each thread does ~252 elements.
- We can use warp shuffle for reduction.
Actually, 1024 blocks might not fully utilize H100 (max 128 SMs * 32 blocks/SM = 4096 blocks, but occupancy matters). 1024 blocks is actually quite good.
But each block doing 252 elements: we can load 252 values, accumulate mean and mean-of-squares, do warp reduction, then block reduction.
But we ALSO need to compute conv in the same kernel! That's the challenge.
If I do two passes:
- Pass 1: compute conv + bias + scale + sigmoid for all pixels. For each block handling a (batch, group), also accumulate mean and variance in shared memory.
- But wait, if I only have 1024 blocks and each handles all spatial pixels for its group, each thread does conv for ~252 pixels * 4 channels. That's ~72k mult-adds per thread. In registers, that might be slow due to register pressure.
Alternatively, I can have the first kernel compute conv + elementwise and write to global memory. Then a second kernel does group norm.
Or: use grid-stride loops. Each block handles a few pixels for many channels... no, group norm groups channels.
OK, let me think about a practical efficient two-kernel approach:
**Kernel 1: Fused Conv + Bias + Scale + Sigmoid**
- Standard direct conv but fused with elementwise ops
- This avoids reading/writing the conv output multiple times
**Kernel 2: Group Norm**
- Standard efficient group norm
Is this enough to beat PyTorch? Let's see what PyTorch does:
- Conv: cuDNN (very fast)
- Bias add: elementwise kernel
- Scale mul: elementwise kernel
- Sigmoid: elementwise kernel
- GroupNorm: probably multiple kernels (reduce, normalize)
So PyTorch has:
1. conv (fast)
2. bias add (memory bound)
3. scale mul (memory bound)
4. sigmoid (memory bound)
5. group norm reduce (memory bound or compute)
6. group norm normalize (memory bound)
My two-kernel approach:
1. fused conv+bias+scale+sigmoid (similar to conv but with extra compute, no extra memory)
2. group norm
The win comes from eliminating 3 memory round trips (bias, scale, sigmoid as separate ops). The conv itself might be slightly slower if my custom conv is worse than cuDNN.
But for 8->32 channels, 3x3 kernel, 254x254 spatial, cuDNN might not be at peak efficiency due to small channel counts. A custom kernel might actually do well.
Let me think about conv implementation:
- Input: [128, 8, 256, 256]
- Output: [128, 32, 254, 254] (no padding)
- Weights: [32, 8, 3, 3]
For direct conv, each output pixel at (n, c, h, w) is:
sum over (ic, kh, kw) of input[n, ic, h+kh, w+kw] * weight[c, ic, kh, kw]
I can implement this with:
- Each thread block handles a batch item and a range of output channels
- Each thread within the block handles multiple output pixels
- Use shared memory for input tiles to reduce global memory reads
Standard tiled conv:
- For a tile of output (TH x TW), we need input tile of (TH+2 x TW+2)
- Load input tile to shared memory
- Compute
For output channels, we can have each thread compute multiple channels.
Let me design Kernel 1 more concretely:
```cuda
__global__ void fused_conv_kernel(
const float* __restrict__ input,
const float* __restrict__ weight,
const float* __restrict__ bias,
const float* __restrict__ scale,
float* __restrict__ output,
int N, int C_in, int C_out, int H, int W, int H_out, int W_out
) {
// Block dimensions: (block_size_x, block_size_y, block_size_c)
// Or simpler: 1D block of threads, each handling a pixel for some channels
int n = blockIdx.z; // batch
int c_start = blockIdx.y * blockDim.y; // output channel start
int h_start = blockIdx.x / tiles_x * tile_h;
int w_start = blockIdx.x % tiles_x * tile_w;
// Load weights for this block's channels into shared memory
// Weight size per channel: 8*3*3 = 72 floats
// For 4 channels: 288 floats ~ 1KB
// ...
}
```
Actually, let me think about a simpler thread organization:
Grid: (N, C_out, spatial_blocks)
Where spatial_blocks = ceil(H_out * W_out / threads_per_block)
Each thread handles one output pixel for one output channel? That would require 128 * 32 * 254 * 254 = 263M threads, which is too many. Better to have each thread handle multiple spatial pixels.
Or: Grid: (N, C_out / C_per_block, spatial_blocks)
Block: (block_size) threads
Let me try: each block handles 1 batch, some output channels, and some spatial pixels.
Option A: Block = (32, 8) = 256 threads
- 32 threads along spatial dimension
- 8 channels per block
- Each thread handles a strip of spatial pixels
Hmm, this gets complicated. Let me look for a simpler approach.
What if I use a 1D grid where:
- Each block has 256 threads
- Grid has N * C_out blocks
- Each block handles all spatial pixels for one (n, c_out) pair
- Each thread handles H_out * W_out / 256 pixels
254 * 254 = 64516. 64516 / 256 = 252 pixels per thread.
Each pixel needs 72 multiply-adds.
Total per thread: 252 * 72 = 18144 mult-adds.
At 67 TFLOPS, that's ~0.27 microseconds per thread, but with memory it will be more.
Actually for 128 * 32 = 4096 blocks, with 256 threads each, that's 1M threads total. Each block is independent. Total compute: 128 * 32 * 254 * 254 * 8 * 3 * 3 * 2 = 4.5 GFLOPS. At 67 TFLOPS, this should take ~67 microseconds if compute bound. But with memory reads it's likely memory bandwidth bound.
Input reads: each input pixel is read by multiple output pixels. The total memory traffic depends on caching.
Actually, for a naive direct conv where each thread reads its own 3x3x8 input values:
- Input reads: 128 * 32 * 254 * 254 * 8 * 3 * 3 * 4 bytes = ~18 GB (but with cache, much less)
- Weight reads: 32 * 8 * 3 * 3 * 4 = 9 KB (negligible, cached)
- Output writes: 128 * 32 * 254 * 254 * 4 = 1 GB
With 3.35 TB/s bandwidth, reading 18GB + writing 1GB = 19GB takes ~5.7ms. That's too slow.
But with L2 cache, input reads are cached. Each input pixel is reused by up to 9 output pixels (for 3x3 kernel). So effective input traffic is ~2GB instead of 18GB.
Total memory: ~3GB read + 1GB write = ~4GB. At 3.35 TB/s, that's ~1.2ms.
For PyTorch reference, let me estimate. The reference does conv + 3 elementwise + groupnorm. If conv takes ~1ms and the rest takes a few hundred microseconds each, total might be 2-3ms.
Actually, let me just implement a reasonable fused kernel and see.
For the conv, I'll use a tiled shared memory approach to reduce memory traffic further.
**Tiled Conv Kernel Design:**
For output tile of size TILE_H x TILE_W:
- Input tile is (TILE_H + 2) x (TILE_W + 2) x C_in
- Block has TILE_H * TILE_W threads
- Load input tile to shared memory
- Each thread computes one output pixel for all C_out channels assigned to the block
Wait, if we want each thread to compute multiple channels, the weight loading gets complex.
Alternative: each block computes a tile of output for a subset of channels. Within the block, threads collaborate to load the input tile, then each thread computes its output pixels for all channels in the block.
For C_out_per_block = 4 (one group), weight is 4 * 8 * 3 * 3 = 288 floats.
For C_out_per_block = 8, weight is 576 floats.
Let me try:
- Block: 16x16 = 256 threads handling a 16x16 spatial tile
- Each block also handles 8 output channels
- Input tile in SM: 18x18x8 = 2592 floats ~ 10KB
- Weights in SM or registers: 8 * 8 * 3 * 3 = 576 floats ~ 2.3KB
- Total SM: ~13KB per block. We can have many blocks per SM.
Each thread computes 1 output pixel for 8 output channels.
Per thread: 8 * 72 = 576 mult-adds.
Total blocks:
- Spatial tiles: ceil(254/16) * ceil(254/16) = 16 * 16 = 256 tiles
- Channel tiles: ceil(32/8) = 4
- Batch: 128
- Total: 128 * 4 * 256 = 131072 blocks
That's plenty of parallelism! 131K blocks with 256 threads each.
But wait, group norm needs to group channels. If I tile channels in groups of 8, but group norm groups are 4 channels, I need to be careful. Actually, group norm with num_groups=8 and 32 channels means 4 channels per group. So handling 8 channels per block spans 2 groups. That's fine for the conv kernel, but group norm needs different grouping.
For the conv kernel, I don't care about groups. I can output [N, C_out, H_out, W_out] in standard NCHW layout, then group norm can read it.
OK so here's my plan:
1. Kernel 1: Tiled direct conv with fused bias, scale, sigmoid
2. Kernel 2: Group norm
But wait, I should check if I can fuse group norm into the conv too. The challenge is the spatial reduction.
One approach for fusing everything:
- Each block handles one group (4 channels) and a spatial tile
- Write partial sums (mean, variance) to a global buffer
- Then a second reduction kernel
- Then a normalization kernel
This is 3 kernels for everything, which might be slower than my 2-kernel approach due to extra kernel launches and synchronization.
Actually, there's another approach for group norm:
- Have each block handle one (batch, group) and compute conv for ALL spatial pixels for that group
- Accumulate mean and variance on the fly
- After block-level reduction, broadcast mean/var and normalize
- Write normalized output
This is a single kernel for everything! Let me check if it's feasible.
For one (batch, group):
- 4 output channels
- 254 * 254 = 64516 spatial pixels
- Block size: 256 threads
- Each thread: 64516 / 256 = 252 pixels
- Per pixel: conv for 4 channels = 4 * 72 = 288 mult-adds
- Total per thread: 252 * 288 = 72600 mult-adds
That's a lot of computation per thread (~73K FMAs). In a block, with 256 threads, we need:
- Each thread iterates over 252 pixels
- For each pixel and each of 4 channels, compute conv
The conv weights (4 channels * 8 input channels * 3 * 3 = 288 floats) can be in shared memory or registers.
The input values need to be loaded from global memory. Each thread loads 252 pixels, but each pixel requires a 3x3 neighborhood, so total unique input loads per thread is about (254 * 254 * 8) / 256 = 2020 floats per thread. With all threads in a block, total input reads per block = 64516 * 8 = 516K floats ~ 2MB for the block. But with shared memory tiling, we can reduce this.
Actually, the issue with this single-kernel approach is that each thread handles 252 spatial pixels that are scattered or strided. If we use a grid-stride loop where each thread handles pixels spaced by blockDim.x * gridDim.x, the memory access pattern is not coalesced.
Better: assign contiguous pixels to each thread. But then shared memory tiling doesn't work well because each thread's pixels are far apart.
Let me reconsider. What if:
- Each warp handles a contiguous row or block of pixels?
- Or use shared memory for a tile, process it, then move to next tile?
For a block handling all 64516 pixels for one (batch, group):
- We can tile spatially: process 16x16 tiles one at a time
- For each tile, load input to shared memory, compute conv for 4 channels, accumulate stats in registers
- At the end, reduce stats and normalize all pixels
This could work! Let me flesh it out:
```cuda
// One block per (batch, group)
// Block: 256 threads
// Process spatial in tiles of e.g. 16x16 = 256 pixels
// Total tiles: ceil(64516 / 256) = 252 tiles
// But 254x254 doesn't divide evenly by 16x16. 16x16 tiles: 16*16 = 256 tiles cover 256*256, but we have 254x254.
Actually 254/16 = 15.875, so 16 tiles in each dimension = 256 tiles.
Each tile has min(16, remaining) x min(16, remaining) pixels.
Wait, 254 = 15*16 + 14. So we need 16 tiles in each dim, but the last tile has 14 pixels.
Total tiles: 16 * 16 = 256 tiles.
256 tiles * 256 threads = perfectly matched if each tile uses all threads!
But we only have 256 threads in a block. If each tile needs 256 threads (one per pixel), and we have 256 tiles... We'd need to launch 256 tiles sequentially within the block. That's a lot.
Actually, each tile is 16x16 = 256 pixels, and we have 256 threads, so one thread per pixel per tile. We process 256 tiles sequentially. That's feasible!
For each tile:
1. Load 18x18x8 input values to shared memory (but wait, we only need 18x18x8 = 2592 floats, but channels are 8 because conv input has 8 channels)
2. Each thread computes conv for 4 output channels at its pixel
3. Accumulate mean and M2 (for variance) in shared memory or registers
Wait, we can accumulate mean and M2 across tiles in registers or shared memory:
- Each thread has a local accumulator for mean and M2
- After all tiles, do warp/block reduction to get global mean and variance
But the issue is: each tile's input is in a different spatial location, and we need to write the normalized output. But we don't know the mean/variance until we've processed all tiles!
So we'd need to either:
1. Store all conv outputs in shared memory or registers (64516 * 4 floats = 258K floats = 1MB - too big for SM)
2. Do two passes over the input: first pass compute stats, second pass compute conv and normalize
3. Store conv output to global memory, then normalize in place
Option 3 is basically my two-kernel approach. For option 2, we do conv twice! That's double compute.
Actually, we can do:
- Pass 1: compute conv + elementwise, store to global memory. During this pass, also atomically accumulate global mean and variance stats.
- Pass 2: read conv output and normalize using the global stats.
But atomic adds for mean and variance across thousands of blocks have contention issues.
Alternatively, use a two-level reduction:
- Pass 1: compute conv, write to global. Each block also writes local mean/var to a global buffer.
- Pass 2: reduce all local mean/var to global mean/var.
- Pass 3: normalize using global mean/var.
That's 3 passes with a global buffer for local stats. The buffer size is num_blocks * num_groups * 2 floats. If we have 128 * 32 * 256 = 1M blocks (from the tiled conv), that's 2M floats for stats. Not too bad.
But wait, we don't need 1M blocks for group norm. Group norm only cares about 128 * 8 = 1024 groups. We can structure the conv so each block writes stats for one group.
Hmm, this is getting complex. Let me just go with the two-kernel approach:
1. Fused conv + bias + scale + sigmoid
2. Single-pass group norm (which internally may do reduce + normalize, but I write it as one efficient kernel)
Actually, I can do group norm in one kernel using warp-level reductions:
```cuda
__global__ void group_norm_kernel(
float* data, int N, int C, int H, int W, int num_groups,
float* gamma, float* beta, float eps
) {
int n = blockIdx.x / num_groups;
int g = blockIdx.x % num_groups;
int c_per_group = C / num_groups;
int c_start = g * c_per_group;
// Each block handles one (n, g) and spatial locations via grid-stride loop
// First pass: compute mean and variance
float sum = 0.0f;
float sum_sq = 0.0f;
int numel = c_per_group * H * W;
for (int idx = threadIdx.x; idx < numel; idx += blockDim.x) {
int c = idx / (H * W) + c_start;
int hw = idx % (H * W);
int h = hw / W;
int w = hw % W;
float val = data[((n * C) + c) * H * W + h * W + w];
sum += val;
sum_sq += val * val;
}
// Warp reduction
for (int offset = 16; offset > 0; offset /= 2) {
sum += __shfl_down_sync(0xffffffff, sum, offset);
sum_sq += __shfl_down_sync(0xffffffff, sum_sq, offset);
}
// Block reduction
__shared__ float shared_sum[32];
__shared__ float shared_sq[32];
int lane = threadIdx.x % 32;
int warp_id = threadIdx.x / 32;
if (lane == 0) {
shared_sum[warp_id] = sum;
shared_sq[warp_id] = sum_sq;
}
__syncthreads();
if (warp_id == 0) {
sum = (threadIdx.x < blockDim.x / 32) ? shared_sum[lane] : 0.0f;
sum_sq = (threadIdx.x < blockDim.x / 32) ? shared_sq[lane] : 0.0f;
for (int offset = 16; offset > 0; offset /= 2) {
sum += __shfl_down_sync(0xffffffff, sum, offset);
sum_sq += __shfl_down_sync(0xffffffff, sum_sq, offset);
}
if (lane == 0) {
shared_sum[0] = sum;
shared_sq[0] = sum_sq;
}
}
__syncthreads();
float mean = shared_sum[0] / numel;
float var = shared_sq[0] / numel - mean * mean;
float inv_std = rsqrtf(var + eps);
// Second pass: normalize
for (int idx = threadIdx.x; idx < numel; idx += blockDim.x) {
int c = idx / (H * W) + c_start;
int hw = idx % (H * W);
int h = hw / W;
int w = hw % W;
int data_idx = ((n * C) + c) * H * W + h * W + w;
float val = data[data_idx];
float normalized = (val - mean) * inv_std;
// Apply gamma and beta if provided
// For PyTorch GroupNorm, we have learnable per-channel gamma and beta
// But in the reference, group_norm is nn.GroupNorm which has weight and bias
data[data_idx] = normalized * gamma[c] + beta[c];
}
}
```
Wait, in the reference model, the group norm is `self.group_norm = nn.GroupNorm(num_groups, out_channels)`. PyTorch's GroupNorm has affine parameters by default (`affine=True`), which means it has per-channel weight and bias.
So yes, I need gamma and beta per channel.
For the two-kernel approach:
- Kernel 1: conv + bias + scale + sigmoid
- Kernel 2: group norm as above
This is clean and likely to be faster than PyTorch.
But I wonder if I can do even better. What about fusing the conv with group norm using a persistent kernel approach? Or using Tensor Cores?
Actually, let me think about using CUTLASS or a simple implicit GEMM for conv. The conv dimensions:
- N = 128
- H = W = 254
- C_in = 8
- C_out = 32
- K = 3
For implicit GEMM, we reshape conv to GEMM:
- M = N * H_out * W_out = 128 * 254 * 254 = 8,260,608
- K = C_in * 3 * 3 = 72
- N = C_out = 32
This is a very imbalanced GEMM: M is huge (~8M), K is small (72), N is small (32).
For such shapes, the optimal approach might be different. The N dimension (C_out=32) is very small. Standard GEMM tile sizes might not work well.
Actually, we can think of this as a batched GEMV (matrix-vector product):
- For each output pixel across all batches, compute weight_matrix @ input_vector
- Where weight_matrix is 32x72 and input_vector is 72x1
This is essentially 8.26M independent GEMVs of size 32x72. That's very parallel.
For such small dimensions, direct convolution or a specialized kernel might be best.
Let me design a kernel specifically for this "many small GEMVs" structure:
```cuda
// Each block handles a tile of output pixels (e.g., 64 pixels)
// Each thread handles one pixel and computes all 32 output channels
// Input: 8 channels * 3x3 = 72 values per pixel
// Weights: 32 * 72 = 2304 floats
__global__ void conv_gemv_kernel(...) {
int pixel_idx = blockIdx.x * blockDim.x + threadIdx.x;
int n = pixel_idx / (H_out * W_out);
int rem = pixel_idx % (H_out * W_out);
int h = rem / W_out;
int w = rem % W_out;
if (n >= N) return;
// Load 72 input values into registers
float input[72];
for (int i = 0; i < 72; i++) {
int ic = i / 9;
int kh = (i % 9) / 3;
int kw = i % 3;
input[i] = in[n, ic, h + kh, w + kw];
}
// Compute 32 output channels
float output[32] = {0};
for (int c = 0; c < 32; c++) {
for (int i = 0; i < 72; i++) {
output[c] += input[i] * weight[c * 72 + i];
}
output[c] = sigmoid((output[c] + bias[c]) * scale[c]);
}
// Write output
for (int c = 0; c < 32; c++) {
out[n, c, h, w] = output[c];
}
}
```
Wait, loading 72 input values per thread means 72 global memory loads. If threads in a warp access coalesced memory, that's good. The input for adjacent pixels is adjacent in memory (for width dimension), but the 3x3 windows overlap.
Actually, for pixel (h, w) and pixel (h, w+1), the input windows are offset by 1 in the W dimension. So the memory access is quite coalesced if we tile properly along W or H*W.
For 256 threads per block handling 256 adjacent pixels in the NHW layout:
- Thread 0 reads window for pixel 0
- Thread 1 reads window for pixel 1
- ...
- For the width dimension, these are adjacent pixels, so their windows overlap significantly.
But each thread independently loads 72 values. There's no shared memory reuse of the overlapping windows. That's 72 * 256 = 18432 float loads per block, while the actual unique data in the 3x3 windows for 256 adjacent pixels is much less.
For 256 pixels arranged in a line (say 16x16), the input region is roughly (16+2)x(16+2) = 18x18 = 324 pixels per channel, times 8 channels = 2592 unique floats. But with the naive approach, we load 18432 floats. That's ~7x redundancy.
A tiled shared memory approach is much better for reusing input data.
Let me reconsider the tile-based conv:
```cuda
// Block: 16x16 threads = 256 threads
// Block handles a 16x16 tile of output
// Each thread computes all 32 output channels at its pixel
// Shared memory:
// Input tile: 18x18x8 floats
// Weights: 32x72 floats
__shared__ float s_input[18*18*8];
__shared__ float s_weight[32*72];
// Load weights (cooperative, all threads participate)
// Load input tile (cooperative)
// Each thread computes conv for all 32 channels
for (int c = 0; c < 32; c++) {
float sum = 0;
for (int i = 0; i < 72; i++) {
sum += s_input[...] * s_weight[c*72 + i];
}
out[c] = sigmoid((sum + bias[c]) * scale[c]);
}
```
This is much better! The input tile (18x18x8 = 2592 floats) is loaded once and reused by all threads. The weights (2304 floats) are also loaded once.
Total SM per block: ~2592 + 2304 = ~4912 floats ≈ 20KB. That's very comfortable.
Number of blocks:
- Spatial tiles: ceil(254/16) * ceil(254/16) = 16 * 16 = 256
- Batch: 128
- Total: 128 * 256 = 32768 blocks
Each block has 256 threads. Total thread count: 8.4M threads.
Each thread computes 32 * 72 = 2304 mult-adds.
Total compute: 32768 * 256 * 32 * 72 * 2 = 38.5 GFLOPS? Wait, let me recalculate.
128 * 32 * 254 * 254 * 8 * 3 * 3 = 128 * 32 * 64516 * 72 = 18,957,467,136 mult-adds ≈ 38 GFLOPs (counting multiply-add as 2 FLOPs).
Actually multiply-add is 1 FMA = 2 FLOPs, but in terms of "operations", a conv with K*C_in*Kh*Kw mult-adds per output element is:
Total FMAs = N * C_out * H_out * W_out * C_in * Kh * Kw
= 128 * 32 * 254 * 254 * 8 * 3 * 3
= 18,957,467,136 FMAs
At 67 TFLOPS FP32, this is ~283 microseconds if purely compute bound. Plus elementwise and memory.
Memory traffic with tiling:
- Input reads: each block loads 18x18x8 floats. Total input reads = 32768 * 2592 * 4 bytes = 340 MB. But with L2 cache reusing overlaps between tiles, actual reads from DRAM will be less. Actually adjacent tiles share edges, but since we process tile by tile, the overlapping regions are reloaded. Total unique input data = 128 * 8 * 256 * 256 * 4 = 256 MB. So ~256 MB input reads (good!).
- Weight reads: 32 * 8 * 3 * 3 * 4 = 9 KB, cached.
- Output writes: 128 * 32 * 254 * 254 * 4 = 1050 MB.
- Total memory: ~1.3 GB. At 3.35 TB/s, ~0.4ms.
Group norm memory traffic:
- Read output + write output = 2 * 1.05 GB = 2.1 GB
- Plus gamma/beta reads (128 B, negligible)
- At 3.35 TB/s, ~0.6ms.
Total theoretical: ~1ms.
PyTorch reference:
- Conv (cuDNN): maybe ~0.3-0.5ms (highly optimized)
- Bias add: read+write 1GB = 0.3ms
- Scale mul: read+write 1GB = 0.3ms
- Sigmoid: read+write 1GB = 0.3ms
- GroupNorm: multiple passes, maybe 0.5-1ms
- Total: ~1.7-2.9ms
So my fused kernel might save ~0.6-0.9ms from eliminating the intermediate reads/writes. The conv might be slightly slower than cuDNN, so net improvement might be modest but hopefully positive.
To do even better, I should try to also eliminate the group norm read by fusing it. Let me think again about single-kernel fusion.
For fusing group norm with conv, the main issue is that group norm needs mean and variance across spatial BEFORE normalization. But we can compute conv and immediately accumulate statistics in a single pass if we do:
1. First pass over all blocks for a (batch, group): compute conv outputs and store to global memory. Also each block writes its local sum and sum-of-squares to a global scratchpad.
2. Second pass: reduce local sums to global mean/var.
3. Third pass: read conv outputs, normalize, write back.
This is 3 kernels and 2 round trips to global memory for the conv output. Not better than 2 kernels.
But wait, there's another approach for group norm when we know all values are produced by the same block... If a single block handles all pixels for a (batch, group), it can compute stats and normalize without global memory!
As I calculated before:
- One (batch, group) has 4 channels * 254 * 254 = 64516 values
- If one block handles all of them, we need 64516 / 256threads ≈ 252 values per thread
- But each value requires conv computation (72 mult-adds per channel * 4 channels = 288 FMAs)
- Total per thread: 252 * 288 = 72576 FMAs
This is feasible! Let's think about the memory:
- Each block loads the entire input region for the 4 output channels.
- For 4 output channels at 254x254 spatial with 3x3 kernel, the input region is 256x256x8 = 524288 floats = 2MB.
- With shared memory tiling: we can process in spatial tiles of 16x16. For each tile, load 18x18x8 input. Over the full spatial area, total unique input loaded ≈ 256x256x8 = 524K floats (since we cover the whole area).
Wait, but each group only produces 4 channels. All 4 channels share the same input! So loading input once and reusing for all 4 channels is very efficient.
Now, the challenge is that one block handling 64516 output pixels with only 256 threads means each thread produces ~252 pixels. But we tile spatially, so:
For a block handling all spatial pixels:
- Tile size: 16x16 = 256 pixels per tile
- Number of tiles: ceil(254/16)^2 = 256 tiles
- Wait, that's 256 tiles total. If each tile takes some time, doing 256 tiles serially in one block is slow.
Hmm, but 256 tiles with 256 threads each = 65536 thread-instances of work. A single block has only 256 threads, so it processes tiles sequentially. If each tile is quick, 256 tiles might still be OK.
Actually, let me compute: for each 16x16 tile:
- Load 18x18x8 = 2592 floats to shared memory
- Each thread computes 4 channels * 72 weights = 288 FMAs
- Total block FMAs per tile: 256 * 288 = 73728 FMAs
- At ~100 GFLOPS per SM (conservative), 73728 FMAs = ~0.74 microseconds per tile
- 256 tiles = ~190 microseconds per block
Plus memory loads. Each tile loads 2592 floats from global memory. 256 tiles * 2592 * 4 bytes = 2.6 MB per block.
At 3.35 TB/s shared across SMs, but if we only have 1024 blocks and they run on 128 SMs, about 8 blocks per SM. Memory bandwidth per SM is ~26 GB/s. So 2.6MB / 26GB/s = 100 microseconds per block.
So total per block might be ~300 microseconds. With 1024 blocks on 128 SMs, that's ~8 blocks per SM, total ~2.4ms if perfectly sequential. But blocks run in parallel, so wall clock might be ~300 microseconds * (1024/128) = ~2.4ms? No, that's if they all run sequentially. Actually all 128 SMs can run blocks in parallel, and with good occupancy, total time might be similar to one block's time if there are enough warps to hide latency. But 1024 blocks / 128 SMs = 8 waves. At 300us per block, total is about 2.4ms.
That's worse than my tiled approach!
So the tiled approach with many small blocks is better for parallelism.
OK, I'll stick with the two-kernel approach:
1. Tiled conv with fused bias/scale/sigmoid
2. Group norm
Now let me write the actual CUDA code.
For Kernel 1 (Conv + bias + scale + sigmoid):
```cuda
#define TILE_H 16
#define TILE_W 16
__global__ void fused_conv_kernel(
const float* __restrict__ input,
const float* __restrict__ weight,
const float* __restrict__ bias,
const float* __restrict__ scale,
float* __restrict__ output,
int N, int C_in, int C_out, int H, int W, int H_out, int W_out
) {
// Block dimensions: (TILE_W, TILE_H) - one thread per output pixel in tile
// gridDim.x = tiles_x * tiles_y? Or separate dims?
int bx = blockIdx.x; // spatial tile index
int by = blockIdx.y; // output channel block
int bz = blockIdx.z; // batch
int tile_x = bx % ((W_out + TILE_W - 1) / TILE_W);
int tile_y = bx / ((W_out + TILE_W - 1) / TILE_W);
int c_out_start = by * C_OUT_PER_BLOCK; // e.g., 8
int tx = threadIdx.x % TILE_W;
int ty = threadIdx.x / TILE_W;
int h = tile_y * TILE_H + ty;
int w = tile_x * TILE_W + tx;
// Load weights to shared memory
// For C_OUT_PER_BLOCK channels, each with C_in*3*3 weights
__shared__ float s_weight[C_OUT_PER_BLOCK * C_in * 3 * 3];
// Cooperative weight loading
int total_weights = C_OUT_PER_BLOCK * C_in * 3 * 3;
for (int i = threadIdx.x; i < total_weights; i += blockDim.x) {
int c = i / (C_in * 3 * 3) + c_out_start;
int rem = i % (C_in * 3 * 3);
s_weight[i] = weight[c * C_in * 3 * 3 + rem];
}
// Load input tile to shared memory
// Tile includes halo: (TILE_H+2) x (TILE_W+2) x C_in
__shared__ float s_input[(TILE_H+2) * (TILE_W+2) * C_in];
int h_in_start = tile_y * TILE_H;
int w_in_start = tile_x * TILE_W;
// Each thread loads multiple input elements
int total_input = (TILE_H + 2) * (TILE_W + 2) * C_in;
for (int i = threadIdx.x; i < total_input; i += blockDim.x) {
int c = i / ((TILE_H+2) * (TILE_W+2));
int rem = i % ((TILE_H+2) * (TILE_W+2));
int ih = rem / (TILE_W+2) + h_in_start;
int iw = rem % (TILE_W+2) + w_in_start;
float val = 0.0f;
if (ih >= 0 && ih < H && iw >= 0 && iw < W) {
val = input[((bz * C_in) + c) * H * W + ih * W + iw];
}
s_input[i] = val;
}
__syncthreads();
// Compute output for valid pixels
if (h < H_out && w < W_out) {
for (int c = 0; c < C_OUT_PER_BLOCK; c++) {
float sum = 0.0f;
int c_out = c_out_start + c;
for (int ic = 0; ic < C_in; ic++) {
for (int kh = 0; kh < 3; kh++) {
for (int kw = 0; kw < 3; kw++) {
int ih = ty + kh;
int iw = tx + kw;
float in_val = s_input[((ic * (TILE_H+2)) + ih) * (TILE_W+2) + iw];
float w_val = s_weight[((c * C_in) + ic) * 9 + kh * 3 + kw];
sum += in_val * w_val;
}
}
}
// Bias, scale, sigmoid
sum = sum + bias[c_out];
sum = sum * scale[c_out];
sum = 1.0f / (1.0f + expf(-sum));
output[((bz * C_out) + c_out) * H_out * W_out + h * W_out + w] = sum;
}
}
}
```
Hmm, I need to be careful about the indexing. Let me simplify by using NCHW layout explicitly.
Actually, I should also consider using `__ldg` for read-only data and making sure the compiler can optimize the inner loops.
One more optimization: unroll the 3x3 loops since they're compile-time constants.
Also, the shared memory for input can be bank-conflict free if arranged properly. But let's start simple and see.
For C_out_per_block, with 32 total output channels, I can use 8 or 16. Let's use 8 to keep shared memory reasonable. Total blocks for channels: 4. Total blocks for spatial: 16*16=256. Total blocks for batch: 128. Total: 128 * 4 * 256 = 131072 blocks. This is plenty.
Wait, grid dimensions in CUDA are limited. gridDim.x * gridDim.y * gridDim.z <= 2^31 - 1 typically, but each dimension is limited. Let me just use 3D grid:
- gridDim.x = num_spatial_tiles (256)
- gridDim.y = num_channel_tiles (4)
- gridDim.z = batch (128)
That's well within limits.
Actually, wait. The block count is 131072, which is large but fine. However, the kernel launch overhead might be noticeable. Alternatively, I can flatten spatial into x and use y for channels, z for batch.
Let me also consider another approach: instead of a 3D grid, use 1D blocks where each block handles more spatial pixels. Or keep 3D but with the tiled approach.
For the group norm kernel:
```cuda
__global__ void group_norm_kernel(
float* __restrict__ data,
const float* __restrict__ gamma,
const float* __restrict__ beta,
int N, int C, int H, int W, int num_groups, float eps
) {
int n = blockIdx.x / num_groups;
int g = blockIdx.x % num_groups;
int c_per_group = C / num_groups;
int c_start = g * c_per_group;
int tid = threadIdx.x;
int block_size = blockDim.x;
float sum = 0.0f;
float sum_sq = 0.0f;
int numel = c_per_group * H * W;
// Grid-stride loop
for (int idx = tid; idx < numel; idx += block_size) {
int c_offset = idx / (H * W);
int c = c_start + c_offset;
int hw = idx % (H * W);
int h = hw / W;
int w = hw % W;
float val = data[((n * C) + c) * H * W + h * W + w];
sum += val;
sum_sq += val * val;
}
// Warp shuffle reduction
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
sum += __shfl_down_sync(0xffffffff, sum, offset);
sum_sq += __shfl_down_sync(0xffffffff, sum_sq, offset);
}
__shared__ float shared_mean[32];
__shared__ float shared_var[32];
int lane = tid % 32;
int warp_id = tid / 32;
int num_warps = (block_size + 31) / 32;
if (lane == 0) {
shared_mean[warp_id] = sum;
shared_var[warp_id] = sum_sq;
}
__syncthreads();
if (warp_id == 0) {
float warp_sum = (tid < num_warps) ? shared_mean[tid] : 0.0f;
float warp_sq = (tid < num_warps) ? shared_var[tid] : 0.0f;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
warp_sum += __shfl_down_sync(0xffffffff, warp_sum, offset);
warp_sq += __shfl_down_sync(0xffffffff, warp_sq, offset);
}
if (lane == 0) {
float mean = warp_sum / numel;
float var = warp_sq / numel - mean * mean;
shared_mean[0] = mean;
shared_var[0] = rsqrtf(var + eps);
}
}
__syncthreads();
float mean = shared_mean[0];
float inv_std = shared_var[0];
// Normalize
for (int idx = tid; idx < numel; idx += block_size) {
int c_offset = idx / (H * W);
int c = c_start + c_offset;
int hw = idx % (H * W);
int h = hw / W;
int w = hw % W;
int data_idx = ((n * C) + c) * H * W + h * W + w;
float val = data[data_idx];
float normalized = (val - mean) * inv_std;
data[data_idx] = normalized * gamma[c] + beta[c];
}
}
```
Wait, the gamma and beta in PyTorch's GroupNorm are learnable parameters. In the reference model, `self.group_norm` defaults to `affine=True`, so it has weight and bias of shape [out_channels].
In my `ModelNew.__init__`, I need to create `self.group_norm_weight` and `self.group_norm_bias` parameters to match. But wait, in PyTorch, `nn.GroupNorm` internally has parameters. If I want to match numerically, I need to copy these parameters over.
Actually, the evaluation probably tests that ModelNew produces the same output as Model when both are initialized the same way, or they copy weights from Model to ModelNew. Let me re-read the problem.
"Replace the forward computation of the PyTorch model below with a custom CUDA kernel."
So the evaluation probably:
1. Creates a Model instance
2. Creates a ModelNew instance with the same constructor args
3. Copies weights from Model to ModelNew (or initializes them identically)
4. Checks outputs match
For this to work, ModelNew needs to have the same parameters that can be copied. But since the reference uses `nn.Conv2d`, `nn.Parameter`, and `nn.GroupNorm`, and I want to use custom kernel, I need to make sure my parameters are compatible.
Actually, I think the simplest approach is:
- Keep `self.conv` as `nn.Conv2d` but only use its weight and bias
- Actually no, I want to replace the whole forward
Wait, I can structure ModelNew like this:
- `self.conv_weight` = copy of conv weight
- `self.conv_bias` - but in the reference, conv has its own bias, and then there's also `self.bias`. Let me check.
In the reference:
```python
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
self.bias = nn.Parameter(torch.randn(bias_shape))
self.scale = nn.Parameter(torch.randn(scale_shape))
self.group_norm = nn.GroupNorm(num_groups, out_channels)
```
`nn.Conv2d` by default has `bias=True`, so it has its own bias parameter! Then `self.bias` is ANOTHER bias, of shape (out_channels, 1, 1).
In `forward`:
```python
x = self.conv(x) # uses conv's internal bias
x = x + self.bias # adds the parameter bias
x = x * self.scale
x = torch.sigmoid(x)
x = self.group_norm(x)
```
So there are actually TWO biases: conv bias and self.bias. Wait, if the evaluation copies weights from Model to ModelNew, I need to be careful.
Actually, I think the evaluation might just call `ModelNew(*get_init_inputs())` and not copy weights. The weights are initialized randomly. For correctness check, they must match exactly.
Hmm. Let me think. The standard approach in these problems is that the evaluator creates both models and either:
1. Initializes them with the same seed
2. Copies parameters by matching names
For approach 2, if I use `self.conv = nn.Conv2d(...)` in ModelNew (but don't use it in forward), the weights would be there but unused. That's bad.
For approach 1, I can create parameters with matching names, or just rely on the same initialization order.
Actually, let me re-read the prompt: "Accept the same constructor arguments as Model." It doesn't say I need to use the same layer names. But for correctness, the evaluator probably tests with a fixed seed or copies parameters.
A safe approach: replicate the exact parameter structure of Model. Use `self.conv = nn.Conv2d(...)`, `self.bias = nn.Parameter(...)`, `self.scale = nn.Parameter(...)`, `self.group_norm = nn.GroupNorm(...)`. But in `forward`, call my custom kernel instead.
The issue is that my custom kernel implements conv + bias + scale + sigmoid + group_norm all together. If I still create `self.conv`, `self.group_norm`, etc., that's fine - the forward just won't use their `.forward()` methods, but will read their `.weight` and `.bias` parameters.
Wait, but if the evaluator copies weights from Model to ModelNew by matching parameter names, this works perfectly!
Let me structure ModelNew:
```python
class ModelNew(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, num_groups, bias_shape, scale_shape):
super().__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
self.bias = nn.Parameter(torch.randn(bias_shape))
self.scale = nn.Parameter(torch.randn(scale_shape))
self.group_norm = nn.GroupNorm(num_groups, out_channels)
# ^ these parameters ensure weight copying works, but we don't call their forward
def forward(self, x):
# Extract parameters
weight = self.conv.weight # [32, 8, 3, 3]
conv_bias = self.conv.bias # [32]
bias = self.bias # [32, 1, 1]
scale = self.scale # [32, 1, 1]
gn_weight = self.group_norm.weight # [32]
gn_bias = self.group_norm.bias # [32]
# Call custom kernel
return custom_forward(x, weight, conv_bias, bias, scale, gn_weight, gn_bias)
```
Wait, but in the reference model, `self.bias` and `self.scale` have shape `(out_channels, 1, 1)` and are added/multiplied directly. In my kernel, I should add `conv_bias` first, then `bias`, then multiply by `scale`, then sigmoid.
Actually, looking at the reference:
```python
x = self.conv(x) # conv includes its own bias
x = x + self.bias
x = x * self.scale
```
So the total bias per channel is `conv_bias[c] + bias[c, 0, 0]`. Since `bias` and `scale` are [32, 1, 1], they broadcast.
In my kernel, I can precompute `total_bias = conv_bias + bias.squeeze()` and pass that, and pass `scale.squeeze()`. Actually, I can do this in Python to simplify the CUDA code.
For group norm, PyTorch's implementation uses `eps=1e-5` by default.
Let me write a complete implementation.
Actually, wait. I should verify whether `nn.Conv2d` default bias is True. Yes, `bias=True` is the default. And the reference code doesn't pass `bias=False`, so conv has bias.
But there's a subtlety: the reference model's `self.bias` has shape `(out_channels, 1, 1)`. In the forward, `x = x + self.bias` broadcasts this over batch and spatial. So the effective bias added is `self.bias[c, 0, 0]` for channel c.
So total bias for channel c is `conv.bias[c] + self.bias[c, 0, 0]`.
For scale: `x = x * self.scale` where `self.scale` is `(out_channels, 1, 1)`, so effective scale is `self.scale[c, 0, 0]`.
In my kernel, I'll precompute:
```python
total_bias = (self.conv.bias + self.bias.squeeze()).contiguous()
scale_flat = self.scale.squeeze().contiguous()
```
And pass these to CUDA.
Now, should I pass conv weight in NCHW or some other layout? PyTorch conv weight is `[C_out, C_in, K, K]`. In CUDA, I should probably keep this layout.
Wait, there's one more thing. If I create `nn.Conv2d` in ModelNew but never call it, its weight initialization will run. But if the evaluator copies weights from Model, that's fine. If the evaluator initializes both with the same seed and then checks outputs, it's also fine because both have the same layers initialized in the same way.
OK, let me write the full code.
Actually, before I write the full CUDA kernel, let me think about whether I can just call PyTorch's group norm but pass the parameters directly. No, we want custom kernel for speed.
Let me consider yet another optimization: can I avoid creating the intermediate buffer entirely? If I do conv in one kernel and group norm in another, I need the intermediate. But maybe I can compute conv and write to a temporary, then group norm reads and overwrites it in place. That's what I planned.
Actually, another thought: if I don't care about the exact same parameter structure and just want speed, I could use PyTorch for conv, then a custom kernel for bias+scale+sigmoid+groupnorm. That might be faster if my custom conv is not as good as cuDNN.
Let me think about this trade-off.
For conv: 8->32 channels, 3x3, 256x256->254x254.
This is quite small. cuDNN might not achieve peak performance, but it's still very optimized.
Actually, the real win in fusing conv with the elementwise ops is:
- We avoid writing conv output and then reading it 3 times
- Adds a tiny bit of extra compute (bias, scale, sigmoid) which is negligible
But if my custom conv is 2x slower than cuDNN, the memory savings from fusing might not compensate.
So maybe a hybrid approach is safer:
1. Use PyTorch's `F.conv2d` or `self.conv` for the convolution
2. Custom kernel that fuses: add bias + scale + sigmoid + group_norm
Wait, but group norm also needs a reduction. Can I fuse bias+scale+sigmoid+group_norm_reduce+group_norm_normalize?
Yes! Here's how:
- Kernel: each block handles one (batch, group)
- Threads grid-stride over all spatial and channel pixels in the group
- First loop: read conv output (after bias/scale/sigmoid), compute and accumulate mean/variance
- Reduce mean/variance across block
- Second loop: read conv output again, normalize, write back
This requires reading the conv output TWICE per pixel (once for stats, once for normalization). But compared to PyTorch which requires:
- conv write
- bias add read+write
- scale mul read+write
- sigmoid read+write
- group norm stats read
- group norm normalize read+write
Total PyTorch memory for conv output: 2 (conv) + 2 (bias) + 2 (scale) + 2 (sigmoid) + 1 (gn stats) + 2 (gn normalize) = 11 round trips! (Actually, some of these might be fused by PyTorch's JIT if it's enabled, but in eager mode they're separate.)
My fused kernel: 2 reads + 1 write for the group norm part. Plus conv's input reads and output write.
So the memory traffic for the "post-conv" part is:
- PyTorch: 7 reads + 5 writes (approx)
- My fused: 2 reads + 1 write (the conv output is read twice and written once)
But wait, PyTorch's group norm implementation might internally fuse some of its operations! PyTorch's native group norm is actually quite optimized and uses a custom kernel that fuses the forward pass. Let me check mentally...
In PyTorch, `nn.GroupNorm` calls `F.group_norm`, which in turn calls a custom ATen kernel. The ATen implementation of group norm might be a single kernel or a few kernels. Even if it's a few kernels, it's likely optimized.
But the bias add, scale mul, and sigmoid are definitely separate eager-mode operations in the reference, so they each read and write memory.
My hybrid approach:
1. PyTorch conv (fast)
2. Custom kernel: reads conv output, applies bias+scale+sigmoid, computes group norm stats, normalizes, writes final output
In this custom kernel, I read conv output twice (once for stats accumulation, once for normalization). But I could also compute bias+scale+sigmoid on first read, store in registers/shared memory, accumulate stats, then wait for global mean/var, then normalize and write.
Wait, I can't store all values in shared memory because 64516 floats * 4 bytes = 258KB > 227KB shared memory limit.
But I can compute bias+scale+sigmoid ON THE FLY in both passes! Since bias, scale, sigmoid are element-wise and deterministic, computing them twice costs a few FLOPs but saves shared memory.
Actually, wait. If I read the conv output, I can immediately apply bias+scale+sigmoid. But if I do this in both the stats pass and the normalize pass, I'm computing sigmoid twice.
Is that OK? Sigmoid is just a few ops. 64516 * 128 batches * 8 groups = 66M sigmoids. At H100's compute, that's negligible.
So:
- Pass 1 (stats): for each pixel in group, compute conv_val + bias + scale, sigmoid, accumulate sum and sum_sq
- Reduce to get mean, var
- Pass 2 (normalize): for each pixel, compute conv_val + bias + scale, sigmoid, normalize, apply gn_weight/gn_bias, write
This reads conv output twice. Memory: 2 reads * 1GB + 1 write * 1GB = 3GB.
PyTorch approach:
- conv: write 1GB
- bias add: read 1GB, write 1GB
- scale mul: read 1GB, write 1GB
- sigmoid: read 1GB, write 1GB
- group norm: read 1GB (for stats), read 1GB (for normalize), write 1GB
- Total: 9GB of memory traffic for the post-conv part!
So my fused kernel saves 6GB of memory traffic. At 3.35 TB/s, that's ~1.8ms saved. Even if my conv is identical to PyTorch's, this should be a big win.
BUT, there's a catch: PyTorch's conv output might be cached, so some of those reads are from L2 cache, not DRAM. Still, there's a lot of traffic.
Also, in practice PyTorch might fuse some operations if the tensors are contiguous and certain conditions are met. But for standard eager mode, these are separate kernels.
OK, I'll go with the hybrid approach:
1. Use `F.conv2d` for convolution (reliable, fast)
2. Custom kernel that fuses: (bias from conv + param bias) * scale -> sigmoid -> group norm
Wait, but `F.conv2d` also takes a bias. So I can compute `total_bias = conv.bias + self.bias` in Python and pass it to `F.conv2d`! That fuses the first two biases into one operation.
Hmm, but `F.conv2d` doesn't accept a separate bias; it uses the bias from the conv module. Actually, I can use `F.conv2d(input, weight, bias=total_bias, ...)`.
This fuses conv + bias_add into one cuDNN call. Then I still have scale, sigmoid, group norm.
Alternatively, I can just do conv without bias, then apply total_bias in my custom kernel along with scale and sigmoid.
Actually, let me think about cuDNN. For inference, `conv2d` with bias might be fused internally by cuDNN. But in PyTorch eager mode, `nn.Conv2d` does the conv and adds bias separately? No, `nn.Conv2d` computes `F.conv2d(input, weight, bias, ...)` in one call, so cuDNN can fuse conv+bias.
If I use `F.conv2d(x, self.conv.weight, self.conv.bias + self.bias.squeeze())`, that's still one call to cuDNN.
Then my custom kernel only needs to do: scale * sigmoid * group_norm.
Wait, the reference is:
```python
x = self.conv(x) # conv + conv_bias
x = x + self.bias # add param bias
x = x * self.scale
x = torch.sigmoid(x)
x = self.group_norm(x)
```
If I do:
```python
x = F.conv2d(x, self.conv.weight, self.conv.bias + self.bias.view(-1))
x = x * self.scale.view(-1, 1, 1)
x = torch.sigmoid(x)
# then custom group norm
```
This is fewer ops but still has scale, sigmoid, and group norm separate.
Alternatively, I could write a custom kernel that reads the conv output, applies scale and sigmoid, and does group norm. This reads conv output twice.
But what if I read the conv output ONCE, store the sigmoid result in a persistent way... I can't because of the group norm reduction.
Unless I use the two-pass approach with recomputation. The cost of two reads + two sigmoids is acceptable compared to PyTorch's approach.
Let me compare:
1. **Custom conv + elementwise + groupnorm (single kernel)**: Don't do this, too complex.
2. **PyTorch conv + custom fused elementwise + groupnorm**:
- PyTorch conv: uses cuDNN, great
- Custom kernel reads conv output, does scale+sigmoid, computes GN stats, normalizes
- Total memory: conv input + conv weights + conv output write + custom kernel reads conv twice + custom kernel writes final output
3. **Custom conv (tiled) + groupnorm**:
- Custom conv writes output
- Group norm reads and writes
- If my custom conv is not as good as cuDNN, this loses
Given the small conv size (8->32), cuDNN might not be dramatically better than a well-tiled custom kernel. But cuDNN is very mature. Let me actually estimate.
For the conv in standard PyTorch: it calls cuDNN. The implicit GEMM is M=8.26M, K=72, N=32. cuDNN might use winograd or direct conv. Winograd for 3x3 is very efficient, reducing multiplications by ~2.25x. cuDNN almost certainly uses winograd here.
Winograd F(2x2, 3x3) transforms the 3x3 kernel to 4x4 in transform space, reducing mults from 9 to 4 per output pixel per channel. cuDNN's implementation is highly optimized.
A custom direct conv would do all 72 mults per output channel without winograd optimization. So my custom conv might be ~2x slower than cuDNN for the conv part alone.
With the fused elementwise saving 3 intermediate passes (~0.9ms), but conv being 2x slower (~0.3ms vs 0.6ms extra), net might still win but it's closer.
Using PyTorch conv + custom fused post-processing:
- Conv: ~0.3ms
- Fused kernel: ~0.9ms (3GB at 3.35TB/s)
- Total: ~1.2ms
PyTorch reference:
- Conv: ~0.3ms
- Bias add: ~0.3ms
- Scale: ~0.3ms
- Sigmoid: ~0.3ms
- GroupNorm: ~0.6ms
- Total: ~1.8ms
Speedup: 1.5x. Good enough.
Custom conv + GN:
- Custom conv (slower): ~0.6ms
- GN: ~0.6ms
- Total: ~1.2ms
Hmm, similar! Let me just do the custom conv + elementwise fusion and see. It's cleaner to have everything in CUDA. But if my conv is bad, it fails.
Actually, I just realized: I can do BOTH. I can use `F.conv2d` inside my custom forward, then immediately call a custom kernel for the rest. This is the safest approach.
But wait, PyTorch's `F.conv2d` will create an intermediate tensor. Then my custom kernel reads it. That's still an extra memory round trip compared to fusing conv with the rest.
Let me go with the **hybrid approach**:
1. `F.conv2d(x, self.conv.weight, self.conv.bias)` - fast
2. Custom kernel that fuses: `+ self.bias`, `* self.scale`, `sigmoid`, `group_norm`
Wait, actually, I can combine `self.conv.bias` and `self.bias` into one bias and pass to `F.conv2d`. So:
```python
combined_bias = (self.conv.bias + self.bias.view(-1)).contiguous()
x = F.conv2d(x, self.conv.weight, combined_bias)
```
Then my custom kernel does: `* scale`, `sigmoid`, `group_norm`.
This is very clean. But can I do even better by also fusing scale and sigmoid into the group norm kernel? Yes! My custom kernel can:
1. Read conv output
2. Multiply by scale
3. Apply sigmoid
4. Compute group norm stats
5. Normalize
6. Apply group norm weight and bias
7. Write output
This requires reading conv output twice (for stats and normalization). But as calculated, it's still much less memory than PyTorch.
Actually, wait! I can avoid reading twice if I do:
1. First kernel: conv + scale + sigmoid, write to intermediate
2. Second kernel: group norm in-place
This is 1 read + 1 write for kernel 1, and 1 read + 1 write for kernel 2. Total 2 reads + 2 writes for the post-conv data.
My "one kernel reads twice" approach: 2 reads + 1 write.
So "one kernel reads twice" saves one write! That's 1GB at 3.35TB/s = 0.3ms. Nice.
But it computes scale+sigmoid twice. The compute cost is minimal.
So: single custom kernel that reads conv output, applies elementwise, computes stats, normalizes, writes final output. This is the best memory-wise.
Let me write this kernel.
```cuda
__global__ void fused_scale_sigmoid_groupnorm_kernel(
const float* __restrict__ conv_out,
const float* __restrict__ scale,
const float* __restrict__ gn_weight,
const float* __restrict__ gn_bias,
float* __restrict__ output,
int N, int C, int H, int W, int num_groups, float eps
) {
int n = blockIdx.x / num_groups;
int g = blockIdx.x % num_groups;
int c_per_group = C / num_groups;
int c_start = g * c_per_group;
int tid = threadIdx.x;
int numel = c_per_group * H * W;
// Pass 1: compute mean and variance
float sum = 0.0f;
float sum_sq = 0.0f;
for (int idx = tid; idx < numel; idx += blockDim.x) {
int c_offset = idx / (H * W);
int c = c_start + c_offset;
int hw = idx % (H * W);
int h = hw / W;
int w = hw % W;
float val = conv_out[((n * C) + c) * H * W + h * W + w];
val = val * scale[c];
val = 1.0f / (1.0f + expf(-val));
sum += val;
sum_sq += val * val;
}
// Warp reduction
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
sum += __shfl_down_sync(0xffffffff, sum, offset);
sum_sq += __shfl_down_sync(0xffffffff, sum_sq, offset);
}
__shared__ float shared_mean[32];
__shared__ float shared_inv_std[32];
int lane = tid % 32;
int warp_id = tid / 32;
int num_warps = (blockDim.x + 31) / 32;
if (lane == 0) {
shared_mean[warp_id] = sum;
shared_inv_std[warp_id] = sum_sq;
}
__syncthreads();
if (warp_id == 0) {
float warp_sum = (tid < num_warps) ? shared_mean[tid] : 0.0f;
float warp_sq = (tid < num_warps) ? shared_inv_std[tid] : 0.0f;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
warp_sum += __shfl_down_sync(0xffffffff, warp_sum, offset);
warp_sq += __shfl_down_sync(0xffffffff, warp_sq, offset);
}
if (lane == 0) {
float mean = warp_sum / numel;
float var = warp_sq / numel - mean * mean;
shared_mean[0] = mean;
shared_inv_std[0] = rsqrtf(var + eps);
}
}
__syncthreads();
float mean = shared_mean[0];
float inv_std = shared_inv_std[0];
// Pass 2: normalize and write
for (int idx = tid; idx < numel; idx += blockDim.x) {
int c_offset = idx / (H * W);
int c = c_start + c_offset;
int hw = idx % (H * W);
int h = hw / W;
int w = hw % W;
float val = conv_out[((n * C) + c) * H * W + h * W + w];
val = val * scale[c];
val = 1.0f / (1.0f + expf(-val));
float normalized = (val - mean) * inv_std;
val = normalized * gn_weight[c] + gn_bias[c];
output[((n * C) + c) * H * W + h * W + w] = val;
}
}
```
Wait, I just realized there's an important issue. For numerical equivalence, I need to make sure the order of operations matches PyTorch. PyTorch does:
1. conv + bias
2. + bias2
3. * scale
4. sigmoid
5. group_norm
For group_norm, PyTorch computes:
- mean and var over the group
- (x - mean) / sqrt(var + eps)
- * weight + bias
My kernel replicates this exactly. But one concern: floating point associativity. Since I'm computing `conv_out * scale` and `1/(1+exp(-x))`, the results should be identical to PyTorch's element-wise operations, as long as I use the same math functions. `torch.sigmoid` in PyTorch uses a custom implementation, but `expf` and division should be very close.
Actually, there might be small numerical differences due to:
1. `expf(-x)` vs PyTorch's exp implementation
2. The order of reductions for mean/var (PyTorch uses Welford or two-pass)
For group norm, the variance reduction can have precision issues. PyTorch's native group norm might use a specific algorithm. I should try to match it as closely as possible.
PyTorch's group norm native implementation uses atomic adds or warp-level reductions in a similar way. But there might be subtle differences. I'll need to test and adjust eps if needed.
Also, in the problem statement, fp32 needs atol=rtol=1e-4. That should be achievable.
Now, one optimization: if I can fuse the conv itself, I save a memory write and read. Let me reconsider doing a custom conv.
Actually, let me look at the problem from a different angle. Can I use PyTorch's conv, but then use a custom kernel for everything else? And if so, is that enough to beat the reference?
PyTorch reference doing:
1. F.conv2d (with conv.bias)
2. x + self.bias
3. x * self.scale
4. torch.sigmoid(x)
5. group_norm(x)
Steps 2-4 each launch a separate element-wise kernel. That's 3 extra kernel launches and 3 extra memory round trips.
My approach:
1. F.conv2d with combined bias (one kernel launch)
2. Custom kernel for scale * sigmoid * group_norm (one kernel launch)
Even if my custom kernel for step 2 is only as efficient as PyTorch's individual kernels combined, I've saved 2 kernel launches and associated memory round trips. Plus PyTorch's group norm might require multiple kernels (reduce + norm), while mine is one kernel.
Wait, actually PyTorch's `F.group_norm` is a single ATen dispatch that might call one or a few kernels. Let me check mentally...
PyTorch native group norm for CUDA calls `group_norm_kernel_cuda` which likely does it in one or two passes. But in eager mode, the preceding operations (scale, sigmoid) are separate kernels.
So yes, just fusing steps 2-5 should give a good speedup.
But wait! I have an even better idea. What if I use `torch.compile` trick? No, the evaluator runs eager mode.
What about using a custom Triton kernel? The problem specifically asks for `torch.utils.cpp_extension.load_inline` with CUDA C++.
Let me now write the full Python code. I'll define a custom CUDA extension with two functions:
1. `fused_conv_scale_sigmoid_groupnorm` - does everything in one kernel
2. `fused_postprocess_groupnorm` - does scale+sigmoid+GN on conv output
Wait, I said I'll use hybrid approach. Let me code that up first.
Actually, you know what, let me try the single custom kernel for everything (conv + scale + sigmoid + groupnorm) first. Even if the conv part is not as optimized as cuDNN, the extreme fusion might win. And if not, I'll fall back to hybrid.
But coding the custom conv is a bit involved. Let me prototype it...
Actually, I realize I should start simple: write the hybrid approach first, compile it, test correctness, measure speed. Then if needed, add custom conv.
Wait, I don't have timing measurements in the loop. I can only check correctness with `run_correctness`, and the final `submit_kernel` gives timing. I need to make sure it's correct first, then submit.
Let me write a hybrid kernel (conv via PyTorch, custom kernel for scale+sigmoid+GN).
Actually, for maximum speed, let me try to also include the conv in CUDA. The fusion should be better.
Here's my design for a single fully-fused kernel:
**Grid**: One block per (batch, group)
**Block**: 256 threads
**Processing**: Use a grid-stride loop over spatial tiles
For each tile of the output:
1. Load input tile (including halo) to shared memory
2. Compute conv for all pixels and all channels in the group
3. Apply scale and sigmoid
4. Accumulate mean and variance locally (in registers)
5. After all tiles, reduce mean/variance across the block
6. Second loop over tiles: recompute conv, scale, sigmoid, normalize with global mean/var, apply GN weight/bias, write to output
Wait, recomputing conv in step 6 means we do conv twice! That's double compute. Not good.
Alternative: store conv+sigmoid outputs in a scratchpad in global memory (per-block, not per-SM). Since each block handles one (batch, group), we could allocate a workspace of size [N, num_groups, C_per_group, H, W] or just [N, C, H, W] and reuse it.
But then we're back to writing an intermediate!
Hmm. What if we process spatial in chunks small enough to fit in shared memory?
For C_per_group = 4 channels, H*W = 64516 pixels. Each pixel after conv+sigmoid is 4 floats. Total = 258K floats = 1MB. Doesn't fit in shared memory (227KB limit).
What if we process one channel at a time? 64516 floats = 258KB per channel. Still doesn't fit.
What if we process a row at a time? 254 floats = 1KB. That fits! But then we do 254 rows * 4 channels = 1016 iterations. Each row needs conv computation.
Actually, for one row of 254 pixels and 4 output channels:
- Input needed: 4 rows (because 3x3 kernel) * 256 width * 8 channels = 8192 floats = 32KB
- Conv output: 254 * 4 = 1016 floats = 4KB
- Total shared memory: ~36KB. Very comfortable!
But we need to process all rows to compute mean/var, then process all rows again to normalize.
With 256 threads, we can process 4 rows at a time or something.
Actually, here's an idea: process spatial in tiles of 64x4 or similar. Let's say 64x4 = 256 pixels per tile. For 4 channels, that's 1024 output values = 4KB of shared memory. Not enough to fit the full tile, but we can process sequentially.
Wait, for a 64x4 spatial tile, the input region is 66x6x8 = 3168 floats = 12.7KB. The conv output is 64x4x4 = 1024 floats = 4KB. Total ~17KB. This fits easily!
If we have 256 threads, each thread handles one spatial pixel for all 4 channels in a 16x16 tile... No, 64x4 = 256 pixels, one thread each.
OK here's a concrete design for the fully-fused single kernel:
```cuda
// One block per (batch, group)
// Each block processes all spatial pixels for its 4 channels
// Tile size: e.g., 64x4 spatial pixels = 256 pixels per tile
// Total tiles for 254x254: ceil(254/64) * ceil(254/4) = 4 * 64 = 256 tiles
// Wait, 254/4 = 63.5, so 64 tiles in H dimension if tile_h=4
// 254/64 = 3.97, so 4 tiles in W dimension if tile_w=64
// Total: 4 * 64 = 256 tiles * 256 pixels = 65536 pixels per block
// But we only have 64516 pixels! Some threads idle in some tiles.
// Actually, better to use 16x16 tiles: 16*16=256 pixels
// Total tiles: ceil(254/16) * ceil(254/16) = 16 * 16 = 256 tiles
// 256 tiles, 256 threads: each thread handles one pixel per tile
// Wait no, 256 threads handle one tile (16x16), and we have 256 tiles.
// So the block loops over 256 tiles. With 256 threads, each tile is handled by all threads.
__shared__ float s_input[(TILE_H+2)*(TILE_W+2)*C_in];
__shared__ float s_weights[C_per_group * C_in * 9];
__shared__ float s_mean_acc[256]; // or use registers
// Local accumulators for mean and variance
float local_sum = 0.0f;
float local_sq = 0.0f;
// Load weights once
for (i = threadIdx.x; i < 4*8*9; i += 256) {
s_weights[i] = weight[...];
}
__syncthreads();
// First pass: compute conv + sigmoid for all tiles, accumulate stats
for (int tile = 0; tile < num_tiles; tile++) {
int tile_y = tile / tiles_x;
int tile_x = tile % tiles_x;
// Load input tile to shared memory
// ...
__syncthreads();
// Each thread handles one pixel in the tile
int ty = threadIdx.x / TILE_W;
int tx = threadIdx.x % TILE_W;
int h = tile_y * TILE_H + ty;
int w = tile_x * TILE_W + tx;
if (h < H_out && w < W_out) {
for (int c = 0; c < 4; c++) {
float sum = conv_pixel(c, ty, tx, s_input, s_weights);
sum = sigmoid(sum * scale[c] + bias[c]);
local_sum += sum;
local_sq += sum * sum;
// CANNOT store to shared memory because we'd need 64516 * 4 floats
}
}
__syncthreads();
}
// Reduce local_sum and local_sq across block
// ...
// Second pass: compute conv + sigmoid again, normalize, write
for (int tile = 0; tile < num_tiles; tile++) {
// same as first pass but also normalize and write to global output
}
```
This recomputes conv twice! 38 GFLOPs * 2 = 76 GFLOPs. At 67 TFLOPS, that's ~1.1ms just for compute. Plus memory loads. Might be slower than the hybrid approach.
So the hybrid approach (PyTorch conv + custom fused rest) seems better.
Wait, unless... I compute the conv ONCE, store the result to global memory (1GB write), then the same block immediately reads it back to compute stats. That's 1 write + 1 read for the intermediate, same as doing conv in one kernel and GN in another.
But if I do that, why not just call PyTorch conv?
Actually, let me think about it this way:
- PyTorch conv writes to a tensor that persists in memory
- My custom kernel reads it twice
This is exactly the same as `F.conv2d` + custom kernel. No need to implement conv in CUDA.
SO, the optimal architecture is:
1. `combined_bias = self.conv.bias + self.bias.squeeze()`
2. `x = F.conv2d(input, self.conv.weight, combined_bias)`
3. `x = fused_scale_sigmoid_groupnorm(x, self.scale, self.group_norm.weight, self.group_norm.bias)`
In step 3, my CUDA kernel reads `x` twice and writes the final output once.
Is there any way to avoid reading `x` twice? Only if I can fit all values in shared memory, which I can't for 254x254x4.
What about reading `x` in chunks and storing in shared memory? I can process a tile at a time:
- Load a tile of `x` (e.g., 64x4x4 = 1024 floats = 4KB) to shared memory
- Compute local sum and sum_sq for those pixels
- Keep a running total in registers
But I still need the GLOBAL mean/var to normalize. I can only normalize after processing ALL pixels and reducing. So I need to either:
1. Store all pixels somewhere (can't fit in SM)
2. Read `x` twice
3. Use global atomics for a running mean/var (tricky with Welford, possible but complex)
Actually, I CAN use global atomics! Here's the idea:
- First pass: each block (handling one batch, one group, one spatial tile) computes local sum and sum_sq
- Atomically add to global sum and sum_sq
- Then each block re-reads x and normalizes
But then I need a synchronization point before normalization. I can do this with a second kernel launch or with grid-level synchronization (not portable).
Two kernel launches:
- Kernel 1a: compute local stats per tile, atomically accumulate to global buffer
- Kernel 1b: read x, normalize using global stats, write output
With atomics, the accumulation has contention. For 1024 groups and maybe 256 tiles per group, that's 262K atomic operations. On H100, atomics are relatively fast but still.
Alternatively, I can do:
- Kernel 1: each block handles one (batch, group) and ALL spatial pixels. No need for atomics if one block per group.
- But then each block has too much work and parallelism suffers.
Wait, I can have multiple blocks per group for the first pass, writing local stats to a scratchpad, then one more block per group reducing the stats. Then another set of blocks for normalization. That's 3 kernels, very complex.
I think reading `x` twice in a single kernel per group is the sweet spot.
Actually, let me reconsider block assignment. What if each block handles one (batch, group) and uses 512 threads?
- 64516 pixels / 512 = 126 pixels per thread
- Two passes: each thread processes 126 pixels
- Total blocks: 128 * 8 = 1024
1024 blocks on H100 (128 SMs) = 8 blocks per SM. If each block runs quickly, we get good parallelism.
With 512 threads per block and grid-stride loop:
- Each thread handles about 126 iterations in loop 1, 126 in loop 2
- 1024 blocks total
Let me compare with PyTorch:
- PyTorch conv launches 1 kernel
- PyTorch bias add: 1 kernel
- Scale: 1 kernel
- Sigmoid: 1 kernel
- Group norm: at least 1 kernel, probably 2
Total: 5-6 kernel launches. My approach: 2 kernel launches (conv + custom). The fewer kernel launches save overhead.
Now let me also consider: can PyTorch's `torch.jit` or something fuse these? In eager mode, no.
Let me also think about memory layout. The conv output from PyTorch is NCHW (standard). My kernel should access it in coalesced manner. In the grid-stride loop:
```cuda
for (int idx = tid; idx < numel; idx += blockDim.x) {
int c_offset = idx / (H * W);
int hw = idx % (H * W);
// c changes slowly, hw changes fast
// In NCHW, channels are outer dimension, so adjacent hw are adjacent in memory
// Within a channel, hw are contiguous
}
```
This is coalesced when threads access adjacent `idx` values, because `hw` changes fastest. Wait, let me check the memory layout:
NCHW: `data[n, c, h, w]` is at offset `((n * C + c) * H + h) * W + w`
For `idx = c_offset * H * W + hw`, where `hw = h * W + w`:
`offset = ((n * C) + c_start + c_offset) * H * W + hw`
Adjacent threads with adjacent `idx` have adjacent `hw` when `c_offset` is the same. When `idx` crosses a channel boundary (hw wraps), the threads are still accessing adjacent memory because we're just moving to the next channel.
So memory access is coalesced! Good.
For the second loop writing output, same indexing.
One more thing: I should use `rsqrtf` for the inverse square root, which is faster than `1.0f / sqrtf`.
Also, PyTorch uses `eps=1e-5` for group norm by default.
Now, regarding sigmoid: PyTorch's `torch.sigmoid(x)` computes `1 / (1 + exp(-x))`. For large positive or negative x, it might use a custom implementation to avoid overflow/underflow. But `expf` in CUDA handles reasonable ranges. The default tolerance is 1e-4 for fp32, so minor differences in the sigmoid tail shouldn't matter.
Actually, PyTorch's sigmoid in CUDA uses `1 / (1 + exp(-x))` with some fast math optimizations or maybe the `__logistic` intrinsic? I should check or just use `1.0f / (1.0f + expf(-val))`.
Let me also think about `blockDim.x`. For group norm, we need block size to be a multiple of 32 for warp shuffle. 256 or 512 are good choices. 512 gives fewer iterations per thread (126 vs 252 for 64516 pixels). Let me use 512.
Actually, 64516 / 512 = 126.01, so almost exactly 126 iterations. That's great.
For warp reduction with 512 threads, we have 16 warps. The block reduction needs to handle 16 warps.
OK, let me now draft the complete code.
One more thing: the grid dimensions. `blockIdx.x` for 128 * 8 = 1024 blocks. That's fine.
Wait, I just realized: can I use cooperative groups for block-level reduction? CUDA has built-in functions, but writing the warp shuffle manually is fine.
Let me write the code.
Actually, I want to make sure I match PyTorch's group norm exactly. In PyTorch's group norm, after computing mean and var, it does:
```python
x = (x - mean) * rsqrt(var + eps)
x = x * weight + bias
```
Where `weight` and `bias` are per-channel.
My kernel does the same.
One more detail: PyTorch uses Bessel's correction for variance? No, group norm uses population variance (divide by n, not n-1). Let me verify mentally... Yes, `nn.GroupNorm` computes population variance for normalization.
OK so variance = E[x^2] - (E[x])^2. I compute:
```cuda
float mean = sum / numel;
float var = sum_sq / numel - mean * mean;
```
This matches.
Now let me write the full Python module.
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline
cuda_source = '''
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void fused_scale_sigmoid_groupnorm_kernel(
const float* __restrict__ input,
const float* __restrict__ scale,
const float* __restrict__ gn_weight,
const float* __restrict__ gn_bias,
float* __restrict__ output,
int N, int C, int H, int W, int num_groups, float eps
) {
int n = blockIdx.x / num_groups;
int g = blockIdx.x % num_groups;
int c_per_group = C / num_groups;
int c_start = g * c_per_group;
int numel = c_per_group * H * W;
int tid = threadIdx.x;
float sum = 0.0f;
float sum_sq = 0.0f;
for (int idx = tid; idx < numel; idx += blockDim.x) {
int c_offset = idx / (H * W);
int c = c_start + c_offset;
int hw = idx % (H * W);
float val = input[((n * C) + c) * H * W + hw];
val = val * scale[c];
val = 1.0f / (1.0f + expf(-val));
sum += val;
sum_sq += val * val;
}
// Warp reduction
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1) {
sum += __shfl_down_sync(0xffffffff, sum, offset);
sum_sq += __shfl_down_sync(0xffffffff, sum_sq, offset);
}
__shared__ float shared_mean[32];
__shared__ float shared_var[32];
int lane = tid % 32;
int warp_id = tid / 32;
int num_warps = (blockDim.x + 31) / 32;
if (lane == 0) {
shared_mean[warp_id] = sum;
shared_var[warp_id] = sum_sq;
}
__syncthreads();
if (warp_id == 0) {
float s = (lane < num_warps) ? shared_mean[lane] : 0.0f;
float sq = (lane < num_warps) ? shared_var[lane] : 0.0f;
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1) {
s += __shfl_down_sync(0xffffffff, s, offset);
sq += __shfl_down_sync(0xffffffff, sq, offset);
}
if (lane == 0) {
float mean = s / numel;
float var = sq / numel - mean * mean;
shared_mean[0] = mean;
shared_var[0] = rsqrtf(var + eps);
}
}
__syncthreads();
float mean = shared_mean[0];
float inv_std = shared_var[0];
for (int idx = tid; idx < numel; idx += blockDim.x) {
int c_offset = idx / (H * W);
int c = c_start + c_offset;
int hw = idx % (H * W);
float val = input[((n * C) + c) * H * W + hw];
val = val * scale[c];
val = 1.0f / (1.0f + expf(-val));
val = (val - mean) * inv_std;
val = val * gn_weight[c] + gn_bias[c];
output[((n * C) + c) * H * W + hw] = val;
}
}
torch::Tensor fused_scale_sigmoid_groupnorm(torch::Tensor input, torch::Tensor scale,
torch::Tensor gn_weight, torch::Tensor gn_bias,
int num_groups, float eps) {
int N = input.size(0);
int C = input.size(1);
int H = input.size(2);
int W = input.size(3);
auto output = torch::empty_like(input);
int total_blocks = N * num_groups;
int threads = 512;
fused_scale_sigmoid_groupnorm_kernel<<<total_blocks, threads>>>(
input.data_ptr<float>(), scale.data_ptr<float>(),
gn_weight.data_ptr<float>(), gn_bias.data_ptr<float>(),
output.data_ptr<float>(), N, C, H, W, num_groups, eps
);
return output;
}
'''
cpp_source = '''
torch::Tensor fused_scale_sigmoid_groupnorm(torch::Tensor input, torch::Tensor scale,
torch::Tensor gn_weight, torch::Tensor gn_bias,
int num_groups, float eps);
'''
fused_ops = load_inline(
name='fused_ops',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['fused_scale_sigmoid_groupnorm'],
extra_cuda_cflags=['-O3'],
verbose=False
)
class ModelNew(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, num_groups, bias_shape, scale_shape):
super(ModelNew, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
self.bias = nn.Parameter(torch.randn(bias_shape))
self.scale = nn.Parameter(torch.randn(scale_shape))
self.group_norm = nn.GroupNorm(num_groups, out_channels)
self.num_groups = num_groups
def forward(self, x):
# Combine conv bias and external bias
combined_bias = (self.conv.bias + self.bias.view(-1)).contiguous()
# Conv with combined bias
x = F.conv2d(x, self.conv.weight, combined_bias)
# Fused scale, sigmoid, group norm
scale_flat = self.scale.view(-1).contiguous()
x = fused_ops.fused_scale_sigmoid_groupnorm(
x, scale_flat, self.group_norm.weight, self.group_norm.bias,
self.num_groups, 1e-5
)
return x
```
Wait, I need to verify the shapes:
- `self.conv.bias`: [out_channels]
- `self.bias`: [out_channels, 1, 1]
- `combined_bias`: [out_channels] - good for `F.conv2d`
- `self.scale`: [out_channels, 1, 1]
- `scale_flat`: [out_channels] - good for kernel, we index by channel
In the kernel, `scale[c]` where `c` is the actual channel index (0 to 31). `scale` is a 1D tensor of size C.
For `gn_weight` and `gn_bias`: these are also [out_channels], same indexing.
One issue: `F.conv2d` with the same arguments as `self.conv` needs `stride=1`, `padding=0`, `dilation=1`, `groups=1`. These are the defaults. But I need to make sure `self.conv` hasn't modified any of these. The constructor is `nn.Conv2d(in_channels, out_channels, kernel_size)` which uses all defaults.
Actually, to be safe, I should pass the same parameters:
```python
x = F.conv2d(x, self.conv.weight, combined_bias,
stride=self.conv.stride, padding=self.conv.padding,
dilation=self.conv.dilation, groups=self.conv.groups)
```
But the defaults are `(1, 1)`, `(0, 0)`, `(1, 1)`, `1`. I'll pass them for safety.
Another issue: I need to handle the case where the input might not be contiguous or might be on CPU. The kernel expects contiguous float tensors on CUDA. I should ensure this.
```python
x = x.contiguous()
```
Actually, `F.conv2d` output should already be contiguous. But I'll add `.contiguous()` before passing to the kernel just to be safe.
Also, I should add CUDA streams or just use the default stream. That's fine.
Let me also consider whether I need to wrap the kernel call in a `torch.cuda.Stream` context or anything. Not necessary for basic correctness.
One more thing: I need to include `#include <torch/extension.h>` in the CUDA source. But `load_inline` might handle this? Let me check typical usage. Actually, `torch/extension.h` provides `torch::Tensor` and is needed for the C++ wrapper. But the CUDA kernel itself doesn't need it.
However, `load_inline` compiles the CUDA source with nvcc, and the C++ wrapper is auto-generated based on `cpp_sources`. Actually, looking at the docs, `cuda_sources` is compiled by nvcc and linked with the C++ wrapper. The wrapper uses the signatures declared in `cpp_sources` and auto-generated pybind11 bindings.
Wait, I think I need to provide the wrapper function in `cpp_sources` that calls the kernel, and the kernel in `cuda_sources`. Let me look at typical `load_inline` usage.
Actually, in `load_inline`:
- `cpp_sources` contains C++ wrapper functions that are exposed to Python
- `cuda_sources` contains CUDA code (kernels and helper functions)
The pybind11 wrapper is auto-generated from `functions=[...]` list.
So I should put the wrapper function in `cpp_sources`:
```cpp
torch::Tensor fused_scale_sigmoid_groupnorm(torch::Tensor input, torch::Tensor scale,
torch::Tensor gn_weight, torch::Tensor gn_bias,
int num_groups, float eps) {
// ...
}
```
And the kernel in `cuda_sources`:
```cuda
__global__ void fused_scale_sigmoid_groupnorm_kernel(...) { ... }
```
But `cpp_sources` needs to reference the kernel, so the kernel should be declared in `cpp_sources` or in a header. Actually, both are compiled together, so it's fine.
Let me structure it properly:
```python
cuda_source = '''
__global__ void fused_scale_sigmoid_groupnorm_kernel(...) { ... }
'''
cpp_source = '''
torch::Tensor fused_scale_sigmoid_groupnorm(torch::Tensor input, ...) {
...
fused_scale_sigmoid_groupnorm_kernel<<<...>>>(...);
return output;
}
'''
```
Wait, but the CUDA source needs to be compiled with nvcc, and the C++ source with the C++ compiler. In `load_inline`, both are combined. The kernel in `cuda_sources` and the wrapper in `cpp_sources` are linked together.
Actually, I think I can just put everything in `cuda_sources` if it compiles as a single CUDA file. Let me check... In `load_inline`, the `cpp_sources` and `cuda_sources` are concatenated. So if I put the kernel in `cuda_sources` and the wrapper in `cpp_sources`, it should work since the wrapper can see the kernel declaration.
But to be safe, I can declare the kernel in `cpp_sources`:
```cpp
__global__ void fused_scale_sigmoid_groupnorm_kernel(...);
```
Or put everything in `cuda_sources`. Actually, `cpp_sources` might not accept `__global__` because it's compiled as C++ initially. But `load_inline` handles the compilation.
Let me look at standard examples. Typically:
```python
cpp_source = """
torch::Tensor my_func(torch::Tensor x) { ... }
"""
cuda_source = """
__global__ void my_kernel(...) { ... }
"""
```
This works because the combined file is valid CUDA C++.
For safety, I'll declare the kernel at the top of `cpp_sources`.
Actually, let me reconsider. I want to make sure `load_inline` works correctly. A common pattern is:
```python
source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void kernel(...) { ... }
torch::Tensor func(torch::Tensor x) {
kernel<<<...>>>(...);
return x;
}
"""
torch.utils.cpp_extension.load_inline(
name='ext',
cpp_sources=source,
functions=['func'],
extra_cuda_cflags=['-O3'],
verbose=True,
)
```
In this case, everything is in one string. But the API wants `cpp_sources` and optionally `cuda_sources`. I'll put the wrapper in `cpp_sources` and kernel in `cuda_sources`.
Let me just write the code and test compilation.
Actually, wait. I realize I should also consider using `torch::empty_like(input)` in the C++ wrapper. This requires `#include <torch/extension.h>`.
Let me draft the complete code now and compile it.
One more thought: should I use 1024 threads per block instead of 512? With 64516 pixels and 1024 threads, each thread handles ~63 pixels. That's fewer loop iterations. But register usage might increase. I'll try 512 first.
Also, I should use `__shared__` for the final mean and inv_std. But they only need one float each. I'll just use `shared_mean[0]` and `shared_inv_std[0]` after the reduction.
Let me write the code properly.
Wait, I also realize that dividing `blockDim.x` by 32 using integer division: `(blockDim.x + 31) / 32` is the number of warps. For 512 threads, that's 16 warps.
In the second warp reduction, I use `lane` to index into `shared_mean`. But `lane` is `tid % 32`, and `tid` ranges from 0 to 511. So `lane` ranges from 0 to 31. I need `lane < num_warps` for the conditional.
In my code:
```cpp
if (warp_id == 0) {
float s = (lane < num_warps) ? shared_mean[lane] : 0.0f;
```
This is correct because the first warp's threads 0 to 15 will read the 16 warps' partial sums.
Actually, for block size = 512, `num_warps = 16`. Thread 0 in warp 0 reads `shared_mean[0]`, thread 1 reads `shared_mean[1]`, ..., thread 15 reads `shared_mean[15]`. Thread 16-31 read 0.
Then the warp shuffle reduces across all 32 threads, but only the first 16 have real data. The shuffle with `0xffffffff` mask includes all threads in the warp. The extra threads (16-31) contribute 0, which is fine.
After reduction, only `lane == 0` writes the final result to `shared_mean[0]`.
All correct.
One more consideration: should I add `__restrict__` to pointers? Yes, for better optimization.
Also, should I compile with `-use_fast_math`? That might affect numerical accuracy. With tolerance 1e-4, maybe it's fine, but let's avoid it to be safe.
Let me write the code now.