We need to write a custom CUDA kernel for mean reduction over a specific dimension. Given the inputs: batch_size=128, dim1=4096, dim2=4095, and the dim to reduce over is 1 (from get_init_inputs). So the input shape is [128, 4096, 4095] and we reduce over dim=1, producing shape [128, 4095]. We need a fast kernel. We can write a simple CUDA kernel that loads elements along the reduction dimension and computes the mean. Since dim1=4096 is large, we can have each thread block handle one (batch, dim2) pair, and threads within the block cooperatively reduce over 4096 elements. Since 4096 is a power of two, we can use shared memory and parallel reduction. However, 4096 elements per block is not huge. We can also vectorize loads (e.g., float4). Since the tensor is contiguous in memory? Let's check memory layout: x shape (128, 4096, 4095). The reference uses torch.mean(x, dim=1). The input is contiguous by default (rand creates contiguous). So memory layout is row-major: stride[0] = 4096*4095, stride[1] = 4095, stride[2] = 1. Reducing over dim=1 means we need to sum over the middle dimension, which has stride 4095. So adjacent elements in the reduction dimension are 4095 apart in memory. That is not contiguous. Actually, if we reduce over dim=1, the elements to sum are separated by stride[1]=4095. So memory access is strided. But we can still load them. Alternatively, maybe we can transpose first? But that would add overhead. Considering the shape: [128, 4096, 4095]. If we reduce over dim=1, we get [128, 4095]. The number of output elements is 128*4095 = 524160. Each output element requires summing 4096 elements with stride 4095. We can assign one thread per output element, and have each thread loop over the reduction dimension. 4096 iterations per thread. That would be 524k threads, but each thread does 4096 loads. Total loads: 128*4096*4095 = 2,147,450,880 ~ 2.1e9 floats ~ 8.6 GB. GPU memory bandwidth is high (e.g., A100 has 1.5-2 TB/s). So 8.6 GB / 1.5 TB/s = ~5.7 ms. But PyTorch likely uses a reduction kernel that's quite optimized. We need to beat it. Better approach: use block-level reduction where each block handles one output element (or multiple). But with 4096 elements, we can have each block of e.g., 1024 threads loop 4 times, coalescing memory? But memory access pattern is strided: each thread loads x[b, i, d] where b=blockIdx.y, d=blockIdx.x, and i varies. The address for thread tid and loop j: x + b*4096*4095 + (j*blockDim.x+tid)*4095 + d. So consecutive threads access consecutive memory (since d is the same, b is same, and i differs by 1). Wait: i is the index along dim=1. The offset is b*stride0 + i*stride1 + d*stride2. stride1 = 4095, stride2 = 1. So for fixed b and d, varying i yields addresses that differ by 4095 floats. So consecutive i values are 4095 * 4 = 16380 bytes apart. That's not contiguous but it's a strided pattern. Within a warp, if we have thread tid handling i values spaced by blockDim.x (e.g., tid=0 handles i=0,4,8; tid=1 handles i=1,5,9), then consecutive threads access i and i+1 which are 16380 bytes apart. This is not coalesced: a warp would fetch 32 separated lines. However, modern GPUs have L1 cache, and each load is a 128-byte cache line. Since the stride is 4095 floats, each load is spread across many cache lines. So memory bandwidth utilization will be poor. But note: PyTorch likely uses a reduction algorithm that might also suffer from strided access. However, PyTorch's reduction over a non-contiguous dimension might use a tensor iterator that does some blocking or transposition implicitly. Actually, PyTorch's reduction kernel for mean over an inner dimension? Wait, dim=1 is the middle dimension. In a 3D tensor, the middle dimension might be reduced using some specialized kernel or a generic reduction that handles arbitrary strides. It might not be super fast. We can improve by transposing the tensor in shared memory or by using a different layout. But transposing 2GB of data might be expensive. However, we can have each block load a tile of data (e.g., a block of size 32x32) into shared memory in a coalesced way and then reduce. Wait, if we reduce over dim=1 (size 4096), the other dimensions are batch (128) and dim2 (4095). We can think of the data as a 2D matrix of size (128*4095) x 4096, but stored with the 128*4095 dimension being the "row" (since stride of dim1 is 4095). Actually, each "column" corresponds to a fixed (b,d) pair and varying i. The distance between consecutive rows (i.e., between (b,d) and (b,d+1)) is 1 float. So if we consider the 128*4095 dimension as the "row" index with stride 1, and dim1 (4096) as the "column" index with stride 4095, then we have a matrix that is stored in column-major order? Let's examine: offset = b * (4096*4095) + i * 4095 + d. If we fix i, then varying (b,d) yields contiguous memory (since d increments by 1, and after d reaches 4094, b increments, but there is a gap? Actually, for a given b, d from 0 to 4094 are contiguous. Then next b starts at offset + 4096*4095. So the segments for each b are contiguous, but the "reduction dimension" is strided by 4095. Thus, if we tile over (b,d) and i, we can load a tile of i (e.g., 32 values) for a contiguous range of d (and maybe b). For a fixed i, a contiguous range of d is contiguous in memory. So if we load across d first, we get coalesced memory. But we need to sum over i. So a block can process a range of d values (e.g., 32 or 64) and a range of i values (e.g., 32). It loads a tile of size i_range x d_range into shared memory, transposing it (since we load d contiguous), and then reduces over i in shared memory. Specifically, consider a thread block of shape (threads_d, threads_i) maybe. Or just 1D block of 256 threads. We assign each block to a tile: tile_d = blockIdx.x * tile_size_d + threadIdx.x, tile_i = blockIdx.y * tile_size_i + ... Actually, we have 3D: batch (128), i (4096), d (4095). The output is (128, 4095). We can have each block process a chunk of output elements for a given batch and a range of d. Alternative: launch a 2D grid: gridDim.x = ceil(4095 / tile_d), gridDim.y = 128. Each block handles one batch and a segment of d. Within the block, we reduce over i=0..4096. We can have each thread loop over i, but to get coalesced loads, we should have consecutive threads load consecutive d for the same i. So for a given i, thread tid loads x[b, i, d_start + tid]. This is coalesced because d is the last dimension. Then each thread accumulates its own sum for its d. After looping over all i, each thread has a partial sum for its d. Then we write the mean to output[b, d_start + tid]. This is extremely simple and memory-friendly! Let's verify: output shape is [128, 4095]. Block size = number of d's per block. But 4095 is not huge; we can have each block handle all d's for a given batch? That would be 4095 threads per block, which is > 1024, not allowed. So we split d into blocks. For example, block size 256 or 512. Then grid.x = ceil(4095 / 256) = 16, grid.y = 128. Total blocks = 2048. Each block has 256 threads. Each thread loops over i=0..4095 (4096 iterations). In each iteration, it loads one float: x[b, i, d]. Since d varies across threads, for a fixed i, the loads are coalesced (consecutive memory). After 4096 iterations, each thread has sum over i. Then compute mean = sum / 4096.0 and write to output[b, d]. This is essentially a reduction where the inner loop is over the reduction dimension, but the loads are coalesced because we parallelize over the non-reduced dimension. This is much better than PyTorch's generic reduction which might parallelize over the reduced dimension. Let's analyze memory traffic: Each load is 4 bytes. Total reads = 128*4096*4095*4 = 8.589 GB. With good coalescing, we should be near memory bandwidth limit. GPU specs: Let's get them.

