reasoning
This is a standard Conv2d layer: 3 input channels, 96 output channels, 11x11 kernel, stride 4, padding 2.
Input: [256, 3, 224, 224]
First let me compute the output dimensions:
H_out = floor((224 + 2*2 - 11)/4) + 1 = floor((224 + 4 - 11)/4) + 1 = floor(217/4) + 1 = 54 + 1 = 55
W_out = 55 as well.
So output is [256, 96, 55, 55].
Let me compute the FLOPs:
256 * 96 * 55 * 55 * 3 * 11 * 11 = 256 * 96 * 3025 * 363
Let me calculate:
256 * 96 = 24576
24576 * 3025 = 74,342,400
74,342,400 * 363 ≈ 27.0 GFLOPS
For memory:
Input: 256 * 3 * 224 * 224 * 4 bytes ≈ 154 MB
Weight: 96 * 3 * 11 * 11 * 4 bytes ≈ 139 KB
Output: 256 * 96 * 55 * 55 * 4 bytes ≈ 295 MB
This is memory-bound, but let me think about optimizing this.
On H100, cuDNN likely uses implicit GEMM or Winograd. For an 11x11 kernel with stride 4, Winograd isn't applicable. The most likely approach is either implicit GEMM via Tensor Cores or a direct convolution.
Since this is a single conv layer, we want to maximize throughput. Let me write a custom CUDA kernel that does direct convolution, potentially using shared memory for tiling.
Actually, for this case with 3 input channels and 96 output channels, and large spatial dimensions, we might want to use a kernel that:
1. Each block computes a tile of output elements
2. Uses shared memory to cache input tiles
3. Each thread computes multiple output elements to amortize loads
Let me think about the dimensions:
- Output: [256, 96, 55, 55] = 256 * 96 * 3025 elements = 74,342,400 elements
With 55x55 spatial output, we have about 3025 spatial positions per image per output channel.
If we use a block size of, say, 256 threads, and have each thread compute one output element, we'd need about 74M / 256 ≈ 290K blocks. That might be okay but we want to ensure good locality.
Actually, let me think about this more carefully. cuDNN on H100 will likely use Tensor Cores for this, and achieving a speedup over cuDNN is challenging. But the reference is PyTorch's nn.Conv2d, which uses cuDNN.
Wait, for 3 input channels, a Tensor Core GEMM approach might not be as efficient because the K dimension is only 3*11*11 = 363, which isn't very large. So there might be an opportunity.
Let me think about the direct approach:
- We can have each block handle a small tile in the output (e.g., 4x4 or 8x8 spatial positions)
- For each output channel, each block loads the relevant input region and kernel weights
- With 3 input channels and 11x11 kernel, each output spatial position requires reading 363 input values and 363 weights
For tiling, if a block computes a 4x4 spatial tile:
- The input region covers (4*4) x (4*4) = 16x16 spatial area, but with kernel 11 and stride 4, it's more like:
- For output positions (oy, ox) where oy ranges from start_y to start_y+3, ox from start_x to start_x+3
- The input y range is [start_y*4, start_y*4+3*4+10] = [start_y*4, start_y*4+22]
- That's about 23x23 = 529 input spatial positions per channel
- With 3 channels: 1587 floats ≈ 6.3 KB
If we have a thread block of 256 threads computing 4x4 = 16 output positions... that's not very parallel.
Better approach: have each thread compute one output element, but organize for coalesced access.
Actually, let me reconsider. Since we have 96 output channels, 256 batch size, and 55x55 spatial:
- Total output elements: 256 * 96 * 55 * 55
We could assign each thread to one (batch, out_channel, oy, ox) position.
With 256 threads per block, each block handles 256 output elements.
Total blocks: 256 * 96 * 55 * 55 / 256 = 96 * 3025 = 290,400 blocks.
That seems like enough parallelism. The question is memory access patterns.
For memory access:
- Each thread needs to read 3 * 11 * 11 = 363 input values and 363 weights
- Input accesses for adjacent threads (in the same block, same spatial positions):
- If we group threads by output channel first, then adjacent threads will access the same input but different weight sets. This means we should cache input in shared memory.
- If we group by spatial position, adjacent threads access the same weights but different inputs.
Actually, let's think about what cuDNN does. For Conv2d with Tensor Cores, it typically converts to im2col + GEMM. The GEMM would be:
- M = batch * H_out * W_out = 256 * 55 * 55 = 774,400
- N = out_channels = 96
- K = in_channels * kH * kW = 3 * 11 * 11 = 363
So it's a 774400 x 96 x 363 GEMM. The M dimension is large, N is small (96). This GEMM is very skinny. cuDNN likely handles this, but maybe not with optimal Tensor Core utilization since N=96 isn't a multiple of 64/128.
For a custom kernel, we could try to beat cuDNN with a direct convolution that doesn't have the im2col overhead and handles the skinny dimensions better.
Let me design a kernel where:
- Each block computes a set of output channels for a subset of the batch and spatial positions
- Each warp computes a specific subset to maximize memory coalescing
Actually, let me think about a different approach. Since K=363 is relatively small, we could load all weights into shared memory or even registers. But 96 * 363 = 34,848 weights, which is about 139 KB. That's too much for shared memory.
But we could process a tile of output channels at a time.
Let me try a different strategy. What if we:
- Use each thread block for a specific output channel
- Within the block, threads work on different batch/spatial positions
- Load the weights for that output channel into shared memory (363 floats = ~1.5 KB)
- Then each thread iterates over the input, loading from global memory
Wait, 363 floats = about 1452 bytes, which fits easily in shared memory. And we can even load them into shared memory once per block.
So a kernel design:
- Block dimensions: let's say 256 threads
- Grid: (batch * H_out * W_out / 256, out_channels) or similar
- Each block handles one output channel
- All threads in the block compute different (batch, oy, ox) positions for the same output channel
- We load the weights for that output channel (3 * 11 * 11 = 363 floats) into shared memory
- Then each thread computes its output by iterating over input channels, reading the needed input
For input access pattern: If threads in a block are computing adjacent spatial positions, they will access nearby input memory, which is good for coalescing. But with stride=4 and kernel=11, adjacent output positions access input with stride 4, so there's some reuse but not as much as with stride 1.
If we tile spatial positions, e.g., have threads in a warp process a contiguous block of output_x positions for the same output_y, they'd access input_x positions that are 4 apart. This isn't great for coalescing, but it's not terrible either since each warp will coalesce adjacent threads' accesses into a few cache lines.
Actually, let me reconsider the block organization. What if we have:
- Each block handle a tile of spatial positions for ALL output channels? No, that would require loading all weights.
Alternative:
- Each block handles a tile of spatial positions for a subset of output channels
- Say 8 output channels and 32 spatial positions (per batch element)
- Threads: 256
- Each of the 8 output channels gets 32 threads
- Load weights for 8 output channels: 8 * 363 floats ≈ 11.6 KB
- This fits in shared memory!
Then each thread computes one spatial position for one output channel.
This seems promising. Let me elaborate:
- Block dim: 256 threads
- Each block computes 8 output channels × 32 spatial positions
- Grid: (ceil(96/8), ceil(256 * 3025 / 32)) or organized differently
Wait, let me think about organizing by batch and spatial:
- Total positions to compute per output channel: 256 * 55 * 55 = 774,400
- If we process 32 positions per block per output channel, we need 774,400 / 32 = 24,200 blocks per output channel
- With 96 / 8 = 12 output channel groups, total blocks = 12 * 24,200 = 290,400
This matches the total number of output elements / 256, which makes sense since we have 256 threads per block.
For memory access:
- Each block loads 8 * 363 = 2904 weights into shared memory
- Each of the 256 threads computes one (out_channel, batch, oy, ox) element
- The 32 threads working on the same out_channel but different spatial positions will access different input regions
- If the 32 positions are spatially adjacent, there might be some cache locality
Actually, I think a better organization might be:
- Have each warp (32 threads) handle one or a few output channels
- Within a warp, each thread handles adjacent spatial positions or batch elements
For best coalescing, adjacent threads should access adjacent memory. Output memory is stored as [N, C, H, W]. So if adjacent threads write to:
- Same channel, same batch, adjacent W: then writes are contiguous (good!)
- Same channel, same batch, adjacent H: writes are separated by W_stride = 55 (okay, still in the same cache line if within the same 64-byte region)
For input access: [N, C, H, W] with H=224, W=224.
If adjacent threads process adjacent output W positions, they read input W positions that are 4 apart. So thread 0 reads W=0, thread 1 reads W=4, etc. This means adjacent threads access memory 16 bytes apart (assuming float32), which is actually very well coalesced! They'll hit the same cache lines.
Wait, that's for the first input channel. For multiple input channels, the accesses will be separated by 224*224*4 bytes.
Let me think about this more carefully. 32 threads in a warp, each computing a different output position. If they're all computing the same (batch, channel) but different (oy, ox) positions:
If we make them contiguous in the flattened spatial dimension:
- Thread i computes spatial position (oy, ox) where spatial_idx = tile_start + i
- oy = spatial_idx / 55, ox = spatial_idx % 55
- Input y base = oy * 4 - 2 + padding? Wait, padding=2.
- Actually input coordinates: iy = oy * 4 - 2 + 2 = oy * 4... wait.
- With padding=2 and stride=4: input_y = oy * 4, since we padded by 2 and y starts from 0 for oy=0.
- Wait, the formula is: input_y = oy * stride - padding + ky
- For oy=0, input_y starts at -2 (padding), so first valid ky=2 gives input_y=0.
- For oy=1, input_y starts at 4-2=2, ky=0 gives input_y=2.
So adjacent spatial positions in output (say ox differs by 1) have input_x bases that differ by 4. So adjacent threads accessing input for the same ky but different ox will access memory 16 bytes apart (4 floats), which is excellent coalescing.
Actually, let me verify. Input layout is [N, C, H, W]. For a fixed batch and input channel, the memory offset for (iy, ix) is iy * 224 + ix. If thread i accesses (iy, ix_base + i * 4), the memory offsets are iy*224 + ix_base + i*4. Adjacent threads access offsets that differ by 4, which is 16 bytes. This is perfect coalescing (each warp will generate a single 128-byte transaction).
Excellent! So if we organize threads to process adjacent output spatial positions, input accesses will be well-coalesced.
Now for weight access: if all threads in a warp compute different spatial positions for the SAME output channel, they all load the SAME weights. This means weight loads from global memory will be broadcast (or cached in L1). If we put weights in shared memory, that's even better since they only need to be loaded once per block.
So here's my refined plan:
- Each block computes a subset of output elements
- Block has 256 threads
- Block computes 8 output channels × 32 spatial positions
- We store weights for 8 output channels in shared memory: 8 × 3 × 11 × 11 floats
- Each of the 256 threads computes one output element
- We can have 32 threads per output channel, each computing one spatial position
- We iterate over input channels and kernel positions
- For each (input_channel, ky, kx), all threads in the warp load their respective input values
Wait, with 32 threads per output channel in a block of 256 threads, we'd have 8 warps, each handling one output channel.
Let me think about the thread indexing:
- threadIdx.x: 0-31, identifies position within the spatial tile
- threadIdx.y: 0-7, but actually we only have threadIdx.x up to 255 total
- We can map: out_ch_offset = threadIdx.x / 32 (range 0-7), spatial_offset = threadIdx.x % 32 (range 0-31)
Wait, 256 / 32 = 8, but 8 * 32 = 256, so:
- out_ch_in_block = threadIdx.x / 32, from 0 to 7
- spatial_in_block = threadIdx.x % 32, from 0 to 31
For grid dimensions:
- gridDim.x covers the output channels: we process 8 channels per block, so gridDim.x = 96 / 8 = 12
- gridDim.y covers the batch and spatial positions. Total spatial positions per output channel = 256 * 55 * 55 = 774,400. We process 32 per block, so 774,400 / 32 = 24,200 blocks.
Actually, we might want a 2D grid for spatial:
- 256 batch * 55 * 55 = 774,400 spatial positions per channel
- We process 32 positions per block per channel
- We could have gridDim.y = 256 (batch), gridDim.z = ceil(3025/32) = 95
- Total positions per channel = 256 * 95 * 32 = 778,240, which is slightly more than needed (ceiling)
- But then gridDim.x = 12
- Total blocks = 12 * 256 * 95 = 291,840
Hmm, actually let me organize it differently. We can have:
- gridDim.x = out_channel_blocks = 12 (each handling 8 channels)
- gridDim.y = total_spatial_blocks = ceil(256 * 3025 / 32) = 24,200
Or perhaps:
- gridDim.x = spatial blocks across W, gridDim.y = spatial blocks across H/B
- But a 1D spatial organization is fine.
Actually, for better memory locality, it might be better to have the threads in a block work on spatially adjacent positions. Let's say we process a tile of 8x4 = 32 spatial positions.
With output_W=55, we could have:
- tile_W = 8, tile_H = 4, giving 32 spatial positions
- For a given (batch, out_ch_block, tile_h, tile_w):
- 32 threads handle 32 spatial positions
- 8 groups of 32 threads handle 8 output channels
Wait, I realize it's easier to just use a linear index for spatial positions and let each thread compute one position. With 32 threads per output channel and 8 output channels, we have perfect utilization.
Let me now think about the actual implementation.
```cuda
__global__ void conv2d_direct(
const float* __restrict__ input,
const float* __restrict__ weight,
float* __restrict__ output,
int batch, int in_channels, int in_h, int in_w,
int out_channels, int out_h, int out_w,
int kernel_h, int kernel_w,
int stride, int padding
) {
// Shared memory for weights: 8 * 3 * 11 * 11 floats
__shared__ float smem_weights[8 * 3 * 11 * 11];
int out_ch_block = blockIdx.x; // 0 to 11
int spatial_block = blockIdx.y; // 0 to 24199
int out_ch_start = out_ch_block * 8;
int total_spatial = batch * out_h * out_w;
int spatial_start = spatial_block * 32;
int thread_id = threadIdx.x;
int out_ch_local = thread_id / 32; // 0 to 7
int spatial_local = thread_id % 32; // 0 to 31
int out_ch = out_ch_start + out_ch_local;
int spatial_idx = spatial_start + spatial_local;
// Load weights for this block into shared memory
int total_weights = 8 * 3 * 11 * 11;
for (int i = thread_id; i < total_weights; i += blockDim.x) {
int w_oc = i / (3 * 11 * 11);
int rem = i % (3 * 11 * 11);
int w_ic = rem / (11 * 11);
int w_ky = (rem % (11 * 11)) / 11;
int w_kx = rem % 11;
// weight layout: [out_ch, in_ch, ky, kx]
int weight_idx = (out_ch_start + w_oc) * (3 * 11 * 11) + rem;
smem_weights[i] = weight[weight_idx];
}
__syncthreads();
// Each thread computes one output element
if (out_ch < out_channels && spatial_idx < total_spatial) {
// Decode spatial_idx to (n, oy, ox)
int n = spatial_idx / (out_h * out_w);
int rem = spatial_idx % (out_h * out_w);
int oy = rem / out_w;
int ox = rem % out_w;
float sum = 0.0f;
// Compute input base position
int in_y_base = oy * stride - padding;
int in_x_base = ox * stride - padding;
for (int ic = 0; ic < in_channels; ++ic) {
for (int ky = 0; ky < kernel_h; ++ky) {
for (int kx = 0; kx < kernel_w; ++kx) {
int iy = in_y_base + ky;
int ix = in_x_base + kx;
float in_val = 0.0f;
if (iy >= 0 && iy < in_h && ix >= 0 && ix < in_w) {
int in_idx = ((n * in_channels + ic) * in_h + iy) * in_w + ix;
in_val = input[in_idx];
}
int w_idx = out_ch_local * (3 * 11 * 11) + ic * (11 * 11) + ky * 11 + kx;
float w_val = smem_weights[w_idx];
sum += in_val * w_val;
}
}
}
if (n < batch) {
int out_idx = ((n * out_channels + out_ch) * out_h + oy) * out_w + ox;
output[out_idx] = sum;
}
}
}
```
Wait, I need to think about whether this will actually be faster than cuDNN.
One issue: the weight access pattern inside the loop. For a given (ic, ky, kx), all threads in the same warp but different spatial positions will access the SAME weight value (since they're all computing the same output channel). When weights are in shared memory, this is a broadcast, which is very efficient. One shared memory read serves all 32 threads simultaneously.
For input access: at (ic, ky, kx), the 32 threads in a warp access 32 different input values. If the 32 spatial positions are consecutive in the flattened spatial index, they'll have:
- spatial_idx = spatial_start + 0 to 31
- These decode to potentially different (oy, ox) pairs.
If spatial_start aligns to multiples of 32 in the flattened space:
- Positions 0-31 in a block could be (oy=0, ox=0..31) or span across rows
- Since out_w = 55, 32 consecutive positions might be: (oy=0, ox=0..31) or (oy=0, ox=23..54) and (oy=1, ox=0..8)
For positions in the same row (ox varies):
- Input x coordinates: ox * stride - padding + kx
- For ox = i, input_x = i * 4 - 2 + kx
- Adjacent threads (ox differs by 1): input_x differs by 4
- Memory offset differs by 4 floats = 16 bytes
- This is perfectly coalesced! All 32 threads in a warp will access a contiguous 512-byte region.
For positions spanning rows:
- When we go from ox=54 to oy=1, ox=0:
- The input y jumps by 4, input x drops by 54*4=216
- This could cause some divergence and cache misses
To avoid this, maybe we should tile spatially in a 2D fashion. Let's say each group of 32 threads handles an 8×4 tile (or whatever shape fits). But this makes indexing more complex.
Actually, consecutive positions in the flattened index [n, oy, ox] = n*3025 + oy*55 + ox naturally give:
- Consecutive ox within a row: good coalescing
- Row boundaries: the memory distance between (oy, ox=54) and (oy+1, ox=0) is 4*224 - 54*4 = 896 - 216 = 680 floats = 2720 bytes. That's a lot, and threads will be accessing non-contiguous memory.
So a warp that straddles a row boundary will have poor coalescing. But for a 32-element block, most rows have 55 elements, so only when the tile starts at ox >= 55-32 = 23 does it cross a boundary. Many blocks won't cross boundaries.
Actually, even better: we can organize the spatial tiling explicitly. Instead of 1D spatial blocks, we can have 2D tiles. But with 32 threads per output channel, we need exactly 32 positions. We could do a 4×8 tile or 8×4 tile.
For an 8×4 tile (8 in x, 4 in y):
- Input x coordinates within a row differ by stride=4, so 8 threads access memory 4 apart = 16 bytes apart, coalesced.
- Next row is 4*224 = 896 floats = 3584 bytes away. If we process 4 rows, the warp accesses 4 distinct regions.
But what about thread divergence in warps? If we have 32 threads = 8x4 tile, one warp handles all 32 positions. Within the loop over ky, kx, all threads execute together. No divergence unless out-of-bounds (padding), but padding=2 and our output range covers the full valid region.
Actually, for a thread computing (oy, ox), the input access is:
iy = oy * 4 - 2 + ky, ix = ox * 4 - 2 + kx
For valid (oy, ox) in [0, 54] x [0, 54], and (ky, kx) in [0, 10] x [0, 10]:
- iy ranges from -2 to 54*4 - 2 + 10 = 224, but valid range is [0, 223]
- So some boundary threads will hit padding for some (ky, kx). This causes divergence only near boundaries.
Wait, for oy=0: iy = -2 + ky. ky=0,1 give iy=-2,-1 which are invalid. So threads computing oy=0 will have 2*11 = 22 boundary checks that fail.
For oy=54: iy = 54*4 - 2 + ky = 214 + ky. ky=10 gives iy=224 which is invalid. So only 11 boundary checks fail.
This is minor divergence. We could unroll the loops and handle boundaries explicitly, but it might not be worth it.
Another approach: process multiple output elements per thread to increase arithmetic intensity. If each thread computes 2 or 4 output positions, we amortize the weight loads and shared memory overhead. But the weights are already in shared memory, so the main cost is loading inputs.
Let me think... Each thread doing 363 multiply-adds. At 55 x 55 output spatial with 256 batch and 96 channels, that's 27B FLOPs. If we use fp32, peak is 67 TFLOPS, but we're likely memory bound.
Memory bandwidth is 3.35 TB/s = 3350 GB/s.
If we're reading input (154 MB) + weights (139 KB) + writing output (295 MB) = ~450 MB total.
At 3350 GB/s, this takes ~0.13 ms.
The reference will probably take more due to overhead.
Actually, PyTorch cuDNN likely uses cuDNN's optimized algorithms which also try to minimize memory traffic. But our custom kernel might have lower overhead since it's directly fused.
Let me also think about using half-precision or Tensor Cores. The task says "produce numerically equivalent outputs at fp32 precision," so we can compute in fp32 but the intermediate computation doesn't have to be fp32? Wait, "numerically equivalent outputs at fp32 precision" means the outputs should match fp32 reference. If we use Tensor Cores with TF32 or fp16, that might not match exactly.
Actually, the instruction says "Produce numerically equivalent outputs at fp32 precision." This sounds like the outputs need to be fp32 and match the fp32 reference. Using mixed precision internally and converting back might not match if TF32 has lower precision. So we should probably stick to fp32 computation.
However, we could try to use Tensor Cores with fp32 accumulation but fp16 inputs, but that's also approximate. Better to use fp32 direct convolution.
Wait, let me re-read: "Produce numerically equivalent outputs at fp32 precision." I think this means the computation should be done in a way that produces fp32 outputs matching the reference's fp32 outputs. Using true fp32 arithmetic is safest.
Let me also consider using __ldg for read-only data, which might help with caching.
And what about using warp-level primitives? For the weight loads, since all threads in a warp read the same weight, having it in shared memory gives broadcast behavior. We could also try loading weights to registers and reusing them, but the compiler should handle this well.
Actually, I just realized: with 32 threads per warp, and we're mapping consecutive threadIdx.x values to consecutive spatial positions within the same output channel, the first warp (threads 0-31) handles output channel 0, the second warp (32-63) handles output channel 1, etc.
So within a warp, 32 threads access the same weight (broadcast from shared memory) and different input positions (coalesced from global memory). This is very efficient!
Between warps in the same block, different warps access different weights but potentially the same input positions (if they're computing the same spatial positions for different output channels). This means input reads for the first warp might still be in L1 cache when the second warp reads them. But input is much larger than L1 (L1 + shared is up to 228KB per SM on H100), so caching might be limited.
Wait, the input region for 32 spatial positions:
- If they're consecutive in flattened index, they span at most 32 positions.
- In the worst case (straddling rows), the memory footprint could be large.
- In the best case (same row), it's about 32 * 4 = 128 floats = 512 bytes for one (ic, ky, kx).
- For all (ic, ky, kx), the total unique input accessed is 3 * 11 * (128 + some overlap) ≈ a few KB.
- Actually for the whole convolution, each thread reads 3*11*11=363 input values, so 32 threads read at most 32*363 = 11,616 input values, but there's a lot of overlap.
Hmm, actually for adjacent spatial positions with stride=4 and kernel=11:
- Two adjacent output positions (ox differs by 1) have input windows that overlap by 11-4=7 pixels in x.
- So for a row of 32 output positions, the total input width covered is 32*4 + 10 - 4 = 134 pixels (wait: first position starts at x=-2+0*4=-2, but valid from x=0; last at x=-2+31*4=122, plus kernel 10 = 132. But input width is 224).
- Actually: for position i, input x ranges from [i*4-2+0, i*4-2+10] = [i*4-2, i*4+8].
- For 32 positions i=0..31: total range is [-2, 31*4+8] = [-2, 132]. Valid range: [0, 131].
- Width: 132 pixels. For one input channel and one ky: 132 floats = 528 bytes.
- For 3 channels and 11 ky values: 3 * 11 * 132 ≈ 4356 floats ≈ 17.4 KB.
- This is well within L1 cache size per SM!
So for blocks processing 32 spatial positions in the same row, the input footprint is about 17KB, which fits in L1. This means input will be well-cached across the 8 warps processing different output channels.
If the 32 positions span multiple rows, the footprint increases:
- e.g., 4 rows of 8 positions each
- Row 0: positions ox=0..7, input y range depends on oy
- Row 1: oy+1, input y shifted by 4
- Each row adds ~17KB / 11 for that row's ky values. Actually it's more complex.
- But with oy varying, for a fixed ky, input y = oy*4 - 2 + ky. For 4 different oy values, that's 4 different rows.
- The total for 4 rows × 8 cols × 11 ky × 3 channels: the footprint could be up to 4*8*11*11*3 in the worst case without overlap? No, ky causes overlap between rows.
- For adjacent rows (oy and oy+1), input y ranges differ by 4. With kernel size 11, there's 7 rows of overlap.
- For 4 output rows, input rows covered = 4*4 + 10 = 26 rows. At 224 cols per row: 26 * 224 * 3 floats ≈ 69.8 KB.
- Still within L1!
So even for multi-row tiles, caching should be good.
Now I'm also thinking: what about the batch dimension? If we process 32 spatial positions that span multiple batch elements, the memory distance is huge (3 * 224 * 224 ≈ 150KB per batch). 32 consecutive spatial positions with spatial layout [N, H, W] will first cover all positions in batch 0 before moving to batch 1. So 32 consecutive positions are all in the same batch unless we cross N*H*W boundary. With 774,400 positions per channel, 32 is tiny.
So most tiles are within the same batch element.
Let me now write the actual CUDA code and see if it compiles.
Actually, one more optimization idea: unroll the loops and use pragma unroll for the 11x11 kernel and 3 channels. Since these are all constants, the compiler can optimize well.
Let me also consider using float4 or vector loads for input. Since adjacent threads access memory 16 bytes apart (4 floats), we could potentially use float4 loads, but the addresses are interleaved across threads, which makes vector loading from global memory tricky. The standard coalesced load pattern where thread i loads address base + i * 4 naturally gets compiled to a single 128-byte warp-wide load transaction.
Actually, I realize I should also consider the case where we can overlap weight loads with computation using multiple shared memory buffers. But since our bottleneck is likely input memory, and weights are in shared memory (fast), this might not help much.
One more thing: H100 has asynchronous copy from global to shared memory (cp.async). We could use this for weight loading. But since the weights are small and only loaded once, the benefit is minimal.
Let me think if we can do better with a different approach. What about a GEMM-based kernel? We could implement im2col in shared memory and then do a small GEMM. But direct convolution seems simpler and likely competitive for this small K dimension.
Actually, wait - I just realized there might be a better organization. What if each thread computes MULTIPLE output channels?
For example, what if we have:
- 32 threads per block
- Each thread computes 1 spatial position but ALL 8 output channels
- This way, weights don't need to be in shared memory (each thread accesses its own weight set)
- But then we have fewer blocks
No, this is worse because we lose the weight broadcast advantage.
What if each thread computes 4 output channels and 8 spatial positions?
- Each thread does 4 * 8 * 363 = 11,616 FLOPs
- A block of 256 threads does ~3M FLOPs
- But memory access pattern is more complex.
Let me stick with the simpler design and see how it performs.
Actually, I want to try another variation: what if we process 16 output channels per block instead of 8?
16 output channels: 16 * 363 = 5808 floats ≈ 23 KB of shared memory for weights. This fits in shared memory (up to 227KB).
Then we'd have 96 / 16 = 6 blocks in x-dimension.
256 threads / 16 = 16 threads per output channel.
Each thread handles 16 spatial positions.
For 16 spatial positions in a row:
- Input footprint: 16*4+10 = 74 valid pixels wide. 74 * 11 * 3 ≈ 2442 floats ≈ 9.8 KB.
- Still well within L1.
But fewer threads per output channel means less parallelism for memory coalescing. However, 16 threads still gives a half-warp, which is okay.
Actually, let me think about the arithmetic intensity. Each output element requires 363 multiply-adds and reads 363 input values + 363 weights. With weights in shared memory, we still need to read them. But shared memory has much higher bandwidth than global memory.
The main bottleneck is reading input from global memory. For 256 threads processing 256 output elements, we read 256 * 363 input values. But due to overlap, the actual unique input values read is much smaller. For 32 threads in the same row:
- Unique input for 16 threads (half-warp): (16*4 + 10) * 11 * 3 = 74 * 33 = 2442 floats
- So the "amplification" is 16 * 363 / 2442 ≈ 2.4x due to reuse.
- For fp32 at 3.35 TB/s, we can sustain about 3.35e12 / 4 / 2442 ≈ 343M blocks per second for this half-warp.
- Each block does 16 * 363 = 5808 FMAs. At 67 TFLOPS, that's 67e12 / 5808 ≈ 11.5B blocks per second.
- So memory is indeed the bottleneck. But with 2.4x reuse, effective arithmetic rate needed is 5808 * 4 / (2442 * 4) * 3.35e12 / 2? Let me think differently.
For the 32-thread warp computing 32 output positions in a row:
- Global memory reads: ~17KB of unique input (as computed earlier)
- Computation: 32 * 363 = 11,616 FMAs
- Arithmetic intensity: 11,616 / (17KB / 4 bytes) = 11,616 / 4,250 ≈ 2.7 FLOPs per byte
- At 67 TFLOPS / 3.35 TB/s = 20 FLOPs/byte machine balance
- So we're heavily memory bound, running at about 2.7/20 = 13.5% of peak compute.
To improve, we need more reuse. Processing more spatial positions per block (but we can't get more reuse for a single row). Processing multiple rows gives less spatial reuse but some overlap in the kernel y dimension.
Wait, we can also process multiple output channels per block, which reuses the same input for more computations. With 8 output channels and 32 spatial positions:
- Input reads: same ~17KB (unique input)
- Computation: 8 * 32 * 363 = 92,928 FMAs
- Arithmetic intensity: 92,928 / 4,250 ≈ 21.9 FLOPs/byte
- Now this is close to the machine balance! So with 8 output channels, we might be compute-bound rather than memory-bound.
Wait, actually. If 8 warps (one per output channel) run sequentially or concurrently on the same SM, and they all read the same input, then input is cached in L1. The first warp reads 17KB from global memory, and subsequent warps hit L1 cache. So total global memory traffic is still ~17KB for input, but we compute 8x more FLOPs.
This is the key insight! By having multiple output channels per block, we amortize the global memory reads across more computation. With 8 output channels, arith intensity jumps from 2.7 to 21.6, making us compute-bound or at least much less memory-bound.
Can we go higher? 16 output channels per block:
- Shared memory weights: 16 * 363 * 4 = 23.2 KB
- Computation: 16 * 32 * 363 = 185,856 FMAs
- Intensity: 185,856 / 4,250 ≈ 43.7 FLOPs/byte
- This is well above machine balance, so we'd be compute-bound.
But 16 output channels means 256 / 16 = 16 threads per output channel. That's only half a warp. If the 16 threads don't perfectly span a warp, we might have some inefficiency.
Wait, a warp is 32 threads. If we have 16 threads per output channel, then each warp handles 2 output channels. So within one warp, threads 0-15 work on output channel 0, and threads 16-31 work on output channel 1.
For threads 0-15:
- They access the same weight (broadcast from shared memory: good)
- They access different input positions: they're computing adjacent spatial positions, so input access is coalesced
- 16 threads produce 2 128-byte transactions (one for threads 0-7, one for threads 8-15)? Actually 16 threads accessing 16 floats spaced 16 bytes apart covers 16*16 = 256 bytes, which is 2 128-byte transactions. This is still efficient.
For threads 16-31 (second output channel):
- Same pattern but different weights
So this should work fine. But is there a downside? The block now computes 16 * 32 = 512 output elements, up from 256. With the same number of total output elements (74.3M), we need half as many blocks. But 74.3M / 512 ≈ 145K blocks, which is plenty for an H100 with 132 SMs.
Let me try with 8 output channels first since it gives a full warp per channel, which is simpler.
Wait, I actually want to reconsider something. The actual machine balance for fp32 FMA is 2 FLOPs per instruction, or 2 operations per cycle per ALU. Let me think about this properly.
Memory bandwidth: 3.35 TB/s
FP32 FLOPS: 67 TFLOPS
Compute-to-memory ratio: 67 / 3.35 ≈ 20 FLOPs/byte
For our kernel with 8 output channels:
- Global memory: ~17KB input + (if we don't count output write) + ~0 weights (shared mem)
- Output write: 8 * 32 * 4 = 1024 bytes
- Global memory traffic per block: ~18KB
- FLOPs: 8 * 32 * 363 * 2 = 185,856 (each FMA = 2 FLOPs)
- FLOPs/byte: 185,856 / 18,000 ≈ 10.3
- This is about half the machine balance, so we're still somewhat memory bound, but much better than the 1.3 FLOPs/byte for 1 output channel.
Actually, I realize I miscounted. With 8 output channels per block:
- 8 warps, each doing 32 * 363 = 11,616 FMAs
- Total FMAs per block: 92,928
- Total FLOPs: 185,856
- Per warp output: 32 elements = 128 bytes
- Total output: 8 * 128 = 1024 bytes
- Input: first warp reads ~17KB, but subsequent warps hit L1 (assuming L1 is large enough and fully hot)
- L1 cache on Hopper: the L1 + shared is 228KB per SM or so. The actual L1 might be smaller.
Actually, on Hopper, L1 cache and shared memory share the same 228KB per SM. If we use 23KB for shared memory (16 channels), that leaves ~200KB for L1. Our input footprint of 17KB fits easily. So subsequent warps should hit L1.
But there's a subtlety: if multiple warps try to access the same L1 cache lines simultaneously, there might be some port contention, but L1 is designed for this.
So with 8 output channels:
- First warp: cold L1, reads ~17KB from global
- Warps 2-8: warm L1, read ~17KB each from L1
- Total global memory: ~17KB input + 1KB output = ~18KB
- Total FLOPs: 185,856
- FLOPs/byte (global): ~10.3
We're still somewhat memory-bound but significantly better.
Can we do 32 output channels per block?
- 256 threads / 32 = 8 threads per output channel
- 8 threads in a warp handling different output channels... wait no.
- If 8 threads per output channel, we could assign 4 output channels per warp (32/8=4).
- But then the 8 threads accessing input would be for different spatial positions, which is good.
- However, threads in the same warp handling DIFFERENT output channels means they access different weights.
- In shared memory, if 4 threads access 4 different weights, that's a 4-way bank conflict if they straddle banks. With 32 banks and float access, bank = address / 4 % 32. If weights for different output channels are stored consecutively, and each channel's weights are 363 floats (not multiple of 32), then the starting bank rotates.
Let me check: channel 0 weights start at bank 0. channel 1 starts at bank 363 % 32 = 11. channel 2 starts at bank 22. channel 3 starts at bank 1. These are all different banks, so no bank conflicts for the first access. But within the 363 floats, accesses to the same offset within different channels will hit different banks since 363 mod 32 = 11, so they cycle through.
Wait, if threads in a warp access weights[oc * 363 + same_offset] for oc = 0,1,2,3:
- Banks: (oc * 363 + offset) % 32
- 363 % 32 = 11
- So banks are offset, offset+11, offset+22, offset+1 (mod 32)
- These are all different, so NO bank conflicts!
This is great! So we could have 32 output channels per block:
- 256 threads / 32 = 8 threads per output channel
- Each warp handles 4 output channels, 8 threads each
- 8 warps total
But wait, 32 output channels means weight shared memory = 32 * 363 * 4 = 46.4 KB. Still fits comfortably.
Total output channels = 96, so we'd have 96/32 = 3 blocks in x-dimension.
Total spatial positions per channel = 774,400
With 8 threads per channel, we process 8 spatial positions per block per channel.
Total blocks = 3 * (774,400 / 8) = 3 * 96,800 = 290,400 blocks.
Wait, block computes 32 channels × 8 positions = 256 output elements, same as before. So total blocks is the same.
But with 32 channels per block:
- We compute 32 * 8 * 363 = 92,928 FMAs per block (same as 8 channels * 32 positions... wait, it's the same!)
- Oh right, total output elements per block is always 256. So total FLOPs per block doesn't change.
- What changes is the input footprint: 8 positions per channel vs 32 positions.
- For 8 positions in a row: width = 8*4 + 10 = 42 pixels. 42 * 11 * 3 ≈ 1,386 floats ≈ 5.5 KB
- So global input per block is ~5.5KB (first warp loads, rest hit L1)
- FLOPs/byte increases to 185,856 / 6,500 ≈ 28.6, exceeding machine balance!
- So we'd be compute-bound!
But is there a downside? Only 8 threads per output channel means we have less parallel coalescing. 8 threads still access memory with 16-byte spacing, producing one or two 128-byte transactions. This is still efficient.
However, one practical issue: with only 3 blocks in the x-dimension (for output channels), and 96,800 blocks in y, the total parallelism might be fine, but the block distribution might not perfectly fill all SMs. But with 290K blocks, this should be okay.
Actually, I realize there's a problem. If we have only 8 threads computing spatial positions for one output channel, and positions are consecutive, then we cover 8 positions. But with total spatial positions per channel = 774,400, the number of spatial blocks is 96,800 per channel group. With 3 channel groups, total blocks = 290,400.
Actually, let me be more careful. Total blocks = (output_channel_groups) × (spatial_groups).
- For 32 channels per block: 96/32 = 3 channel groups
- 774,400 spatial pos / 8 pos per block = 96,800 spatial groups
- Total blocks: 3 * 96,800 = 290,400
For 8 channels per block:
- 96/8 = 12 channel groups
- 774,400 / 32 = 24,200 spatial groups
- Total blocks: 12 * 24,200 = 290,400
Same total blocks! The difference is in the dimensions of the grid. With 32 channels per block, grid.x = 3, grid.y = 96,800. With 8 channels, grid.x = 12, grid.y = 24,200.
A "tall and skinny" grid (large y, small x) might be fine. CUDA handles large grids well.
The key difference is the input footprint per block. Smaller footprint means better L1 cache behavior. With 8 positions in a row, footprint is 5.5KB. With 32 positions, it's 17KB. Both fit in L1, but 5.5KB means more cache lines stay resident. Also, with 32 channels per block, we have more warps (up to 32 channels would mean... actually, for 32 channels with 256 threads, if we use 8 threads per channel, that's 32 warps of 8 threads each? No, 256 threads / 32 = 8 per channel. A warp has 32 threads, so each warp handles 4 channels (4 * 8 = 32 threads). There are 8 warps total.
Wait, I need to think about warp divergence. If a warp has 8 threads for channel 0, 8 for channel 1, etc., and all execute the same instruction (loop over ic, ky, kx), there's no divergence. Good.
At each (ic, ky, kx) step:
- 32 threads in a warp each load one input value and one weight value
- Input loads: threads in the same channel group access offsets 0, 4, 8, 12, 16, 20, 24, 28 within the row. These are 16 bytes apart, so threads 0 and 1 at offset 0 and 4 (64 bytes apart) would be in the same 128-byte line. But actually float spacing is 4 bytes per float, and threads access 16 bytes apart = 4 floats apart. So thread 0 at addr, thread 1 at addr+16, thread 2 at addr+32, ..., thread 7 at addr+112. This spans 128 bytes, which is one transaction. Actually, if the base address is 128-byte aligned, threads 0-7 fit exactly in one 128-byte transaction. But threads 8-15 (next channel's spatial positions) access a completely different memory region (different channel's weights determine different... no, actually they access the same input region but different weights.
Wait, all threads in the block work on the same spatial positions (or overlapping spatial tiles) for different output channels. If spatial positions are the same for all 8 threads in a channel group, but different channel groups also work on the same spatial positions, then all 32 threads in a warp access the same 8 input positions? No, that would mean each channel group of 8 threads all computes the same spatial positions, which would mean the 8 threads in a group are redundant. That's not what I said.
I said 8 threads per output channel, each computing different spatial positions. So channel 0's 8 threads compute spatial positions 0-7, channel 1's 8 threads compute spatial positions 0-7 (same positions but different output channel), etc.
In a warp (32 threads):
- Threads 0-7: channel 0, spatial positions 0-7
- Threads 8-15: channel 1, spatial positions 0-7
- Threads 16-23: channel 2, spatial positions 0-7
- Threads 24-31: channel 3, spatial positions 0-7
For input access at a given (ic, ky, kx):
- Threads 0-7 access input for spatial positions 0-7. Since positions are adjacent in x, they access addresses: base + 0*16, base + 1*16, ..., base + 7*16 bytes. This spans 128 bytes = 1 transaction.
- Threads 8-15 access the exact same input addresses as threads 0-7! Because they're computing the same spatial positions, just for a different output channel.
- This means 16 threads are reading from the same 128 bytes. But each thread in the warp gets the full cache line broadcast. Actually, in a single warp, all threads requesting the same cache line means only 1 DRAM/L1 transaction serves all of them. This is EXCELLENT!
And threads 16-23 and 24-31 do the same. So for the 4 channel groups in a warp, they all access the same 128 bytes of input. The warp only needs 1 memory transaction for all 32 threads!
Then the next warp in the block (if any) handles channels 4-7, also computing the same spatial positions 0-7. When this second warp runs, it also needs the same 128 bytes of input, which are now in L1 cache.
So global memory traffic:
- For each spatial tile of 8 positions: ~1.5KB input (loaded once into L1 by the first warp)
- Plus output writes: 32 channels * 8 positions * 4 bytes = 1KB
- FLOPs: 32 * 8 * 363 * 2 = 185,856 FLOPs
- Ratio: 185,856 / 2,500 ≈ 74 FLOPs/byte!
This is way above machine balance. We'd be fully compute-bound.
Wait, but actually the input footprint for 8 positions in a row is:
- Width covered: 8*4 + 10 = 42 pixels
- For 3 channels and 11 ky: 42 * 11 * 3 = 1,386 floats ≈ 5.5KB
- But all warps need this same 5.5KB, and it stays in L1.
- Global input traffic: 5.5KB
- Output: 1KB
- Total: 6.5KB
- Ratio: 185,856 / 6,500 ≈ 28.6 FLOPs/byte
Still above machine balance of 20. So compute-bound!
This is a very efficient design. But we need to make sure the 5.5KB fits in L1. On H100, L1 is part of the 228KB shared/L1 partition. If we use ~47KB for shared memory (weights for 32 channels), that leaves ~180KB for L1. 5.5KB easily fits.
Wait, one issue: with 32 channels per block, we have 8 warps, each computing the same spatial positions. These 8 warps might run concurrently on an SM (H100 can handle many warps). They all compete for the same L1 cache lines, but since the footprint is small, it should be fine.
Actually, there's a subtle point: if 8 warps access the same L1 cache lines simultaneously, there's bank/port contention in L1. But L1 is designed for concurrent access from multiple warps.
Let me also verify the output channel count. We have 96 output channels. 96 / 32 = 3, which divides evenly. Great!
Now, what about the grid structure with 32 channels per block?
- Grid.x = 3 (output channel groups)
- Grid.y = 96,800 (spatial groups, where each group is 8 positions)
Actually, wait. Each block computes 8 positions for 32 channels. But in a block, the warps all compute the same 8 positions (for different channels). What about the next block in the spatial dimension? It computes the next 8 positions for the same 32 channels.
So grid.y should be: total_spatial_positions / 8 = 774,400 / 8 = 96,800.
But 774,400 / 8 = 96,800 exactly? Let me check: 774,400 / 8 = 96,800. Yes, it divides evenly!
Wait, 774,400 = 256 * 55 * 55 = 256 * 3025. Is 3025 divisible by 8? 3025 / 8 = 378.125. No! 3025 = 55*55 = 5*11*5*11 = 25*121 = 3025. Not divisible by 8.
So 256 * 3025 = 774,400. 774,400 / 8 = 96,800. But the issue is that 3025 is not divisible by 8. However, 774,400 is divisible by 8. Let me verify: 774,400 / 8 = 96,800 exactly. Yes.
So we can just divide the flattened spatial index by 8 without any remainder issues at the batch×H×W level. But when we decode spatial_idx to (n, oy, ox), the tiling doesn't naturally align to rows.
For example, spatial_idx = 0..7 might span the end of one row and the beginning of the next. With row size 55, positions 0..54 are row 0, 55..109 are row 1, etc.
A block processing 8 positions starting at spatial_idx = 48 would process positions 48-55. Position 48 is (n=0, oy=0, ox=48), 49=(0,0,49), ..., 54=(0,0,54), 55=(0,1,0).
For positions 48-54 (7 positions in row 0): input x ranges from 48*4-2=190 to 54*4+8=224. Position 55 (row 1, col 0): input x ranges from 0-2+0=-2 to -2+10=8, and input y = 4-2=2 to 4+8=12.
So the memory footprint for this block spans two disjoint regions. The threads that hit row 1 might cause cache pollution. But this is still fine since the total footprint for 8 positions is small.
Actually, it might be better to tile in 2D to ensure spatial locality. Let me reconsider.
For a 2D tile, say 4x2 = 8 positions:
- Tile (oy=0..3, ox=0..1): input y ranges from -2 to 3*4+8=20, input x from -2 to 1*4+8=12
- Footprint: roughly 23 * 15 * 3 * 11 with overlap... actually more like (3*4+10) * (1*4+10) * 3 = 22 * 14 * 3 = 924 floats ≈ 3.7KB.
This is even smaller than the 1D case! And spatial locality is better.
But a 4x2 tile means the 8 threads in a channel group compute a non-contiguous set in memory (since output is stored as [N, C, H, W], and the tile spans multiple W values and H values). Writing output would require strided accesses.
For output writes, if we have a 4x2 tile:
- Thread 0 writes (oy=0, ox=0): offset n*C*H*W + ch*H*W + 0*W + 0
- Thread 1 writes (oy=0, ox=1): same + 1
- Thread 2 writes (oy=1, ox=0): same + W
- Thread 3 writes (oy=1, ox=1): same + W + 1
- Thread 4 writes (oy=2, ox=0): same + 2*W
- etc.
These writes are not all contiguous. Within a warp, some threads write adjacent memory but others write 55 floats away. This means multiple cache lines are touched, possibly causing divergent stores.
If we use a row-major 1D tile of 8 positions, the writes are contiguous (for positions within the same row). If the tile straddles a row boundary, the last few positions write to the start of the next row.
But stores are less critical than loads for performance. Let me stick with the 1D tile.
Actually, I realize with 8 positions, the probability of crossing a row boundary is: each row has 55 positions, so blocks of 8 cross boundaries at positions 48, 56, 64, ... wait, no. A block starting at 48 covers 48-55, which crosses the boundary at 55. Starting at 56 covers 56-63 which is entirely in row 1.
Wait, 55 / 8 = 6.875. So blocks 0-6 are within row 0 (positions 0-55). Block 6 starts at 48, covers 48-55, crossing into row 1. Block 7 starts at 56, covers 56-63, all in row 1.
So out of ~13,672 blocks per batch element (3025/8 ≈ 378.125... actually ceil(3025/8) = 379 blocks per image, and 379 * 256 = 97,024 blocks total. Wait, but I said 96,800 earlier.
Let me recalculate: 774,400 / 8 = 96,800 exactly. So no partial blocks at the global level. But at the per-image level, each image has 3025 positions, so 3025/8 = 378.125. If we use simple 1D blocking, we'd need 379 blocks per image, where the last block has some invalid positions.
Actually, I was thinking of grid.y = 774,400/8 = 96,800, which requires that we process exactly the right amount. But if we just let grid.y = 96,800 and each block handles spatial_start = blockIdx.y * 8, the last block for each image would have spatial positions 3024 to 3031, of which only 3024 is valid.
So I need to add bounds checking. This is fine.
But there's another issue: with 96,800 blocks per channel group and 3 channel groups, total is 290,400 blocks. But each block only computes 8 positions for 32 output channels. The grid organization matters.
Actually, if each block handles 8 positions and 32 channels, and grid is (3, 96,800), then:
- blockIdx.x in 0..2 selects the output channel group
- blockIdx.y in 0..96799 selects the spatial block
For decoding spatial position:
- total_spatial = batch * out_h * out_w = 256 * 55 * 55 = 774,400
- spatial_start = blockIdx.y * 8
- position i in thread: spatial_start + i, where i = threadIdx.x % 8
For the last few blocks per image, some positions exceed 3024. We need to handle this with bounds checking.
Wait, I need to also handle batch correctly. If we treat spatial_idx as n * 3025 + oy * 55 + ox, then:
- n = spatial_idx / 3025
- pos_in_img = spatial_idx % 3025
- oy = pos_in_img / 55
- ox = pos_in_img % 55
For spatial_start = blockIdx.y * 8, and blockIdx.y ranges from 0 to 96,799:
The mapping to (n, oy, ox) is straightforward.
For thread 7 in a block starting at spatial_start=3024: spatial_idx = 3024 + 7 = 3031.
n = 3031 / 3025 = 0 (first image, since 3031 < 3025? Wait, 3031 > 3025, so n = 1, pos_in_img = 6)
Actually: 3024 + 7 = 3031. 3031 / 3025 = 0 with remainder 3031? No, 3031 >= 3025, so n = 1, pos_in_img = 6.
Hmm wait, 3024 is the last position in image 0 (since positions 0..3024 are 3025 positions). So thread 0 computes image 0, position 3024 (oy=54, ox=54 - wait, 54*55+54 = 2970+54 = 3024). Yes.
Thread 1 computes image 0, position 3025 = image 1, position 0 (oy=0, ox=0).
Oh! So a single block can span batch elements. The last thread in one image computes the first pixel of the next image!
Is this a problem? For memory access: image 0, position 3024 (oy=54, ox=54) and image 1, position 0 (oy=0, ox=0) are very far apart in memory (about 150KB apart). The 8 threads in a warp accessing these memory locations would generate two very different memory regions. This is bad for coalescing.
To avoid this, we should ensure blocks don't span batch boundaries. We should handle batch as a separate dimension or align spatial blocks to image boundaries.
Let me redesign: process each image independently.
- For each image, we have 3025 spatial positions.
- We want tiles of 8 positions: ceil(3025 / 8) = 379 tiles per image.
- Total spatial tiles: 256 * 379 = 97,024.
Alternatively, we could use a 2D grid:
- gridDim.z = batch = 256
- gridDim.y = ceil(out_h * out_w / tile_spatial) = 379
- gridDim.x = out_channel_groups = 3
This ensures all threads in a block work on the same image.
Total blocks = 256 * 379 * 3 = 291,072.
Let me verify: 379 * 8 = 3,032 positions per image, but only 3,025 are valid. So the last tile has 7 valid positions. We need bounds checking, but spatial locality within a tile is preserved.
In the last tile of each image, which covers positions 3016 to 3023 (8 positions, all valid? 3016+7=3023, and 3024 is also valid... wait, 3025 total positions, so 0..3024. 3016..3023 covers 8 positions. The next would be 3024..3031, where only 3024 is valid.)
Actually: 3025 / 8 = 378 remainder 1. So tiles 0..377 cover 378*8 = 3024 positions (0..3023). Tile 378 covers positions 3024..3031, where only 3024 is valid.
So there are 379 tiles per image, with tile 378 having 7 threads doing useless work. That's fine.
With this 3D grid:
- blockIdx.x = 0..2 (output channel group)
- blockIdx.y = 0..378 (spatial tile)
- blockIdx.z = 0..255 (batch)
Each block: 256 threads, computing 32 output channels × 8 spatial positions.
Thread mapping:
- threadIdx.x: 0..255
- out_ch_local = threadIdx.x / 8 = 0..31
- spatial_local = threadIdx.x % 8 = 0..7
Wait, if we use 32 channels and 8 positions, and 256 threads:
- thread 0: ch=0, pos=0
- thread 1: ch=0, pos=1
- ...
- thread 7: ch=0, pos=7
- thread 8: ch=1, pos=0
- ...
- thread 255: ch=31, pos=7
This packs the spatial positions for each channel consecutively in thread index. But for coalescing, adjacent threads should access adjacent memory. Here, threads 0-7 all access the same channel but adjacent spatial positions - good for input coalescing but they write to different output positions (adjacent in x, or at least contiguous in flattened index if in same row).
For output writes: threads 0-7 write output channel 0, positions 0-7. Since output is [N, C, H, W], these are adjacent in memory! Excellent.
Threads 8-15 write output channel 1, positions 0-7. These are also adjacent in memory but offset by H*W from threads 0-7.
Within a warp (threads 0-31):
- Threads 0-7: channel 0, positions 0-7 -> writes to addr+0 through addr+7
- Threads 8-15: channel 1, positions 0-7 -> writes to addr+H*W through addr+H*W+7
- These are not contiguous in memory. They differ by 3025 floats = 12,100 bytes.
- A warp writing these would generate two separate 32-byte transactions (for threads 0-7) and two separate 32-byte transactions (for threads 8-15), etc.
- Actually, with 32 threads writing to 4 different regions (channels 0-3), each spaced 12KB apart, that's 4 different cache lines. The store unit can handle this but it's not as efficient as a single contiguous write.
An alternative mapping: interleave channels.
- thread 0: ch=0, pos=0
- thread 1: ch=1, pos=0
- thread 2: ch=2, pos=0
- ...
- thread 31: ch=31, pos=0
Then threads 32-63: ch=0..31, pos=1
And so on.
With this mapping, in a warp (32 threads):
- Thread i computes channel i, position 0
- All threads write the SAME spatial position but DIFFERENT channels
- Output layout: [N, C, H, W]. For fixed N, H, W, adjacent channels are contiguous!
- So threads 0-31 write to addresses: base + 0*H*W, base + 1*H*W, ..., base + 31*H*W
- These are spaced by 3025 floats = 12,100 bytes
- Not contiguous! Same issue.
Hmm, no matter how we map, writing 32 different channels to the same spatial position means strided writes. And writing 8 different positions in the same channel means contiguous writes.
Which is better? In our kernel, we care more about input reads than output writes (since we read much more input than we write output). Input reads are better when threads in a warp read adjacent spatial positions, which happens when we assign consecutive threads to adjacent positions in the same channel.
For output, we can use the same mapping as input. The writes won't be perfectly coalesced across channels, but stores are generally less of a bottleneck.
Actually, I realize: if we process all 32 channels for the same 8 positions, we write 32*8=256 output floats. In the first mapping (threads 0-7: ch0 pos0-7, threads 8-15: ch1 pos0-7, etc.), a warp writes:
- Channels 0-3, positions 0-7 = 32 floats total
- These occupy 4 separate regions, each 8 floats contiguous.
- The memory system handles this with 4 separate 32-byte stores.
In the second mapping (threads 0-31: ch0-31 pos0, threads 32-63: ch0-31 pos1, etc.), a warp writes:
- All 32 channels for position 0 = 32 floats
- These are spaced by H*W = 3025 floats = 12,100 bytes.
- Also 4 separate 32-byte stores? Actually each float is 4 bytes, so 32 floats = 128 bytes, but spaced far apart. This might be 32 separate 4-byte stores (uncoalesced!).
So the first mapping is much better for output writes.
Let me verify the first mapping output write pattern more carefully:
- Warp 0 (threads 0-31): ch0 pos0-7, ch1 pos0-7, ch2 pos0-7, ch3 pos0-7
- Threads 0-7 write contiguous 8 floats at offset ch0*H*W + pos
- Threads 8-15 write contiguous 8 floats at offset ch1*H*W + pos
- etc.
Each group of 8 threads generates one 32-byte store transaction (8 floats = 32 bytes). The 4 groups in a warp generate 4 transactions. Since they're independent, the memory system can handle them in parallel or pipeline them. This is fine.
And for the second warp (threads 32-63): ch4 pos0-7, ch5 pos0-7, ch6 pos0-7, ch7 pos0-7. Same pattern.
So the first mapping is good. Let me stick with it.
Now, for input reads in the first mapping:
- Warp 0 (threads 0-31): ch0 pos0-7, ch1 pos0-7, ch2 pos0-7, ch3 pos0-7
- At (ic=0, ky=0, kx=0):
- Threads 0-7 (ch0, pos0-7): all read the same input positions. Thread i reads input for pos i.
- Threads 8-15 (ch1, pos0-7): read the exact same input positions as threads 0-7!
- So all 32 threads in the warp access only 8 unique input values.
- These 8 values are at addresses: base + 0*4, base + 1*4, ..., base + 7*4 (where 4 is the stride in floats).
- Wait, for adjacent output positions (ox differs by 1), input_x differs by stride=4. So the addresses are base, base+4, base+8, ..., base+28 floats.
- In bytes: base*4, base*4+16, base*4+32, ..., base*4+112.
- These span 128 bytes, perfectly aligned for one 128-byte warp transaction!
- And since all 32 threads access these same 8 values, the hardware broadcasts them to all threads.
This is incredibly efficient! One 128-byte transaction serves all 32 threads in the warp.
And this happens for every (ic, ky, kx) step. So each warp needs only one transaction per (ic, ky, kx) to load all input it needs. The total global memory transactions per warp = 3 * 11 * 11 = 363 transactions.
Wait, but the 8 unique values are loaded once and broadcast to all 32 threads. But actually, the L1 cache handles this. The warp issues loads, and the L1 cache services them. With all threads accessing the same cache line, the hardware is smart enough to only fetch it once.
So global memory traffic per warp is just the 128 bytes per step (well, actually 8 floats = 32 bytes per step, since we only need 8 unique values). But memory transactions are at least 32 bytes. So it's 32 bytes per step.
Wait, 8 floats = 32 bytes. A warp transaction can be 32 bytes minimum. But actually, modern GPUs do 128-byte transactions per warp. Even if only 32 bytes are needed, the full 128-byte line is loaded into L1.
Hmm, but if 8 floats are exactly 32 bytes, and they're contiguous with 4-float spacing? Wait, adjacent threads access addresses 16 bytes apart. So the 8 threads access 8 addresses: A, A+16, A+32, ..., A+112. These cover the range [A, A+115], which is 116 bytes, fitting in a 128-byte line. But the accessed bytes are not contiguous! They're strided within the line.
For a 128-byte line starting at A_aligned (where A_aligned <= A < A_aligned + 128):
- If A is 128-byte aligned, then accesses are at offsets 0, 16, 32, 48, 64, 80, 96, 112 in the line.
- The 128-byte line contains all these.
- A warp load transaction brings the whole 128-byte line into L1.
- So we effectively read 128 bytes to get 32 bytes of useful data.
This is still good! 25% utilization of the loaded cache line, but it's the minimum transaction size for a warp, and we can't do better with this strided access pattern. Actually, we could try to pack 16 positions instead of 8 to get 64 bytes of useful data per 128-byte transaction, but that reduces our compute-to-memory ratio advantage.
Wait, let me reconsider. If we process 16 positions per channel group:
- 16 threads per channel, accessing addresses A, A+16, ..., A+240.
- This spans 256 bytes, requiring two 128-byte transactions.
- Useful data: 64 bytes.
- Utilization: 25%, same as before.
- But now each warp has 16 threads per channel, so 2 channels per warp.
- Total channels per block: 256/16 = 16 channels.
With 16 channels per block:
- Shared memory: 16 * 363 * 4 = 23,232 bytes ≈ 23KB
- Still fine.
- Channels per warp: 32/16 = 2
- Warps per block: 256/32 = 8
But 16 positions per block means fewer blocks: 774,400 / 16 = 48,400 spatial blocks per channel group. With 96/16 = 6 channel groups: total blocks = 6 * 48,400 * 256 = wait, let me use the 3D grid.
With 16 channels per block:
- grid.x = 96 / 16 = 6
- tiles per image = ceil(3025 / 16) = 189
- grid.y = 189
- grid.z = 256
Total blocks: 6 * 189 * 256 = 290,304.
Per block, first warp loads input for channels 0-1, positions 0-15.
- 16 positions accessed by 16 threads.
- Span: A to A+240 bytes = 256 bytes, but we're accessing every 16 bytes.
- Two 128-byte transactions needed.
Hmm, same utilization but now 16 positions gives more compute per input. However, we also have more output channels sharing the same input (16 channels vs 32 channels). With 16 channels, only 2 channel groups in a warp access the same input, vs 4 with 32 channels. But all warps in the block still share the cached input via L1.
Actually, the key advantage of more channels per block is more compute per block, but the total compute per output element doesn't change. What matters is global memory traffic.
With 32 channels, first warp loads 128 bytes for 8 positions. Remaining 7 warps hit L1. Total global input per block: ~128 bytes per (ic, ky, kx) step? Wait, no - per step we need different input based on ky, kx.
For 8 positions in a row:
- At (ic, ky, kx), we read the same row of input for all ky? No, ky changes the input row.
- For ky=0, kx=0: input row = oy*4 - 2 + 0, columns = ox*4 - 2 + 0 for ox = pos0..pos7.
- For ky=0, kx=1: same row, columns shifted by 1.
- For ky=1, kx=0: row shifted by 1, same columns.
So for all 3 * 11 * 11 = 363 steps:
- Each step reads 128 bytes (one cache line) for the 8 positions.
- Total global input traffic for the block (assuming perfect L1 caching across warps): 363 * 128 bytes = 46,464 bytes ≈ 45.4 KB.
- Output writes: 32 channels * 8 positions * 4 bytes = 1,024 bytes.
- Total FLOPs: 32 * 8 * 363 * 2 = 185,856.
- FLOPs / global byte: 185,856 / 46,464 ≈ 4.0.
Hmm, that's much lower than I thought! The reason is that each step loads a full 128-byte cache line but only uses 8 floats (32 bytes) from it. The other 96 bytes are wasted.
Wait, I need to reconsider. If we have 8 positions with stride 4 in x, and we're reading input at (iy, ix) for varying ky, kx within the kernel:
At (ky, kx), the input x coordinates are: pos0_x*4 - 2 + kx, pos1_x*4 - 2 + kx, ..., pos7_x*4 - 2 + kx.
If pos0_x through pos7_x are consecutive output x values, then input x values are: base, base+4, base+8, ..., base+28.
These span 28 floats = 112 bytes. A cache line is 128 bytes. So we need one 128-byte line per step.
For 363 steps, that's 363 cache lines. But many cache lines overlap! Because as kx varies from 0 to 10, the columns shift by 1 each time. So kx=0 might load columns [base, base+4, ..., base+28] which are in cache line starting at base. kx=1 needs [base+1, base+5, ..., base+29]. Some of these might be in the same cache line, some in the next.
For example, base=0 (aligned):
- kx=0: accesses offsets 0,4,8,12,16,20,24,28. All in line [0,127] bytes.
- kx=1: accesses offsets 1,5,9,13,17,21,25,29. All in line [0,127] bytes.
- ...
- kx=0 to kx=7: all accesses fit in line [0,127] (max offset = 28+7=35 floats = 140 bytes? Wait, 35 floats = 140 bytes. So kx=7 accesses up to offset 35, which is in [0,127]).
- kx=8: accesses up to offset 36 = 144 bytes, which exceeds 128 bytes. Some accesses in line [0,127], some in [128,255].
- kx=9: accesses up to offset 37.
- kx=10: accesses up to offset 38.
Hmm wait, base=0 means first access is at float 0. For kx=10, accesses are at 10, 14, 18, 22, 26, 30, 34, 38. Offsets 38*4=152 bytes, which is in line [128,255]. So only kx=8,9,10 need a second cache line.
Actually, I should think about this differently. For a fixed (ic, ky), as kx varies from 0 to 10, the set of 8 accessed columns spans from base+kx to base+kx+28. The total span is [base, base+38] = 39 floats = 156 bytes. This requires at most 2 cache lines.
For 11 kx values, we might need ~11 cache lines (but with overlap). Actually, consecutive kx values shift by 1 float = 4 bytes, so consecutive accesses are mostly in the same cache line, with a line boundary crossed at most once.
Similarly, for ky varying from 0 to 10, each ky shifts the input row by 1. For each (ic, ky), we need roughly 11 cache lines for the kx sweep, but consecutive ky values touch different rows (each row is 224*4=896 bytes apart). So for different ky, there's no overlap in cache lines. 11 ky values = 11 rows, each needing ~11 cache lines = ~121 cache lines. For 3 ic = 363 cache lines.
Each cache line = 128 bytes. Total input traffic = 363 * 128 = 46,464 bytes per block.
This is the first warp's traffic. Subsequent warps should hit L1 cache, so total global traffic is ~46KB per block.
185,856 FLOPs / 46,464 bytes = 4.0 FLOPs/byte. This is well below the 20 FLOPs/byte machine balance. So we're STILL memory bound, even with 32 channels sharing the same input!
Wait, that can't be right. Let me recalculate.
185,856 FLOPs per block, 46,464 bytes global input per block. But 46,464 bytes is a LOT for just 8 positions. Let's think about what a direct convolution actually needs.
For 8 output positions in a row:
- Input region: width = 8*4 + 10 = 42, height = 11 (for one ky... no, for all ky it's 11)
- Total input floats: 42 * 11 * 3 = 1,386 floats = 5,544 bytes.
But we argued that due to cache line granularity, we read 363 * 128 = 46,464 bytes. That's 8.4x the useful data! Because of the strided access pattern and cache line size.
Can we improve this? What if we use float4 or some vector type to load multiple contiguous values at once? But our access pattern is strided (every 4th float), so we can't use vector loads directly.
What if we reorganize so threads load contiguous memory and then shuffle? This is getting complicated.
Alternatively, what if we process 32 positions in a row instead of 8?
- 32 positions: span = 32*4 + 10 = 138 floats.
- For a fixed (ic, ky), kx=0 accesses 32 columns at offsets 0,4,8,...,124. Span = 124 floats = 496 bytes, requiring 4 cache lines.
- For 11 kx values: total span = [base, base+134] = 135 floats. That's 5 cache lines per row if aligned, maybe 6.
- For 11 ky × 3 ic: 33 rows × 6 cache lines = 198 cache lines.
- Total input bytes = 198 * 128 = 25,344 bytes.
- FLOPs: 32 channels * 32 positions * 363 * 2 = 743,424.
- FLOPs/byte = 743,424 / 25,344 ≈ 29.3. This is above machine balance!
Wait, but 32 positions with 32 channels means only 1 position per channel if we use 256 threads. Let me check: 256 threads / 32 channels = 8 threads per channel. 32 positions / 8 threads = 4 positions per thread. So each thread computes 4 output elements.
If each thread computes 4 positions serially:
- The thread does a loop over its 4 positions, accumulating each independently or reusing registers.
- But this complicates the kernel.
Alternatively, with 256 threads and 32 channels: threads per channel = 8. If we want 32 positions, we'd need 32/8 = 4 blocks in some sense. Or we could increase block size to 1024 threads, but that's 32 channels * 32 positions = 1024 threads. Then each thread computes exactly 1 output element.
Block size 1024:
- 32 channels × 32 positions = 1024 elements.
- Threads per channel = 32. A warp has 32 threads, so 1 warp per channel, 32 warps total.
- All 32 warps compute the same 32 positions but different channels.
- First warp loads input for 32 positions. How many cache lines?
- 32 positions span 32*4+10 = 138 floats = 552 bytes.
- At kx=0: accesses at 0,4,8,...,124 offsets = span 124 floats = 496 bytes. With 128-byte cache lines, that's ceil(496/128) = 4 lines? Wait, the accesses are at every 16 bytes within 496 bytes. They touch 4 cache lines.
- For 11 kx values: total span 135 floats = 540 bytes = 5 cache lines (if aligned well, maybe 5).
- For 11 ky × 3 ic = 33 rows, each needing ~5 lines: 165 lines.
- Total: 165 * 128 = 21,120 bytes.
- FLOPs: 32 * 32 * 363 * 2 = 743,424.
- FLOPs/byte: 743,424 / 21,120 ≈ 35.2. Compute bound!
But 1024 threads per block means fewer total blocks: 96/32 = 3 channel groups, 774,400/32 = 24,200 spatial groups. Total blocks: 3 * 24,200 * 256 (batch dimension if separate, but with 3D grid it's 3 * ceil(3025/32) * 256 = 3 * 95 * 256 = 72,960 blocks).
Wait, if we use 3D grid: grid.x=3, grid.y=ceil(3025/32)=95, grid.z=256.
Total blocks = 3 * 95 * 256 = 72,960.
That's a bit low. 72,960 blocks for an H100 with 132 SMs. At full occupancy of say 16 blocks per SM, that's only 2,112 blocks active at once. We need to queue up blocks, but 72,960 blocks should be enough.
Actually, each block has 1024 threads, which is the maximum per block. With 1024 threads and 32 warps, and H100 having up to 1536 threads per SM (?), we can only run 1 block per SM at a time (since 1024 < 1536 but 2*1024 > 1536). So we'd have 132 blocks active at once. 72,960 / 132 ≈ 552 waves. This should be fine, but we might not fully utilize the SM if we have only 1 block resident.
H100 has more than 1536 threads per SM. Let me check: Actually, the maximum threads per SM on Hopper is 2048. With 1024 threads per block, we can have 2 blocks per SM. So 264 blocks active at once. 72,960 / 264 ≈ 276 waves.
Wait, but with 1024 threads, each thread only computes 1 output element (32 channels * 32 positions = 1024 output elements). The total FLOPs per thread is only 363 * 2 = 726. That's not much register pressure.
Actually, let me think about this more carefully. On H100, a block of 1024 threads with 32 warps will run 32 warps concurrently on the SM (or as many as the warp scheduler can handle). The warps from the same block might run somewhat serially depending on resource availability.
For the first warp (channel 0), it loads input from global memory. The next 31 warps for channels 1-31 will read the same input from L1 cache. But the first warp needs to complete its loads and have the data in L1 before the other warps benefit. If warps run concurrently, there might be some contention.
Actually, on modern GPUs, all 32 warps of a block can be resident on an SM simultaneously. The warp schedulers pick ready warps each cycle. Memory loads cause warps to stall, so the scheduler switches to other warps.
For our block with 32 warps:
1. Warp 0 issues loads for (ic=0, ky=0, kx=0). It stalls.
2. Warps 1-31 also issue loads for the same step. Since they access the same cache lines... actually, if warp 0 hasn't fetched the data yet, warps 1-31 will also miss L1 and issue global load requests for the same lines.
3. The memory coalescing unit de-duplicates these requests. So only one request per unique cache line goes to DRAM/L2.
4. Once the data arrives, all 32 warps can proceed.
This is actually fine! The hardware handles this well.
OK, let me revisit whether 1024 threads per block is the right choice. Let me also try 512 threads.
With 512 threads per block:
- Could do 16 channels × 32 positions = 512 elements per block.
- Threads per channel = 32. A warp is 32 threads, so 1 warp per channel, 16 warps total.
- grid.x = 96 / 16 = 6
- grid.y = ceil(3025 / 32) = 95
- grid.z = 256
- Total blocks = 6 * 95 * 256 = 145,920.
- Per block: 16 channels * 32 positions * 363 * 2 = 371,712 FLOPs.
- Input footprint for 32 positions: 32*4+10 = 138 floats wide. Per (ic, ky): ~5 cache lines. For 33 rows: ~165 cache lines = 21,120 bytes.
- FLOPs/byte: 371,712 / 21,120 ≈ 17.6. Slightly below machine balance.
Hmm, 17.6 is close to 20. We might be near the balance point. If actual global traffic is slightly more due to misalignment, we could be more memory bound. But 1024 threads gives 35.2, which is comfortably compute-bound.
However, 1024 threads per block might limit occupancy. With 2048 threads max per SM, we get 2 blocks per SM. If each block uses a lot of registers, we might only fit 1 block.
Let's estimate register usage per thread:
- A few integers for indices
- A few floats for accumulation
- Loop variables
- Probably ~20-30 registers
At 32 registers per thread, 1024 threads use 32,768 registers. Each SM has 64K registers, so we could fit 2 blocks (65,536 registers). If we use more like 64 registers, we can only fit 1 block.
With 1 block per SM and 1024 threads, our parallelism is limited. But with 72,960 blocks total and 132 SMs, we have ~552 blocks per SM in total, which at 1 at a time gives 552 sequential blocks per SM. If each block takes a short time, total runtime is still fast.
Let me also consider the alternative: 256 threads per block gave us 290,400 blocks. With 256 threads, we could fit more blocks per SM (up to 8 if register limits allow). But the memory efficiency was worse.
I think I'll try both approaches: a 256-thread version and a 1024-thread version, and see which is faster.
Actually, before writing the kernel, let me think about another optimization. What if we load input into shared memory as well? Since the input footprint for a tile is small (e.g., ~21KB for 32 positions), we could preload all needed input into shared memory, and then all warps read from shared memory instead of global memory.
For 32 positions in a row:
- Input width: 32*4 + 10 = 138 floats
- Input height: 11 (for all ky)
- For 3 channels: 138 * 11 * 3 = 4,554 floats = 18,216 bytes.
This fits in shared memory! We have 228KB per SM. If we use shared memory for both input and weights:
- Input: 18,216 bytes
- Weights: 32 channels * 363 floats * 4 bytes = 46,464 bytes
- Total: 64,680 bytes ≈ 64KB
This leaves plenty of shared memory. Only 1-2 blocks per SM due to shared memory limits, but that's okay.
With input in shared memory:
- First, load input from global to shared memory. This can be done by all threads collaboratively.
- Then load weights from global to shared memory.
- Then compute using only shared memory.
The global memory reads for input would be more efficient because we can load contiguous regions:
- For each input channel and each ky row, we need ~138 contiguous floats.
- We can load these with vectorized loads (float4).
- Total global input traffic: 4,554 floats * 4 bytes = 18,216 bytes. This is much closer to the actual useful data!
But loading input into shared memory requires careful synchronization and handling of strided memory. However, the input regions we need are contiguous in memory (for each row in each channel). We just need to figure out which rows to load.
Actually, for output positions (oy, ox), the input access is:
iy = oy * 4 - 2 + ky
ix = ox * 4 - 2 + kx
For a tile of output positions spanning 32 in x and 1 in y (since 256 threads / 32 channels = 8 threads per channel, which means only 8 positions per channel if each thread handles 1 position... wait, I'm mixing things up).
Let me be concrete. With 256 threads and 32 channels:
- 8 threads per channel.
- If we want 32 positions, each thread must compute 4 positions.
- Or we use 1024 threads: 32 threads per channel, 32 positions per channel.
With 1024 threads (32 channels, 32 positions), the input tile for one output row is:
- Output positions: (oy, ox_start) through (oy, ox_start+31)
- Input y for ky=0: oy*4 - 2 through oy*4 - 2 + 10 = [oy*4-2, oy*4+8]
- All ky values: [oy*4-2, oy*4+18]
- But wait, we only have 1 output row in the tile (32 consecutive ox for the same oy).
- Input x: [ox_start*4-2, (ox_start+31)*4-2+10] = [ox_start*4-2, ox_start*4+126]
- Width: 129 floats (approximately).
Hmm, 129 floats * 11 rows * 3 channels = 4,257 floats = 17KB. But if we want to cover all output positions in the tile, we need input for various oy values. Wait, with 32 threads per channel and 1 position per thread, all 32 positions are for threads in the SAME warp, which all compute the SAME output channel but DIFFERENT output positions. If the 32 positions are all in the same output row (same oy), then input rows needed are only for the 11 ky values of that single oy.
But if the 32 positions span multiple output rows... Let me see. 32 consecutive flattened positions with out_w=55 could span at most 1 row boundary (since 32 < 55). But most of the time, they might span a boundary.
If we restrict positions to be within the same row (e.g., tile_w = 32, tile_h = 1), then:
- For batch element n and output row oy, positions are (n, oy, ox_start) through (n, oy, ox_start+31), but clipped to out_w=55.
- Input needed: rows [oy*4-2, oy*4+8] for ky=0..10, that's 11 rows.
- And columns [ox_start*4-2, ox_start*4+126]. Width = 129 floats.
- For 3 channels: 129 * 11 * 3 = 4,257 floats ≈ 17KB.
This fits nicely in shared memory. And the input is mostly contiguous in global memory (row by row). We could load it efficiently.
But then our spatial tiles are 32 wide by 1 tall. With out_h=55, out_w=55, we'd have ceil(55/32) = 2 in x and 55 in y = 110 tiles per image. With 256 batch: 28,160 spatial tiles. With 3 channel groups: 84,480 blocks.
Each block: 1024 threads, 32 channels, up to 32 positions.
Compute: 32 * 32 * 363 * 2 = 743,424 FLOPs per block.
Total blocks: 3 * ceil(55/32) * 55 * 256 = 3 * 2 * 55 * 256 = 84,480.
Total FLOPs: 84,480 * 743,424 ≈ 62.8 TFLOPs.
At 67 TFLOPS, this would take ~0.94 ms if compute-bound... but we have global memory loads for input still.
Actually, if we preload input into shared memory, the compute part is very fast because shared memory has higher bandwidth than global memory.
Hmm, but writing a kernel with shared memory for both input and weights is getting complex. Let me first try the simpler version without input shared memory and see if it passes correctness and what performance we get.
Actually, let me step back and consider whether we need a custom kernel at all for such a simple conv operation. cuDNN is highly optimized. But the task specifically asks for a custom CUDA kernel, and the reference is PyTorch's nn.Conv2d. PyTorch likely calls into cuDNN, which might not be perfectly optimized for this specific case (small N=96, skinny GEMM).
Also, PyTorch/cuDNN has overhead for algorithm selection, workspace allocation, etc. A custom kernel might have lower launch overhead.
But for fairness, PyTorch's conv2d on H100 for a single layer at batch=256 should be quite fast. Let me try to implement a reasonable kernel and iterate.
Let me start with a simple 256-thread direct convolution kernel with weights in shared memory.
Actually, I just had another idea. What if we use im2col in registers? We can unroll the kernel loops and have each thread compute 1 output element. The input accesses are strided but the compiler might optimize them. With weights in shared memory and input from global memory with __ldg, this might be competitive.
Let me write the kernel and test.
Wait, let me also consider using tensor cores via wmma or mma. For this GEMM (M=774400, N=96, K=363), tensor cores would require specific shapes. The K dimension of 363 is awkward - not a multiple of 16. We'd need padding and the setup overhead might not be worth it for a single conv layer. I'll skip tensor cores for now.
Let me write a simple but well-optimized direct kernel.