reasoning
Let me analyze the reference model:
1. Conv2d: in_channels=8, out_channels=64, kernel_size=3
2. Divide by 2
3. LeakyReLU with negative_slope=0.01
The input is [128, 8, 128, 128].
A Conv2d with these parameters can be fused with the subsequent elementwise operations (division and LeakyReLU) into a single custom CUDA kernel for better performance.
For the conv2d, this is a standard 2D convolution. With batch_size=128, in_channels=8, out_channels=64, H=W=128, k=3.
Output dimensions: [128, 64, 126, 126] (since default padding=0, stride=1).
I should write a fused conv+divide+leaky_relu kernel. The naive approach would be a standard conv2d followed by elementwise ops, but fusing them saves a memory round-trip.
For a custom CUDA kernel, I need to implement the convolution. Given the small kernel size (3x3) and moderate sizes, I can use a shared memory tiled approach. Each thread block can load a tile of input and compute a tile of output.
Let me think about the tiling strategy:
- Input: [128, 8, 128, 128]
- Weights: [64, 8, 3, 3]
- Output: [128, 64, 126, 126]
Given out_channels=64 and output spatial size=126x126, we have 128*64*126*126 ≈ 128M output elements.
For tiling, I can tile over the output channels and spatial dimensions. One approach:
- Each thread computes one or more output elements.
- Use shared memory for input tiles.
A simpler approach that might work well for H100:
- Use im2col style in shared memory or direct convolution.
- For each output pixel, a thread loads 8*3*3 = 72 input values and computes 64 output channels.
Actually, let me think more carefully. With in_channels=8, k=3, each output pixel requires 8*3*3 = 72 multiplies per output channel, times 64 output channels = 4608 FMAs per output pixel.
Total output pixels = 128 * 126 * 126 ≈ 2M.
Total FMAs ≈ 2M * 4608 ≈ 9.4 GFLOPs.
At 67 TFLOPS, this should take ~0.14 ms if compute bound. But this small conv might be memory bandwidth bound.
Memory traffic if not fused:
- Input: 128 * 8 * 128 * 128 * 4 bytes ≈ 64 MB
- Weights: 64 * 8 * 3 * 3 * 4 bytes ≈ 18 KB
- Output: 128 * 64 * 126 * 126 * 4 bytes ≈ 512 MB
- Total read+write ≈ 576 MB + ~512 MB for intermediate = ~1 GB
- At 3.35 TB/s = ~0.3 ms.
But with fusion, we write output once and avoid reading/writing intermediate tensors.
Memory traffic with fusion: ~64 MB input + ~0.018 MB weights + ~512 MB output = ~576 MB → ~0.17 ms.
Actually wait, the weights can stay in L2 cache. Input is 64MB, weights are tiny. Output is 512MB. The compute is ~9.4 GFLOPs.
At 3.35 TB/s, 576 MB takes ~0.17 ms. At 67 TFLOPS, 9.4 GFLOPs takes ~0.14 ms. So it's roughly balanced.
I should write a fused kernel. Let me design it:
For direct convolution with fusion:
- Grid over (batch, out_channels/OC_BLOCK, oh/ OH_BLOCK, ow/ OW_BLOCK)
- Each block handles a tile.
Let me think about a reasonable block size. On H100, we want good occupancy.
Given in_channels=8, I can have each thread compute 1 output pixel for multiple output channels.
Option 1: 16x16 thread block computing a 16x16 spatial tile for a subset of output channels.
- For 16x16 spatial tile: each thread loads input for its pixel plus the halo for the 3x3 kernel.
- With in_channels=8, each thread block needs (16+2)*(16+2)*8 = 2592 floats ≈ 10KB in shared memory.
- Computes 16*16 = 256 output pixels.
- If each pixel computes all 64 output channels, that's 256*64 = 16384 output values per block.
- Each thread would compute 64 output channels for 1 pixel = 64 values.
But then each thread needs 64 registers for accumulation, which is fine. Shared memory for input tile is ~10KB, weights (64*8*3*3 = 4608 floats = 18KB) can be in shared memory too, or loaded from global/L2.
Actually, with 18KB weights + 10KB input = 28KB shared memory per block, and up to 32 blocks per SM, total would exceed available shared memory. But occupancy isn't necessarily 32 blocks per SM.
Let me think about a more register-friendly approach.
For small in_channels (8) and small kernel (3x3), another approach is:
- Have each warp compute one (or few) output channels for a tile of output pixels.
- Load weights for one output channel into registers (8*3*3 = 72 floats).
- Load input tile into shared memory.
- Each thread in the warp computes one output pixel.
Actually, for better memory coalescing and compute efficiency, let me consider:
- Block size: e.g., 256 threads
- Each thread computes C_OUT_PER_THREAD output channels for one output pixel.
- The block computes a tile of, say, 8x8 or 16x16 output pixels across all batches and output channels.
Let me be more concrete. Let's say we have a 2D grid:
- blockDim = (16, 16, 1) = 256 threads
- Each thread computes 1 output pixel for all 64 output channels? No, that would need 64 accumulators per thread, which is fine.
Wait, 256 threads * 64 accumulators = 16384 floats = 64KB registers per block. With 64K registers per SM, that's 1 block per SM at most. Not great for occupancy.
Better: have the block compute a subset of output channels.
- Block computes 16x16 spatial tile and, say, 4 output channels.
- Each thread computes 1 pixel for 4 output channels = 4 accumulators.
- Total accumulators per block: 256 * 4 = 1024.
- But wait, each pixel needs 72 input values. With 16x16 tile + halo = 18x18*8 ≈ 2600 floats in shared memory.
- Weights: 4 * 8 * 3 * 3 = 288 floats.
This seems reasonable. But then to cover all 64 output channels, we'd need 16 passes or 16x more blocks.
Actually, a simpler approach that usually works well for small kernels is to use a 1D thread block where each thread computes all output channels for one output pixel.
- Thread block size: 256
- Each block computes 256 output pixels (distributed across batch and spatial dims).
- Each thread loops over all 64 output channels.
For one output pixel at position (n, oh, ow):
- Load input 8x3x3 values into registers.
- For each output channel c: dot product with weights[c, :, :, :].
- Apply /divisor and leaky_relu.
- Store to output.
But loading 72 input floats per thread means 72*4 = 288 bytes per thread, read from global memory. With 256 threads, that's ~73KB per block. The memory access pattern might not be coalesced well because neighboring threads read neighboring pixels, but for conv they need overlapping regions.
Actually, using shared memory for the input tile is much better for data reuse.
Let me design a kernel with:
- Thread block: 16x16 threads for a 16x16 output tile.
- Each thread computes all 64 output channels for its pixel.
- Shared memory for input tile: 18x18x8 floats = 2592 floats = ~10KB.
- Weights stay in global memory/L2 (only 18KB total, easily cached).
For the 64 output channels, each thread does 64 dot products of 72 elements each.
- 64 * 72 = 4608 FMAs per thread.
- For 256 threads per block: 1.18M FMAs per block.
- Total blocks: 128 * (126/16) * (126/16) ≈ 128 * 8 * 8 = 8192 blocks (rounding up).
Wait, 126/16 = 7.875, so 8 in each dimension. 128 * 8 * 8 = 8192 blocks.
Each block 256 threads = ~2M threads total.
2M threads * 4608 FMAs = 9.4 GFLOPs. ✓
Register usage: each thread needs 64 accumulators + 72 input values + some temporaries.
64 + 72 = 136 floats. At 4 bytes each, that's 544 bytes. At 255 registers max (1020 bytes), this should fit. But the compiler may not lay it out perfectly.
Actually, 136 floats might use quite a few registers, potentially causing spilling or low occupancy. Let me check: 136 floats * 4 bytes / 4 bytes per register = 136 registers. That's over the 128 register threshold for max threads per block, but under 255. This would limit to maybe 512 or 1024 threads per SM.
For H100, the warp scheduler can handle this, but occupancy might be around 50%. That's probably fine for such a compute-heavy kernel.
An alternative: process output channels in groups. Say each thread computes 8 output channels.
- 64/8 = 8 passes per block.
- Each pass loads 8 sets of weights (8*72 = 576 floats) into registers or shared memory.
- Actually, if we load weights into shared memory per pass: 8*8*3*3 = 576 floats = 2.3KB.
- Thread accumulators: 8 floats per thread.
- Total shared memory: 10KB input + 2.3KB weights = 12.3KB.
- Each thread does 8 * 72 = 576 FMAs per pass, 8 passes = 4608 FMAs total.
- This uses fewer registers per thread (8 instead of 64 accumulators).
This seems like a much better trade-off! But it requires iterating over output channels in passes, which adds loop overhead and repeated reads of input from shared memory.
Actually, once input is in shared memory, reading it repeatedly is cheap. And weights in shared memory are also cheap to read.
Let me think about another approach: have each thread compute multiple pixels for fewer output channels.
Or perhaps use a GEMM-based approach. We could implement im2col in shared memory and then use a small GEMM. But that seems overly complex.
Let me try a hybrid approach:
- 1D thread block of 256 threads.
- The block computes a 2D tile of output, say 16x16 = 256 pixels.
- Each thread computes one pixel for a group of output channels.
- We tile over output channels in chunks of, say, 16 or 32.
Hmm, let me look at what cuDNN likely does. For a conv like this, it probably uses implicit GEMM with Tensor Cores. Since we're supposed to be faster than PyTorch (which uses cuDNN), we need to be clever.
Wait, can we actually beat cuDNN on a standard conv? PyTorch's reference uses cuDNN, which is highly optimized. However, we can fuse the divide and LeakyReLU, which PyTorch might not do in a single kernel. That could give us an edge due to reduced memory bandwidth.
PyTorch might:
1. Run cudnnConvolutionForward
2. Run division kernel
3. Run leaky_relu kernel
If cudnn doesn't fuse the epilogue, then we have extra memory traffic for the 512MB intermediate tensor. That's our opportunity.
Output is 128*64*126*126 = 128,024,576 floats = 512MB.
At 3.35 TB/s, reading and writing 512MB takes ~0.3 ms. Avoiding that could be a nice win.
Can we write a kernel that's competitive with cuDNN conv + 2 elementwise ops? Yes, likely.
But actually, modern cuDNN supports fused conv+bias+activation. Does PyTorch use it? PyTorch tries to use cudnn_convolution_relu when possible. But dividing by a constant might prevent fusion in the standard nn.Sequential or manual operations.
Actually, in the reference:
```python
x = self.conv(x)
x = x / self.divisor
x = torch.nn.functional.leaky_relu(x, negative_slope=0.01)
```
PyTorch's graph executor or eager mode might not fuse these. In eager mode, it's definitely 3 separate kernels. So our fused kernel has a good chance.
For the kernel design, I should think about whether to do direct conv or something else.
Given the small sizes (C_in=8, k=3), direct convolution is appropriate.
Let me design a kernel with the following parameters:
- Block: 16x16 threads (256 threads)
- Each block computes a 16x16 output spatial tile
- Each thread computes all 64 output channels for its pixel
- Shared memory for input: (16+2) * (16+2) * 8 floats = 2592 floats
Wait, 64 accumulators per thread might cause register pressure. But on H100 with 64K registers per SM and max 2048 threads per SM, 64K/2048 = 32 registers per thread minimum. But we want higher registers per thread.
If block is 256 threads and we use ~128 registers per thread, that's 32K registers per block. With 64K per SM, that's 2 blocks per SM. At 2 blocks * 256 threads = 512 threads per SM. That's 16 warps, which is enough to hide latency on H100.
Actually, let me calculate more carefully. If each thread uses 128 registers, a block of 256 threads uses 256*128 = 32768 registers. 64K / 32768 = 2 blocks per SM concurrently. That's actually decent. 2 blocks * 256 = 512 threads per SM, which is enough for latency hiding.
But with 64 accumulators + 72 inputs + temp registers... let me think. If the compiler keeps 64 accumulators in registers and loads input values one-by-one or reuses them, it might not need all 72 inputs in registers at once.
Actually, for a dot product of 72 elements, the compiler can do:
- Initialize 64 accumulators to 0.
- For each weight position (i, j) and each input channel k:
- Load input[k, oh+i, ow+j]
- For each output channel c: accumulator[c] += input_val * weight[c, k, i, j]
This only needs 1 input value in a register at a time, plus the 64 accumulators, plus maybe a few temporaries for loading weights. So ~64 + a few = ~70 registers. That's very manageable!
Wait, but loading weight[c, k, i, j] for each output channel requires either loading them sequentially or having them in registers. If weights are in L2 cache, the compiler can load them one at a time:
- Load weight_0, compute accumulator[0] += input * weight_0
- Load weight_1, compute accumulator[1] += input * weight_1
- ...
This only needs 1 accumulator active at a time? No, the compiler will keep all 64 accumulators in registers and use FMAs or MADs.
Actually, for a loop:
```
for c in range(64):
acc[c] += inp * weight[c]
```
The compiler needs 64 accumulators, 1 input value, and loads weight[c] sequentially. That's about 66 registers. Very reasonable.
But for the best performance, we want the weights to be cached well. Since all threads in a block need the same weights for a given (c, k, i, j), and weights are tiny, they should stay in L2 nicely.
Let me reconsider the loop order for best memory access:
```
for c in range(out_channels):
for k in range(in_channels):
for i in range(3):
for j in range(3):
inp = input_shared[n, k, oh+i, ow+j] // different per thread
w = weight[c, k, i, j] // same for all threads
acc[c] += inp * w
```
Wait, if we loop c outermost, then for each c, all threads access the same weight but different inputs. This seems good.
But actually, with c outermost, we load weights for c=0 first (all threads load weight[0,0,0,0] - same address, cached), then c=1, etc. This has good cache behavior.
However, looping over 64 channels means writing to 64 different accumulators. That's fine.
Let me think about vectorization. If I process 4 output channels at a time, I can potentially use vectorized loads for weights, but it's probably not necessary for such a small kernel.
Another consideration: should I process batch elements in parallel within a block? The natural blocking is:
- Each block handles some batch elements and a spatial tile.
I think the simplest is:
- Grid: (batch, oh_tiles, ow_tiles)
- Block: (16, 16) threads mapping to a 16x16 spatial tile.
- Each thread handles one output pixel for all output channels.
Actually, since batch_size=128 and output is 126x126, we have plenty of parallelism. Let me make the block 8x8x4 or something...
Actually, let me try a simpler block size first: 8x8=64 threads per block, each thread computes all 64 output channels. Then each block computes 64 pixels. Number of blocks = 128 * 16 * 16 = 32768 blocks. That's plenty. And register pressure is lower.
With 64 threads per block, even if each thread uses 128 registers, that's 8192 registers per block, so 8 blocks per SM. 8*64 = 512 threads per SM. Good occupancy. But each block does less work, so block scheduling overhead might matter.
Let me go with 16x16 threads per block. 256 threads * ~70 registers = ~18K registers per block. That's 3-4 blocks per SM. 768-1024 threads per SM. Good.
Actually, I realize we can potentially use warp-level primitives or just rely on the compiler.
Let me draft the CUDA kernel:
```cuda
__global__ void fused_conv_div_leakyrelu(
const float* __restrict__ input,
const float* __restrict__ weight,
const float* __restrict__ bias,
float* __restrict__ output,
int batch_size,
int in_channels,
int out_channels,
int in_h,
int in_w,
int out_h,
int out_w,
int kernel_size,
float divisor,
float negative_slope
) {
int ow = blockIdx.x * blockDim.x + threadIdx.x;
int oh = blockIdx.y * blockDim.y + threadIdx.y;
int n = blockIdx.z;
if (ow >= out_w || oh >= out_h || n >= batch_size) return;
float acc[64]; // For out_channels=64 specifically, or use a loop
// Actually, to make it general but also fast, let me think...
}
```
Hmm, using a fixed-size array `float acc[64]` with a literal 64 might not work if out_channels can vary. But we need to support the constructor arguments. I could use a dynamically sized array, but CUDA doesn't support VLAs. I could allocate from shared memory or unroll with a max size.
Wait, the kernel needs to work for the given sizes but also be general enough for the constructor arguments. However, for performance, knowing C_out=64 at compile time would help. I can use template parameters or just loop.
Since C_out isn't a template parameter, I'll need a loop. But the compiler might not unroll it well.
Alternative: pass C_out as a template parameter or use a `constexpr` if I generate code. But with `load_inline`, generating templated code is possible.
Actually, let me look at it differently. I can use shared memory for the accumulators if register pressure is a concern. But shared memory is slower than registers.
Or I can process output channels in chunks of, say, 8 or 16 per thread to balance register usage.
Let me try this: process output channels in fixed-size chunks.
```cuda
#define OC_BLOCK 16
__global__ void fused_conv_div_leakyrelu(
const float* __restrict__ input,
const float* __restrict__ weight,
float* __restrict__ output,
...
int out_channels
) {
...
int oc_block_idx = blockIdx.w; // Need 4D grid or handle differently
...
}
```
Using a 4D grid is awkward. Better to have the thread block handle a subset of output channels.
Let me redesign:
- Block: (16, 16, 1) = 256 threads
- Grid over (spatial_tile_x, spatial_tile_y, batch)
- Each thread block handles a subset of output channels, say 8.
- Wait, then I'd need another dimension or blocks iterating over output channels.
Actually, I can map blockIdx.z to:
- blockIdx.z = batch * num_oc_tiles + oc_tile
Where num_oc_tiles = (out_channels + OC_TILE - 1) / OC_TILE.
Each block handles:
- A spatial tile
- A subset of output channels (e.g., 8 channels)
Then:
- Shared memory for input: (tile_h + 2) * (tile_w + 2) * in_channels
- Shared memory for weights: OC_TILE * in_channels * 3 * 3
- Each thread computes its pixel for OC_TILE output channels.
For OC_TILE=8:
- Thread accumulators: 8 floats.
- Weights in shared memory: 8 * 8 * 3 * 3 * 4 bytes = 2.3 KB.
- Input in shared memory: 18 * 18 * 8 * 4 bytes = 10.4 KB.
- Total shared memory: ~13 KB. Well within limits.
For H100 with 228KB per SM, we could fit ~17 blocks per SM if shared memory is the limit. With 256 threads per block, that's 4352 threads. But we only need ~512-1024 for good latency hiding. So shared memory won't be the bottleneck.
With 8 accumulators per thread, register usage is low, so we get good occupancy. We loop over OC_TILE * KH * KW * IC = 8 * 3 * 3 * 8 = 576 FMAs per thread for the tile, and do this 8 times to cover all 64 output channels.
Wait, each thread computes its assigned pixel. It accumulates 8 output channels. The loop is:
```
for c = 0 to OC_TILE:
for k = 0 to IC:
for i = 0 to 3:
for j = 0 to 3:
acc[c] += input_sm[k][oh+i][ow+j] * weight_sm[c][k][i][j]
```
Wait, this needs 4 nested loops. The inner loop over j(3) and i(3) can be unrolled or just kept as small loops.
Let me think about vectorization for the 8 output channels. If I use float4 or float2, I might get better instruction throughput.
For the input, since IC=8, I can load 2 float4s per position to get 8 channels.
Let me write out a more concrete plan:
Grid and block:
```cuda
dim3 block(16, 16, 1); // 256 threads
dim3 grid((out_w + 15) / 16, (out_h + 15) / 16, batch_size * ((out_channels + 7) / 8));
```
Each block:
- Determines: n = blockIdx.z / num_oc_tiles, oc_tile = blockIdx.z % num_oc_tiles
- Computes output channels [oc_tile * 8, oc_tile * 8 + 8)
- Computes spatial tile starting at (blockIdx.y * 16, blockIdx.x * 16)
Wait, in CUDA, blockIdx.x is usually the fastest varying. Let me use:
- blockIdx.x for ow_tile
- blockIdx.y for oh_tile
- blockIdx.z for batch * oc_tiles
Actually, for memory coalescing, having threads in a warp access consecutive memory addresses is important for the output writes. If thread (tx, ty) writes to output[n, c, oh, ow], then for consecutive threads in x, ow varies by 1, so they write to consecutive addresses. Good.
For the shared memory layout of input:
A 2D tile of threads (tx, ty) loads a 2D tile of input plus halo.
- The spatial extent of input needed: tile_h + 2 = 18, tile_w + 2 = 18.
- For each position, we need in_channels=8 values.
I should load the input tile cooperatively:
```cuda
__shared__ float input_sm[18][18][8]; // or [8][18][18] for better access
```
Actually, for the convolutional access pattern where we iterate channels innermost, [8][18][18] might be better? Let me think.
In the compute loop:
```
for c in 0..7:
for k in 0..7:
for i in 0..2:
for j in 0..2:
float inp = input_sm[k][oh+i][ow+j];
float w = weight_sm[c][k][i][j];
acc[c] += inp * w;
```
The access `input_sm[k][oh+i][ow+j]` for different threads in a warp:
- oh and ow differ by threadIdx.x/y.
- If threads in a warp have the same oh but different ow, then `input_sm[k][oh+i][ow+j]` varies in the last dimension.
- In row-major layout, consecutive ow means consecutive memory. Good.
For `weight_sm[c][k][i][j]`, all threads in the block read the same address. Good for broadcast.
So the layout `[in_channels][tile_h][tile_w]` should work well.
Now, for the shared memory loading:
We need 18*18*8 = 2592 floats. With 256 threads, that's ~10 floats per thread. Each thread can load 10 floats, but the arrangement might be tricky.
A simple approach: flatten the shared memory and load with strided access.
```cuda
float* input_sm_flat = (float*)input_sm;
int total = 18 * 18 * 8;
for (int idx = threadIdx.x + threadIdx.y * blockDim.x; idx < total; idx += blockDim.x * blockDim.y) {
// Map idx to (k, ih, iw)
// ...
}
```
This ensures coalesced loads from global memory.
Alternatively, I can design the loading more carefully:
```cuda
for (int k = 0; k < in_channels; k++) {
for (int ih = threadIdx.y; ih < tile_h + 2; ih += blockDim.y) {
for (int iw = threadIdx.x; iw < tile_w + 2; iw += blockDim.x) {
int in_ih = blockIdx.y * tile_h + ih;
int in_iw = blockIdx.x * tile_w + iw;
if (in_ih < in_h && in_iw < in_w) {
input_sm[k][ih][iw] = input[...];
}
}
}
}
```
Wait, with blockDim.x=16, blockDim.y=16, for iw we'd have threads 0-15 load iw=0-15, then some threads load iw=16 and 17 if tile_w+2=18. Thread 0 would load 0 and 16, thread 1 loads 1 and 17. That's good for coalescing.
For ih=0..17, threads 0-15 load 0, then 16, then... wait, threadIdx.y ranges 0-15. So for ih: 0, 16 for each thread. But 16 is the last row and 17 isn't loaded by any thread? Oh wait, 17 would need a thread with threadIdx.y=17 which doesn't exist. So we'd need another iteration. With ih += blockDim.y:
- ih = 0, 16 for threadIdx.y=0
- ih = 1, 17 for threadIdx.y=1
- ...
But not all combinations are covered because 18 rows and 16 threads means rows 0-15 are covered in first pass, rows 16-17 in second pass by threads 0-1.
Actually this does work. For thread (x, y):
- First iteration: ih = y, iw = x and x+16
- Second iteration: ih = y + 16, but only if y < 2 (since y+16 < 18)
- Wait, no. The loop is `for (ih = threadIdx.y; ih < 18; ih += 16)`. So:
- y=0: ih = 0, 16
- y=1: ih = 1, 17
- y=2: ih = 2
- ...
- y=15: ih = 15
- And for each ih, `for (iw = threadIdx.x; iw < 18; iw += 16)`:
- x=0: iw = 0, 16
- x=1: iw = 1, 17
- x=2: iw = 2
- ...
- x=15: iw = 15
So each thread handles a unique combination. Wait, does this cover all 18*18=324 positions?
Total positions = 324.
Threads per block = 256.
Some threads handle 2 positions (those with x<2 and y<2), others handle 1.
Total handled = (2*2)*2 + (2*14)*1 + (14*2)*1 + (14*14)*1 = 8 + 28 + 28 + 196 = 260? No that's not 324.
Wait, for each thread, the number of positions is:
- (num_ih_values) * (num_iw_values)
- For y=0: ih in {0, 16} → 2 values
- For y=1: ih in {1, 17} → 2 values
- For y=2..15: ih = y → 1 value
- Same for iw.
So total positions loaded =
(threads with 2x2 positions): 4 threads * 4 = 16
(threads with 2x1 positions): 28 threads * 2 = 56
(threads with 1x2 positions): 28 threads * 2 = 56
(threads with 1x1 positions): 196 threads * 1 = 196
Total = 324. ✓
But for the in_channels dimension, we need to loop over k=0..7. So each thread loads 324 * 8 / 256 = ~10.1 values on average. Some loads to the same position from different threads? No, because the (ih, iw) assignment is unique per thread per k.
Wait, no. If we put k in the outer loop:
```cuda
for (int k = 0; k < 8; k++) {
for (int ih = threadIdx.y; ih < 18; ih += 16) {
for (int iw = threadIdx.x; iw < 18; iw += 16) {
input_sm[k][ih][iw] = ...;
}
}
}
```
Then each thread loads ~10 values per k iteration, total ~80 values. But the load pattern within the k loop is the same for all k, just from different addresses in global memory.
Actually, this should be fine. Global memory loads will be coalesced because consecutive threads (in x) access consecutive iw values.
For the weights, all threads load the same weights. We can have thread 0 load all weights into shared memory:
```cuda
__shared__ float weight_sm[8][8][3][3]; // for OC_TILE=8
if (threadIdx.x == 0 && threadIdx.y == 0) {
for (int c = 0; c < oc_tile_size; c++) {
for (int k = 0; k < 8; k++) {
for (int i = 0; i < 3; i++) {
for (int j = 0; j < 3; j++) {
weight_sm[c][k][i][j] = weight[...];
}
}
}
}
}
__syncthreads();
```
Or better, have all threads participate in loading weights:
```cuda
int total_weights = oc_tile_size * 8 * 9;
for (int idx = threadIdx.x + threadIdx.y * blockDim.x; idx < total_weights; idx += 256) {
// map idx to (c, k, i, j)
weight_sm[...] = weight[...];
}
__syncthreads();
```
This is faster and more balanced.
Now let's think about the specific parameter values:
- batch_size = 128
- in_channels = 8
- out_channels = 64
- height = width = 128
- kernel_size = 3
- divisor = 2
Output height/width = 128 - 3 + 1 = 126.
For the weights loading, with OC_TILE=8, total_weights=8*8*9=576. 576 elements loaded by 256 threads means most threads load 2 elements, some load 3. Very efficient.
For input loading, per block: 18*18*8 = 2592 elements, loaded by 256 threads = ~10 elements each. Fine.
Compute per thread: 8 accumulators * 8 * 9 = 576 FMAs.
Total FMAs per block: 256 threads * 576 FMAs = 147,456 FMAs.
Number of blocks: 128 * 8 * 8 * 8 = 65,536 blocks?
Wait: spatial tiles = ceil(126/16)^2 = 8*8 = 64. OC tiles = 64/8 = 8. Batch = 128.
Total blocks = 128 * 64 * 8 = 65,536 blocks.
Total FMAs = 65,536 * 147,456 = 9.66 billion FMAs (each FMA is 2 FLOPs) = ~19.3 GFLOPs.
Actually, the conv has:
- output elements: 128 * 64 * 126 * 126
- MACs per output element: 8 * 3 * 3 = 72
- Total MACs = 128 * 64 * 126 * 126 * 72 ≈ 9.3 billion
- Total FLOPs = 2 * 9.3B = 18.6B. ✓
Wait, I double counted. A MAC (multiply-accumulate) is 1 FMA = 2 FLOPs.
So 9.3B MACs = 9.3B FMAs = 18.6 GFLOPs. My block calculation had 65,536 blocks * 147,456 FMAs per block. Let me check: 65,536 * 147,456 = 9,663,676,416. That's ~9.66B FMAs. But the total should be 128*64*126*126*72 = 9,290,XXX,XXX. Let me recalculate.
128 * 64 = 8192
126 * 126 = 15876
8192 * 15876 = 130,xxx,xxx... wait.
8192 * 15876 = 8192 * (16000 - 124) = 131,072,000 - 1,015,808 = 130,056,192
130,056,192 * 72 = 9,364,045,824
My block count assumes every block does a full 16x16 tile. But edge blocks do less because 126 isn't divisible by 16.
- 126 / 16 = 7 remainder 14. So some tiles are smaller (14 wide).
However, with the guard conditions, threads outside the boundary don't compute. But they still participate in shared memory loading, which is fine.
Let me recalculate with full tiles:
Actually, for spatial, we'd have tiles:
- ow tiles: ceil(126/16) = 8. Last tile width = 126 - 7*16 = 126 - 112 = 14.
- oh tiles: similarly 8, last height = 14.
So some blocks process fewer pixels. The actual total FMAs computed by all threads is the same as the mathematical total (since we have if-guards), but the number of FMA instructions issued might be slightly higher for edge tiles due to inactive threads.
Anyway, this looks reasonable.
One concern: with 65,536 blocks and an H100 having ~132 SMs (I think H100 has 132 or 144 SMs for the SXM version), that's ~500 blocks per SM. That's plenty for good work distribution.
Wait, H100 SXM5 has 132 SMs. PCIe has 114. Let me check... Actually, for H100 80GB, it's likely 132 SMs or 144 SMs. With 65,536 blocks / 132 SMs ≈ 496 blocks per SM. At maybe 3-4 concurrent blocks per SM, that's ~100+ waves. Fine.
Now, is there a risk that this kernel is slower than cuDNN? cuDNN uses Tensor Cores and highly optimized algorithms. For this small conv, cuDNN might also use implicit GEMM with direct loading. Our advantage is fusion and not writing the intermediate.
Actually, let me reconsider whether PyTorch fuses these operations. Let me look at the operations:
- conv2d
- division by scalar
- leaky_relu
In PyTorch eager mode, these are 3 kernels. In torch.compile mode, it might fuse them. But the reference says we're replacing the forward computation, so we're competing against whatever PyTorch does by default.
One more thing: the Conv2d layer has a bias by default? Let me check the reference:
```python
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
```
In PyTorch, Conv2d has `bias=True` by default! So there IS a bias term. I need to include bias in the kernel.
The reference Conv2d op computes: conv(x, weight) + bias.
Then divides by divisor, then leaky_relu.
So the full fused op is:
leaky_relu((conv(x, weight) + bias) / divisor, 0.01)
I need to add bias loading. Bias is small: 64 floats.
For our kernel, we can have thread 0 load the bias for the 8 output channels into shared memory:
```cuda
__shared__ float bias_sm[8];
if (threadIdx.x < 8 && threadIdx.y == 0) {
bias_sm[threadIdx.x] = bias[oc_start + threadIdx.x];
}
__syncthreads();
```
Or all threads participate.
Now let me think about weight access patterns for better L2 cache. If block 0 loads weights for oc=0..7, block 1 loads for oc=0..7 (but different spatial tile), block 2 as well... then all blocks in a wave need the same weights. With 64KB L1 cache per SM and L2 cache being much larger, this should be fine since weights are only 64*8*9*4 = 18KB + 64*4 bytes bias. Very small.
Wait, actually each block only loads 8*8*9*4 = 2.3KB of weights at a time. With 8 OC tiles, each wave of blocks on an SM processes weights for one oc_tile at a time? Not necessarily, because blocks are scheduled dynamically.
Actually, since weights are tiny, they'll be comfortably in L2 cache. Reading them from global memory should effectively be L2 cache hits after the first wave.
Alright, let me draft the kernel more carefully.
```cuda
__global__ void fused_conv_div_leakyrelu_kernel(
const float* __restrict__ input,
const float* __restrict__ weight,
const float* __restrict__ bias,
float* __restrict__ output,
int batch_size,
int in_channels,
int out_channels,
int in_h,
int in_w,
int out_h,
int out_w,
int kernel_size,
float inv_divisor,
float negative_slope
) {
const int TILE_X = 16;
const int TILE_Y = 16;
const int OC_TILE = 8;
int ow = blockIdx.x * TILE_X + threadIdx.x;
int oh = blockIdx.y * TILE_Y + threadIdx.y;
int n = blockIdx.z / ((out_channels + OC_TILE - 1) / OC_TILE);
int oc_tile = blockIdx.z % ((out_channels + OC_TILE - 1) / OC_TILE);
int oc_start = oc_tile * OC_TILE;
int oc_end = min(oc_start + OC_TILE, out_channels);
int oc_count = oc_end - oc_start;
if (n >= batch_size) return;
__shared__ float input_sm[8][18][18]; // [in_channels][tile_h+2][tile_w+2]
__shared__ float weight_sm[8][8][3][3]; // [OC_TILE][in_channels][3][3]
__shared__ float bias_sm[8]; // [OC_TILE]
// Load input tile cooperatively
int in_h_start = blockIdx.y * TILE_Y;
int in_w_start = blockIdx.x * TILE_X;
for (int k = 0; k < in_channels; k++) {
for (int ih = threadIdx.y; ih < TILE_Y + 2; ih += TILE_Y) {
for (int iw = threadIdx.x; iw < TILE_X + 2; iw += TILE_X) {
int in_ih = in_h_start + ih;
int in_iw = in_w_start + iw;
if (in_ih < in_h && in_iw < in_w) {
input_sm[k][ih][iw] = input[((n * in_channels + k) * in_h + in_ih) * in_w + in_iw];
} else {
input_sm[k][ih][iw] = 0.0f;
}
}
}
}
// Load weights cooperatively
int total_weights = oc_count * in_channels * kernel_size * kernel_size;
for (int idx = threadIdx.x + threadIdx.y * TILE_X; idx < total_weights; idx += TILE_X * TILE_Y) {
int tmp = idx;
int j = tmp % kernel_size; tmp /= kernel_size;
int i = tmp % kernel_size; tmp /= kernel_size;
int k = tmp % in_channels; tmp /= in_channels;
int c = tmp;
weight_sm[c][k][i][j] = weight[((oc_start + c) * in_channels + k) * kernel_size * kernel_size + i * kernel_size + j];
}
// Load bias
if (threadIdx.x < oc_count && threadIdx.y == 0) {
bias_sm[threadIdx.x] = bias[oc_start + threadIdx.x];
}
__syncthreads();
if (ow >= out_w || oh >= out_h) return;
// Compute
float acc[8];
for (int c = 0; c < oc_count; c++) {
acc[c] = bias_sm[c];
}
for (int c = 0; c < oc_count; c++) {
for (int k = 0; k < in_channels; k++) {
for (int i = 0; i < kernel_size; i++) {
for (int j = 0; j < kernel_size; j++) {
float inp = input_sm[k][threadIdx.y + i][threadIdx.x + j];
float w = weight_sm[c][k][i][j];
acc[c] += inp * w;
}
}
}
}
// Apply divisor and leaky_relu
for (int c = 0; c < oc_count; c++) {
acc[c] *= inv_divisor;
acc[c] = (acc[c] > 0.0f) ? acc[c] : (acc[c] * negative_slope);
int out_idx = (((n * out_channels + oc_start + c) * out_h + oh) * out_w + ow);
output[out_idx] = acc[c];
}
}
```
Wait, there's a subtle bug. The loops `for (int ih = threadIdx.y; ih < TILE_Y + 2; ih += TILE_Y)` and `for (int iw = threadIdx.x; iw < TILE_X + 2; iw += TILE_X)` have `TILE_Y` and `TILE_X` as the step, but the conditions should be `< TILE_Y + 2` and `< TILE_X + 2`. With TILE_Y=16, threadIdx.y ranges 0-15. The values of ih would be:
- y=0: 0, 16
- y=1: 1, 17
- y=2: 2
- ...
- y=15: 15
This correctly covers 0-17 (18 values). Same for iw.
But wait, the loop condition `ih < TILE_Y + 2` is correct (ih < 18), but the second iteration for y=0 is ih = 0 + 16 = 16, and the third would be 32 which is not < 18. So y=0 gets 0 and 16. That's correct.
Actually, there is a subtle issue: TILE_Y is used both as the block dimension (16) and as the step. This is correct.
Now, one more issue: I need to be careful about the case where `TILE_Y + 2 = 18` but the actual needed height might be less for boundary tiles. My code loads into shared memory based on block position and checks global bounds, but writes 0 for out-of-bounds. The compute threads check `if (ow >= out_w || oh >= out_h) return;` so boundary issues are handled. But inactive threads still participate in `__syncthreads()`, which is fine since we only have `__syncthreads()` before the compute.
Wait, actually I have `__syncthreads()` after loading, and then the return for out-of-bounds threads. If some threads return before `__syncthreads()`, that would deadlock! Oh no, that's a critical bug.
I need to move the boundary check. The threads that are outside the output spatial bounds should still participate in shared memory loading, or at least not return before `__syncthreads()`.
Actually, I can just not have the early return at all, and instead guard the compute and store:
```cuda
__syncthreads();
bool valid = (ow < out_w && oh < out_h);
float acc[8];
...
if (valid) {
for (int c = 0; c < oc_count; c++) {
acc[c] *= inv_divisor;
...
output[...] = acc[c];
}
}
```
But loops over c and k etc would still run for invalid threads, wasting some work. Since most tiles are full (7 out of 8 tiles in each dimension are full 16x16, only the last is partial), this is acceptable.
Actually, let me compute: out_w=126. Tiles: 0-15, 16-31, ..., 112-125. That's 8 tiles. The last tile has width 14. So in the last column of tiles, threads with threadIdx.x >= 14 are invalid. That's 2 out of 16 columns. Similarly for rows. So only 4 threads out of 256 are invalid in the last row+col tile, and 16 threads invalid in last row tiles, 16 in last col tiles.
Most tiles are fully utilized. So the waste is minimal. We can avoid the early return and just guard the store.
Actually, I should structure the compute like this to avoid unnecessary compute for invalid threads:
```cuda
__syncthreads();
if (ow < out_w && oh < out_h) {
float acc[8];
for (int c = 0; c < oc_count; c++) acc[c] = bias_sm[c];
#pragma unroll
for (int c = 0; c < oc_count; c++) {
#pragma unroll
for (int k = 0; k < in_channels; k++) {
#pragma unroll
for (int i = 0; i < kernel_size; i++) {
#pragma unroll
for (int j = 0; j < kernel_size; j++) {
acc[c] += input_sm[k][threadIdx.y + i][threadIdx.x + j] * weight_sm[c][k][i][j];
}
}
}
}
for (int c = 0; c < oc_count; c++) {
float val = acc[c] * inv_divisor;
val = fmaxf(val, val * negative_slope); // or conditional
output[...] = val;
}
}
```
No `__syncthreads()` after this, so the early return for invalid threads is safe here.
But wait, I need to make sure all threads participate in loading shared memory. The loading code runs before the `__syncthreads()` and before any early returns, so all threads participate. Good.
Let me also think about whether to process the bias differently. For the invalid threads, we don't need to compute. Good.
One more optimization: instead of `if (val > 0) val = val; else val = val * negative_slope`, we can use:
```cuda
val = (val > 0.0f) ? val : val * negative_slope;
```
Or on newer GPUs, there's a `leaky_relu` instruction or something, but the ternary is fine.
Actually, I should also think about vectorizing the loads and stores. For the output, if a thread writes oc_count values, these are not consecutive in memory (they're separated by out_h*out_w). So we can't easily vectorize the store per thread. But consecutive threads writing the same channel for consecutive ow positions means their writes are consecutive. This is already coalesced.
For the input loading, consecutive threads load consecutive iw values, which is coalesced.
Let me think about unrolling. Since kernel_size=3 and in_channels=8 and oc_count<=8, the compiler can unroll these small loops if we add `#pragma unroll`. I should add that.
There's another potential optimization: loop reordering. Currently:
```
for c:
for k:
for i:
for j:
acc[c] += input_sm[k][ty+i][tx+j] * weight_sm[c][k][i][j]
```
For a given (k, i, j), all threads in the block read weight_sm[c][k][i][j] for the same c but the weight is different for each c. For a fixed c, all threads read the same weight. The inner loop over j means we access input_sm with varying j (consecutive in memory for fixed k, ty+i). This is good.
Actually, for better instruction-level parallelism, I could reorder:
```
for k:
for i:
for j:
float inp = input_sm[k][ty+i][tx+j];
for c:
acc[c] += inp * weight_sm[c][k][i][j];
```
This way, `inp` is loaded once and reused across all output channels. For OC_TILE=8, this saves 7 loads from shared memory per (k,i,j). Since shared memory is fast but not free, this is better!
However, the inner loop over c might prevent the compiler from keeping acc[c] in registers efficiently, or it might cause execution dependencies. But with only 8 accumulators, it should be fine.
Actually, let me think about register pressure. If we loop c innermost:
```
for k, i, j:
float inp = ...;
acc[0] += inp * w[0];
acc[1] += inp * w[1];
...
acc[7] += inp * w[7];
```
Here, `acc[0..7]` stay in registers throughout, and we just load one `inp` and 8 weights per iteration. Total registers: 8 acc + 1 inp + temp for address calc. ~12 registers. Very low! But we need the compiler to not spill.
This also allows instruction-level parallelism: the 8 FMAs can be pipelined.
Wait, actually I need to verify: does CUDA allow variable-length arrays? `float acc[8]` where 8 is a constant should be fine. But `oc_count` might be less than 8 for the last tile. I'll just use `OC_TILE` as a constant and ignore the extra values, or use a pragma unroll with a constant.
Let me set `OC_TILE = 8` as a constant. Since out_channels=64 is divisible by 8, oc_count will always be 8. For generality, we should handle cases where out_channels is not divisible by 8, but for this problem we know it's 64.
In the kernel, I can declare `float acc[OC_TILE];` and only use the first `oc_count` of them. The compiler will likely optimize away the unused ones.
But wait, `oc_count` is not known at compile time (it depends on `blockIdx.z`). So I should either:
1. Always compute all OC_TILE and only store the relevant ones
2. Use a compile-time constant or template
For robustness, I'll use a loop with `oc_count` but declare `float acc[OC_TILE];`. The compiler might not unroll the loop with a variable bound. If performance is critical, I can templatize the kernel with `OC_TILE`.
Actually, for this specific problem, I know `out_channels=64`, so I can just process all 64 output channels per thread without tiling. But that means all threads need the full weight matrix in registers or L1/L2.
Let me revisit the single-thread-all-64-channels approach:
- `float acc[64];`
- For each input position, load 1 input value.
- For each output channel, load weight and compute: `acc[c] += inp * weight[c];`
- 64 accumulators in registers. The compiler might actually handle this.
- No output channel tiling needed. Fewer blocks (128 * 8 * 8 = 8192 blocks instead of 65536).
- Each block does more work: 256 threads * 64 * 72 = 1.18M FMAs per block.
- Total FMAs: 8192 * 1.18M = 9.66B. Same total.
With fewer blocks, we might have better block-level occupancy but fewer opportunities for parallelism. 8192 blocks is still plenty for 132 SMs (62 blocks per SM). With maybe 2-3 concurrent blocks per SM, that's 20-30 waves. Still enough.
The advantage is fewer loops and no output channel tiling overhead. Let me think about register pressure again.
If each thread uses 64 accumulators + 1 input + 1 weight + a few loop vars + output address calc. Maybe 70-80 registers. At 80 registers per thread and 256 threads per block, that's 20,480 registers per block. 64K / 20K = ~3 blocks per SM concurrently. 768 threads per SM. Good latency hiding.
Actually, on H100, the register file is 64K 32-bit registers per SM. If a thread uses 80 registers, a warp uses 2560 registers. Max warps per SM is 32 for compute capability 9.0? Actually, H100 has max 64 warps per SM. Let's look it up.
Wait, the specs say "Maximum number of threads per SM" for H100. Ampere was 2048. Hopper might be 2048 as well. With 2048 threads and 64K registers, if each thread uses 80 registers, 2048*80 = 163,840 > 64,000. So at 80 registers, max threads would be 64K/80 = 800 threads. So we'd be limited to maybe 2-3 blocks per SM.
Is 800 threads per SM enough? On Ampere/Hopper, 1 warp = 32 threads, so 800/32 = 25 warps. If max is 64 warps, that's ~39% occupancy. For a compute-bound kernel, this might be okay but not ideal.
However, we have plenty of parallelism (8192 blocks * 256 threads). The GPU can keep all SMs busy.
Let me consider the actual register count. If I use `float acc[64];`, will the compiler actually allocate 64 registers? It depends on the loop structure.
```cuda
float acc[64];
for (int c = 0; c < 64; c++) acc[c] = bias[c];
for (int k = 0; k < 8; k++)
for (int i = 0; i < 3; i++)
for (int j = 0; j < 3; j++) {
float inp = input_sm[k][ty+i][tx+j];
for (int c = 0; c < 64; c++) {
acc[c] += inp * weight_sm[c][k][i][j];
}
}
```
For the inner loop, the compiler might:
1. Keep `inp` in one register.
2. Load `weight_sm[c]` into a register sequentially.
3. Update `acc[c]` in place.
4. Since `acc[c]` is only used in this one statement and later in the store, the compiler must keep all 64 values live. Yes, 64 registers for `acc`.
Plus maybe a few for loop indices and addresses. Total ~70.
Alternatively, if the compiler decides to spill some accumulators to local memory, we'd lose performance. But 70 registers is below 128, so it probably won't spill.
Actually, I should check: what is the max registers per thread before occupancy drops significantly?
- If max is 255, then 70 is fine.
- The occupancy calculator:
- 70 registers -> max 4 blocks of 256 threads = 1024 threads per SM. If SM limit is 2048, that's 50% occupancy. But 1024 threads is plenty.
Wait, let me recalculate. 64K registers per SM. 70 registers per thread * 256 threads per block = 17,920 registers per block. 64,000 / 17,920 = 3.57 blocks per SM. So 3 blocks concurrently = 768 threads. 50% occupancy of 1536 or ~37% of 2048.
For a kernel that does 4608 FMAs per thread, this should be fine. The arithmetic intensity is high enough that we care about throughput, not latency hiding. 768 threads * 4608 FMAs = 3.5M FMAs in flight per SM. At clock ~2GHz, an SM can do maybe 128 FMAs per clock (with Tensor Cores, much more). But even 128 FMAs/clock * 2GHz = 256 GFMA/s per SM. 3.5M / 256G = 0.014 ms per SM. But total work is ~128 SMs, so actually it's much less.
Actually, let me think about this differently. Total time = total work / throughput. Total FMAs = 9.3B. At 67 TFLOPS = 67,000 GFLOPS. Time = 9.3B / 67,000B * 2 = wait, 67 TFLOPS = 33.5 TFMA/s. Time = 9.3 / 33,500 seconds? No.
67 TFLOPS for FP32 on H100. 1 FMA = 2 FLOPs. So 33.5 TFMA/s.
Total FMAs = 9.3B.
Time = 9.3B / 33.5T = 0.000277 seconds = 0.277 ms.
Plus memory. At 50% occupancy, throughput is maybe 50% of peak, so ~0.55 ms.
If we avoid the 2 extra memory passes (read+write 512MB each), we save:
- 512MB read + 512MB write for division = 1GB
- 512MB read + 512MB write for leaky_relu = 1GB
- Total saved: ~2GB memory traffic.
Wait, that's not quite right:
1. Conv writes 512MB output
2. Division reads 512MB, writes 512MB
3. LeakyReLU reads 512MB, writes 512MB
Total writes = 512 + 512 + 512 = 1536MB
Total reads = 512 + 512 = 1024MB
(in addition to conv reading input and weights)
With fusion, we only write 512MB once. The reads are the same for conv (plus input). So we save ~1GB of memory traffic.
At 3.35 TB/s, 1GB = 0.3 ms. So fusion alone could save 0.3 ms. If conv takes 0.3 ms and elementwise ops take 0.2 ms each, total reference is ~0.7 ms, and our kernel might do it in ~0.4 ms. That's a significant win.
OK let me try both approaches and see which is faster. I'll start with the all-64-channels-per-thread approach since it's simpler and might have less loop overhead.
One thing to be careful about: shared memory bank conflicts. With `input_sm[8][18][18]`, when threads in a warp access `input_sm[k][ty+i][tx+j]` with varying `tx`, they access consecutive elements in the last dimension. With 32 banks and 4 bytes per float, each bank holds every 32nd float. So consecutive floats go to consecutive banks. No bank conflicts for the read!
When loading weights, all threads read the same address, so broadcast, no conflicts.
Great.
Now, one more optimization: prefetching or using `__ldg` for global memory loads. But the compiler likely does this automatically.
Another thought: can I use the `mma` instructions or Tensor Cores? For a 3x3 conv with C_in=8, C_out=64, the GEMM is quite small. Using Tensor Cores would require careful data packing and might not be worth it for such small dimensions. The overhead of setting up TC operations might exceed the benefit.
Let me also consider using `constexpr` for the loop bounds to help the compiler:
```cuda
const int K = kernel_size;
#pragma unroll
for (int i = 0; i < K; i++) ...
```
Actually, `kernel_size` is passed as a parameter but we know it's 3. If I hardcode 3, the compiler can unroll completely. For generality, I'll pass it but maybe add `constexpr int K = 3` in the kernel body or just use the parameter with `#pragma unroll`.
For a parameter, `#pragma unroll` with a small value should unroll if the compiler can determine it's small. Actually, `pragma unroll` without a number tells the compiler to fully unroll if possible. But if `kernel_size` is a runtime parameter, the compiler doesn't know it's 3. I should pass it as a template parameter or just hardcode based on the known values.
Since the constructor takes `kernel_size`, I should support variable kernel_size. But for performance, I'll write the kernel with a `kernel_size` parameter and rely on the compiler to optimize small loops. Or I can use a macro.
Actually, the simplest approach: write the kernel for arbitrary kernel_size but add `#pragma unroll 1` or nothing for the weight loops, and trust the compiler. For kernel_size=3, even without unrolling, the loop overhead is negligible compared to 64*8*9=4608 iterations.
Wait, without unrolling, there's loop overhead (branch instructions, index updates). But with 4608 iterations of the innermost body, and loops of size 64, 8, 3, 3:
- Inner j loop: 3 iterations, some overhead but small.
- i loop: 3 iterations.
- k loop: 8 iterations.
- c loop: 64 iterations.
Total iterations: 64*8*3*3 = 4608. Each iteration does 1 FMA. The loop overhead for 4608 iterations is minimal. But the branch instructions might cause some divergence.
Actually, let me just fully unroll the 3x3 spatial loops since they're tiny. I can do this by using `constexpr int K = 3` or just relying on `#pragma unroll`. But if `kernel_size` is a runtime variable, the compiler can't fully unroll it.
Solution: use template parameter for kernel_size.
```cuda
template<int K>
__global__ void fused_conv_div_leakyrelu(...) {
#pragma unroll
for (int i = 0; i < K; i++) ...
}
```
Then in PyTorch: instantiate with `fused_conv_div_leakyrelu<3>`.
But with `load_inline`, calling a templated function from Python is tricky. I'd need to instantiate all possible templates and call the right one. Since kernel_size is a constructor argument, it's known at module init time. So I can generate the kernel string with the kernel_size hardcoded.
Actually, an easier approach: just hardcode `3` in the kernel but make it a parameter that's only used for bounds checking. Or generate the CUDA code dynamically in Python based on the kernel_size.
Since the problem only tests with `kernel_size=3`, I can safely hardcode it, but for correctness with other sizes I should handle them. Let me write a general kernel but with `#pragma unroll` and hope the compiler optimizes small constants.
Actually, a neat trick: use a `switch` on `kernel_size` with cases 1, 2, 3, etc., each calling a templated device function. But that's getting complicated.
For simplicity and correctness, I'll write a general kernel and trust the compiler for the small loop. I can also add `if (kernel_size == 3)` specialization within the kernel:
```cuda
if (kernel_size == 3) {
#pragma unroll
for (int i = 0; i < 3; i++) ...
} else {
for (int i = 0; i < kernel_size; i++) ...
}
```
This lets the compiler optimize the common case.
OK, actually I just realized: in PyTorch, `bias` defaults to `True` for `nn.Conv2d`, but the code doesn't show bias being explicitly set. Let me verify:
```python
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
```
Yes, `bias=True` by default in PyTorch. So there is definitely a bias parameter. I need to make sure my ModelNew extracts and uses the bias from the conv layer.
Now let me write the Python code.
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
# CUDA source
cuda_source = r'''
__global__ void fused_conv_div_leakyrelu_kernel(
const float* __restrict__ input,
const float* __restrict__ weight,
const float* __restrict__ bias,
float* __restrict__ output,
int batch_size,
int in_channels,
int out_channels,
int in_h,
int in_w,
int out_h,
int out_w,
int kernel_size,
float inv_divisor,
float negative_slope
) {
const int TILE_X = 16;
const int TILE_Y = 16;
int ow = blockIdx.x * TILE_X + threadIdx.x;
int oh = blockIdx.y * TILE_Y + threadIdx.y;
int n = blockIdx.z;
if (n >= batch_size) return;
int in_h_start = blockIdx.y * TILE_Y;
int in_w_start = blockIdx.x * TILE_X;
__shared__ float input_s[8][18][18];
__shared__ float weight_s[64][8][3][3];
__shared__ float bias_s[64];
// Load input tile
for (int k = 0; k < in_channels; k++) {
for (int ih = threadIdx.y; ih < TILE_Y + 2; ih += TILE_Y) {
for (int iw = threadIdx.x; iw < TILE_X + 2; iw += TILE_X) {
int in_ih = in_h_start + ih;
int in_iw = in_w_start + iw;
if (in_ih < in_h && in_iw < in_w) {
input_s[k][ih][iw] = input[((n * in_channels + k) * in_h + in_ih) * in_w + in_iw];
} else {
input_s[k][ih][iw] = 0.0f;
}
}
}
}
// Load weights
int total_weights = out_channels * in_channels * kernel_size * kernel_size;
for (int idx = threadIdx.x + threadIdx.y * TILE_X; idx < total_weights; idx += TILE_X * TILE_Y) {
int tmp = idx;
int j = tmp % kernel_size; tmp /= kernel_size;
int i = tmp % kernel_size; tmp /= kernel_size;
int k = tmp % in_channels; tmp /= in_channels;
int c = tmp;
weight_s[c][k][i][j] = weight[((c * in_channels + k) * kernel_size + i) * kernel_size + j];
}
// Load bias
for (int idx = threadIdx.x + threadIdx.y * TILE_X; idx < out_channels; idx += TILE_X * TILE_Y) {
bias_s[idx] = bias[idx];
}
__syncthreads();
if (ow >= out_w || oh >= out_h) return;
float acc[64];
for (int c = 0; c < out_channels; c++) {
acc[c] = bias_s[c];
}
if (kernel_size == 3) {
#pragma unroll
for (int c = 0; c < out_channels; c++) {
#pragma unroll
for (int k = 0; k < in_channels; k++) {
float w00 = weight_s[c][k][0][0];
float w01 = weight_s[c][k][0][1];
float w02 = weight_s[c][k][0][2];
float w10 = weight_s[c][k][1][0];
float w11 = weight_s[c][k][1][1];
float w12 = weight_s[c][k][1][2];
float w20 = weight_s[c][k][2][0];
float w21 = weight_s[c][k][2][1];
float w22 = weight_s[c][k][2][2];
float i00 = input_s[k][threadIdx.y + 0][threadIdx.x + 0];
float i01 = input_s[k][threadIdx.y + 0][threadIdx.x + 1];
float i02 = input_s[k][threadIdx.y + 0][threadIdx.x + 2];
float i10 = input_s[k][threadIdx.y + 1][threadIdx.x + 0];
float i11 = input_s[k][threadIdx.y + 1][threadIdx.x + 1];
float i12 = input_s[k][threadIdx.y + 1][threadIdx.x + 2];
float i20 = input_s[k][threadIdx.y + 2][threadIdx.x + 0];
float i21 = input_s[k][threadIdx.y + 2][threadIdx.x + 1];
float i22 = input_s[k][threadIdx.y + 2][threadIdx.x + 2];
acc[c] += i00 * w00 + i01 * w01 + i02 * w02
+ i10 * w10 + i11 * w11 + i12 * w12
+ i20 * w20 + i21 * w21 + i22 * w22;
}
}
} else {
for (int c = 0; c < out_channels; c++) {
for (int k = 0; k < in_channels; k++) {
for (int i = 0; i < kernel_size; i++) {
for (int j = 0; j < kernel_size; j++) {
acc[c] += input_s[k][threadIdx.y + i][threadIdx.x + j] * weight_s[c][k][i][j];
}
}
}
}
}
for (int c = 0; c < out_channels; c++) {
float val = acc[c] * inv_divisor;
val = (val > 0.0f) ? val : (val * negative_slope);
output[((n * out_channels + c) * out_h + oh) * out_w + ow] = val;
}
}
''';
# But wait, in the unrolled version, I'm doing out_channels * in_channels iterations,
# and in each I load 9 weights and 9 inputs. That's a lot of registers.
# Let me see: for out_channels=64 and in_channels=8, the c loop is 64, k loop is 8.
# Total iterations: 512. In each, I load 18 floats (9 weights + 9 inputs) into temporaries.
# That would need a lot of registers. But the compiler should NOT keep all 18 across iterations.
# Wait, I'm declaring them inside the c loop, so they're local to that scope.
# But the compiler might not reuse registers well across the inner k loop.
# Actually, with the current structure, for each c and k, I load 18 values.
# 18 values * 4 bytes = 72 bytes. Plus 64 accumulators = 256 bytes.
# Total ~330 bytes = ~83 registers. Still manageable.
# But actually, the weights are in shared memory and inputs are in shared memory.
# The compiler can load them directly in the FMA instruction if they're in registers or shared memory.
# Having explicit temporaries might help the compiler schedule loads early.
# Hmm, but there's a bigger issue. In the unrolled code above, for each (c, k), I compute:
# acc[c] += sum of 9 products.
# This means for each c, I need to load 9 weights.
# With 64 values of c, that's 576 weights loaded per k.
# Since the weights don't fit in registers, they get loaded from shared memory each time.
# Shared memory has high bandwidth, so this should be fine.
# But register pressure might be high because the compiler might try to keep some weights around.
# Actually, a cleaner unrolled version:
# for c in 0..63:
# float sum = 0;
# for k in 0..7:
# sum += input[k][ty+0][tx+0] * weight[c][k][0][0] + ... + input[k][ty+2][tx+2] * weight[c][k][2][2];
# acc[c] = sum + bias[c];
# This loads 9 inputs per k, but inputs can be reused across c!
# But with the current code, I'm loading inputs inside the c loop, so they're reloaded 64 times.
# That's bad! 64*8*9 = 4608 input loads from shared memory per thread.
# If I restructure to load inputs once and reuse across c, that would be better.
# Let me restructure:
# float i[8][3][3]; // 72 registers - too many
# Actually, for each k, I can load the 9 input values once, then loop over c:
# for k = 0..7:
# load input[k][ty+i][tx+j] into local variables or just use directly
# for c = 0..63:
# acc[c] += input_val * weight[c][k][i][j];
# But to avoid reloading inputs, I need the 9 values in registers for the c loop.
# With 9 input values in registers * 8 channels = 72 register loads total, but only 9 at a time per k.
# For each k:
# float i00 = input_s[k][ty][tx];
# ...
# float i22 = input_s[k][ty+2][tx+2];
# for c = 0..63:
# acc[c] += i00 * weight_s[c][k][0][0] + ... + i22 * weight_s[c][k][2][2];
#
# Wait, even better, for each spatial position (i, j), load one input and accumulate across c:
# for k = 0..7:
# for i = 0..2:
# for j = 0..2:
# float inp = input_s[k][ty+i][tx+j];
# for c = 0..63:
# acc[c] += inp * weight_s[c][k][i][j];
#
# This loads each input value only once, and each weight once.
# Total shared memory reads: 72 inputs + 4608 weights = 4680 reads.
# The previous version loaded inputs 64 times: 4608 inputs + 4608 weights = 9216 reads.
# This saves half the shared memory reads!
#
# Register usage: 64 accumulators + 1 input + 1 weight + loop vars.
# ~70 registers. Great!
#
# And instruction-level parallelism: the inner loop over c does 64 FMAs with the same input.
# The weight loads are broadcast to all threads in the warp? No, each thread reads weight_s[c] for its own c loop.
# Actually, consecutive threads have the same c loop, but they access the same weight_s values at different times.
# Wait, all threads in a warp execute the same instruction. So when thread 0 does c=0, all threads in the warp do c=0 simultaneously. They all read weight_s[0][k][i][j]. This is a broadcast! Perfect.
# Then c=1: all read weight_s[1][k][i][j]. Broadcast again.
# This is extremely efficient for shared memory.
#
# So the optimal loop order is:
# for k, i, j:
# float inp = input_s[k][ty+i][tx+j];
# for c:
# acc[c] += inp * weight_s[c][k][i][j];
#
# Let me rewrite the kernel with this loop order.
```
Wait, I should be careful about the weight indexing with this loop order. Let me make sure it's correct.
Actually, if I do:
```cuda
for (int c = 0; c < out_channels; c++) {
acc[c] = bias_s[c];
}
for (int k = 0; k < in_channels; k++) {
for (int i = 0; i < kernel_size; i++) {
for (int j = 0; j < kernel_size; j++) {
float inp = input_s[k][threadIdx.y + i][threadIdx.x + j];
for (int c = 0; c < out_channels; c++) {
acc[c] += inp * weight_s[c][k][i][j];
}
}
}
}
```
Yes, this is correct. Each input is used for all output channels.
But with `out_channels=64`, the inner loop is 64. Will the compiler unroll it? If I add `#pragma unroll`, it might try to unroll 64 iterations, which would create a huge basic block and might hurt instruction cache. Better to not unroll the c loop, or partially unroll it.
Actually, for 64 iterations, the loop overhead is tiny (one integer add and compare per iteration). The bigger issue is whether the compiler keeps all 64 accumulators in registers and can schedule the FMAs well.
With H100's high issue rate, a loop of 64 FMAs should be well-pipelined.
Let me think about the unrolled version for `kernel_size == 3`:
```cuda
#pragma unroll
for (int k = 0; k < in_channels; k++) {
float i00 = input_s[k][threadIdx.y + 0][threadIdx.x + 0];
float i01 = input_s[k][threadIdx.y + 0][threadIdx.x + 1];
// ... 9 inputs
for (int c = 0; c < out_channels; c++) {
acc[c] += i00 * weight_s[c][k][0][0] + i01 * weight_s[c][k][0][1] + ...;
}
}
```
This keeps 9 inputs in registers for the duration of the c loop. That's 9 registers. Plus 64 accumulators. Plus 1 for the sum of products if the compiler optimizes that way. Total ~75 registers.
But now, the compiler might fuse the 9 FMAs into a tree or just emit them sequentially. With 64 iterations, that's 576 FMAs per k iteration, times 8 k iterations = 4608 FMAs. Plus the extra adds to combine the 9 products. Actually, for each (c, k), we need to compute:
`acc[c] += i00*w00 + i01*w01 + ... + i22*w22;`
If the compiler generates 9 FMAs and 8 adds per (c,k), that's 17 ops per (c,k), not 9. But FMA does multiply-add, so it can be:
```
t1 = i00*w00;
t2 = t1 + i01*w01;
t3 = t2 + i02*w02;
...
acc[c] = acc[c] + t9;
```
Actually, FMA computes `acc += a*b`. So each product can directly accumulate or use a temporary. For 9 products added together and then added to acc[c]:
```
float sum = i00*w00 + i01*w01 + ... + i22*w22;
acc[c] += sum;
```
This is 9 FMAs (or MULs) + 8 ADDs + 1 ADD = 18 ops, but the compiler might optimize differently.
If instead we accumulate directly:
```
acc[c] += i00*w00;
acc[c] += i01*w01;
...
acc[c] += i22*w22;
```
That's 9 FMAs! Much better. Let's write it this way.
Wait, actually `acc[c] += i00*w00` is `acc[c] = acc[c] + i00*w00`. The compiler can use FMA: `acc[c] = fma(i00, w00, acc[c])`. Yes, FMA is `d = a*b + c`. So it computes `acc[c] = i00*w00 + acc[c]`. Perfect! So 9 FMAs per (c,k).
For 64*8 = 512 (c,k) pairs, that's 4608 FMAs total. ✓
So the code should be:
```cuda
for (int k = 0; k < in_channels; k++) {
float i00 = input_s[k][ty][tx];
float i01 = input_s[k][ty][tx+1];
// etc.
for (int c = 0; c < out_channels; c++) {
acc[c] = fmaf(i00, weight_s[c][k][0][0], acc[c]);
acc[c] = fmaf(i01, weight_s[c][k][0][1], acc[c]);
// ...
}
}
```
Using `fmaf` explicitly might help the compiler use the FMA instruction. But modern compilers should automatically recognize `a += b * c;` as FMA. However, with strict floating point, `a = a + b * c` might not be contracted to FMA unless `-ffast-math` or similar is enabled. PyTorch compiles with `-ftz=true` etc but I'm not sure about contractions.
Actually, to match PyTorch's numerical results exactly, using explicit FMA might cause slight differences because FMA computes `a*b+c` with a single rounding instead of two. If PyTorch uses separate MUL and ADD due to separate operations, the results might differ slightly.
PyTorch's conv2d uses cuDNN, which definitely uses FMA instructions. The elementwise division and leaky_relu are exact. So using FMA should be fine and actually more accurate (only 1 rounding instead of 2). But to be safe and match exactly, I should avoid explicit FMA and just do `acc[c] += inp * w;` and let the compiler decide.
Wait, but PyTorch uses tf32 for cudnn conv on Ampere/Hopper by default for some convs? No, for fp32, cudnn should use fp32 FMAs. But the issue is: does cudnn accumulate in fp32 or fp64 or something else? Standard cudnn conv with fp32 accumulates in fp32. And our kernel accumulates in fp32.
For numerical equivalence, as long as we do the same operations in the same order, we should match. But since conv is a reduction over 72 elements, the order matters for floating point. cuDNN might use different reduction trees or split-K algorithms. So it's unlikely we'd match exactly regardless of our loop order.
The tolerance is `atol=rtol=1e-4` for fp32, which is quite loose. A conv with 72 terms should have error around 1e-6 to 1e-5 relative. 1e-4 tolerance should be easily achievable with any reasonable summation order.
So I don't need to worry too much about exact bit-matching. I'll just write a natural, correct kernel.
Let me reconsider the loop unrolling approach. If use `#pragma unroll` on the k loop and the spatial loops, and keep the c loop rolled, that might be the sweet spot.
Actually, I want to maximize instruction-level parallelism. If the c loop is rolled, the processor can pipeline iterations. With 64 iterations and 9 FMAs per iteration, that's plenty of independent operations to hide latency.
Let me try a more general but still efficient approach:
```cuda
__global__ void fused_conv_div_leakyrelu_kernel(
const float* __restrict__ input,
const float* __restrict__ weight,
const float* __restrict__ bias,
float* __restrict__ output,
int batch_size,
int in_channels,
int out_channels,
int in_h,
int in_w,
int out_h,
int out_w,
int kernel_size,
float inv_divisor,
float negative_slope
) {
const int TILE_X = 16;
const int TILE_Y = 16;
int ow = blockIdx.x * TILE_X + threadIdx.x;
int oh = blockIdx.y * TILE_Y + threadIdx.y;
int n = blockIdx.z;
if (n >= batch_size) return;
int in_h_start = blockIdx.y * TILE_Y;
int in_w_start = blockIdx.x * TILE_X;
__shared__ float input_s[8][18][18];
__shared__ float weight_s[64][8][3][3];
__shared__ float bias_s[64];
// Load input tile
for (int k = 0; k < in_channels; k++) {
for (int ih = threadIdx.y; ih < TILE_Y + 2; ih += TILE_Y) {
for (int iw = threadIdx.x; iw < TILE_X + 2; iw += TILE_X) {
int in_ih = in_h_start + ih;
int in_iw = in_w_start + iw;
float val = 0.0f;
if (in_ih < in_h && in_iw < in_w) {
val = input[((n * in_channels + k) * in_h + in_ih) * in_w + in_iw];
}
input_s[k][ih][iw] = val;
}
}
}
// Load all weights (small, 18KB total)
int total_weights = out_channels * in_channels * kernel_size * kernel_size;
for (int idx = threadIdx.x + threadIdx.y * TILE_X; idx < total_weights; idx += TILE_X * TILE_Y) {
int tmp = idx;
int j = tmp % kernel_size; tmp /= kernel_size;
int i = tmp % kernel_size; tmp /= kernel_size;
int k = tmp % in_channels; tmp /= in_channels;
int c = tmp;
weight_s[c][k][i][j] = weight[((c * in_channels + k) * kernel_size + i) * kernel_size + j];
}
// Load bias
for (int idx = threadIdx.x + threadIdx.y * TILE_X; idx < out_channels; idx += TILE_X * TILE_Y) {
bias_s[idx] = bias[idx];
}
__syncthreads();
if (ow >= out_w || oh >= out_h) return;
float acc[64];
for (int c = 0; c < out_channels; c++) {
acc[c] = bias_s[c];
}
for (int k = 0; k < in_channels; k++) {
for (int i = 0; i < kernel_size; i++) {
for (int j = 0; j < kernel_size; j++) {
float inp = input_s[k][threadIdx.y + i][threadIdx.x + j];
for (int c = 0; c < out_channels; c++) {
acc[c] += inp * weight_s[c][k][i][j];
}
}
}
}
for (int c = 0; c < out_channels; c++) {
float val = acc[c] * inv_divisor;
val = (val > 0.0f) ? val : (val * negative_slope);
output[((n * out_channels + c) * out_h + oh) * out_w + ow] = val;
}
}
```
Wait, there's a subtle issue: `out_channels` is a runtime variable, but I declare `float acc[64]`. If `out_channels > 64`, this is a buffer overflow. In the constructor, `out_channels` can be any value.
For a robust solution, I should either:
1. Use dynamic shared memory for accumulators (slow)
2. Use a smaller fixed max and process in tiles
3. Just process a chunk of output channels per thread
For generality, I should use the tiled approach with `OC_TILE=8` or similar. This way, `acc[8]` works for any `out_channels`.
Let me think about this more carefully. The problem says "Accept the same constructor arguments as Model." So `out_channels` could be different. Using `acc[64]` is only safe if `out_channels <= 64`. For arbitrary `out_channels`, I need tiling.
I implemented the tiled approach earlier. Let me go back to that but with the optimized loop order (k,i,j outer, c inner).
```cuda
__global__ void fused_conv_div_leakyrelu_kernel(
const float* __restrict__ input,
const float* __restrict__ weight,
const float* __restrict__ bias,
float* __restrict__ output,
int batch_size,
int in_channels,
int out_channels,
int in_h,
int in_w,
int out_h,
int out_w,
int kernel_size,
float inv_divisor,
float negative_slope
) {
const int TILE_X = 16;
const int TILE_Y = 16;
const int OC_TILE = 8;
int ow = blockIdx.x * TILE_X + threadIdx.x;
int oh = blockIdx.y * TILE_Y + threadIdx.y;
int n = blockIdx.z / ((out_channels + OC_TILE - 1) / OC_TILE);
int oc_tile = blockIdx.z % ((out_channels + OC_TILE - 1) / OC_TILE);
int oc_start = oc_tile * OC_TILE;
int oc_end = min(oc_start + OC_TILE, out_channels);
int oc_count = oc_end - oc_start;
if (n >= batch_size) return;
int in_h_start = blockIdx.y * TILE_Y;
int in_w_start = blockIdx.x * TILE_X;
__shared__ float input_s[8][18][18];
__shared__ float weight_s[8][8][3][3];
__shared__ float bias_s[8];
// Load input tile
for (int k = 0; k < in_channels; k++) {
for (int ih = threadIdx.y; ih < TILE_Y + 2; ih += TILE_Y) {
for (int iw = threadIdx.x; iw < TILE_X + 2; iw += TILE_X) {
int in_ih = in_h_start + ih;
int in_iw = in_w_start + iw;
float val = 0.0f;
if (in_ih < in_h && in_iw < in_w) {
val = input[((n * in_channels + k) * in_h + in_ih) * in_w + in_iw];
}
input_s[k][ih][iw] = val;
}
}
}
// Load weights for this tile
int total_weights = oc_count * in_channels * kernel_size * kernel_size;
for (int idx = threadIdx.x + threadIdx.y * TILE_X; idx < total_weights; idx += TILE_X * TILE_Y) {
int tmp = idx;
int j = tmp % kernel_size; tmp /= kernel_size;
int i = tmp % kernel_size; tmp /= kernel_size;
int k = tmp % in_channels; tmp /= in_channels;
int c = tmp;
weight_s[c][k][i][j] = weight[(((oc_start + c) * in_channels + k) * kernel_size + i) * kernel_size + j];
}
// Load bias
if (threadIdx.x < oc_count && threadIdx.y == 0) {
bias_s[threadIdx.x] = bias[oc_start + threadIdx.x];
}
__syncthreads();
if (ow >= out_w || oh >= out_h) return;
float acc[OC_TILE];
for (int c = 0; c < oc_count; c++) {
acc[c] = bias_s[c];
}
for (int k = 0; k < in_channels; k++) {
for (int i = 0; i < kernel_size; i++) {
for (int j = 0; j < kernel_size; j++) {
float inp = input_s[k][threadIdx.y + i][threadIdx.x + j];
for (int c = 0; c < oc_count; c++) {
acc[c] += inp * weight_s[c][k][i][j];
}
}
}
}
for (int c = 0; c < oc_count; c++) {
float val = acc[c] * inv_divisor;
val = (val > 0.0f) ? val : (val * negative_slope);
output[(((n * out_channels + (oc_start + c)) * out_h + oh) * out_w + ow)] = val;
}
}
```
This is much more robust. `acc[8]` works for any `out_channels`. Register usage is low (8-12 accumulators + a few temporaries = ~20 registers). Occupancy will be very high.
But wait, with this tiling, each block only computes 8 output channels for its spatial tile. For `out_channels=64`, we need 8x more blocks. Total blocks = 128 * 8 * 8 * 8 = 65,536.
With 65,536 blocks and ~132 SMs, that's ~496 blocks per SM. If each block uses ~20 registers and 256 threads, that's ~5K registers per block. 64K / 5K = ~12 blocks per SM concurrently. 12 * 256 = 3072 threads per SM. But the max threads per SM on H100 is probably 2048. So we're limited by threads, not registers. ~8 blocks per SM = 2048 threads. That's 100% occupancy!
Wait, actually I'm not sure about H100 max threads. Let me check: Ampere (sm_80) has max 2048 threads per SM. Hopper (sm_90) also has 2048 threads per SM typically.
So 8 blocks * 256 threads = 2048 threads. 100% occupancy. Great!
But we also use shared memory. Each block uses:
- input_s: 8 * 18 * 18 * 4 = 10,368 bytes
- weight_s: 8 * 8 * 3 * 3 * 4 = 2,304 bytes
- bias_s: 8 * 4 = 32 bytes
- total: ~12.7 KB
With 228 KB per SM and 8 blocks: 8 * 12.7 = 101.6 KB. That's under 228 KB, so shared memory won't limit occupancy.
With 100% occupancy and good memory access patterns, this kernel should perform very well.
But will it beat cuDNN? cuDNN is extremely optimized. However, our advantage is fusion plus no extra memory passes. For such a small conv, cuDNN's overhead plus PyTorch's 3-kernel launch overhead might be significant. Our kernel does everything in one launch with one memory pass for output.
Let me also think about global memory coalescing for output writes. In this kernel, threads in a warp have:
- Same n, same oh, varying ow, same c (within the oc_tile).
Wait, no. `ow` varies by `threadIdx.x`, so threads in a warp have consecutive `ow`. And they write output for `oc_start + c` where `c` is the loop index. For a fixed `c`, consecutive threads write to consecutive memory addresses. This is perfectly coalesced.
For input loading, threads load consecutive `iw` values. Coalesced.
For weight loading, threads load consecutive elements from global memory (because `idx` is linear in thread index). Coalesced.
Everything looks good!
One more thing: the grid size calculation. For `blockIdx.z`, we need:
```python
num_oc_tiles = (out_channels + OC_TILE - 1) // OC_TILE
grid_z = batch_size * num_oc_tiles
```
For out_channels=64, num_oc_tiles=8, grid_z=1024.
Grid dimensions: (8, 8, 1024) or more precisely:
- grid_x = (out_w + 15) // 16 = 8
- grid_y = (out_h + 15) // 16 = 8
- grid_z = batch_size * num_oc_tiles = 1024
Total blocks: 8 * 8 * 1024 = 65,536.
Now, the CUDA grid has max limits. For Hopper, `gridDim.x` and `y` max is 2^31-1, `z` max is 65535. Oops! `grid_z = 1024` is fine since it's < 65535.
Actually wait, CUDA grid max is:
- x: 2^31 - 1
- y: 65535
- z: 65535
So grid_z = 1024 is well within the 65535 limit. Good.
For larger batch sizes or out_channels, grid_z could exceed 65535. For example, batch=128 and out_channels=512 would give 128*64=8192. Still fine. Batch=1024 and out_channels=512 would give 65536, which exceeds the limit.
For robustness, I might want to restructure the grid if z could be large. But for this problem, the given dimensions are small. I'll add a check and handle it by moving some dimensions to x or y if needed, but for simplicity I'll assume it fits.
Actually, I can just use a 1D grid and compute block indices manually. Or use 2D grid for spatial and iterate over batch+oc in z. For safety, I'll keep it as is but note that for very large z it might fail.
One more optimization: I should consider half-precision or other optimizations. But the input and weights are fp32, and we need fp32 output. So fp32 it is.
Let me also think about whether I should use warp-level tiled matrix multiplication for the conv. With `mma` instructions, I could use Tensor Cores. But the problem says the model performs a simple conv + div + leaky_relu. Using Tensor Cores for a 3x3 conv with C_in=8 is overkill and the dimensions don't align well with mma shapes (e.g., 16x16x16). We could do it, but it's a lot of complexity for probably minimal gain.
Actually, wait. For implicit GEMM with Tensor Cores:
- We can view the conv as a GEMM: output = weight_matrix * im2col_matrix.
- weight_matrix: [64, 72]
- im2col_matrix: [72, 128*126*126]
- output_matrix: [64, 128*126*126]
This is a GEMM of size 64x128*126*126x72. That's tiny! Using Tensor Cores for a 64x72 GEMM would have very low utilization. So direct conv is definitely better.
OK, let me write the full Python module.
One more consideration: `load_inline`. I need to make sure the code compiles. I'll use a simple wrapper function that the pybind11 auto-generated module can call.
```python
import torch
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void fused_conv_div_leakyrelu_kernel(...) {
// ... kernel code ...
}
torch::Tensor fused_conv_div_leakyrelu(
torch::Tensor input,
torch::Tensor weight,
torch::Tensor bias,
float divisor,
float negative_slope
) {
int batch_size = input.size(0);
int in_channels = input.size(1);
int in_h = input.size(2);
int in_w = input.size(3);
int out_channels = weight.size(0);
int kernel_size = weight.size(2);
int out_h = in_h - kernel_size + 1;
int out_w = in_w - kernel_size + 1;
auto output = torch::empty({batch_size, out_channels, out_h, out_w}, input.options());
const int TILE_X = 16;
const int TILE_Y = 16;
const int OC_TILE = 8;
dim3 block(TILE_X, TILE_Y);
dim3 grid((out_w + TILE_X - 1) / TILE_X, (out_h + TILE_Y - 1) / TILE_Y, batch_size * ((out_channels + OC_TILE - 1) / OC_TILE));
fused_conv_div_leakyrelu_kernel<<<grid, block>>>(
input.data_ptr<float>(),
weight.data_ptr<float>(),
bias.data_ptr<float>(),
output.data_ptr<float>(),
batch_size,
in_channels,
out_channels,
in_h,
in_w,
out_h,
out_w,
kernel_size,
1.0f / divisor,
negative_slope
);
return output;
}
"""
functions = ["fused_conv_div_leakyrelu"]
fused_conv = load_inline(
name="fused_conv",
cuda_sources=cuda_source,
functions=functions,
extra_cuda_cflags=["-O3", "--use_fast_math"],
)
```
Wait, I need to be careful. The wrapper function `fused_conv_div_leakyrelu` is the one listed in `functions=[...]`. It takes torch Tensors and returns a torch Tensor. The auto-generated pybind11 bindings will handle the conversion.
I don't need to include `<torch/extension.h>` from the kernel file if I'm using `load_inline`, but it's fine to include it.
Actually, for `load_inline`, I should provide the CUDA source and function names. The CUDA source can contain the kernel and the wrapper. `load_inline` will compile it into a shared library and expose the functions.
One issue: `load_inline` caches by default. If I change the code, I might need to set `verbose=True` or remove the cache. But for the submission, the code only runs once.
Let me write the complete ModelNew class.
```python
class ModelNew(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, divisor):
super(ModelNew, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
self.divisor = divisor
def forward(self, x):
weight = self.conv.weight
bias = self.conv.bias
return fused_conv.fused_conv_div_leakyrelu(x, weight, bias, self.divisor, 0.01)
```
Wait, but if `self.conv.bias` is None (if bias=False was somehow passed), we'd crash. But the reference uses default `bias=True`. For safety, I should handle it, but it's not needed.
One more issue: `fused_conv` is loaded at module level. If the compiled module has issues, it might fail. I should make sure the code is correct.
Also, I need to include `import torch.nn as nn` in the module.
Let me review the weight indexing in the kernel. The weights are stored as [out_channels, in_channels, kernel_size, kernel_size] in PyTorch. In my kernel, I access:
```cuda
weight[(((oc_start + c) * in_channels + k) * kernel_size + i) * kernel_size + j]
```
Is this right?
- For weight[c][k][i][j] in row-major (which CUDA uses):
- stride for c: in_channels * kernel_size * kernel_size
- stride for k: kernel_size * kernel_size
- stride for i: kernel_size
- stride for j: 1
So index = c * (in_channels * K * K) + k * (K * K) + i * K + j.
My formula: `(((oc_start + c) * in_channels + k) * kernel_size + i) * kernel_size + j`
= `(oc_start + c) * in_channels * kernel_size * kernel_size + k * kernel_size * kernel_size + i * kernel_size + j`. ✓
For bias: `bias[oc_start + c]`. ✓
For input: `input[((n * in_channels + k) * in_h + in_ih) * in_w + in_iw]`
= `n * in_channels * in_h * in_w + k * in_h * in_w + in_ih * in_w + in_iw`. ✓
For output: `output[(((n * out_channels + (oc_start + c)) * out_h + oh) * out_w + ow)]`
= `n * out_channels * out_h * out_w + (oc_start + c) * out_h * out_w + oh * out_w + ow`. ✓
Everything looks correct.
Now let me think about potential numerical issues. The kernel computes the same operations as the reference:
1. conv: sum of input * weight + bias
2. divide by divisor
3. leaky_relu
The order of summation in the conv might differ from cuDNN. cuDNN probably uses some optimized reduction tree or split-K, while we do a straightforward nested loop. This could lead to slightly different results due to floating point associativity.
However, the tolerance is 1e-4 for fp32. For a sum of ~72 terms with values ~1, the rounding error is at most ~72 * epsilon * ~1 ≈ 72 * 1e-7 = 7e-6, well within 1e-4. Even with different associativity, the error should be within 1e-5 or so.
Actually, wait. The inputs are random ~Uniform(0,1), weights are initialized by PyTorch (Kaiming uniform, std ~ 1/sqrt(fan_in) = 1/sqrt(8*3*3) ≈ 0.068). So weights are small. Inputs are ~[0,1]. Products are ~[0, 0.068]. Sum of 72 such products: mean ~ 72 * 0.034 = 2.4, std ~ sqrt(72) * 0.068 / sqrt(3) ≈ 0.33. Max maybe a few. So values are modest.
Bias is initialized similarly small. Division by 2 makes values smaller. So everything is well-behaved.
Numeric difference due to associativity should be O(epsilon * sum) ≈ 1e-7 * 2 ≈ 2e-7, which is << 1e-4. We should easily pass correctness.
One more thing: I should pass `-O3` to the compiler for optimization. I'll also use `--use_fast_math` which enables faster but slightly less accurate math functions. We only need basic arithmetic here, so fast-math shouldn't affect much.
Actually, I should double check: does `--use_fast_math` affect FMA contraction? It might enable more aggressive optimizations. But since we only need 1e-4 tolerance, it should be fine.
Wait, I just realized: the problem states "Use torch.utils.cpp_extension.load_inline to compile and bind CUDA source at module load time". So the compilation happens when the module is loaded.
Also, I should not write a `PYBIND11_MODULE` block because load_inline auto-generates one.
Let me write the complete code now.
Actually, one more optimization: since `kernel_size` is small, I can precompute `out_h = in_h - kernel_size + 1` and `out_w = in_w - kernel_size + 1`. In the wrapper, I compute these. But the kernel receives them as parameters.
I should also think about whether to use `__restrict__` on all pointers. Yes, this helps the compiler optimize.
Let me also add `-allow-unsupported-compiler` or similar if needed? Usually `load_inline` handles the compilation flags automatically.
Let me draft the code and compile it.
```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>
__global__ void fused_conv_div_leakyrelu_kernel(
const float* __restrict__ input,
const float* __restrict__ weight,
const float* __restrict__ bias,
float* __restrict__ output,
int batch_size,
int in_channels,
int out_channels,
int in_h,
int in_w,
int out_h,
int out_w,
int kernel_size,
float inv_divisor,
float negative_slope
) {
const int TILE_X = 16;
const int TILE_Y = 16;
const int OC_TILE = 8;
int ow = blockIdx.x * TILE_X + threadIdx.x;
int oh = blockIdx.y * TILE_Y + threadIdx.y;
int n = blockIdx.z / ((out_channels + OC_TILE - 1) / OC_TILE);
int oc_tile = blockIdx.z % ((out_channels + OC_TILE - 1) / OC_TILE);
int oc_start = oc_tile * OC_TILE;
int oc_end = min(oc_start + OC_TILE, out_channels);
int oc_count = oc_end - oc_start;
if (n >= batch_size) return;
int in_h_start = blockIdx.y * TILE_Y;
int in_w_start = blockIdx.x * TILE_X;
__shared__ float input_s[8][18][18];
__shared__ float weight_s[8][8][3][3];
__shared__ float bias_s[8];
// Cooperative loading of input tile
for (int k = 0; k < in_channels; k++) {
for (int ih = threadIdx.y; ih < TILE_Y + 2; ih += TILE_Y) {
for (int iw = threadIdx.x; iw < TILE_X + 2; iw += TILE_X) {
int in_ih = in_h_start + ih;
int in_iw = in_w_start + iw;
float val = 0.0f;
if (in_ih < in_h && in_iw < in_w) {
val = input[((n * in_channels + k) * in_h + in_ih) * in_w + in_iw];
}
input_s[k][ih][iw] = val;
}
}
}
// Cooperative loading of weights for this output channel tile
int total_weights = oc_count * in_channels * kernel_size * kernel_size;
for (int idx = threadIdx.x + threadIdx.y * TILE_X; idx < total_weights; idx += TILE_X * TILE_Y) {
int tmp = idx;
int j = tmp % kernel_size; tmp /= kernel_size;
int i = tmp % kernel_size; tmp /= kernel_size;
int k = tmp % in_channels; tmp /= in_channels;
int c = tmp;
weight_s[c][k][i][j] = weight[(((oc_start + c) * in_channels + k) * kernel_size + i) * kernel_size + j];
}
// Loading bias
if (threadIdx.x < oc_count && threadIdx.y == 0) {
bias_s[threadIdx.x] = bias[oc_start + threadIdx.x];
}
__syncthreads();
if (ow >= out_w || oh >= out_h) return;
float acc[OC_TILE];
for (int c = 0; c < oc_count; c++) {
acc[c] = bias_s[c];
}
for (int k = 0; k < in_channels; k++) {
for (int i = 0; i < kernel_size; i++) {
for (int j = 0; j < kernel_size; j++) {
float inp = input_s[k][threadIdx.y + i][threadIdx.x + j];
for (int c = 0; c < oc_count; c++) {
acc[c] += inp * weight_s[c][k][i][j];
}
}
}
}
for (int c = 0; c < oc_count; c++) {
float val = acc[c] * inv_divisor;
val = (val > 0.0f) ? val : (val * negative_slope);
output[(((n * out_channels + (oc_start + c)) * out_h + oh) * out_w + ow)] = val;
}
}
torch::Tensor fused_conv_div_leakyrelu(
torch::Tensor input,
torch::Tensor weight,
torch::Tensor bias,
float divisor,
float negative_slope
) {
int batch_size = input.size(0);
int in_channels = input.size(1);
int in_h = input.size(2);
int in_w = input.size(3);
int out_channels = weight.size(0);
int kernel_size = weight.size(2);
int out_h = in_h - kernel_size + 1;
int out_w = in_w - kernel_size + 1;
auto output = torch::empty({batch_size, out_channels, out_h, out_w}, input.options());
const int TILE_X = 16;
const int TILE_Y = 16;
const int OC_TILE = 8;
dim3 block(TILE_X, TILE_Y);
dim3 grid((out_w + TILE_X - 1) / TILE_X,
(out_h + TILE_Y - 1) / TILE_Y,
batch_size * ((out_channels + OC_TILE - 1) / OC_TILE));
fused_conv_div_leakyrelu_kernel<<<grid, block>>>(
input.data_ptr<float>(),
weight.data_ptr<float>(),
bias.data_ptr<float>(),
output.data_ptr<float>(),
batch_size,
in_channels,
out_channels,
in_h,
in_w,
out_h,
out_w,
kernel_size,
1.0f / divisor,
negative_slope
);
return output;
}
"""
fused_conv = load_inline(
name="fused_conv_ext",
cuda_sources=cuda_source,
functions=["fused_conv_div_leakyrelu"],
extra_cuda_cflags=["-O3", "--use_fast_math"],
)
class ModelNew(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, divisor):
super(ModelNew, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
self.divisor = divisor
def forward(self, x):
return fused_conv.fused_conv_div_leakyrelu(
x, self.conv.weight, self.conv.bias, self.divisor, 0.01
)
```
Wait, I need to make sure the code handles `out_channels` that aren't multiples of `OC_TILE`. My code does handle it with `oc_end = min(...)`. Good.
But there's a potential issue: `dim3 grid(...)` with the z dimension being `batch_size * ((out_channels + OC_TILE - 1) / OC_TILE)`. If this exceeds 65535, it will fail. For the given problem, this is 128 * 8 = 1024. Fine.
Another issue: I need to make sure `input`, `weight`, `bias` are contiguous and on CUDA. PyTorch's `data_ptr<float>()` works for contiguous tensors. If they're non-contiguous, we might have issues. The reference uses nn.Conv2d which expects standard NCHW tensors. So it should be fine.
Also, I assume `input` is on CUDA. The reference calls it on CUDA in practice (since we're competing with CUDA reference), so we should probably call `.cuda()` or just assume it. But the wrapper doesn't check device. It should be fine.
One thing I want to add: since `self.conv` has parameters that need gradients, but we're bypassing the PyTorch autograd for the forward pass. The problem doesn't mention training or backward pass. It just says "Replace the forward computation". So we only need inference. Great.
Wait, but if someone tries to call `.backward()`, PyTorch won't know how to backprop through our custom op because it doesn't have a `torch.autograd.Function` wrapper. The problem only asks to replace forward, so inference is the only concern.
Hmm, actually there's a caveat. The reference code is a `nn.Module`, and in PyTorch, modules with parameters usually support backward if the operations are differentiable. But the problem says "Replace the forward computation". It doesn't require backward support. So a custom CUDA kernel is fine.
However, the evaluation might call `model.eval()` and do inference only. Let me assume so.
One more optimization: I can save `1.0f / divisor` as `inv_divisor` and multiply instead of divide. Division is slower than multiplication. My code already does this.
Another thing: I should check if `input` is contiguous. If not, accessing via raw pointer will be wrong. For safety, I can call `input.contiguous()` in the wrapper. But contiguous() might create a copy. For the default case, inputs are already contiguous.
Actually, `torch.rand(...)` produces a contiguous tensor. And even if not, the standard conv2d handles it. To be safe:
```python
input = input.contiguous()
```
But since we're already competing for speed, an unnecessary copy would be bad. I'll trust that inputs are contiguous.
Let me now actually compile this using the tool to check for errors.