reasoning
Let me analyze the problem. We need to implement a custom CUDA kernel for 2D convolution that is faster than PyTorch's native implementation.
Given the test parameters:
- batch_size = 16
- in_channels = 16
- out_channels = 128
- kernel_size = 3
- width = height = 1024
- stride = 1
- padding = 0
- dilation = 1
- groups = 1
- bias = False
This is a reasonably large convolution. Let me think about the best approach.
For a Conv2d with these parameters on H100:
- Input: [16, 16, 1024, 1024]
- Output: [16, 128, 1022, 1022]
- Weights: [128, 16, 3, 3]
This is mostly compute-bound since input/output channels are relatively large.
On H100, we have tensor cores with FP32 support. We could use implicit GEMM via cutlass or write a simple direct convolution kernel. However, since PyTorch already uses cuDNN which is highly optimized, we need to find an edge case where we can do better.
A few observations:
1. PyTorch's reference may use the implicit GEMM path in cuDNN.
2. Since kernel_size is small (3x3) and in_channels=16, out_channels=128, this could be amenable to a direct convolution with shared memory tiling.
Actually, wait - on H100 with FP32, PyTorch's cuDNN is extremely optimized. To beat it, we might need to exploit specific characteristics:
- Small kernel size (3x3)
- Input channel = 16 which is quite small
- No padding, no bias
- We could potentially fuse operations or use a very carefully tiled direct convolution
Let me think about whether im2col + GEMM approach could be faster. cuDNN already does this implicitly. But maybe we can save memory traffic by doing direct convolution and keeping data in shared memory.
Another approach: Use CUTLASS or a custom implicit GEMM kernel directly in CUDA.
Let me think about what specific configuration could be faster:
- We have in_channels=16, out_channels=128, k=3.
- The total MACs per output pixel: 16 * 128 * 3 * 3 = 18,432
- Total output pixels: 16 * 1022 * 1022 = ~16.7M
- Total FLOPs: ~307 GFLOPs
At 67 TFLOPS FP32 (non-TC) or 989 TFLOPS with TC, this should take between ~0.3 ms to ~4.6 ms theoretical.
Actually, wait - PyTorch might not be using Tensor Cores for FP32 unless TF32 is enabled. Also, for FP32 on TC, it's really TF32 Tensor Cores, not true FP32.
For FP32 on this small channel count (16), cuDNN might not achieve great utilization. Let me design a custom direct convolution kernel that:
1. Uses shared memory for input tiles and weights
2. Processes multiple output channels per thread block
3. Uses vectorized loads (float4)
4. Carefully tuned for 16 input channels, 128 output channels, 3x3 kernel
One key insight: in_channels=16 means we can load an entire input channel set for a 3x3 kernel efficiently. With 16 float channels, 3x3 = 9 spatial positions, that's 144 floats = 576 bytes per sample. This fits very nicely in shared memory.
Let me design a kernel:
- Each thread block handles a tile of output.
- We can have each thread compute multiple output pixels to amortize load costs.
A common direct convolution approach for small kernels:
- Block dimensions: 16x16 threads (256 threads)
- Each thread block computes an output tile of, say, 16x16 output pixels × some number of output channels
- Since out_channels=128, we can divide by 16 or 32 output channels per block
Let me think more carefully. The input to shared memory should be (tile_width + kernel_size - 1) × (tile_height + kernel_size - 1) × in_channels.
For tile=16×16 with kernel 3x3, input tile is 18×18×16 = 5184 floats. At 4 bytes each, that's 20.7KB. With cooperative loading, this fits nicely in shared memory (up to 227KB per block).
But we also need weights in shared memory: if we do 16 output channels per block, weights = 16 × 16 × 3 × 3 = 2304 floats = 9KB.
Total shared memory = ~30KB, well within limits.
But there's another consideration - register usage. We need registers for accumulating. Each thread would compute 16x16 output pixels × (some output channels). That's a lot.
Let me revise:
- Thread block: 16x16 threads (256 threads)
- Each thread computes 1 output pixel
- Each block computes 16 output channels
- So each block computes a 16×16×16 output tile
For this:
- Input shared memory: 18×18×16 = 5184 floats
- Weight shared memory: 16×16×3×3 = 2304 floats
- Total: ~30KB
Each thread loads some portion of input and weights, then does MACs over in_channels × k × k = 16 × 3 × 3 = 144 operations.
144 FMAs per thread is a bit low. We want more compute per thread to amortize memory loads.
Let me try a different configuration:
- Thread block: 16x16 threads
- Each thread computes 2x2 output pixels = 4 pixels
- Tile of output: 32x32 pixels
- Output channels per block: 8
- Total output per block: 32×32×8 = 8192 pixels
Input tile: 34×34×16 = 18496 floats = ~74KB shared memory. That's too much.
Let me try:
- Thread block: 16x8 = 128 threads
- Each thread computes 4 output pixels (e.g., 2x2)
- Output tile: 32x16 pixels
- Output channels per block: 8
- Input tile: 34x18x16 = 9792 floats = ~39KB. Still high but maybe OK.
Actually, let me think about this differently. For H100 with 3.35 TB/s memory bandwidth and the convolution is compute-bound, the key is to maximize compute throughput.
Another alternative: use CUTLASS-based implicit GEMM. But that requires headers.
Given the constraints of load_inline, we can't easily include CUTLASS headers since they might not be available in the environment.
So let me write a custom direct convolution kernel. But actually, let me consider if we can use torch.matmul in a clever way or if a pure CUDA kernel is needed.
For speedup, I think a well-tuned CUDA kernel that avoids workspace allocation (which cuDNN im2col requires) could win.
Let me design an implicit GEMM kernel or direct convolution kernel.
Actually, I think the most promising approach for a kernel_size=3, in_channels=16, out_channels=128 is to use a specialized direct convolution kernel that exploits:
1. in_channels=16 is small, so we can unroll the inner loop
2. k=3 is small, so we can unroll spatial loop too
Here's a kernel design:
- Each block handles (out_ch_per_block) output channels and a spatial tile
- Each thread computes multiple output pixels
Let me try this specific configuration:
```
#define BLOCK_H 16
#define BLOCK_W 16
#define TILE_OC 8
```
Thread block = 16 × 16 = 256 threads
Each thread computes output for 1 spatial position
Block computes 16×16 spatial × 8 output channels = 2048 output values
Input shared: (BLOCK_H + KH - 1) × (BLOCK_W + KW - 1) × IC = 18 × 18 × 16 = 5184 floats = 20.7 KB
Weight shared: TILE_OC × IC × KH × KW = 8 × 16 × 3 × 3 = 1152 floats = 4.6 KB
Total shared: ~25 KB ✓
Each thread does: 8 (output ch) × 16 (input ch) × 3 × 3 = 1152 MACs
1152 FMAs per thread at 256 threads/block = 295K FMAs per block
For ~16.7M output pixels / 2048 per block = ~8192 blocks
Total FMAs = 8192 × 295K = ~2.4B FMAs = matches total compute
Wait, let me recalculate total output size:
Batch 16, OC 128, H 1022, W 1022 = 16 × 128 × 1022 × 1022 ≈ 2.14B output pixels
Each pixel requires 16 × 3 × 3 = 144 MACs
Total MACs = 2.14B × 144 ≈ 308B MACs = 616 GFLOPs (since each MAC is 2 FLOPs)
Hmm wait - Output size with no padding, kernel 3, stride 1: 1024 - 3 + 1 = 1022. Yes.
16 * 128 * 1022 * 1022 * 16 * 3 * 3 = let me recalc:
16 * 128 = 2048
1022 * 1022 ≈ 1,044,484
2048 * 1,044,484 ≈ 2,139M output positions
Each needs 16 * 3 * 3 = 144 mult-adds
Total mult-adds = 308B mult-adds = 616 GFLOPs. Yes that's right.
At 67 TFLOPS that's ~9.2 ms. At TF32 Tensor Core, ~0.62 ms.
Each thread computing 1152 MACs is not a lot. We'd launch 16*128*1022*1022 / 2048 = 8,189,952 blocks? No wait.
Output is [16, 128, 1022, 1022]. Each block computes [?, 8, 16, 16] in (oc, h, w).
So blocks needed = 16 * (128/8) * (1022/16) * (1022/16).
1022/16 = 63.875, so 64 tiles in each dimension with some padding.
Blocks = 16 * 16 * 64 * 64 = 1,048,576 blocks
1M blocks * 256 threads = 256M threads total
Each thread 1152 MACs = total 295B MACs... wait, that's off by a factor.
Let me recalculate: 1,048,576 blocks × 256 threads × 1152 MACs/thread = 309B MACs. Yes that's right.
1M thread blocks is actually a lot for H100. H100 has ~132 SMs, so that's ~8K blocks per SM, which is fine for latency hiding.
But 1152 MACs per thread might not be enough to saturate compute. Let me increase the work per thread.
What if we do vectorized loads and have each thread compute 2x2 output tiles?
Thread block = 16x16 = 256
Each thread computes 2x2 spatial outputs
Tile spatial = 32x32
Tile OC = 4
Output per block = 32×32×4 = 4096 values
Input shared: 34×34×16 = 18,496 floats = 74 KB. Too much for shared memory.
Let me try:
Thread block = 16x8 = 128
Each thread computes 2x2 spatial outputs
Tile = 32x16 spatial
Tile OC = 8
Output per block = 32×16×8 = 4096 values
Input shared: 34×18×16 = 9792 floats = 39 KB. Getting closer to limit.
Actually, max shared memory per block on H100 is 227KB, but we also need to stay within 255 registers per thread. Let me not be too conservative.
Hmm, but the occupancy matters. If we use too much shared memory or registers, we limit occupancy.
Let me try a different approach: What if we don't put input in shared memory but instead use a register file approach or directly load from global memory with good cache behavior?
For the specific case of in_channels=16, k=3:
- Each output pixel needs 16×3×3 = 144 input values
- With 16 output channels per block, each thread can cache the input in registers or shared memory
Actually, for this problem, since the input is reused heavily across output channels, the best is shared memory for input and weights.
Let me go back to the earlier design but increase tile size:
Thread block = 16x16
Tile OC = 16
Each thread computes 1 output pixel
Block output = 16×16×16 = 4096 values
Input shared: 18×18×16 = 5184 floats = 20.7 KB
Weight shared: 16×16×3×3 = 2304 floats = 9.2 KB
Total: ~30 KB
Each thread computes 16 output channels × 1 pixel = 16 values, each requiring 144 MACs = 2304 MACs per thread.
Blocks = 16 * (128/16) * ceil(1022/16) * ceil(1022/16) = 16 * 8 * 64 * 64 = 524,288 blocks.
524K blocks is more reasonable. Each thread does 2304 MACs which is better.
But can we do more? What about having each thread compute a 2x2 block of output for 8 output channels?
Block dims: 16x8 threads
Each thread: 2x2 outputs × 8 channels = 32 values
Spatial tile: 32x16
Total output: 32×16×8 = 4096 values
Input shared: 34×18×16 = 9792 floats = 39 KB
Weight shared: 8×16×3×3 = 1152 floats = 4.6 KB
Total: ~44 KB
Each thread: 32 outputs × 144 MACs = 4608 MACs.
This is getting better. But 34×18 = 612 input pixels × 16 channels = 9792 floats to load. With 128 threads, each thread loads 76 floats. With memory coalescing, this should be fine.
Actually, let me reconsider. H100 benefits from Tensor Cores. If we structure the computation as a matmul, we could get much higher throughput.
For a Conv2d with k=3, in_ch=16, out_ch=128:
The natural im2col would create a matrix of size [out_height*out_width, in_channels*k*k] = [1022*1022, 144]
And weight matrix of [144, 128]
This is [~1M, 144] × [144, 128] = fairly small inner dimension.
With such a small K (144), Tensor Cores might not be very efficient. This might be why a direct convolution could beat cuDNN's implicit GEMM if cuDNN doesn't have a specialized kernel for this case.
So let's go with a direct convolution kernel. Let me try to optimize it further.
Actually, for maximum speed on H100, I should think about:
1. Using vectorized loads (float4)
2. Avoiding shared memory bank conflicts
3. Maximizing register-level reuse
With in_channels=16, we have nice alignment for float4 if we process 4 channels at a time.
Let me design a kernel using float4 for input channels:
```cuda
__global__ void conv2d_direct(const float* __restrict__ input,
const float* __restrict__ weight,
float* __restrict__ output,
int N, int IC, int OC, int H, int W,
int K, int OH, int OW) {
// Block compute: (tile_oc, tile_oh, tile_ow)
// Thread within block maps to (thread_oh, thread_ow, thread_oc_group)
}
```
Actually, let me look at this from another angle. What if we use the implicit GEMM approach but in a custom kernel using shared memory, treating it as matrix multiply where:
- A matrix: [M = OH*OW, K = IC*K*K] - the "im2col" expanded input
- B matrix: [K = IC*K*K, N = OC] - the weight
- C matrix: [M, N] - the output
For this specific case:
- M = 1022*1022 = 1,044,484
- K = 16*3*3 = 144
- N = 128
Since K=144 is small, a standard GEMM might not be efficient. But with a custom kernel that:
1. Loads K dimension into shared memory in chunks
2. Each block computes a tile of C
Actually, for K=144, if we tile K by say 72, each thread block does:
- Load A tile [M_tile, 72] and B tile [72, N_tile]
- This is quite small.
For small-K GEMM, one approach is to not use shared memory for K but instead stream through K with direct loads, or use a specialized kernel.
Let me think about this differently. PyTorch likely uses cuDNN which has thousands of kernels. For this specific shape, it probably picks a good one. We need to find where we can win.
One possibility: cuDNN may have launch/algorithm selection overhead, or workspace allocation overhead. A fused custom kernel that avoids extra memory allocation could be faster, especially for fp32.
Another possibility: The reference implementation in PyTorch might not be the absolute fastest cuDNN path for fp32 if TF32 is not enabled or if the heuristics are off.
Wait, I just realized - the test might run on fp32 but with TF32 disabled, or it might run true fp32. cuDNN for true fp32 on H100 might use a non-Tensor Core path since FP32 Tensor Cores are TF32, not true fp32. So the reference might be running at ~67 TFLOPS or less.
For our custom kernel, we can also just do true fp32.
Let me try a fairly aggressive direct convolution kernel. I'll target:
- Use shared memory for input feature maps
- Load weights into registers or shared memory
- Each thread computes multiple output pixels across multiple output channels
Here's my refined design:
```cuda
// Configuration
#define TILE_H 16
#define TILE_W 16
#define TILE_OC 8
// Block: 16 x 16 threads = 256
// Grid: (OC/TILE_OC, OW/TILE_W, OH/TILE_H) per batch
// Actually: (OW/TILE_W, OH/TILE_H, OC/TILE_OC) with batch as outer loop
// Each thread computes 1 spatial position for TILE_OC output channels
// Input shared: (TILE_H + 2) * (TILE_W + 2) * IC = 18 * 18 * 16 = 5184 floats
// Weight: each thread reads from global directly or via shared
```
If weights are in global memory with L2 cache, and each warp accesses the same weights, broadcast should work well. We can cache some weights in shared memory or rely on L2.
Actually, with out_channels=128 and in_channels=16, kernel=3x3:
Total weight size = 128 * 16 * 9 = 18,432 floats = 73.7 KB
This fits in L1/L2 cache! So we could load weights directly from global memory.
Let me design a kernel where weights stay in global memory (or L2) and only input is in shared memory:
```cuda
template<int TILE_H, int TILE_W, int TILE_OC>
__global__ void conv_direct(const float* __restrict__ input,
const float* __restrict__ weight,
float* __restrict__ output,
int N, int IC, int OC, int H, int W,
int K, int OH, int OW) {
int n = blockIdx.z; // or some other mapping
int oc_base = blockIdx.x * TILE_OC;
int oh_base = blockIdx.y * TILE_H;
int ow_base = blockIdx.z * TILE_W; // Need to think about mapping
// Shared memory for input tile
__shared__ float s_input[(TILE_H + K - 1) * (TILE_W + K - 1) * IC];
int tx = threadIdx.x;
int ty = threadIdx.y;
int tid = ty * blockDim.x + tx;
// Load input tile into shared memory
// ...
__syncthreads();
// Compute output
float accum[TILE_OC] = {0};
for (int ic = 0; ic < IC; ++ic) {
for (int kh = 0; kh < K; ++kh) {
for (int kw = 0; kw < K; ++kw) {
float inp = s_input[...];
#pragma unroll
for (int oc = 0; oc < TILE_OC; ++oc) {
accum[oc] += inp * weight[(oc_base + oc) * IC * K * K + ic * K * K + kh * K + kw];
}
}
}
}
// Store output
for (int oc = 0; oc < TILE_OC; ++oc) {
int oh = oh_base + ty;
int ow = ow_base + tx;
if (oh < OH && ow < OW) {
output[((n * OC + oc_base + oc) * OH + oh) * OW + ow] = accum[oc];
}
}
}
```
The issue with this is that the inner loop over TILE_OC is small (8) and each iteration reads from global memory. But since all threads in a warp read the same weight (for the same oc), they should be broadcast from L2.
Actually, let me reconsider. H100 has very good L2 cache. If we iterate with threadIdx.y over output channels and threadIdx.x over spatial, then adjacent threads read adjacent weights which is bad for broadcast.
Better approach: Have all threads in a warp compute the same output channel group, so they all read the same weights. This means the weight accumulator loop should be unrolled with broadcast reads.
So:
- BlockDim.x = TILE_W (say 16)
- BlockDim.y = TILE_H (say 16)
- All threads compute TILE_OC output channels
Wait, but then each thread computes 1 pixel × TILE_OC channels. The weights for different oc are read by all threads but different oc values. Since adjacent threads read the same weight for a given oc (broadcast), and we iterate over oc, this is actually good IF we have enough warps to hide the latency.
Actually, no - for a given (ic, kh, kw), all threads in the block need all TILE_OC weights. These can stay in registers after being loaded once, and shared across threads via the memory subsystem.
Hmm, this is getting complex. Let me think about using shared memory for weights too.
If we put weights in shared memory for TILE_OC=16:
Weight shared = 16 * 16 * 3 * 3 * 4 bytes = 9.2 KB
Combined with input shared 20.7 KB = ~30 KB total.
This is very manageable. And then weights are read from shared memory with no bank conflicts (since all threads read the same weight for a given oc, or if we arrange correctly, threads read different oc weights).
Wait, if threads are computing different oc channels simultaneously, they'll read different weight addresses. If thread i reads weight[oc_i] and thread j reads weight[oc_j], these could conflict.
To avoid conflicts, we should have threads compute the SAME oc but DIFFERENT spatial positions. Then all threads read the same weight (broadcast, no bank conflict in shared memory). But then we need sequential execution over oc, which is fine.
So the right approach is:
- Threads within a block differ in spatial position (oh, ow)
- All threads compute the same set of oc channels sequentially
- Or better: each thread computes a subset of oc channels
Actually, for shared memory weights, let's say we tile by TILE_OC=8. We load 8 output channel weights into shared memory. Then all threads compute these 8 channels. For this to be efficient:
- Thread (tx, ty) computes output at spatial position (oh_base + ty, ow_base + tx)
- For oc = 0 to 7, thread computes output for that oc
In this case, at a given inner loop iteration over oc, all threads are reading the SAME weight for their spatial position. Wait no - for a fixed (ic, kh, kw), each oc has a different weight. So if all threads compute oc=0 simultaneously, they all read weight for oc=0, which is a broadcast. Good.
Then for oc=1, all threads read weight for oc=1. Another broadcast. Good.
So sequential over oc with all threads computing the same oc is actually ideal for weight sharing.
But then each thread needs an accumulator for each oc. With TILE_OC=8, each thread has 8 accumulators. That's fine.
Wait, but I said all threads compute the SAME oc at a time. So each thread computes 8 output values (1 per oc). Yes.
The loop structure would be:
```
for oc in 0..TILE_OC-1:
accum[oc] = 0
for ic in 0..IC-1:
for kh in 0..K-1:
for kw in 0..K-1:
float w = s_weight[oc * IC * K * K + ic * K * K + kh * K + kw];
float inp = s_input[...];
accum[oc] += w * inp;
```
In this case, all threads read the same `w` (broadcast, good), and read different `inp` (coalesced, good) because threads correspond to different spatial positions.
This is the ideal pattern for both shared memory and cache performance!
Let me refine:
```cuda
#define TILE_OC 8
#define TILE_H 16
#define TILE_W 16
__global__ void conv2d_kernel(const float* __restrict__ input,
const float* __restrict__ weight,
float* __restrict__ output,
int N, int IC, int OC, int H, int W,
int K, int OH, int OW) {
int n = blockIdx.z;
int oc_block = blockIdx.x;
int oh_base = blockIdx.y * TILE_H;
int ow_base = (blockIdx.z / N) * TILE_W; // This mapping is wrong, let me fix
```
Actually I need a clean block mapping. Let me use:
- blockIdx.x: oc_block (OC / TILE_OC)
- blockIdx.y: spatial_h_block (OH / TILE_H)
- blockIdx.z: batch * spatial_w_block + spatial_w_block
Or:
- blockIdx.x: w_block
- blockIdx.y: h_block
- blockIdx.z: n * (OC/TILE_OC) + oc_block
This keeps the z dimension large enough for good occupancy.
Actually, a simpler mapping:
- blockIdx.x: oc_block
- blockIdx.y: hw_block (combined h and w blocks)
- blockIdx.z: n
But CUDA grids have a limit on each dimension. For our problem:
- OC/TILE_OC = 16 blocks
- OH/TILE_H = 64 blocks
- OW/TILE_W = 64 blocks
- N = 16
Total blocks = 16 * 64 * 64 * 16 = 1,048,576
Grid dim limits are fine (up to 2^31-1 for x and y, up to 65535 for z before recent CUDA, but newer CUDA might have higher limits).
Actually, let me check - max gridDim.z is typically 65535 even on modern CUDA for some compatibility reasons. So if I put N in z, and N=16, and OC/TILE_OC=16, I could combine them: z = n * 16 + oc_block.
But then y needs to handle 64*64 = 4096. gridDim.y limit is also typically large (65535 or more).
Let me use:
- blockIdx.x: w_block (0 to 63)
- blockIdx.y: h_block (0 to 63)
- blockIdx.z: n * (OC/TILE_OC) + oc_block (0 to 255)
Wait, gridDim.z might be limited to 65535, and 256 is fine.
Actually, I need to be careful about warp scheduling. On H100, having z dimension up to 256 might be fine.
But actually, to maximize spatial locality and coalescing, it's often better to put spatial dimensions in x and y.
Let me use:
- blockIdx.x: w_block
- blockIdx.y: h_block
- blockIdx.z: combined n and oc_block
And within the kernel:
```
int n = blockIdx.z / (OC / TILE_OC);
int oc_block = blockIdx.z % (OC / TILE_OC);
int oh_base = blockIdx.y * TILE_H;
int ow_base = blockIdx.x * TILE_W;
```
Now for loading input into shared memory:
Input tile covers [oh_base, oh_base + TILE_H + K - 1) vertically and [ow_base, ow_base + TILE_W + K - 1) horizontally, for all input channels.
The padded tile size is: (TILE_H + K - 1) × (TILE_W + K - 1) × IC = 18 × 18 × 16 = 5184 elements.
With 256 threads, each thread loads 5184 / 256 = 20.25 elements. That means some threads load 21 elements. We can handle this with a loop.
For coalescing, we'd want adjacent threads to load adjacent elements. We can linearize the 3D input tile and have each thread load a contiguous chunk. But the input in global memory is NCHW, so we need to think about strides.
Input format: [N, IC, H, W]
So input[n, ic, h, w] is at offset ((n * IC + ic) * H + h) * W + w
The input tile spans multiple channels and spatial positions. To coalesce loads, we should have threads in a warp load across the W dimension first, then H, then IC.
Actually, let me think about the layout in shared memory. I'll store it as [IC][tile_h][tile_w] to match the access pattern in computation, or [tile_h][tile_w][IC] to match NHWC-like access.
For the computation, we access by (ic, local_h, local_w). If we store in row-major as [ic * tile_h * tile_w + h * tile_w + w], then for a fixed ic and h, adjacent w are contiguous. This is good for accessing multiple w values across a warp if threads map to different w.
But in our kernel, different threads map to different (oh, ow) spatial positions. For a given (ic, kh, kw), thread (ty, tx) reads input at (oh_base + ty + kh, ow_base + tx + kw). So adjacent tx threads read adjacent w positions. This is great if we lay out shared memory as [local_h][local_w] for each ic.
Actually, for the inner loop over kh, kw, each thread accesses a single spatial position. The coalescing isn't critical here since each thread is reading a single float. But if we want to use vectorized loads (float4), we could load 4 w-contiguous elements.
Wait, with IC=16, we could process 4 channels at a time using float4. Let me explore this.
If we use float4, treating 4 floats as one unit, then for IC=16, we have 4 float4 units per spatial position. When loading input, each spatial position (h, w) has 4 float4 values representing 16 channels.
For the computation, we'd do:
```
float4 inp4 = s_input4[local_h * tile_w + local_w];
// But we need to access individual channels to multiply with weights
```
Actually, for convolution, we need:
out[oc] += sum over ic of in[ic] * w[oc, ic, kh, kw]
If we have ic as vector of 4, we'd compute partial dot products. But since weight access is per (oc, ic), and ic is sequential, using float4 for input and weights could work.
Weight for oc is [IC][K][K] in shape. With IC=16, that's 16 floats per (oc, kh, kw). We could load 4 floats at a time.
So the computation becomes:
```
for kh:
for kw:
for ic_group in 0..3:
float4 in4 = s_input[(ic_group * 4) * tile_h * tile_w + local_h * tile_w + local_w];
for oc in 0..TILE_OC:
float4 w4 = s_weight[oc * 4 + ic_group]; // Need proper indexing
accum[oc] += in4.x * w4.x + in4.y * w4.y + in4.z * w4.z + in4.w * w4.w;
```
This is getting complicated. Let me keep it simple first with float, and optimize if needed.
Simple kernel first, then we can iterate on optimization.
Actually, here's another important insight: For this specific problem, the most efficient kernel might actually leverage the fact that kernel_size=3, in_channels=16, and fully unroll those loops. The compiler can unroll small constant loops very well.
Let me write a kernel that:
1. Has the loop over kh and kw fully unrolled (since K=3 is constant, we can use template)
2. Has IC=16 known at compile time (we can use template parameter if we want, but since it's passed as arg, compiler might still unroll with pragma)
Actually, we know the test config but the kernel needs to be general. However, using template or #define for K=3 and expecting compiler optimization is reasonable.
Let me just write a clean kernel and see how it performs.
Wait, another approach: Could we just use torch's MatMul in a clever way without writing CUDA C?
For im2col:
- Use unfold to extract [N, IC, KH, KW, OH, OW]
- Reshape to [N, IC*KH*KW, OH*OW]
- Weight reshape to [OC, IC*KH*KW]
- matmul: [N, OH*OW, IC*KH*KW] × [N or broadcast, IC*KH*KW, OC]? Actually [N, M, K] × [N, K, OC] using bmm
- Reshape back
But unfold creates a large intermediate tensor. And for the sizes here:
- Unfold output: [16, 16, 3, 3, 1022, 1022] = 16 * 16 * 9 * 1022 * 1022 * 4 bytes ≈ 9.4 GB
This is huge and slow due to memory bandwidth.
So custom CUDA kernel is definitely the way to go.
Let me think about optimization strategies again:
1. **Shared Memory for Input Only**: Load input tile to shared memory, read weights from global memory with L2 caching. Total weight size is 73KB, which fits in L2 cache.
2. **Shared Memory for Both Input and Weights**: Load both to shared memory. Higher shared memory usage but better latency.
3. **Register Blocking**: Each thread computes multiple output pixels to increase compute intensity.
For H100, since L2 is large (50MB+) and the weights are only 73KB, they will easily stay in L2. So loading weights from global might be almost as fast as shared memory.
I think the simplest high-performance kernel is:
- Load input to shared memory
- Keep weights in global/L2
- Each thread computes a 2x2 block of output pixels for multiple output channels
- Use float4 for vectorized loads where possible
Let me try a specific kernel configuration:
```cuda
// Block: 16x8 threads
// Each thread: computes 2x2 outputs for 4 output channels
// Tile: 32x16 spatial, 4 output channels
// Output per block: 32*16*4 = 2048
__shared__ float s_input[18 * 18 * 16]; // 18*18*16 = 5184 floats
// Actually for 32x16 spatial tile with K=3, input is 34x18
__shared__ float s_input[34 * 18 * 16]; // 34*18*16 = 9792 floats = 39 KB
```
Wait, if I use 16x8 threads and each thread computes 2x2 outputs:
- Spatial tile: (16*2) × (8*2) = 32 × 16
- For batch dimension and output channels, I'll need to figure out grid
Actually, let me do this:
- Block: 16x16 threads
- Each thread computes a 1x1 output for 8 output channels
- Tile: 16x16 spatial, 8 output channels
- This is simpler and I know the shared memory fits
With 524K blocks and 256 threads each doing 2304 MACs, this might be decent. But 2304 MACs per thread is still not a lot.
What if we compute 1x1 output for 16 output channels? Then 4608 MACs per thread. With 262K blocks.
What if we compute 2x2 outputs for 4 output channels? Still 2304 MACs per thread, but fewer memory loads per MAC since spatial reuse.
Actually, spatial reuse is key! If a thread computes a 2x2 block of outputs, the input region it needs is 4x4 (for K=3) or more precisely, outputs at (h, w), (h, w+1), (h+1, w), (h+1, w+1) need inputs at (h to h+2, w to w+2) and (h to h+2, w+1 to w+3), etc. Shared memory helps with this.
For a 2x2 output block with K=3, the required input region is 4x4 (overlapping). So 4 outputs share input, giving better compute-to-load ratio.
But using shared memory already provides this reuse across all threads in a block. So whether each thread computes 1 pixel or 4 pixels doesn't change the shared memory loading - it's done once per block.
What does change is the register load from shared memory. If each thread computes 4 pixels, it loads 4 values from shared memory and accumulates. The shared memory bandwidth is very high, so this might not be the bottleneck.
The real bottleneck is likely the instruction throughput for the multiply-adds. H100 can do a lot of FMAs, but 4608 FMAs per thread with 262K blocks should be fine.
Let me try with 16x16 threads, each computing 16 output channels for 1 pixel. Then each thread does:
- 16 accumulators in registers
- For ic in 0..15:
- For kh in 0..2:
- For kw in 0..2:
- Load 1 input from shared
- For oc in 0..15: load weight, do FMA
Wait, if we have 16 accumulators and the inner loop is unrolled, compiler should keep accumulators in registers. But loading weights sequentially from shared memory or global memory might not saturate the FMA units.
A better approach for instruction-level parallelism: Have each thread compute multiple pixels, so it can interleave loads and FMAs.
Actually, let me reconsider the "weights in shared memory" approach with all threads computing the same oc at a time:
```cuda
__shared__ float s_input[18][18][16]; // [local_h][local_w][ic] or [ic][local_h][local_w]
__shared__ float s_weight[16][3][3][16]; // [oc][kh][kw][ic] - wait this is 16*9*16 = 2304 floats
```
If we store weights as [oc][kh][kw][ic], then for a fixed oc, kh, kw, the ic dimension is contiguous. All threads read the same weight[oc][kh][kw][ic] for each ic. Since all threads in the warp read the same address, this is a broadcast - no bank conflicts.
And for each ic, all threads read different input positions. If input is stored as [ic][local_h][local_w], then for fixed ic, local_h, local_w, threads reading different local_w are reading contiguous memory - good coalescing.
Wait, but if all threads read the same weight simultaneously, and different inputs, then:
```
for oc = 0..15:
accum[oc] = 0
for ic = 0..15:
for kh = 0..2:
for kw = 0..2:
float w = s_weight[oc * 9 * 16 + kh * 3 * 16 + kw * 16 + ic]; // Same for all threads
float in = s_input[ic * 18 * 18 + (thread_h + kh) * 18 + (thread_w + kw)]; // Coalesced across warp
accum[oc] += w * in;
```
This is perfect! The weight read is broadcast, and the input reads are coalesced.
With unrolled loops (since K=3 and IC=16 are small), compiler can optimize this very well.
Each thread computes 16 output values (16 channels for 1 spatial position) = 16 * 16 * 3 * 3 = 2304 FMAs.
Actually, with pragma unroll on the ic, kh, kw loops, the compiler will generate straight-line code with 2304 FMA instructions (plus loads).
On H100, each SM can issue a lot of FMAs per cycle. The throughput should be limited by either:
1. Shared memory bandwidth (but with broadcast for weights and coalesced for input, it's very efficient)
2. Instruction issue rate
3. Register dependency latency
With 2304 independent accumulators? No, we have 16 accumulators. The FMAs for different ic/kh/kw values have a dependency on the same accumulator. So the next FMA for accum[oc] can't start until the previous one completes.
This is a chain of 2304/16 = 144 dependent FMAs per accumulator. FMA latency on H100 is ~4 cycles for FP32. So the minimum time per thread is 144 * 4 = 576 cycles. With 256 threads, that's 576 cycles per block.
At 2 GHz, 576 cycles = 288 ns per block. For 262K blocks... but blocks run in parallel across SMs. With 132 SMs and say 2 blocks per SM concurrently (depending on occupancy), that's ~1000 blocks running at once. Total time ~262K / 1000 * 288ns ≈ 75ms. That seems too long.
Wait, my math is wrong. Actually, H100 has much higher throughput. Each SM can execute many threads simultaneously with latency hiding.
For throughput: 132 SMs * ~100 FMAs/cycle/SM (rough estimate) = ~13,200 FMAs/cycle total.
Total FMAs = 262K blocks * 256 threads * 2304 FMAs/thread = 154B FMAs.
At 13,200 FMAs/cycle and 2GHz = 26.4 TFMA/s = 154B / 26.4B/s = 5.8 seconds? That can't be right.
Hmm wait. H100 FP32 throughput is 67 TFLOPS. Each FMA is 2 FLOPs, so 33.5 TFMA/s.
154B FMAs / 33.5B/s = 4.6 seconds. WAY too slow.
What's wrong? Oh! The total output is [16, 128, 1022, 1022] = 2.14B values.
Each output value needs 144 MACs = 288 FLOPs.
Total FLOPs = 2.14B * 288 = 616 GFLOPs.
At 67 TFLOPS = 9.2 ms. Yes, that's the theoretical minimum.
Let me recalculate number of blocks:
If each block computes 16x16x16 = 4096 output values:
Blocks needed = 2.14B / 4096 = 523,264 blocks.
Each block 256 threads. Each thread computes 16 outputs.
FMAs per block = 4096 outputs * 144 MACs/output = 589,824 MACs = 589K FMAs? No wait, each output needs 144 MACs (mult-add = 1 FMA).
Total FMAs = 2.14B outputs * 144 MACs = 308B FMAs.
At 67 TFLOPS for FMA (which counts as 2 FLOPs), actually FP32 throughput of 67 TFLOPS means ~33.5 TFMA/s.
Wait, I'm confusing FMA and FLOP. An FMA is usually counted as 2 FLOPs.
H100 FP32: 67 TFLOPS = 33.5 TFMA/s.
Total FMAs = 308B.
Time = 308B / 33.5B/s = 9.2 ms. Yes.
Now, each block with 256 threads doing 2304 FMAs (16 outputs × 144 MACs) would take:
256 threads * 2304 FMAs = 589,824 FMAs per block.
Blocks: ~524K (we need to account for padding to 1024x1024 tiles, but close enough).
Total FMAs: 524K * 589K = 308B. Correct.
With 132 SMs and say we can run 4 blocks per SM (1024 threads roughly, depending on occupancy), we have 528 blocks in flight.
Total time ≈ (524K / 528) * block_time.
Block time with 256 threads * 2304 FMAs. At 4 FMAs/cycle per SM per thread? Actually each SM has much higher FMA throughput.
If block_time is roughly the time for all threads in the block to complete their FMAs, and with good latency hiding, block_time ≈ (256 threads * 2304 FMAs / threads_per_SM / FMA_throughput_per_SM).
This is getting complicated. The point is, 9.2 ms is the theoretical minimum for pure compute. We need to also account for memory loads.
For memory: Each block loads:
- Input: 18*18*16*4 bytes = 20.7 KB
- Weight: 16*16*9*4 bytes = 9.2 KB
Total per block: ~30 KB
For 524K blocks = 15.7 GB of memory traffic
At 3.35 TB/s, 15.7 GB takes 4.7 ms. But this is from shared memory, which is much faster.
Global memory reads are:
- Input: 16 * 16 * 1024 * 1024 * 4 = 1 GB (read once, since tiles don't overlap much with padding)
- Weight: 128 * 16 * 9 * 4 = 73 KB
- Output: 16 * 128 * 1022 * 1022 * 4 = 8.5 GB
Total global: ~9.5 GB
At 3.35 TB/s = 2.8 ms for memory-bound.
Given that compute is 9.2 ms and memory is 2.8 ms, this is compute-bound. The kernel needs to maximize FMA throughput.
Now, with the sequential accumulator dependency, the throughput is limited by the latency of the dependency chain. For 16 accumulators, each with 144 dependent ops, if FMA latency is 4 cycles, then to fully utilize the FMA units, we need at least enough warps to cover this latency.
On H100, each SM has many warp schedulers. If each thread has 16 independent accumulators (for 16 output channels), and the inner loop of 144 iterations can be pipelined... actually no, all 144 iterations target the SAME 16 accumulators. So there's a 4-cycle dependency between consecutive iterations.
But wait, the compiler can unroll and interleave! If pragma unroll is used, the compiler might interleave multiple FMAs targeting different accumulators.
Actually, since the inner loop over kh, kw, ic gives us 144 iterations, and each iteration does 16 FMAs to 16 different accumulators. The 16 FMAs within an iteration are independent! So:
For iteration i: FMA(accum[0], in[i], w[i,0]), FMA(accum[1], in[i], w[i,1]), ..., FMA(accum[15], in[i], w[i,15])
These 16 FMAs are independent of each other. But iteration i+1 depends on iteration i's results for the same accumulators.
So the dependency chain is only within each accumulator. With 16 independent accumulators, we can overlap the latency.
FMA latency is ~4 cycles. We'd need 4 * 2 = 8 independent FMAs in flight to fully utilize 1 FMA pipe per warp? Actually, modern GPUs have multiple FMA units per SM and within a warp, multiple independent FMAs can be issued to different execution units.
For FP32 on H100, a warp of 32 threads can issue many FMAs. If each thread does 16 FMAs, that's 512 FMAs per warp. With sufficient instruction-level parallelism, throughput should be good.
However, the key issue is the loop-carried dependency on accumulators. For throughput, we'd want many independent operations. One solution: Have each thread compute MULTIPLE spatial pixels, creating more independent accumulators.
If each thread computes 2x2 spatial pixels for 4 output channels:
- 2x2 * 4 = 16 accumulators, same as before
- But now the inner loop still has 144 iterations
- The 16 accumulators are still independent within each iteration
So the number of independent accumulators is the same. The difference is that for 2x2 pixels, the input values differ but the weights are the same. So the FMAs are still independent.
Wait, actually the number of accumulators doesn't change - we still have the same ILP. Unless we compute multiple pixels for the SAME output channel, which creates more accumulators for that channel.
If a thread computes 4 pixels for 4 channels: 16 accumulators.
If a thread computes 1 pixel for 16 channels: 16 accumulators.
Same ILP.
But for 2x2 pixels, each pixel needs a different input but same weights. The FMAs targeting different pixels are independent. So with mix of channel and pixel parallelism, we have the same number of independent accumulators.
To increase ILP, we need each thread to compute MORE outputs. Like 1 pixel for 32 channels (32 accumulators) or 2x2 pixels for 8 channels (32 accumulators).
But adding more accumulators uses more registers. H100 allows 255 registers per thread, and 32 floats use 32 registers - very reasonable.
Let me consider computing more output channels per thread.
If tile_oc=32, each thread computes 32 outputs. But then weight shared memory = 32*16*9*4 = 18.4 KB, and input = 20.7 KB, total = 39 KB. Still OK.
Each thread then does 32 * 144 = 4608 FMAs with 32 independent accumulators.
Blocks = 16 * (128/32) * 64 * 64 = 131,072 blocks.
At 32 accumulators, we have better ILP. Let me go with this!
Tile config:
- TILE_H = 16
- TILE_W = 16
- TILE_OC = 32
Block = 16x16 = 256 threads
Each thread computes 1 pixel for 32 output channels.
Shared memory:
- Input: 18*18*16*4 = 20,736 bytes
- Weight: 32*16*9*4 = 18,432 bytes
- Total: 39,168 bytes = 38.2 KB < 227 KB limit ✓
Register estimate:
- 32 accumulators
- A few temporaries for loading
- Total maybe ~40 registers, well under 255 ✓
Number of blocks:
- N=16, OC/TILE_OC=4, OH/TILE_H=64, OW/TILE_W=64
- Total: 16 * 4 * 64 * 64 = 262,144 blocks
- Blocks per SM: 262,144 / 132 ≈ 1986 blocks/SM
This should provide excellent occupancy and latency hiding.
Actually, can we fit even more?
TILE_OC = 64:
- Weight shared = 64*16*9*4 = 36,864 bytes
- Input shared = 20,736 bytes
- Total = 57,600 bytes = 56.25 KB. Still OK.
Blocks = 16 * 2 * 64 * 64 = 131,072 blocks
Each thread has 64 accumulators = 64 registers just for that.
Plus other regs, maybe ~80 total. Still under 255.
TILE_OC = 128 (all output channels):
- Weight shared = 128*16*9*4 = 73,728 bytes
- Input shared = 20,736 bytes
- Total = 94,464 bytes = 92.25 KB. Still well under limit.
- Blocks = 16 * 1 * 64 * 64 = 65,536 blocks
- Each thread 128 accumulators. That's 128 registers. Plus some for indexes/loads. Maybe ~150 total. Still OK.
65,536 blocks for 132 SMs = ~496 blocks/SM. This is getting lower but should still be enough.
If we do TILE_OC=128, each block computes a 16x16x128 output tile = 32,768 values.
Each thread computes 128 values.
FMAs per thread = 128 * 144 = 18,432
With 128 accumulators, we have tons of ILP. The FMA units should be very well utilized.
Wait, but then each block would need to load 128*16*9 = 18,432 weights into shared memory. With 256 threads, each thread loads 72 weights. That's a lot of loading overhead compared to compute.
Also, shared memory for weights is 73KB. If we have a 16x16 thread block computing 16x16x128 outputs, that's a large tile.
Actually, let me reconsider the occupancy constraint more carefully. If shared memory is 92KB per block, and each SM has 227KB shared memory, then at most 2 blocks can run on an SM simultaneously. This could limit occupancy.
With 2 blocks per SM * 132 SMs = 264 blocks concurrently. 65,536 blocks / 264 = ~248 rounds of block execution. If each block takes some time, total execution time might be significant.
With TILE_OC=32 and shared memory 38KB, we could fit up to 5 or 6 blocks per SM (227/38 = 5.97), but limited by other factors. At 4 blocks per SM, that's 528 blocks concurrently. 262,144 blocks / 528 = ~497 rounds.
Actually, the number of rounds isn't the right metric. What matters is whether there are enough concurrent warps to hide latency.
For TILE_OC=32 with 256 threads/block (8 warps), and 4 blocks/SM = 32 warps/SM. H100 can handle hundreds of warps per SM, so 32 might not be enough to hide all latency.
For TILE_OC=128 with 256 threads/block, and 2 blocks/SM = 16 warps/SM. Even fewer warps, but each thread has much more compute, so latency hiding is less critical.
With more ILP (128 independent accumulators), the kernel naturally needs less thread-level parallelism to keep the FMA units busy.
I think TILE_OC=64 is a good sweet spot:
- Weight shared = 36.9 KB
- Input shared = 20.7 KB
- Total = 57.6 KB
- Blocks/SM = up to 3 or 4
- Warps/SM = 24 to 32
- Each thread has 64 accumulators = good ILP
- Total blocks = 131,072
Let me go with this configuration and see how it performs.
Actually, let me also consider increasing the spatial tile to reduce the number of blocks. With TILE_H=32 and TILE_W=32:
- Input shared = 34*34*16*4 = 73,728 bytes = 72 KB
- Weight shared = 64*16*9*4 = 36,864 bytes = 36 KB
- Total = 108 KB. This limits to 2 blocks per SM.
Block = 32x32 = 1024 threads (the max per block is 1024).
Each thread computes 1 pixel for 64 output channels.
FMAs per thread = 64 * 144 = 9,216.
Blocks = 16 * 2 * 32 * 32 = 32,768 blocks.
This is fewer blocks but more work per block.
Hmm, 32,768 blocks / (132 SMs * 2 blocks) = 124 rounds. If each block takes ~0.1ms, that's 12.4ms total. But this depends heavily on block time.
Actually, 9,216 FMAs per thread at 4 cycles FMA latency with 64 independent accumulators gives good throughput. The issue is whether 1024 threads per block causes issues.
Let me try a middle ground. TILE_H=16, TILE_W=16 with TILE_OC=64 seems good.
Wait, let me re-examine. PyTorch likely runs this in ~5-10 ms (rough estimate). To beat it, we need something close to theoretical. And cuDNN is very well tuned.
Another insight: For this specific case where IC=16 is small, maybe a specialized kernel that treats IC as float4 (4 channels) would be much faster. But only if we can get good vectorization.
With IC=16, reading 16 channels as 4 float4 values is natural. And weight can also be read as float4. The dot product of 4 float4 vectors for input and weight can be computed efficiently.
For a single output channel and a single spatial position:
out = dot(in[0..15], w[0..15]) for each (kh, kw)
With float4: 4 dot-products of float4 = 4 * (4 fma) per (kh, kw).
For K=3, that's 9 * 16 = 144 FMAs = 36 float4 operations.
But the issue is that weights for different output channels are different, so the float4 dot-product needs to be done per output channel.
For 64 output channels, we'd do this 64 times.
Actually, let me think about this in terms of matrix multiply. The operation is:
- A matrix: [OH*OW, IC*K*K] - implicit im2col
- B matrix: [IC*K*K, OC] - weight
- C matrix: [OH*OW, OC] - output
For K=3, IC=16, this is [M, 144] × [144, 128] where M ≈ 1M.
In CUTLASS terms, this is a GEMM with small K. For small-K GEMM, specialized kernels that load the K dimension once per tile and maximize register reuse are used.
We could write a kernel where:
- Each block computes a [M_tile, OC_tile] output tile
- The K=144 dimension is loaded once into shared memory or registers
- Actually K=144 is small enough to keep in shared memory or even registers
If M_tile = 64 (e.g., 8x8 spatial) and OC_tile = 64:
- C tile = 64x64 = 4096 values
- A tile = 64x144 = 9,216 floats = 36.9 KB
- B tile = 144x64 = 9,216 floats = 36.9 KB
Both fit in shared memory! And then it's a straightforward GEMM in shared memory.
For M=1,044,484 and M_tile=64, we need ~16,320 tiles in M.
For OC=128 and OC_tile=64, we need 2 tiles in OC.
Total blocks = 16 * 16,320 * 2 = 522,240 blocks.
Each block does a 64x64x144 GEMM in shared memory = 1.18M FMAs.
256 threads = ~4608 FMAs per thread.
This is very similar to our direct convolution!
The advantage of GEMM formulation is we can optimize the inner loop more easily.
But direct convolution might be more efficient because:
1. We don't need to materialize the full im2col A matrix
2. We can load the input tile once instead of loading duplicated values
Actually in the GEMM approach with shared memory, the A tile corresponds to an im2col tile. We'd need to load input and expand it to im2col format on the fly. This "on-the-fly" expansion is what implicit GEMM does.
For a direct convolution kernel, we're implicitly doing the GEMM by loading the input and accessing the correct elements.
I think direct convolution is the right approach. Let me write the kernel.
One more consideration: PyTorch by default uses NCHW format. But tensor cores often prefer NHWC. However, for our custom kernel, we stick with NCHW since that's what PyTorch gives us.
Let me write a kernel with these parameters:
- TILE_H = 8
- TILE_W = 32
- TILE_OC = 32
Why TILE_W=32? This gives better coalescing for output writes. With 32 threads in x dimension, a warp writes 32 contiguous output values.
Block = 8x32 = 256 threads.
Each thread computes 1 pixel for 32 output channels.
Input shared = (8+2)*(32+2)*16 = 10*34*16 = 5440 floats = 21.8 KB
Weight shared = 32*16*9 = 4608 floats = 18.4 KB
Total = 40.2 KB
Blocks:
- OH/TILE_H = 1022/8 = 128 tiles
- OW/TILE_W = 1022/32 = 32 tiles
- N=16, OC/TILE_OC=4
- Total = 16 * 128 * 32 * 4 = 262,144 blocks.
Each thread FMAs = 32 * 144 = 4608
This looks good. TILE_W=32 gives clean warp coalescing. TILE_H=8 might be OK.
Alternatively, TILE_H=16, TILE_W=16 is more balanced. Let me stick with 16x16 since it's simpler.
Actually, let me consider warp shape. With block=16x16, the warp shape depends on how threads are grouped. Typically, threads are linearized as threadIdx.y * blockDim.x + threadIdx.x, so warp has 32 consecutive threads. If blockDim.x=16, then a warp covers threadIdx.x=0..15 and threadIdx.y=0..1 (2 rows).
For output writes, threads in a warp write 2 rows × 16 columns, which is 32 contiguous floats per row but only 16 per row per warp. With 16x16 outputs and TILE_OC=64, output format is [N, OC, OH, OW]. For a fixed n and oc, OH×OW is row-major. So writing 16 consecutive W values per row is coalesced.
But writing different OC channels: output[oc0][h][w] and output[oc1][h][w] are not contiguous (they're OC apart). So we should write all OC values for one pixel sequentially.
Each thread computes pixel (h, w) for 64 OCs. If we store these 64 values in local array, then we write them one by one. The writes are at stride OH*OW in output memory. So threads in a warp writing their first OC value are at addresses:
`base + (oc0) * OH * OW + h * OW + w` for each thread where w varies by thread.
Wait, output is NCHW: [N, OC, OH, OW]
Position (n, oc, h, w) is at offset ((n * OC + oc) * OH + h) * OW + w
For fixed n, h, and varying w across a warp, and fixed oc, the addresses are contiguous (differ by 1). Good!
For varying oc, the addresses differ by OH*OW.
So in our kernel, for a given (h, w), we should write oc=0, then oc=1, etc. Each write for a specific oc is coalesced across the warp (since w varies).
But writing 64 OCs means 64 writes per thread. If we have 32 accumulators and write sequentially, that's fine.
Actually, I want to minimize these writes. With TILE_OC=64, each thread writes 64 floats. 64 writes per thread = 64 * 256 = 16,384 writes per block = 64KB written. At high bandwidth this is fine, but it's more instruction overhead.
What if we write with float4? We can write 4 OCs at a time. But as noted, OCs are not contiguous for the same pixel. float4 doesn't help unless we transpose.
Between TILE_OC=64 (more ILP but more writes) and TILE_OC=32 (less ILP but fewer writes), let me try 64 and see.
Actually, I just realized: H100 supports warp-level primitives. We could potentially use tensor cores or warp-matrix-multiply if we structure this right. But given the small K=144, standard FMA might be better.
Let me also think about whether we can use `wmma` or use CUTLASS. Actually, for small K, tensor cores are not efficient because they operate on larger tiles (e.g., 16x16x16).
OK, let me write the direct convolution kernel with template parameters for tile sizes, making it easy to tune.
Here's my plan:
1. Write a CUDA kernel with TILE_H, TILE_W, TILE_OC as template parameters
2. Use shared memory for input and weights
3. Map blocks to (ow_block, oh_block, batch_and_oc_block)
4. Load input tile and weight tile cooperatively
5. Compute with all threads computing the same oc sequentially (broadcast weights)
6. Write output
Let me draft the kernel:
```cuda
template <int TILE_H, int TILE_W, int TILE_OC>
__global__ void conv2d_forward_kernel(
const float* __restrict__ input,
const float* __restrict__ weight,
float* __restrict__ output,
int N, int IC, int OC, int H, int W,
int K, int OH, int OW) {
// Block mapping
int ow_block = blockIdx.x;
int oh_block = blockIdx.y;
int n = blockIdx.z / (OC / TILE_OC);
int oc_block = blockIdx.z % (OC / TILE_OC);
int ow_base = ow_block * TILE_W;
int oh_base = oh_block * TILE_H;
int oc_base = oc_block * TILE_OC;
int tx = threadIdx.x;
int ty = threadIdx.y;
int tid = ty * blockDim.x + tx;
int nthreads = blockDim.x * blockDim.y;
// Shared memory
// Layout: s_input[IC][tile_h][tile_w]
// s_weight[TILE_OC][IC][K][K]
extern __shared__ float s_mem[];
float* s_input = s_mem;
float* s_weight = s_mem + IC * (TILE_H + K - 1) * (TILE_W + K - 1);
// Load input tile
int input_tile_h = TILE_H + K - 1;
int input_tile_w = TILE_W + K - 1;
int input_tile_size = IC * input_tile_h * input_tile_w;
for (int idx = tid; idx < input_tile_size; idx += nthreads) {
int ic = idx / (input_tile_h * input_tile_w);
int rem = idx % (input_tile_h * input_tile_w);
int ih_local = rem / input_tile_w;
int iw_local = rem % input_tile_w;
int ih = oh_base + ih_local;
int iw = ow_base + iw_local;
// Load from global memory
float val = 0.0f;
if (n < N && ic < IC && ih < H && iw < W) {
val = input[((n * IC + ic) * H + ih) * W + iw];
}
s_input[idx] = val;
}
// Load weights
int weight_size = TILE_OC * IC * K * K;
for (int idx = tid; idx < weight_size; idx += nthreads) {
int oc = idx / (IC * K * K);
int rem = idx % (IC * K * K);
int ic = rem / (K * K);
rem = rem % (K * K);
int kh = rem / K;
int kw = rem % K;
float val = 0.0f;
if (oc_base + oc < OC) {
val = weight[((oc_base + oc) * IC + ic) * K * K + kh * K + kw];
}
s_weight[idx] = val;
}
__syncthreads();
// Compute
float accum[TILE_OC];
#pragma unroll
for (int i = 0; i < TILE_OC; i++) {
accum[i] = 0.0f;
}
int oh = oh_base + ty;
int ow = ow_base + tx;
if (oh < OH && ow < OW) {
#pragma unroll
for (int kh = 0; kh < K; kh++) {
#pragma unroll
for (int kw = 0; kw < K; kw++) {
int ih_local = ty + kh;
int iw_local = tx + kw;
#pragma unroll
for (int ic = 0; ic < IC; ic++) {
float inp = s_input[(ic * input_tile_h + ih_local) * input_tile_w + iw_local];
#pragma unroll
for (int oc = 0; oc < TILE_OC; oc++) {
float w = s_weight[((oc * IC + ic) * K + kh) * K + kw];
accum[oc] += w * inp;
}
}
}
}
}
// Store output
if (oh < OH && ow < OW) {
#pragma unroll
for (int oc = 0; oc < TILE_OC; oc++) {
output[(((n * OC + oc_base + oc) * OH) + oh) * OW + ow] = accum[oc];
}
}
}
```
Wait, there are several issues:
1. **Shared memory bank conflicts**: For `s_input`, if layout is [ic][h][w] and we access with fixed ic, h_local, varying w_local, then adjacent threads (varying tx) read adjacent w values. In shared memory, these are contiguous and map to different banks (assuming 32 banks, each 4 bytes). So bank 0 gets w=0, bank 1 gets w=1, etc. No conflicts!
But for `s_weight`, all threads read the same weight for a given (oc, ic, kh, kw) - that's a broadcast, also no conflicts. But different iterations of the `oc` loop read different addresses. This sequential access by all threads is also a broadcast for each oc.
Wait, but we're iterating `oc` in the inner loop. For a given (ic, kh, kw), all threads read `w[oc=0]`, then `w[oc=1]`, etc. Each read is a broadcast. But we're reading 64 values sequentially. Since all threads read the same address, no bank conflicts.
Actually, there is a problem: in the `for (int oc = 0; oc < TILE_OC; oc++)` loop inside `for (int ic = 0; ic < IC; ic++)`, all threads are reading the same weight values but at different times. This is fine - it's sequential broadcast, no conflicts.
2. **Loop order**: The inner loops are kh, kw, ic, oc. For fixed ic, kh, kw, we read one input and iterate over all oc weights. This is good because the input value stays in a register, and weights are read from shared memory (broadcast).
Wait, actually if we restructure slightly, the input register can be reused across all oc:
```cpp
for kh {
for kw {
for ic {
float in = s_input[...];
for oc = 0..TILE_OC {
accum[oc] += in * s_weight[oc][ic][kh][kw];
}
}
}
}
```
This keeps `in` in a register and reuses it for all output channels. Much better!
3. **Register pressure**: With TILE_OC=64, `accum[64]` uses 64 registers. Plus `in`, plus loop variables. This is maybe ~70 registers. Under 255, good.
4. **Output writes**: We write TILE_OC values per thread. If TILE_OC=64, that's 64 stores. But output is only 8.5GB total, and peak write bandwidth is ~1.675 TB/s (half of total), so memory is not the bottleneck.
Let me refine the kernel with the loop reordering.
Also, I should add vectorized loads where possible. For input loading, since IC * tile_h * tile_w = 16 * 18 * 18 = 5184, we could load using float4 if the indexing allows. But since the test is general (IC might not be multiple of 4 in general), maybe float4 is tricky. For our test, IC=16, so we COULD use float4 for channels, but that's an optimization for later.
Let me also check: can we make TILE_OC a template parameter and instantiate it with 64? Yes.
Actually, there might be an issue with `float accum[TILE_OC]` when TILE_OC is large. CUDA allows up to 256 bytes of local memory per thread without spilling if compiler decides, but with 64 floats = 256 bytes, this fits in registers if compiler is smart. Actually 64 floats = 256 bytes. H100 registers are 4 bytes each. 64 registers is fine.
But with `accum` as local array, the compiler should put it in registers if there's no dynamic indexing. Here we index with constant/unrolled `oc`, so it should be in registers.
Wait, I used `#pragma unroll` on the `oc` loop, but dynamic array indexing inside an unrolled loop is resolved at compile time. So `accum[oc]` for each unrolled iteration refers to a specific register. Good.
Let me also think about writing output. If TILE_OC=64, writing output sequentially does 64 stores. Could we use float4 for output if we rearrange? No, output for different oc of the same pixel are not contiguous.
Actually, we could use warp shuffles to reorder and write float4, but that's probably overkill.
One issue with large TILE_OC: Each block only computes one spatial tile, so for oh=1022, ow=1022, with TILE_H=16, TILE_W=16:
- Tile count = ceil(1022/16) = 64 in each dimension
- That's 64*64 = 4096 spatial tiles
- N=16, OC=128, TILE_OC=64 → 2 oc tiles
- Total blocks = 16 * 4096 * 2 = 131,072 blocks
Actually that's only 131K blocks. For 132 SMs with shared mem limiting to maybe 3 blocks each = 396 concurrent blocks. 131K / 396 = 331 rounds. If each round takes ~15 microseconds, that's ~5 ms total. But this is very rough.
For comparison, the reference might take ~8-12 ms. So we have a chance!
But PyTorch's cuDNN might already be highly optimized for this. To beat it, we need to be very efficient.
Let me think about another optimization: Since H=1024, W=1024, OH=1022, OW=1022, we have many threads doing nothing in the last tile of each row/column (since 1022 is not divisible by 16, the last tile has some invalid positions).
Actually, 1022 / 16 = 63.875, so 64 tiles. The last tile has positions 1008 to 1023, but valid outputs are only up to 1021. So 1024 - 1022 = 2 rows/cols are invalid in the last tiles. That's a small inefficiency (~3%).
Now, the input loading: for the last tiles, we load pixels up to 1023 (which is valid since H=W=1024), but only use them for outputs up to 1021. The extra 2 rows/cols are loaded but not all used. Small inefficiency.
Let me write this up. But I want to also consider if we should tune TILE_H, TILE_W, TILE_OC. Let me make a few versions or start with one and iterate.
I'll start with TILE_H=16, TILE_W=16, TILE_OC=64.
One concern: 131K blocks might not fully saturate the GPU. On H100, having many concurrent thread blocks is important. Let me see if I can increase the number of blocks while keeping good work per thread.
What if TILE_OC=32 and TILE_H=8, TILE_W=8?
- Blocks = 16 * 4 * 128 * 128 = 1,048,576 blocks. Lots of blocks!
- Input shared = 10*10*16*4 = 6.4 KB
- Weight shared = 32*16*9*4 = 18.4 KB
- Total = 25 KB. Can fit 8 blocks per SM.
- Each thread: 32 * 144 = 4608 FMAs.
This has more blocks and higher occupancy. But each block has only 64 threads, meaning only 2 warps. A block with 2 warps might not be enough to keep an SM busy, but with 8 such blocks per SM, that's 16 warps/SM. Not great.
If TILE_H=8, TILE_W=16:
- Block = 128 threads = 4 warps
- Input shared = 10*18*16*4 = 11.5 KB
- Weight shared = 32*16*9*4 = 18.4 KB
- Total = 30 KB
- Blocks = 16 * 4 * 128 * 64 = 524,288 blocks
- Work per thread = 4608 FMAs
Blocks per SM at 30KB = ~7, giving 28 warps/SM. This is decent.
Actually, I think a healthy balance is important. Let me try TILE_H=16, TILE_W=16, TILE_OC=32 or 64.
Wait, I can also make block smaller but have each thread compute multiple spatial pixels. Like block=8x8=64 threads, each thread computes 2x2 pixels for 32 output channels.
Let me try:
- Block = 8x8 = 64 threads
- Each thread computes 2x2 spatial pixels for 32 output channels = 128 accumulators
- Spatial tile = 16x16
- Input shared = 18*18*16*4 = 20.7 KB
- Weight shared = 32*16*9*4 = 18.4 KB
- Total = 39.1 KB
Blocks = 16 * 4 * 64 * 64 = 262,144 blocks
Each thread FMAs = 128 * 144 = 18,432 FMAs
This gives lots of ILP (128 accumulators), many blocks, and reasonable shared memory. But 64 threads = 2 warps per block might have scheduling inefficiency.
What about block=16x8=128 threads, each thread computes 2x1 spatial pixels (2 along H) for 16 output channels = 32 accumulators.
- Spatial tile = 32x8
- Input shared = 34*10*16*4 = 21.8 KB
- Weight shared = 16*16*9*4 = 9.2 KB
- Total = 31 KB
- Blocks = 16 * 8 * 32 * 128 = 524,288 blocks
- Each thread FMAs = 32 * 144 = 4608 FMAs
With 128 threads = 4 warps per block. Blocks per SM ~7. 28 warps/SM.
Actually, I think a key insight is: for H100, having at least 4 warps per block and many blocks is good. And having enough ILP to saturate FMA units.
Let me go with a safe first attempt: block=16x16=256 threads, TILE_H=16, TILE_W=16, TILE_OC=64. Each thread has 64 accumulators and does 64*144=9216 FMAs. 131K blocks is sufficient.
Wait, 131K blocks for 132 SMs = ~1000 blocks per SM. With maybe 2 blocks concurrently per SM (limited by shared memory), that's 500 rounds. But each block does a LOT of compute (9216 FMAs per thread * 256 threads = 2.36M FMAs). With H100 FMA throughput, 2.36M FMAs should take maybe 20-30 microseconds? Then 500 rounds * 25 us = 12.5 ms. Hmm, that's starting to look slow.
Wait, I think I'm double counting. If 2 blocks run concurrently on an SM, they share the FMA units. The throughput is per SM, not per block. So 2 blocks on the same SM take longer than 1 block alone, but the total throughput of the SM is fixed.
If each SM can do X FMAs per cycle, and we have 132 SMs, then total FMAs / total throughput = time. The number of rounds doesn't directly multiply block time.
Actually, 132 SMs * 2 blocks/SM = 264 blocks running at once. 131K blocks / 264 = 496 "waves". Each wave takes time T_wave. Total time = 496 * T_wave.
T_wave is the time for a block to finish given shared SM resources. With 256 threads and 64 accumulators, each thread has high ILP, so the block should be compute-bound and take time proportional to FMAs / throughput.
If an SM has 128 FP32 FMA units running at 2GHz, that's 256 GFMA/s per SM = 256B FMAs per second.
Total FMAs per block per SM = 2 blocks * 256 threads * 9216 FMAs = 4.7M FMAs.
Time = 4.7M / 256B/s = 18.4 microseconds.
Total time = 496 * 18.4us = 9.1 ms. Very close to theoretical compute limit!
But that's assuming perfect utilization and no memory overhead. Realistically maybe 70% = 13 ms. And cuDNN might already achieve ~10-12 ms.
Hmm, this might be close. Let me check if there's more parallelism we can exploit.
Actually, if we do TILE_OC=32, we get 262K blocks and maybe 4 blocks per SM (shared mem ~30KB each, SM has 227KB = ~7 blocks, but other resources limit to 4). Then 262K / (132*4) = 496 waves again. Same number of waves. But each wave has less work: 256 threads * 4608 FMAs * 4 blocks/SM = 4.7M FMAs per SM per wave. Same compute! So same time.
Wait, I see. The number of waves is roughly the same because total work / SM capacity is the same. But with more blocks per SM, latency hiding is better.
But with TILE_OC=32 and 4 blocks/SM, if total FMAs per wave per SM is the same, and wave count is the same, total time is the same. But if latency hiding is better with 4 blocks vs 2, the actual time per wave might be shorter.
Actually, the real question is: what's the FMA throughput utilization? With 64 independent accumulators (TILE_OC=64) vs 32 (TILE_OC=32), but 2 blocks vs 4 blocks, which gives better utilization?
With 2 blocks * 256 threads * 64 accumulators = 32,768 independent FMAs in flight.
With 4 blocks * 256 threads * 32 accumulators = 32,768 independent FMAs in flight.
Same ILP!
So I should pick the configuration that minimizes overhead (loads, synchronizations, etc.) and maximizes actual throughput.
Let me think about this differently. What if each thread computes 2x2 spatial pixels for 16 output channels?
- 4 * 16 = 64 accumulators
- Block = 16x16, Tile = 32x32 spatial, Tile_OC = 16
- Input shared = 34*34*16*4 = 73.7 KB
- Weight shared = 16*16*9*4 = 9.2 KB
- Total = 82.9 KB
- Blocks per SM = 2 (227/83 = 2.7)
- Blocks = 16 * 8 * 32 * 32 = 131,072
- FMAs per thread = 64 * 144 = 9216
Independent FMAs in flight = 2 blocks * 256 threads * 64 accumulators = 32,768. Same.
But now each block handles 32x32 = 1024 spatial pixels. Total spatial tiles = (1022/32)^2 = 1024. With 16 batches and 8 oc tiles, total 131K blocks. And each block loads 34*34*16 = 18,496 input values but computes 1024*16 = 16,384 output values. The input loading overhead is amortized over more outputs.
Wait, loading 73KB vs computing 1024*16 outputs with 144 MACs each. It's actually the same compute as before but with a lot more input loaded per block.
For 32x32 spatial, the block boundary issue is: last tile has 1022 mod 32 = 30 valid outputs (out of 32). So ~6% waste. Not bad.
With 73KB of shared memory, fitting 2 blocks per SM = 2 * 256 * 64 accumulators = 32,768 in flight. This seems like a good sweet spot.
But actually, loading input for 34x34 region when we only compute 32x32 outputs (per channel) gives a "useful compute" to "loaded data" ratio of:
Outputs: 1024 * 16 * 144 = 2.36M MACs
Loaded: 34*34*16 + 16*16*9 = 18,496 + 2,304 = 20,800 values (not all used for edge tiles)
For TILE_H=16, TILE_W=16, TILE_OC=64:
Outputs: 256 * 64 * 144 = 2.36M MACs
Loaded: 18*18*16 + 64*16*9 = 5,184 + 9,216 = 14,400 values
The 32x32x16 tile loads more input relative to output (because the halo is 34x34 vs 32x32 = 12.5% overhead, while 18x18 vs 16x16 = 26% overhead). Wait, no:
Effective input region per output:
- For 32x32 output with K=3: input is 34x34. Ratio = 34*34 / (32*32) = 1.128
- For 16x16 output with K=3: input is 18x18. Ratio = 18*18 / (16*16) = 1.266
So the 32x32 tile is actually MORE efficient in terms of input reuse! The overhead of loading the halo is smaller.
And weight loading:
- 32x32x16 loads 16*16*9 = 2,304 weights for 1024*16 = 16,384 outputs. Ratio = 0.14
- 16x16x64 loads 64*16*9 = 9,216 weights for 256*64 = 16,384 outputs. Ratio = 0.56
The 32x32 tile uses weights much more efficiently! This is because the larger spatial tile amortizes the weight loading cost.
This is an important insight. We want LARGER spatial tiles to better amortize weight loading.
With block=16x16 threads and each thread computing 2x2 pixels:
- We get 32x32 spatial tile
- 256 threads * 4 pixels = 1024 pixels total
- For 16 output channels, that's 16,384 outputs
- FMAs = 16,384 * 144 = 2.36M per block
If we do 32 output channels with 2x2 pixels:
- Block computes 32x32x32 = 32,768 outputs
- But input shared = 34*34*16*4 = 73.7 KB
- Weight shared = 32*16*9*4 = 18.4 KB
- Total = 92.1 KB
- Blocks per SM = 2 (limited by shared memory)
- FMAs per block = 32,768 * 144 = 4.72M
Blocks = 16 * 4 * 32 * 32 = 65,536
Total blocks / concurrent blocks = 65,536 / (132*2) = 248 waves.
FMAs per wave per SM = 2 * 4.72M = 9.44M
Time per wave = 9.44M / (SM throughput)
Total time = 248 * time_per_wave.
This might be good since there are fewer waves, and each wave has more compute.
Actually, let me look at this differently. What limits the minimum time?
- Total FMAs = 308B
- Peak FMA throughput = 33.5B FMA/s (at 67 TFLOPS)
- Theoretical min = 9.2 ms
Can any configuration beat this? No, it's the absolute limit for FMA throughput.
Now, with the weight reuse consideration, can we get closer to 9.2 ms?
If each block spends some time loading shared memory, that's overhead. With larger spatial tiles, the compute time increases while the load time doesn't increase as much.
For example, 32x32x32 block:
- Load 73KB input + 18KB weight = 91KB from global to shared
- Time to load 91KB at L2 bandwidth (say 1TB/s) = 91ns. But with 256 threads and coalescing, maybe ~2 microseconds.
- Compute 4.72M FMAs at 67 TFLOPS equivalent = ~141 microseconds per block for the whole GPU, but per SM it's different.
If each SM does 2 blocks * 4.72M FMAs = 9.44M FMAs. At 67 TFLOPS / 132 SMs = 508 GFLOPS/SM = 254 GFMA/s per SM. Time = 9.44M / 254B = 37 microseconds.
Load time ~2 microseconds is small compared to compute time. So compute dominates. Good.
For 16x16x64 block:
- Load 20KB input + 36KB weight = 56KB
- Time ~1.2 microseconds
- Compute 2.36M FMAs. Per SM with 2 blocks = 4.72M FMAs. Same compute, same time.
- But load overhead is proportionally higher: 1.2us / 37us = 3.2% vs 2us / 37us = 5.4%? Wait no, 2us / 37us = 5.4%, 1.2us / 37us = 3.2%. Actually smaller spatial has less load overhead in absolute terms.
Hmm, but smaller spatial means more waves, and each wave has more overhead from kernel launch/scheduling.
Let me just go with block=16x16 threads, each computing 2x2 spatial for 16 output channels. This gives:
- Spatial tile: 32x32
- Output channels per block: 16
- Input shared: 34*34*16*4 = 73.7 KB
- Weight shared: 16*16*9*4 = 9.2 KB
- Total: 82.9 KB
- Can fit 2 blocks per SM
Actually, wait. I want to compute 2x2 spatial, but block is 16x16 threads. That means:
- Thread (ty, tx) computes outputs at (oh_base + ty*2, ow_base + tx*2), (oh_base + ty*2, ow_base + tx*2 + 1), (oh_base + ty*2 + 1, ow_base + tx*2), (oh_base + ty*2 + 1, ow_base + tx*2 + 1)
- But only if these positions are within OH and OW.
Each thread needs to read input at 4 different positions for each (kh, kw):
- For output (h, w), input is at (h+kh, w+kw)
- For 4 outputs, we need inputs at:
- (ty*2+kh, tx*2+kw)
- (ty*2+kh, tx*2+1+kw)
- (ty*2+1+kh, tx*2+kw)
- (ty*2+1+kh, tx*2+1+kw)
For a given (kh, kw), these 4 input positions form a 2x2 region in the input tile. We can load them individually or use vectorized loads. But using 4 separate loads from shared memory is fine since shared memory has very high bandwidth.
Let me write this kernel carefully.
Actually, another idea: What if each thread computes a 4x1 strip (4 pixels along W) for 8 output channels?
- Block = 16x16
- Tile = 16x64 spatial, 8 OC
- Input shared = 18x66x16 = 19,008 floats = 76 KB. Too much.
Or block = 8x16, each thread computes 2x2 spatial for 8 OC:
- Tile = 16x32 spatial, 8 OC
- Input shared = 18x34x16 = 9,792 floats = 39 KB
- Weight shared = 8*16*9*4 = 4.6 KB
- Total = 43.6 KB
- FMAs per thread = 4 * 8 * 144 = 4608
- Blocks = 16 * 16 * 64 * 32 = 524,288
This gives many blocks and reasonable shared memory!
Let me go with: block = 16x16, each thread computes 1 pixel for 32 or 64 output channels. It's the simplest and most likely to compile correctly.
Actually let me also consider: can we do a 1D convolution kernel? No, 2D conv needs both dimensions.
Alright, I'll write the kernel with these parameters:
- TILE_H = 16
- TILE_W = 16
- TILE_OC = 64
- blockDim = (16, 16)
- Each thread computes 1 pixel for 64 output channels
And use template specialization for K=3 to help compiler unroll.
Wait, I just want to check if there's an even simpler approach that might already be close to optimal. Could we just use the built-in implicit GEMM approach but avoid workspace allocation by writing our own?
One trick: For K=3, IC=16, we can manually "unfold" the input into registers and then do a matrix multiply.
Actually, here's an even simpler kernel idea:
Instead of using shared memory for the entire input tile, what if we stream the computation? Each warp computes a set of output pixels, and for each (kh, kw), we load the required input values and broadcast-multiply with weights.
Specifically:
- A warp of 32 threads computes 32 output pixels (spanned across spatial positions)
- For each output channel group (say 32 OCs), for each (ic, kh, kw):
- Load 32 input values (one per thread, coalesced)
- Load 32 weights (broadcast)
- FMA
This is very similar to what cuDNN likely does for small K implicit GEMM.
But without tensor cores, it's hard to beat cuDNN.
Let me just implement the shared-memory version and tune iteratively.
Let me refine the shared memory kernel one more time:
```cuda
template<int TILE_H, int TILE_W, int TILE_OC>
__global__ void conv2d_kernel(const float* __restrict__ input,
const float* __restrict__ weight,
float* __restrict__ output,
int N, int IC, int OC, int H, int W, int K, int OH, int OW) {
const int in_tile_h = TILE_H + K - 1;
const int in_tile_w = TILE_W + K - 1;
int n = blockIdx.z / (OC / TILE_OC);
int oc_block = blockIdx.z % (OC / TILE_OC);
int oh_base = blockIdx.y * TILE_H;
int ow_base = blockIdx.x * TILE_W;
int oc_base = oc_block * TILE_OC;
int tx = threadIdx.x;
int ty = threadIdx.y;
extern __shared__ float s_mem[];
float* s_input = s_mem;
float* s_weight = s_mem + IC * in_tile_h * in_tile_w;
// Cooperative loading of input tile
int input_size = IC * in_tile_h * in_tile_w;
int tid = ty * blockDim.x + tx;
int nthreads = blockDim.x * blockDim.y;
for (int idx = tid; idx < input_size; idx += nthreads) {
int ic = idx / (in_tile_h * in_tile_w);
int rem = idx % (in_tile_h * in_tile_w);
int ih = oh_base + rem / in_tile_w;
int iw = ow_base + rem % in_tile_w;
float val = 0.0f;
if (n < N && ic < IC && ih < H && iw < W) {
val = input[((n * IC + ic) * H + ih) * W + iw];
}
s_input[idx] = val;
}
// Cooperative loading of weights
int weight_size = TILE_OC * IC * K * K;
for (int idx = tid; idx < weight_size; idx += nthreads) {
int idx2 = idx;
int oc = idx2 / (IC * K * K);
idx2 %= IC * K * K;
int ic = idx2 / (K * K);
idx2 %= K * K;
int kh = idx2 / K;
int kw = idx2 % K;
float val = 0.0f;
if (oc_base + oc < OC) {
val = weight[(((oc_base + oc) * IC + ic) * K + kh) * K + kw];
}
s_weight[idx] = val;
}
__syncthreads();
int oh = oh_base + ty;
int ow = ow_base + tx;
float accum[TILE_OC];
#pragma unroll
for (int i = 0; i < TILE_OC; i++) {
accum[i] = 0.0f;
}
if (oh < OH && ow < OW) {
#pragma unroll
for (int kh = 0; kh < K; kh++) {
#pragma unroll
for (int kw = 0; kw < K; kw++) {
int ih = ty + kh;
int iw = tx + kw;
float inp = s_input[(ih * in_tile_w + iw) * IC + 0]; // Wait, layout?
```
Here's a key design decision: How to lay out shared memory for input?
Option 1: [IC][in_tile_h][in_tile_w] (channel-major)
Option 2: [in_tile_h][in_tile_w][IC] (spatial-major)
For our access pattern: For a given spatial position (ih, iw), we iterate over all IC channels. So spatial-major gives contiguous IC values. This means for a fixed (ih, iw), channels are contiguous.
But in the computation, we do:
```
for kh:
for kw:
for ic:
in = s_input[...ic...]
for oc:
accum[oc] += in * weight[oc][ic][kh][kw]
```
Here, for fixed (kh, kw), we iterate over ic. If layout is [in_tile_h][in_tile_w][IC], then for fixed (ih, iw), ic is contiguous. Good for reading sequential values. But different threads read different (ih, iw), so shared memory access is coalesced across threads.
Actually, wait. For fixed (kh, kw), thread (ty, tx) reads input at (ty+kh, tx+kw). So all threads in a block read different (ih, iw) positions. If layout is [h][w][IC], then thread accesses are at different h*w positions, strided by IC. If IC=16, then thread 0 reads at offset0, thread1 at offset16, thread2 at offset32, etc. These map to different banks (since 16 floats = 64 bytes = bank 0, bank 16, bank 0, bank 16...). Wait, 16 floats * 4 bytes = 64 bytes. Shared memory has 32 banks, each 4 bytes. So offset 0 = bank 0, offset 64 = bank 0. So threads 0 and 2 would both access bank 0!
This is a bank conflict!
To avoid bank conflicts with channel-major layout [IC][h][w]:
Thread (ty, tx) reads at ic * tile_h * tile_w + (ty+kh) * tile_w + (tx+kw).
For fixed ic, kh, kw: thread accesses are at (ty*tile_w + tx). With tile_w=16, thread 0 (0,0) at 0, thread 1 (0,1) at 1, etc. Adjacent threads access adjacent banks. No conflict! This is ideal.
But for iterating over ic, we go from ic=0 to ic=15. Each ic is offset by tile_h*tile_w*4 bytes = 18*18*4 = 1296 bytes. 1296 / 128 = 10.125? Let me recalculate.
Actually, bank = (address_in_bytes / 4) % 32.
For [ic][h][w] layout, address = (ic * tile_h * tile_w + h * tile_w + w) * 4.
bank = (ic * tile_h * tile_w + h * tile_w + w) % 32.
For fixed h, w and varying ic: bank changes by tile_h * tile_w mod 32 = 18*18 mod 32 = 324 mod 32 = 4.
So consecutive ic values are in banks 4 apart (0, 4, 8, 12, 16, 20, 24, 28, 0...). No conflicts for consecutive ic! Good.
For different threads at fixed ic, h=ty+kh, w=tx+kw:
bank difference between thread (ty, tx) and (ty, tx+1) = 1. Adjacent banks, no conflict.
This layout [ic][h][w] seems excellent.
What about [h][w][ic] layout?
address = (h * tile_w * IC + w * IC + ic) * 4.
For fixed h, w and adjacent threads (tx, tx+1): address difference = IC*4 = 64 bytes = bank difference = 16.
Wait, no: adjacent tx means w differs by 1. address difference = IC * 4 = 64 bytes. bank difference = 16. So thread 0 is bank 0, thread 1 is bank 16, thread 2 is bank 0 (since IC=16, 2*IC=32 ≡ 0 mod 32). So threads 0 and 2 conflict!
This is BAD for [h][w][ic] layout with IC=16.
So [ic][h][w] is definitely the right layout.
Now for weights: [oc][ic][kh][kw] or [oc][kh][kw][ic].
We access weights as: for fixed oc_block, all threads iterate over oc=0..TILE_OC-1, and within each oc, iterate over ic, kh, kw.
If layout is [oc][ic][kh][kw], then for fixed oc, ic, kh, kw, all threads read the same address = broadcast, no conflict.
For fixed kh, kw, ic, and iterating over oc: addresses differ by IC*K*K elements. IC*K*K = 16*9 = 144. 144 mod 32 = 16. So consecutive oc are 16 banks apart. For TILE_OC=64, we get oc=0 at bank X, oc=1 at bank X+16, oc=2 at bank X, oc=3 at bank X+16. For a warp reading these sequentially, half the threads in warp read same bank at a time? Actually all threads read the same address for a given oc. Then for next oc, all threads read the next address. No conflict since it's broadcast each time, just with a stride.
Actually, since ALL threads read the SAME address for each (oc, ic, kh, kw), it's always a broadcast. No bank conflicts regardless of layout!
But for loading weights from global memory to shared memory, we want coalescing. For [oc][ic][kh][kw], adjacent threads in the loading loop read adjacent memory if they load adjacent (oc, ic, kh, kw). Weight in global memory is stored as [OC][IC][K][K]. So loading sequentially is naturally coalesced.
OK, so layout [ic][h][w] for input and [oc][ic][kh][kw] for weights.
Let me rewrite the kernel with these layouts.
For input:
```
s_input[ic * in_tile_h * in_tile_w + ih * in_tile_w + iw]
```
For weights:
```
s_weight[oc * IC * K * K + ic * K * K + kh * K + kw]
```
And in the compute loop:
```
for (int kh = 0; kh < K; kh++) {
for (int kw = 0; kw < K; kw++) {
int ih = ty + kh;
int iw = tx + kw;
for (int ic = 0; ic < IC; ic++) {
float in = s_input[ic * in_tile_h * in_tile_w + ih * in_tile_w + iw];
for (int oc = 0; oc < TILE_OC; oc++) {
float w = s_weight[oc * IC * K * K + ic * K * K + kh * K + kw];
accum[oc] += w * in;
}
}
}
}
```
Wait, I should move `ic` to the outermost of the inner loops to maximize `in` reuse:
```
for (int kh = 0; kh < K; kh++) {
for (int kw = 0; kw < K; kw++) {
int ih = ty + kh;
int iw = tx + kw;
for (int ic = 0; ic < IC; ic++) {
float in = s_input[ic * in_tile_h * in_tile_w + ih * in_tile_w + iw];
#pragma unroll
for (int oc = 0; oc < TILE_OC; oc++) {
float w = s_weight[oc * IC * K * K + ic * K * K + kh * K + kw];
accum[oc] += w * in;
}
}
}
}
```
Actually, a better loop order for register reuse of weights? No, weights for different oc are different, so they can't be reused in a register. But `in` is reused for all oc. So keeping `in` in a register across the oc loop is optimal.
Can we swap the ic and kh, kw loops for better memory locality?
If we iterate ic outermost:
```
for (int ic = 0; ic < IC; ic++) {
for (int kh = 0; kh < K; kh++) {
for (int kw = 0; kw < K; kw++) {
float in = s_input[...];
for (oc...) ...
}
}
}
```
For a given ic, we touch 9 spatial locations. With [ic][h][w] layout, these 9 locations are close together (within a 3x3 window). But they span multiple rows. The locality depends on the cache structure of shared memory, which is divided into banks but otherwise linear.
Actually, for kh and kw as inner loops, for each (ic, kh, kw), we load one float. The total working set for ic=0 is 3x3=9 floats = 36 bytes. This fits in a small number of cache lines. So locality is good either way.
Let me keep kh, kw outer to exploit the possibility of unrolling K=3.
Actually, since K=3 is small, `#pragma unroll` on kh and kw should unroll the 9 iterations. This gives us:
```
#pragma unroll
for (kh = 0; kh < 3; kh++) {
#pragma unroll
for (kw = 0; kw < 3; kw++) {
#pragma unroll
for (ic = 0; ic < 16; ic++) { // if IC=16 and we can specialize
float in = s_input[...];
#pragma unroll
for (oc = 0; oc < TILE_OC; oc++) {
accum[oc] += w[oc][ic][kh][kw] * in;
}
}
}
}
```
With K=3 and IC=16 fully unrolled, this is a straight line of 9*16*(TILE_OC+1) operations.
For TILE_OC=64, that's 144 * 65 = 9360 operations per thread (including loads and FMAs). All loads from shared memory. This is very efficient.
Wait, if IC=16 is unrolled, the compiler might inline the loads. But we pass IC as a runtime parameter. To force unrolling, we'd need either a template parameter or to hardcode 16.
For now, I'll let the compiler handle it with `#pragma unroll` and hope it unrolls small loops. But with a runtime variable, `#pragma unroll` without a number might not fully unroll IC.
I could add:
```
if (IC == 16 && K == 3) {
// Specialized code with hardcoded 16 and 3
} else {
// Generic code
}
```
Let me do this to help the compiler generate optimal code for the test parameters.
Actually, a cleaner approach: write the kernel so it works for any IC, K, but use `#pragma unroll` aggressively. For IC=16, the compiler might unroll it if it knows it's small. But the compiler only sees IC as a variable.
For template parameters, I can do:
```
template<int TILE_H, int TILE_W, int TILE_OC, int K_CONST, int IC_CONST>
__global__ void conv2d_kernel(...)
```
But we'd only instantiate when K_CONST == K and IC_CONST == IC. We can check at runtime and fall back to generic or instantiate common cases.
For simplicity in iteration, let me first write a generic kernel and see if it's fast enough. If not, I'll specialize.
Actually, I realize that `#pragma unroll` with a runtime loop bound might not work. The standard syntax is `#pragma unroll N` where N is a compile-time constant. But for a loop `for (int i = 0; i < K; i++)`, if K is passed as a parameter, `#pragma unroll` (without N) might fully unroll very small loops, but it's compiler-dependent.
To guarantee unrolling, I need to use templates or macros. Let me use a helper that the compiler can inline.
For now, I'll do:
```cpp
#pragma unroll
for (int kh = 0; kh < K; kh++) {
```
And hope the compiler unrolls it when K=3 is passed as a small constant. Actually, K is passed as an `int` parameter to the kernel. The compiler cannot know its value at compile time. So this WON'T unroll!
Ah, this is important. For the loop to unroll, the compiler needs to know K is 3 at compile time.
Solutions:
1. Template parameter: `template<int K>` and instantiate `conv2d_kernel<3>`. But we also need IC, so `template<int K, int IC>`.
2. Use a macro or specialized function.
Let me use template parameters for K and IC, and instantiate with <3, 16>:
```cpp
template<int TILE_H, int TILE_W, int TILE_OC, int K_VAL, int IC_VAL>
__global__ void conv2d_kernel(...) {
// Use K_VAL and IC_VAL in the loops
#pragma unroll
for (int kh = 0; kh < K_VAL; kh++) {
...
}
}
```
Then in host code:
```cpp
if (K == 3 && IC == 16) {
conv2d_kernel<16, 16, 64, 3, 16><<<grid, block, smem>>>(...);
}
```
This way, K_VAL and IC_VAL are compile-time constants, and the compiler can fully unroll these small loops.
Let me implement this.
But wait, I also need to handle general cases as fallback. The problem says the kernel should work for the model parameters, not necessarily arbitrary sizes. But to be safe, I can add a generic fallback path.
Actually, the task says "Replace the forward computation... must accept same constructor arguments and implement forward with same signature". So the kernel needs to handle the general case of any parameters. But our test is with specific parameters. I'll add a specialized path and a generic fallback.
For the host function, I'll have:
```cpp
torch::Tensor conv2d_forward(torch::Tensor input, torch::Tensor weight, int stride, int padding, int dilation, int groups) {
// ... compute dimensions ...
if (K == 3 && stride == 1 && padding == 0 && dilation == 1 && groups == 1) {
// Use specialized kernel
} else {
// Fall back to generic kernel (or we could call ATen, but the task says to write our own)
}
}
```
Wait, falling back to ATen might give correct results but wouldn't use our custom kernel. For the generic case, I should also write a generic CUDA kernel or just call at::conv2d. But the task says to write a custom CUDA kernel that replaces the reference. It's probably OK to fall back to at::conv2d for unsupported cases, as long as our kernel handles the common case.
Actually, the safest is to write a generic CUDA kernel that handles any valid parameters.
Alright, let me write a comprehensive kernel. I'll include:
1. A generic template kernel
2. Or just one kernel with `#pragma unroll` that uses template params and a generic fallback
Let me structure the Python file:
```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>
template<int K_VAL, int IC_VAL>
__global__ void conv2d_kernel(const float* __restrict__ input, ...) {
...
#pragma unroll
for (int kh = 0; kh < K_VAL; kh++) {
#pragma unroll
for (int kw = 0; kw < K_VAL; kw++) {
#pragma unroll
for (int ic = 0; ic < IC_VAL; ic++) {
...
}
}
}
}
// Generic kernel for arbitrary parameters
__global__ void conv2d_kernel_generic(const float* __restrict__ input, ...) {
// Same logic but without template params and unroll
}
torch::Tensor conv2d_forward(torch::Tensor input, torch::Tensor weight, int stride, int padding, int dilation, int groups) {
// ... compute output dimensions ...
// Dispatch to specialized or generic kernel
if (K == 3 && IC == 16) {
// Launch specialized kernel with TILE_H=16, TILE_W=16, TILE_OC=64
} else {
// Launch generic kernel
}
}
'''
class ModelNew(nn.Module):
def __init__(self, ...):
super().__init__()
self.weight = nn.Parameter(torch.randn(...))
if bias: self.bias = nn.Parameter(torch.zeros(...))
...
def forward(self, x):
return conv2d_forward(x, self.weight, self.stride, self.padding, self.dilation, self.groups)
```
Wait, I need to initialize the weight from the model! Since we're replacing nn.Conv2d, we need to match its weight initialization. But in the test code, the model is instantiated and then used. The `ModelNew` class might need to accept the same init pattern.
The test calls:
```python
model = ModelNew(*get_init_inputs())
```
And `get_init_inputs()` returns `[in_channels, out_channels, kernel_size]`.
So `ModelNew` receives these and should initialize its conv2d-like parameters. We need to create a weight parameter. In the reference `Model`, `nn.Conv2d` initializes weights using Kaiming uniform or similar. But for the test, they create a random input and just compare output. The model parameters are initialized to defaults.
Wait, the weights are initialized by nn.Conv2d automatically. For our ModelNew, we need to create the weight with the same initialization. But actually, in the actual test environment, the evaluator will instantiate Model and ModelNew with the same arguments and compare outputs. For the outputs to match, both need the same weights.
But Model is defined as `super(Model, self).__init__()` and creates `self.conv2d = nn.Conv2d(...)`. The evaluator might copy weights from Model to ModelNew? Or instantiate them independently?
Actually, from the problem statement: "Replace the forward computation... " and "produce numerically equivalent outputs". If both models are initialized independently, they might have different random weights, so outputs wouldn't match.
I suspect the evaluator does something like:
1. Instantiate Model
2. Instantiate ModelNew
3. Copy weights/buffers from Model to ModelNew (or vice versa)
4. Run both with the same input
5. Compare outputs
OR, they instantiate one, get its weights, and then set them on the other.
So our ModelNew needs to have compatible parameter names and shapes. Looking at the reference: `self.conv2d` contains parameters `weight` and `bias`.
For ModelNew, we should either:
1. Also have `self.conv2d = nn.Conv2d(...)` but override forward
2. Or create our own `self.weight` and `self.bias` parameters
If the evaluator copies state_dict, it might copy `conv2d.weight` and `conv2d.bias`. For compatibility, we could keep `self.conv2d` and just not use it in forward, or create parameters manually.
Actually, the safest is to create a dummy `nn.Conv2d` for parameter storage, and in forward, call our custom CUDA kernel using the parameters:
```python
class ModelNew(nn.Module):
def __init__(self, *args, **kwargs):
super().__init__()
self.conv2d = nn.Conv2d(*args, **kwargs)
# We could also store the parameters directly
self.in_channels = args[0]
self.out_channels = args[1]
...
def forward(self, x):
return conv2d_forward(x, self.conv2d.weight, self.conv2d.bias, self.conv2d.stride, ...)
```
This way, parameters are exactly compatible with `Model`. And in forward, we call our custom kernel with `self.conv2d.weight` and optionally `self.conv2d.bias`.
For bias, we need to add it after convolution. We can do this in the CUDA kernel or in PyTorch after. Adding in PyTorch is probably easier unless we want to fuse it.
Actually, fusing bias in the kernel is better for performance. But for simplicity and correctness, we can add it in PyTorch.
Let me write the code now.
One more thing: output padding and transposed conv are not in the args, so we don't need them. The args are: in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=False.
For computing output dimensions:
```
OH = (H + 2 * padding - dilation * (kernel_size - 1) - 1) / stride + 1
OW = (W + 2 * padding - dilation * (kernel_size - 1) - 1) / stride + 1
```
In C++ host code:
```cpp
int N = input.size(0);
int IC = input.size(1);
int H = input.size(2);
int W = input.size(3);
int OC = weight.size(0);
int K = weight.size(2);
int OH = (H + 2 * padding - dilation * (K - 1) - 1) / stride + 1;
int OW = (W + 2 * padding - dilation * (K - 1) - 1) / stride + 1;
```
Actually, the kernel_size from weight could be different, but is (kernel_size, kernel_size) in the model.
Now, let me think about the grid dimensions. I'll use:
- blockDim.x = TILE_W
- blockDim.y = TILE_H
- gridDim.x = (OW + TILE_W - 1) / TILE_W
- gridDim.y = (OH + TILE_H - 1) / TILE_H
- gridDim.z = N * ((OC + TILE_OC - 1) / TILE_OC)
This should give a clean mapping.
For the dispatch:
```cpp
dim3 block(TILE_W, TILE_H);
dim3 grid((OW + TILE_W - 1) / TILE_W, (OH + TILE_H - 1) / TILE_H, N * ((OC + TILE_OC - 1) / TILE_OC));
```
Note that gridDim.z can be large. For N=16, OC=128, TILE_OC=64: z=16*2=32. This is well within limits.
Wait, let me verify CUDA grid dimension limits. For modern CUDA:
- gridDim.x: up to 2^31-1
- gridDim.y: up to 2^16-1 = 65535
- gridDim.z: up to 2^16-1 = 65535
So z=32 is fine.
Now, one issue: the grid size calculation. With TILE_H=16, OH=1022:
(1022 + 15) / 16 = 1037 / 16 = 64 (integer division). Correct.
For the specialized kernel call, I'll need to pass all the parameters. Let me make sure the weight pointer is correct.
For torch tensors in C++: `input.data_ptr<float>()`, `weight.data_ptr<float>()`.
If no bias, pass `nullptr`. If bias, add it.
Actually, let me think about whether adding bias in CUDA or in PyTorch is better. In PyTorch:
```python
out = conv2d_forward(x, self.conv2d.weight, self.conv2d.stride, ...)
if self.conv2d.bias is not None:
out = out + self.conv2d.bias.view(1, -1, 1, 1)
```
This might launch an extra kernel for bias addition. For maximum performance, we could add bias in our CUDA kernel. But it's a minor overhead compared to conv itself. Let me add it to the kernel for cleanliness.
Actually, bias addition is trivial and might be negligible. Let me not complicate the kernel with bias and just do it in Python. But if the reference includes bias and we do extra kernel launches, it might add ~10-50 microseconds, which for a 5ms kernel is 0.1-1%. Not a big deal, but let's add it to the kernel.
In the kernel, after computing accum[oc], if bias is not nullptr:
```cpp
accum[oc] += bias[oc_base + oc];
```
Easy to add.
OK let me write the full Python file now. I'll include a bias pointer.
One more consideration: for the specialized kernel, should I also handle the case where OC is not divisible by TILE_OC? In the test, OC=128, TILE_OC divides evenly. But for generality, in the last oc_block, we should only compute valid OCs. Our kernel handles this because we guard against `oc_base + oc < OC` when loading weights (we load 0 for invalid weights), and the computation just adds 0. For output writing, we also guard.
Actually, our kernel writes all TILE_OC outputs. We check boundary conditions only when writing? Wait, in my draft, I check `if (oh < OH && ow < OW)` for computation and writing. But what about `oc_base + oc < OC`? In the specialized version, we load 0s for invalid weights, so accum will be 0 for invalid OCs. We should not write them.
In the generic version, the same applies. We need to guard output writes with `oc_base + oc < OC`.
Let me update:
```cpp
if (oh < OH && ow < OW && oc_base + oc < OC) {
output[...] = accum[oc];
}
```
Actually, since invalid weights are 0 and we might not want to write 0 to potentially existing output, we should guard. But output is freshly allocated, so writing 0 wouldn't be wrong. Still, let's guard to be safe.
Wait, if `oc_base + oc >= OC`, the weight was loaded as 0, so accum[oc] = 0. Writing 0 to output is correct since output should be 0 in those positions (no one else writes there). But if another block could write there, we'd have a race condition. Since we partition OC into non-overlapping oc_blocks, each output position is written by exactly one block. So writing 0 for invalid OCs is fine.
But it's cleaner to just guard. Let me guard both computation and writing.
For computation: if `oc_base + oc < OC`, use the real weight; else weight=0. We already loaded 0 for invalid weights, so computation is correct.
For writing: we should only write if `oh < OH && ow < OW && oc_base + oc < OC`.
OK, this is getting detailed. Let me write the kernel.
Actually I want to double check one thing. For the specialized kernel with K=3 and IC=16, could using constant memory for weights be beneficial?
Constant memory is cached and broadcasts well. The total weights are 73KB. On H100, constant memory size is 64KB per kernel. So 73KB won't fit entirely in constant memory. But if we process output channels in tiles, we could load each tile into constant memory? Actually, that's what shared memory does, and shared memory is faster than constant memory for this use case.
Shared memory bandwidth on H100 is TB/s per SM, much higher than constant memory. So shared memory is good.
Alright, here's my final kernel structure for the first attempt:
```cuda
template<int TILE_H, int TILE_W, int TILE_OC, int K_VAL, int IC_VAL>
__global__ void conv2d_kernel_specialized(
const float* __restrict__ input,
const float* __restrict__ weight,
const float* __restrict__ bias,
float* __restrict__ output,
int N, int OC, int H, int W, int OH, int OW) {
const int in_tile_h = TILE_H + K_VAL - 1;
const int in_tile_w = TILE_W + K_VAL - 1;
int n = blockIdx.z / ((OC + TILE_OC - 1) / TILE_OC);
int oc_block = blockIdx.z % ((OC + TILE_OC - 1) / TILE_OC);
int oh_base = blockIdx.y * TILE_H;
int ow_base = blockIdx.x * TILE_W;
int oc_base = oc_block * TILE_OC;
int tx = threadIdx.x;
int ty = threadIdx.y;
int tid = ty * blockDim.x + tx;
int nthreads = blockDim.x * blockDim.y;
__shared__ float s_input[IC_VAL * in_tile_h * in_tile_w];
__shared__ float s_weight[TILE_OC * IC_VAL * K_VAL * K_VAL];
// Load input
#pragma unroll
for (int idx = tid; idx < IC_VAL * in_tile_h * in_tile_w; idx += nthreads) {
int ic = idx / (in_tile_h * in_tile_w);
int rem = idx % (in_tile_h * in_tile_w);
int ih = oh_base + rem / in_tile_w;
int iw = ow_base + rem % in_tile_w;
float val = 0.0f;
if (n < N && ih < H && iw < W) {
val = input[((n * IC_VAL + ic) * H + ih) * W + iw];
}
s_input[idx] = val;
}
// Load weights
#pragma unroll
for (int idx = tid; idx < TILE_OC * IC_VAL * K_VAL * K_VAL; idx += nthreads) {
int oc = idx / (IC_VAL * K_VAL * K_VAL);
int rem = idx % (IC_VAL * K_VAL * K_VAL);
int ic = rem / (K_VAL * K_VAL);
rem = rem % (K_VAL * K_VAL);
int kh = rem / K_VAL;
int kw = rem % K_VAL;
float val = 0.0f;
if (oc_base + oc < OC) {
val = weight[(((oc_base + oc) * IC_VAL + ic) * K_VAL + kh) * K_VAL + kw];
}
s_weight[idx] = val;
}
__syncthreads();
int oh = oh_base + ty;
int ow = ow_base + tx;
float accum[TILE_OC];
#pragma unroll
for (int i = 0; i < TILE_OC; i++) {
accum[i] = 0.0f;
}
if (oh < OH && ow < OW) {
#pragma unroll
for (int kh = 0; kh < K_VAL; kh++) {
#pragma unroll
for (int kw = 0; kw < K_VAL; kw++) {
int ih = ty + kh;
int iw = tx + kw;
#pragma unroll
for (int ic = 0; ic < IC_VAL; ic++) {
float in_val = s_input[ic * in_tile_h * in_tile_w + ih * in_tile_w + iw];
#pragma unroll
for (int oc = 0; oc < TILE_OC; oc++) {
float w_val = s_weight[oc * IC_VAL * K_VAL * K_VAL + ic * K_VAL * K_VAL + kh * K_VAL + kw];
accum[oc] += w_val * in_val;
}
}
}
}
}
// Store output
if (oh < OH && ow < OW) {
#pragma unroll
for (int oc = 0; oc < TILE_OC; oc++) {
if (oc_base + oc < OC) {
float val = accum[oc];
if (bias != nullptr) {
val += bias[oc_base + oc];
}
output[(((n * OC + oc_base + oc) * OH + oh) * OW + ow)] = val;
}
}
}
}
```
Wait, there's a problem with this: the `#pragma unroll` on the load loops with `idx += nthreads` won't fully unroll because the trip count depends on `nthreads` but the loop has a runtime exit condition. The compiler might not unroll a loop with a variable upper bound unless it can figure out that the loop body is small and trip count is bounded.
Actually, for the loading loops, unrolling isn't critical. The main compute loop should be unrolled. But the compute loop doesn't have variable trip count - `K_VAL` and `IC_VAL` are compile-time constants, so `#pragma unroll` should work.
One concern: `s_weight` is declared as `__shared__ float s_weight[TILE_OC * IC_VAL * K_VAL * K_VAL];` with template parameters. For TILE_OC=64, IC_VAL=16, K_VAL=3, that's 64*16*9 = 9216 floats = 36.8KB. Combined with `s_input` of 18*18*16*4 = 20.7KB, total = 57.5KB.
We cannot declare two `__shared__` arrays with sizes from template parameters in a way that sums them? Actually we can, as long as total shared memory per block is under the limit. The compiler will compute the sizes.
But there might be an issue: the total shared memory per block on H100 is 227KB, and 57KB is fine.
However, for generality, some devices have only 48KB shared memory per block by default. But H100 has 227KB, and we can set max dynamic shared memory if needed. Static shared memory (declared as `__shared__`) is limited to 48KB by default on older architectures, but on H100 (compute capability 9.0?), it might be higher.
Actually, for CC 9.0 (Hopper), the maximum static shared memory is 48KB unless using dynamic shared memory and requesting more. But actually with `__device__` attributes or compiler flags, you can go up to 227KB. With `extern __shared__` you can request up to 227KB but need to call `cudaFuncSetAttribute` with `cudaFuncAttributeMaxDynamicSharedMemorySize`.
Wait, is that true for static shared memory too? Let me think...
On Ampere/Hopper, the physical limit per block is 164KB or 227KB, but the default CUDA limit for static shared memory is still 48KB unless you compile with specific options or use dynamic shared memory. Actually, I think on CUDA 11+ and SM80+, static shared memory can be up to the device limit if you don't exceed it, but there might be a compile-time warning or error.
To be safe, I should use `extern __shared__ float s_mem[];` and then slice it into s_input and s_weight using offsets:
```cpp
extern __shared__ float s_mem[];
float* s_input = s_mem;
float* s_weight = s_mem + IC_VAL * in_tile_h * in_tile_w;
```
Then in the kernel launch, pass the shared memory size: `kernel<<<grid, block, shared_mem_bytes>>>(...)`.
This is the standard way to handle dynamically sized shared memory. And it allows up to 227KB on H100.
Let me use this approach.
Actually, there's a subtlety: with dynamic shared memory, `cudaFuncSetAttribute` might still be needed for >48KB. But I believe on recent CUDA versions and H100, cudaOccupancyMaxActiveBlocksPerMultiprocessor automatically handles it. Let me look this up.
Actually, for dynamic shared memory, if we request >48KB at launch, the kernel might fail to launch unless we've set the attribute. But on Hopper/SM90, the default max is 48KB for backward compatibility, and we need to explicitly request more using `cudaFuncSetAttribute`.
However, in PyTorch's cpp_extension, we might not have easy access to set this attribute for each kernel launch. But we could do it in the C++ wrapper function.
Let me include a call to `cudaFuncSetAttribute` for the specialized kernel if shared mem > 48KB:
```cpp
size_t shared_mem = ...;
if (shared_mem > 48 * 1024) {
cudaFuncSetAttribute(conv2d_kernel_specialized<...>, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_mem);
}
```
Actually, the attribute is set per function, not per launch. But since we use template instantiations, we'd need to set it for each instantiation. This is getting complicated.
Can we keep shared memory under 48KB without sacrificing much performance?
If we use TILE_H=16, TILE_W=16, TILE_OC=32:
- s_input = 18*18*16*4 = 20,736 bytes
- s_weight = 32*16*9*4 = 18,432 bytes
- Total = 39,168 bytes = 38.2 KB < 48 KB!
This fits under 48KB! And we don't need any special attributes.
With TILE_OC=32:
- FMAs per thread = 32 * 144 = 4,608
- Blocks = 16 * 4 * 64 * 64 = 262,144 blocks
- Each thread has 32 accumulators
This should have good ILP and fit in 48KB.
If we want more ILP, we could have each thread compute 2x2 spatial pixels for 8 OCs:
- Block = 16x16 = 256 threads
- Tile = 32x32 spatial, 8 OC
- s_input = 34*34*16*4 = 73.7 KB > 48KB. Too much.
But we can try Tile = 16x32 spatial, 8 OC with block 8x16=128 threads:
- s_input = 18*34*16*4 = 39.1 KB
- s_weight = 8*16*9*4 = 4.6 KB
- Total = 43.7 KB. Fits!
- FMAs per thread = 2*8*144 = 2304 per pixel? Wait, each thread computes 2x2=4 pixels for 8 OCs = 32 accumulators. Total FMAs = 4 * 8 * 144 = 4608.
- Blocks = 16 * 16 * 64 * 32 = 524,288 blocks
This has more blocks but less warps per block (128 threads = 4 warps).
I think 38KB for TILE_H=16, TILE_W=16, TILE_OC=32 is a good sweet spot under 48KB.
Can we do TILE_H=16, TILE_W=16, TILE_OC=48?
- s_weight = 48*16*9*4 = 27,648 bytes
- s_input = 20,736 bytes
- Total = 48,384 bytes = slightly over 48KB. Probably fails.
TILE_OC=40:
- s_weight = 40*16*9*4 = 23,040 bytes
- s_input = 20,736 bytes
- Total = 43,776 bytes. Fits!
TILE_OC=40 is a bit odd but works. Let me check if 40 divides 128 evenly. 128/40 = 3.2, so no. Last block wastes 8 OCs. Not ideal.
TILE_OC=32 divides 128 evenly (4 blocks). Good.
Could we do TILE_H=8, TILE_W=32, TILE_OC=32?
- s_input = 10*34*16*4 = 21,760 bytes
- s_weight = 32*16*9*4 = 18,432 bytes
- Total = 40,192 bytes. Fits!
- Blocks = 16 * 4 * 128 * 32 = 262,144 blocks
- Each thread = 32 * 144 = 4608 FMAs
Block = 8x32 = 256 threads. A warp would cover a row of 32 threads. Coalescing for both loads and stores is excellent since W dimension is 32.
This might actually be better than 16x16 because:
1. x dimension of block is 32, which is exactly a warp size. This means one warp covers an entire row.
2. For output writes, a warp writes 32 contiguous floats in W dimension. Perfect coalescing.
3. For input loads during cooperative loading, the mapping can be arranged for coalescing too.
Let me use TILE_H=8, TILE_W=32, TILE_OC=32.
Actually wait, for cooperative input loading, we have 10*34*16 = 5440 elements. With 256 threads, threads 0-5439 load one element each, thread 5440+ don't participate in the first round. But 5440 / 256 = 21.25, so some threads do 22 loads.
For coalescing during input loading: if we linearize as ic, ih, iw and have thread idx load s_input[idx], then adjacent threads load adjacent elements. Since s_input is contiguous, this is fully coalesced. But we also need to load from global memory. The global memory offset is:
```cpp
int ic = idx / (in_tile_h * in_tile_w);
int rem = idx % (in_tile_h * in_tile_w);
int ih = oh_base + rem / in_tile_w;
int iw = ow_base + rem % in_tile_w;
offset = ((n * IC + ic) * H + ih) * W + iw;
```
Adjacent threads have the same ic, ih, and adjacent iw. So global memory loads are contiguous (differ by 1 float = 4 bytes). Perfect coalescing!
Wait, but iw goes from 0 to 33 (TILE_W + K - 1 = 32 + 2 = 34? Actually K=3, so TILE_W + K - 1 = 32 + 2 = 34). And for iw=0 to 33, with block covering ow_base. Yes, adjacent threads load adjacent iw values, which are contiguous in memory. Excellent.
For weight loading, similar logic applies. Good.
So I'll go with TILE_H=8, TILE_W=32, TILE_OC=32.
Now, about the block mapping:
- gridDim.x = (OW + TILE_W - 1) / TILE_W = (1022 + 31) / 32 = 32
- gridDim.y = (OH + TILE_H - 1) / TILE_H = (1022 + 7) / 8 = 128
- gridDim.z = N * (OC / TILE_OC) = 16 * 4 = 64
Total blocks = 32 * 128 * 64 = 262,144 blocks.
256 threads per block = 67M threads.
Each thread 4608 FMAs = 308B total FMAs. Correct.
67M threads is a lot but H100 can handle it. 262K blocks across 132 SMs = ~2000 blocks/SM. With 256 threads and 38KB shared memory, we can probably run 4-5 blocks per SM concurrently (limited by ~1K threads per SM? No, H100 can run many more).
Actually, the occupancy limit on H100 is: each SM can have up to 2048 threads. With 256 threads per block, max 8 blocks per SM. With 38KB shared memory, max 227/38 = 5.9 blocks per SM. So we might get 5 or 6 blocks per SM.
5 blocks * 132 SMs = 660 concurrent blocks.
262K blocks / 660 = 397 waves.
Each wave: 660 blocks * 256 threads * 4608 FMAs = 778M FMAs per wave.
At 67 TFLOPS = 33.5B FMA/s, time per wave = 778M / 33.5B = 23.2 ms? No wait, 67 TFLOPS is for the WHOLE GPU, not per wave.
If 660 blocks run concurrently across 132 SMs, that's 5 blocks per SM.
Each SM handles 5 * 256 = 1280 threads.
At 67 TFLOPS for the whole GPU, and 1280 threads * 132 SMs = 168,960 threads.
Each thread does 4608 FMAs = 778M FMAs per wave.
Time per wave = 778M / 33.5B = 23.2 ms.
397 waves * 23.2 ms = 9200 ms. That's way too high!
I'm clearly making a calculation error. Let me reconsider.
67 TFLOPS = 67 * 10^12 operations per second = 67B operations per second if we count FMA as 1 op, or 33.5B FMAs per second if FMA = 2 ops.
Actually, the standard is TFLOPS counts FMA as 2 FLOPs. So 67 TFLOPS = 33.5B FMAs per second = 33.5 * 10^9.
778M FMAs = 0.778 * 10^9.
Time = 0.778 / 33.5 = 0.0232 seconds = 23.2 milliseconds.
And 397 waves * 23.2 ms = 9.2 seconds? That can't be right.
OH WAIT. 778M is the FMAs per wave for the WHOLE GPU (660 blocks). Not per SM. The whole GPU can do 33.5B FMAs/second. So one wave of 778M FMAs takes 778M / 33.5B/s = 0.0232 seconds = 23.2 MILLISECONDS.
Total time = 397 waves * 23.2 ms = 9.2 SECONDS. That's absurd!
Let me recalculate total FMAs:
Output size: N * OC * OH * OW = 16 * 128 * 1022 * 1022 = 2,140,218,112
Each output requires IC * K * K multiply-adds = 16 * 9 = 144 MACs.
Total MACs = 2.14B * 144 = 308B MACs.
Each MAC is 1 FMA?
In NVIDIA terminology, 1 FMA = 2 FLOPs (multiply + add). But 1 MAC = 1 Multiply-ACcumulate = 1 FMA.
So total FMAs = 308B.
At 67 TFLOPS, which counts each FMA as 2 FLOPs, the FMA throughput is 33.5 TFMA/s.
Time = 308B / 33.5B = 9.19 seconds.
Wait, that CAN'T be right. PyTorch does this convolution in tens or hundreds of milliseconds, not seconds!
Oh no, I see the issue. H100 FP32 TFLOPS is 67 TFLOPS. But 1 FMA instruction does 1 multiply and 1 add = 2 FLOPs. So FMA throughput is 33.5 billion FMAs per second.
308 billion FMAs / 33.5 billion FMA/s = 9.19 seconds.
But that would mean a single conv2d of this size takes 9 seconds? That's impossible. Let me recheck the numbers.
Actually, the H100 spec says 67 TFLOPS for FP32. For fp32 on Hopper, each SM can issue 128 FP32 FMAs per cycle? Let me think.
H100 has 132 SMs. At ~2 GHz, that's 264B cycles/second total.
67 TFLOPS / 264B cycles = 0.254 FLOPs per cycle per SM?
Wait, 67 TFLOPS = 67 * 10^12 FLOPs/s.
67 * 10^12 / 132 SMs / 2 * 10^9 cycles/s = 253.8 FLOPs per cycle per SM.
Each FMA is 2 FLOPs, so 126.9 FMAs per cycle per SM.
Each warp has 32 threads, so 32 FMAs per instruction. 126.9 / 32 = 3.96 warp FMAs per cycle per SM.
That seems plausible - each SM has 4 warp schedulers, each can issue 1 FMA instruction per cycle, giving 4 * 32 = 128 FMAs per cycle per SM. Yes.
So 132 SMs * 128 FMAs/cycle/SM * 2 * 10^9 cycles/s = 33.78 * 10^9 FMAs/second? Wait:
132 * 128 = 16,896 FMAs per cycle for the whole GPU.
At 2 GHz = 2 * 10^9 cycles/s.
Total FMA/s = 16,896 * 2 * 10^9 = 33.792 * 10^12 FMA/s = 33.792 TFMA/s.
Each FMA = 2 FLOPs, so 67.58 TFLOPS. Matches the spec.
Total FMAs for conv = 308B.
Time = 308B / 33.792B/s = 9.12 seconds.
But this can't be right for a real-world conv2d. There must be something wrong with my MAC count.
Let me recalculate:
N=16, OC=128, OH=1022, OW=1022 = 16 * 128 * 1022 * 1022.
1022 * 1022 = 1,044,484.
16 * 128 = 2,048.
2,048 * 1,044,484 = 2,139,103,232 output pixels.
Each output pixel is computed from a 3x3 kernel across 16 input channels.
For each output pixel: 16 * 3 * 3 = 144 multiplications and 143 additions? Or 144 multiply-adds?
In PyTorch Conv2d, each multiply is paired with an accumulation. So it's 144 multiply-accumulate operations = 144 FMAs per output pixel.
Total FMAs = 2.139B * 144 = 308B FMAs. That's correct.
At 33.8 TFMA/s, time = 9.1 seconds.
But wait, a 1024x1024 image conv with 16->128 channels is a very large compute! In practice, on an H100, if I run this in PyTorch:
Actually, let me reconsider whether my understanding of H100 throughput is correct. Tensor Cores for FP32 are much faster if using TF32 (989 TFLOPS). But cuDNN might use TF32 for fp32 convolutions if TF32 is enabled by default in PyTorch.
Actually, in recent PyTorch, `torch.backends.cudnn.allow_tf32` defaults to True. So cuDNN might use TF32 Tensor Cores, which gives ~989 TFLOPS with sparsity or ~495 TFLOPS without sparsity.
If the reference uses TF32 Tensor Cores, then the reference FMA throughput for fp32 is much higher than 67 TFLOPS! The effective throughput could be ~250 TFLOPS or ~500 TFLOPS depending on the math mode.
Wait, does TF32 count as fp32 for MAC? TF32 uses 19-bit mantissa and does matmul on 4x4x4 tiles in tensor cores. NVIDIA reports TF32 Tensor Core throughput as 494/989 TFLOPS.
If cuDNN uses TF32 TC for this conv, the throughput is much higher, and our custom fp32 kernel (without TC) might be much slower!
But the problem says "fp32 precision" for tolerance. This doesn't necessarily mean the reference uses true fp32 arithmetic. It could use TF32.
Actually, in PyTorch, `allow_tf32` defaults to True for cuDNN, but `torch.matmul` has its own setting. For conv2d, I believe cuDNN uses TF32 by default on Ampere and Hopper.
If the reference uses TF32 Tensor Cores, then to beat it, we either:
1. Also use TF32 or FP16/BF16 Tensor Cores (but output must match fp32 tolerance - TF32 might match since it's ~1e-3 relative precision, and our tolerance is 1e-4... wait, TF32 tolerance might be around 1e-3 to 1e-4 depending on accumulation. The problem says fp32 uses atol=rtol=1e-4. TF32 roundoff error is about 2^-10 ≈ 0.001, so TF32 accumulation might not meet 1e-4 tolerance for large reductions... but 144 multiplies isn't that many.
Actually, for the default PyTorch behavior, I need to check: does cuDNN use fp32 or tf32?
From PyTorch docs: `torch.backends.cudnn.allow_tf32` defaults to True. This means cuDNN is allowed to use TF32.
But for our custom kernel, we must produce numerically equivalent output. If the reference is using TF32, then doing true fp32 might give slightly different results (outside 1e-4 tolerance).
Actually, let me think about whether TF32 vs FP32 matters for 1e-4 tolerance. For a single multiply, TF32 rounds inputs to ~1e-3 relative, so product error can be ~1e-3. Accumulating 144 such products, error could be ~0.1 relative, which is way above 1e-4.
But wait, TF32 Tensor Cores accumulate in fp32! So each multiply is TF32, but accumulation is fp32. The rounding error comes from the input precision.
For random inputs in [0,1], the typical error of TF32 vs FP32 for a dot product of 144 elements should be well below 1e-4? Let me think. Each element has ~1e-3 relative error. The dot product of two random [0,1] vectors of length 144 has expected value ~36 (since each product has mean 0.25). The error per product is ~1e-3, so total error is ~144 * 1e-3 = 0.14 relative to the sum? That's 0.14/36 = 0.004 relative. That's borderline for 1e-4 absolute if output values are ~1.
Hmm, for 1e-4 absolute tolerance, if output values are say 4 (sum of 16 positive products), absolute error of 0.004 would fail.
Actually, the exact tolerance depends. But it's possible that TF32 output doesn't match true fp32 to 1e-4.
Wait, but PyTorch's tolerance for fp32 is 1e-4. If they run the reference in TF32 mode, the reference itself has TF32 error. And our kernel using true fp32 would have different error. But the comparison checks if our output matches the reference. If reference uses TF32 and we use fp32, we might not match.
To match the reference, we should use the same arithmetic precision. If PyTorch uses TF32 by default for conv, then we need to account for that.
However, there's a catch: the evaluation might actually disable TF32 to ensure deterministic testing. Or they might use the "autocast" behavior.
Looking at the reference code, it's a plain nn.Conv2d. In many benchmark environments, TF32 is enabled by default for speed on Ampere+. But for numerical comparison, they might either:
1. Run reference in fp32 mode (no TF32) and expect fp32 results
2. Or accept whatever the reference produces
I suspect they run exactly the reference code as given. With default PyTorch settings, TF32 is enabled for cuDNN on Ampere/Hopper.
To match the reference numerically, the safest approach is to also use some form of tensor cores or TF32, OR to directly call cuDNN from our kernel, OR to ensure we produce the same numerical results as whatever PyTorch produces.
But wait - the problem says "produce numerically equivalent outputs at fp32 precision" and tolerance is 1e-4. If the reference uses TF32, then "fp32 precision" here refers to the data type (float32 tensors), not the compute precision. The reference output might have TF32-level error.
For a custom kernel to match a TF32 reference output, we could:
1. Use TF32 ourselves
2. Use CUTLASS with TF32
3. Or do true fp32 and hope the error is within 1e-4
Actually, the tolerance of 1e-4 for fp32 is quite tight. For TF32 accumulation of 144 products, the error might exceed this.
Let me think more carefully. The reference is `nn.Conv2d` with default initialization (random weights with some distribution). The inputs are `torch.rand` which gives [0, 1). Output values for a 3x3 conv with 16 channels could be in range [0, 16) roughly. An absolute error of 1e-4 is quite small.
If PyTorch uses TF32, the error per MAC is about 1 ulp of TF32 ≈ 2^-10 ≈ 0.001. For 144 MACs, the accumulation error could be ~0.001 * sqrt(144) ≈ 0.012 with random sign, or 0.001 * 144 = 0.144 with same sign. But since weights are random (both positive and negative), errors partially cancel. Expected absolute error is probably closer to 0.001 * sqrt(144) * std(input*weight) ≈ 0.01 or so.
An absolute error of 0.01 at output values of ~1-4 would give relative error of 0.25%-1%, which is well above 1e-4.
Wait, but I think TF32 tensor cores accumulate in full fp32! The multiplication inputs are rounded to TF32, but the multiply and add are exact in fp32. So the error only comes from the multiplication inputs being rounded.
For random inputs in [0,1] and random weights ~N(0, std), the product has some variance. The rounding of inputs to TF32 introduces an error of at most 0.5 ulp(TF32) in each operand. The relative error in each operand is <= 2^-10. For a product, relative error is about 2 * 2^-10 = 2^-9. The absolute error in each product is then about val * 2^-9.
For values around 1, absolute error per product is ~0.002. Sum of 144 such products with random sign: the error variance is 144 * (0.002)^2 * variance_of_sign. The standard deviation of the total error is sqrt(144) * 0.002 = 12 * 0.002 = 0.024. So we expect absolute errors around 0.02-0.03.
0.03 is way above 1e-4! So if the reference uses TF32, a true fp32 custom kernel would produce very different results.
Hmm, but wait. Does `nn.Conv2d` really use TF32 by default? Let me double check.
Actually, looking at cuDNN documentation and PyTorch source: `torch.backends.cudnn.allow_tf32` controls whether cuDNN is allowed to use Tensor Core TF32 ops. It defaults to True. So yes, cuDNN will use TF32 if available and heuristically chosen.
But for the tolerance check to pass at 1e-4, the test must either:
1. Not use TF32 (set allow_tf32=False)
2. Or compare with a tolerance that accounts for TF32 error
Since the problem explicitly states fp32 tolerance is 1e-4, and if TF32 produces ~0.01 error, the test would fail against a true fp32 reference. So the test environment must disable TF32 or use a reference that matches fp32.
Actually, there are two possibilities:
1. The test runs both Model and ModelNew with the same PyTorch settings (allow_tf32=True by default). In this case, both the reference and our output would be compared. But our custom kernel does true fp32, so we'd differ from the TF32 reference. We'd fail.
2. The test specifically sets fp32 mode (no TF32) to compare exact fp32 implementations.
I think option 2 is more likely, or option 1 but the evaluator copies weights from the reference to ModelNew and expects exact match regardless.
Actually, I can avoid this issue entirely by using a custom kernel that matches the reference numerically. If the reference uses TF32, we fail. So to be safe, I should make our kernel produce the same results as the reference.
One way: if we're not sure about TF32, we could call cudnn ourselves from C++, or we could use an ATen fallback for the general case, or we could just do true fp32 and hope.
Actually, you know what, let me check if there's a way to verify. I can't actually run the test to check whether allow_tf32 is on. But I expect that the test environment either disables TF32 or has tolerance that works with either.
Wait, re-reading the problem: "Produces numerically equivalent outputs at fp32 precision". The fact that they say "at fp32 precision" and give 1e-4 tolerance suggests they're testing exact fp32 math, or at least the reference is in whatever precision cuDNN chooses.
But here's the key: the reference is `self.conv2d(x)` which is PyTorch's standard conv2d. If PyTorch chooses TF32, that's the reference. We need to match it.
For our custom kernel to match TF32 numerically, we'd need to round our inputs to TF32, which is very tricky. If we're not sure, maybe we should call cuDNN from our C++ code rather than writing a custom CUDA kernel.
Wait, but the point of the problem is to write a custom CUDA kernel. If we just call cuDNN or at::conv2d, we'd be using the exact same code as the reference and would only be as fast, never faster (and potentially slower due to Python overhead).
Another thought: the tolerance might actually allow TF32-level differences. The problem says fp32 uses atol=rtol=1e-4, but for a conv output of magnitude ~10, 1e-4 relative + 1e-4 absolute is actually quite tight.
Actually, for this specific config: in_channels=16, kernel_size=3, batch=16, oc=128, the output values would be roughly:
- Input in [0, 1), 16 channels
- Conv2d weight initialized by default: `kaiming_uniform_` with a = sqrt(5), fan_in = 16*9 = 144. The weight scale is roughly 1/sqrt(fan_in/3) ≈ 1/sqrt(48) ≈ 0.14. But the actual range might be [-0.2, 0.2] or so.
- Output: sum of 144 products of [0,1] and ~[-0.2, 0.2]. Expected value around 0. Mean is 0. Variance is 144 * Var(input*weight) ≈ 144 * 0.01 = 1.44. Std dev ≈ 1.2.
- So typical output values are in range [-4, 4], with mean near 0.
- 1e-4 absolute tolerance is very tight for outputs near 0, but for outputs around 4, relative tolerance of 1e-4 gives 4e-4 absolute.
If TF32 introduces ~0.02 error, that would fail 1e-4 tolerance.
So I strongly suspect the test disables TF32 or has some fallback that ensures exact fp32.
Therefore, let's proceed with a true fp32 custom kernel.
Now, regarding whether we can beat cuDNN without TF32:
If the reference uses true fp32 cuDNN, its throughput is bounded by ~67 TFLOPS. The theoretical minimum is ~9.2 seconds? No, that can't be. There's definitely a mistake because conv2d of this size on GPU takes milliseconds, not seconds.
Let me recalculate. 308B FMAs? No, that's wrong!
Wait. Output pixels: 16 * 128 * 1022 * 1022.
16 * 128 = 2048.
1022 * 1022 = 1,044,484.
2048 * 1044484.
Let me compute: 2,000 * 1,000,000 = 2,000,000,000.
48 * 1,044,484 = 50,135,232.
Total = 2,050,135,232. Yes, about 2.05B output pixels.
FMAs per output pixel: IC * KH * KW = 16 * 3 * 3 = 144.
Total FMAs = 2.05B * 144 = 295.2B. Or approximately 300B.
300 billion FMAs.
Oh... I see the issue now. The H100 FP32 non-TC throughput is 67 TFLOPS = 67 * 10^12 operations per second = 67 billion operations per second.
Wait, TERA = 10^12. 67 TFLOPS = 67 * 10^12 FLOPs per second.
67 * 10^12 = 67,000,000,000,000 = 67 trillion.
Oh my god. 67 TFLOPS is 67 TRILLION FLOPs per second, not 67B.
I was off by 1000!
67 TFLOPS = 67 * 10^12 FLOPs/s.
Each FMA = 2 FLOPs, so 33.5 * 10^12 FMAs/s = 33.5 trillion FMAs/s.
308 billion FMAs / 33.5 trillion FMAs/s = 308 * 10^9 / 33.5 * 10^12 = 0.0092 seconds = 9.2 milliseconds.
Phew! Yes, 9.2 ms. That's a very reasonable time for a large conv.
So total FMAs = 308 billion (3.08 * 10^11).
Throughput = 3.35 * 10^13 FMA/s.
Time = 9.19 ms.
cuDNN for fp32 might achieve ~70-85% of peak = ~6.5-8 ms.
If TF32 is enabled, throughput is ~250-500 TFLOPS, and time could be ~1-2 ms.
If the reference uses TF32, we would need TF32 or some other optimization to beat it.
But cuDNN might also choose a non-TF32 algorithm if TF32 is enabled but doesn't improve performance for this shape, or if the heuristic chooses otherwise.
OK, with 9.2 ms theoretical for true fp32, and cuDNN likely achieving around 7-10 ms, we need a custom kernel that achieves similar or better.
For TILE_H=8, TILE_W=32, TILE_OC=32:
- Total FMAs = 308B
- FMAs per block = 256 threads * 4608 FMAs = 1.18M FMAs... wait, 256 * 4608 = 1,179,648 FMAs per block.
- With 262,144 blocks, total FMAs = 262,144 * 1,179,648 = 309,237,748,736 ≈ 309B. Matches.
Number of waves = total blocks / (concurrent blocks).
With 38KB shared memory, we can fit 5-6 blocks per SM on H100 (max shared 227KB).
132 SMs * 5 blocks = 660 concurrent blocks.
Waves = 262,144 / 660 = 397 waves.
Time = 397 * time_per_wave.
Time per wave: 660 blocks * 1.18M FMAs = 778M FMAs per wave.
At 33.5T FMA/s = 33,500B FMA/s.
Time per wave = 778M / 33,500B/s = 23.2 microseconds.
Total time = 397 * 23.2us = 9.21 ms.
This matches the theoretical limit! But this assumes:
1. Perfect FMA throughput utilization
2. No overhead from shared memory loads, output writes, synchronization, etc.
3. Perfect occupancy
In practice, there will be overhead, so actual time might be 10-15 ms. If cuDNN achieves 7-8 ms, we would be slower.
Can we reduce overhead? The time_per_wave calculation assumes the FMA units are the bottleneck. But we also have:
- Input loading: ~20KB per block * 660 blocks = 13.2MB per wave. At 3TB/s = 4.4 microseconds.
- Weight loading: ~18KB per block * 660 = 11.9MB per wave. At 3TB/s = 4 microseconds.
- Output writes: ~32 values * 4 bytes * 256 threads * 660 blocks = 21.5MB per wave. At 1.5TB/s = 14.3 microseconds.
Wait, output writes are 21.5MB per wave! At "half" bandwidth (since writes share bandwidth with reads), that's ~7 microseconds.
So memory overhead per wave: 4.4 + 4 + 7 = 15.4 microseconds.
Compute per wave: 23.2 microseconds.
Total per wave with overlapping: max(23.2, 15.4) = 23.2 if perfectly overlapped? But memory and compute compete for resources.
Actually, FMA units and memory units operate concurrently. The question is whether the SM can hide the memory latency. With 5 blocks per SM and 256 threads each = 1280 threads = 40 warps. H100 SMs can have hundreds of warps in flight, so 40 warps might not fully hide latency.
If not fully hidden, effective time per wave = compute_time + (memory_time that can't be hidden).
If we say 15.4us memory vs 23.2us compute, and 60% of memory is hidden, then overhead = 6.2us.
Total = 29.4us per wave. Total = 397 * 29.4us = 11.7 ms.
Hmm, we need to be more competitive.
What if each thread computes MORE work to amortize the loading overhead?
If each thread computes 1x4 spatial pixels for 8 output channels (using vectorized output writes):
- Block = 16x8 = 128 threads
- Tile spatial = 16x32
- Tile OC = 8
- Input shared = 18x34x16 = 9,792 floats = 39 KB
- Weight shared = 8*16*9 = 1,152 floats = 4.6 KB
- Total = 43.6 KB
- Number of concurrent blocks: 5 per SM
- FMAs per thread: 4 * 8 * 144 = 4,608
- FMAs per block: 128 * 4608 = 589,824
- Blocks: N * (OC/8) * ceil(OH/16) * ceil(OW/32) = 16 * 16 * 64 * 32 = 524,288 blocks
- Concurrent: 132 * 5 = 660
- Waves: 524,288 / 660 = 794 waves
- Compute per wave: 660 * 589,824 = 389M FMAs = 11.6us at peak
- Memory per wave: input 128*5*39KB = 25MB at 3TB/s = 8.3us; weight ~3MB = 1us; output 128*5*4*8*4 = 10.2MB at 1.5TB/s = 6.8us
- Total time ≈ 794 * (11.6 + some overhead) ≈ 794 * 15us = 11.9 ms
Not better.
What about using TILE_H=16, TILE_W=16, TILE_OC=64 but using warp-level primitives? Or what if we go bigger:
Actually, what if we don't use shared memory for weights at all, and keep them in L2 cache? Weight is only 73KB total. Then we can use ALL shared memory for input.
TILE_H=32, TILE_W=32, input shared = 34x34x16x4 = 73.7 KB. This is exactly at the 48KB limit... wait, it exceeds it.
What about TILE_H=32, TILE_W=16? Input shared = 34x18x16x4 = 39KB.
Block = 32x16 = 512 threads. But max threads per block is 1024, so OK.
Each thread computes 1 pixel for some output channels.
Wait, let's try TILE_H=16, TILE_W=32:
- Input shared = 18*34*16*4 = 39.1 KB
- Block = 16x32 = 512 threads
- If TILE_OC=16: weight shared = 16*16*9*4 = 9.2 KB. Total = 48.3 KB. Slightly over.
If TILE_OC=8: weight = 4.6 KB. Total = 43.7 KB.
Block computes 16x32x8 = 4096 outputs.
FMAs per thread = 1 * 8 * 144 = 1152. This is too small.
What if each thread computes MORE per thread but we keep block size smaller?
TILE_H=32, TILE_W=16, block=16x16=256 threads, each thread computes 2x1 spatial for 16 OCs:
- Input shared = 34*18*16 = 9,792 floats = 39.1 KB
- Weight shared = 16*16*9 = 2,304 floats = 9.2 KB
- Total = 48.3 KB. Slightly over 48KB.
If we reduce to 8 OCs: weight = 4.6 KB. Total = 43.7 KB.
Each thread computes 2x1 * 8 = 16 outputs.
FMAs per thread = 16 * 144 = 2304.
Blocks = 16 * 16 * 64 * 64 = 1,048,576 blocks.
1M blocks is good for saturation.
Hmm, 2304 FMAs per thread is still a bit low. But with 1M blocks and 5 per SM, that's 200K waves. Compute per wave = 660 blocks * 256 threads * 2304 = 389M FMAs = 11.6us. Total = 200K * 11.6us... wait, no:
1,048,576 blocks / 660 = 1588 waves.
1588 * 11.6us = 18.4 ms. That's worse than before.
The issue is that each block has fewer FMAs.
So we want FEWER blocks but MORE FMAs per block, while keeping shared memory under 48KB and maximizing occupancy.
What if we use threads to compute 4x1 spatial for 16 OCs with block=16x8=128:
- Tile = 64x8 spatial, 16 OC
- Input shared = 66*10*16 = 10,560 floats = 42.2 KB
- Weight shared = 16*16*9 = 2,304 floats = 9.2 KB
- Total = 51.4 KB > 48KB.
If 8 OCs: weight = 4.6 KB. Total = 46.8 KB. Fits!
Tile = 64x8 spatial, 8 OC.
FMAs per thread = 4 * 8 * 144 = 4608.
Blocks = 16 * 16 * 16 * 128 = 524,288.
Not bad, but spatial is very elongated (64x8). Might have poor locality.
Let me get back to a balanced tile and just try to beat cuDNN by having low overhead.
Actually, here's another idea: What if we use L2 cache for weights and don't load them to shared memory? Then we can use all shared memory for input and get a larger spatial tile, which means fewer blocks and less overhead.
TILE_H=32, TILE_W=32:
- Input shared = 34*34*16*4 = 73.7 KB > 48KB. Won't work with static shared.
What if we manually manage caching? We could load input directly to registers and avoid shared memory for some of it. But that's complex.
Actually, with dynamic shared memory, we can go up to 227KB. But we need to call `cudaFuncSetAttribute`. In PyTorch's load_inline, we might not have easy access to CUDA APIs.
Wait, can we use `cudaFuncSetAttribute` in inline C++? Yes! We just need to `#include <cuda_runtime.h>` and call it.
For the specialized kernel, we can do:
```cpp
size_t shared_mem_size = ...;
auto fn = conv2d_kernel_specialized<...>;
if (shared_mem_size > 48 * 1024) {
cudaFuncAttributes attr;
cudaFuncGetAttributes(&attr, fn);
if (attr.maxDynamicSharedSizeBytes < shared_mem_size) {
cudaFuncSetAttribute(fn, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_mem_size);
}
}
fn<<<grid, block, shared_mem_size>>>(...);
```
Wait, `cudaFuncSetAttribute` expects the first argument as a function pointer. For a template kernel, we'd need to specify the template args. This should work.
But we also need to handle the case where the kernel is launched for the first time. Setting attributes before the first launch is fine.
Actually, there's another requirement: the kernel must be launched with `block` and `grid` compatible with the occupancy. But `load_inline` should handle compilation.
Let me try using dynamic shared memory > 48KB. Actually, I want to see if we can get close to 80KB or even 100KB shared memory for larger tiles.
But for the first attempt, let me keep it under 48KB to avoid issues.
TILE_H=16, TILE_W=16, TILE_OC=32 gave 38KB.
With TILE_H=16, TILE_W=16, TILE_OC=40, it's 43.8KB. But 40 doesn't divide 128.
TILE_H=16, TILE_W=16, TILE_OC=48: 48.4KB. Slightly over.
What about TILE_H=8, TILE_W=32, TILE_OC=48?
Input shared = 10*34*16*4 = 21,760 bytes.
Weight shared = 48*16*9*4 = 27,648 bytes.
Total = 49,408 bytes. Over 48KB.
TILE_OC=40: weight = 23,040. Total = 44,800 bytes. Under 48KB.
Blocks = 16 * 4 (ceil(128/40)=4) * 128 * 32 = 262,144.
FMAs per thread = 40 * 144 = 5,760.
Output per block = 8*32*40 = 10,240 values.
Blocks per SM: shared 44KB, max ~5. So 132 * 5 = 660 concurrent.
Waves = 262,144 / 660 = 397.
Compute per wave = 660 * 256 * 5760 = 973M FMAs = 29us at peak.
Memory per wave: input ~660*22KB = 14.5MB. Weight ~15.1MB. Output ~660*10KB*4 = 26.4MB. Total read+write ~56MB. At 3TB/s (bidirectional) = 18.7us.
Total ≈ 47.7us per wave. Total = 397 * 47.7 = 19 ms. Worse.
Hmm, increasing TILE_OC doesn't help if it increases weight loading proportionally.
Wait, output memory scales with TILE_OC too. And output writes at 26.4MB per wave is significant.
What if we process output channels sequentially per block but keep them in registers? That's what we're doing. The output write bandwidth is proportional to how many outputs we compute per block.
Actually, each output value is 4 bytes and must be written once. Total output bytes = 2.14B outputs * 4 bytes = 8.56 GB.
At peak bandwidth 3.35 TB/s, writing output takes 2.55 ms.
Reading input (1 GB) and weights (73KB) adds ~0.3 ms.
Total memory time = ~2.9 ms.
Compute time = 9.2 ms.
So the operation is compute-bound. The exact memory overhead per wave is relatively small compared to compute.
In my wave calculation, the memory bandwidth per wave is for the blocks in that wave, and total memory is limited by the GPU's total BW. So total memory time for the whole kernel is ~2.9 ms, and compute time is ~9.2 ms. The kernel is compute-bound, so total time ≈ 9.2 ms + small overhead.
The key to beating cuDNN is maximizing FMA unit utilization!
For this, we need:
1. Many independent FMAs in flight per SM
2. Good warp scheduling
3. Minimal stalls
With 38KB shared memory and 256 threads, we can potentially fit 5+ blocks per SM, which gives 30+ warps. H100 can schedule many warps.
But to keep FMA units busy, each warp needs independent FMAs. If a thread has 32 accumulators and does a loop of 32 FMAs per iteration, then within a warp there are 32 threads * 32 FMAs = 1024 independent FMAs per iteration. That's quite a lot! Each cycle, a warp can issue multiple FMAs. On H100, a warp scheduler can issue an FMA instruction every cycle. With 4 warp schedulers, that's 4 * 32 = 128 FMAs per cycle per SM.
To saturate 128 FMAs/cycle, we need 128 independent FMAs ready each cycle. With 1024 independent FMAs per iteration, we have 8 cycles worth. This should be sufficient to keep the pipeline full.
So I think TILE_H=8, TILE_W=32, TILE_OC=32 should be able to achieve near-peak FMA throughput.
What if we do TILE_H=16, TILE_W=16, TILE_OC=32 with the same analysis?
- Block = 256 threads
- Each thread 4608 FMAs with 32 accumulators
- Same ILP, same total compute
- Memory patterns differ slightly
I think both should perform similarly.
One consideration: The 16x16 spatial tile has 18x18 input. The 8x32 tile has 10x34 input. The 8x32 tile has input aspect ratio 1:3.4, which might spread accesses across more cache lines when loading ic dimension. But in shared memory, layout is [ic][h][w], so for fixed ic, contiguous w is contiguous memory. With w=34, that's 34*4=136 bytes per row. Adjacent threads in a warp (32 threads along w) read 32 contiguous floats = 128 bytes, which is 2 cache lines (assuming 64-byte lines). Good.
For 16x16 tile, 16 contiguous floats = 64 bytes = 1 cache line. Also good.
For output writes, 8x32 block writes 32 contiguous floats per row at w=32. That's 128 bytes = 2 cache lines, or with 32 threads in a warp... wait.
If block is 8x32 (ty=0..7, tx=0..31), then:
- Thread (ty, tx) computes pixel (oh_base + ty, ow_base + tx)
- Warp covers threads with consecutive tid. If blockDim.x=32, then tid = ty*32 + tx. A warp of 32 threads has ty=0, tx=0..31 or ty=1, tx=0..31 depending on how warps are formed.
- Actually, warps are formed from consecutive thread IDs. So warp 0: ty=0, tx=0..31. Warp 1: ty=1, tx=0..31, etc.
- Each warp writes 32 contiguous floats in the W dimension. Perfect coalescing!
For 16x16 block: warp 0 covers ty=0, tx=0..15 and ty=1, tx=0..15? No, with blockDim.x=16, warp 0 = ty=0..1, tx=0..15? Let me check:
tid = ty * 16 + tx.
tid 0..15: ty=0, tx=0..15
tid 16..31: ty=1, tx=0..15
So a warp covers 2 rows x 16 columns.
For output writes, the warp writes 16 values in row 0 (contiguous) and 16 values in row 1 (contiguous but offset by OW). A warp can issue 2 stores, each 16 floats. Not as clean as 1 store of 32 floats, but still coalesced.
For 8x32 block, each warp writes exactly one row of 32 contiguous floats. This is slightly better.
For shared memory loading of weights in 8x32 block:
- A warp of 32 threads loads weights. If we map consecutive threads to consecutive weights, that's fully coalesced.
For input loading:
- tid = ty * 32 + tx. With 256 threads, tid 0..255.
- For 10*34*16 = 5440 elements, thread i loads element i, i+256, i+512, etc.
- Elements are laid out as [ic][h][w]. So element 0 is (ic=0, h=0, w=0), element 1 is (ic=0, h=0, w=1), ..., element 33 is (ic=0, h=0, w=33), element 34 is (ic=0, h=1, w=0).
- Thread 0 loads elements 0, 256, 512, 768, 1024, 1280, 1536, 1792, 2048, 2304, 2560, 2816, 3072, 3328, 3584, 3840, 4096, 4352, 4608, 4864, 5120, 5376.
- These are spaced 256 apart in the array. With 32 banks, 256 floats = 1024 bytes. 1024 / 128 = 8, so they access banks 0, 8, 16, 24, 0, 8, 16, 24...?
Wait, bank conflict when writing TO shared memory during loading is irrelevant (each thread writes its own location). For reading FROM global memory, we need coalesced global loads.
Global memory access pattern:
Thread 0 reads global at offset for (n, ic=0, h=oh_base+0, w=ow_base+0)
Thread 1 reads global at offset for (n, ic=0, h=oh_base+0, w=ow_base+1)
...
Thread 33 reads global at offset for (n, ic=0, h=oh_base+1, w=ow_base+0)
But thread 0 loads global at: base + ((n*IC + 0) * H + oh_base) * W + ow_base
Thread 256 loads global at: base + ((n*IC + ???) * H + h) * W + w
If element 256 corresponds to ic=256/340=0 (since 10*34=340), h=256/34=7, w=256%34=18. Wait:
ic = idx / 340. idx=256: ic=0, rem=256, h=256/34=7, w=256%34=18.
So thread 0 loads (ic=0, h=0, w=0).
Thread 256 loads (ic=0, h=7, w=18).
These are NOT contiguous in global memory. In fact, they're far apart!
This means the loop `for (idx = tid; idx < total; idx += nthreads)` with elements mapped to [ic][h][w] gives poor coalescing for global loads!
To fix this, we need to assign threads to contiguous elements in global memory.
Global memory layout is [N][IC][H][W]. For a fixed n, elements are contiguous along W, then H, then IC.
For the input tile, we want to load from [n][ic][oh_base + ih][ow_base + iw].
Contiguous global memory corresponds to varying iw, then ih, then ic.
So for coalesced loads, adjacent threads should load adjacent iw values.
With 256 threads, we want thread i to load iw = (i % 34), ih = (i / 34) % 10, ic = i / 340.
But 256 threads = 7 full rows (7*34=238) + 18 remaining.
Warp 0: threads 0..31 load iw=0..31, ih=0, ic=0. That's 32 contiguous floats. Perfect coalescing.
Warp 1: threads 32..63 load iw=32..33 (2 floats), then ih=1, iw=0..29 (30 floats). Not fully contiguous but still within a small range.
Actually, global memory coalescing works within 128-byte segments. So as long as threads in a warp access addresses within a 128-byte window, it's coalesced.
For threads 32..63:
- thread 32: (ic=0, h=0, w=32)
- thread 33: (ic=0, h=0, w=33)
- thread 34: (ic=0, h=1, w=0)
...
Address of (ic, h, w) = base + ic*H*W + h*W + w.
thread 32: base + 0 + 0 + 32 = base + 32*4 = base + 128 bytes.
thread 33: base + 33*4 = base + 132 bytes.
thread 34: base + 1*W + 0 = base + 1024*4 = base + 4096 bytes.
Oh no, thread 34 is at base + 4096, which is not contiguous with thread 33 at base + 132. This means warp 1 (threads 32-63) does NOT coalesce! Threads 32 and 33 are close, but thread 34 jumps to a different row.
This is a major issue. To fix global memory coalescing, we need adjacent threads to load adjacent W coordinates.
One approach: Have each warp load a contiguous 1D slice of the 3D input tile. Map thread 0 to w=0, h=0, ic=0; thread 1 to w=1, h=0, ic=0; etc. When w exceeds tile_w, increment h. When h exceeds tile_h, increment ic.
Formally:
```
int idx = tid; // 0 to nthreads-1
int total_per_ic = in_tile_h * in_tile_w = 340;
int ic = tid / total_per_ic;
int rem = tid % total_per_ic;
int h = rem / in_tile_w;
int w = rem % in_tile_w;
```
But with 256 threads and 340 elements per ic, warp 0 covers threads 0-31 which are within ic=0, w=0..31, h=0. Contiguous. Good.
Warp 1 covers threads 32-63: ic=0, w=32..33 and h=1, w=0..29.
Address of thread 32: base + 32*4 = base+128.
Address of thread 33: base + 33*4 = base+132.
Address of thread 34: base + 1*1024*4 + 0*4 = base+4096.
These are far apart. 4096 - 132 = 3964 bytes = 31 cache lines. This is NOT coalesced.
To make it coalesced, we should process the entire input tile by having each thread load along the W dimension for a specific (ic, h). But each (ic, h) pair has only 34 elements, and we have 256 threads. So we need to load multiple (ic, h) positions.
Better: Have warps load one (ic, h) at a time. Each warp loads 32 contiguous w values for one (ic, h).
- Warp 0: (ic=0, h=0, w=0..31)
- Warp 1: (ic=0, h=0, w=32..33) and then (ic=0, h=1, w=0..29)
But warp 1 still has a gap.
Actually, for (ic=0, h=0, w=32..33): that's only 2 elements. We could have warp 1 do (ic=0, h=1, w=0..31). Then the remaining elements (ic=0, h=0, w=32..33) are loaded by thread 32 and 33 of warp 0? No, warp 0 already loaded 0..31.
Wait, warp 0 loads 32 elements. So 2 elements (w=32,33) remain for h=0. We could load them with threads 0 and 1 of some warp. Not clean.
Let's see: 340 elements per ic. Let's divide by warp size 32.
340 / 32 = 10.625 warps per ic. Not an integer.
For the configuration with TILE_H=16, TILE_W=16:
Input tile = 18*18 = 324 elements per ic.
324 / 32 = 10.125 warps. Still not integer.
For TILE_H=8, TILE_W=16:
Input tile = 10*18 = 180 elements per ic.
180 / 32 = 5.625 warps per ic.
Hmm.
What if we use float4 vectorized loads? With float4 (16 bytes), each thread loads 4 elements. 340 / 4 = 85 float4 loads per ic. With 256 threads / 4 = 64 float4 operations per iteration... no, 256 threads each doing float4 loads = 256*4 = 1024 elements per iteration.
If we do:
```cpp
for (int idx = tid; idx < total_elems; idx += nthreads) { ... }
```
With nthreads=256 and total_elems=5440, we get 5440/256 = 21.25 iterations.
Thread 0: elements 0, 256, 512, ..., 5120.
Thread 1: elements 1, 257, 513, ..., 5121.
Adjacent threads load adjacent elements! This is NOT contiguous in global memory because of the big stride, but...
Wait, for global memory coalescing, what matters is whether threads in a warp access contiguous memory. With `idx = tid + iter * nthreads`:
In iteration 0, threads 0..31 load elements 0..31. These are contiguous. Coalesced!
In iteration 1, threads 0..31 load elements 256..287. Are these contiguous?
Element 256: ic=0, h=7, w=18
Element 257: ic=0, h=7, w=19
...
Element 273: ic=0, h=8, w=0 (wait, 7*34+18=256, 7*34+19=257... 7*34+33=271, then 8*34+0=272).
So elements 256..271 are in row 7, columns 18..33. Elements 272..287 are in row 8, columns 0..15.
Are these contiguous in global memory?
Element (ic, h, w) = base + ic*H*W + h*W + w.
h=7, w=18: base + 7*1024 + 18 = base + 7186
h=7, w=19: base + 7187
...
h=7, w=33: base + 7201
h=8, w=0: base + 8192
Oops! Element 272 (h=8, w=0) is at base + 8192, while element 271 (h=7, w=33) is at base + 7201. There's a jump of 991 elements! That's H - 34 = 1024 - 34 = 990. So elements 271 and 272 are NOT contiguous in global memory.
So iteration 1 for warp 0 loads:
thread 0: base + 7186
thread 1: base + 7187
...
thread 15: base + 7201
thread 16: base + 8192
thread 17: base + 8193
...
This is a TERRIBLE access pattern. The warp would issue 2 transactions: one for 0..15 and one for 16..31. Actually, because coalescing works on 128-byte segments, thread 15 is at base+7201 (+15*4 = +60 bytes from thread 0's address). That's still within one 128-byte segment (base+7168 to base+7295). Thread 16 is at base+8192, which is a different segment. So the warp issues at least 2 loads, and threads 0..15 are in one segment, threads 16..31 in another. But these two segments are far apart.
Actually, modern GPUs can handle up to 32 transactions per warp. But each transaction to a different cache line has overhead.
This access pattern is really bad! We'd have many cache lines loaded per warp, and most of each cache line would be unused.
To fix this, we need adjacent threads to load adjacent memory in global memory. The simplest way is to map threads to contiguous W coordinates, and use a strided mapping that respects the global layout.
For the tile [ic][h][w], global memory has w contiguous. So within a row (fixed ic, h), w is contiguous. Across rows (fixed ic, h+1), there's a gap of W - in_tile_w.
For a single warp to load contiguous memory, it should load a single row of up to 32 elements, or multiple rows if they're contiguous (they aren't).
One strategy: process row by row. Each warp loads elements from one row at a time.
- For row h, there are in_tile_w = 34 elements.
- Warp has 32 threads. It loads w=0..31.
- Remaining 2 elements (w=32,33) are loaded by threads 0..1 of the next warp (or same warp in another iteration).
- Then move to row h+1.
To implement this cleanly, we can have threads iterate over rows:
```cpp
for (int idx = tid; idx < in_tile_h * in_tile_w; idx += nthreads) {
int h = idx / in_tile_w;
int w = idx % in_tile_w;
int ic = ...;
}
```
No, this maps adjacent threads to adjacent (h, w) in the tile. For the first warp:
thread 0: h=0, w=0
thread 1: h=0, w=1
...
thread 31: h=0, w=31
thread 32: h=0, w=32
Global addresses for threads 0-31: base + h*W + w = base + 0*1024 + 0 to 31. Contiguous! Coalesced!
Thread 32: base + 0*1024 + 32 = base + 32. Still contiguous with thread 31! If warp has 32 threads, thread 31 is at base+31*4=base+124. Thread 32 would be in the next warp. So warp 0 (threads 0-31) loads 32 contiguous floats = 128 bytes = perfect coalescing (2 cache lines of 64 bytes each).
Warp 1 (threads 32-63):
thread 32: h=0, w=32
thread 33: h=0, w=33
thread 34: h=1, w=0
thread 35: h=1, w=1
...
Thread 32: base + 32. Thread 33: base + 33. Thread 34: base + 1024. Thread 35: base + 1025.
These are NOT contiguous. Warp 1 loads base+32, base+33, base+1024, base+1025, etc. This will require many transactions.
Hmm, the issue is that we're mixing the tail of one row with the start of the next row within the same warp.
To fix this, we can process each row with as many full warps as possible, and handle the tail separately.
For row length 34:
- 32 elements can be loaded by one warp (w=0..31).
- Remaining 2 elements (w=32,33) loaded by threads 0..1 of some warp. Not efficient but only 2 elements.
So for each row h:
```cpp
int w_start = tid % 32; // or something
```
Actually, a cleaner approach: map threads so that each warp handles a full 32-width chunk if possible.
For in_tile_w = 34:
- Elements 0..31: loaded by warp 0
- Elements 32..33: loaded by threads 0..1 of warp 1
- Elements 34..65: loaded by warp 1... wait, this overlaps.
This is getting complicated. Let me instead choose TILE_W such that in_tile_w is a multiple of 32!
If K=3, then in_tile_w = TILE_W + 2.
For in_tile_w to be a multiple of 32:
TILE_W + 2 = 32, 64, 96...
TILE_W = 30, 62, 94...
If TILE_W = 30:
Block = 8x30 = 240 threads.
in_tile_w = 32. Perfect multiple!
in_tile_h = 10.
s_input = 10*32*16*4 = 20,480 bytes.
s_weight = 32*16*9*4 = 18,432 bytes.
Total = 38,912 bytes = 38 KB.
Blocks:
- grid.x = (1022 + 29) / 30 = 1052 / 30 = 35
- grid.y = (1022 + 7) / 8 = 128
- grid.z = 16 * 4 = 64
- Total = 35 * 128 * 64 = 286,720 blocks
FMAs per thread = 32 * 144 = 4,608.
This is very similar to 8x32 config but with 286K blocks instead of 262K, and better coalescing!
For input loading:
Total elements = 10*32*16 = 5,120.
With 240 threads, each thread loads 5,120/240 = 21.33 elements.
Thread i loads elements i, i+240, i+480, ..., i+5040.
Element mapping: ic = idx / (10*32) = idx / 320, rem = idx % 320, h = rem / 32, w = rem % 32.
For iteration 0, thread i (i in {0..31}) has:
idx = i, so ic=0, h=0, w=i. Global address = base + w. Contiguous! Coalesced!
For iteration 1, idx = i + 240.
ic = 240/320 = 0, rem = 240, h = 240/32 = 7, w = 240%32 = 16.
So thread i loads (ic=0, h=7, w=16+i).
For i=0..15: w=16..31, row 7. Contiguous! Coalesced!
For i=16..31: w=32..47, but w only goes up to 31! So w wraps around? No, 16+i for i=16 is 32, but in_tile_w=32, so w=32 is out of bounds!
Wait, I miscalculated. idx=256 would be ic=0, h=8, w=0. idx=257: h=8, w=1, etc. So for i=16 (idx=256), h=8, w=0. This is contiguous in tile index but not contiguous in global memory (since h=7,w=31 to h=8,w=0 jumps by 1024-31 = 993 elements).
But if we look at the elements loaded in iteration 1 (idx=240 to 479):
Thread 0: idx=240, ic=0, h=7, w=16
Thread 1: idx=241, ic=0, h=7, w=17
...
Thread 15: idx=255, ic=0, h=7, w=31
Thread 16: idx=256, ic=0, h=8, w=0
Thread 17: idx=257, ic=0, h=8, w=1
...
For warp 0 (threads 0-31):
Address of thread 15: base + 7*1024 + 31 = base + 7168 + 31 = base + 7199
Address of thread 16: base + 8*1024 + 0 = base + 8192
Gap = 8192 - 7199 = 993 elements = 3972 bytes.
This spans about 62 cache lines of 64 bytes. NOT coalesced at all!
So even with TILE_W=30, the problem remains: within a single warp loading elements with `idx = tid + k * nthreads`, the elements jump across rows.
The fundamental issue with this mapping is: `tid` to `idx` gives contiguous idx, but contiguous idx in [ic][h][w] layout can span row boundaries, causing non-contiguous global memory accesses.
The fix: Process the input tile by iterating over complete [ic][h] slices, where each slice has contiguous w values.
For each (ic, h), there are 32 contiguous w values. We can load each (ic, h) slice with 32 threads (one warp), and each warp loads 32 contiguous floats.
Implementation:
```cpp
int total_slices = IC * in_tile_h; // Each slice is one row of in_tile_w elements
int slice_len = in_tile_w; // = 32 for TILE_W=30
for (int slice_idx = tid / 32; slice_idx < total_slices; slice_idx += nthreads / 32) {
int ic = slice_idx / in_tile_h;
int h = slice_idx % in_tile_h;
int lane = tid % 32;
int w = lane; // Each lane loads one w
if (w < slice_len) {
int ih = oh_base + h;
int iw = ow_base + w;
float val = 0.0f;
if (n < N && ih < H && iw < W) {
val = input[((n * IC + ic) * H + ih) * W + iw];
}
s_input[(ic * in_tile_h + h) * in_tile_w + w] = val;
}
}
```
In this implementation:
- Threads 0-31 form warp 0. tid/32 = 0, lane = 0..31. They load slice 0: (ic=0, h=0), w=0..31. Global addresses: base + 0*1024 + 0 to base + 0*1024 + 31. Contiguous across the warp. Coalesced!
- Threads 32-63 form warp 1. tid/32 = 1, lane = 0..31. They load slice 1: (ic=0, h=1), w=0..31. Global addresses: base + 1*1024 + 0 to base + 1*1024 + 31. Contiguous across the warp. Coalesced!
This is beautiful! Each warp loads exactly one row of 32 contiguous floats.
For a block with 240 threads, we have 240/32 = 7.5 warps. Actually 240 isn't a multiple of 32. 240 = 7*32 + 16. So 7 full warps + 1 partial warp of 16 threads. The partial warp loads slice 7: (ic=0, h=7), w=0..15. That's still fully coalesced (only 16 threads active).
For weight loading, similar approach:
Weight is [OC][IC][K][K]. For our tile, we load [TILE_OC][IC][K][K].
For fixed (oc, ic, kh), kw varies from 0..2. That's only 3 elements, not worth a whole warp.
For weight, total elements = TILE_OC * IC * K * K = 32 * 16 * 9 = 4608.
We can just linearize and have each thread load contiguous elements from global memory. Since weights in global memory are [OC][IC][K][K], and for fixed (oc, ic, kh), kw=0..2 are contiguous, then if we linearize in the same order, contiguous indices map to contiguous memory for the most part.
Actually, with 4608 elements and 240 threads, each thread loads 19.2 elements. If we linearize with `idx = tid + k * nthreads`, thread 0 loads elements 0, 239, 478, ... which are NOT contiguous. But if the layout matches...
Weight global layout: [oc][ic][kh][kw]. So (oc=0, ic=0, kh=0, kw=0) is at offset 0, kw=1 at offset 1, kw=2 at offset 2. Then kh=1 at offset 3, etc. So within a (oc, ic) block, 9 elements are contiguous.
For warp coalescing, we want 32 threads to load 32 contiguous elements. But each (oc, ic) has only 9 elements. So 32 threads would span 3-4 (oc, ic) groups. The memory would be contiguous (since groups are adjacent in global memory), but there might be small gaps if we don't fill the last group completely.
For simplicity, we can load weights the same way: each warp loads 32 contiguous elements using `float` loads. Since 4608 / 32 = 144 slices, and 4608 is divisible by 32, we have exactly 144 slices of 32 contiguous elements.
```cpp
int total_weight_elems = TILE_OC * IC * K * K;
for (int slice = tid / 32; slice < (total_weight_elems + 31) / 32; slice += nthreads / 32) {
int lane = tid % 32;
int idx = slice * 32 + lane;
if (idx < total_weight_elems) {
// decode idx to (oc, ic, kh, kw)
int oc = idx / (IC * K * K);
int rem = idx % (IC * K * K);
int ic = rem / (K * K);
rem = rem % (K * K);
int kh = rem / K;
int kw = rem % K;
float val = 0.0f;
if (oc_base + oc < OC) {
val = weight[(((oc_base + oc) * IC + ic) * K + kh) * K + kw];
}
s_weight[idx] = val;
}
}
```
Each warp loads 32 contiguous floats from global memory (assuming weight layout matches our linearization). This is fully coalesced as long as there are no gaps. Since the total elements per (oc, ic) is 9, slice 0 covers (oc=0, ic=0, kh=0..2, kw=0..2) = 9 elements, then (oc=0, ic=0, kh=... wait, 9 elements for one (oc, ic). So slice 0 covers 3 full (oc, ic) groups + 5 elements of the 4th group. The groups are adjacent in memory, so a 32-element read spans ~4 groups with no gaps (each group is 9 contiguous floats). The warp loads a contiguous 128-byte region. Perfect!
Actually wait, if idx=0 corresponds to (oc=0,ic=0,kh=0,kw=0) and idx=8 is (oc=0,ic=0,kh=2,kw=2), then idx=9 is (oc=0,ic=1,kh=0,kw=0). These ARE contiguous in memory! Because weight is stored as [oc][ic][kh][kw], so (oc=0, ic=0, all kh, kw) are followed by (oc=0, ic=1, all kh, kw). There's no gap.
So yes, any contiguous range of idx maps to contiguous memory. The warp loads are fully coalesced.
Great! This "slice" approach fixes coalescing for both input and weight loading.
Now, for output writing: Each thread writes its computed output values. For TILE_W=30, in a 8x30 block:
- Thread (ty, tx) writes output at (ow_base + tx, oh_base + ty).
- A warp covers threads with consecutive tid. BlockDim.x = 30, so 30 threads per row. But warp size is 32!
- Warp 0: threads 0..29 (ty=0, tx=0..29) and thread 30 (ty=1, tx=0).
- This means the warp writes 30 values in row 0 (contiguous) and 1 value in row 1. This is NOT a clean coalescing pattern.
To fix this, blockDim.x should be a multiple of 32! So 32, 64, etc.
TILE_W = 30 gives blockDim.x = 30, which is not a multiple of 32. This breaks output coalescing.
So let me go back to TILE_W = 32:
Block = 8x32 = 256 threads. BlockDim.x = 32, which is exactly warp size.
Warp 0: threads 0..31 → ty=0, tx=0..31. Writes row 0, columns 0..31. Perfect coalescing.
Warp 1: threads 32..63 → ty=1, tx=0..31. Writes row 1. Perfect.
For input loading with in_tile_w = 34:
Using the slice approach, each warp loads one row of 32 elements (we ignore lanes >= 32). Wait, in_tile_w = 34, so a row has 34 elements. A warp of 32 threads can load 32 elements. The remaining 2 elements (w=32, 33) need another warp or partial warp.
So we can have:
```cpp
int total_slices = IC * in_tile_h * ((in_tile_w + 31) / 32);
for (int slice = tid / 32; slice < total_slices; slice += nthreads / 32) {
int lane = tid % 32;
int w_group = slice % ((in_tile_w + 31) / 32);
int slice_in_ic_h = slice / ((in_tile_w + 31) / 32);
int ic = slice_in_ic_h / in_tile_h;
int h = slice_in_ic_h % in_tile_h;
int w = w_group * 32 + lane;
if (w < in_tile_w) {
// load from global to s_input
}
}
```
For in_tile_w=34:
- w_group can be 0 or 1.
- For w_group=0, w=0..31. Warp loads 32 contiguous floats. Coalesced.
- For w_group=1, w=32..33. Only lanes 0..1 are active. Lanes 2..31 don't participate. This wastes some threads but is still coalesced (the active lanes access contiguous memory).
For each (ic, h), we have 2 slices (32 + 2 elements).
Total slices = 16 * 10 * 2 = 320 slices.
With 8 warps (256 threads), 320/8 = 40 iterations. Each warp handles one slice per iteration.
But each slice is a single global memory transaction (or 2 for w_group=0 since 32 floats = 128 bytes = 2 cache lines). For w_group=1, it's a tiny transaction.
Actually, this seems pretty good. Let me write it out.
Total elements loaded per block = 16*10*34 + 32*16*9 = 5440 + 4608 = 10,048 floats.
With 256 threads, warps = 8.
For input slices: 320 slices. Each warp does 320/8 = 40 slices. But 320 is divisible by 8. Actually, 320 slices total, 8 warps, each warp does 40 slices over the iterations. Good.
For weight slices: 4608 elements / 32 = 144 slices. 144/8 = 18 slices per warp. Good.
So total slices = 320 + 144 = 464. Each warp does 58 slices. With 58 slices per warp and 2 transactions per slice (average), that's 116 transactions per warp. At ~100 cycles per transaction (very rough), that's 11,600 cycles. At 2GHz, that's 6 microseconds. This is the loading time.
Compute time per block: 256 threads * 4608 FMAs = 1.18M FMAs. At FMA throughput... wait, per BLOCK, not per warp.
A block has 8 warps, each with 32 threads doing 4608 FMAs = 147,456 FMAs per warp.
Total 8 * 147,456 = 1,179,648 FMAs per block.
Time per block ≈ 1.18M FMAs / (128 FMAs/cycle/SM / 8 warps * 32... I'm getting confused with per-SM vs per-block).
Actually, per block we have 8 warps. Each warp can execute independently. An SM has 4 schedulers, so it can execute 4 warps simultaneously in terms of instruction issue. But each warp executes on its own data.
If all 8 warps are running on the same SM, they share 4 schedulers, so each scheduler handles 2 warps. Each warp has 4608 FMAs / (some IPC). With ILP of 32, each warp can issue FMAs back-to-back. So each scheduler can issue 1 FMA instruction per cycle (1 for each warp, alternating). 4 schedulers * 32 FMAs = 128 FMAs per cycle per SM.
8 warps * 4608 FMAs = 36,864 FMAs.
At 128 FMAs/cycle, time for the SM to finish all 8 warps = 36,864 / 128 = 288 cycles.
At 2GHz, 288 cycles = 144 ns per block! That's tiny!
Wait, but 1.18M FMAs for the block at 128 FMAs/cycle = 9,216 cycles = 4.6 microseconds. I was counting per-warp before.
9,216 cycles for compute.
Load time: ~6 microseconds = 12,000 cycles.
Total = ~21,000 cycles = 10.5 microseconds per block.
This is much faster than my earlier estimate! If each block takes ~10us, and we run many blocks...
Wait, total FMAs = 308B.
At 128 FMAs/cycle/SM * 132 SMs = 16,896 FMAs/cycle for the whole GPU.
At 2GHz, throughput = 33.79 TFMA/s.
Time = 308B / 33.79B/s = 9.12 ms.
But per block compute time was 4.6us? Let's check:
9.12 ms / 262,144 blocks = 34.8 ns per block if run serially. But blocks run in parallel.
Concurrent blocks on the GPU: let's say 4 per SM * 132 SMs = 528 blocks.
Total waves = 262,144 / 528 = 497 waves.
If each wave takes T time, total = 497 * T.
Time per wave: 528 blocks * 1.18M FMAs = 623M FMAs per wave.
At 33.79 TFMA/s = 33,790M FMAs/s.
Time per wave = 623 / 33,790 seconds = 0.0184 milliseconds = 18.4 microseconds.
But per block we computed ~4.6us (just compute). And load might be another ~6us. With 4 blocks on the same SM sharing resources, the compute time for the wave would be longer than 4.6us. Actually, 4 blocks on one SM have 4 * 8 = 32 warps. At 128 FMAs/cycle per SM, 32 warps * 1.18M FMAs / 4 blocks = 9.4M FMAs per SM per wave. Time = 9.4M / (128 * 2B/s/cycle... wait).
SM throughput for FMAs: 128 per cycle = 128 * 2GHz = 256B FMAs/s = 256M FMAs per microsecond.
Per SM per wave: 4 blocks * 1.18M = 4.72M FMAs.
Time = 4.72M / 256M per us = 18.4 microseconds. Matches!
If memory load is 6us per block and can't be fully hidden: 4 blocks * 6us = 24us memory time, but much of it overlaps with compute.
If 50% hidden: effective overhead = 12us per wave per SM.
Total per wave = max(18.4us, 12us) with overlap = maybe 20-25us.
497 waves * 22us = 10.9 ms.
This is getting into the range where we might be competitive with cuDNN.
But to beat cuDNN, we might need to get down to ~8 ms or less. CuDNN has highly tuned kernels and might use tensor cores or other optimizations.
One big optimization we're missing: using Tensor Cores! H100 has TF32 tensor cores with ~500-1000 TFLOPS. If we could use tensor cores, we would drastically beat the reference if it doesn't use them, or match it if it does.
For TF32 tensor cores, the operation is a matrix multiply. Our conv needs to be reformulated as GEMM. With K=3, IC=16, we need to multiply [M, 144] × [144, 128].
For a small inner dimension of 144, traditional tensor cores might not be optimal because TC operate on tiles like 16x16x8. The inner dimension 144 is not a multiple of 8. But it is a multiple of 16? No, 144 / 16 = 9.
Actually, for Ampere/Hopper, TC can do m16n8k8. So K=144 = 18 * 8. This is a good fit!
M = OH * OW = 1022 * 1022 = 1,044,484
N = OC = 128
K = 144
For TC GEMM:
- M is huge, N=128 is moderate, K=144 is small but divisible by 8.
We can write a simple wmma kernel or use CUTLASS. But `load_inline` might not have CUTLASS headers. We could try to include them, but it's risky.
Alternatively, we can write our own custom implicit GEMM that uses `mma.sync.aligned.m16n8k8.row.col.tf32.tf32.f32` instructions.
This is more advanced. For the first attempt, let me try the direct convolution approach. If it's not fast enough, I'll add tensor core support.
Actually, one more thought: if cuDNN uses TF32 TC, it might run at ~500 TFLOPS, giving ~1.2 ms for this conv. To beat that, we'd need TC too.
But if the problem sets strict fp32 tolerance of 1e-4, AND the reference uses TF32 (which has ~1e-3 relative error), then the tolerance would fail. Unless... the tolerance they gave is what they use to compare, and it might be loose enough for TF32. But 1e-4 is pretty tight.
Actually, in PyTorch, `torch.allclose` with `atol=rtol=1e-4` is the stated tolerance. For TF32 outputs, the difference between TF32 and true FP32 might be around 1e-3 to 1e-2 for this reduction size. So the reference must be using true FP32 for the comparison to make sense.
Or... the test environment disables TF32.
I think the safest assumption is: the reference uses true FP32 arithmetic (or at least something precise enough for 1e-4 tolerance). cuDNN without TF32 or TC achieves maybe 50-100 GFLOPS/SM for FP32, giving total ~6.6-13.2 TFLOPS? No, the theoretical non-TC FP32 is 67 TFLOPS. CuDNN probably achieves 60-80% of this for direct conv, or ~40-55 GFLOPS for the whole GPU.
Wait, non-TC fp32 on H100 IS 67 TFLOPS peak. This is the same as the peak tf32 TC throughput divided by ~7.5. So direct non-TC kernels can use the same FMA units but without the tensor core boost.
If cuDNN uses a non-tensor-core algorithm for this conv for some reason (e.g., small K doesn't benefit from TC), then our custom kernel competes on the same 67 TFLOPS playing field.
Actually, for small K=144, TC might not be as efficient because of the overhead of loading A and B matrices into TC registers. The effective throughput for small K GEMM on TC is often lower than peak. For M=1M, N=128, K=144, the M and N are large enough but K is small.
I think a highly optimized direct convolution or small-K GEMM without TC can be competitive.
Let me also explore another major optimization: processing multiple images per block.
If batch=16, we process 16 images. Currently each block handles 1 image (n is in gridDim.z). What if we process multiple n values in the same block?
- We could load weights once and reuse them across images.
- But input changes per image, so shared memory for input would increase.
- Processing 2 images per block: input shared doubles to 80KB for TILE_H=8, TILE_W=32. Too much.
For N=16, with our current config (z=64), we have plenty of blocks.
Actually, with a tile of 8x32 spatial and 32 OCs, each block handles 8*32*32=8192 outputs. For N=16, that's 131K output channel-spatial tiles. Total outputs = 2.14B. Blocks = 2.14B / 8192 = 261,718. Close to our 262,144.
Let me now write the actual kernel code. I'll use dynamic shared memory with `extern __shared__`. And I'll add the slice-based coalesced loading.
For the tile dimensions, I'll start with TILE_H=8, TILE_W=32, TILE_OC=32. This gives us 256 threads, good coalescing, and 38KB shared memory.
Actually, one more concern: with `extern __shared__`, the compiler doesn't know the size of s_input and s_weight at compile time for indexing optimizations. Using static shared memory might give better performance. But with total under 48KB, I could declare them statically:
```cpp
#define MAX_SINPUT (16 * 10 * 34)
#define MAX_SWEIGHT (32 * 16 * 9)
__shared__ float s_input[MAX_SINPUT];
__shared__ float s_weight[MAX_SWEIGHT];
```
Wait, but MAX_SINPUT would be 5440, MAX_SWEIGHT = 4608. This is fine for static shared memory under 48KB.
With template parameters, I can even make the sizes template-dependent:
```cpp
template<int TILE_H, int TILE_W, int TILE_OC, int K_VAL, int IC_VAL>
__global__ void conv2d_kernel(...) {
const int IN_TILE_H = TILE_H + K_VAL - 1;
const int IN_TILE_W = TILE_W + K_VAL - 1;
__shared__ float s_input[IC_VAL * IN_TILE_H * IN_TILE_W];
__shared__ float s_weight[TILE_OC * IC_VAL * K_VAL * K_VAL];
...
}
```
For TILE_H=8, TILE_W=32, TILE_OC=32, K_VAL=3, IC_VAL=16:
s_input = 16*10*34 = 5440 floats = 21,760 bytes
s_weight = 32*16*9 = 4608 floats = 18,432 bytes
Total static shared = 40,192 bytes = 39.25 KB.
Is this under the default static shared memory limit? For Hopper, I believe the limit might be 48KB. Let me verify.
On Ampere (SM80), the shared memory limit per block is 164KB or 228KB with `cudaFuncSetAttribute`. But some sources say the default without setting attributes is 48KB. Actually, I think the architecture supports up to 164/228KB, but the default max is 48KB for backward compatibility unless you explicitly increase it.
For static shared memory, if the compiler sees `__shared__ float arr[10000]`, and 10000*4 = 40KB, I think it compiles without issues but the kernel might fail to launch if total shared memory exceeds the limit. On newer architectures, there might not be a 48KB compile-time limit, but rather a runtime limit based on occupancy.
To be safe, I could use `extern __shared__` and pass the size at launch. But then I need `cudaFuncSetAttribute` for >48KB. At 39KB, I'm under 48KB, so it should be fine with static shared memory.
Let me try static and see if it compiles. Actually, I should use `extern` to be safe and since dynamic shared memory allows more flexibility.
For the first attempt, I'll use `extern __shared__ float s_mem[];` and slice it.
Now, I also need a generic kernel as fallback for cases where K != 3 or IC != 16. I'll write a simple generic kernel that uses the same approach but without unrolling and without compile-time constants.
Let me draft the complete C++ code.
In the wrapper function `conv2d_forward`:
```cpp
torch::Tensor conv2d_forward(torch::Tensor input, torch::Tensor weight, c10::optional<torch::Tensor> bias,
int64_t stride, int64_t padding, int64_t dilation, int64_t groups) {
// input: [N, IC, H, W]
// weight: [OC, IC/groups, K, K]
int N = input.size(0);
int IC = input.size(1);
int H = input.size(2);
int W = input.size(3);
int OC = weight.size(0);
int K = weight.size(2);
// Calculate output dims
int OH = (H + 2 * padding - dilation * (K - 1) - 1) / stride + 1;
int OW = (W + 2 * padding - dilation * (K - 1) - 1) / stride + 1;
// ... handle stride, padding, dilation
auto output = torch::empty({N, OC, OH, OW}, input.options());
if (K == 3 && stride == 1 && padding == 0 && dilation == 1 && groups == 1 && IC == 16) {
const int TILE_H = 8;
const int TILE_W = 32;
const int TILE_OC = 32;
dim3 block(TILE_W, TILE_H);
dim3 grid((OW + TILE_W - 1) / TILE_W, (OH + TILE_H - 1) / TILE_H, N * ((OC + TILE_OC - 1) / TILE_OC));
size_t smem = (IC * (TILE_H + K - 1) * (TILE_W + K - 1) + TILE_OC * IC * K * K) * sizeof(float);
conv2d_kernel_specialized<TILE_H, TILE_W, TILE_OC, 3, 16><<<grid, block, smem>>>(
input.data_ptr<float>(), weight.data_ptr<float>(),
bias.has_value() ? bias.value().data_ptr<float>() : nullptr,
output.data_ptr<float>(), N, OC, H, W, OH, OW);
} else {
// Generic fallback - we could call at::conv2d or implement a generic kernel
output = at::conv2d(input, weight, bias, {stride, stride}, {padding, padding}, {dilation, dilation}, groups);
}
return output;
}
```
Wait, actually to match the constructor args, we might want to just call at::conv2d for the generic case. But the task says to write a custom CUDA kernel. It's probably fine to fall back to at::conv2d for unsupported parameters.
Also, for the test, the parameters are K=3, stride=1, padding=0, dilation=1, groups=1, IC=16. So our specialized kernel will be used.
Now let me think about whether padding will be applied in our kernel. The test has padding=0.
If padding > 0, we need to handle it in the window coordinates. For the generic fallback, at::conv2d handles it. For the specialized kernel, we can add support or fall back.
Since our test has padding=0, and we're focusing on the test case, I'll add padding support in the generic fallback but for the specialized kernel, I'll just compute the reference if padding is non-zero.
Same for stride and dilation.
Let me write the kernel more carefully.
First, the specialized kernel with slice-based coalesced loading:
```cpp
template<int TILE_H, int TILE_W, int TILE_OC, int K_VAL, int IC_VAL>
__global__ void conv2d_specialized(
const float* __restrict__ input,
const float* __restrict__ weight,
const float* __restrict__ bias,
float* __restrict__ output,
int N, int OC, int H, int W, int OH, int OW) {
const int IN_TILE_H = TILE_H + K_VAL - 1;
const int IN_TILE_W = TILE_W + K_VAL - 1;
int n = blockIdx.z / ((OC + TILE_OC - 1) / TILE_OC);
int oc_block = blockIdx.z % ((OC + TILE_OC - 1) / TILE_OC);
int oh_base = blockIdx.y * TILE_H;
int ow_base = blockIdx.x * TILE_W;
int oc_base = oc_block * TILE_OC;
int tx = threadIdx.x;
int ty = threadIdx.y;
int tid = ty * blockDim.x + tx;
int nthreads = blockDim.x * blockDim.y;
int num_warps = nthreads / 32;
int warp_id = tid / 32;
int lane_id = tid % 32;
extern __shared__ float s_mem[];
float* s_input = s_mem;
float* s_weight = s_mem + IC_VAL * IN_TILE_H * IN_TILE_W;
// Load input with slice-based coalescing
int input_slices_per_row = (IN_TILE_W + 31) / 32;
int total_input_slices = IC_VAL * IN_TILE_H * input_slices_per_row;
for (int slice = warp_id; slice < total_input_slices; slice += num_warps) {
int tmp = slice / input_slices_per_row;
int w_group = slice % input_slices_per_row;
int ic = tmp / IN_TILE_H;
int ih_local = tmp % IN_TILE_H;
int iw_local = w_group * 32 + lane_id;
if (iw_local < IN_TILE_W) {
int ih = oh_base + ih_local;
int iw = ow_base + iw_local;
float val = 0.0f;
if (n < N && ih < H && iw < W) {
val = input[((n * IC_VAL + ic) * H + ih) * W + iw];
}
s_input[(ic * IN_TILE_H + ih_local) * IN_TILE_W + iw_local] = val;
}
}
// Load weights
int weight_total = TILE_OC * IC_VAL * K_VAL * K_VAL;
int weight_slices = (weight_total + 31) / 32;
for (int slice = warp_id; slice < weight_slices; slice += num_warps) {
int idx = slice * 32 + lane_id;
if (idx < weight_total) {
int oc = idx / (IC_VAL * K_VAL * K_VAL);
int rem = idx % (IC_VAL * K_VAL * K_VAL);
int ic = rem / (K_VAL * K_VAL);
rem = rem % (K_VAL * K_VAL);
int kh = rem / K_VAL;
int kw = rem % K_VAL;
float val = 0.0f;
if (oc_base + oc < OC) {
val = weight[(((oc_base + oc) * IC_VAL + ic) * K_VAL + kh) * K_VAL + kw];
}
s_weight[idx] = val;
}
}
__syncthreads();
int oh = oh_base + ty;
int ow = ow_base + tx;
float accum[TILE_OC];
#pragma unroll
for (int i = 0; i < TILE_OC; i++) {
accum[i] = 0.0f;
}
if (oh < OH && ow < OW) {
#pragma unroll
for (int kh = 0; kh < K_VAL; kh++) {
#pragma unroll
for (int kw = 0; kw < K_VAL; kw++) {
int ih_local = ty + kh;
int iw_local = tx + kw;
int in_idx = ih_local * IN_TILE_W + iw_local;
#pragma unroll
for (int ic = 0; ic < IC_VAL; ic++) {
float in_val = s_input[ic * IN_TILE_H * IN_TILE_W + in_idx];
#pragma unroll
for (int oc = 0; oc < TILE_OC; oc++) {
float w_val = s_weight[oc * IC_VAL * K_VAL * K_VAL + ic * K_VAL * K_VAL + kh * K_VAL + kw];
accum[oc] += w_val * in_val;
}
}
}
}
}
// Store output
if (oh < OH && ow < OW) {
#pragma unroll
for (int oc = 0; oc < TILE_OC; oc++) {
if (oc_base + oc < OC) {
float val = accum[oc];
if (bias != nullptr) {
val += bias[oc_base + oc];
}
output[(((n * OC + oc_base + oc) * OH + oh) * OW + ow)] = val;
}
}
}
}
```
Wait, if IC_VAL and K_VAL are compile-time constants, then the compiler CAN unroll the loops. Good.
One issue: `weight_slices` computation at runtime. The loop `for (int slice = warp_id; slice < weight_slices; slice += num_warps)` can't be unrolled because `weight_slices` is a runtime constant derived from template parameters... but actually `weight_slices = (TILE_OC * IC_VAL * K_VAL * K_VAL + 31) / 32` which IS a compile-time constant if all are template parameters!
Yes, since TILE_OC, IC_VAL, K_VAL are all compile-time constants, `weight_total` and `weight_slices` are also compile-time constants. The compiler can unroll these loops too!
Similarly, `total_input_slices = IC_VAL * IN_TILE_H * input_slices_per_row` where `IN_TILE_H = TILE_H + K_VAL - 1` and `input_slices_per_row = (IN_TILE_W + 31) / 32`. These are all compile-time constants. So the loading loops can be unrolled too!
This is excellent. The compiler should generate unrolled code for the entire loading and compute phases.
One more optimization: for output writing, we currently do:
```cpp
for (int oc = 0; oc < TILE_OC; oc++) {
output[...] = val;
}
```
Since `oc_base + oc < OC` is usually true for all oc (except possibly the last block if OC is not divisible by TILE_OC), and in our test OC=128, TILE_OC=32, so it's always true. The compiler might optimize away the branch.
But what if we write output with fewer stores? We can't use vectorized stores easily since different oc are not contiguous. But we could process oc=0..31 and store sequentially. The compiler will probably generate 32 `STG` instructions.
One thought: instead of having each thread compute all 32 OCs, we could have warps specialize in subsets of OCs. For example, warp 0 computes oc=0..7, warp 1 computes oc=8..15, etc. Then each thread has fewer accumulators but we could interleave compute across warps.
But with 32 accumulators already, the FMA pipeline should be well-utilized.
Let me also consider output layout. Output is NCHW: [N][OC][OH][OW]. For fixed n, oc, the spatial pixels are contiguous along W. We write pixel (oh, ow). Since adjacent threads in a warp (same warp for blockDim.x=32) write adjacent W for the same OC, each store is coalesced.
Now, the only issue with 32 stores per thread is instruction overhead. 32 STG instructions per thread. With 256 threads, that's 8192 stores per block. For 262K blocks, about 2.1B stores. Sounds like a lot, but it's the correct number since we have 2.14B output values. Each output value is written exactly once.
Wait, 2.14B outputs / 256 threads = 8.4M stores per thread? No, each thread writes 32 values, but over many blocks. Total stores = 2.14B. The total store bandwidth needed = 8.56GB. At 3.35TB/s, that's negligible compared to compute.
OK, this all seems reasonable. Let me write the complete code.
One thing I forgot: in the weight loading, the layout in shared memory is [oc][ic][kh][kw]. But in the compute loop, we access as:
```cpp
s_weight[oc * IC_VAL * K_VAL * K_VAL + ic * K_VAL * K_VAL + kh * K_VAL + kw]
```
This matches the linearization. Good.
Let me double-check input indexing. For s_input layout [ic][ih][iw]:
```cpp
s_input[ic * IN_TILE_H * IN_TILE_W + ih_local * IN_TILE_W + iw_local]
```
And we store at that index during loading. During compute:
```cpp
int in_idx = ih_local * IN_TILE_W + iw_local;
float in_val = s_input[ic * IN_TILE_H * IN_TILE_W + in_idx];
```
This is correct.
Now, for the generic fallback, I can just use `at::conv2d`. But to be more complete, let me also write a generic kernel without templates.
Actually, if we're writing a custom CUDA kernel, using `at::conv2d` as fallback is fine.
One more consideration: `c10::optional` for bias. In the C++ API, bias might be an empty optional. Let me use `at::Tensor bias_ = {}` or similar.
Looking at PyTorch C++ API, `at::conv2d` takes optional bias. In our custom wrapper, we can accept `at::Tensor` or `c10::optional<at::Tensor>`.
For simplicity in the Python binding:
```cpp
torch::Tensor conv2d_forward(torch::Tensor input, torch::Tensor weight, c10::optional<torch::Tensor> bias, ...)
```
Actually, in pybind11/torch binding, optional tensors are tricky. I can just accept `torch::Tensor bias` and check `bias.defined()`. Or use `at::OptionalTensorRef`.
Let me use:
```cpp
torch::Tensor conv2d_forward(
torch::Tensor input,
torch::Tensor weight,
c10::optional<torch::Tensor> bias,
int64_t stride,
int64_t padding,
int64_t dilation,
int64_t groups)
```
And in the kernel call, pass `bias.has_value() ? bias.value().data_ptr<float>() : nullptr`.
Wait, but `load_inline` with `functions=["conv2d_forward"]` will bind a function that matches this signature. From Python, we'd call it as `conv2d_forward(input, weight, bias, stride, padding, dilation, groups)`. If bias is None, torch handles optional tensors gracefully.
In ModelNew:
```python
return conv2d_forward(x, self.conv2d.weight, self.conv2d.bias, self.conv2d.stride[0], ...)
```
Actually, `self.conv2d.stride` is a tuple like (1, 1). We need `self.conv2d.stride[0]`.
For padding, same thing.
Now, what about `groups`? In the test, groups=1. We pass `self.conv2d.groups`.
For the generic fallback, `at::conv2d` handles groups. In our specialized kernel, we assume groups=1.
Also, we need to handle the case where input is not contiguous, or not on CUDA. Let's assume input is contiguous CUDA float tensor, which is the case for the test.
Actually, `at::conv2d` in fallback handles all cases.
One more thing: AMP or autocast. The input might be fp32, and that's what we're targeting.
Let me now write the complete Python module.
I think I'll start with compiling and running to see if it works, then optimize.