The reference implementation computes `torch.mean(x, dim=self.dim)` where `x` has shape `[128, 4096, 4095]` and `dim` can be 0, 1, or 2. Let me analyze the memory access patterns: 1. If `dim=0`: reducing 128 elements across 4096*4095 = ~16M reduction operations. Each output element reads 128 floats. Total memory read: 128*4096*4095*4 bytes ≈ 8.5 GB. Output: 4096*4095*4 bytes ≈ 67 MB. This is memory bandwidth bound. 2. If `dim=1`: reducing 4096 elements across 128*4095 = ~525K reduction operations. Total memory read: same 8.5 GB. Output: 128*4095*4 bytes ≈ 2 MB. 3. If `dim=2`: reducing 4095 elements across 128*4096 = ~524K reduction operations. Total memory read: same 8.5 GB. Output: 128*4096*4 bytes ≈ 2 MB. The mean operation is simply sum / count. We want to write a custom kernel that is faster than PyTorch's native mean. For the H100, we have 3.35 TB/s bandwidth. PyTorch's mean operation should be quite optimized, but we can potentially do better by: 1. Using a more efficient thread block configuration 2. Coalesced memory access 3. Better use of shared memory / warp shuffle Let me think about the optimal approach. For mean along a dimension, we need to handle arbitrary dimensions. Let me write a general-purpose reduction kernel. For the most common case given in the benchmark (dim=1), the input shape is [128, 4096, 4095]. Reduction along dim=1 means: - Output shape: [128, 4095] - Each output element is the mean of x[b, :, d] for b in [0,128), d in [0,4095) To maximize memory throughput, we should: 1. Have each thread block work on a contiguous chunk of output elements 2. Use vectorized loads if possible 3. Use warp shuffle for intra-warp reduction Actually, let me think about memory coalescing more carefully. The input is [128, 4096, 4095]. - If dim=0: Access pattern for reading elements that contribute to a single output is strided by 4096*4095. Very strided, not coalesced. - If dim=1: For output[b, d], we read x[b, :, d]. This is contiguous in memory since dim=1 elements are contiguous. Great! - If dim=2: For output[b, d1], we read x[b, d1, :]. This is contiguous too. So dim=1 and dim=2 have contiguous access, while dim=0 is very strided. For the strided dim=0 case, we could use a different strategy. Let me write a fairly general kernel. For contiguous reductions (dim=1 or dim=2), we can have each thread read contiguous elements and reduce them. I'll write a kernel that handles the permutations of dimensions properly. Actually, let me first check if PyTorch's mean is already close to memory bound. But let me try writing a kernel and see. For contiguous reduction along a dimension, here's a good approach: - Use a 2D grid where one dimension maps to the batch of independent reductions - Each thread block reduces a chunk using parallel reduction in shared memory Let me write a kernel optimized for the given input shapes. Since dim can vary, I need to handle it generally. For shape [B, M, N] where we reduce over dim=1: - We have B*N independent reductions, each of length M - Each reduction reads M contiguous floats A good strategy: - Launch B*N threads total, group them into thread blocks - Each thread reads multiple elements (if M is large) and does partial sum - Then reduce within warp using shuffle, then within block using shared memory Wait, if we have 128*4095 = 524,160 independent reductions each of length 4096, and each block handles one reduction, we'd need 524,160 blocks. That's actually fine on H100 with many SMs. But a better approach: Each block handles multiple reductions to amortize block startup costs, or each block handles one reduction with many threads. For reduction of 4096 elements: - Use 1024 threads per block (one thread block per reduction) - Each thread sums 4 elements (4096/1024 = 4) - Warp shuffle reduction - This gives good memory coalescing Let me write this. I'll write CUDA code that handles generic reductions by transposing the problem into "rows" and reducing each row. Actually, to make it even more efficient, I should check what load_dim gives us. For contiguous dimension reduction, we can just treat the tensor as a collection of "rows". Here's my plan: Write a templated kernel that handles 1D, 2D, 3D tensor reductions. The kernel will find the contiguous dimension and optimize for it. Actually, let me just write a simple but efficient kernel for up to 3D tensors which handles any dimension. ```cuda __global__ void mean_reduce_kernel(const float* input, float* output, int dim, int ndim, int dim0, int dim1, int dim2, int numel, int reduce_size) { // We treat each "row" (independent reduction) as having a unique index // For 3D tensor [D0, D1, D2]: // dim=0: row_idx = d1 * D2 + d2, within-row offset = d0 // dim=1: row_idx = d0 * D2 + d2, within-row offset = d1 // dim=2: row_idx = d0 * D1 + d1, within-row offset = d2 } ``` Actually, let me think about this differently. For performance on the benchmark, I should optimize for the known shapes but keep it general enough. Let me write a kernel using a simple grid-stride loop approach within each thread, followed by block-level reduction. For 3D case [B, M, N] reducing over dim=1: - Total output elements = B * N - Launch one thread per output element, but since the reduction is 4096 elements long, we should probably have multiple threads work together. Let me try a two-level approach: 1. Each block handles one or more output positions 2. Threads within a block cooperatively reduce the corresponding input 3. Store result For the benchmark shapes, if dim=1: - Output has 128*4095 = 524160 elements - We'd want maybe 256 threads per block, each block handling one output - This gives us ~2048 blocks, which is enough to saturate H100 Wait, 524160 outputs, if each block handles one, that's plenty of parallelism. But 256 threads reducing 4096 elements means each thread only processes 16 elements. That might be okay. Actually, I think the best approach is: - Use a block size that matches the reduction size or is a power of 2 - For reduction size 4096, use 1024 threads (4 elements per thread) - For reduction size 4095, use 1024 threads (about 4 elements per thread) - For reduction size 128, use 128 threads Let me write a flexible kernel. Since we don't know dim ahead of time, I should handle it in the kernel. But actually, to be more efficient, I can use C++ templates or just handle it with if-statements in the host wrapper. Let me write host code that dispatches based on dim. For the actual kernel, I'll write a general approach: 1. Compute the output index from blockIdx 2. Compute the base input address for that output 3. Compute the stride to the next element in the reduction dimension 4. Have all threads in the block load elements and reduce Let me write this: ```cuda template <typename scalar_t> __global__ void mean_kernel_3d(const scalar_t* __restrict__ input, scalar_t* __restrict__ output, int dim, int D0, int D1, int D2) { // Each block handles one output element int out_idx = blockIdx.x; int total_out; int b, d1, d2; // coordinates in input int reduce_size; int base_idx, stride; if (dim == 0) { total_out = D1 * D2; if (out_idx >= total_out) return; d1 = out_idx / D2; d2 = out_idx % D2; reduce_size = D0; base_idx = d1 * D2 + d2; stride = D1 * D2; } else if (dim == 1) { total_out = D0 * D2; if (out_idx >= total_out) return; b = out_idx / D2; d2 = out_idx % D2; reduce_size = D1; base_idx = b * D1 * D2 + d2; stride = D2; } else { // dim == 2 total_out = D0 * D1; if (out_idx >= total_out) return; b = out_idx / D1; d1 = out_idx % D1; reduce_size = D2; base_idx = b * D1 * D2 + d1 * D2; stride = 1; } // Block reduction scalar_t sum = 0; for (int i = threadIdx.x; i < reduce_size; i += blockDim.x) { sum += input[base_idx + i * stride]; } // Reduce within block __shared__ scalar_t sdata[1024]; sdata[threadIdx.x] = sum; __syncthreads(); for (int s = blockDim.x / 2; s > 0; s >>= 1) { if (threadIdx.x < s) { sdata[threadIdx.x] += sdata[threadIdx.x + s]; } __syncthreads(); } if (threadIdx.x == 0) { output[out_idx] = sdata[0] / reduce_size; } } ``` Wait, this has a problem: for dim=0, the stride is large (D1*D2) and memory access is very uncoalesced. Let me think about how to handle that. For dim=0: To read x[d0, d1, d2] for fixed d1,d2 and varying d0, the address is d0*(D1*D2) + d1*D2 + d2. This is strided by D1*D2 = ~16M elements. So consecutive threads reading consecutive d0 values would access memory with a huge stride, which is terrible for coalescing. The solution: For strided reduction (dim=0 in our case), we need to transpose the access pattern. Instead of having the threads in a block collectively reduce one output, we should have each thread handle one output independently, and use grid-stride loops. Actually, for dim=0 with 4096*4095 outputs and only 128 elements each, having one thread per output is actually better! 524160 threads, each reading 128 strided floats. Each thread's reads are to the same cache line or nearby, and different threads are actually reading entirely different cache lines. With 128 threads per block, that's 4096 blocks, which is okay. Let me revise: for small reductions (e.g., reduce_size <= 128), use one thread per output. For large contiguous reductions, use one block per output with cooperative reduction. Actually, to keep it simple and efficient, let me just write separate kernels or a unified kernel that handles all cases reasonably well. For the strided case (dim=0), the most efficient is to have each thread loop over the reduction dimension independently. Since there are many outputs (524160), we launch a large grid. For contiguous cases (dim=1 or dim=2), cooperative block reduction is better. Let me write a unified kernel that uses a hybrid approach: ```cuda template <typename scalar_t> __global__ void mean_reduce_kernel(const scalar_t* __restrict__ input, scalar_t* __restrict__ output, int D0, int D1, int D2, int dim) { // For each independent reduction (each output element) int idx = blockIdx.x * blockDim.x + threadIdx.x; int total_outputs; int reduce_size; if (dim == 0) { total_outputs = D1 * D2; reduce_size = D0; } else if (dim == 1) { total_outputs = D0 * D2; reduce_size = D1; } else { total_outputs = D0 * D1; reduce_size = D2; } if (idx >= total_outputs) return; // Compute input coordinates and base index int b, d1, d2; int base_idx, stride; if (dim == 0) { d1 = idx / D2; d2 = idx % D2; base_idx = d1 * D2 + d2; stride = D1 * D2; } else if (dim == 1) { b = idx / D2; d2 = idx % D2; base_idx = b * D1 * D2 + d2; stride = D2; } else { b = idx / D1; d1 = idx % D1; base_idx = b * D1 * D2 + d1 * D2; stride = 1; } scalar_t sum = 0; #pragma unroll for (int i = 0; i < reduce_size; i++) { sum += input[base_idx + i * stride]; } output[idx] = sum / reduce_size; } ``` Wait, for dim=1 case: reduce_size = 4096, total_outputs = 524160. If we use 1 thread per output, we're launching 524160 threads. On H100, that's fine. Each thread does a loop of 4096 loads, all contiguous. This is actually memory bandwidth bound, which is good. But could cooperative reduction be better? Let's think: - With 1 thread per output: each thread does 4096 loads and stores 1 result. Loads are fully coalesced within a warp (since consecutive threads process consecutive outputs, and reads within a thread are contiguous... wait). Actually, for dim=1: - Thread 0 processes output[0] = mean of x[0, :, 0] - Thread 1 processes output[1] = mean of x[0, :, 1] - Thread 0 loads: x[0, 0, 0], x[0, 1, 0], x[0, 2, 0], ... - Thread 1 loads: x[0, 0, 1], x[0, 1, 1], x[0, 2, 1], ... At iteration i, warp reads x[0, i, 0], x[0, i, 1], x[0, i, 2], ..., x[0, i, 31]. These are contiguous! So memory access is perfectly coalesced. For dim=0: - Thread 0 processes output[0] = mean of x[:, 0, 0] - Thread 1 processes output[1] = mean of x[:, 0, 1] - Thread 0 loads: x[0, 0, 0], x[1, 0, 0], x[2, 0, 0], ... with huge stride - Thread 1 loads: x[0, 0, 1], x[1, 0, 1], x[2, 0, 1], ... At iteration i, warp reads x[i, 0, 0], x[i, 0, 1], x[i, 0, 2], ..., x[i, 0, 31]. These are contiguous too! The stride between x[i,0,j] and x[i,0,j+1] is 1. So actually dim=0 is ALSO coalesced when each thread handles one output independently! Wait, let me verify. For input [128, 4096, 4095], address(b, m, n) = b*4096*4095 + m*4095 + n. For dim=0, for a fixed output (d1, d2), we iterate over b. Address = b*4096*4095 + d1*4095 + d2. Consecutive b differ by 4096*4095. But when we have: - Thread 0: output(0,0), reads x[0,0,0], x[1,0,0], x[2,0,0], ... - Thread 1: output(0,1), reads x[0,0,1], x[1,0,1], x[2,0,1], ... - Thread 31: output(0,31), reads x[0,0,31], x[1,0,31], x[2,0,31], ... At iteration i (reading the i-th element in each thread's reduction): Warp reads x[i,0,0], x[i,0,1], ..., x[i,0,31]. These addresses are consecutive. So yes, perfectly coalesced within a warp! So the simple approach of one thread per output with a loop over the reduction dimension actually gives perfect coalescing for all dimensions! The key is that consecutive threads in a warp process consecutive output positions. The only issue is when the reduction size is large, a single thread accumulates in a register. For reduce_size=4096, each thread does 4096 iterations. This is actually fine for memory-bound code. But we need to be careful about instruction overhead. Actually, since this is memory bound (we're reading the entire tensor once), the key metric is memory bandwidth utilization. The kernel should be reading close to peak bandwidth. Let me think about performance. For reduce_size=4096, each thread does a loop of 4096. But a warp has 32 threads, each doing 4096 loads = 131072 loads per warp. At 32 loads per cycle (roughly), that's a lot of instructions. But the loads are memory-bound anyway. However, having a lot of instruction overhead per load could hurt. Cooperative reduction where a block jointly reduces the dim=1 might have better instruction efficiency because it processes contiguous chunks and can use wider vector loads or just have less loop overhead per memory transaction. Let me think about it more carefully. With one thread per output: - Total threads = total_outputs = 524160 for dim=1 - Warps = 16380 - Each warp performs 4096 iterations, each iteration doing 32 loads - Total loads = 524160 * 4096 ≈ 2.15 billion loads - But total unique floats read = 128 * 4096 * 4095 ≈ 2.15 billion - So each float is read exactly once. Good. With cooperative block reduction (e.g., 1024 threads per block, 1 block per output): - Total blocks = 524160 - Each block of 1024 threads cooperatively loads 4096 floats - Block 0 reads x[0, :, 0], block 1 reads x[0, :, 1], etc. - In each block iteration, threads load contiguous floats. Perfect coalescing. - Total loads is the same. But we have more efficient reduction via shared memory. For bandwidth-bound operations, both should perform similarly. But cooperative reduction might have lower instruction overhead per element because each thread only does 4 loads (4096/1024) and some fast shared memory operations, while the 1-thread approach does 4096 loads with 4096 loop iterations. Wait, actually both have the same number of loads per element. The difference is in reduction overhead. In the 1-thread approach, there's no reduction overhead, just accumulation. In the cooperative approach, you have the overhead of shared memory reduction. But for 4096 elements, the reduction is amortized. Let me think: 1-thread approach has loop overhead of 4096 iterations per thread. For 524160 threads, that's a lot of loop overhead instructions. Cooperative approach: 524160 blocks, each with 1024 threads doing 4 loads and then shared memory reduction. Loop overhead is much smaller (only 4 iterations). But block launch overhead might be higher. Actually, having 524160 blocks is a lot but H100 handles it well. But 524160 blocks might mean some SMs process many blocks sequentially, which is fine. Hmm, let me try a vectorized approach. For contiguous reduction dimensions, we can use float4 loads. For dim=1 or dim=2, consecutive elements in the reduction are contiguous, so we can load 4 at a time. This helps reduce instruction overhead. For dim=0, consecutive elements in the reduction are strided, so float4 doesn't help directly. But we can still do regular loads. I think the best approach is: 1. For contiguous reductions (dim=1, dim=2): use cooperative block reduction with vectorized loads 2. For strided reduction (dim=0): use 1 thread per output (still coalesced) Wait, actually for dim=0, consecutive output elements (read by consecutive threads) differ by 1 in the last dimension. So at reduction step i, threads read x[i, d1, d2], x[i, d1, d2+1], ... which are contiguous. So dim=0 is ALSO perfectly coalesced for the 1-thread-per-output approach! Let me reconsider. Could we also use vectorized loads for dim=0? No, because within a single thread's reduction, elements are strided by D1*D2. What about cooperative reduction for dim=0? That would require threads in a block to read strided elements, which is terrible for coalescing. So for dim=0, 1-thread-per-output is optimal. For dim=1,2, either approach works, but cooperative reduction with vectorized loads might be better. Actually, for a general solution, let me just write one kernel that handles all cases with 1 thread per output, and see how it compares to PyTorch. If it's already faster, great. If not, I'll optimize further. Actually though, I know PyTorch's reduction is extremely well optimized. For mean along a dimension, it likely uses cutlass or custom reduction kernels that are highly tuned. To beat it, I'll need something quite clever. Let me think about what PyTorch might do: - It likely detects the reduction axis and uses optimized kernels - For contiguous reductions, it may use warp-level or block-level cooperative reduction - It might also use tensor cores or other optimizations for large reductions One thing I can do is use half-precision intermediate accumulations or other tricks, but we need fp32 output. Actually, another idea: since mean = sum / n, and the input is fp32, maybe I can use mixed precision or use more threads. But that won't help bandwidth. Actually, on H100, memory bandwidth is 3.35 TB/s. Reading 8.5 GB at peak would take 2.5ms. Actually that's not right... 8.5 GB / 3.35 TB/s = 2.5ms? No: 8.5 / 3350 = 0.0025s = 2.5ms. Hmm, that seems like a lot. Let me recalculate. 8.5 GB = 8.5 * 10^9 bytes. 3.35 TB/s = 3.35 * 10^12 bytes/s. Time = 8.5e9 / 3.35e12 = 0.002537s = 2.54 ms. Actually, that seems like a reasonable time for a large reduction. But PyTorch's mean is probably quite optimized for this exact operation. Wait, I should also consider that we're reading the input and writing the output. If the input is in GPU memory and output is also in GPU memory, we read the entire input once and write a small output. So the operation is memory-bandwidth bound by the input size. To beat PyTorch, we probably need to: 1. Achieve closer to peak memory bandwidth 2. Reduce overhead (fewer CUDA API calls, etc.) One thing we can do is use `__ldg` (load through read-only cache/texture cache) for the input. This might help with caching. Or use `load_inline` with appropriate compiler flags. Actually, another consideration: PyTorch's `torch.mean` may be doing more than just summing and dividing. It might be creating intermediate tensors, doing two-pass algorithms, or using less optimal memory access patterns depending on the version. For our custom kernel, we can make sure to read each element exactly once and use efficient memory access. Let me try writing a simple kernel first and see how it does. Actually, I just realized: for contiguous reductions on large dimensions, the fastest approach on modern GPUs is often to use a tree reduction with vectorized loads. Let me design that. For dim=1 with shape [128, 4096, 4095]: - We have 128 * 4095 = 524160 independent reductions - Each reduction is over 4096 contiguous floats Optimized kernel for dim=1: ```cuda template <int BLOCK_SIZE> __global__ void mean_reduce_dim1(const float* input, float* output, int B, int M, int N) { // Each block handles one output element (b, n) int out_idx = blockIdx.x; int b = out_idx / N; int n = out_idx % N; if (b >= B) return; const float* row = input + b * M * N + n; float sum = 0; // Each thread loads BLOCK_SIZE elements with stride BLOCK_SIZE // But use vectorized loads: each thread loads 4 floats at a time // Actually, let's use a block-stride loop for (int m = threadIdx.x; m < M; m += BLOCK_SIZE) { sum += row[m * N]; } // Block reduction using shared memory __shared__ float sdata[BLOCK_SIZE]; sdata[threadIdx.x] = sum; __syncthreads(); for (int s = BLOCK_SIZE / 2; s > 32; s >>= 1) { if (threadIdx.x < s) { sdata[threadIdx.x] += sdata[threadIdx.x + s]; } __syncthreads(); } if (threadIdx.x < 32) { float val = sdata[threadIdx.x]; for (int s = 16; s > 0; s >>= 1) { val += __shfl_down_sync(0xffffffff, val, s); } if (threadIdx.x == 0) { output[out_idx] = val / M; } } } ``` Wait, `row[m * N]` means consecutive threads read `row[0], row[N], row[2N], ...` which are NOT contiguous if N > 1. For dim=1, contiguous elements in the reduction dimension are separated by N. Actually wait -- let me recheck. Input layout: x[b, m, n] is at address b*M*N + m*N + n. For fixed b and n, varying m: addresses are b*M*N + n, b*M*N + N + n, b*M*N + 2N + n, ... Consecutive elements in the reduction dimension are strided by N! Oh no, I had this wrong. The default memory layout in PyTorch is row-major, which for a 3D tensor [D0, D1, D2] means: x[i, j, k] is at i*D1*D2 + j*D2 + k. So for dim=1 reduction, elements that are summed together are at positions differing by D2 (the last dimension size). For dim=2 reduction, elements differ by 1. That's contiguous! For dim=1: strided by 4095, NOT contiguous. For dim=0: strided by 4096*4095, very strided. For dim=2: contiguous! Wait, PyTorch default is row-major (C-style), so the rightmost dimension is contiguous. So: - dim=2: contiguous reduction (stride 1) - dim=1: strided reduction (stride 4095) - dim=0: strided reduction (stride 4096*4095) Hmm, but earlier I said dim=1 was contiguous. That was wrong. Let me reconsider. For dim=1, the elements being reduced are x[:, m, :] for fixed outer dims... no wait, we reduce along dim=1 with fixed dim=0 and dim=2. Actually, re-reading: `torch.mean(x, dim=1)` means for each combination of other dimensions (dim=0 and dim=2), we average over dim=1. So for output[b, n], input elements are x[b, 0, n], x[b, 1, n], ..., x[b, M-1, n]. Address of x[b, m, n] = b*M*N + m*N + n. Consecutive m values differ by N in address. So yes, stride = N = 4095. For dim=0: output[m, n], input elements are x[0, m, n], x[1, m, n], ..., x[B-1, m, n]. Address = b*M*N + m*N + n. Consecutive b values differ by M*N = 4096*4095. For dim=2: output[b, m], input elements are x[b, m, 0], x[b, m, 1], ..., x[b, m, N-1]. Address = b*M*N + m*N + n. Consecutive n values differ by 1. Contiguous! So only dim=2 gives contiguous reduction. This changes things. For dim=2, we can use vectorized loads. For dim=1 and dim=0, we have strided access. Now, for the strided cases, having many threads work on the same output (cooperative reduction) means: - Thread 0 reads x[b, 0, n], x[b, 32, n], x[b, 64, n], ... (with block size 32, stride N) - Thread 1 reads x[b, 1, n], x[b, 33, n], x[b, 65, n], ... For a warp at iteration i, we read x[b, i, n], x[b, i+1, n], ..., x[b, i+31, n]. Wait, no. If threadIdx.x corresponds to the starting offset in the reduction dimension: In iteration 0: - Thread 0 reads x[b, 0, n] = b*M*N + 0*N + n - Thread 1 reads x[b, 1, n] = b*M*N + 1*N + n = b*M*N + N + n These differ by N! For N=4095, consecutive threads read addresses 4095 apart. That's NOT coalesced! For coalesced access, we want consecutive threads to read consecutive memory. So for dim=1: - To read x[b, m, n] coalesced, we'd want threads with consecutive n (not m). - But then reducing over m requires synchronization across different "lanes" of n. This is getting complicated. Let me think about what PyTorch does. PyTorch likely handles this by either: 1. Transposing the tensor when beneficial (so the reduction dim becomes contiguous) 2. Using specialized reduction kernels that handle strided access For dim=1 with stride 4095, consecutive threads reading consecutive m don't coalesce. But if we have each thread handle one (b, n) pair and iterate over m, then at step m: - Thread processing (b=0, n=0): reads x[0, m, 0] - Thread processing (b=0, n=1): reads x[0, m, 1] - Thread processing (b=0, n=2): reads x[0, m, 2] - ... - Thread 31 (b=0, n=31): reads x[0, m, 31] These ARE contiguous! So if we launch 1 thread per output and the threads are ordered by output index, then at each step of the loop, consecutive threads read consecutive memory addresses. This gives perfect coalescing! So the 1-thread-per-output approach actually gives perfect coalescing for ALL dimensions, as long as output elements are laid out in natural order and we loop over the reduction dimension! Let me verify for dim=0: - Thread 0: output(0, 0), reads x[0, 0, 0], x[1, 0, 0], x[2, 0, 0], ... - Thread 1: output(0, 1), reads x[0, 0, 1], x[1, 0, 1], x[2, 0, 1], ... - At step b: thread j reads x[b, 0, j]. These are contiguous! Perfect coalescing. For dim=1: - Thread 0: output(0, 0), reads x[0, 0, 0], x[0, 1, 0], x[0, 2, 0], ... - Thread 1: output(0, 1), reads x[0, 0, 1], x[0, 1, 1], x[0, 2, 1], ... - At step m: thread j reads x[0, m, j]. These are contiguous! Perfect coalescing. For dim=2: - Thread 0: output(0, 0), reads x[0, 0, 0], x[0, 0, 1], x[0, 0, 2], ... - Thread 1: output(0, 1), reads x[0, 1, 0], x[0, 1, 1], x[0, 1, 2], ... - At step n: thread j reads x[0, j, n]. Wait, are these contiguous? x[0, j, n] = 0*M*N + j*N + n Consecutive j differ by N = 4095. NOT contiguous! Oh! For dim=2, the 1-thread-per-output approach does NOT give contiguous access! Thread 0 and thread 1 read addresses that differ by N, not 1. So: - dim=0: coalesced (addresses differ by 1 at each step) - dim=1: coalesced (addresses differ by 1 at each step) - dim=2: NOT coalesced (addresses differ by N at each step) For dim=2, we need cooperative reduction. Block i handles output(b, m). Threads 0..N-1 (or fewer) read x[b, m, 0], x[b, m, 1], ..., which ARE contiguous within the warp. So optimal strategy: - For dim=0 or dim=1: 1 thread per output, loop over reduction - For dim=2: cooperative block reduction over contiguous dimension Wait, let me recheck dim=1. Output index mapping: for dim=1, total_outputs = B * N = 128 * 4095. Output (b, n) is at flat index b * N + n. Thread idx (b*N + n) reads x[b, m, n] in iteration m. At iteration m, warp threads 0..31 read: x[b0, m, n0], x[b0, m, n0+1], ..., x[b0, m, n0+31] Wait, (b, n) and (b, n+1) are consecutive output indices. So threads 0 and 1 could be output(0,0) and output(0,1). At step m, thread 0 reads x[0, m, 0] and thread 1 reads x[0, m, 1]. Addresses: 0*M*N + m*N + 0 = m*N and m*N + 1. These ARE contiguous! But what about thread N and thread N+1? Thread N is output(1, 0), thread N+1 is output(1, 1). At step m, thread N reads x[1, m, 0] and thread N+1 reads x[1, m, 1]. Addresses: M*N + m*N and M*N + m*N + 1. Also contiguous. So yes, within a warp, all threads read contiguous memory at each step. Perfect coalescing for dim=1. For dim=0: At step b, thread with output(m, n) reads x[b, m, n]. Consecutive outputs: output(m, n) and output(m, n+1) read x[b, m, n] and x[b, m, n+1], contiguous. Good. For dim=2: At step n, thread with output(b, m) reads x[b, m, n]. Consecutive outputs: output(b, m) and output(b, m+1) read x[b, m, n] and x[b, m+1, n]. Addresses: b*M*N + m*N + n and b*M*N + (m+1)*N + n. Difference is N = 4095. NOT contiguous. So for dim=2, cooperative block reduction is better. Now, for dim=2, with B*M = 128*4096 = 524288 independent reductions each over N=4095 contiguous elements. Using 1024 threads per block (one block per reduction): - Block processes x[b, m, threadIdx.x], x[b, m, threadIdx.x + 1024], etc. - But wait, N=4095, so thread 0 reads x[b, m, 0] and x[b, m, 1024], etc. - Thread i reads x[b, m, i] and x[b, m, i+1024]. These differ by 1024. - Within a warp, threads 0..31 read x[b,m,0..31], which are contiguous! Perfect. And we can use float4 loads for the contiguous dim=2! Actually, we can use float4 for block-stride cooperative reduction too. OK let me write a comprehensive kernel. I'll handle up to 3D tensors with generic indexing, but optimize the kernel dispatch for different dims. Actually, let me think about this more carefully for the benchmark. The benchmark might only test dim=1 (from get_init_inputs()), but I should handle all dims correctly. Let me write a kernel that handles any tensor shape up to some number of dimensions, using proper stride calculation. Actually, wait. Looking at get_init_inputs, it returns [1], which means dim=1 is what gets benchmarked. But we should still handle all cases correctly. For dim=1 with shape [128, 4096, 4095]: - Total outputs: 128 * 4095 = 524160 - Reduction size: 4096 - Memory access pattern: coalesced with 1-thread-per-output Each thread loops 4096 times. 524160 threads. That's a LOT of loop iterations but good memory coalescing. For 524160 threads, with block size 256, we have ~2048 blocks. Each SM can handle ~2048 threads, so about 32 blocks per SM concurrently if registers allow. Total warps = 16380. H100 has 132 SMs? Actually H100 SXM has 132 SMs or 144? Let me check... The specs say "132 SMs" for H100 SXM5. Actually it might vary, but let's say ~100+ SMs. 16380 warps / 132 SMs ≈ 124 warps per SM. Max warps per SM on Hopper is 64 (with 2048 threads / 32 = 64 warps per SM). So we'd need about 2 waves of execution. This is acceptable. But is the instruction overhead from 4096 loop iterations per thread significant? Each iteration does 1 FMA + 1 load + loop overhead. 4096 iterations is a lot. Total instruction count is high. For cooperative reduction with 1024 threads per block: - 524160 blocks total - 524160 blocks / 132 SMs = ~3970 blocks per SM - At 32 blocks per SM max concurrency, that's ~124 waves. That's a LOT of waves. - Block launch overhead could be significant. Actually, 524160 blocks of 1024 threads each is 536 million thread instances, but each thread only runs 4 iterations (4096/1024). Total "thread-iterations" is 2.1 billion in both cases. But launching 524160 blocks has much more block overhead than launching 2048 blocks. For small reductions (like dim=0 with reduce_size=128), 1-thread-per-output is better because block overhead is high. For large reductions (like dim=1 with reduce_size=4096), the tradeoff depends on block launch overhead. Actually, let me reconsider block counts. For dim=1 with 1-thread-per-output and block size 256: - Blocks = 524160 / 256 = 2047.5 ≈ 2048 blocks. - Each block has 256 threads, each doing 4096 iterations. - Total warps across all blocks = 524160 / 32 = 16380 warps. - Blocks per SM (on 132 SMs) = 2048 / 132 ≈ 15.5 blocks per SM. - Can run concurrently: ~8 blocks per SM (2048 threads / 256 = 8). Wait, max 32 blocks per SM, max 2048 threads per SM. So with 256 threads per block, max 8 blocks per SM concurrently. - Waves = 15.5 / 8 ≈ 2 waves. Good! For cooperative reduction (1024 threads per block, 1 block per output): - Blocks = 524160. - Blocks per SM = 524160 / 132 ≈ 3970 blocks per SM. - With 1024 threads per block, max 2 blocks per SM concurrently (2048 threads max). - Waves = 3970 / 2 ≈ 1985 waves. That's a huge number of waves! Block launch overhead on each wave could be significant. So 1-thread-per-output seems better for dim=1! Let me reconsider: maybe for dim=2 with N=4095, cooperative reduction with larger blocks per output is also bad. But for dim=2, we don't have good coalescing with 1-thread-per-output. Let's see: For dim=2 with 1-thread-per-output (block size 256): - Output shape: [128, 4096] = 524288 outputs - Blocks = 524288 / 256 = 2048 - Each thread does 4095 iterations - At step n, warp reads x[b, m, n] for 32 consecutive (b, m) pairs. - Addresses differ by N = 4095. Not coalesced but within a cache line... no, 32 elements spaced 4095 apart = 32*4095*4 bytes = ~524KB, which is way more than a cache line or even an L2 cache line. This is completely uncoalesced. For dim=2 with cooperative reduction (1024 threads per block, 1 block per output): - Blocks = 524288 blocks - Same huge number of blocks problem. Hmm, maybe there's a middle ground. For dim=2, we could have each block handle multiple outputs, or use warp-level reductions. Actually, for dim=2, here's a better approach: - Process multiple outputs within the same block - Have a 2D thread block where one dimension iterates over the reduction (n) and the other dimension iterates over different (b, m) outputs. For example: - BlockDim.x = 1024 (covers the reduction dimension 4095, with some threads doing 4, some doing 5 elements) - BlockDim.y = 2 (covers 2 different outputs) - Total threads per block = 2048 (max for SM) - Total blocks = 524288 / 2 = 262144 Still a lot of blocks. But within a block, threads with same threadIdx.y read contiguous memory? No, actually thread (x, 0) reads x[b0, m0, x], thread (x, 1) reads x[b1, m1, x]. These are far apart in memory. Not coalesced within a warp unless we arrange it carefully. To coalesce, we want consecutive threads in a warp to access consecutive memory. A warp is 32 consecutive threads in x, then moving to y. So if we arrange x as the "output" dimension and y as the reduction dimension: - Thread (0, n) reads x[b, m, n] - Thread (1, n) reads x[b+1, m, n] if output index is (b, m) ... wait, consecutive outputs (b,m) and (b,m+1) are not contiguous in memory for dim=2. Actually, consecutive outputs for dim=2: output index = b*M + m. Output (b, m) and (b, m+1) correspond to reading x[b, m, :] and x[b, m+1, :]. The starting addresses differ by N. If we want contiguous access, we need threads that read the same n but different m. At n, x[b, m, n] and x[b, m+1, n] differ by N. Still not contiguous. Hmm, actually for any output element (b, m) and reduction step n, we read x[b, m, n]. The only way consecutive threads read consecutive memory is if consecutive threads differ by 1 in n (with fixed b, m). That means the reduction dimension must be mapped to consecutive threads. So for dim=2, the cooperative approach where block handles one (b, m) and threads 0..1024 read n=0..1024 gives contiguous reads. This requires one block per output. Given the large number of outputs, maybe we do need one block per output for dim=2. Alternatively, we could transpose the input so dim=2 becomes dim=0 or dim=1, then use the 1-thread approach. But transposing has its own cost. For the benchmark, since dim=1 is the tested case, let me focus on that. But I still need correctness for all dims. For dim=2, PyTorch's built-in mean is probably extremely well optimized and might use tensor cores or complex memory access patterns. Beating it for dim=2 might be hard. But maybe PyTorch's generic reduction isn't as good as a custom kernel. Actually, on further thought, PyTorch dispatches to ATen's reduction kernels, which are highly optimized. For contiguous axis reduction, ATen likely uses very good kernels. For strided reduction, it may transpose or use multi-stage reductions. But let me think about whether we can still beat it. One advantage of a custom kernel is avoiding framework overhead and potentially doing exactly what we need without generality. Let me try a simple but carefully written kernel first. I'll make sure it's correct, then optimize. For dim=1, my 1-thread-per-output kernel should have: - 2048 blocks of 256 threads - Each thread loops 4096 times - Perfect coalescing within each warp Key optimizations: 1. Use `#pragma unroll` or manual unrolling 2. Process multiple elements per iteration (unroll by 4 or 8) 3. Use `const` and `__restrict__` qualifiers Actually, since memory is the bottleneck, unrolling the loop might not help much, but it can reduce instruction overhead. Wait, there's a better idea: process 4 elements per loop iteration and sum them up. Since we know the reduction size is 4096 (power of 2), we can do: ```cuda float sum = 0; for (int m = 0; m < M; m += 4) { sum += input[... + m * stride]; sum += input[... + (m+1) * stride]; sum += input[... + (m+2) * stride]; sum += input[... + (m+3) * stride]; } ``` This reduces loop overhead by 4x. But the compiler might already unroll loops. Another optimization: use warp-level implicit synchrony. Since all threads in a warp execute the same instruction, and memory loads have latency, having dependent operations like FMA immediately after load might cause stalls. But with enough warps per SM, we can hide latency. For 2048 blocks of 256 threads with 8 blocks per SM, that's 2048 threads per SM. Wait, max threads per SM is 2048, so exactly 100% occupancy with 8 blocks of 256 threads. That's actually great for latency hiding! Hmm, actually that's 256 threads per block * 8 blocks = 2048 threads per SM. But each SM can only handle 32 blocks maximum, and 2048 threads maximum. So with 256 threads/block, we can indeed run 8 blocks concurrently. Total blocks = 2048. On 132 SMs, we need ~15.5 blocks per SM total. With 8 blocks concurrently, that's 2 waves. After the first wave of 8 blocks finishes, the second wave of 7.5 blocks starts. Wait, this doesn't add up. Total blocks = 2048, so 2048/132 = 15.5 blocks per SM. With 8 blocks running concurrently, we have 2 waves: first wave of 8, second wave of 7.5. Total time for a block = time to process 4096 iterations * 256 threads. Latency hiding should be excellent with 2048 threads per SM (64 warps). For compute: each thread does 4096 FMAs. 256 threads per block * 2048 blocks = 524288 threads. Total FLOPs ≈ 524288 * 4096 ≈ 2.15 billion FMAs (but actually additions, so 2.15 GFLOPs). H100 does 67 TFLOPS, so this is COMPLETELY memory bound. The FLOPs are negligible. Memory bandwidth needed: reading 2.15B floats = 8.6 GB. Writing 524288 floats = 2 MB. Total ≈ 8.6 GB. At 3.35 TB/s, this is 2.57 ms. But we only read once, so memory efficiency matters. Wait, each float is read exactly once by exactly one thread. If we achieve near-peak bandwidth, the kernel should take ~2.5-3 ms. PyTorch's mean might do the same or it might do extra operations. Let me just write the kernel and see. Actually, I realize that PyTorch's reduction might use a different layout or might have to handle arbitrary tensor shapes/strides, which could add overhead. Our custom kernel can be optimized for exactly the shapes we care about. Let me write the kernel. I'll handle Nd tensors but optimize for 3D. For generality, let me use a compute_strides approach in the host, and pass the actual strides to the kernel. Actually, PyTorch tensors can be non-contiguous. But for this benchmark, they seem to be contiguous. Let me make sure to use tensor.strides(). For a 3D tensor x with shape [B, M, N] and strides [s0, s1, s2]: - x[b, m, n] is at b*s0 + m*s1 + n*s2. - For contiguous tensor, s0 = M*N, s1 = N, s2 = 1. To handle arbitrary dims and strides, the kernel needs: - Input strides for each dimension - Output shape and strides But to keep the kernel simple, I can compute everything on the host and pass the linear index formula. For reduction along dim `d` in a tensor with `ndim` dimensions: - The output shape removes dimension d. - For each output element, we need to know its base linear index in the input and the stride along dimension d. The base index = sum over i != d of (output_coord[i'] * input_stride[i]), where i' maps output dimensions to input dimensions. The reduction advances by input_stride[d] each step. This is easy to compute on the host for each output element, but to do it in parallel on GPU, we need each thread to compute its own coordinates. For a 3D tensor, this is straightforward. For general dims, it's a bit more work but still doable. Let me write a general Nd kernel using dynamic indexing, but optimize the common cases. Actually, let me write a kernel that handles up to 3D (since that's what the benchmark uses) and handles different dims with switch statements. ```cpp __global__ void mean_kernel_1d(const float* input, float* output, int size) { // trivial } __global__ void mean_kernel_2d_dim0(const float* input, float* output, int D0, int D1) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= D1) return; float sum = 0; for (int i = 0; i < D0; i++) { sum += input[i * D1 + idx]; } output[idx] = sum / D0; } __global__ void mean_kernel_2d_dim1(const float* input, float* output, int D0, int D1) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= D0) return; float sum = 0; for (int i = 0; i < D1; i++) { sum += input[idx * D1 + i]; } output[idx] = sum / D1; } ``` For 3D: ```cpp __global__ void mean_kernel_3d_dim0(const float* input, float* output, int D0, int D1, int D2) { int idx = blockIdx.x * blockDim.x + threadIdx.x; int total = D1 * D2; if (idx >= total) return; int d1 = idx / D2; int d2 = idx % D2; float sum = 0; for (int i = 0; i < D0; i++) { sum += input[i * D1 * D2 + d1 * D2 + d2]; } output[idx] = sum / D0; } __global__ void mean_kernel_3d_dim1(const float* input, float* output, int D0, int D1, int D2) { int idx = blockIdx.x * blockDim.x + threadIdx.x; int total = D0 * D2; if (idx >= total) return; int d0 = idx / D2; int d2 = idx % D2; float sum = 0; for (int i = 0; i < D1; i++) { sum += input[d0 * D1 * D2 + i * D2 + d2]; } output[idx] = sum / D1; } __global__ void mean_kernel_3d_dim2(const float* input, float* output, int D0, int D1, int D2) { // For dim=2, need cooperative reduction // Each block handles one output (d0, d1) int out_idx = blockIdx.x; int total = D0 * D1; if (out_idx >= total) return; int d0 = out_idx / D1; int d1 = out_idx % D1; float sum = 0; for (int i = threadIdx.x; i < D2; i += blockDim.x) { sum += input[d0 * D1 * D2 + d1 * D2 + i]; } __shared__ float sdata[1024]; sdata[threadIdx.x] = sum; __syncthreads(); for (int s = blockDim.x / 2; s > 32; s >>= 1) { if (threadIdx.x < s) sdata[threadIdx.x] += sdata[threadIdx.x + s]; __syncthreads(); } if (threadIdx.x < 32) { float val = sdata[threadIdx.x]; val += __shfl_down_sync(0xffffffff, val, 16); val += __shfl_down_sync(0xffffffff, val, 8); val += __shfl_down_sync(0xffffffff, val, 4); val += __shfl_down_sync(0xffffffff, val, 2); val += __shfl_down_sync(0xffffffff, val, 1); if (threadIdx.x == 0) output[out_idx] = val / D2; } } ``` Wait, but for dim=2, having one block per output with 524288 blocks is a lot. Is there a better way? What if for dim=2, I also use 1-thread-per-output but with a different thread ordering? The issue was that consecutive threads reading consecutive outputs don't access contiguous memory. But what if I process the outputs in a different order? For dim=2: each output is (d0, d1), reading x[d0, d1, n] for n in [0, D2). If I launch threads with index based on (d1, d0) instead of (d0, d1), then: - Thread 0: output(0, 0), reads x[0, 0, n] - Thread 1: output(1, 0), reads x[1, 0, n] - Thread D0: output(0, 1), reads x[0, 1, n] At step n, thread i reads x[i, 0, n], which is at address i*D1*D2 + 0*D2 + n = i*D1*D2 + n. Consecutive threads differ by D1*D2, which is still huge. Any ordering of (d0, d1) pairs will have consecutive threads read addresses that differ by multiples of D1*D2 or D2. Since D2=4095, even the smallest difference is 4095. So we can't get coalescing for dim=2 with 1-thread-per-output unless we interleave outputs and reduction steps in a complex way. So cooperative reduction is indeed needed for dim=2. But wait, what about a hybrid approach? Launch blocks where each block handles multiple outputs, but within a warp, threads work on the SAME output but different n (so same b, m, different n). This gives contiguous access. And each block can handle a few outputs sequentially. For example, blockDim.x = 1024, and each block handles 1 output. That's 524288 blocks. Actually, is 524288 blocks that bad? H100 with 132 SMs. Each SM can process many blocks over time. On Ampere/Hopper, block launch overhead is very low. And each block does 4095/1024=4 loads per thread plus a quick warp shuffle reduction. But wait, with block size = 1024 and 524288 blocks, total threads = 536 million. The kernel launch overhead is small but not zero. Maybe I can use a larger block size or have each block handle multiple outputs. But 1024 is already the max practical. Actually, total threads is just a scheduling concept. What matters is whether we can saturate the memory bandwidth. With 1024 threads per block and 64 warps per SM, we'd have 32 blocks running concurrently per SM... wait, 1024 threads per block means max 2 blocks per SM concurrently (2048 threads max / 1024 per block). So 524288 blocks / 2 = 262144 waves of 2 blocks per SM. With 132 SMs, that's only ~1985 waves total, but each wave is only 264 blocks. Wait, that IS a lot of waves. Actually, each "wave" on all SMs processes 132 * 2 = 264 blocks. Total waves = 524288 / 264 ≈ 1986 waves. If each wave takes, say, 100 cycles (mostly memory bound), that's 198600 cycles ≈ 0.1 ms at 2 GHz. Hmm, actually that's not bad. Plus the actual memory transfer time. Wait, I should also consider occupancy. With 1024 threads per block and 2 blocks per SM, we get 2048 threads per SM, which is 100% occupancy. So latency hiding should be excellent. The issue is just whether there are enough concurrent memory transactions. Actually, on modern GPUs, the memory controller can handle many concurrent requests. With 2048 threads per SM doing loads, and 132 SMs, that's 270,336 concurrent threads. At one load per thread every few cycles, that's plenty for the memory controller. For dim=2 with cooperative reduction, the memory access is: each warp loads 32 contiguous floats = 128 bytes, which is typically 2-4 memory transactions (depending on cache line size). This is efficient. Total memory reads = 8.5 GB. At 3.35 TB/s = 3350 GB/s. Time = 8.5/3350 = 2.54 ms. For 1986 waves at, say, 1000 cycles per wave (very rough estimate), and 2 GHz clock, time = 1986 * 1000 / 2e9 = 0.993 ms. That's faster than 2.54 ms, so memory bandwidth should be the bottleneck, not wave count. Actually, each wave processes a small amount of data: 264 blocks * 1024 threads * 4 loads per thread * 4 bytes = 4.3 MB per wave. At 3350 GB/s, reading 4.3 MB takes 1.28 microseconds. With 1986 waves, total time = 2.54 ms. So the bandwidth calculation checks out. The concern is whether the block launch overhead adds significantly. On H100, block launch is pipelined and very fast. Probably not a huge issue. For dim=1 with 1-thread-per-output: - 2048 blocks of 256 threads, each doing 4096 loads. - 2048 blocks / 132 SMs = 15.5 blocks per SM. - With 256 threads per block, max 8 blocks per SM concurrently. - 15.5 / 8 = ~2 waves. - Total memory = 8.5 GB. At peak bandwidth = 2.54 ms. - With 2 waves, each wave processes 1056 blocks * 256 threads * 4096 loads * 4 bytes = 4.4 GB. - Time per wave = 1.3 ms. Total = 2.6 ms. So both approaches are memory-bound and should take similar time. Alright, let me write the code. I'll implement: 1. A generic Nd dispatch function 2. Kernels for 1D, 2D, 3D optimized for each dim Actually, to make it really efficient, I should also consider using `float` accumulation. PyTorch might use a two-pass algorithm for numerical stability, especially for fp32. But our benchmark just checks tolerance of 1e-4, and summing 4096 floats with similar magnitudes should be fine with single-pass. Actually, for numerical stability, when summing many floats, pairwise or tree reduction is better than sequential accumulation. This is because adding a small number to a large accumulated sum can lose precision. For dim=1 with reduce_size=4096, a single thread sequentially adding 4096 floats could accumulate error. But 4096 is not THAT many, and with random data in [0, 1], the sum is around 2048. The smallest numbers we're adding are ~0, and the precision of fp32 is about 1e-7 relative, so absolute error around 2e-4 at most. The tolerance is 1e-4. Hmm, this might be borderline. PyTorch likely uses Kahan summation or tree reduction for numerical accuracy. To match PyTorch exactly, I should also use a more accurate reduction. Options: 1. Tree reduction using shared memory (but requires one block per output for strided dims, which is expensive) 2. Kahan summation in each thread 3. Two-pass algorithm (first compute partial sums, then combine) For the strided dims (dim=0, dim=1), tree reduction with one block per output doesn't coalesce memory. But we can do a two-level reduction: - First, each thread does a local sum, but with a tree-like pattern or Kahan summation - Then reduce across the block Actually, for 1-thread-per-output with loop unrolling, we can use Kahan summation: ```cuda float sum = 0; float c = 0; for (int i = 0; i < reduce_size; i++) { float y = input[idx + i * stride] - c; float t = sum + y; c = (t - sum) - y; sum = t; } ``` This is more accurate but adds some overhead (a few extra operations per iteration). Alternatively, for dim=1 where reduce_size=4096, we can chunk into groups and sum each group, then sum the groups. For example, sum 64 groups of 64 floats each. This creates a tree-like structure within a single thread. ```cuda float sum = 0; for (int i = 0; i < reduce_size; i += 64) { float group_sum = 0; #pragma unroll for (int j = 0; j < 64; j++) { group_sum += input[...]; } sum += group_sum; } ``` This reduces accumulated error. But does PyTorch use such techniques? Let me think. Actually, for fp32 tolerance of 1e-4 and typical random data in [0,1], simple summation of 4096 elements should be fine. The maximum error from sequential summation of 4098 [0,1] random numbers is roughly sqrt(N) * eps * mean ≈ 64 * 1e-7 * 0.5 ≈ 3e-6. That's well within 1e-4. For worst-case data (all 1s except one 1e-6), error could be larger but still around 1e-6. Actually wait, that's not right either. The error analysis for summation is: worst-case error bound is N * eps * max(|x_i|). For N=4096, eps≈1e-7, max=1, bound is 4e-4. That's slightly above 1e-4. But random data typically does much better. Still, to be safe, I should use a slightly more accurate summation or match PyTorch's behavior. PyTorch's mean on CPU uses pairwise summation for better accuracy. On GPU, it might use a similar strategy. Actually, looking at ATen's code, PyTorch's sum on GPU often uses a block-level parallel reduction for contiguous dimensions and might use multiple passes for strided dimensions. The exact algorithm depends on the tensor layout. Since we're matching within 1e-4, and the benchmark generates random data, simple summation should be fine. But if PyTorch uses a higher-accuracy algorithm, our results might differ slightly. Let me just try a simple kernel and see if it passes correctness. If not, I'll add Kahan summation. Actually, one more thought: for maximum performance, I should also consider using `__restrict__` and reading through `__ldg()` (load global through read-only data cache). `__ldg()` can help with caching and coalescing. OK, let me write the kernel now. Here's my plan for the code: ```python import torch from torch.utils.cpp_extension import load_inline cuda_source = ''' #include <torch/extension.h> #include <cuda_runtime.h> // Utility to compute linear index from coordinates // ... __global__ void mean_kernel_3d_dim0(const float* __restrict__ input, float* __restrict__ output, int D0, int D1, int D2) { int idx = blockIdx.x * blockDim.x + threadIdx.x; int total = D1 * D2; if (idx >= total) return; int d1 = idx / D2; int d2 = idx % D2; int base = d1 * D2 + d2; int stride = D1 * D2; float sum = 0.0f; #pragma unroll 4 for (int i = 0; i < D0; i++) { sum += input[base + i * stride]; } output[idx] = sum / D0; } __global__ void mean_kernel_3d_dim1(const float* __restrict__ input, float* __restrict__ output, int D0, int D1, int D2) { int idx = blockIdx.x * blockDim.x + threadIdx.x; int total = D0 * D2; if (idx >= total) return; int d0 = idx / D2; int d2 = idx % D2; int base = d0 * D1 * D2 + d2; int stride = D2; float sum = 0.0f; #pragma unroll 4 for (int i = 0; i < D1; i++) { sum += input[base + i * stride]; } output[idx] = sum / D1; } template <int BLOCK_SIZE> __global__ void mean_kernel_3d_dim2(const float* __restrict__ input, float* __restrict__ output, int D0, int D1, int D2) { int out_idx = blockIdx.x; int total = D0 * D1; if (out_idx >= total) return; int d0 = out_idx / D1; int d1 = out_idx % D1; int base = d0 * D1 * D2 + d1 * D2; float sum = 0.0f; for (int i = threadIdx.x; i < D2; i += BLOCK_SIZE) { sum += input[base + i]; } __shared__ float sdata[BLOCK_SIZE]; sdata[threadIdx.x] = sum; __syncthreads(); // Tree reduction #pragma unroll for (int s = BLOCK_SIZE / 2; s > 32; s >>= 1) { if (threadIdx.x < s) { sdata[threadIdx.x] += sdata[threadIdx.x + s]; } __syncthreads(); } // Warp reduction if (threadIdx.x < 32) { float val = sdata[threadIdx.x]; val += __shfl_down_sync(0xffffffff, val, 16); val += __shfl_down_sync(0xffffffff, val, 8); val += __shfl_down_sync(0xffffffff, val, 4); val += __shfl_down_sync(0xffffffff, val, 2); val += __shfl_down_sync(0xffffffff, val, 1); if (threadIdx.x == 0) { output[out_idx] = val / D2; } } } // 1D and 2D kernels similarly... torch::Tensor mean_reduce_cuda(torch::Tensor input, int dim) { int ndim = input.dim(); auto sizes = input.sizes(); auto output_sizes = std::vector<int64_t>(); for (int i = 0; i < ndim; i++) { if (i != dim) output_sizes.push_back(sizes[i]); } auto output = torch::empty(output_sizes, input.options()); if (ndim == 1) { // ... } else if (ndim == 2) { // ... } else if (ndim == 3) { int D0 = sizes[0]; int D1 = sizes[1]; int D2 = sizes[2]; const float* input_ptr = input.data_ptr<float>(); float* output_ptr = output.data_ptr<float>(); if (dim == 0) { int total = D1 * D2; const int block_size = 256; int grid = (total + block_size - 1) / block_size; mean_kernel_3d_dim0<<<grid, block_size>>>(input_ptr, output_ptr, D0, D1, D2); } else if (dim == 1) { int total = D0 * D2; const int block_size = 256; int grid = (total + block_size - 1) / block_size; mean_kernel_3d_dim1<<<grid, block_size>>>(input_ptr, output_ptr, D0, D1, D2); } else { int total = D0 * D1; const int block_size = 1024; int grid = total; mean_kernel_3d_dim2<1024><<<grid, block_size>>>(input_ptr, output_ptr, D0, D1, D2); } } else { // fallback to ATen or generic kernel // Actually, let's implement a generic Nd kernel } return output; } ''' functions = ['mean_reduce_cuda'] ``` Wait, I need to make sure I handle the case gracefully. PyTorch tensors can be any number of dimensions, but the benchmark uses 3D. Let me write a generic kernel too, or just use PyTorch as fallback. But we shouldn't use PyTorch as fallback in the kernel - we should implement it ourselves. For a generic Nd kernel, I can compute the strides and pass them to the kernel. Or I can write a kernel that handles up to, say, 6D with precomputed information. Actually, for a generic kernel, here's an approach: - Precompute `output_stride` for each output dimension (mapping to input dimension) - For each output index, compute the base input address and then loop adding `input_stride[dim]` ```cpp __global__ void mean_kernel_nd(const float* input, float* output, int ndim, int dim, int reduce_size, int total_outputs, const int64_t* input_strides, const int64_t* output_sizes, const int64_t* output_strides) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= total_outputs) return; // Compute output coordinates int64_t base = 0; int tmp = idx; for (int i = ndim - 2; i >= 0; i--) { int coord = tmp / output_strides[i]; tmp %= output_strides[i]; int input_dim = i >= dim ? i + 1 : i; base += coord * input_strides[input_dim]; } float sum = 0; int64_t stride = input_strides[dim]; for (int i = 0; i < reduce_size; i++) { sum += input[base + i * stride]; } output[idx] = sum / reduce_size; } ``` This handles any dimensionality. But it does coordinate computation in the kernel which adds overhead. For the benchmark shapes, this overhead might be acceptable. But for maximum performance, I should special-case 3D. Also, for the generic Nd case, memory coalescing depends on how we map threads to outputs. If we map thread idx to output linear index in natural order, then at each step of the loop, consecutive threads read consecutive memory IF the innermost output dimension corresponds to the innermost input dimension (excluding the reduced dimension). This is generally true for contiguous tensors. Actually, let's think about this more carefully. For a contiguous tensor, input strides are [D1*D2*..., D2*D3*..., ..., 1]. When we reduce over dim `d`, the output has the same structure but with dim `d` removed. For consecutive output indices, they differ by 1 in the rightmost non-reduced dimension. Consecutive threads reading these outputs will access input elements that differ by 1 in the rightmost dimension if the reduced dimension is not the rightmost. If the reduced dimension IS the rightmost, then consecutive outputs differ by 1 in the second-rightmost dimension, and the stride is D_last, so consecutive threads access memory D_last apart. This is the dim=-1 case, which is NOT coalesced. So the generic case has the same issue: dim=-1 is not coalesced for 1-thread-per-output. For correctness and simplicity, let me implement: 1. Special 3D kernels for dim=0,1,2 2. Generic Nd kernel as fallback But actually, I'll just rely on 3D for the benchmark and make the generic kernel correct. Let me also consider whether I should use smaller block sizes for better occupancy. For the 1-thread-per-output kernels, smaller blocks (e.g., 128) allow more blocks per SM concurrently, which means better latency hiding between warps of the same block? Actually, a block of 128 threads has 4 warps. Four blocks of 128 threads on an SM gives 16 warps. A block of 256 threads has 8 warps, and two blocks on an SM also gives 16 warps. So occupancy is the same (2048 threads / 128 = 16 blocks max, but limited by 32 per SM, so 16 blocks max of 128 threads). Actually wait: 2048 / 128 = 16 blocks. And max blocks per SM is 32. So 16 blocks of 128 threads = 2048 threads. Wait, 2048 is the thread limit per SM. So 16 blocks of 128 = 2048 threads. 8 blocks of 256 = 2048 threads. 4 blocks of 512 = 2048 threads. 2 blocks of 1024 = 2048 threads. For latency hiding, having more blocks is better because they can be scheduled independently. But with memory-bound kernels, the execution is mostly serialized by memory bandwidth. The difference between 128 and 256 thread blocks might be minimal. I'll use 256 for 1-thread-per-output (good balance) and 1024 for cooperative reduction. One more optimization: For dim=2, instead of 1024 threads, I could use 512 threads to allow 4 blocks per SM, or 256 to allow 8 blocks per SM. But for dim=2, each block handles one output, so more blocks per SM means handling more outputs concurrently, which is good. But with 512 threads, we still need D2/512 = 8 loads per thread. With 256 threads, D2/256 = 16 loads per thread. The total memory bandwidth is the same, but fewer blocks = less block overhead. Wait, with 1024 threads per block for dim=2: 524288 blocks. That's too many. What if each block handles MULTIPLE outputs? For example, 256 threads per block, each warp (32 threads) handles one output. Then 256/32 = 8 outputs per block. For dim=2: - Block size = 256 threads = 8 warps - Warp 0 handles output 0, warp 1 handles output 1, etc. - Each warp independently reduces D2 elements using warp shuffle - Block handles 8 outputs - Total blocks = 524288 / 8 = 65536 This is much more manageable! And within a warp, threads read contiguous memory. Let me write this: ```cpp __global__ void mean_kernel_3d_dim2(const float* __restrict__ input, float* __restrict__ output, int D0, int D1, int D2) { int block_out_base = blockIdx.x * 8; // 8 outputs per block int lane_id = threadIdx.x % 32; // position within reduction int warp_id = threadIdx.x / 32; // which output this thread handles int out_idx = block_out_base + warp_id; int total = D0 * D1; if (out_idx >= total) return; int d0 = out_idx / D1; int d1 = out_idx % D1; int base = d0 * D1 * D2 + d1 * D2; float sum = 0.0f; for (int i = lane_id; i < D2; i += 32) { sum += input[base + i]; } // Warp-level reduction sum += __shfl_down_sync(0xffffffff, sum, 16); sum += __shfl_down_sync(0xffffffff, sum, 8); sum += __shfl_down_sync(0xffffffff, sum, 4); sum += __shfl_down_sync(0xffffffff, sum, 2); sum += __shfl_down_sync(0xffffffff, sum, 1); if (lane_id == 0) { output[out_idx] = sum / D2; } } ``` This is elegant! 8 outputs per block, each warp handles one output. Warp shuffle reduction is very fast. Total blocks = 524288 / 8 = 65536. Blocks per SM = 65536 / 132 ≈ 496. With 256 threads per block, max 8 blocks per SM concurrently. Waves = 496 / 8 = 62 waves per SM. That's still a bit high but much better than 1986! And within each warp, memory access is perfectly coalesced: lane i reads base + i, base + i + 32, base + i + 64, etc. For dim=2, this should be very efficient. Can I apply a similar idea to dim=0 and dim=1? For those, 1-thread-per-output is already coalesced. But I can also use warp-level reduction if it helps. For the 1-thread-per-output approach: - Each thread reduces the entire dimension independently - No synchronization needed - But sequential accumulation of 4096 elements per thread For dim=1 with reduce_size=4096 and block size 256, each block has 8 warps. But each thread is independent. There's no warp-level cooperation. Actually, I can also use warp cooperation for dim=1: - Have a warp jointly reduce ONE output - But for dim=1, each output needs 4096 elements, and a warp has 32 threads. Each thread handles 128 elements. - Block of 256 threads = 8 warps = 8 outputs per block. - Total blocks = 524160 / 8 = 65520. This is very similar to the dim=2 approach! And memory access: for output (d0, d2), warp reads x[d0, lane_id, d2], x[d0, lane_id+32, d2], etc. At the first iteration, warp reads x[d0, 0, d2], x[d0, 1, d2], ..., x[d0, 31, d2]. These addresses are d0*D1*D2 + 0*D2 + d2, d0*D1*D2 + 1*D2 + d2, ... which differ by D2 = 4095. NOT contiguous! So for dim=1, having a warp jointly reduce one output does NOT coalesce memory! The 1-thread-per-output approach where consecutive threads handle consecutive outputs DOES coalesce memory. So for dim=1, 1-thread-per-output is better. For dim=2, warp-level cooperative reduction is better. For dim=0, 1-thread-per-output is better (each thread reads strided, but consecutive threads read contiguous). Let me verify dim=0 with 1-thread-per-output: Consecutive threads handle output(m, n) and output(m, n+1) (or output(m, n+1) and output(m, n+2)). Actually, for dim=0, outputs are [D1, D2] in shape. Flat index = m*D2 + n. Consecutive threads: thread i handles output with flat index i. At step b, thread i reads x[b, m_i, n_i] where m_i*D2 + n_i = i. Address = b*D1*D2 + m_i*D2 + n_i = b*D1*D2 + i. So consecutive threads read consecutive memory (addresses differ by 1). Perfect coalescing! Great, so the memory access patterns are: - dim=0: 1-thread-per-output, coalesced - dim=1: 1-thread-per-output, coalesced - dim=2: warp-cooperative, coalesced Now, one concern about dim=1 with 1-thread-per-output: sequential summation of 4096 floats per thread. Is this numerically stable enough? For 4096 random floats in [0,1], the sum is around 2048. Adding a float around 0.5 to an accumulator of 2048: the spacing between representable floats at 2048 is 2^11 * 2^-23 = 2^-12 ≈ 2.4e-4. So each addition could lose some precision. Total error might be on the order of 4096 * 2.4e-4 ≈ 1.0. Hmm, that's for the worst case. But with random numbers that average to ~0.5, the partial sums grow gradually, so the error per addition is smaller on average. PyTorch's mean of 4096 random numbers... let me think about what result we'd get. Actually, in fp32, adding 0.5 to 2048 gives exactly 2048.5 (I think). The spacing at 2048 is 2^-12 ≈ 0.00024, so numbers up to 0.00024 away from 2048 get rounded to 2048. But 0.5 is much larger than the spacing, so 2048 + 0.5 = 2048.5 exactly representable in fp32? 2048 = 2^11, 0.5 = 2^-1. 2048.5 = 2^11 + 2^-1. In fp32, the mantissa has 23 bits, so 2^11 + 2^-1 = (1 + 1/2 + ...) * 2^11. The fractional part is 2^-1 / 2^11 = 2^-12. The mantissa can represent numbers with precision 2^-23 relative to the exponent, so at exponent 11, it can represent 2^-23 * 2^11 = 2^-12. So 2^-12 is exactly representable. 2048.5 is exact. Actually, the issue is more subtle. At sum = 2048, the spacing between floats is 2^-12 ≈ 0.00024. Numbers being added are around 0.5, which is much larger than the spacing, so additions are exact for a while. Once the sum reaches ~4096, spacing is 2^-11 ≈ 0.00049. Still small. Actually, for additions a + b where |b| <= |a|, the rounding error is at most 0.5 * ulp(a). If a is growing, the absolute error grows. For summing N numbers of similar magnitude, the worst-case error is about N * eps * |sum| / 2 or something like that. For N=4096, eps=1.2e-7, sum≈2048, worst-case error ≈ 4096 * 1.2e-7 * 2048 ≈ 1.0. But that's truly worst-case (alternating signs or adversarial data). For random positive numbers, error is much smaller, roughly sqrt(N) * eps * sum ≈ 64 * 1.2e-7 * 2048 ≈ 0.016. Tolerance is 1e-4 absolute. Hmm, 0.016 >> 1e-4. Wait, that can't be right. Let me recalculate. The standard result for error in summation is that the computed sum S_hat satisfies: |S_hat - S| <= N * eps * sum(|x_i|) For x_i in [0,1], sum(|x_i|) ≈ 2048. N=4096, eps=1.2e-7. Error <= 4096 * 1.2e-7 * 2048 = 4096 * 2.4576e-4 ≈ 1.007. Hmm, that's the worst-case bound, which is quite loose. For random data, error is typically O(sqrt(N) * eps * sum(|x_i|)) ≈ 64 * 1.2e-7 * 2048 ≈ 0.016. 0.016 is still much larger than 1e-4. But this is the accumulated sum error. The MEAN is sum / 4096. Error in mean = error_in_sum / 4096 ≈ 0.016 / 4096 ≈ 3.9e-6. That's well below 1e-4! So the mean should be accurate enough with simple sequential summation. The error in the mean is what matters, and dividing by N reduces the absolute error. Actually, PyTorch likely computes sum first, then divides by N. So as long as our sum is close to PyTorch's sum within 4096 * 1e-4 = 0.41, we're fine. For random data, sequential summation should be well within that. Great, so numerical accuracy shouldn't be an issue. Let me also think about performance vs PyTorch. PyTorch's reduction for dim=1 on 3D tensors might not be as optimized as dim=2 because dim=2 is contiguous. PyTorch might: 1. For dim=2, use a highly optimized block-level reduction 2. For dim=1, either transpose and reduce, or use a less-optimized strided reduction kernel If PyTorch transposes for dim=1, it has to: 1. Allocate and fill a transposed tensor (extra memory allocation + copy) 2. Reduce over the now-contiguous dimension 3. The transpose itself is memory bandwidth bound, so total time ≈ 2x memory bound. If PyTorch doesn't transpose, it might use a multi-stage reduction or block-level reduction with strided access. Strided block reduction has poor coalescing. So our custom kernel for dim=1, which achieves perfect coalescing without transposing, might be significantly faster than PyTorch! This is a key insight. Let me write the kernel carefully. For the actual code, one more thing: I should pass the actual strides, not assume contiguous layout. PyTorch tensors can have arbitrary strides. For dim=1 with shape [B, M, N], strides [s0, s1, s2]: output[b, n] = mean_m x[b, m, n] = mean_m input[b*s0 + m*s1 + n*s2] For our 1-thread-per-output kernel: ```cpp int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= B * N) return; int b = idx / N; int n = idx % N; int base = b * s0 + n * s2; float sum = 0; for (int m = 0; m < M; m++) { sum += input[base + m * s1]; } output[idx] = sum / M; ``` Memory coalescing: at step m, thread idx reads input[b*s0 + m*s1 + n*s2]. Thread idx+1 reads input[b'*s0 + m*s1 + (n+1)*s2] where b' = (idx+1)/N, n+1 = (idx+1)%N. If n+1 < N, then b' = b, and the addresses differ by s2. If s2=1 (contiguous), addresses are consecutive. Good coalescing. If n+1 == N, then b' = b+1, n+1 = 0. Thread idx+1 reads input[(b+1)*s0 + m*s1]. The previous thread read input[b*s0 + m*s1 + (N-1)*s2]. The difference is s0 - (N-1)*s1. For contiguous tensor, s0 = M*N, s1 = N, s2 = 1. Difference = M*N - (N-1)*N = N*(M - N + 1). For M=4096, N=4095: difference = 4095 * 2 = 8190. That's a big gap! So within a warp, some threads might have a discontinuity. A warp of 32 threads will have at most one such discontinuity if N is a multiple of 32. But N=4095 which is not a multiple of 32 (4095 = 128*32 - 1). So every 32 threads, there will be some warp that straddles the boundary between different b values. Actually, let's be more precise. Threads with output indices 0 to 31: - Thread 0: (b=0, n=0) - Thread 1: (b=0, n=1) - ... - Thread 31: (b=0, n=31) They all have b=0, n=0 to 31. Addresses at step m: m*s1 + n*s2. Consecutive threads differ by s2. If s2=1, perfect coalescing. Threads 32 to 63: - Thread 32: (b=0, n=32) - ... - Thread 63: (b=0, n=63) Also b=0, perfect coalescing. ... Threads N-1 and N: - Thread N-1 = 4094: (b=0, n=4094) - Thread N = 4095: (b=1, n=0) Addresses: m*s1 + 4094*s2 and 1*s0 + m*s1 + 0*s2. Difference: s0 - 4094*s2. For contiguous: M*N - 4094*1 = 4096*4095 - 4094 = huge. But threads N-1 and N are in different warps (since N=4095, thread 4094 is in warp 127, thread 4095 is in warp 128). Actually, 4095 / 32 = 127.96, so thread 4094 is in warp 127 (threads 4064 to 4095), and thread 4095 is in warp 128 (threads 4096 to 4127). Wait, warp 127 has threads 4064...4095. So thread 4094 is in warp 127, thread 4095 is also in warp 127! Threads in warp 127: indices 4064 to 4095. That's all b=0, n=4064 to 4095. Still all b=0. Thread 4096 would be in warp 128, with b=1, n=0. So warp 127 and warp 128 are in different memory regions, but that's fine because they're separate warps. Within each warp, all threads access consecutive memory. So the only issue is within a warp that spans the boundary. Warp 127 doesn't span the boundary (indices 4064-4095 all have b=0). Warp 128 has indices 4096-4127, all with b=1, n=0-31. Also contiguous. What about the last warp of b=0 and first warp of b=1? They're different warps, so they don't need to coalesce with each other. Memory coalescing is within a warp. So for any b, n, as long as a warp doesn't span across different b values, we're fine. And since warps are aligned to 32-thread boundaries, a warp spans different b values only if N is not a multiple of 32, AND the boundary between b and b+1 falls within a warp. N = 4095. 4095 / 32 = 127.96875. So the warp boundaries are at multiples of 32: 0, 32, 64, ..., 4064, 4096, 4128, ... The output index boundary between b=0 and b=1 is at index 4095. Warp 127: indices 4064-4095 (all b=0). Warp 128: indices 4096-4127 (all b=1). The boundary is between warps, not within a warp! Because 4095 is odd and 4096 is a multiple of 32. Wait, 4096 / 32 = 128 exactly. So the boundary at 4096 is exactly on a warp boundary. N = 4095. For b=0, n ranges 0 to 4094. Warp 127 covers n=4064 to 4095, but n=4095 doesn't exist for b=0. So warp 127 for b=0 actually has only thread 4064 to 4094 (31 threads), and thread 4095 is b=1, n=0. Oh! So within warp 127: - Threads 4064-4094 (31 threads): b=0, n=4064-4094. Consecutive memory. - Thread 4095: b=1, n=0. Address jumps by s0 - 4094*s2. This is a discontinuity WITHIN the warp! Thread 4094 and thread 4095 are in the same warp. Their memory accesses differ by a huge amount. This breaks coalescing for that warp. How bad is it? For warp 127 at step m: - Threads 4064-4094 read m*s1 + (4064 to 4094)*s2. That's a contiguous range of 31 elements. - Thread 4095 reads s0 + m*s1 + 0*s2. This is far away. The warp issues two memory transactions: one for the 31 contiguous elements, and one for the single element. The single element transaction wastes bandwidth but is unavoidable. For other warps, they're fully coalesced. Out of 524160 / 32 = 16380 warps, how many have this issue? The boundary between b and b+1 occurs at every multiple of N=4095. So warps that contain both index 4095*k - 1 and 4095*k will have the discontinuity. 4095 / 32 = 127.96875. Since the boundary shifts relative to warp alignment: - Boundary 0: between indices -1 and 0 (not applicable) - Boundary 1: between 4094 and 4095. 4094 is in warp 127, 4095 is in warp 127 (since 4095 = 127*32 + 31). Actually 127*32 = 4064. 4064 to 4095 is warp 127. So yes, both in warp 127. - Boundary 2: between 8189 and 8190. 8189 / 32 = 255.90625. So 255*32 = 8160. Warp 255: 8160-8191. Both in warp 255. - In general, 4095*k mod 32 = (4095 mod 32) * k mod 32 = 31*k mod 32. For k=1: 31. Boundary at index 4095, which is in warp 127 (127*32 + 31 = 4095). Index 4094 is also in warp 127. Discontinuity within warp. For k=2: 62 mod 32 = 30. Boundary at 8190. 8190 = 255*32 + 30, 8189 = 255*32 + 29. Both in warp 255. So every boundary falls within some warp, causing a minor coalescing issue. Out of 128 warps per b (4096/32), typically 1 warp has a discontinuity. So about 1/128 ≈ 0.8% of warps are affected. For the affected warp, instead of 1 transaction of 32 elements, we have 1 transaction of 31 elements and 1 transaction of 1 element. The overhead is small. The overall impact on bandwidth is negligible. So 1-thread-per-output with block size 256 works well for dim=1. Great, I'm now quite confident in the approach. Let me write the complete kernel code. One more optimization: I can use `grid-stride` loop for the 1-thread-per-output kernels to handle arbitrary numbers of outputs. But that doesn't change much. Actually, for dim=1 with block size 256 and total outputs 524160, grid = 2048. This is fine. And for dim=2 with warp-cooperative and 8 outputs per block, total blocks = 524288 / 8 = 65536. Also fine. Let me now write the actual Python code. I need: 1. CUDA source code with the kernels 2. Host dispatch function 3. Python wrapper class ModelNew Also, I should compile with optimization flags. `load_inline` uses nvcc with `-O2` or `-O3` by default? I should pass extra flags. ```python extra_cuda_cflags=['-O3', '--use_fast_math'] ``` Wait, `--use_fast_math` might change numerical results slightly. The tolerance is 1e-4, so it might be okay, but let me be careful. For simple addition and division, `--use_fast_math` doesn't change anything. It mainly affects transcendental functions. So it should be safe. Let me also add `-maxrregcount=64` or something to limit register usage, but `load_inline` might not support all flags easily. Here's my complete plan for the kernel code: ```python cuda_code = r''' #include <torch/extension.h> #include <cuda_runtime.h> // 1D kernel __global__ void mean_reduce_1d(const float* __restrict__ input, float* __restrict__ output, int size) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx > 0) return; // only one output float sum = 0.0f; for (int i = 0; i < size; i++) { sum += input[i]; } *output = sum / size; } // 2D kernels __global__ void mean_reduce_2d_dim0(const float* __restrict__ input, float* __restrict__ output, int D0, int D1) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= D1) return; float sum = 0.0f; for (int i = 0; i < D0; i++) { sum += input[i * D1 + idx]; } output[idx] = sum / D0; } __global__ void mean_reduce_2d_dim1(const float* __restrict__ input, float* __restrict__ output, int D0, int D1) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= D0) return; float sum = 0.0f; for (int i = 0; i < D1; i++) { sum += input[idx * D1 + i]; } output[idx] = sum / D1; } // 3D kernels __global__ void mean_reduce_3d_dim0(const float* __restrict__ input, float* __restrict__ output, int D0, int D1, int D2) { int idx = blockIdx.x * blockDim.x + threadIdx.x; int total = D1 * D2; if (idx >= total) return; int d1 = idx / D2; int d2 = idx % D2; float sum = 0.0f; #pragma unroll 4 for (int i = 0; i < D0; i++) { sum += input[i * D1 * D2 + d1 * D2 + d2]; } output[idx] = sum / D0; } __global__ void mean_reduce_3d_dim1(const float* __restrict__ input, float* __restrict__ output, int D0, int D1, int D2) { int idx = blockIdx.x * blockDim.x + threadIdx.x; int total = D0 * D2; if (idx >= total) return; int d0 = idx / D2; int d2 = idx % D2; float sum = 0.0f; #pragma unroll 4 for (int i = 0; i < D1; i++) { sum += input[d0 * D1 * D2 + i * D2 + d2]; } output[idx] = sum / D1; } __global__ void mean_reduce_3d_dim2(const float* __restrict__ input, float* __restrict__ output, int D0, int D1, int D2) { int warp_id = threadIdx.x >> 5; // threadIdx.x / 32 int lane_id = threadIdx.x & 31; // threadIdx.x % 32 int out_base = blockIdx.x * 8 + warp_id; int total = D0 * D1; if (out_base >= total) return; int d0 = out_base / D1; int d1 = out_base % D1; int base = d0 * D1 * D2 + d1 * D2; float sum = 0.0f; for (int i = lane_id; i < D2; i += 32) { sum += input[base + i]; } // Warp shuffle reduction #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { sum += __shfl_down_sync(0xffffffff, sum, offset); } if (lane_id == 0) { output[out_base] = sum / D2; } } // Generic Nd kernel __global__ void mean_reduce_nd(const float* __restrict__ input, float* __restrict__ output, int ndim, int reduce_dim, int reduce_size, int num_outputs, const int64_t* __restrict__ input_strides, const int64_t* __restrict__ output_strides) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= num_outputs) return; // Compute base index from output index int64_t base = 0; int tmp = idx; for (int d = ndim - 2; d >= 0; d--) { int coord = tmp / output_strides[d]; tmp %= output_strides[d]; int input_d = d >= reduce_dim ? d + 1 : d; base += coord * input_strides[input_d]; } int64_t stride = input_strides[reduce_dim]; float sum = 0.0f; for (int i = 0; i < reduce_size; i++) { sum += input[base + i * stride]; } output[idx] = sum / reduce_size; } torch::Tensor mean_reduce_cuda(torch::Tensor input, int dim) { int ndim = input.dim(); TORCH_CHECK(dim >= 0 && dim < ndim, "dim out of range"); auto sizes = input.sizes().vec(); std::vector<int64_t> output_sizes; for (int i = 0; i < ndim; i++) { if (i != dim) output_sizes.push_back(sizes[i]); } auto output = torch::empty(output_sizes, input.options()); // Handle empty tensor if (input.numel() == 0) { return output; } const float* input_ptr = input.data_ptr<float>(); float* output_ptr = output.data_ptr<float>(); if (ndim == 1) { const int block_size = 256; mean_reduce_1d<<<1, 1>>>(input_ptr, output_ptr, sizes[0]); } else if (ndim == 2) { int D0 = sizes[0]; int D1 = sizes[1]; const int block_size = 256; if (dim == 0) { int grid = (D1 + block_size - 1) / block_size; mean_reduce_2d_dim0<<<grid, block_size>>>(input_ptr, output_ptr, D0, D1); } else { int grid = (D0 + block_size - 1) / block_size; mean_reduce_2d_dim1<<<grid, block_size>>>(input_ptr, output_ptr, D0, D1); } } else if (ndim == 3) { int D0 = sizes[0]; int D1 = sizes[1]; int D2 = sizes[2]; if (dim == 0) { int total = D1 * D2; const int block_size = 256; int grid = (total + block_size - 1) / block_size; mean_reduce_3d_dim0<<<grid, block_size>>>(input_ptr, output_ptr, D0, D1, D2); } else if (dim == 1) { int total = D0 * D2; const int block_size = 256; int grid = (total + block_size - 1) / block_size; mean_reduce_3d_dim1<<<grid, block_size>>>(input_ptr, output_ptr, D0, D1, D2); } else { int total = D0 * D1; int grid = (total + 7) / 8; mean_reduce_3d_dim2<<<grid, 256>>>(input_ptr, output_ptr, D0, D1, D2); } } else { // Generic fallback int num_outputs = output.numel(); int reduce_size = sizes[dim]; // Compute output strides std::vector<int64_t> output_strides(ndim - 1); int64_t stride = 1; for (int d = ndim - 2; d >= 0; d--) { output_strides[d] = stride; int input_d = d >= dim ? d + 1 : d; stride *= sizes[input_d]; } auto input_strides_tensor = torch::tensor(input.strides().vec(), torch::dtype(torch::kLong).device(input.device())); auto output_strides_tensor = torch::tensor(output_strides, torch::dtype(torch::kLong).device(input.device())); const int block_size = 256; int grid = (num_outputs + block_size - 1) / block_size; mean_reduce_nd<<<grid, block_size>>>( input_ptr, output_ptr, ndim, dim, reduce_size, num_outputs, input_strides_tensor.data_ptr<int64_t>(), output_strides_tensor.data_ptr<int64_t>() ); } return output; } ''' ``` Wait, for the generic fallback, I'm creating tensors inside the dispatch function. These tensors need to remain alive until the kernel finishes. Since this is CUDA, the kernel launch is asynchronous, so the tensors might go out of scope before the kernel runs if we're not careful. But the data_ptr() is copied to the kernel arguments, and CUDA copies the arguments synchronously before launching. So the actual pointer values are passed to the kernel. The tensors themselves can be destroyed, but the memory they allocated on the GPU remains until freed. Since we're using the default CUDA stream and the kernel is enqueued, the memory won't be reused until the kernel completes. Actually, PyTorch's allocator might recycle the memory... So I should use a different approach. Better: use `cudaMemcpyAsync` or just put the strides in a small array on the stack. But for arbitrary ndim, stack arrays are fine. Actually, I can avoid tensors altogether and use `cudaMalloc`/`cudaMemcpy` or just copy to constant memory. Actually, for simplicity, I can just allocate host arrays and copy them to device with `cudaMemcpyToSymbol` or `cudaMalloc` + `cudaMemcpy`. But this adds complexity. Since the generic case is just a fallback and might not be used in the benchmark, let me simplify. Alternative: use `at::Tensor` with proper stream synchronization. In PyTorch, when you create a tensor and pass its `.data_ptr()` to a kernel, the tensor will be kept alive as long as there's a reference to it. But creating tensors in the host function and letting them go out of scope at the end of the function... PyTorch uses reference counting for tensors on the Python side, but on C++ side, `at::Tensor` is a smart pointer. When it goes out of scope, it decrements the refcount. If the refcount reaches 0, the memory is freed. But since the kernel might still be running, this could be a use-after-free. Actually, in PyTorch's C++ API, tensor deallocation goes through the caching allocator, which tracks stream usage. When a tensor is freed, its memory is returned to the pool but marked as "in use" on the current stream until all pending operations complete. So it SHOULD be safe. But to be extra safe, I can avoid the generic case or use a different approach. For the generic Nd case, instead of creating tensors, I can compute the base index on the CPU for each output. But that's too much data to copy. Actually, I can just avoid the generic case entirely and only support up to 3D (which the benchmark uses). If someone passes a 4D tensor, I can fall back to PyTorch. But wait, I can't use PyTorch inside the kernel file because we'd get recursion or it would defeat the purpose. Actually, for this benchmark, I know the inputs are 3D. Let me just handle 1D, 2D, 3D. That's sufficient. But wait, `torch.mean` handles arbitrary dimensions, and our `ModelNew` should be a drop-in replacement. If the benchmark tests different dims, I need to handle them. Looking at the code, `get_init_inputs` returns `[1]`, so dim=1. And the tensor is 3D. So I only need to handle 3D with dim=1 correctly. For completeness, I'll handle 1D, 2D, 3D with any dim, and for >3D I'll use a generic kernel but in a safer way. Or I can just not support >3D since the benchmark doesn't test it. Actually, the problem says "Accept the same constructor arguments as Model" and "Implement forward() with the same signature". It doesn't say it has to work for ALL possible inputs, just that it should produce the same output for the tested inputs. But to be safe, I'll handle common cases. Let me write the complete code now. One more thing: I need to make sure to handle negative dims. PyTorch supports `dim=-1`, etc. Let me normalize `dim` in Python before passing to CUDA. ```python if dim < 0: dim += x.dim() ``` Also, I should check that the input is on CUDA. The benchmark seems to use CPU tensors in get_inputs(), but the model might be moved to CUDA. Let me check: the benchmark creates tensors with `torch.rand()`, which are on CPU by default. But for GPU kernel benchmarks, the evaluation likely moves them to GPU or creates them on GPU. Actually, looking at the benchmark code, `get_inputs()` returns `[x]` where x is on CPU. But the evaluation framework likely moves tensors to the GPU device before calling forward(). Let me assume the input is on CUDA when our forward is called. But to be safe, I can check: ```python if not x.is_cuda: return torch.mean(x, dim=self.dim) ``` Actually, the evaluation probably calls the model on GPU. Let me not add a CPU fallback. Also, I need to make sure the kernel handles the data type correctly. The benchmark uses fp32. I'll check that the input is float. ```python if x.dtype != torch.float32: return torch.mean(x, dim=self.dim) ``` This is a safety fallback. The benchmark says "produce numerically equivalent outputs at fp32 precision", so we should be fine. Alright, let me write the complete code and test it. Also, I want to add compiler optimization flags: ```python mean_reduce_cuda = load_inline( name='mean_reduce_cuda', cpp_sources='', cuda_sources=cuda_code, functions=['mean_reduce_cuda'], extra_cuda_cflags=['-O3', '--use_fast_math'], verbose=False ) ``` Wait, `load_inline` syntax: - `name`: module name - `cpp_sources`: C++ source (host code) - `cuda_sources`: CUDA source (device code) - `functions`: list of function names to expose - `extra_cuda_cflags`: CUDA compiler flags For `cpp_sources`, I can put the host wrapper there, or put everything in `cuda_sources`. Actually, `cpp_sources` is for CPU-only code and `cuda_sources` is for CUDA code. But for inline CUDA with pybind11, we typically put everything in `cuda_sources` and leave `cpp_sources` empty. Wait, looking at PyTorch docs, `load_inline` signature is: ```python torch.utils.cpp_extension.load_inline(name, cpp_sources, cuda_sources='', functions=None, extra_cuda_cflags=None, extra_cflags=None, ...) ``` The auto-generated pybind11 module will expose functions marked with `m.def()` or just functions declared in `functions=[]`. Since we don't write our own pybind11 module, we just list the function names in `functions`. One thing: the C++ function `mean_reduce_cuda` needs to be in a header or declared before the auto-generated pybind11 code. `load_inline` handles this, I think. Let me check by looking at examples. In many examples: ```python module = load_inline( name='my_module', cpp_sources='void my_func() { ... }', functions=['my_func'] ) ``` For CUDA: ```python module = load_inline( name='cuda_module', cpp_sources='', cuda_sources='__global__ void kernel(...) { } void my_func(...) { kernel<<<...>>>(...); }', functions=['my_func'] ) ``` This should work. Now let me write the final code. I'll make sure: 1. No custom PYBIND11_MODULE 2. Correct function names in `functions=` 3. Class is called `ModelNew` 4. Constructor takes `dim` argument Also, for correctness, I need to handle negative dims: ```python dim = self.dim if dim < 0: dim = x.dim() + dim ``` And I should check that when dim is negative, it normalizes correctly. One more thing: I should handle the case where the input has 0 elements. torch.mean returns nan for empty reductions. Let me match that behavior. In my kernels, if `reduce_size == 0`, dividing by zero would give inf or nan. Actually, in CUDA, `0.0f / 0` gives nan. So it should match PyTorch if we compute sum/0. But I need to make sure sum starts at 0. Actually, PyTorch's mean on empty tensor returns nan. I'll make sure my kernel handles this. For reduce_size=0, the loop doesn't run, sum=0, and 0/0 = nan in IEEE 754. So that works. Wait, in CUDA, integer division by zero would crash, but we're doing float division. `0.0f / 0` should give `nan` in ieee 754. Let me not worry about it since the benchmark doesn't test empty tensors. Let me write the final kernel code and test it with compile and correctness. I'll write the kernel file now. Wait, I should also consider that PyTorch's `torch.mean(x, dim=1)` on a tensor of shape [128, 4096, 4095] might use intermediate float64 accumulation for higher precision, and our fp32 accumulation might not match exactly. But for random data and 1e-4 tolerance, it should be fine. Actually, one more consideration: the fact that PyTorch's native mean might be vectorized using SIMD or use tensor cores. H100 has very fast memory bandwidth, so our memory-bound kernel should still be competitive if not faster. Let me also think about whether I can use tensor cores for this. Tensor cores do matrix multiplications, not reductions. So they're not useful here. I could use async copy (TMA) on Hopper, but that's complex and probably not worth it for this simple reduction. Alright, let me write the code. ```python import torch import torch.nn as nn from torch.utils.cpp_extension import load_inline CUDA_SOURCE = """ #include <torch/extension.h> #include <cuda_runtime.h> __global__ void mean_reduce_3d_dim0(const float* __restrict__ input, float* __restrict__ output, int D0, int D1, int D2) { int idx = blockIdx.x * blockDim.x + threadIdx.x; int total = D1 * D2; if (idx >= total) return; int d1 = idx / D2; int d2 = idx % D2; float sum = 0.0f; #pragma unroll 4 for (int i = 0; i < D0; i++) { sum += input[i * D1 * D2 + d1 * D2 + d2]; } output[idx] = sum / D0; } __global__ void mean_reduce_3d_dim1(const float* __restrict__ input, float* __restrict__ output, int D0, int D1, int D2) { int idx = blockIdx.x * blockDim.x + threadIdx.x; int total = D0 * D2; if (idx >= total) return; int d0 = idx / D2; int d2 = idx % D2; float sum = 0.0f; #pragma unroll 4 for (int i = 0; i < D1; i++) { sum += input[d0 * D1 * D2 + i * D2 + d2]; } output[idx] = sum / D1; } __global__ void mean_reduce_3d_dim2(const float* __restrict__ input, float* __restrict__ output, int D0, int D1, int D2) { int warp_id = threadIdx.x >> 5; int lane_id = threadIdx.x & 31; int out_base = blockIdx.x * 8 + warp_id; int total = D0 * D1; if (out_base >= total) return; int d0 = out_base / D1; int d1 = out_base % D1; int base = d0 * D1 * D2 + d1 * D2; float sum = 0.0f; for (int i = lane_id; i < D2; i += 32) { sum += input[base + i]; } #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { sum += __shfl_down_sync(0xffffffff, sum, offset); } if (lane_id == 0) { output[out_base] = sum / D2; } } torch::Tensor mean_reduce_cuda(torch::Tensor input, int dim) { int ndim = input.dim(); TORCH_CHECK(dim >= 0 && dim < ndim, "dim out of range"); auto sizes = input.sizes().vec(); std::vector<int64_t> output_sizes; for (int i = 0; i < ndim; i++) { if (i != dim) output_sizes.push_back(sizes[i]); } auto output = torch::empty(output_sizes, input.options()); if (input.numel() == 0) { return output; } const float* input_ptr = input.data_ptr<float>(); float* output_ptr = output.data_ptr<float>(); if (ndim == 3) { int D0 = sizes[0]; int D1 = sizes[1]; int D2 = sizes[2]; if (dim == 0) { int total = D1 * D2; const int block_size = 256; int grid = (total + block_size - 1) / block_size; mean_reduce_3d_dim0<<<grid, block_size>>>(input_ptr, output_ptr, D0, D1, D2); } else if (dim == 1) { int total = D0 * D2; const int block_size = 256; int grid = (total + block_size - 1) / block_size; mean_reduce_3d_dim1<<<grid, block_size>>>(input_ptr, output_ptr, D0, D1, D2); } else { int total = D0 * D1; int grid = (total + 7) / 8; mean_reduce_3d_dim2<<<grid, 256>>>(input_ptr, output_ptr, D0, D1, D2); } } else { // Fallback to ATen for unsupported dimensions // Actually, we shouldn't fallback in evaluation if it introduces bias // But for this benchmark, we only need 3D. // Use a generic approach output = torch::mean(input, dim); } return output; } """ # Compile inline mean_reduce_cuda = load_inline( name='mean_reduce_cuda', cpp_sources='', cuda_sources=CUDA_SOURCE, functions=['mean_reduce_cuda'], extra_cuda_cflags=['-O3', '--use_fast_math'], verbose=False ) class ModelNew(nn.Module): def __init__(self, dim: int): super().__init__() self.dim = dim def forward(self, x: torch.Tensor) -> torch.Tensor: if x.dim() != 3 or x.dtype != torch.float32 or not x.is_cuda: return torch.mean(x, dim=self.dim) dim = self.dim if dim < 0: dim += x.dim() return mean_reduce_cuda.mean_reduce_cuda(x, dim) ``` Wait, there's an issue with the fallback in C++: `output = torch::mean(input, dim);` - this would recursively call into our own code if we overloaded something. But since we're in a raw C++ function, it calls ATen's mean. However, `load_inline` might have issues with using ATen functions inside the kernel file. Actually, it should be fine. But wait, `load_inline` auto-generates pybind11 bindings. If I use `torch::mean` inside the C++ code, it links against libtorch which is available. So it should work. However, I want to avoid the fallback if possible since it might be slower. For the benchmark, I know the inputs, so the 3D path is what matters. Actually, there's one more issue: `load_inline` compiles the C++ code into a shared library. The C++ code uses `torch::Tensor` and ATen APIs, which are linked through the torch library. This should work. Let me also think about whether I can improve dim=2 further. What if I use vectorized (float4) loads for dim=2? ```cpp for (int i = lane_id; i < D2; i += 32) { sum += input[base + i]; } ``` Could become: ```cpp const float4* input_vec = reinterpret_cast<const float4*>(input + base); for (int i = lane_id; i < D2 / 4; i += 32) { float4 val = input_vec[i]; sum += val.x + val.y + val.z + val.w; } for (int i = (D2 / 4) * 4 + lane_id; i < D2; i += 32) { sum += input[base + i]; } ``` This would reduce the number of loads by 4x and improve memory bandwidth. But we need to ensure `base` is aligned to 16 bytes for float4. base = d0*D1*D2 + d1*D2. For D2=4095, if D2 is not a multiple of 4, then consecutive outputs (different d1) might not be aligned. For a single output, `base + lane_id` needs to be aligned for float4 loads starting at lane_id. If lane_id=0 and base is aligned, then float4 loads work. But for lane_id=1, base+1 is not aligned. Actually, in the cooperative warp approach, lane 0 reads base+0, base+32, base+64, ... Lane 1 reads base+1, base+33, base+65, ... For lane 0, if base is aligned, we can do float4 loads at base+0, base+4, base+8, etc. For lane 1, base+1 is not aligned. To use float4, all lanes in a warp would need to load float4 values. But lane 1 loads from base+1 which is not 16-byte aligned. Unaligned float4 loads are slower. What if we restructure: each warp handles 4 outputs (or some multiple), and lanes within a warp read aligned float4 for different outputs? Hmm, that gets complicated. For simplicity, I'll stick with float loads for now. If needed, I can optimize later. One more thought: for dim=2, I could also use the 1-thread-per-output approach if I reorder the memory access. But since the reduction dimension is contiguous, consecutive outputs have non-contiguous access. However, if I iterate over the reduction dimension in the inner loop: Wait, that's what the cooperative approach does. Alternatively, what if I use a block-level approach similar to dim=2 but handle 8 outputs per block? That's what I already have. Let me think about whether I should use shared memory + tree reduction vs warp shuffle. Warp shuffle is faster because it avoids shared memory access. My current approach uses warp shuffle within each warp, and each warp handles a different output. This is very efficient. For dim=1, could I also use warp-level cooperation? The issue was memory coalescing. But what if I have each warp reduce over the strided dimension by having lane i read element i, i+32, i+64, etc.? Then lane 0 reads base+0, base+32*stride, ... and lane 1 reads base+stride, base+33*stride, ... Within a warp at iteration 0, we read base+0, base+stride, base+2*stride, ... For stride=D2=4095, these are not contiguous. Bad. So no, warp cooperation doesn't work for dim=1. Let me also consider the compiler optimization. `#pragma unroll 4` tells the compiler to unroll the loop by a factor of 4. This can reduce branch overhead and allow better instruction scheduling. For a loop of 4096 iterations, unrolling by 4 reduces loop overhead by 4x. Actually, I realize that for dim=1 with block size 256, I could also consider using `__ldg` for cached loads. But on H100, regular loads through L1 cache should be fine for read-only access. OK, I think the code is good. Let me write it out properly and compile it.

The kernel passed correctness. Now let me run a quick static check before submitting. Let me also think about whether I can optimize further. The input shape is: batch_size=128, dim1=4096, dim2=4095, and dim=1 So the reduction is over a dimension of size 4096, and the output has shape [128, 4095]. My dim1 kernel handles reading per output element, where every block handles a contiguous chunk of D0 * D2 elements. Each thread handles 1 output element. For dim=1, each thread reads 4096 float values along dimension 1, which are strided by D2=4095 in memory. Wait, that's a huge stride! If each thread reads: input[d0 * D1 * D2 + i * D2 + d2], this is accessing memory with stride D2=4095. 4095 * 4 bytes = 16380 bytes, which is close to 16384. This is bad for memory coalescing because consecutive threads within a warp access: Thread 0: d0=0, d2=0 -> addresses 0, 4095, 8190, ... Thread 1: d0=0, d2=1 -> addresses 1, 4096, 8191, ... These are 1 float apart, so they DO coalesce within the warp (consecutive threads read consecutive floats). But each thread then jumps by 4095 floats. Wait, let me think about this more carefully. Actually, within the for loop, iteration i reads at offset (i * D2). So for consecutive threads t and t+1: Thread t (d2 = t): reads at i*D2 + t Thread t+1 (d2 = t+1): reads at i*D2 + t + 1 Within a warp, thread t reads address base + i*D2 + t. Since t ranges from 0-31, these are 32 consecutive floats. So this IS coalesced. But the stride between iteration i and i+1 for the same thread is D2 floats = 4095 * 4 bytes = 16380 bytes. So a warp reads 32 consecutive floats at offset i*D2, then at offset (i+1)*D2 = i*D2 + 16380. 16380 / 128 = 128 cache lines approximately? No, 16380 / 128 = 128-ish? Let's see: 16380 / 128 = 127.97. Hmm, that's almost exactly 128 cache lines. Actually, a cache line is 128 bytes on H100, which is 32 floats. So D2=4095 floats corresponds to 4095/32 ≈ 127.97, which means each iteration reads from a different cache line (mostly). But consecutive threads in the warp read from adjacent floats, so for a 32-thread warp reading 32 floats, that is exactly one cache line per iteration. But consecutive iterations access cache lines that are 4095 floats apart = 127.97 cache lines apart. Since the cache doesn't hold that many unique lines, this means the reads are not cached well. Each iteration for a warp accesses 1 new cache line, and since D0=4096, each warp accesses 4096 cache lines, or about 128 cache lines per warp (for 32 threads... wait, each warp accesses 4096 * 32 floats total, which is 4096 cache lines). Actually, per thread: D1=4096 iterations, 1 float per iteration = 4096 floats = 128 cache lines. Per warp (32 threads): 32 * 128 = 4096 cache lines = 512 KB. That's a LOT of cache lines! The L2 cache can hold this, but it's not ideal. A better approach for dim=1 reduction would be to have each thread block handle a tile of output elements where threads collectively load a strip of input and compute partial sums in shared memory or via warp reductions. Actually, looking at the problem again: D2 = 4095. That's very close to 4096. And D1 = 4096. So when reducing dim=1, each output element [d0, d2] = mean over d1 of input[d0, d1, d2]. If I transpose the tensor so that d1 is the last dimension, then I could use the efficient last-dim reduction. But transposing costs extra memory operations. Alternatively, I can use vectorized loads. If I have each thread read multiple consecutive elements along d1... but the stride is D2=4095, so consecutive elements along d1 are separated by 4095 floats. That's not consecutive in memory for a single thread. Wait, I think I should restructure the dim=1 reduction kernel. Let me think about this differently. For dim=1 reduction over D1=4096: - Output shape: [D0=128, D2=4095] - Total output elements: 128 * 4095 = 524160 If I use one block per output row (d0), and within a block have threads reduce along d2? No, we don't reduce along d2. Actually, let's use a different tiling strategy: - Each block handles one d0 value - Within a block, each thread handles one d2 value - Threads collectively load input in tiles along d1 dimension But the memory layout is row-major: input[d0, d1, d2] is at input[d0 * D1 * D2 + d1 * D2 + d2]. For a fixed d0, as d1 changes, the base address changes by D2. So for fixed d0 and d2, reading different d1 values means striding by D2=4095 floats. For a fixed d0, the memory layout is: [d1=0, d2=0], [d1=0, d2=1], ..., [d1=0, d2=4094] [d1=1, d2=0], [d1=1, d2=1], ..., [d1=1, d2=4094] ... So within a row (fixed d0, fixed d1), elements are consecutive in d2. Between rows (different d1), the offset is D2=4095. For dim=1 reduction, we want to read down columns (fixed d0, fixed d2, varying d1). These columns have stride D2. Actually, 4095 is just 1 less than 4096. So consecutive columns (d2 and d2+1) are offset by 1 in memory for the same d1. So: input[d0, d1, d2] = base + d1 * 4095 + d2 input[d0, d1, d2+1] = base + d1 * 4095 + d2 + 1 For a block of 256 threads, if I assign: thread tx handles d2 = tx, then all threads read consecutive memory within the same d1 iteration: warp reads addresses base + d1*4095 + (0..31) = consecutive 32 floats. This IS coalesced. So the access pattern in my current kernel IS coalesced! The issue is cache efficiency for the same thread across iterations: Thread tx reads: base + 0*4095 + tx, base + 1*4095 + tx, base + 2*4095 + tx, ... = tx, 4095+tx, 8190+tx, 12285+tx, ... The distance between consecutive reads for the same thread is 4095 floats = 16380 bytes = ~128 cache lines. So after reading element at offset tx, the next element at 4095+tx is 128 cache lines away. The L1 cache line size is 128 bytes = 32 floats. The L1 can hold about 128 cache lines (16 KB / 128 bytes = 128). But the working set for a single thread is tiny (just reusing one cache line). The issue is that 128 cache lines is a lot. Actually, for the same warp, iteration 0 reads cache line at offset k, iteration 1 reads cache line at offset k+128, etc. After 128 iterations, iteration 128 reads cache line k+128*128 = k+16384, and iteration 0's cache line k has been evicted by then (assuming LRU). So after that, every read is a cache miss. With 4096 iterations, we have 4096/128 = 32 cycles through the cache. So most reads will be L1 misses. But they might hit in L2. L2 cache is much larger. Each warp accesses 4096 cache lines = 512 KB. With 132 SMs and maybe 32 warps per SM, total working set is large but maybe L2 (50 MB on H100) can hold some of it. Actually, wait. We're launching 524160 threads total in my kernel (each handling one output element). With block size 256, that's 2048 blocks. Each block handles 256 output elements. The total memory read is input size = 128 * 4096 * 4095 * 4 bytes = ~8.5 GB. The output is 128 * 4095 * 4 = ~2 MB. Hmm, the H100 has 3.35 TB/s bandwidth. Reading 8.5 GB at this rate takes about 2.5 ms. The reference probably does the same. A simple reduction over a strided dimension like this is bandwidth-bound. The question is whether my kernel can achieve close to peak bandwidth. Let me think about how PyTorch implements mean over dim=1. It probably uses specialized reduction kernels with vectorized loads. PyTorch's reduction kernels are typically very optimized with unrolled loops, vectorized loads, etc. I think the key to beating PyTorch is: 1. Vectorized loads (load 4 floats at a time) 2. Reordering iterations for better L1 cache reuse 3. Using warp shuffle for the reduction sum For dim=1 with D2=4095, since D2 is not aligned to 4, vectorized float4 loads might be tricky. But the PyTorch mean also has this issue. Looking at the last dimension (dim=2) reduction, my kernel uses warp shuffle. That should be efficient. For dim=1, maybe I can restructure as follows: - Process output elements in tiles - For a tile of D2 values (e.g., 32 consecutive d2 values), have the warp collectively read a strip of memory Actually, a better approach: let's use float4 vectorized loads. For a block assigned to output elements [d0, d2_start:d2_start+3], each thread handles 4 consecutive d2 values. But since 4095 is not divisible by 4, we'd need to handle the tail. Wait, 4095 / 4 = 1023.75. So it's close but not clean. Hmm, let me think about the dim=1 reduction differently. Actually, looking at the get_init_inputs() = [1], so dim=1 in the test. That means our main kernel should be the dim1 one. Let me try a completely different approach for dim=1: - Instead of each thread computing one output element, have each warp compute multiple output elements simultaneously. - Or use shared memory to preload a tile of data and then reduce. Consider a 2D tile: - Each block has (warp_count, 32) threads - Block handles a 2D tile: d0 dimension and d2 dimension - We load input in tiles along d1, but each tile is a D0_block × D2_block tile of the input at a specific d1. Wait, for a fixed d1, the input slice is a D0 × D2 matrix stored in row-major order (with large stride between d0 values but consecutive within d0 row). Actually, for fixed d1, input[..., d1, :] has shape [D0, D2] and is stored with a stride of D1*D2 between consecutive d0 values. So it's not very coherent if we process multiple d0 values. Actually, within a single d0 value, for fixed d1, the d2 values are consecutive in memory. So if we process one d0 at a time and have 256 threads each handle 16 d2 values (float4 loads), we could be very efficient. Let me design a kernel: - Grid: D0 blocks (128 blocks), each block handles one d0 - Block: 256 threads - Each thread handles (D2 + 255) / 256 d2 values, or about 16 d2 values per thread - But we do float4 vectorized loads along d2 to read 4 consecutive floats Actually, better yet: have each block handle one d0, and threads within a block collectively load strips of the input and accumulate in registers. For block d0: - Thread tx handles a subset of d2 values - For each d1 from 0 to D1-1: - Load input[d0, d1, tx_start : tx_end] - Accumulate in registers After all d1, divide by D1 and write output. If each thread uses float4 loads, it handles 4 d2 values at once. With 256 threads, each block handles 1024 d2 values. But D2=4095, so we'd need 4 blocks per d0, or each block handles a segment. Wait, let me reconsider. If we have block (d0, d2_block), where d2_block ranges from 0 to ceil(D2/256): - Grid: [D0, ceil(D2/256)] - Block: 256 threads - Each thread handles one d2, but only for d2 values in its block With D0=128 and D2=4095, ceil(D2/256) = 16. So grid is 128x16 = 2048 blocks. Each block has 256 threads, but many threads are idle for the last block since 4095 = 16*256 - 1, so the last block only has 4095 - 15*256 = 4095 - 3840 = 255 active threads. Each active thread reads D1=4096 floats with stride D2=4095 between consecutive reads (same d1 dimension access pattern as before). If we use float4, we can read 4 floats at a time. But the addresses need to be aligned. For d2 values divisible by 4, float4 is aligned at input[d0, d1, d2] if (d0 * D1 * D2 + d1 * D2 + d2) % 4 == 0. Since D2=4095 is odd and D1=4096 is even: - For d0 even: d0 * D1 * D2 is even * even * odd = even. d1 * D2 = even * odd = even. So address parity = d2 % 4. For d2=0, address is divisible by 4, so float4 is aligned. - For d0 odd: similar analysis, address parity = d2 % 4. So for all d0 and d1, the address is divisible by 4 if and only if d2 is divisible by 4. So we can do 4 consecutive float reads starting from d2 values divisible by 4. But actually, using float4 might not help much if each thread's accesses across d1 are still strided with no cache reuse. Hmm, let me think about what makes PyTorch's reduction fast. PyTorch's reduction over an inner dimension typically has excellent memory access patterns. For dim=1, it's still reducing over an intermediate dimension. PyTorch would likely: 1. Use vectorized loads when possible 2. Process elements in a way that maximizes cache line utilization 3. Use shared memory for cross-warp reductions when needed Actually, I think the main issue with my current dim=1 kernel is that each thread reads D1=4096 elements with a large stride, which means each thread generates many cache misses. But since the access is coalesced within a warp, the overall memory bandwidth utilization might still be reasonable. The warp-wise access pattern: - Warp 0, iteration 0: reads addresses [0, 1, 2, ..., 31] (32 floats, 1 cache line) - Warp 0, iteration 1: reads addresses [4095, 4096, ..., 4126] (32 floats, 2 cache lines since 4095..4126 spans across a cache line boundary) Actually, consecutive iterations for the same warp read 32 floats starting at i*D2. Since D2=4095 = 4096 - 1 = 128*32 - 1, and a cache line is 32 floats: - Iteration 0 starts at offset 0, ends at 31. Fits in cache lines 0. - Iteration 1 starts at 4095, ends at 4126. 4095/32 = 127.96..., so starts in cache line 128 and ends in cache line 129 (since 4095 + 31 = 4126, and 4126/32 = 128.93). So spans 2 cache lines. - Iteration 2 starts at 8190, ends at 8221. 8190/32 = 255.93..., so cache lines 256 and 257. So each iteration reads 2 cache lines (except when perfectly aligned). Since D2=4095 is 1 less than 4096, there's always a misalignment. Each warp iteration reads 2 cache lines but only uses 32 floats = 1 cache line worth of data. So 50% cache line utilization. For 4096 iterations: warp reads ~8192 cache lines = 1 MB. With many warps running, the L2 cache might struggle. But the memory bandwidth is still the bottleneck. Actually, 50% cache line utilization for a strided access is actually pretty good. And since the total data is 8.5 GB, we're still mostly bandwidth-bound. Maybe I should try float2 or float4 loads to improve this. If each thread loads 4 consecutive floats (along d2), then consecutive iterations would be 4*D2 apart? No, if each thread handles 4 d2 values and loops over d1 with stride D2: addresses: i*D2 + tx*4 + {0, 1, 2, 3} Wait, that's still strided by D2. But we could instead have the inner loop iterate over d2 (vectorized) and have the outer loop over d1. That way, each thread reads 4 consecutive floats from the same d1, then increments d1. Actually, we should iterate: for d1 in range(D1): for d2_vector in range(d2_values_per_thread): load input[d0, d1, d2_start + d2_vector*4 : d2_start + d2_vector*4 + 4] This accesses consecutive memory within the inner loop, but the d1 loop has large stride. The memory access pattern within a warp: Thread tx handles d2 values [tx*4, tx*4+3]. At d1=0, it reads [tx*4, tx*4+1, tx*4+2, tx*4+3]. Consecutive threads read [0..3], [4..7], [8..11], etc. That's coalesced. With float4, each warp handles 256*4 = 1024 d2 values. But D2=4095, so we'd need 4 warps or restructure. Actually, let me think about this as processing in tiles. For each d1, we read a [d0, d2] tile. If we fix d0 and process multiple d2 values, the memory access for reading all d2 values of one d1 is exactly D2 consecutive floats, which is 4095 floats ≈ 128 cache lines. For dim=1 reduction, one approach is: 1. For each d0 and each chunk of d2: 2. Load the entire [d1, d2_chunk] tile into shared memory 3. Reduce along d1 But loading an entire [4096, 4095] tile into shared memory is too big (16M floats = 64 MB). A smaller tile: [tile_d1, tile_d2]. E.g., tile_d1=32, tile_d2=256. Size = 32*256*4 = 32 KB shared memory. For each d0: For each tile_d2: Initialize accumulator registers for each d2 in tile_d2 (say 256 floats in registers) For tile_d1_start in range(0, D1, tile_d1): Load tile [tile_d1, tile_d2] from input[d0, tile_d1_start:tile_d1_start+tile_d1, tile_d2_start:tile_d2_start+tile_d2] Accumulate into registers Divide by D1 and write output This is essentially transposing the access order so that consecutive reads are in memory. Let me design this: - Each block handles 1 d0 value and 1 tile of d2 values - Block size: 256 threads - Each thread handles 1 or more d2 values within the tile - All threads collectively load input by reading along d2 (consecutive memory) For tile_d2=256, tile_d1 can be any value. We loop over d1, and for each d1, we load 256 consecutive floats (the d2 tile), which is 8 cache lines and coalesces perfectly. With 256 threads, each thread loads 1 float per d1 iteration. After loading, we can accumulate in registers or shared memory. Since each thread handles a specific d2 value, it just adds the loaded value to its accumulator. Wait, this is exactly what my current kernel does! My current kernel: - Each thread handles 1 output element (1 d2 value for 1 d0) - The thread iterates over d1 and accumulates The difference is in the execution configuration. My current kernel assigns blocks based on total output elements, while this tile-based approach would be the same since each thread has one accumulator. So the memory access pattern is already the same. The question is whether float4 vectorization helps. Actually, with float4 loads: - Each thread handles 4 d2 values - Block of 256 threads handles 1024 d2 values - For each d1, block loads 1024 consecutive floats (32 cache lines), perfectly coalesced - Each thread accumulates 4 values For D2=4095, we'd need ceil(4095/1024) = 4 tiles per d0. - Tile 0: d2 = 0..1023 - Tile 1: d2 = 1024..2047 - Tile 2: d2 = 2048..3071 - Tile 3: d2 = 3072..4094 (only 1023 values) Grid: D0 * num_tiles = 128 * 4 = 512 blocks. Block: 256 threads. Each thread accumulates 4 floats over D1=4096 iterations = 4096*4 float loads = 16 loads of float4. This is quite nice! Let me implement this. For float4 loads, we need alignment. The address of input[d0, d1, d2] is: ptr = input + d0 * D1 * D2 + d1 * D2 + d2 For float4 load, we need ptr aligned to 16 bytes (4 floats). d0 * D1 * D2 + d1 * D2 + d2 = d2 + D2 * (d0 * D1 + d1) Since D2=4095 is odd, and d0 * D1 + d1 can be any integer. So: ptr % 4 = d2 % 4 (since D2 ≡ 1 (mod 4), actually D2=4095=1023*4+3, so D2 ≡ 3 (mod 4)) Wait, let me recalculate: 4095 / 4 = 1023 remainder 3. So D2 = 4k + 3. D2 ≡ 3 (mod 4). So ptr = d2 + (4k+3) * m where m = d0 * D1 + d1. ptr mod 4 = d2 + 3m (mod 4). For m=0 (d0=0, d1=0): ptr mod 4 = d2 mod 4. So for d2=0, ptr is aligned to 4. For m=1 (d0=0, d1=1): ptr mod 4 = d2 + 3 (mod 4) = d2 - 1 (mod 4). For d2=1, ptr is aligned to 4. So alignment depends on both d2 and d1! This means for a fixed d2 tile and varying d1, the float4 alignment changes with d1. That's problematic. Hmm, since D2 ≡ 3 (mod 4), the alignment of address for input[d0, d1, d2] with d2=0 is: - d1 even: aligned to multiples of ... let's see: d2=0, so ptr = (4k+3)*m. For m=0: 0. For m=1: 4095, 4095/4 = 1023.75. Not aligned. For m=2: 8190, 8190/4 = 2047.5. Not aligned. Actually, 4095*m mod 4 = 3*m mod 4, which is 0 when m=0, 3 when m=1, 2 when m=2, 1 when m=3, 0 when m=4, etc. So for d2=0, alignment cycles as d1 changes. This makes vectorized float4 loads very difficult for this particular shape. Maybe float2 (8-byte alignment)? D2 = 4095 is odd, so 4095*m mod 2 = m mod 2. So for d2=0: - d1 even: address even (aligned to 2) - d1 odd: address odd (not aligned to 2) Actually, any address is aligned to 2 bytes if it's even. An odd address is not aligned to 2 bytes. So float2 loads are also problematic. Hmm, that means even float2 vectorization is nontrivial for this problem. Let me try a different approach: process the input in a way that improves cache reuse. Since D2 is close to a power of 2, we get bad cache behavior because of cache line conflicts. If we transpose the computation or process in a different order, maybe we can do better. Actually, I just realized: if I can't easily beat PyTorch's reference with a custom kernel, maybe my current simple kernel is already competitive or even slower. But the reference might not be that fast for this specific shape. Let me think about what PyTorch does: PyTorch's at::mean would dispatch to a generic reduction kernel. For dim=1 over a non-contiguous dimension, it probably uses TensorIterator with a specialized reduction algorithm. TensorIterator has optimizations for reductions, including: 1. Reordering dimensions for better memory access 2. Vectorized stores on the output side 3. Unrolled inner loops But for a 3D tensor reducing over the middle dimension with D2=4095, the memory access pattern is fundamentally strided, and the performance ceiling is set by DRAM bandwidth with poor cache utilization. Wait, I have an idea. What if I reorder the computation to process multiple d1 values per warp in a way that leverages L1 cache? For instance, process d1 in chunks and have all warps in a block work on the same chunk: For a block handling multiple d0 values and/or multiple d2 values: 1. Load a strip of memory: input[d0, d1:d1+chunk, d2_start:d2_start+block_size] 2. Have threads accumulate partial sums Actually, this is getting complicated. Let me think about whether I can use cuDNN / cuBLAS or other primitives. But no, for a simple mean reduction, there's no direct cuDNN primitive. Alternative: use cub::BlockReduce for better reduction performance. But CUB might not be easily available. Let me try one more thing: what if I launch fewer threads with each thread doing more work, using the GPU registers more effectively? Another idea: For the dim=1 case specifically, since D1=4096 and D2=4095, what if I do shared memory transpose? I could load tiles like: - Tile size: [32, 32] in [d1, d2] space - Load 32x32 tile into shared memory (transposed or not) - Then each thread reduces along d1 within shared memory Wait, but to reduce along d1, I need all d1 values for a fixed d2. So within a tile [tile_d1=32, tile_d2=32]: - Thread (tx, ty) loads input[d0, tile_d1_start+ty, tile_d2_start+tx] into shared_mem[ty][tx] - After loading the tile, each thread tx reduces shared_mem[0..31][tx] along the d1 dimension For 256 threads, I can have 8 warps. Each warp can load a 32x8 tile (32 in d2, 8 in d1). After loading, each thread reduces its column. But the trick is: when loading the tile, we want coalesced memory access. If thread tx loads input[d0, d1, d2] with d2 varying and d1 fixed within a warp, that's coalesced. So if we have a warp load a row of 32 d2 values for a fixed d1, that's coalesced. Let me formalize: - Block: 8 warps × 32 threads = 256 threads - Grid: D0 blocks - Each block handles all d2 values for one d0 - Tile size in d1: 8 (one per warp) - Within a warp, 32 threads load 32 consecutive d2 values For a fixed d0 and for each chunk of 8 d1 values: For warp w in 0..7: d1 = chunk_start + w For thread tx in 0..31: load input[d0, d1, tx*4 : tx*4+3] as float4? No, each warp handles 32 d2 values, so tx loads one d2 value. Actually, if warp handles 32 d2 values: thread tx loads input[d0, d1, d2_start + tx] Then each thread tx has 8 values (one from each d1 in the chunk) and accumulates them. Wait, but this means thread tx handles d2 = d2_start + tx for ALL d1 values. After processing all chunks, thread tx has accumulated the sum over all d1 for its specific d2. This is exactly the same as my current kernel! The difference is: - Current: threads within a block are assigned to arbitrary d2 values (based on global thread ID: d2 = idx / D0, d0 = idx % D0). For my current dim1 kernel: idx = blockIdx.x * blockDim.x + threadIdx.x total = D0 * D2 d0 = idx / D2, d2 = idx % D2 So consecutive threads handle consecutive d2 values for the same d0, or when d2 wraps around, the next d0. Oh wait! In my current kernel: ```cuda int d0 = idx / D2; int d2 = idx % D2; ``` So consecutive threads handle consecutive d2 values. Within a warp (32 threads), they handle d2 values [k, k+1, ..., k+31] all for the same d0. For a given d1, these read input[d0, d1, k:k+31], which are 32 CONSECUTIVE floats! This is perfectly coalesced. So my current kernel is already reading perfectly coalesced within a warp for each d1 iteration. The issue is just that consecutive d1 iterations are spaced D2 floats apart. To improve cache reuse across d1 iterations: for a 32-thread warp handling d2=k..k+31, iteration i reads addresses i*D2 + [k, k+31]. The next iteration (i+1) reads (i+1)*D2 + [k, k+31]. The distance is D2 floats = 4095 floats. If we instead process the d1 dimension in chunks and have multiple warps cooperate: say warp 0 handles d1=0..7, warp 1 handles d1=8..15, etc. But this doesn't improve cache locality since each warp still has large strides. Actually, I think the real question is: will my simple kernel beat PyTorch? In many cases, a custom kernel can be competitive with PyTorch for simple operations. If the kernel achieves close to memory bandwidth, it might be hard to beat PyTorch significantly. But sometimes PyTorch has overhead from generic dispatch that we can avoid. Let me also consider that my dim2 kernel is already quite optimal with warp shuffle. For the actual test, dim=1, so I should focus on optimizing the dim1 kernel. Let me try adding float2 or float4 vectorized loads where possible. Actually, looking at alignment again: For my dim1 kernel, a thread handling d2 handles: input[d0, d1, d2] for d1 = 0, 1, 2, ... The address is base = d0*D1*D2 + d1*D2 + d2. For fixed d0 and d2, as d1 varies, addresses are d2, d2+D2, d2+2*D2, ... If d2=0, addresses are 0, 4095, 8190, ... Not float4 aligned for all d1. If d2=1, addresses are 1, 4096, 8191, ... Here 4096 is divisible by 16, so address 4096 is float4 aligned. But address 1 is not. So for a single thread, its access pattern across d1 is not aligned for vector loads. BUT, what if we have a thread read multiple d2 values simultaneously and increment d1 once per group? Instead of: for d1 in 0..D1-1: sum += input[d0, d1, d2] We do: for d1 in 0..D1-1: sum0 += input[d0, d1, d2] sum1 += input[d0, d1, d2+1] sum2 += input[d0, d1, d2+2] sum3 += input[d0, d1, d2+3] For consecutive d2 values within the same d1, these are consecutive in memory. We could read them as a float4! But wait, the alignment requirement: address = base + d1*D2 + d2 For float4, we need address divisible by 16 (4 floats). So we need (d1*D2 + d2) % 4 == 0. As analyzed before, since D2=4095 ≡ 3 (mod 4), this means: (d1 * 3 + d2) % 4 == 0. For d1=0: d2 % 4 == 0. So if we start at d2=0, the first read is aligned. For d1=1: (3 + d2) % 4 == 0 => d2 % 4 == 1. But at d1=1, we're reading the same d2 values as starting point, not d2=1. So for d1=1, reading input[d0, 1, 0:3] gives addresses 4095, 4096, 4097, 4098. 4095 is not divisible by 4. Oh, but I can start reading from a different d2 for each d1 to maintain alignment! If for d1=1, I start at d2 such that (3 + d2) % 4 == 0, i.e., d2 ≡ 1 (mod 4). But I need to compute the mean over the exact d2 values from 0 to D2-1, so I can't skip or reorder. This alignment issue is tricky with D2=4095. Actually, what if I just use non-aligned float4 loads? Modern GPUs support vector loads that aren't aligned (they translate into multiple instructions), but they might still be more efficient than 4 separate scalar loads. Alternatively, I could use the `__builtin_assume_aligned` or just do 2 float2 loads. For addresses like 4095, 4096, 4097, 4098: - I can read float2 at 4096 (aligned) and float2 at 4094 (not aligned well) or just read 4095 and 4097 as scalar + float2 at 4096. Actually, this gets messy. Let me just try to submit my current kernel and see what speedup we get. If it's not sufficient, I'll optimize further. Wait, let me also consider whether my dim2 kernel is actually the right one to optimize, since dim=1 is what's used in the benchmark. Let me try one more optimization for dim1: loop tiling / blocking along d1 to improve L1 cache hit rate. If we process d1 in tiles of, say, 32 values, and process all d2 values for that tile before moving to the next tile, we can potentially get better cache locality. But I'm not sure how to implement this without shared memory. Actually, I think the L1 cache on H100 is 256 KB per SM. With 4096 threads running per SM (128 threads per block, 32 blocks, though limited by other factors), each thread handles some d2. Let me just try running the static check and then submitting. If performance is bad, I'll iterate. Actually, wait. Let me reconsider the dim=2 kernel too, since we need all dimensions to work. The test uses dim=1, but correctness test may test all dims. For dim=2 with D2=4095, my kernel launches grid = (D0*D1 + 7) / 8 with 256 threads per block (8 warps). Each warp handles 1 output element and reduces 4095 values using lane_id loop with shuffle. This is efficient because: 1. Warp accesses are strided by 1 (consecutive lanes read consecutive memory for the same output element) 2. Shuffle reduction is fast For 4095 values with 32 lanes: each lane reads ceil(4095/32) = 128 values. The last few iterations (for lane 31) only read up to 4095 - 31 = value index, so 4095 - 3968 + 1 = 128th value is at index 3968+31=3999, wait: lane 31 reads: 31, 63, 95, ..., 31 + 127*32 = 31 + 4064 = 4095. Yes, exactly 128 values. And 4095 is the last index, so it works perfectly! Wait, D2=4095, so indices go from 0 to 4094. lane_id=31: 31, 63, ..., 31 + k*32 <= 4094 k_max = (4094 - 31) / 32 = 4063 / 32 = 126.96... So lane 31 reads 127 values: indices 31, 63, ..., 31+126*32 = 31+4032 = 4063. Wait, 4063 <= 4094. What about 4063 + 32 = 4095 > 4094. So lane 31 reads 127 values. lane_id=0: 0, 32, ..., 0+127*32 = 4064 <= 4094. 4064+32=4096 > 4094. So lane 0 reads 128 values. Some lanes read 128, some read 127. The loop: for (int i = lane_id; i < D2; i += 32) For lane_id=0: reads 128 values For lane_id=31: reads 127 values (since 4095 > 4094) Actually, 31 + 127*32 = 4063, and 4063 < 4095, so the condition i < D2 (4095) passes. Next: 4063 + 32 = 4095, and 4095 < 4095 is false. So lane 31 reads 128 values? Wait: i = 31: valid (31 < 4095) i = 63: valid ... i = 4063: valid i = 4095: NOT valid How many values: (4063 - 31) / 32 + 1 = 4032/32 + 1 = 126 + 1 = 127. For lane_id=0: i = 0, 32, ..., 4064. 4064 < 4095. Next = 4096, not valid. (4064 - 0) / 32 + 1 = 127 + 1 = 128 values. So some lanes do 128 iterations, some do 127. This imbalance might cause a tiny bit of divergence but it's minor since the difference is just 1 iteration out of 128. Actually, the loop `for (int i = lane_id; i < D2; i += 32)` means that some lanes finish one iteration earlier. In a SIMD execution model, all lanes execute the same number of loop iterations (the maximum needed by any active lane), and lanes that are done are masked out. So having some lanes do 127 and some 128 means the warp executes 128 iterations, with lane 31 masked out on the last one. This is perfect - no divergence cost. OK, I'm fairly confident in the kernels. Let me do a static check and then submit. Actually, I just realized I might want to verify static check before submitting. Let me do that.