reasoning
We need to write an efficient CUDA kernel for 3D Average Pooling. The reference is `nn.AvgPool3d` with kernel_size=3, stride=2, padding=1 on inputs of shape (16, 32, 128, 128, 256). Output shape: Let's compute. Dout = floor((128 + 2*1 - 3)/2) + 1 = floor(127/2)+1 = 63+1 = 64. Hout = same = 64. Wout = floor((256 + 2*1 - 3)/2)+1 = floor(255/2)+1 = 127+1 = 128. So output shape is (16, 32, 64, 64, 128). Total output elements = 16*32*64*64*128 ≈ 268.4 million elements. That's quite large. Each output element averages 3*3*3 = 27 input elements. So total input reads per naive approach: 268M * 27 ≈ 7.2B reads. In fp32 that's ~28.8 GB of reads. The kernel needs to be efficient.
Since stride=2, kernel_size=3, padding=1, this is like a 3x3x3 pooling with stride 2. The input shape is (16, 32, 128, 128, 256). We need to handle the boundaries due to padding. Padding is only 1 on each side.
We can write a tile-based kernel that loads a tile of input into shared memory and computes the averages. However, with large spatial dims and relatively small channels (32), maybe a simple per-element kernel iterating over the kernel region directly from global memory is fine, but we need to beat PyTorch's optimized implementation. PyTorch's AvgPool3d on H100 might be decent but not necessarily fully optimized for these specific sizes. Since output is ~268M elements, we spawn enough threads.
But global memory reads are the bottleneck. Each output reads 27 values. However, there is significant overlap in input reads between adjacent outputs because stride=2 and kernel=3. For example, output (d,h,w) uses input pixels from d*2-1 to d*2+1, etc. Adjacent outputs share a lot of input. So a tile-based approach with shared memory could reduce global memory bandwidth roughly by a factor of kernel_size^3 / stride^3 = 27/8 = 3.375x. It might be beneficial.
However, managing 3D tiles in shared memory is tricky because of the large dimensions. Let's see memory size: If we process a tile of outputs of size, say, 8x8x8, then the input tile needed is (8*2+1) x (8*2+1) x (8*2+1) = 17x17x17, times 32 channels. In fp32, 17*17*17*32*4 = 17^3*128 = 4913*128 ≈ 629KB. That's too large for shared memory per block (max 227KB). We could process fewer outputs or fewer channels per block.
Alternatively, we could process one channel at a time or a small number of channels (e.g., 8). Let's try tile of 8x8x16 outputs. Input tile: (8*2+1)=17 in D and H, (16*2+1)=33 in W. 17*17*33 ≈ 9537. Times 8 channels * 4 bytes ≈ 305KB, still maybe too high. If we use 4 channels: 153KB, okay. But then each block only processes 8*8*16*4 = 4096 outputs, which is not many. With 268M outputs, we'd need ~65536 blocks. That's fine.
Alternatively, we can use registers and vectorized loads, or simply compute directly from global memory with coalesced reads. Given the H100's high memory bandwidth (3.35 TB/s), the naive kernel might be memory-bound but still fast. 28.8 GB of reads and writes (~1 GB of output) gives total ~30 GB. At 3.35 TB/s, that's ~9 ms. That's not super fast but decent. But PyTorch's implementation likely also does similar. To beat it, we need to be clever.
Actually, with 3D average pooling, we can think of it as two operations: average over depth, then 2D average pool over HxW? Actually AvgPool3d with kernel_size K and stride S is not separable in the same way because it's just averaging over a 3D window. But we can write the kernel to be separable: first average over depth (1D), then over spatial (2D). But that requires intermediate storage. For a single kernel call, it's not separable unless we use a nested loop.
Wait, `nn.AvgPool3d` on GPU likely uses a CUDA kernel similar to `avg_pool3d_out_cuda` in PyTorch. We can check if there are known inefficiencies, but we don't have source access here. We need to write a kernel that is faster.
Let's think about the memory access pattern.
Input shape: (N, C, D, H, W) = (16, 32, 128, 128, 256)
Output shape: (N, C, D_out, H_out, W_out) = (16, 32, 64, 64, 128)
For each output element (n, c, d, h, w):
sum over kd=0..2, kh=0..2, kw=0..2:
id = d*stride + kd - padding = d*2 + kd - 1
ih = h*2 + kh - 1
iw = w*2 + kw - 1
if inside bounds, add x[n,c,id,ih,iw], and count valid neighbors.
Because padding=1, for output (0,0,0): id in {-1,0,1}. -1 is out of bounds. So count is 18 (missing one plane of 9 pixels) for corner cases. So we need to handle boundary conditions. PyTorch's AvgPool3d divides by the count of valid elements (including padding), not kernel_size^3. The default `count_include_pad=True` actually counts padding zeros in the denominator? Wait, `nn.AvgPool3d` default is `count_include_pad=True`. Let's verify. Actually PyTorch's `AvgPool3d` documentation: `count_include_pad` defaults to True. That means when computing the average, padding elements are included in the count. Wait, if `count_include_pad=True`, does it mean the denominator is always kernel_size^3? Actually, with padding, the count is kernel_size^3 if the window overlaps any padding? Let's check. From PyTorch docs: "When True, will include the zero-padding in the averaging calculation." So yes, the denominator is `kernel_size[0] * kernel_size[1] * kernel_size[2]`, which is 27, for all output elements, even if some input indices are out of bounds (the zeros from padding are included). Wait, but if the pad is smaller than kernel_size, and the window extends beyond the padded region? No, padding is applied to the input, so the padded input has size D+2*pad etc. The indices never go out of bounds of the original input? Wait, AvgPool3d pads the input first. So the actual input to pooling is padded. If count_include_pad=True, the denominator is kernel_size^3. If False, it's the count of non-padding elements. But wait, the reference code uses default `nn.AvgPool3d(...)` with no `count_include_pad` argument, so default is True. Yes.
So denominator is always 27. That's great! It simplifies the kernel. But note: if padding=1 and stride=2 and output size 64, does the padded input size become 130? The last output index d=63 gives id = 63*2 + kd -1 = 125 + kd. With kd max=2, id=127. Original D=128, so id ranges 0..127. So with padding=1, the needed input indices after padding are from 0..129. But the window only accesses indices -1 to 127 before padding? Wait, in PyTorch, padding is applied to the input, meaning the input is conceptually padded with zeros, and then the pooling operates on the padded tensor. But in the standard implementation, the kernel directly accesses the original tensor with boundary checks. If the index is out of bounds, it is treated as zero. And with count_include_pad=True, the denominator is always 27.
Thus in our kernel, for each output, we sum x[n,c,id,ih,iw] if in bounds, else add 0. Then divide by 27.
This simplifies things because we don't need to count valid elements. However, we still need to check bounds. But since padding=1, only the first and last planes have out-of-bounds accesses. For depth: d=0 => id=-1,0,1 (one invalid). d>=1 and d<=62 => id in [2*d-1, 2*d+1] which is >=1 and <=125, all valid. d=63 => id=125,126,127 (all valid). Actually D=128, indices 0..127. So for d=0: id=-1 is invalid, 0 and 1 valid. For d=1: id=1,2,3 valid. So only d=0 has an invalid depth index. Similarly for h=0: ih=-1 invalid. For w=0: iw=-1 invalid. For the last ones: d=63, id=125,126,127 all valid. h=63, ih=125,126,127 all valid. w=127, iw=253,254,255 all valid. Wait w=127: iw=127*2+kw-1 = 254,255,256. kw=2 => iw=256 which is out of bounds (W=256, indices 0..255). So w=127 has one invalid index. Similarly for h=63: ih=125+2-1=126? Wait h=63 => ih = 63*2 + kh - 1 = 126 + kh - 1 = 125 + kh. For kh=0:125, kh=1:126, kh=2:127. H=128, indices 0..127, so 127 is valid. So h=63 is all valid. d=63: id = 63*2 + kd -1 = 126 + kd -1 = 125+kd. kd=0:125, kd=1:126, kd=2:127. D=128, valid. So only the first plane (d=0,h=0,w=0) in each dimension has invalid indices for the negative side. For the positive side, w=127 has one invalid (256). Let's check: w_max = floor((W+2*padding - kernel_size)/stride) = floor((256+2-3)/2) = floor(255/2)=127. So w ranges 0..127. For w=127: iw_start=127*2-1=253. iw=253,254,255. All valid (W=256). Wait, I earlier computed kw index: iw = w*stride + kw - padding = 127*2 + kw - 1 = 253 + kw. So kw=0->253, kw=1->254, kw=2->255. All valid. So why did I think 256? That's wrong. Because padding=1, stride=2, w*stride=254. 254+2-1=255. All valid.
Let's check h: h ranges 0..63. h=63: ih=63*2+kh-1 = 125+kh. Max 127. Valid. d=0: id=-1,0,1 (one invalid). h=0: ih=-1,0,1 (one invalid). w=0: iw=-1,0,1 (one invalid).
So the only out-of-bounds accesses are on the "front" face of the 3D grid (d=0 or h=0 or w=0). Actually all three boundaries at the low end have an invalid plane. No invalid accesses at the high end because output dimensions are chosen such that the last window fits exactly.
Wait, D_out = floor((128+2-3)/2)+1 = 64. So d=0..63. The last valid id for d=63 is 127. So the high end is exactly covered. H same. W: 256+2-3 = 255. floor(255/2) = 127. 127+1=128. w=0..127. iw max = 127*2 + 2 - 1 = 255. W=256, max index 255. Perfect fit. So only low-side boundaries have invalid indices.
This means we can have a fast path for the general case where we know all indices are valid, and only handle boundaries for d=0 or h=0 or w=0. Since output is 64x64x128, the boundary elements are a tiny fraction: d=0 is 1*64*128 = 8192 out of 64*64*128 = 524288 per sample/channel. That's ~1.56%. h=0 is 64*1*128 = 8192. w=0 is 64*64*1 = 4096. Overlap at edges. So overall, less than 5% of elements are boundary. Most elements can skip bounds checks! That's a great optimization.
For the general case (d>=1, h>=1, w>=1), the read indices are always in bounds. So we can write a kernel that checks if the thread coordinates are in the "fast" region and then do unrolled reads from global memory directly.
Given the kernel is small (3x3x3), we can unroll the loops.
Let's compute arithmetic intensity: Each output does 27 loads, 26 adds, 1 multiply (by 1/27). In fp32. So about 27 FLOPs per element? Actually adds are FLOPs too. 26 adds + 1 mul = 27 FLOPs. Memory traffic: 27 loads * 4 bytes = 108 bytes read per output, 1 store * 4 bytes = 4 bytes write. The loads are not cached well if each thread reads a unique 27-element neighborhood with stride 2. Since adjacent outputs share many input elements, caching in L2 might help, but still, 108 bytes per output is high. So the kernel is heavily memory bandwidth bound. To improve performance, we need to leverage shared memory or register tiling to reduce global memory bandwidth.
Given that the kernel size is 3x3x3=27 and stride is 2, each input element is reused by up to 8 output elements (since stride=2 in each dimension, each input is in at most 2x2x2=8 windows? Actually let's think: an input at (id,ih,iw) contributes to outputs at d such that d*2 + kd -1 = id => kd = id - d*2 +1. Since kd in {0,1,2}, there are up to 2 values of d for each id? Let's solve: d = floor((id+1)/2) or ceil((id-1)/2)? Actually for a fixed id, the output d must satisfy d*2 <= id+1 <= d*2+2. So there are at most 2 possible d values. Similarly for h and w. So each input is used by at most 2*2*2 = 8 outputs. So there is an 8x reuse factor. If we read each input from global memory only once and share among outputs in a tile, we reduce reads by ~8x. But naive global memory approach reads each input multiple times; however, the L2 cache might capture some reuse. But because we have many output threads, locality depends on thread scheduling.
A shared memory tile approach can explicitly exploit this reuse. However, the complexity of 3D shared memory tiling is high. We need to load an input tile that corresponds to an output tile. For a 3D output tile of size (Td, Th, Tw), the input tile size is (2*Td+1, 2*Th+1, 2*Tw+1). If we choose Td=4, Th=4, Tw=8, input tile is 9x9x17 = 1377. For 8 channels, that's 1377*8*4=44KB. That's fine. Output tile size is 4*4*8 = 128 per channel. Wait, 128*8 channels = 1024 outputs per block. Each thread could handle one output. 1024 threads per block is okay.
But the output count per block is small. For 16*32*64*64*128 = 268,435,456 outputs, with 1024 outputs per block, we need 262,144 blocks. That's a lot but manageable. But each block only processes a small tile. The overhead of loading the input tile into shared memory and synchronizing might be non-negligible. However, H100 has high shared memory bandwidth, so it might be okay.
Alternatively, we can process a larger output tile per block by using more threads or having each thread compute multiple outputs. E.g., each thread computes 2 outputs, blockDim = 512, output tile = 8x8x8 = 512 per channel? Actually 512 threads * 2 outputs = 1024 outputs. Let's think of using block size 256 threads, each thread computes 8 outputs. Output tile = 256*8 = 2048. If we process C=4 channels per block, that's 8192 outputs. But then shared memory for input tile: if output tile is, say, 16x8x16 (?) Hard to map.
Let me think of a 1D thread layout. Each thread computes one output element (n,c,d,h,w). The kernel launches with enough threads to cover all outputs. Total outputs = 16*32*64*64*128 = 2^4 * 2^5 * 2^6 * 2^6 * 2^7 = 2^28 = 268,435,456. That's 2^28. With block size 256 threads, that's 1,048,576 blocks. That's a lot but within limits.
Each thread reads 27 floats. If threads in a warp are adjacent in the last dimension (w), then their input windows overlap significantly. For example, thread for w and w+1: windows are [2w-1, 2w+1] and [2w+1, 2w+3]. They share 2 elements in the w dimension. In the h and d dimensions, if they are adjacent in h or d, windows overlap by 2. But with a 1D block layout mapping to (n,c,d,h,w) in row-major order, adjacent threads will likely be adjacent in w, then h, then d, etc. So for a warp of 32 threads, they will cover consecutive w. So w windows will be [2w-1, 2w+1], [2w+1, ...], etc. So adjacent threads share a lot of w elements. But reads are not perfectly coalesced because each thread reads 27 scattered values? Let's see: The first read for each thread is at iw = 2*w - 1 (if no padding). For w varying by 1, iw varies by 2. So adjacent threads read at addresses offset by 2 floats. That's still coalesced (but not fully sequential). Actually if threads are adjacent in w, the addresses they access are contiguous or strided by 2. The memory controller on H100 can handle that well. But the issue is that each thread accesses 27 addresses, many of which are far apart in memory (different h and d planes). The memory access pattern is strided across h and d dimensions, which might be less efficient.
However, given the large size, I suspect PyTorch's built-in kernel uses a similar approach. To beat it, we might need to either:
1. Use tensor cores or vectorized loads.
2. Use a custom algorithm that fuses channels or uses registers more efficiently.
3. Use the fact that we can avoid bounds checks for most elements.
Actually, looking at the problem, PyTorch's AvgPool3d might be implemented generically and not highly optimized for specific small kernel sizes. A custom kernel unrolled for kernel_size=3, stride=2, padding=1 could be much faster because the compiler can optimize the unrolled loops and we can skip unnecessary indexing logic.
Also, the input size is fairly large: 16*32*128*128*256*4 bytes = 16*32*2^7*2^7*2^8*4 = 16*32*2^22*4 = 512*2^22*4? Let's compute: 128*128*256 = 2^7 * 2^7 * 2^8 = 2^22 = 4,194,304. Times 32 = 134,217,728. Times 16 = 2,147,483,648. Times 4 bytes = 8,589,934,592 bytes = 8 GB. That's a large input. Writing a kernel that processes it efficiently is important.
Wait, the batch size is 16, channels 32, depth 128, height 128, width 256. That's 16*32*128*128*256 = 2,147,483,648 elements. In fp32, that's 8 GB. The output is 16*32*64*64*128 = 268,435,456 elements = 1 GB.
The H100 HBM3 bandwidth is 3.35 TB/s. So reading 8GB and writing 1GB takes about 2.7 ms if perfectly utilizing bandwidth. But naive kernel reads each element multiple times. If each input element is reused ~8 times by L2 cache, then effective read traffic is still ~8GB (if cache hits perfectly). In practice, L2 cache might not perfectly capture all reuse because the working set is huge. The input for one sample/channel is 128*128*256*4 = 16MB. L2 cache on H100 is 50MB. So all channels for one sample: 32*16MB = 512MB. That's larger than L2. But the kernel processes channels in order. The working set for a block of threads might be small enough to fit in L2, but not L1. Since each input is reused by ~8 outputs, and outputs are computed in order, the L2 cache might capture some reuse but not all.
A shared memory tile approach could reduce traffic to DRAM. Let's explore that.
**Shared Memory Tiling Approach**
We can process the 5D tensor with 3D thread blocks. Let's map:
- Each block handles a tile of (C_tile, D_tile, H_tile, W_tile) outputs.
- Threads are arranged in a 3D grid inside the block: maybe (TILE_C, TILE_D, TILE_H, TILE_W) but that's 4D. We can reduce by having each thread compute multiple W outputs.
Better: process channels sequentially (or in small groups) and use 3D spatial tiling. For each block, we fix the batch index n and a channel range c0..c0+C_tile-1. Then the block processes a spatial tile (d0..d0+Td-1, h0..h0+Th-1, w0..w0+Tw-1).
If we set C_tile = 1 to keep shared memory low. Then each thread can compute one (d,h,w) output. Block size = Td * Th * Tw.
Shared memory needed: Input tile for these channels: (2*Td+1) * (2*Th+1) * (2*Tw+1) * C_tile floats. If we set Td=4, Th=4, Tw=16:
Input tile = 9*9*33 = 2673 floats. *4 bytes = ~10.7KB. Block size = 4*4*16=256 threads. That's comfortable. For each sample and channel, we would launch (64/4)*(64/4)*(128/16) = 16*16*8 = 2048 blocks per channel. Times 32 channels * 16 samples = 1,048,576 blocks. That's exactly the same number of blocks as the naive 1D approach with block size 256. But now each block loads its input tile into shared memory once and reuses it.
Wait, output per block: 4*4*16 = 256 outputs. Each output uses 27 inputs. Total input accesses if not cached = 256*27 = 6912 loads. But the input tile has 2673 unique inputs. So the "theoretical" reuse factor is 6912/2673 ≈ 2.59x. That's not huge because the tile is small and the kernel size is large relative to stride. If we make Tw larger, say Tw=32, then input tile = 9*9*65 = 5265 floats. Output = 4*4*32 = 512. 512*27 = 13824 loads. Ratio = 13824/5265 ≈ 2.62x. Similar. If we make Td=8, Th=8, Tw=8: input = 17*17*17 = 4913. Output = 512. Ratio = 512*27/4913 = 13824/4913 ≈ 2.81x. If Td=8,Th=8,Tw=16: input=17*17*33=9537. Output=1024. Ratio=1024*27/9537≈2.9x. The theoretical max reuse factor is 8 (since each input is used by up to 8 outputs). To get closer to 8, the tile must be large enough that interior inputs are used by 8 outputs. For Td=8, the interior outputs (not on the boundary of the tile) are those with d in [1,6], h in [1,6], w in [1,14] (for Tw=16). The number of interior outputs = 6*6*14 = 504. Inputs for interior outputs are fully reused 8 times? Wait, each input contributes to up to 8 outputs. In the tile, an input in the interior is needed by outputs in all 8 surrounding windows. If we have enough outputs, all 8 uses are within the block. So the ratio could approach 8 for large tiles. But our tile is limited by shared memory. For Td=8,Th=8,Tw=32: input=17*17*65=18785 floats. *4=~75KB. Output=8*8*32=2048. 2048*27/18785 ≈ 2.94x. Not great because the boundary is large relative to volume.
To get better reuse, we need very large tiles. But SM shared memory is only 227KB. So using larger C_tile might be better than larger spatial.
Alternatively, we can use a different strategy: **Register blocking + warp-shuffle** or process multiple outputs per thread to increase the arithmetic intensity and let the L2 cache handle reuse. Actually, the most practical approach to beat a generic kernel is to write a highly specialized, fully unrolled kernel with minimal indexing overhead and no dynamic loops, and process vectorized loads.
Since kernel_size=3, we can unroll the kd, kh, kw loops. And we can use float4 vectorized loads for the width dimension. Width is 256, which is divisible by 4. For a given (n,c,d,h), the 3x3x3 window consists of 3 depth planes, each with 3 height rows, each with 3 width elements. We can load 3 floats per row. If adjacent threads in a warp handle adjacent w, then their memory accesses for a given (kd,kh) are at iw = 2*w + kw - 1. For kw=0,1,2, threads read at addresses: (base + 2*w -1), (base + 2*w), (base + 2*w +1). For w varying by 1 across warp, the addresses for kw=0 are strided by 2, kw=1 strided by 2, kw=2 strided by 2. We could use vector loads if we reorganize.
But maybe a simpler approach: Use the fact that the kernel is basically a strided convolution with a constant kernel (all ones / 27). We could implement it by first applying a 1x1x3 moving sum along W, then 1x3x1 along H, then 3x1x1 along D? Actually average pooling with kernel 3 and stride 2 is not separable in the sense of standard convolution because of the stride. Wait, can we decompose stride-2 average pool into separable operations? A 3D moving average with stride 2 can be written as: first compute a sliding window sum along W with window size 3 (stride 1), then along H with window size 3 (stride 1), then along D with window size 3 (stride 1), then subsample by stride 2 in all dimensions. But that's not correct because the sum over the 3D cube is the same as sum over D of sum over H of sum over W. So we can do:
1. Sum along W: compute `tmp1[n,c,d,h,w] = x[n,c,d,h,w-1] + x[n,c,d,h,w] + x[n,c,d,h,w+1]`, for valid w. This is a 1D convolution with stride 1, kernel [1,1,1], padding=1.
2. Sum along H: `tmp2[n,c,d,h,w] = tmp1[n,c,d,h-1,w] + tmp1[n,c,d,h,w] + tmp1[n,c,d,h+1,w]`.
3. Sum along D: `tmp3[n,c,d,h,w] = tmp2[n,c,d-1,h,w] + tmp2[n,c,d,h,w] + tmp2[n,c,d+1,h,w]`.
4. Subsample: `out[n,c,d,h,w] = tmp3[n,c,2*d,2*h,2*w] / 27`.
This is a sequence of 3 separable passes, each with stride 1 and then a final stride-2 subsample. The advantage is that each intermediate pass has much better memory locality and can be implemented with very simple 1D convolution kernels. The total number of operations is: W pass: each element reads 3 values, adds them => 2 adds per element. H pass: 3 reads, 2 adds. D pass: 3 reads, 2 adds. Total reads = 9 per original element, not 27. Plus writes for intermediate buffers. And finally a subsample read + divide.
But we need intermediate storage. tmp1, tmp2, tmp3 each have size of padded input: (N,C,D,H,W) with padding 1 in the dimension being summed? Actually for sum along W with padding 1, output size W (since stride 1). But we can preserve the same spatial dims (128,128,256) for all intermediates if we pad. Padding 1 in W: W becomes 258? Actually if we pad by 1 on each side, size becomes W+2 = 258. But if we just handle boundaries by zero, we can keep the same shape. Let's see:
Pass 1 (sum W): For each (n,c,d,h), compute sum over w' = w-1, w, w+1. This gives a tensor of shape (N,C,D,H,W). It uses 3 loads and 2 adds per element.
Pass 2 (sum H): For each (n,c,d,w) on the result of pass1, sum over h-1,h,h+1. Another (N,C,D,H,W) tensor.
Pass 3 (sum D): For each (n,c,h,w) on result of pass2, sum over d-1,d,d+1. Another (N,C,D,H,W) tensor.
Pass 4 (subsample): Read every 2nd element in each dim and divide by 27.
Total memory traffic:
- Read input: 8GB
- Write tmp1: 8GB (plus pad?) actually same size. But wait, input is (16,32,128,128,256) = 2^31 bytes = 2GB? Wait, I miscalculated earlier! Let's recompute: 16*32 = 512. 128*128 = 16384. 16384*256 = 4,194,304. Multiply by 512 = 2,147,483,648 elements. Times 4 bytes = 8,589,934,592 bytes = 8GB. Yes, 8GB. The output is 1/8 of that? Let's see: output D*H*W = 64*64*128 = 524,288. Input D*H*W = 128*128*256 = 4,194,304. Ratio = 524288/4194304 = 1/8. So output = 8GB / 8 = 1GB. Yes.
Intermediate tmp1, tmp2, tmp3 would each be 8GB if full spatial dims. That's 24GB of allocations, which might be heavy. We could also do the passes in-place? No, because each pass depends on the previous. But we can reduce memory by fusing passes or only padding in one dimension at a time.
Actually we don't need to store the full 8GB intermediate. For W-pass, we can compute line by line and store. But launching multiple kernels has overhead.
Alternatively, we can fuse the W and H passes. For each (n,c,d), load a 2D slice of HxW, compute W-sums and then H-sums on the fly. Then store a (H,W) plane of tmp2. Then do D-pass on tmp2 slices. This reduces intermediates.
But maybe there's an even simpler approach: Since the kernel weights are all 1, we can compute the output by loading 3x3x3 elements directly. And since the kernel is small, unrolling is the best optimization. The main cost is global memory reads. Perhaps we can improve cache efficiency by processing the tensor in a block that traverses W contiguously. For each thread, we compute multiple W outputs (e.g., 4 or 8) so that the loads for the first W output are reused for the next? Not exactly, but adjacent W outputs share 2 out of 3 W elements. If a thread serially computes 4 consecutive W outputs, it can keep some values in registers across the loop. Let's explore this.
Suppose a thread computes outputs for fixed (n,c,d,h) and a range of w values, say w_start to w_start+W_unroll-1. For w, the input window in W is [2w-1, 2w, 2w+1]. For w+1, it's [2w+1, 2w+2, 2w+3]. So the loads for the next output share 2 elements (at 2w+1 and maybe 2w+2?). Actually w window: A=w-1, B=w, C=w+1 (relative to 2w). For w+1: A'=(w+1)-1 = w, B'=w+1, C'=w+2. So it shares B and C with the next? Let's be precise: indices are absolute: for output w: [2w-1, 2w, 2w+1]. For output w+1: [2w+1, 2w+2, 2w+3]. So they share exactly one element: 2w+1. So serializing across W gives 2 new loads per next output instead of 3. That's a 33% saving in W-dimension loads. Similarly across H: output h uses [2h-1, 2h, 2h+1]. Output h+1 shares 2h+1. So if we process a 2D tile of HxW outputs per thread, we can reuse some loads across H as well. Across D: output d uses [2d-1, 2d, 2d+1]. Output d+1 shares 2d+1.
If we process a small 3D tile of outputs per thread (e.g., 2x2x2 = 8 outputs), we can keep many values in registers and reuse loads. Each output requires 27 loads. Without reuse, 8 outputs = 216 loads. With reuse, the 8 outputs correspond to a region of input of size (2*2+1)=5 in D, H, W? Actually for outputs d in [d0,d0+1], h in [h0,h0+1], w in [w0,w0+1], the input region is [2d0-1, 2d0+3] in D (size 5), same for H and W: total 125 inputs. Each input is used by up to 8 outputs (2x2x2). So total unique loads = 125. Compare to naive: 8*27 = 216. So savings = 125/216 ≈ 42% reduction. Not as good as full tile in shared memory, but register-level reuse can be efficient and avoid shared memory overhead.
But perhaps the most effective optimization is to write the kernel to process the W dimension with vectorized loads and use warp-level primitives to share data. However, that gets complex.
Let's step back. The problem asks for a custom CUDA kernel via `load_inline`. It doesn't have to be the absolute fastest, just faster than the reference. On H100, PyTorch's generic `AvgPool3d` might have moderate performance. A well-written custom kernel with unrolled loops and careful indexing might already be faster.
But we should also consider the possibility that PyTorch's `AvgPool3d` is already quite optimized. Let's check typical performance: For 3D avg pool on large tensors, PyTorch uses the ATen library which has vectorized kernels. However, generic kernels often have loop overhead for arbitrary kernel sizes. Since our kernel size=3, stride=2, we can specialize and unroll completely.
Let's design the kernel.
**Kernel Design Option 1: Naive unrolled per-element**
Launch one thread per output element. The thread calculates its (n,c,d,h,w) indices. Then do:
```cpp
float sum = 0.0f;
#pragma unroll
for (int kd = 0; kd < 3; ++kd) {
int id = d * 2 + kd - 1;
bool vd = (id >= 0 && id < 128);
#pragma unroll
for (int kh = 0; kh < 3; ++kh) {
int ih = h * 2 + kh - 1;
bool vh = (ih >= 0 && ih < 128);
#pragma unroll
for (int kw = 0; kw < 3; ++kw) {
int iw = w * 2 + kw - 1;
bool vw = (iw >= 0 && iw < 256);
if (vd && vh && vw) {
sum += in[INDEX(n,c,id,ih,iw)];
}
}
}
}
out[idx] = sum / 27.0f;
```
But we can skip bounds checks for the fast path. We can check if n,c,d,h,w are in the interior: d >= 1 && h >= 1 && w >= 1. Wait, we need to check the input bounds. Input depth is 128, height 128, width 256. Output d ranges 0..63. For d=0, id can be -1 (invalid). For d>=1, id >= 1 (since id = 2*d + kd -1, min kd=0 => 2*d-1 >= 1). Max id for d=63: 126+2-1=127. So for d>=1, id is always valid. For h>=1, ih valid. For w>=1, iw valid. So fast path: if (d > 0 && h > 0 && w > 0) { unrolled loads without bounds checks }. Else { handle edges }.
But `d > 0`: d is an int in [0,63]. So for the vast majority (63/64 * 63/64 * 127/128 ≈ 96.6%) of elements, we can skip bounds checks. Actually d>0 and h>0 and w>0 includes d in 1..63 (63 values), h 1..63 (63), w 1..127 (127). Fraction = (63*63*127)/(64*64*128) ≈ 0.966. So ~96.6% fast path, 3.4% slow path.
In the fast path, we can do fully unrolled loads and adds. The compiler might be able to vectorize some loads if the addresses are aligned.
But contiguous threads reading with stride 2 in W may cause some divergence. However, addressing is deterministic.
Actually, we can also do a small optimization: reorder loops to iterate W innermost for better memory coalescing. But since each thread handles a unique output, memory coalescing for output writes is perfect if we map threads to contiguous output memory. For input reads, each thread reads 27 values scattered across memory. Not great.
**Kernel Design Option 2: Shared memory tile**
Let's calculate occupancy for Option 1. Block size 256 threads, registers per thread maybe ~32. Occupancy high. It's simple.
For Option 2, block size maybe 256 threads, shared memory maybe 20KB. Occupancy also high.
Let's try Option 2 with a modest tile size. The challenge is writing a correct kernel. Let's try to write a 3D shared memory tile.
We can map:
- blockDim.x = Tw (e.g., 16)
- blockDim.y = Th (e.g., 4)
- blockDim.z = Td (e.g., 4)
Total threads = 4*4*16 = 256.
Each thread computes one output (d,h,w) for a given (n,c). We process channels in the outer loops: gridDim.x = W_out / Tw, gridDim.y = H_out / Th, gridDim.z = D_out / Td. But we also need to iterate over batch and channels. We can map them to additional grid dimensions or use loops inside the kernel. Since CUDA supports up to 3 grid dimensions, we can map:
- blockIdx.x -> tile in W
- blockIdx.y -> tile in H
- blockIdx.z -> tile in D
But then we need loops for n and c. Alternatively, we can have:
- blockIdx.x -> n
- blockIdx.y -> c
- blockIdx.z -> linearized spatial tile index.
Or we can use a 1D block index for spatial and compute n,c from it.
Let's think about memory layout. Input is NCDHW. Within a channel, data is stored as D*H*W contiguous with W innermost. So for a fixed (n,c,d,h), the W dimension is contiguous. That means if threads in a block differ in w, they will access contiguous memory for a given (d,h). This is good. In our block layout, if we set blockDim.x = Tw = 32 (warp size), then for a fixed (d,h), 32 threads read contiguous W elements. Since the input window for W is [2w-1, 2w, 2w+1], threads with w varying will access strided memory (stride 2). But still within a cache line (128 bytes = 32 floats). So a single warp reading kw=0 will access 32 floats with stride 2, which might require 2 cache lines, but still coalesced into 2 memory transactions. That's okay.
If we use shared memory, we want threads to cooperatively load the input tile. Since W is contiguous, loading a block of W is efficient. Let's design a tile where the block loads a 3D input region into smem.
Pseudo-code:
```
#define TW 16
#define TH 4
#define TD 4
__shared__ float smem[(2*TD+1)*(2*TH+1)*(2*TW+1)];
int n = blockIdx.x; // or derived
int c = blockIdx.y;
int tile_d = blockIdx.z / ((H_out/TH)*(W_out/TW)); // not ideal
// Let's use 1D grid for spatial tiles
int spatial_tiles = (D_out/TD) * (H_out/TH) * (W_out/TW);
int tile_idx = blockIdx.z;
int tile_w = tile_idx % (W_out/TW);
int tile_h = (tile_idx / (W_out/TW)) % (H_out/TH);
int tile_d = tile_idx / ((W_out/TW)*(H_out/TH));
// thread indices
int tx = threadIdx.x; // 0..TW-1
int ty = threadIdx.y; // 0..TH-1
int tz = threadIdx.z; // 0..TD-1
int w_out = tile_w * TW + tx;
int h_out = tile_h * TH + ty;
int d_out = tile_d * TD + tz;
// Compute smem dimensions
#define S_W (2*TW+1)
#define S_H (2*TH+1)
#define S_D (2*TD+1)
// Load input tile into smem
// We need to load (2*TD+1)*(2*TH+1)*(2*TW+1) elements.
// Use all threads to load.
int total_smem = S_D * S_H * S_W;
for (int i = tz*TH*TW + ty*TW + tx; i < total_smem; i += blockDim.x*blockDim.y*blockDim.z) {
// compute local smem coords (sd, sh, sw)
int sw = i % S_W;
int tmp = i / S_W;
int sh = tmp % S_H;
int sd = tmp / S_H;
int id = tile_d * TD * 2 + sd - 1;
int ih = tile_h * TH * 2 + sh - 1;
int iw = tile_w * TW * 2 + sw - 1;
float val = 0.0f;
if (id >= 0 && id < D && ih >= 0 && ih < H && iw >= 0 && iw < W) {
val = in[INDEX(n,c,id,ih,iw)];
}
smem[INDEX_S(sd,sh,sw)] = val;
}
__syncthreads();
// Now compute output
float sum = 0.0f;
int base_sd = tz * 2 + 1 - 1; // because d_out = tile_d*TD + tz, and id = d_out*2 + kd - 1 = tile_d*TD*2 + tz*2 + kd - 1.
// But in smem, sd corresponds to id - (tile_d*TD*2) + 1. Let's verify: smem was loaded with id = tile_d*TD*2 + sd - 1. So sd = id - tile_d*TD*2 + 1.
// For our output d_out, id = d_out*2 + kd - 1 = tile_d*TD*2 + tz*2 + kd - 1.
// So sd = (tile_d*TD*2 + tz*2 + kd - 1) - tile_d*TD*2 + 1 = tz*2 + kd.
// Similarly sh = ty*2 + kh, sw = tx*2 + kw.
#pragma unroll
for (int kd = 0; kd < 3; ++kd) {
for (int kh = 0; kh < 3; ++kh) {
for (int kw = 0; kw < 3; ++kw) {
sum += smem[(tz*2+kd)*S_H*S_W + (ty*2+kh)*S_W + (tx*2+kw)];
}
}
}
out[INDEX(n,c,d_out,h_out,w_out)] = sum / 27.0f;
```
This is clean. Shared memory size: S_D=9, S_H=9, S_W=33. Total = 9*9*33=2673 floats = ~10.7KB. Block size = 4*4*16=256 threads. That's good. Occupancy: 256 threads per block, 10.7KB shared memory. On H100, max 2048 threads per SM, 227KB shared memory. So we can have up to 8 blocks per SM (shared memory limited to 21 blocks? Actually 227/10.7 ≈ 21, but threads limit is 2048/256=8). So 8 blocks per SM, full occupancy.
Now each block loads ~10.7KB from global memory (plus some boundary threads load less because of zero-padding). Each block computes 256 outputs. Without shared memory, 256 outputs would require 256*27 = 6912 float loads ≈ 27.6KB. With shared memory, we load 10.7KB. That's a 2.6x reduction in global memory reads. This should significantly improve performance.
But wait, the shared memory tile includes the boundaries. The input tile overlaps between adjacent blocks. For example, adjacent tiles in W share an input "column" because of the 2*TW+1 width. Specifically, block tile_w loads W input from `tile_w*TW*2 - 1` to `tile_w*TW*2 + 2*TW`. The next block `tile_w+1` loads from `(tile_w+1)*TW*2 - 1` to `(tile_w+1)*TW*2 + 2*TW`. The overlap is 1 element? Let's see: Block 0 loads sw=0..32 -> iw=-1..31. Block 1 loads sw=0..32 -> iw=31..63? Wait, tile_w=1: iw = 1*32*2 + sw -1 = 63 + sw. So sw=0 -> 63, sw=32 -> 95. Block 0 max iw = 31. So no overlap. Actually because S_W = 2*TW+1 = 33, and stride between output tiles is TW=16 in output space, which corresponds to 32 in input space. The input tile width is 33. So adjacent tiles overlap by 1 element. That's fine. The overlap means some global memory reads are duplicated between blocks. But the total number of reads per block is still much smaller than per output reads.
Wait, is there no overlap? Let's compute carefully. Output tile covers w_out = tile_w*TW + tx, where tx=0..TW-1. The input needed for these outputs is iw = w_out*2 + kw -1 = (tile_w*TW + tx)*2 + kw -1. For tx=0..TW-1, kw=0..2: min iw = tile_w*TW*2 -1. Max iw = (tile_w*TW + TW-1)*2 + 2 -1 = tile_w*TW*2 + 2*TW -2 +1 = tile_w*TW*2 + 2*TW -1. So input width needed = 2*TW+1. The starting index is tile_w*TW*2 - 1. The next tile starts at (tile_w+1)*TW*2 -1 = tile_w*TW*2 + 2*TW - 1. So the previous tile ended at tile_w*TW*2 + 2*TW -1. So the next tile starts exactly at the previous tile's end. So there is overlap of exactly 1 element. For TW=16, block0 reads iw=-1..31 (33 elements). Block1 reads iw=31..63 (33 elements). Overlap at iw=31. That's 1 element out of 33. Very small. So each input element is loaded by exactly one block (except boundary overlap). Great.
Same for H and D. So total global memory reads per output = 10.7KB / 256 outputs = 42.8 bytes per output. Without shared memory, 27*4 = 108 bytes per output. So we save 2.5x memory bandwidth. Since the kernel is memory bound, we expect ~2.5x speedup. This should be more than enough to beat PyTorch.
But we need to be careful: The kernel must handle arbitrary batch, channel, depth, height, width, but we can specialize for these sizes or write a generic kernel. Since `ModelNew` accepts same constructor args and we know the input shapes from the problem, we could hardcode sizes? The evaluation likely uses the exact sizes from the problem for correctness checks, but the kernel should be generic enough to handle any sizes matching the constructor parameters. However, the tiling requires dimensions to be multiples of tile sizes. Here D_out=64, H_out=64, W_out=128. 64 is divisible by 4, 128 by 16. So it's fine. But to be safe, we can add code to handle remainders or just use generic sizes with the tiling approach and handle boundary conditions.
Actually, we don't need to tile perfectly; we can launch enough tiles to cover the output, and have threads check if their output coordinates are within bounds. For the shared memory load loop, we can load the required region for the tile, including out-of-bounds padding. Since the tile's output region is partially out of bounds, the SMEM load region might need to be smaller. But that's complicated.
Simpler: Since the problem provides exact input sizes, and the constructor arguments are fixed (kernel_size=3, stride=2, padding=1), we can write a kernel that works for any input sizes but assumes output dimensions are multiples of the tile sizes? Actually we can just write a generic kernel without shared memory but with unrolled loops and fast path. That might be easier to get correct. The performance gain from shared memory might be unnecessary if PyTorch is also not using shared memory. But PyTorch might be using a highly vectorized implementation from cuDNN. Wait, does PyTorch use cuDNN for AvgPool3d? I'm not sure. PyTorch's AvgPool3d has a native CUDA implementation. It might not be as optimized as cuDNN pooling. But for 3D, there might be a cuDNN path. Actually PyTorch's `avg_pool3d` on CUDA calls into ATen's `avg_pool3d_out_cuda` which uses a custom kernel, not necessarily cuDNN. For 3D, cuDNN pooling might be used for `cudnn_avg_pool3d` only if certain conditions are met. I think the default `nn.AvgPool3d` uses the native implementation. So a custom kernel could beat it.
But implementing shared memory correctly is a bit more code but manageable. Let's design the kernel with shared memory tiling but also generic enough.
However, there's a catch: the memory layout is NCDHW. For a given (n,c), the D*H*W block is contiguous. The channel dimension is fairly small (32). We could assign each block to a single (n,c) and a spatial tile. That would be 16*32 = 512 blocks per spatial grid. Spatial grid: D_out/TD * H_out/TH * W_out/TW = 16 * 16 * 8 = 2048. Total blocks = 512 * 2048 = 1,048,576. That's fine. Each block handles 4*4*16=256 outputs.
But we need to compute the linear spatial tile index. Since we have 3D spatial grid of tiles, we can compute:
```
int tile_d = blockIdx.z / ((H_out/TH) * (W_out/TW));
int rem = blockIdx.z % ((H_out/TH) * (W_out/TW));
int tile_h = rem / (W_out/TW);
int tile_w = rem % (W_out/TW);
```
This assumes blockIdx.z is the spatial tile index. But gridDim.z has a max limit (65535). 2048 < 65535, so okay.
Alternatively, we can combine n and c into blockIdx.x and spatial into blockIdx.y and blockIdx.z, but 3D grid is limited. Actually we can do:
- gridDim.x = N * C
- gridDim.y = D_out / TD
- gridDim.z = (H_out / TH) * (W_out / TW)
Then inside kernel:
```
int n = blockIdx.x / C;
int c = blockIdx.x % C;
int tile_d = blockIdx.y;
int tmp = blockIdx.z;
int tile_h = tmp / (W_out/TW);
int tile_w = tmp % (W_out/TW);
```
This uses 3D grid. gridDim.x = 512, gridDim.y = 16, gridDim.z = 128. All within 65535. Good.
But we need to compute `W_out`, `H_out`, `D_out` inside the kernel or pass them as arguments.
Actually, since the constructor receives kernel_size, stride, padding, we can compute output sizes in Python and pass them to the kernel along with the dimensions.
Even better, we can compute output size in C++ as well.
Let's write the CUDA source.
We need to be careful about the number of registers. Unrolling deeply with shared memory might use more registers but should be fine.
Let's define the tile sizes as constants. For generality, we can template them.
Since we only need to support kernel_size=3, stride=2, padding=1, we can hardcode those in the kernel. But it's cleaner to pass them or just hardcode because they're fixed for this problem. However, if the grader tests different sizes, the kernel must still work. The problem says "Replace the forward computation ... Accept the same constructor arguments". So the kernel must support arbitrary kernel_size, stride, padding. But the performance is measured on the given configuration. So we can write a generic kernel but maybe with compile-time constants for the given case? We can generate CUDA code dynamically in Python using `load_inline` with `cpp_sources` built from a string; we can embed the parameters into the CUDA source at module load time. But `ModelNew.__init__` gets the arguments; we can compile the kernel inside `__init__` with the specific parameter values. However, `load_inline` compiles at module load time (when the Python file is executed), not per instance. Actually the code says "Use `torch.utils.cpp_extension.load_inline` to compile and bind CUDA source at module load time". Module load time is when `exec()` runs the Python file. At that time, we don't have the constructor arguments yet. We could compile a generic kernel at module load time. Or we can compile inside the constructor, but the instruction says "at module load time". It might be acceptable to compile inside the constructor as long as we use `load_inline`. But to be safe, let's write a generic kernel that handles arbitrary kernel_size, stride, padding.
For a generic kernel with tiling, handling arbitrary kernel_size, stride is tricky because the shared memory input tile size depends on them: `S_D = tile_d * stride + kernel_size`, etc. Also the mapping from output tile to input tile uses stride and padding. It's doable if we pass them as template parameters or runtime variables.
But for simplicity and performance, we can assume that the generic case is used only for these parameters, but code them as runtime variables. The compiler might not unroll loops based on runtime variables. However, we can still use `for` loops and hope the compiler optimizes them. But the main cost is the loops themselves. With kernel_size up to maybe 7 or so, loops are fine.
Given that the only performance-sensitive case is kernel_size=3, we can write a generic kernel but with a specialized path when `kernel_size==3 && stride==2 && padding==1`. Actually, we can even write a completely generic kernel with loops, and rely on the fact that PyTorch's generic kernel is also loop-based, and ours might be faster due to shared memory.
But writing a generic shared memory tile for arbitrary kernel_size and stride is more complex. Let me think if we can do a semi-generic approach: The shared memory tile approach works for any kernel_size K and stride S. The output tile is (TD, TH, TW). The input tile is (TD*S + K - S + ...)? Wait, the input needed for an output tile of size T in dimension X with stride S and kernel K is: The first output uses indices from `out_start*S` to `out_start*S + K - 1` (considering padding offset). Actually with padding P, the input index for output o and kernel offset k is `o*S + k - P`. For outputs o in [tile_start, tile_start + T - 1], the input indices range from `tile_start*S - P` to `(tile_start + T - 1)*S + K - 1 - P`. The size of this range is `(T-1)*S + K`. So the shared memory dimension is `(T-1)*S + K`. For T=TD, S=2, K=3: size = (TD-1)*2 + 3 = 2*TD +1. Same as before. For K=3, S=1: size = (TD-1)*1 + 3 = TD+2.
So we can define `S_D = (TD-1)*stride + kernel_size`, etc. The threads in the block need to load `S_D * S_H * S_W` elements. The mapping from output thread (tz,ty,tx) to smem indices for summation is: smem_d = tz*stride + kd, smem_h = ty*stride + kh, smem_w = tx*stride + kw.
This is generic. We can make TD, TH, TW template parameters or constants for compilation. The kernel_size, stride, padding can be runtime integers, but `S_D`, etc., would be computed at runtime. However, the shared memory array size must be a compile-time constant if declared as `__shared__ float smem[...]`. But we can declare a dynamic shared memory array: `extern __shared__ float smem[];` and allocate size in the kernel launch via `shared_mem_size` argument. However, in `load_inline`, the kernel launch is done in Python via `torch.ops`? Actually we will write a host wrapper function in C++ that launches the kernel. We can pass `shared_mem_size` to `<<<grid, block, shared_mem_size>>>`. So we can make it fully generic.
But the assignment says we should call the compiled extension from `ModelNew.forward`. The host wrapper receives the input and output tensors, and other parameters (kernel_size, stride, padding). We compute output sizes in the wrapper and launch the kernel.
Dynamic shared memory is fine.
Let's design the generic host wrapper and kernel:
```
torch::Tensor avgpool3d_cuda(torch::Tensor input, int kernel_size, int stride, int padding) {
// compute output sizes
int N = input.size(0);
int C = input.size(1);
int D = input.size(2);
int H = input.size(3);
int W = input.size(4);
int D_out = (D + 2*padding - kernel_size) / stride + 1;
int H_out = (H + 2*padding - kernel_size) / stride + 1;
int W_out = (W + 2*padding - kernel_size) / stride + 1;
auto output = torch::empty({N, C, D_out, H_out, W_out}, input.options());
// define tile sizes
const int TW = 16;
const int TH = 4;
const int TD = 4;
// smem size per block
int S_W = (TW-1)*stride + kernel_size;
int S_H = (TH-1)*stride + kernel_size;
int S_D = (TD-1)*stride + kernel_size;
size_t smem_size = S_D * S_H * S_W * sizeof(float);
dim3 block(TW, TH, TD);
dim3 grid(N*C, (D_out + TD - 1) / TD, ((H_out + TH - 1) / TH) * ((W_out + TW - 1) / TW));
// But we need a 3D grid. Here we use:
// grid.x = N*C
// grid.y = num_tiles_d
// grid.z = num_tiles_h * num_tiles_w
int num_tiles_d = (D_out + TD - 1) / TD;
int num_tiles_h = (H_out + TH - 1) / TH;
int num_tiles_w = (W_out + TW - 1) / TW;
avgpool3d_kernel<<<grid, block, smem_size>>>(
input.data_ptr<float>(), output.data_ptr<float>(),
N, C, D, H, W,
D_out, H_out, W_out,
kernel_size, stride, padding,
num_tiles_h, num_tiles_w
);
return output;
}
```
In the kernel:
```
__global__ void avgpool3d_kernel(const float* in, float* out,
int N, int C, int D, int H, int W,
int D_out, int H_out, int W_out,
int K, int S, int P,
int num_tiles_h, int num_tiles_w) {
int n = blockIdx.x / C;
int c = blockIdx.x % C;
int tile_d = blockIdx.y;
int tile_hw = blockIdx.z;
int tile_h = tile_hw / num_tiles_w;
int tile_w = tile_hw % num_tiles_w;
int tx = threadIdx.x;
int ty = threadIdx.y;
int tz = threadIdx.z;
int w_out = tile_w * blockDim.x + tx;
int h_out = tile_h * blockDim.y + ty;
int d_out = tile_d * blockDim.z + tz;
int S_W = (blockDim.x - 1) * S + K;
int S_H = (blockDim.y - 1) * S + K;
int S_D = (blockDim.z - 1) * S + K;
extern __shared__ float smem[];
// Load input tile
int total_smem = S_D * S_H * S_W;
for (int i = tz * blockDim.y * blockDim.x + ty * blockDim.x + tx; i < total_smem; i += blockDim.x * blockDim.y * blockDim.z) {
int sw = i % S_W;
int tmp = i / S_W;
int sh = tmp % S_H;
int sd = tmp / S_H;
int id = tile_d * blockDim.z * S + sd - P;
int ih = tile_h * blockDim.y * S + sh - P;
int iw = tile_w * blockDim.x * S + sw - P;
float val = 0.0f;
if (id >= 0 && id < D && ih >= 0 && ih < H && iw >= 0 && iw < W) {
val = in[((n * C + c) * D + id) * H * W + ih * W + iw];
}
smem[i] = val;
}
__syncthreads();
if (w_out < W_out && h_out < H_out && d_out < D_out) {
float sum = 0.0f;
#pragma unroll
for (int kd = 0; kd < K; ++kd) {
int sd = tz * S + kd;
#pragma unroll
for (int kh = 0; kh < K; ++kh) {
int sh = ty * S + kh;
#pragma unroll
for (int kw = 0; kw < K; ++kw) {
int sw = tx * S + kw;
sum += smem[(sd * S_H + sh) * S_W + sw];
}
}
}
out[((n * C + c) * D_out + d_out) * H_out * W_out + h_out * W_out + w_out] = sum / (K*K*K);
}
}
```
Wait, `#pragma unroll` with runtime `K` might not work. Actually `#pragma unroll` requires compile-time loop bounds. If `K` is a `const int` passed as template parameter, it works. But as a run-time parameter, the compiler may ignore it or generate a runtime loop. For small K like 3, a runtime loop is fine. But for generic K, we can't unroll.
However, we can specialize the kernel for `K==3` using a template or separate function. But since the performance is measured on K=3, and the kernel is generic, the runtime loop might still be okay. But to be safe, we can create a templated kernel with `int K` as a template parameter, and in the host wrapper, switch on `K`. Since `K` is known at runtime (from constructor), we can generate the appropriate template instantiation. But with `load_inline`, we need to have the C++ source compile at module load time. We can include a switch statement that dispatches to templated kernels for common K values (e.g., 1..7). That is doable.
Alternatively, simpler: we compute `float inv = 1.0f / (K*K*K);` and just use a runtime loop. On H100, a small loop of 3 is likely unrolled by the compiler anyway if `K` is used as a loop bound and is in a register. Actually the compiler might not know it's small. We can add `switch(kernel_size)` inside the kernel? That would be expensive per thread. Better to template.
But we can also just not worry about unrolling for now, write a generic kernel, and see if it's fast enough.
Wait, the problem says the kernel must handle generic constructor arguments. But grid and tile sizes need to be chosen. For large kernel sizes, shared memory might overflow. We could have a fallback: if the shared memory needed exceeds a threshold, use a naive kernel. But for this problem's parameters, a tile-based kernel works.
Let's think about the maximum possible shared memory size. If kernel_size can be any integer up to, say, 10, and stride any, then S_D = (TD-1)*stride + K. If TD=4, TH=4, TW=16, and K=10, stride=10, S_D = 3*10+10 = 40. S_H=40, S_W=3*10+10=40. Total smem = 40*40*40*4 = 256KB > 227KB limit. So for arbitrary parameters, we might not have enough shared memory.
Options:
1. Choose tile sizes adaptively based on K and S. But that's complex in C++.
2. Use a naive kernel for large K/S, and shared memory kernel for small K/S. Since the reference is generic, our naive kernel might still be competitive for large K.
3. Since the problem says "Simple model that performs 3D Average Pooling" and the test uses kernel_size=3, stride=2, padding=1, we can optimize for that case and write a fallback generic path that still works. For example, if `kernel_size == 3 && stride == 2`, use the optimized tiled kernel; else use a generic naive kernel.
This is the most pragmatic approach.
So we can write a host function that checks the parameters and launches the appropriate kernel.
Let's design the code structure:
CUDA source includes two kernels:
1. `avgpool3d_ks3_s2_p1_kernel` specialized for kernel_size=3, stride=2, padding=1, with shared memory tiling.
2. `avgpool3d_generic_kernel` for arbitrary parameters, simple per-element with loops.
Host wrapper:
```cpp
torch::Tensor avgpool3d_forward(torch::Tensor x, int kernel_size, int stride, int padding) {
// compute output shape
...
auto out = torch::empty(...);
if (kernel_size == 3 && stride == 2 && padding == 1) {
// launch optimized kernel
} else {
// launch generic kernel
}
return out;
}
```
Then in `ModelNew.forward`, we call this function.
Actually, to be even more robust, we can make the optimized kernel work for any sizes (multiples of tile sizes or with boundary checks). Our generic grid handles boundaries with `if (w_out < W_out...)`. For the shared memory load, we load the entire tile region even if some outputs are out of bounds. That's fine.
Wait, the `num_tiles` formula: `(D_out + TD - 1) / TD`. For D_out=64 and TD=4, it's 16. For D_out=65, it's 17. The last tile will have some threads with `d_out >= D_out`. Those threads will still participate in shared memory loading (which is fine; they load extra input data that might be beyond the output bounds but we don't care). They just won't write output. The `__syncthreads()` ensures all threads finish loading before any compute. Good.
One concern: The shared memory load uses `blockDim.x`, `blockDim.y`, `blockDim.z`. In our specialized kernel, we can define them as compile-time constants for clarity, or just use the runtime values. Since the kernel is called with fixed block dimensions, the compiler can optimize them if we pass them as templates. But using `blockDim.x` inside the kernel won't be treated as compile-time constant by the compiler (as it's a special register). However, the loop `i += blockDim.x * blockDim.y * blockDim.z` will just use the register value. That's fine.
But to get maximum performance, compile-time constants for tile sizes help loop unrolling and indexing. So we can define `constexpr int TW = 16, TH = 4, TD = 4;` inside the kernel and launch with exactly those dimensions.
Let's do that for the specialized kernel.
**Optimized kernel details:**
```cpp
template <typename scalar_t>
__global__ void avgpool3d_ks3_s2_kernel(
const scalar_t* __restrict__ in,
scalar_t* __restrict__ out,
int N, int C, int D, int H, int W,
int D_out, int H_out, int W_out,
int padding) {
constexpr int K = 3;
constexpr int S = 2;
constexpr int TW = 16;
constexpr int TH = 4;
constexpr int TD = 4;
constexpr int S_W = (TW - 1) * S + K; // 33
constexpr int S_H = (TH - 1) * S + K; // 9
constexpr int S_D = (TD - 1) * S + K; // 9
int n = blockIdx.x / C;
int c = blockIdx.x % C;
int tile_d = blockIdx.y;
int tile_hw = blockIdx.z;
int num_tiles_w = (W_out + TW - 1) / TW;
int tile_h = tile_hw / num_tiles_w;
int tile_w = tile_hw % num_tiles_w;
int tx = threadIdx.x;
int ty = threadIdx.y;
int tz = threadIdx.z;
int w_out = tile_w * TW + tx;
int h_out = tile_h * TH + ty;
int d_out = tile_d * TD + tz;
int base_id = tile_d * TD * S - padding;
int base_ih = tile_h * TH * S - padding;
int base_iw = tile_w * TW * S - padding;
__shared__ scalar_t smem[S_D * S_H * S_W];
// Cooperative load of input tile
int thread_idx = tz * TH * TW + ty * TW + tx;
int total_smem = S_D * S_H * S_W;
#pragma unroll
for (int i = thread_idx; i < total_smem; i += TD * TH * TW) {
int sw = i % S_W;
int tmp = i / S_W;
int sh = tmp % S_H;
int sd = tmp / S_H;
int id = base_id + sd;
int ih = base_ih + sh;
int iw = base_iw + sw;
scalar_t val = static_cast<scalar_t>(0);
if (id >= 0 && id < D && ih >= 0 && ih < H && iw >= 0 && iw < W) {
val = in[((n * C + c) * D + id) * H * W + ih * W + iw];
}
smem[i] = val;
}
__syncthreads();
if (w_out < W_out && h_out < H_out && d_out < D_out) {
scalar_t sum = static_cast<scalar_t>(0);
#pragma unroll
for (int kd = 0; kd < K; ++kd) {
int sd = tz * S + kd;
#pragma unroll
for (int kh = 0; kh < K; ++kh) {
int sh = ty * S + kh;
#pragma unroll
for (int kw = 0; kw < K; ++kw) {
int sw = tx * S + kw;
sum += smem[(sd * S_H + sh) * S_W + sw];
}
}
}
out[((n * C + c) * D_out + d_out) * H_out * W_out + h_out * W_out + w_out] = sum / static_cast<scalar_t>(K * K * K);
}
}
```
Wait, we used `padding` as a runtime variable. But the kernel is specialized for stride=2, kernel=3. For this specific case, padding is passed as an argument. The boundary handling in the load checks `id >= 0 && id < D`, etc. Because padding can be any value, the base index includes `-padding`. If padding=1, base = tile_start*S - 1. That's correct. If padding=0, base = tile_start*S. Also works. But if padding > 0, the low-side boundaries have invalid indices. We handle them by zeroing out. The denominator is always K*K*K because count_include_pad=True. So dividing by 27 is correct regardless of padding.
However, note that with `padding > 0`, the high side might also have invalid indices if the output size isn't perfectly fitted. The `if` check handles that generically. So the kernel is correct for any padding.
But is the output size formula always integer division floor? PyTorch's AvgPool3d uses `output_size = floor((input_size + 2*padding - kernel_size) / stride) + 1`. The kernel launch must ensure we cover all outputs. Yes.
One issue: The denominator K*K*K assumes kernel_size is the same in all dimensions. The constructor takes a single `kernel_size` int, implying cubic kernel. So K*K*K is fine.
But what about `count_include_pad`? The reference uses default `True`. Our kernel divides by 27 always. That's correct.
What if `stride` is not 2? The optimized kernel is only launched for `stride==2 && kernel_size==3`. But we need to pass `stride` from the constructor. The host wrapper will check and dispatch.
Also, for the generic kernel, we need to handle the general case. Let's write a generic naive kernel:
```cpp
template <typename scalar_t>
__global__ void avgpool3d_generic_kernel(
const scalar_t* __restrict__ in,
scalar_t* __restrict__ out,
int N, int C, int D, int H, int W,
int D_out, int H_out, int W_out,
int K, int S, int P) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = N * C * D_out * H_out * W_out;
if (idx >= total) return;
int tmp = idx;
int w = tmp % W_out; tmp /= W_out;
int h = tmp % H_out; tmp /= H_out;
int d = tmp % D_out; tmp /= D_out;
int c = tmp % C; tmp /= C;
int n = tmp;
scalar_t sum = static_cast<scalar_t>(0);
for (int kd = 0; kd < K; ++kd) {
int id = d * S + kd - P;
for (int kh = 0; kh < K; ++kh) {
int ih = h * S + kh - P;
for (int kw = 0; kw < K; ++kw) {
int iw = w * S + kw - P;
if (id >= 0 && id < D && ih >= 0 && ih < H && iw >= 0 && iw < W) {
sum += in[((n * C + c) * D + id) * H * W + ih * W + iw];
}
}
}
}
out[idx] = sum / static_cast<scalar_t>(K * K * K);
}
```
This generic kernel uses a 1D grid. For the specific case, the tiled kernel should be much faster.
But we also need to consider that the input might not be contiguous. In PyTorch, `.data_ptr()` gives a pointer to the underlying storage, but strides might not be contiguous. However, for standard NCDHW tensors created by `torch.rand`, they are contiguous. The reference implementation operates on contiguous memory. Our kernel assumes contiguous layout. To be safe, we might want to call `.contiguous()` in the host wrapper or just assume it. Actually, `nn.AvgPool3d` handles non-contiguous inputs by accessing via strides. But for this problem, we can just ensure the input is contiguous before launching. But that might involve a copy. However, the inputs are generated by `torch.rand`, so they're contiguous. Let's add a check or call `x.contiguous()`.
Better: In the host wrapper, we can do:
```cpp
auto x_contig = x.contiguous();
```
But since `x` is already contiguous, this is a no-op.
Now, the host wrapper:
```cpp
torch::Tensor avgpool3d_forward(torch::Tensor x, int kernel_size, int stride, int padding) {
int N = x.size(0);
int C = x.size(1);
int D = x.size(2);
int H = x.size(3);
int W = x.size(4);
int D_out = (D + 2 * padding - kernel_size) / stride + 1;
int H_out = (H + 2 * padding - kernel_size) / stride + 1;
int W_out = (W + 2 * padding - kernel_size) / stride + 1;
auto out = torch::empty({N, C, D_out, H_out, W_out}, x.options());
auto x_contig = x.contiguous();
if (kernel_size == 3 && stride == 2) {
constexpr int TW = 16;
constexpr int TH = 4;
constexpr int TD = 4;
int num_tiles_d = (D_out + TD - 1) / TD;
int num_tiles_h = (H_out + TH - 1) / TH;
int num_tiles_w = (W_out + TW - 1) / TW;
dim3 block(TW, TH, TD);
dim3 grid(N * C, num_tiles_d, num_tiles_h * num_tiles_w);
avgpool3d_ks3_s2_kernel<float><<<grid, block>>>(
x_contig.data_ptr<float>(), out.data_ptr<float>(),
N, C, D, H, W, D_out, H_out, W_out, padding);
} else {
int total = N * C * D_out * H_out * W_out;
int threads = 256;
int blocks = (total + threads - 1) / threads;
avgpool3d_generic_kernel<float><<<blocks, threads>>>(
x_contig.data_ptr<float>(), out.data_ptr<float>(),
N, C, D, H, W, D_out, H_out, W_out,
kernel_size, stride, padding);
}
return out;
}
```
Actually, for the optimized kernel, the template type `float` is fine because the reference is fp32 and the prompt says fp32 precision. We can also support double if needed, but inputs are float.
Wait, should we handle `stride=None`? In the constructor, `stride=None` defaults to `kernel_size`. In the host wrapper, we need to handle that. But we are given `get_init_inputs()` returning `[kernel_size, stride, padding]`. So `stride` is an integer (2). In `ModelNew.__init__`, we can replicate the same logic: `if stride is None: stride = kernel_size`. Then pass `stride` to the kernel.
One more thing: We need to compile the code with `load_inline`. We need to include `<torch/extension.h>`. The host wrapper must be listed in `functions=[...]`.
Let's write the Python module.
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = r'''
#include <torch/extension.h>
#include <cuda_runtime.h>
template <typename scalar_t>
__global__ void avgpool3d_ks3_s2_kernel(
const scalar_t* __restrict__ in,
scalar_t* __restrict__ out,
int N, int C, int D, int H, int W,
int D_out, int H_out, int W_out,
int padding) {
constexpr int K = 3;
constexpr int S = 2;
constexpr int TW = 16;
constexpr int TH = 4;
constexpr int TD = 4;
constexpr int S_W = (TW - 1) * S + K;
constexpr int S_H = (TH - 1) * S + K;
constexpr int S_D = (TD - 1) * S + K;
int n = blockIdx.x / C;
int c = blockIdx.x % C;
int tile_d = blockIdx.y;
int tile_hw = blockIdx.z;
int num_tiles_w = (W_out + TW - 1) / TW;
int tile_h = tile_hw / num_tiles_w;
int tile_w = tile_hw % num_tiles_w;
int tx = threadIdx.x;
int ty = threadIdx.y;
int tz = threadIdx.z;
int w_out = tile_w * TW + tx;
int h_out = tile_h * TH + ty;
int d_out = tile_d * TD + tz;
int base_id = tile_d * TD * S - padding;
int base_ih = tile_h * TH * S - padding;
int base_iw = tile_w * TW * S - padding;
__shared__ scalar_t smem[S_D * S_H * S_W];
int thread_idx = tz * TH * TW + ty * TW + tx;
int total_smem = S_D * S_H * S_W;
#pragma unroll
for (int i = thread_idx; i < total_smem; i += TD * TH * TW) {
int sw = i % S_W;
int tmp = i / S_W;
int sh = tmp % S_H;
int sd = tmp / S_H;
int id = base_id + sd;
int ih = base_ih + sh;
int iw = base_iw + sw;
scalar_t val = static_cast<scalar_t>(0);
if (id >= 0 && id < D && ih >= 0 && ih < H && iw >= 0 && iw < W) {
val = in[((n * C + c) * D + id) * H * W + ih * W + iw];
}
smem[i] = val;
}
__syncthreads();
if (w_out < W_out && h_out < H_out && d_out < D_out) {
scalar_t sum = static_cast<scalar_t>(0);
#pragma unroll
for (int kd = 0; kd < K; ++kd) {
int sd = tz * S + kd;
#pragma unroll
for (int kh = 0; kh < K; ++kh) {
int sh = ty * S + kh;
#pragma unroll
for (int kw = 0; kw < K; ++kw) {
int sw = tx * S + kw;
sum += smem[(sd * S_H + sh) * S_W + sw];
}
}
}
out[((n * C + c) * D_out + d_out) * H_out * W_out + h_out * W_out + w_out] = sum / static_cast<scalar_t>(K * K * K);
}
}
template <typename scalar_t>
__global__ void avgpool3d_generic_kernel(
const scalar_t* __restrict__ in,
scalar_t* __restrict__ out,
int N, int C, int D, int H, int W,
int D_out, int H_out, int W_out,
int K, int S, int P) {
int total = N * C * D_out * H_out * W_out;
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= total) return;
int tmp = idx;
int w = tmp % W_out; tmp /= W_out;
int h = tmp % H_out; tmp /= H_out;
int d = tmp % D_out; tmp /= D_out;
int c = tmp % C; tmp /= C;
int n = tmp;
scalar_t sum = static_cast<scalar_t>(0);
for (int kd = 0; kd < K; ++kd) {
int id = d * S + kd - P;
for (int kh = 0; kh < K; ++kh) {
int ih = h * S + kh - P;
for (int kw = 0; kw < K; ++kw) {
int iw = w * S + kw - P;
if (id >= 0 && id < D && ih >= 0 && ih < H && iw >= 0 && iw < W) {
sum += in[((n * C + c) * D + id) * H * W + ih * W + iw];
}
}
}
}
out[idx] = sum / static_cast<scalar_t>(K * K * K);
}
torch::Tensor avgpool3d_forward(torch::Tensor x, int kernel_size, int stride, int padding) {
int N = x.size(0);
int C = x.size(1);
int D = x.size(2);
int H = x.size(3);
int W = x.size(4);
if (stride == 0) stride = kernel_size;
int D_out = (D + 2 * padding - kernel_size) / stride + 1;
int H_out = (H + 2 * padding - kernel_size) / stride + 1;
int W_out = (W + 2 * padding - kernel_size) / stride + 1;
auto out = torch::empty({N, C, D_out, H_out, W_out}, x.options());
auto x_contig = x.contiguous();
if (kernel_size == 3 && stride == 2) {
constexpr int TW = 16;
constexpr int TH = 4;
constexpr int TD = 4;
int num_tiles_d = (D_out + TD - 1) / TD;
int num_tiles_h = (H_out + TH - 1) / TH;
int num_tiles_w = (W_out + TW - 1) / TW;
dim3 block(TW, TH, TD);
dim3 grid(N * C, num_tiles_d, num_tiles_h * num_tiles_w);
avgpool3d_ks3_s2_kernel<float><<<grid, block>>>(
x_contig.data_ptr<float>(), out.data_ptr<float>(),
N, C, D, H, W, D_out, H_out, W_out, padding);
} else {
int total = N * C * D_out * H_out * W_out;
const int threads = 256;
int blocks = (total + threads - 1) / threads;
avgpool3d_generic_kernel<float><<<blocks, threads>>>(
x_contig.data_ptr<float>(), out.data_ptr<float>(),
N, C, D, H, W, D_out, H_out, W_out,
kernel_size, stride, padding);
}
return out;
}
'''
# Compile
avgpool3d_cuda = load_inline(
name="avgpool3d_cuda",
cpp_sources="torch::Tensor avgpool3d_forward(torch::Tensor x, int kernel_size, int stride, int padding);",
cuda_sources=cuda_source,
functions=["avgpool3d_forward"],
extra_cuda_cflags=["-O3", "--use_fast_math"],
verbose=False,
)
class ModelNew(nn.Module):
def __init__(self, kernel_size: int, stride: int = None, padding: int = 0):
super(ModelNew, self).__init__()
self.kernel_size = kernel_size
self.stride = stride if stride is not None else kernel_size
self.padding = padding
def forward(self, x: torch.Tensor) -> torch.Tensor:
return avgpool3d_cuda.avgpool3d_forward(x, self.kernel_size, self.stride, self.padding)
```
I need to be careful with `load_inline` arguments. The signature is:
```python
load_inline(
name,
cpp_sources,
cuda_sources,
functions,
extra_cflags=None,
extra_cuda_cflags=None,
extra_ldflags=None,
extra_include_paths=None,
build_directory=None,
verbose=False,
)
```
`cpp_sources` should contain C++ host code or just declarations? Actually `cpp_sources` is where you define the C++ wrapper functions that will be bound by pybind11. It can also include the CUDA source if there's no separate `cuda_sources`? Wait, `load_inline` has both `cpp_sources` and `cuda_sources`. The `cuda_sources` contain the CUDA code (kernels), and `cpp_sources` contain the C++ code that calls the kernels. But actually, if we put everything in `cuda_sources` and also include the host wrapper there, does it work? Usually `cuda_sources` is appended after `cpp_sources`. The host wrapper uses `<<< >>>` which is CUDA C++ syntax, so it must be in the `.cu` file, not `.cpp`. But `cpp_sources` is compiled with the C++ compiler (unless it's a `.cu` file). Actually `load_inline` generates a `.cpp` from `cpp_sources` and a `.cu` from `cuda_sources`. The host wrapper with `<<< >>>` must be in the `.cu` file (or compiled with nvcc). However, `cpp_sources` can contain declarations that are implemented in `cuda_sources`? Or we can put the host wrapper in `cuda_sources` because it's compiled with nvcc. The `functions=[...]` argument tells pybind11 which functions to bind. It looks for them in the combined translation unit.
But there is a subtlety: we need to provide a declaration of the function in `cpp_sources` if it's defined in `cuda_sources`, because pybind11 binding code is inserted into `cpp_sources`. Actually, `load_inline` processes `cpp_sources` and `cuda_sources` and compiles them together into a shared library. The `functions` list is used to generate a `PYBIND11_MODULE` block in the C++ source. The bound functions must be visible to the C++ compiler. If the function is defined in the `.cu` file (compiled with nvcc), the C++ compiler sees it as an `extern` function if declared in `cpp_sources`. So standard practice with `load_inline` is:
- `cpp_sources`: contains `#include <torch/extension.h>` and the forward declaration of the host function.
- `cuda_sources`: contains the actual host function definition (with `<<< >>>`) and the kernels.
Wait, actually I think `load_inline` just concatenates them. Let me check typical usage.
Typical usage (from PyTorch docs):
```python
cuda_source = '''
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void kernel(...) { ... }
torch::Tensor my_op(torch::Tensor x) {
...
kernel<<<...>>>(...);
return y;
}
'''
cpp_source = 'torch::Tensor my_op(torch::Tensor x);'
load_inline(
name='my_op',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['my_op'],
...
)
```
So `cpp_sources` just has declarations. The definitions are in `cuda_sources`. That works because both are compiled together eventually, but the pybind11 wrapper is generated in the C++ part and calls the function declared there, which is defined in the CUDA part.
But does the C++ compiler need to see the definition at link time? Since it's a single shared library built from both sources, the linker resolves it. The declaration in `cpp_sources` is enough for the generated binding code to compile. So this pattern works.
However, I recall that `load_inline` might also accept `cpp_sources` containing the actual wrapper and `cuda_sources` containing the kernels. Let's follow the pattern above: Put declarations in `cpp_sources` and definitions in `cuda_sources`.
Let's format the code carefully.
One more detail: `torch::Tensor` by value in host wrapper is fine. Inside the kernel, we don't need `.data_ptr()` for `out` until after it's created.
Also, consider half precision: The problem says produce numerically equivalent outputs at fp32 precision. So inputs are fp32, outputs fp32. We can just use `float`.
Another potential issue: The shared memory array `__shared__ scalar_t smem[S_D * S_H * S_W];` requires `S_D*S_H*S_W` to be a compile-time constant, which it is because `TW, TH, TD, K, S` are constexpr. For `scalar_t = float`, size = 9*9*33 = 2673 floats. That's well within limits.
Wait, let's double-check the formula: `(TW-1)*S + K`. TW=16, S=2, K=3 -> (15)*2 + 3 = 33. TH=4 -> (3)*2+3 = 9. TD=4 -> 9. 9*9*33 = 2673. Yes.
The block size is 4*4*16 = 256 threads. The load loop stride is 256. 2673 / 256 = ~10.44, so 11 iterations for some threads. That's fine.
But the `#pragma unroll` on the load loop might not be effective because the upper bound is a variable `total_smem` (though it's constexpr). Actually `total_smem` is a local variable initialized with a constexpr expression. The compiler might treat it as a constant. But the loop variable `i` starts at `thread_idx` and increments by `TD*TH*TW` (256), and condition `i < total_smem`. Since the trip count is small and constant, the compiler might unroll it. We can just write `for (int i = thread_idx; i < S_D * S_H * S_W; i += TD * TH * TW)` to be sure it's compile-time constant.
Potential issue: `blockIdx.z` has a maximum of 65535. For our grid, `num_tiles_h * num_tiles_w`. With W_out=128, H_out=64, TW=16, TH=4: num_tiles_w = 8, num_tiles_h = 16. Product = 128. That's fine. With larger sizes, e.g., W_out=256, H_out=256: num_tiles_w=16, num_tiles_h=64, product=1024. Fine. If W_out=512, H_out=512: product=1024*1024=1,048,576 > 65535. That would exceed gridDim.z. But the problem's reference size is modest. To be safe for larger sizes, we can split into a 2D grid for spatial tiles but map differently. Or we can just handle it by using a 1D block index for spatial tiles. Let's make it robust.
We can change the launch to use a 1D block index for the entire 3D spatial grid, and compute tile_d, tile_h, tile_w from it. Or we can use a 2D grid: blockIdx.y for D tiles, blockIdx.x for combined H*W tiles and N*C combined. But blockIdx.x is often used for the widest dimension. Actually we can do:
- `gridDim.x = N*C`
- `gridDim.y = D_tile * H_tile * W_tile`
Then inside kernel:
```
int spatial_tile = blockIdx.y;
int tile_w = spatial_tile % num_tiles_w;
int tmp = spatial_tile / num_tiles_w;
int tile_h = tmp % num_tiles_h;
int tile_d = tmp / num_tiles_h;
```
This uses `gridDim.y` for spatial tiles. gridDim.y max is 65535. For very large spatial grids, we might still exceed. But we can also put everything in `gridDim.x` if needed.
Actually, the maximum grid dimensions for CUDA are:
- x: 2^31-1
- y: 65535
- z: 65535
For a huge output, say W_out=1024, H_out=1024, D_out=1024, with TW=16,TH=4,TD=4: num_tiles = 64*256*256 = 4,194,304 > 65535. This would exceed both y and z.
Better to linearize everything into `blockIdx.x` and perhaps `blockIdx.y` for N*C.
We can do:
- `gridDim.x = (total_spatial_tiles + max_blocks_per_grid - 1) / max_blocks_per_grid`? No, we want to maximize parallelism.
Actually the simplest robust approach is:
- `blockIdx.x` covers N*C
- `blockIdx.y` covers D tiles
- `blockIdx.z` covers H tiles
- And inside the kernel, each block loops over W tiles? No, we'd lose parallelism.
Alternatively, we can map N*C*D_tiles*H_tiles*W_tiles into a single linear block index. Since CUDA supports 1D, 2D, 3D. We can use 2D:
- `blockIdx.x` = linear index for (N, C, D_tile, H_tile, W_tile) with large range.
- `blockIdx.y` = 1.
But we can only do this if we use a single dimension. `gridDim.x` can be up to 2^31-1, which is enough.
Let's redesign the launch to use 1D blocks:
```
int total_spatial_tiles = num_tiles_d * num_tiles_h * num_tiles_w;
dim3 block(TW, TH, TD);
// Each CUDA block = one (n, c, spatial_tile) combination
int total_blocks = N * C * total_spatial_tiles;
// But we need a 3D grid? Actually CUDA blocks can be 1D with large x.
dim3 grid(total_blocks, 1, 1);
```
Wait, but `dim3` has x, y, z all unsigned int. So `grid(total_blocks, 1, 1)` is fine. Then inside kernel:
```
int block_id = blockIdx.x;
int spatial_tile = block_id % total_spatial_tiles;
int nc = block_id / total_spatial_tiles;
int n = nc / C;
int c = nc % C;
int tile_w = spatial_tile % num_tiles_w;
int tmp = spatial_tile / num_tiles_w;
int tile_h = tmp % num_tiles_h;
int tile_d = tmp / num_tiles_h;
```
This is robust and simple. `total_blocks` = 512 * 2048 = 1,048,576 for our case, well within 2^31-1.
So let's change the kernel launch to 1D grid.
Now rewrite the kernel signature accordingly.
**Revised Optimized Kernel:**
```cpp
template <typename scalar_t>
__global__ void avgpool3d_ks3_s2_kernel(
const scalar_t* __restrict__ in,
scalar_t* __restrict__ out,
int N, int C, int D, int H, int W,
int D_out, int H_out, int W_out,
int padding,
int num_tiles_d, int num_tiles_h, int num_tiles_w,
int total_spatial_tiles) {
constexpr int K = 3;
constexpr int S = 2;
constexpr int TW = 16;
constexpr int TH = 4;
constexpr int TD = 4;
constexpr int S_W = (TW - 1) * S + K;
constexpr int S_H = (TH - 1) * S + K;
constexpr int S_D = (TD - 1) * S + K;
int block_id = blockIdx.x;
int spatial_tile = block_id % total_spatial_tiles;
int nc = block_id / total_spatial_tiles;
int n = nc / C;
int c = nc % C;
int tile_w = spatial_tile % num_tiles_w;
int tmp = spatial_tile / num_tiles_w;
int tile_h = tmp % num_tiles_h;
int tile_d = tmp / num_tiles_h;
int tx = threadIdx.x;
int ty = threadIdx.y;
int tz = threadIdx.z;
int w_out = tile_w * TW + tx;
int h_out = tile_h * TH + ty;
int d_out = tile_d * TD + tz;
int base_id = tile_d * TD * S - padding;
int base_ih = tile_h * TH * S - padding;
int base_iw = tile_w * TW * S - padding;
__shared__ scalar_t smem[S_D * S_H * S_W];
int thread_idx = tz * TH * TW + ty * TW + tx;
#pragma unroll
for (int i = thread_idx; i < S_D * S_H * S_W; i += TD * TH * TW) {
int sw = i % S_W;
int tmp2 = i / S_W;
int sh = tmp2 % S_H;
int sd = tmp2 / S_H;
int id = base_id + sd;
int ih = base_ih + sh;
int iw = base_iw + sw;
scalar_t val = static_cast<scalar_t>(0);
if (id >= 0 && id < D && ih >= 0 && ih < H && iw >= 0 && iw < W) {
val = in[((n * C + c) * D + id) * H * W + ih * W + iw];
}
smem[i] = val;
}
__syncthreads();
if (w_out < W_out && h_out < H_out && d_out < D_out) {
scalar_t sum = static_cast<scalar_t>(0);
#pragma unroll
for (int kd = 0; kd < K; ++kd) {
int sd = tz * S + kd;
#pragma unroll
for (int kh = 0; kh < K; ++kh) {
int sh = ty * S + kh;
#pragma unroll
for (int kw = 0; kw < K; ++kw) {
int sw = tx * S + kw;
sum += smem[(sd * S_H + sh) * S_W + sw];
}
}
}
out[((n * C + c) * D_out + d_out) * H_out * W_out + h_out * W_out + w_out] = sum / static_cast<scalar_t>(K * K * K);
}
}
```
This is robust.
Sum indexing: `(sd * S_H + sh) * S_W + sw`. Since `S_H=9`, `S_W=33`, the index = sd*297 + sh*33 + sw. But `sd` max is `tz*S + kd`. tz max = 3, S=2, kd=2 => 8. `S_D=9`, so sd in [0,8]. `sh` in [0,8]. `sw` in [0,32]. So index max = 8*297 + 8*33 + 32 = 2376 + 264 + 32 = 2672. Array size = 2673. Good.
Now the generic kernel can also be 1D grid.
```cpp
template <typename scalar_t>
__global__ void avgpool3d_generic_kernel(
const scalar_t* __restrict__ in,
scalar_t* __restrict__ out,
int N, int C, int D, int H, int W,
int D_out, int H_out, int W_out,
int K, int S, int P) {
int total = N * C * D_out * H_out * W_out;
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= total) return;
int tmp = idx;
int w = tmp % W_out; tmp /= W_out;
int h = tmp % H_out; tmp /= H_out;
int d = tmp % D_out; tmp /= D_out;
int c = tmp % C; tmp /= C;
int n = tmp;
scalar_t sum = static_cast<scalar_t>(0);
for (int kd = 0; kd < K; ++kd) {
int id = d * S + kd - P;
for (int kh = 0; kh < K; ++kh) {
int ih = h * S + kh - P;
for (int kw = 0; kw < K; ++kw) {
int iw = w * S + kw - P;
if (id >= 0 && id < D && ih >= 0 && ih < H && iw >= 0 && iw < W) {
sum += in[((n * C + c) * D + id) * H * W + ih * W + iw];
}
}
}
}
out[idx] = sum / static_cast<scalar_t>(K * K * K);
}
```
Host wrapper:
```cpp
torch::Tensor avgpool3d_forward(torch::Tensor x, int kernel_size, int stride, int padding) {
int N = x.size(0);
int C = x.size(1);
int D = x.size(2);
int H = x.size(3);
int W = x.size(4);
if (stride == 0) stride = kernel_size;
int D_out = (D + 2 * padding - kernel_size) / stride + 1;
int H_out = (H + 2 * padding - kernel_size) / stride + 1;
int W_out = (W + 2 * padding - kernel_size) / stride + 1;
auto out = torch::empty({N, C, D_out, H_out, W_out}, x.options());
auto x_contig = x.contiguous();
if (kernel_size == 3 && stride == 2) {
constexpr int TW = 16;
constexpr int TH = 4;
constexpr int TD = 4;
int num_tiles_d = (D_out + TD - 1) / TD;
int num_tiles_h = (H_out + TH - 1) / TH;
int num_tiles_w = (W_out + TW - 1) / TW;
int total_spatial_tiles = num_tiles_d * num_tiles_h * num_tiles_w;
dim3 block(TW, TH, TD);
int total_blocks = N * C * total_spatial_tiles;
avgpool3d_ks3_s2_kernel<float><<<total_blocks, block>>>(
x_contig.data_ptr<float>(), out.data_ptr<float>(),
N, C, D, H, W, D_out, H_out, W_out,
padding, num_tiles_d, num_tiles_h, num_tiles_w, total_spatial_tiles);
} else {
int total = N * C * D_out * H_out * W_out;
const int threads = 256;
int blocks = (total + threads - 1) / threads;
avgpool3d_generic_kernel<float><<<blocks, threads>>>(
x_contig.data_ptr<float>(), out.data_ptr<float>(),
N, C, D, H, W, D_out, H_out, W_out,
kernel_size, stride, padding);
}
return out;
}
```
One concern: `stride == 0` check. In PyTorch, `stride=None` is handled in Python before passing to C++. In our host wrapper, `stride` is an `int`, so we should handle it. In Python `ModelNew.__init__`, we will set `self.stride = stride if stride is not None else kernel_size`. Then pass `self.stride` to the C++ function. So the `stride == 0` check in C++ is a safety net and might be unnecessary. But it's fine.
Wait, but `torch::empty({N, C, D_out, H_out, W_out}, x.options())` creates a tensor. For large sizes, this is fine.
Now, `cpp_sources` should declare the function:
```cpp
torch::Tensor avgpool3d_forward(torch::Tensor x, int kernel_size, int stride, int padding);
```
`cuda_sources` contains all the kernel code and the host wrapper definition.
But `load_inline` requires that the C++ source be compiled before the CUDA source? Or they are compiled together. Typically the declaration in `cpp_sources` is enough for the generated binding code. But the generated binding code calls `avgpool3d_forward`. The definition is in `cuda_sources` which is compiled with nvcc and linked into the same shared library. So it should work.
However, I want to make sure that `avgpool3d_forward` in `cuda_sources` is visible (not `static`). Since it's a template-free function, it's fine.
Also, I should add `#include <torch/extension.h>` at the top of `cuda_sources`. Actually `cpp_sources` will also have includes for the binding, but for the CUDA source we need it for `torch::Tensor` and `.data_ptr()`.
Another issue: The compiler might complain about using `blockIdx.x` with a 1D grid and `dim3 block`. `dim3 total_blocks_dim(total_blocks, 1, 1);` or just `<<<total_blocks, block>>>` works because the first argument can be an integer for 1D grid.
Wait, `avgpool3d_ks3_s2_kernel<<<total_blocks, block>>>` expects the first argument to be `dim3` or `int`. Passing `int` is fine for 1D grid.
But inside the kernel, I used `blockIdx.x`. If the grid is 1D, that's correct.
One more thing: The total number of blocks is `N*C*total_spatial_tiles`. For N=16, C=32, total_spatial=2048, total_blocks=1,048,576. That's a lot of blocks, but each block is small (256 threads). The GPU can handle it. With 132 SMs on H100, each SM can handle up to 32 blocks (concurrent). So total concurrent blocks = ~4224. With 1 million blocks, each block processes 256 outputs. Total work is high, but the scheduler will distribute them. No issue.
Wait, H100 has 132 SMs? Actually H100 SXM has 132 SMs (for H100 80GB). PCIe has 114. Either way, many.
Now, should we worry about the generic kernel using 1D grid but `blockDim.x` being 256? It's fine.
Potential correctness issue: The denominator in AvgPool3d with `count_include_pad=True` is always `kernel_size**3`, which we do. With `padding`, the output size formula might produce negative sizes if invalid, but PyTorch would error too. So fine.
Another concern: The optimized kernel uses `float` only. The generic kernel also uses `float`. What if the input is double? The reference is `torch.rand` which defaults to float32. And the problem says "fp32 precision". So we are fine.
But `x.options()` includes dtype. If input is float, `out` is float. Good.
Let's finalize the Python code. But before that, we need to test compilation and correctness. Let's write the full Python source and compile.