reasoning
The kernel is correct. Now let me optimize it further for performance. The bottleneck here is memory bandwidth since we're doing a sum reduction. Key optimizations:
1. Use vectorized loads (float4) to read 4 floats at a time
2. Better register usage and instruction scheduling
3. Possibly use shared memory for cooperative reduction when reduce_size is large
Given the problem dimensions: batch=128, dim1=4096, dim2=4095, reduce_dim=1
- outer_size = 128
- reduce_size = 4096
- inner_size = 4095
So each thread sums 4096 elements along the reduce dimension. Each output element is computed by one thread. The memory access pattern is strided by inner_size (4095).
To optimize:
1. Use float4 loads to reduce memory instructions
2. Use grid-stride to allow fewer blocks with more work per thread
3. Ensure good occupancy
Let me also consider using 1D grid with grid-stride loop for better GPU utilization. The current kernel launches ~128 * 16 = 2048 blocks, which is small but might be okay. Each block has 256 threads, so total threads = 524288. Each thread processes one output element and sums 4096 elements.
Actually, wait - inner_size = 4095, and we're using 256 threads per block. So we need ceil(4095/256) = 16 blocks_x per outer. Total blocks = 128 * 16 = 2048 blocks. On H100 with many SMs, this should be okay.
But the issue is that memory accesses are strided by 4095, which is not great for L2 cache. Each thread reads at stride 4095*4 = 16380 bytes, which means threads in the same warp access very different cache lines. This could be bad for memory coalescing.
Wait, let me reconsider the layout:
- Memory layout is [outer, reduce, inner]
- blockIdx.y = outer_idx (fixed for all threads in a block)
- threadIdx.x + blockIdx.x*blockDim.x = inner_idx
For a given warp, inner_idx values are consecutive (0, 1, 2, ..., 31). The elements accessed are at:
base + r * inner_size + inner_idx
For r=0, warp accesses base, base+1, base+2, ..., base+31 - this is coalesced!
For r=1, warp accesses base+inner_size, base+inner_size+1, ..., base+inner_size+31 - also coalesced!
So memory accesses ARE coalesced within a warp since consecutive threads access consecutive inner elements. This is good.
But the strided pattern means that consecutive memory loads are far apart in memory. We could use persistent threads or prefetching.
Actually, let me think about vectorization. float4 loads would help here because:
- Each thread loads 4 consecutive floats (inner_idx, inner_idx+1, inner_idx+2, inner_idx+3)
- But the next load is at inner_idx + 4*inner_size away
Wait, I'm confused. Let me trace through:
offset = outer_idx * reduce_size * inner_size + inner_idx
For inner_idx = 0:
r=0: offset=0
r=1: offset=inner_size
r=2: offset=2*inner_size
So a single thread reads at offsets 0, inner_size, 2*inner_size, 3*inner_size, etc.
Consecutive threads in a warp read at 0, 1, 2, 3,... for r=0, then inner_size, inner_size+1, ... for r=1.
Float4 loading wouldn't work well here per thread because a single thread would read 0, 1, 2, 3 but then next is inner_size, not 4. So float4 per thread would read non-consecutive memory.
Actually, to do vectorization, we'd need to have a thread load elements that are consecutive in memory. But in our case, consecutive elements in the reduce dimension are inner_size apart. So a float4 load from offset_t = outer_idx * reduce_size * inner_size + r * inner_size + inner_idx would load 4 consecutive floats starting at that offset. Each float4 load would correspond to summing 4 consecutive inner elements for the same r value... but we sum across r, not across inner.
Actually wait, for a single output element we fix outer_idx and inner_idx, and sum over r. So float4 doesn't naturally apply unless reduce_size is a multiple of 4 and we process 4 r values at a time - but then those 4 values are at offsets inner_idx, inner_idx + inner_size, inner_idx + 2*inner_size, inner_idx + 3*inner_size, which are NOT consecutive. So float4 won't work for a single thread's reduction.
However, we can have each thread process multiple inner elements using vector operations to increase compute intensity. For instance, a thread could compute 4 output elements at once:
sum0 += row_ptr[r * inner_size + 0];
sum1 += row_ptr[r * inner_size + 1];
sum2 += row_ptr[r * inner_size + 2];
sum3 += row_ptr[r * inner_size + 3];
But these are strided by inner_size, not consecutive. So float4 doesn't help directly.
Actually, consecutive output elements (along inner dimension) ARE consecutive in memory for a fixed outer and r. So if we want to load from the same r but 4 consecutive inner values, that IS a float4:
row_ptr[r * inner_size] gives 4 floats at positions r*inner_size, r*inner_size+1, etc.
But wait, a single thread computing one output element only needs one value per r. If we have a thread compute 4 output elements (inner_idx, inner_idx+1, inner_idx+2, inner_idx+3), then for each r it needs those 4 values, which ARE consecutive in memory! So yes, we can use float4 loads.
Let me design:
- Each thread computes 4 output elements
- Thread idx i computes inner_idx = blockIdx.x * blockDim.x * 4 + threadIdx.x * 4
- For each r, load float4 at offset = outer_idx * reduce_size * inner_size + r * inner_size + inner_idx
- Sum the 4 components separately
This gives us 4x fewer loads (each load is float4) and better memory bandwidth utilization.
Actually, since consecutive threads in a warp load consecutive float4s, the memory accesses are perfectly coalesced (128 bytes per warp, optimal).
Let me try this approach. Also need grid-stride since inner_size=4095, and we use 256 threads * 4 = 1024 elements per block, so we need ceil(4095/1024) = 4 blocks_x per outer. Total blocks = 128 * 4 = 512 blocks. This might be a bit small, but should be fine.
Actually, I'll also use a 1D grid with a more sophisticated mapping.
Wait, inner_size=4095. If each thread handles 4 elements, total elements per outer = 4095. Threads per block = 256, elements per block = 1024. So blocks needed = ceil(4095/1024) = 4. Total blocks = 128 * 4 = 512. That's low.
Let me also have each thread process multiple groups with grid-stride. For example, with 256 threads and 4 inner elements each, step = 1024. After processing 1024 elements, we'd move by gridDim.x * 1024.
For better GPU utilization, I should increase the number of blocks.
Alternative: use 1D grid where each block processes one outer element and loops over inner elements with block stride.
Actually, let me try to just make a more aggressive version with vectorized loads and see if it's faster. I'll also make sure we have enough blocks.
Or I can process multiple output elements per thread in a grid-stride fashion. This is already essentially what we're doing if we make each thread compute 4 elements.
Let me create a new version. I'll also try to add __builtin_assume or other hints, and use restrict properly.
Another consideration: I should check if the tensor is contiguous and if not, I might need to handle it. But contiguous should be fine.
Let me also think about using cooperative groups or warp-level primitives. Since each thread independently sums its own elements, and there's no reduction across threads, warp-level primitives don't apply.
So the key optimization is vectorized memory access.
Actually, let me verify my understanding of float4. In CUDA:
float4 vec = reinterpret_cast<const float4*>(ptr)[idx];
This loads 16 bytes (4 floats) in one instruction. This halves the number of load instructions and can significantly improve bandwidth.
But there's a constraint: the pointer must be 16-byte aligned. For float4, we need alignment of 16 bytes.
For our case:
offset = outer_idx * reduce_size * inner_size + r * inner_size + inner_idx
For float4 loads, inner_idx must be a multiple of 4. If we assign 4 consecutive inner elements to each thread (inner_idx = tid * 4), then for each thread, the offset is aligned as long as inner_size itself is aligned properly... wait, that's not the issue. Each thread's own float4 needs to be aligned. So outer_idx * reduce_size * inner_size + r * inner_size + tid * 4 needs to be a multiple of 4, which it is as long as the base is aligned. But float4 requires 16-byte alignment, so the address must be a multiple of 16 bytes = 4 floats.
Base = outer_idx * reduce_size * inner_size. Since outer_idx varies, the alignment depends on whether reduce_size * inner_size is a multiple of 4. In our case, reduce_size=4096, inner_size=4095. 4096 * 4095 = 16773120, which is divisible by 4. But when inner_size is not divisible by 4, the offset for r=1 is base + inner_size, which may not be 16-byte aligned even if base is.
Wait, this is important. The input tensor is contiguous in memory with shape [128, 4096, 4095]. The base pointer is 16-byte aligned (PyTorch ensures this for CUDA tensors).
For outer_idx=0:
r=0, inner_idx=0: offset 0, aligned
r=0, inner_idx=4: offset 4, not 16-byte aligned!
Wait, float4 at offset 4 means address = base + 4 * sizeof(float) = base + 16, which IS 16-byte aligned if base is.
Ah yes, offset 4 in floats = 16 bytes. So inner_idx=0, 4, 8, etc. would all be 16-byte aligned.
But if I assign thread 0 to inner_idx 0, 1, 2, 3 (processing 4 elements but not via float4), that doesn't use float4.
For float4 loading: thread 0 loads inner_idx 0,1,2,3 as a float4. Address = base + 0, which is 16-byte aligned. Good.
Thread 1 loads inner_idx 4,5,6,7. Address = base + 16, also aligned. Good.
Thread k loads inner_idx 4k, 4k+1, 4k+2, 4k+3. Address = base + 16k, aligned. Good.
For r=1:
Thread 0 loads at base + inner_size + 0 = base + 4095.
Is base + 4095 16-byte aligned? base is 16-byte aligned, 4095 * 4 = 16380, which is divisible by 4 but not by 16. 16380 / 16 = 1023.75. So base + 16380 is NOT 16-byte aligned!
This means float4 loads from r=1 would be misaligned, which can cause crashes or performance issues.
To handle this, I would need to use aligned loads for r=0 and something else (like two float2 or four float loads) for other r values. This complicates things.
But actually, we can check if inner_size % 4 == 0. If it is, all float4 loads are aligned. If not, we can fall back to scalar loads, or we can use different strategies.
For the test case, inner_size = 4095, which is NOT divisible by 4. So we can't safely use float4 for all r.
However, we can still optimize by having each thread process 4 consecutive inner elements, loading them with 4 scalar loads which are coalesced across the warp. Or we can use float2 when inner_size % 2 == 0.
Wait, for r=1 with inner_size=4095: offset = base + 4095 * 4 = base + 16380. base is aligned to 16. 16380 mod 16 = 16380 - 16*1023 = 16380 - 16368 = 12. So offset is at alignment 12, which means a float4 starting here would span across a 16-byte boundary. This is bad.
But what if we reorder the computation? Instead of iterating r from 0 to reduce_size-1, what if we load all r values for a group of 4 inner elements and accumulate?
For a group of 4 inner elements starting at inner_idx_g = blockIdx.x * blockDim.x * 4 + threadIdx.x * 4:
We need to read base + r * inner_size + inner_idx_g for all r, which gives us a 4-stride pattern.
Hmm, float4 doesn't help with the r dimension.
Another approach: transpose the problem. If inner_size is 4095 and reduce_size is 4096, we could load a tile of the input into shared memory and then reduce. For example, a block could load a [reduce_size, tile_width] tile into shared memory, where tile_width is small (like 4 or 8). Since threads in a warp load consecutive elements along inner dimension for a fixed r, this is coalesced. Then once in shared memory, each thread can sum along the reduce dimension for its assigned inner element.
But the problem is reduce_size = 4096 is quite large. If we use 256 threads and each loads 16 elements (4 float4s), we'd need 4096 * tile_width elements. If tile_width = 4, that's 16384 floats = 64KB, which fits in shared memory. Then 256 threads can each process one or more inner elements by summing 4096 values.
But actually, this approach doesn't change the fundamental memory access pattern - we're still reading the same data, just buffered through shared memory. Since the access is already coalesced, shared memory might not help much and could hurt due to the extra writes/reads.
Let me think about what else could be done:
1. Use grid-stride loops where each thread computes multiple output positions, increasing compute per thread
2. Use float2 when possible
3. Use warp specialization or other advanced techniques
Actually, let me look at a different optimization: using 128 threads per block instead of 256 to reduce divergence, or increasing occupancy. On H100, 256 threads per block should still give good occupancy.
Let me try another approach: if we can get higher occupancy and ILP (instruction level parallelism), we might see better performance. We can do this by having each thread compute multiple output elements (say N elements spaced by gridDim.x * blockDim.x), which allows better instruction pipelining.
Also, since each output element is computed independently, there's no synchronization needed. This is embarrassingly parallel.
Let me try a version with:
- 128 threads per block (for higher SM occupancy and more blocks)
- Grid-stride: each thread computes multiple output elements along inner dimension with stride gridDim.x * blockDim.x
- Unroll the reduction loop more aggressively
Actually, with batch=128 and inner_size=4095, and reduce_size=4096, the total number of output elements is 128 * 4095 = 524160. If I use 256 threads and a reasonable number of blocks (say 32 blocks_x per outer, giving 128*32=4096 total blocks), that's 4096 * 256 = 1,048,576 threads, which is about 2x the output elements. So average 0.5 output elements per thread.
If I instead use fewer blocks and have grid-stride: say 8 blocks_x per outer, 128*8=1024 blocks, 1024*256=262144 threads for 524160 output elements, average 2 output elements per thread. This gives better ILP.
Let me try this approach with more careful tuning.
Also, I can use float4 for r=0 and scalar for other r... but that's messy.
Let me try a simpler improvement: increase ILP by processing 2 or 4 output elements per thread using grid-stride. Also use 128 threads per block and more blocks_x.
Let me write a version with grid-stride loop and see if performance improves.
Actually, let me also consider the kernel I've already written. It passed correctness and should already be fairly optimized. Let me see if we can do better.
One thing I notice: the reference PyTorch kernel is probably already highly optimized using CUTLASS or similar. To beat it, we might need very specific optimizations.
For sum reduction over a dimension, PyTorch likely uses a custom reduction kernel with:
1. Multiple stages for large reductions
2. Shared memory for inter-thread reduction when the dimension is large
3. Vectorized loads
For dim=1 with shape [128, 4096, 4095]:
- PyTorch might split the reduction into tiles and reduce each tile in shared memory
Given that we can't easily vectorize due to alignment issues, let me try to use shared memory for cooperative reduction.
Actually, another insight: for each fixed (outer, inner) pair, we sum reduce_size elements. This is like a 1D reduction of size 4096. For a 1D reduction, standard techniques use:
1. Load a tile into shared memory
2. Reduce in shared memory using tree reduction
3. Store result
But since each (outer, inner) pair is independent, we can have each block handle multiple pairs.
For example, a block of 256 threads could:
1. Each thread loads 16 elements from the reduce dimension (stride 256 in reduce dim). But in our memory layout, consecutive threads access consecutive inner elements for the same r, not consecutive r values.
Wait, let me reconsider the mapping. If we have a block process ONE (outer, inner) pair:
- 256 threads need to sum 4096 elements
- Each thread sums 4096/256 = 16 elements with stride 256 in r
- But in memory, r varies and inner is fixed. Elements at offset + (0, 256, 512, ...) * inner_size.
Memory accesses: thread 0 reads at offset + 0, offset + 256*inner_size, offset + 512*inner_size, ...
thread 1 reads at offset + 0, offset + 256*inner_size, ...
All threads in a warp read the SAME offset for a given block iteration! This is terrible for memory coalescing.
So PyTorch's reduction is probably mapped differently. They likely use a strided reduction where:
- Thread i in a warp reads elements i, i+32, i+64, ... along the reduce dimension
- For our memory layout, this means thread i reads offset + i*inner_size, offset + (i+32)*inner_size, ...
- Consecutive threads i and i+1 read offset + i*inner_size and offset + (i+1)*inner_size
- These are inner_size apart, NOT coalesced if inner_size > 1!
Wait, on modern GPUs, consecutive threads should read consecutive memory for coalescing. For our layout:
- Thread 0, r=0: address = base + 0*inner_size + inner_idx
- Thread 1, r=0: address = base + 0*inner_size + inner_idx + 1
So for a fixed r, consecutive threads read consecutive floats - this is coalesced. This is exactly what my kernel does. The issue is that consecutive LOAD instructions from the same thread read memory separated by inner_size floats.
In my kernel, thread inner_idx does:
load(base + inner_idx)
load(base + inner_size + inner_idx)
load(base + 2*inner_size + inner_idx)
...
The loads are separated by 4095*4 = 16380 bytes. For an L1 cache line of 128 bytes, each load brings in 32 floats, but the thread only uses 1. So 31/32 of the loaded cache line is wasted for THAT thread. However, threads in the same warp with consecutive inner_idx values (e.g., 0, 1, 2, ..., 31) will have their loads in the same cache line for a given r.
Wait, actually for warp of threads with inner_idx = 0 to 31 and r=0:
Thread 0 reads base + 0
Thread 1 reads base + 1
...
Thread 31 reads base + 31
These are all in the same 128-byte cache line (covers base+0 to base+31). So the cache line is fully utilized across the warp. For r=1:
Thread 0 reads base + 4095
Thread 1 reads base + 4096
...
Thread 31 reads base + 4126
These span addresses base+4095 to base+4126. Are they in the same cache line? 4126 - 4095 = 31 floats = 124 bytes. Base+4095 to base+4126 is 32 floats = 128 bytes. But cache lines are 128 bytes and aligned to 128 bytes. If base is aligned, base+4095*4 = base+16380. Is this in the same cache line as base+16380+124 = base+16504?
16380 / 128 = 127.96875. So 16380 is at byte offset 124 in cache line 127 (0-indexed). Base+16380 to base+16503 spans cache lines 127 and 128. So the warp's accesses for r=1 span TWO cache lines.
Actually, 32 floats = 128 bytes = 1 cache line. But since the starting address is at offset 124, the 32-float span straddles two cache lines. For r=1, half the warp uses cache line 127 and half uses cache line 128. So effectively, each r iteration loads 2 cache lines per warp. But each cache line contains 32 floats, and only 16 or so are used by this warp.
For 32 threads, on r=1:
Thread 0-3 use cache line 127 (bytes 124-139)
Thread 28-31 use cache line 128 (bytes 240-255... wait, let me recalculate)
Actually base + 4095*4 + 0 = base + 16380
base + 4095*4 + 31*4 = base + 16380 + 124 = base + 16504
If base=0 for simplicity:
16380 / 128 = 127 with remainder 124. So cache lines are 127 (bytes 16256-16383) and 128 (bytes 16384-16511).
Bytes 16380-16383 are in line 127 (4 bytes for thread 0)
Bytes 16384-16503 are in line 128 (120 bytes for threads 0-... wait)
16384 corresponds to thread 1 (float at byte 4): 16380 + 4 = 16384, yes thread 1 starts at line 128.
Actually thread i starts at byte 16380 + 4*i.
Thread 0: 16380
Thread 1: 16384 (start of line 128)
...
Thread 31: 16504
So line 127 holds only thread 0's float (4 bytes at 16380-16383).
Line 128 holds threads 1-31 (124 bytes at 16384-16507).
Line 129 holds thread 31's remaining? No, 16504 < 16512, so all of thread 31 is in line 128.
So for r=1, we load cache line 127 (only 4 bytes used) and cache line 128 (124 bytes used). 4/128 = 3.1% utilization for line 127. That's terrible!
For r=0:
base+0 to base+124 are all in line 0 (bytes 0-127). Full utilization.
For r=2:
base + 8190 to base + 8314. 8190/128 = 63 with remainder 126.
Line 63: bytes 8064-8191, holds thread 0's first 2 bytes (8190-8191)
Line 64: bytes 8192-8319, holds threads 0's remaining 2 bytes and threads 1-31.
Again mostly in line 64, only 2 bytes in line 63.
The pattern continues: for most r values, the warp straddles two cache lines, with one cache line being barely utilized.
This means our memory bandwidth utilization is much worse than theoretical. For inner_size=4095, only about 15/16 of the bandwidth is useful (since on average, 1 out of 32 floats in the loaded cache lines is wasted due to misalignment). Wait, actually it's worse because for many r values, we load two lines when one might suffice if aligned.
This is a key insight! The alignment of inner_size affects cache line utilization.
To fix this, we can load elements in a way that better utilizes cache lines. One approach:
- Have each warp process a contiguous chunk of 32 inner elements
- Load 32 floats as a float4 from 8 threads or 32 floats from all 32 threads
- But consecutive r values are offset by inner_size, which causes the misalignment issue
Actually, the misalignment is inherent to the memory layout because inner_size = 4095 is odd. There's not much we can do about the layout without transposing.
However, we can:
1. Use 4 consecutive float loads per thread instead of scattered ones, but that doesn't help since each thread still needs different r values.
Wait, let me reconsider: what if we transpose the tensor in shared memory? A block loads a tile of [reduce_chunk, inner_chunk] into shared memory in a coalesced way, then threads read from shared memory to compute their sums.
For example:
- Inner chunk = 32
- Reduce chunk = 32
- A block with 256 threads loads 32x32 = 1024 floats.
Each thread in the block handles ONE (outer, inner) pair, summing 4096 reduce elements.
We process 4096 reduce elements in chunks of 32. For each chunk:
1. All 256 threads collaboratively load a [32, 32] tile into shared memory
- Actually we need 32 inner elements and 32 reduce elements.
- Thread i loads at some position in the tile.
2. Each thread reads 32 elements from shared memory for its (inner) position
3. Sum these 32 values
4. Move to next reduce chunk
This way, the global memory loads are coalesced (reading tiles), and the shared memory reads are conflict-free (each thread reads a row of the tile).
For a tile of 32x32 floats = 4KB, we can use shared memory effectively.
Total shared memory per block: if we use one tile, it's 4KB. With inner_chunk=64 and reduce_chunk=32, tile = 8KB. This easily fits in the 228KB per SM.
Let me design this more carefully:
For a given outer_idx and a block of threads:
- We need to sum reduce_size=4096 elements for each inner_idx
- Number of inner positions to compute: inner_size=4095
- Threads per block: 256
If we process in tiles of [reduce_chunk=32, inner_chunk=32]:
- Number of tiles along reduce = 4096/32 = 128
- Number of tiles along inner = ceil(4095/32) = 128
- Total tiles = 128 * 128 = 16384 tiles per outer
- Number of output elements per outer = 4095
Hmm, this doesn't map well because we have many tiles but each thread only computes one output.
Let me think differently. A block of 256 threads can compute 256 output elements. If we assign inner_chunk = 256 / something.
Actually, let's say each block computes 256 inner elements for one outer. Then:
- Shared memory approach: for each reduce chunk (say 32 reduce elements):
1. All 256 threads load 32 floats from global memory to shared memory... wait, we need a [reduce_chunk, inner_block] tile in shared memory.
2. If inner_block = 256 and reduce_chunk = 8, tile = 8 * 256 = 2048 floats = 8KB.
3. To load this tile: for each reduce element r in chunk, all 256 threads load their inner value at once. This is coalesced.
4. Once tile is in shared memory, each thread sums the 8 reduce values for its inner position.
So per outer, we process 128 chunks (since 4096/8 = 512... wait, 4096/32 = 128 if reduce_chunk=32).
Let's use reduce_chunk = 32, inner_block = 256:
- Tile size in shared mem = 32 * 256 * 4 bytes = 32KB. This is acceptable.
- Number of chunks along reduce = 4096/32 = 128.
- Number of blocks along inner = ceil(4095/256) = 16.
- Total blocks = 128 outer * 16 inner = 2048 blocks.
For each block:
1. Initialize sum[256] = 0 in registers/locally.
2. For each reduce chunk c from 0 to 127:
a. Load tile of [32, 256] into shared memory from global memory.
- For r within chunk and inner within block, load input[outer][c*32+r][block_start+inner]
- 256 threads, each loads... we need to assign threads to positions.
- Actually we have 256 threads and 32*256 = 8192 elements to load.
- Each thread loads 32 elements. We can have thread i load all elements where (r * 256 + inner) % 256 = i, i.e., thread i loads elements (r, i) for all r. Or we can be smarter.
b. Synchronize.
c. Each thread j reads 32 elements from shared memory at positions (0, j), (1, j), ..., (31, j) and accumulates.
d. Synchronize.
But there's an issue: the loads in step 2a need to be coalesced. If thread j loads (0, j), (1, j), ... these are at global addresses separated by inner_size*4 = 16380 bytes. This is strided and NOT coalesced.
To make loads coalesced, we should have thread j load consecutive memory addresses. In our layout, consecutive addresses correspond to consecutive inner values for the same r. So thread j should load (r, j) for some fixed r.
Wait, for a fixed r, addresses are base + r*inner_size + j for j=0,1,...,255. These are consecutive and coalesced!
But we have 32 r values and 256 inner values. We can assign 256 threads to load the first 256 elements of the first r, then next, etc. But that means each thread loads 32 elements per r... no wait.
Actually, each thread loads one element per instruction. For 32 r values and 256 inner values = 8192 elements. With 256 threads, we need 32 load rounds. In round i, all 256 threads load their element for r=i.
If thread j loads input[outer][c*32+i][block_start+j], then for fixed i, all threads load consecutive addresses: base + (c*32+i)*inner_size + block_start + j. These are consecutive across j, so coalesced!
After loading, we store in shared memory. We should store in shared memory such that thread j's element for r=i is at shared[i][j]. Then in the accumulation phase, thread j reads shared[0][j], shared[1][j], ..., shared[31][j] and sums them.
This shared memory access pattern: thread j reads column j. In shared memory layout shared[32][256], column j means addresses &shared[0][j], &shared[1][j], etc. Consecutive accesses are separated by 256 floats = 1024 bytes. If each warp (32 threads) reads from the same column but different rows, consecutive threads j and j+1 read &shared[i][j] and &shared[i][j+1]. These are 4 bytes apart! Since consecutive threads access consecutive shared memory addresses, there's no bank conflict.
Wait, let me verify: thread j reads shared[i][j], thread j+1 reads shared[i][j+1]. Are these in the same bank? On H100, shared memory has 32 banks. Address bank = (byte_address / 4) % 32.
For shared[i][j]: byte_offset = (i * 256 + j) * 4. Bank = (i * 256 + j) % 32.
For shared[i][j+1]: bank = (i * 256 + j + 1) % 32 = (bank_of_j + 1) % 32.
Since j and j+1 are in different banks, consecutive threads access different banks. For a warp of 32 threads, j=0 to 31, they access banks 0 to 31. No bank conflicts!
Great! So this approach should work well.
But wait, in step 2a, after thread j loads the global memory element, it stores to shared memory. The store should also be conflict-free. Thread j stores to shared[i][j], which is at byte offset (i*256+j)*4. As we just showed, consecutive threads store to different banks. No conflicts!
One thing: in PyTorch, the output shape keeps the reduced dimension as size 1. So output is [128, 1, 4095]. Our blocks compute inner elements and store to output[outer][0][inner].
Let me refine the kernel design:
```cuda
#define REDUCE_CHUNK 32
#define INNER_BLOCK 256
__global__ void sum_reduce_kernel(const float* input, float* output,
int outer_size, int reduce_size, int inner_size) {
int outer_idx = blockIdx.y;
int inner_base = blockIdx.x * INNER_BLOCK;
int tid = threadIdx.x;
float sum = 0.0f;
__shared__ float tile[REDUCE_CHUNK][INNER_BLOCK];
for (int c = 0; c < reduce_size; c += REDUCE_CHUNK) {
// Load tile into shared memory
for (int i = 0; i < REDUCE_CHUNK; ++i) {
int r = c + i;
int inner_idx = inner_base + tid;
if (outer_idx < outer_size && r < reduce_size && inner_idx < inner_size) {
tile[i][tid] = input[outer_idx * reduce_size * inner_size + r * inner_size + inner_idx];
} else {
tile[i][tid] = 0.0f;
}
}
__syncthreads();
// Accumulate from shared memory
#pragma unroll
for (int i = 0; i < REDUCE_CHUNK; ++i) {
int inner_idx = inner_base + tid;
if (inner_idx < inner_size) {
sum += tile[i][tid];
}
}
__syncthreads();
}
int inner_idx = inner_base + tid;
if (outer_idx < outer_size && inner_idx < inner_size) {
output[outer_idx * inner_size + inner_idx] = sum;
}
}
```
Wait, there's a bug in the load loop. With 256 threads and REDUCE_CHUNK=32, the loop `for (int i = 0; i < REDUCE_CHUNK; ++i)` means each thread loads 32 elements. But each load is for a different r, and the loads are strided by inner_size. This defeats the purpose!
What we want is: in each iteration, all 256 threads load for the SAME r but different inner values. For 32 r values, we need 32 iterations where in iteration i, all threads load for r = c + i.
So the inner loop should iterate over r within the chunk. But then each thread loads 32 separate values, one per r. The loads for a single r are coalesced, but they're separated by other loads.
Actually, this can still be okay. In iteration i:
- Thread tid loads address = base + (c+i)*inner_size + inner_base + tid
- This is coalesced for all 256 threads
- We immediately store to shared[i][tid]
This gives us 32 coalesced loads per chunk, one for each r. ILP can help here: after issuing one load, the GPU can work on another warp while waiting.
But we can do even better by loading multiple r values in one instruction using vector loads... but due to alignment, we're limited.
Let me think if we can improve:
- For r=0: loads are aligned as long as inner_base is aligned
- For other r: depends on inner_size alignment
For our case, inner_base = blockIdx.x * INNER_BLOCK * sizeof(float). Since INNER_BLOCK = 256, inner_base is always a multiple of 1024 bytes, so 16-byte aligned.
For r=0: offset = outer_idx * reduce_size * inner_size + inner_base. If this is 16-byte aligned, we can use float4.
But is it? outer_idx * reduce_size * inner_size * 4 + inner_base.
reduce_size * inner_size = 4096 * 4095 = 16773120, which is divisible by 4. inner_base is divisible by 1024. So yes, offset is 16-byte aligned for r=0.
For r=1: offset + inner_size * 4. inner_size = 4095, so inner_size * 4 = 16380, which is not divisible by 16. So not float4 aligned.
So for r=0, we could use float4 loads. For other r, scalar loads. This gives a modest improvement.
Alternatively, we could load 4 consecutive floats using float4 for each r. But that would load inner values that belong to different threads...
Actually, if we use float4, each thread loads 4 consecutive inner values. For r=0:
- Thread tid loads float4 at inner_base + tid*4
- This loads inner values [inner_base + tid*4, inner_base + tid*4 + 3]
- We need to store these in shared memory at [0][tid*4], [0][tid*4+1], [0][tid*4+2], [0][tid*4+3]
But wait, with 256 threads and each loading float4, that's 1024 inner values loaded. Since INNER_BLOCK = 256, we'd only need 64 threads to load the entire block. So we could reduce threads.
Let me reconsider: if we use float4 loads for r=0, and we need INNER_BLOCK = 256 inner values, we can use 64 threads per block, each loading 4 floats. Then for r=1 to 31, we also use 64 threads with scalar loads... actually, with 64 threads, each r needs 64 scalar loads = 64 instructions per r. For 32 r values, that's 2048 instructions. With 256 threads and scalar loads, it's 256 instructions per r * 32 = 8192 instructions but each instruction is for 1/4 the threads, so total throughput is similar.
Actually, using 256 threads with scalar loads or 64 threads with float4 both issue 64*4=256 loads per r, so the same number of memory transactions. But float4 is more efficient because each transaction is 16 bytes instead of 4 bytes, reducing instruction overhead.
Hmm, but with 64 threads, we only have 64 warps, meaning 1 warp per block (on some architectures, 64 threads = 2 warps). With fewer threads, we have fewer warps to hide latency. But we can increase the number of blocks to compensate.
Actually, let's think about occupancy. With 64 threads per block and 32 registers per thread, an SM can hold many blocks. But the parallelism might be limited.
I think for the first optimization attempt, let me use shared memory with tile loading in a coalesced fashion. The key improvement over my first kernel should be:
1. Better cache utilization since we load a tile once and access it multiple times from shared memory
2. Strided accesses happen only for the tile load, and from shared memory they're fast
Wait, actually in my original kernel, strided accesses happen in global memory. With the tile approach, strided accesses happen in shared memory, which is much faster. But we pay the cost of writing to and reading from shared memory.
Let me count the memory operations:
Original: 4096 reads from global memory per output element.
Tile: 4096 reads from global memory per output element (coalesced), plus 4096 reads from shared memory, plus 4096 writes to shared memory.
So the tile approach adds shared memory traffic. It might not be faster unless the improved coalescing significantly helps. In the original kernel, the global memory accesses are already coalesced across the warp... so why would the tile be faster?
In the original kernel, the loads within a warp are coalesced, BUT each load brings a cache line that is partially wasted because of misalignment (as I analyzed earlier). With the tile approach, loads for a single r bring fully utilized cache lines because we load 256 consecutive floats (1024 bytes = 8 cache lines), all of which are used. Then in shared memory, each thread reads its 32 values without any cache line waste.
So the original kernel has poor cache line utilization due to strided-misaligned accesses, while the tile approach has perfect global memory utilization but adds shared memory overhead.
Given that shared memory is much faster than global memory on H100, the tile approach should be faster.
Let me try implementing the tile-based approach.
Actually, I need to be careful about the inner loop. If I have:
```cuda
for (int i = 0; i < REDUCE_CHUNK; ++i) {
int r = c + i;
int inner_idx = inner_base + tid;
tile[i][tid] = input[...];
}
```
This loop has 32 iterations, each with a load. The compiler might not fully unroll this, causing loop overhead. I should explicitly unroll it or use `#pragma unroll`.
Also, note that `tile[i][tid]` access in the store is transposed compared to the load pattern. If we store as `tile[i][tid]` where i is r index and tid is inner index, then in the accumulation phase, thread tid reads `tile[0][tid]`, `tile[1][tid]`, etc. These are separated by INNER_BLOCK floats in shared memory = INNER_BLOCK * 4 bytes.
Since INNER_BLOCK = 256, separation = 1024 bytes. Bank = (1024/4) % 32 = 256 % 32 = 0.
Wait, this is bad! Consecutive accesses by the same thread are in the SAME bank!
Thread tid reads:
- tile[0][tid]: bank = (0 * 256 + tid) % 32 = tid % 32
- tile[1][tid]: bank = (1 * 256 + tid) % 32 = (256 + tid) % 32 = tid % 32
So tile[i][tid] and tile[i+1][tid] are in the same bank! This creates a 32-way bank conflict because all 32 threads in a warp try to read from the same bank simultaneously.
To fix this, we need to pad the shared memory or change the layout.
With padding of 1 float:
```cuda
__shared__ float tile[REDUCE_CHUNK][INNER_BLOCK + 1];
```
Then tile[i][tid] has bank = (i * (INNER_BLOCK + 1) + tid) % 32.
For i=0: bank = tid % 32
For i=1: bank = (257 + tid) % 32 = (1 + tid) % 32
For i=2: bank = (514 + tid) % 32 = (2 + tid) % 32
These are all different banks! With 32 rows, banks cycle through 0-31. Since a warp also accesses 32 different banks, we need to check the specific access pattern.
For warp reading tile[i][tid] where tid varies across 0-31:
Bank = (i * 257 + tid) % 32 = (i + tid) % 32.
Threads access banks tid, tid+1, ..., tid+31 (mod 32). These are all different! No conflict.
Great, adding a padding of 1 float fixes the bank conflicts.
Actually, with 256 threads and INNER_BLOCK=256, a block has 8 warps. Warp k accesses tid = k*32 to k*32+31. For these threads reading tile[i][tid], banks = (i + k*32 + j) % 32 = (i + j) % 32 for j=0 to 31. These are all different, so no conflict within a warp. And multiple warps access the same banks but at different times (they're serialized by the warp scheduler), which is fine.
With padding, the shared memory size becomes 32 * 257 * 4 = 32,896 bytes, which is fine.
Let me write the padded shared memory version.
Also, I realize there's another optimization: instead of having each thread load one element per r, we can have each warp load multiple r values cooperatively. But that might complicate things. Let me start with the simpler version.
Wait, I want the loads to the shared memory to be coalesced and not have too many instructions. Currently each thread executes 32 loads. A total of 256*32 = 8192 loads per block, 2048 blocks, so 16.7 million loads total. But these are all coalesced.
Let me also consider if I can reduce the shared memory reads by accumulating within the load loop. But then I'd be reading from shared memory in a strided way during the store phase... no, the issue is that the store to shared must be coalesced.
Actually, what if we don't use shared memory at all and instead have each thread accumulate directly from global memory, but use vectorized loads for the first r and possibly others?
Given the complexity and the fact that my original kernel might already be competitive (or even faster than PyTorch for this specific shape), let me measure the performance by submitting.
But wait, I should at least try the shared memory approach first since it's a common optimization for reductions and we identified the cache line utilization issue.
Let me implement both approaches and submit the better one. Since we can't benchmark here, I'll need to submit to measure.
Let me try the shared memory tile approach first.
Wait, I also want to try another technique: for this problem, since each output is computed independently and memory accesses are coalesced, maybe using CUDA streams or processing multiple outputs per thread more efficiently would help.
Actually, let me reconsider whether alignment really matters. On H100, L2 cache line is 128 bytes. My original kernel loads:
- r=0: warp loads bytes 0-127 (full cache line, aligned)
- r=1: warp loads bytes 16380-16503. This spans cache line 127 (16256-16383, only uses 4 bytes) and cache line 128 (16384-16511, uses 124 bytes).
- r=2: bytes 32760-32883. Spans cache line 255 and 256.
- ...
For every r, one or two cache lines are loaded. The fully utilized cache lines (128 bytes all used) occur when the 32-float block is fully within one cache line. This happens when inner_idx*4 mod 128 = 0, i.e., inner_idx is a multiple of 32.
For inner_idx not a multiple of 32, the warp's accesses straddle two cache lines, and one of them is underutilized.
What's the average utilization? For a random inner_idx, the probability that inner_idx*4 mod 128 = 0 is 1/32. In all other cases, the accesses span two lines. On average, the underutilized line has about 64 bytes used (uniform distribution over 0-124 bytes from line start).
So average useful bytes per warp per r:
- With probability 1/32: 128 bytes (fully aligned)
- With probability 31/32: on average, (64 + 128) bytes loaded but only 128 bytes used. Wait, that's not right.
Let me be more precise. For a warp at inner_idx:
Start byte = inner_idx * 4.
Start cache line = floor(start_byte / 128).
End byte = start_byte + 124 (31 floats * 4).
End cache line = floor(end_byte / 128).
If start_byte % 128 <= 4 (i.e., 0 or 4), then end_byte = start_byte + 124 <= 128, so fully within one cache line.
Actually, if start_byte % 128 <= 4, then start_byte + 124 <= 128, so yes, one cache line.
If start_byte % 128 > 4, then start_byte + 124 > 128, so spans two lines.
Probability one line: if inner_idx mod 32 <= 1, i.e., inner_idx mod 32 is 0 or 1? Wait:
start_byte = inner_idx * 4
start_byte % 128 = (inner_idx * 4) % 128
For this to be <= 4: inner_idx * 4 ≡ 0 or 4 (mod 128), so inner_idx ≡ 0 or 1 (mod 32).
Ah! So only when inner_idx mod 32 is 0 or 1 do we get a single cache line. That's 2/32 = 1/16 of the cases.
For the other 30/32 cases:
- We load 2 cache lines, but only 128 bytes are useful.
- Total bytes loaded = 256.
- Efficiency = 128/256 = 50%.
Wait, is that right? Let's check: if inner_idx*4 % 128 = 8, then we load bytes 8-131. Cache line 0: bytes 0-127, we use 8-127 = 120 bytes. Cache line 1: bytes 128-255, we use 128-131 = 4 bytes. Total useful = 124 bytes + 4 bytes = 128 bytes. Total loaded = 256 bytes. Efficiency = 50%.
Hmm, but what if the cache lines are already in cache from another warp? On H100, adjacent warps might have overlapping accesses. But generally, each warp's loads are independent.
So average efficiency:
- For 2/32 of warps: 100% efficiency (128 useful / 128 loaded)
- For 30/32 of warps: 50% efficiency (128 useful / 256 loaded)
Average useful bandwidth = (2/32)*100% + (30/32)*50% = 6.25% + 46.875% = 53.125%.
So we're only getting about 53% effective memory bandwidth! That's a huge waste.
With shared memory tiles, we can get nearly 100% efficiency on global memory loads, which should more than compensate for the extra shared memory traffic.
Let me implement the shared memory approach.
Actually wait, I had an error in my analysis. Multiple warps in the same block may have overlapping accesses that hit the same cache lines. For example, warp 0 (inner_idx 0-31) and warp 1 (inner_idx 32-63) don't overlap in cache lines. But warp k and warp k+1 might have some overlap if inner_size causes wrapping.
Actually for fixed outer and r, different warps access different cache lines (since 32 floats = 128 bytes = 1 cache line). There's no overlap. So each warp independently loads its cache lines.
My 53% efficiency estimate is correct.
For shared memory tiles:
- A block loads INNER_BLOCK=256 floats for a given r. That's 8 cache lines, all fully utilized.
- For REDUCE_CHUNK=32 r values, loads 32*256 = 8192 floats = 64 cache lines, all fully utilized.
- Effective global bandwidth = 100%.
- We pay 8192 shared memory reads and 8192 shared memory writes per block.
Shared memory bandwidth on H100 is ~19 TB/s, while HBM is ~3.35 TB/s. So shared memory is about 5.7x faster. The penalty of shared memory traffic is small compared to the gain in HBM efficiency.
Let's compare total data movement:
Original: reads 128 * 4096 * 4095 floats = 2,147,450,880 floats = 8.59 GB.
At 53% efficiency, actual HBM traffic ≈ 16.2 GB.
Tile: reads same 8.59 GB at 100% efficiency + writes 8.59 GB to shared + reads 8.59 GB from shared.
Shared mem traffic = ~17.18 GB at 19 TB/s ≈ 0.9 ms (negligible).
HBM traffic = 8.59 GB at 3.35 TB/s ≈ 2.56 ms.
Original HBM traffic = 16.2 GB at 3.35 TB/s ≈ 4.83 ms.
So tile approach should be about 1.9x faster, ignoring overheads. This is significant!
Let me implement it.
One more optimization: I can have each thread accumulate its values directly during the shared memory load phase, avoiding the separate accumulation loop. But this requires reading from shared memory in a strided pattern during the store, which would cause bank conflicts.
Actually, we can use a different shared memory layout. Instead of storing as [r][inner] and reading as thread j reading column j, we could store as [inner][r] or some other layout.
Wait, what if we accumulate directly in registers without shared memory? We can still get coalesced loads by having warps load full cache lines.
Consider: warp k (32 threads) processes inner values k*32 to k*32+31. For a given r, all 32 threads load coalesced. The issue is that for different r values, the loads are to different cache lines with poor utilization.
What if a warp processes 32 r values for ONE inner element, using all 32 threads to cooperatively sum those 32 r values? Then:
- Thread j loads input[outer][r_base+j][inner_idx]
- These are at addresses separated by inner_size floats = 16380 bytes.
- Not coalesced at all! Each thread accesses a completely different cache line.
So that's worse.
What if a warp processes one inner element but uses all threads to load consecutive r values via shared memory... but then we're back to the tile approach.
OK let me just go with the tile approach.
Here's my refined kernel:
```cuda
#define REDUCE_CHUNK 32
#define INNER_BLOCK 256
__global__ void sum_reduce_kernel(const float* __restrict__ input,
float* __restrict__ output,
int outer_size,
int reduce_size,
int inner_size) {
int outer_idx = blockIdx.y;
int inner_base = blockIdx.x * INNER_BLOCK;
int tid = threadIdx.x;
__shared__ float tile[REDUCE_CHUNK][INNER_BLOCK + 1]; // +1 to avoid bank conflicts
float sum = 0.0f;
for (int c = 0; c < reduce_size; c += REDUCE_CHUNK) {
// Load a tile into shared memory
#pragma unroll
for (int i = 0; i < REDUCE_CHUNK; ++i) {
int r = c + i;
int inner_idx = inner_base + tid;
if (outer_idx < outer_size && r < reduce_size && inner_idx < inner_size) {
int input_idx = outer_idx * reduce_size * inner_size + r * inner_size + inner_idx;
tile[i][tid] = input[input_idx];
} else {
tile[i][tid] = 0.0f;
}
}
__syncthreads();
// Accumulate from shared memory
#pragma unroll
for (int i = 0; i < REDUCE_CHUNK; ++i) {
sum += tile[i][tid];
}
__syncthreads();
}
int inner_idx = inner_base + tid;
if (outer_idx < outer_size && inner_idx < inner_size) {
output[outer_idx * inner_size + inner_idx] = sum;
}
}
```
Wait, we need to also consider the output. The output is [128, 1, 4095]. In contiguous memory, this is 128 * 4095 floats. The output for outer_idx and inner_idx is at output[outer_idx * inner_size + inner_idx].
There's actually a potential issue: the loads within the `#pragma unroll` loop will have divergent branches because `r < reduce_size` and `inner_idx < inner_size` might not be true. But these are outer conditions - inside the loop, r changes but outer_idx and inner_idx are fixed. So for a given thread, the branch outcome is the same for all i within the loop (ignoring r < reduce_size, which transitions at the last chunk).
On the last chunk, some values of i have r >= reduce_size, so some threads take one branch and some take another. But this only happens for the last chunk.
Actually, there's a bigger issue: we're storing to `tile[i][tid]` where `i` varies but `tid` is fixed. In the accumulation phase, `tile[i][tid]` for varying `i` are in different banks because of the padding. Good.
But wait, there's a subtlety: half-warps or warps diverging. On modern GPUs, the `#pragma unroll` loop with 32 iterations should be fine as long as the compiler unrolls it.
Also: do we need `__syncthreads()` between the load and accumulation? Technically yes, because all threads need to finish writing to shared memory before any thread reads from it. Same between accumulation and next load - though we're not writing to shared memory in accumulation, the next load will overwrite shared memory, so we need to ensure all threads have finished reading. __syncthreads() is needed both times.
Performance impact of __syncthreads(): for 128 chunks and 2 syncs per chunk, that's 256 syncs per block. This is a lot. Each sync could cost some cycles.
Can we avoid some syncs? Not easily with this approach.
Alternative: use warp-level shared memory or reduce within a warp before writing to shared memory. But since each output element is computed by one thread and no two threads compute the same output, there's no in-warp reduction possible unless we load extra data.
Actually, what if we process multiple inner elements per thread and accumulate them in registers?
For example, each thread could compute 4 consecutive inner elements:
sum0, sum1, sum2, sum3
For each r, the thread loads 4 floats. In unrolled form:
```cuda
for (int r = 0; r < reduce_size; ++r) {
float val0 = input[base + r * inner_size + inner_idx + 0];
float val1 = input[base + r * inner_size + inner_idx + 1];
float val2 = input[base + r * inner_size + inner_idx + 2];
float val3 = input[base + r * inner_size + inner_idx + 3];
sum0 += val0; sum1 += val1; sum2 += val2; sum3 += val3;
}
```
This gives ILP: 4 loads can be in flight simultaneously, and 4 FMAs can follow.
But the loads are strided by inner_size, so each load is to a different cache line with poor utilization.
However, consecutive threads in a warp load:
Thread 0: base + r*inner_size + 0, 1, 2, 3
Thread 1: base + r*inner_size + 4, 5, 6, 7
For a given r and a given component (say component 0), threads 0, 1, ..., 63 load base+r*inner_size+0, base+r*inner_size+4, ..., base+r*inner_size+252.
These are spaced by 4 floats = 16 bytes. A warp of 32 threads loads at positions 0, 4, 8, ... 124, which are all in the same 128-byte cache line! Wait, these are spaced by 16 bytes. Cache line is 128 bytes. So 32 threads load 32 values at offsets 0, 16, 32, ..., 496 (relative to base+r*inner_size).
496 bytes span 4 cache lines (0-127, 128-255, 256-383, 384-511). Each cache line holds 8 values. So 4 cache lines are loaded, all fully utilized.
This is better than the scalar case! In fact, each warp fully utilizes 4 cache lines per component. But there are 4 components, so 16 cache lines per warp per r. Thread 0 loads from cache lines 0, 0, 0, 0 (for components 0,1,2,3 at the same r). Wait no, for a single r:
Thread 0 loads at offsets 0, 1, 2, 3 (all in cache line 0).
Thread 1 loads at offsets 4, 5, 6, 7 (all in cache line 0).
Thread 7 loads at offsets 28, 29, 30, 31 (all in cache line 0).
Thread 8 loads at offsets 32, 33, 34, 35 (cache line 1).
So the 4 loads by each thread are in the same or adjacent cache lines for nearby threads. More importantly, the loads are coalesced.
For warp with threads 0-31:
Load 0: offsets 0, 4, 8, ..., 124 (spans 128 bytes, fully within cache line 0 for first 32 values... wait, 124+3=127, so yes, all in cache line 0).
Load 1: offsets 1, 5, 9, ..., 125 (all in cache line 0).
Load 2: offsets 2, 6, 10, ..., 126 (all in cache line 0).
Load 3: offsets 3, 7, 11, ..., 127 (all in cache line 0).
So all 4 loads by all 32 threads are in cache line 0! That's 128 floats all in one cache line. But a cache line only holds 32 floats. How can 128 threads * 4 components = 128 loads all be in one cache line?
Oh I see the issue. If each thread handles 4 inner elements, and we have 256 threads, then we're processing 1024 inner elements. But warp 0 only processes inner elements 0-127 (since each of 32 threads handles 4). The first component loads offsets 0, 4, 8, ... 124, which are 32 offsets spaced by 4. These are bytes 0, 16, 32, ..., 496 relative to base. That's 496/128 = 3.875 cache lines. So it spans 4 cache lines.
But wait, I said "threads 0-31, component 0" loads at offsets 0, 4, 8, ..., 124. That's 32 floats at 32 different offsets. These offsets range from 0 to 124, which span 125 floats = 500 bytes. 500 bytes definitely span multiple cache lines.
On H100, memory coalescing works differently. A warp's loads are coalesced into as few transactions as possible. For offsets 0, 4, 8, ..., 124 (all components 0):
- These are 32 floats spaced by 4 floats = 16 bytes apart.
- They span 500 bytes.
- With 128-byte cache lines, we'd need ceil(500/128) = 4 cache lines.
- All 4 cache lines are fully utilized (each holds exactly 8 floats from this warp).
So 128 bytes of useful data are spread across 4 cache lines = 512 bytes loaded. Efficiency = 25%.
Wait, that's worse than before! Earlier with scalar loads, we had about 53% efficiency.
The issue is that with spaced accesses (stride 4), each warp loads 4 cache lines but only uses half the data in each (8 out of 32 floats).
Hmm, so this approach is actually worse for cache line utilization.
Let me reconsider. For scalar loads with inner_idx consecutive (0, 1, 2, ..., 31):
Offsets: 0, 1, 2, ..., 31 (32 floats = 128 bytes)
All in cache line 0. Efficiency for aligned case: 100%.
For unaligned (inner_idx=1): offsets 1, 2, ..., 32. Span 132 bytes = 2 cache lines. Efficiency: ~50% (64 bytes used per line on average, but one line has 124 bytes used and the other 4 bytes).
Average ~53%. Yes.
For the 4-wide approach where each thread loads 4 consecutive elements but the threads are spaced by 4:
Actually, I think I confused myself with the mapping. Let me clarify:
If thread j handles inner_idx = j*4, j*4+1, j*4+2, j*4+3:
Then thread j's loads for component k are at base + r*inner_size + j*4 + k.
For component 0: thread j loads at base + r*inner_size + j*4.
Warp 0 (j=0 to 31): loads at offsets 0, 4, 8, ..., 124. Span = 125 floats = 500 bytes = 4 cache lines.
Utilization per cache line: 8 floats = 32 bytes. Efficiency = 32/128 = 25% per line.
So component 0 has 25% efficiency.
Components 1, 2, 3 reuse the same cache lines (since offsets 1, 5, 9, ... are mostly in cache line 0; offsets 2, 6, 10, ... mostly in cache line 0; offsets 3, 7, 11, ... mostly in cache line 0).
Wait, for component 1: offsets 1, 5, 9, ..., 125. Offset 125 is in cache line 0 (bytes 500-503? No, 125*4=500, cache line 3 starts at 384). Actually:
Cache line 0: bytes 0-127, offsets 0-31
Cache line 1: bytes 128-255, offsets 32-63
Cache line 2: bytes 256-383, offsets 64-95
Cache line 3: bytes 384-511, offsets 96-127
Offset 125 is in cache line 3. So component 1 spans cache lines 0, 1, 2, 3. Same 4 lines.
Component 2: offsets 2, 6, ..., 126. Same 4 lines.
Component 3: offsets 3, 7, ..., 127. Same 4 lines.
So for all 4 components, we load 4 cache lines per warp per r, with 32 floats useful per cache line (8 per component). Total useful = 128 floats = 512 bytes. Total loaded = 512 bytes. Efficiency = 100%!
Wait, that changes things. Let me recalculate:
Component 0: offsets 0, 4, 8, ..., 124. These are 32 offsets in [0, 124].
Cache line 0: offsets 0-31, component 0 uses offsets 0, 4, 8, 12, 16, 20, 24, 28 = 8 floats.
Cache line 1: offsets 32-63, component 0 uses offsets 32, 36, 40, 44, 48, 52, 56, 60 = 8 floats.
Cache line 2: offsets 64-95, component 0 uses 64, 68, 72, 76, 80, 84, 88, 92 = 8 floats.
Cache line 3: offsets 96-127, component 0 uses 96, 100, 104, 108, 112, 116, 120, 124 = 8 floats.
Total useful for component 0: 32 floats = 128 bytes.
Cache lines loaded: 4 lines = 512 bytes.
Efficiency for component 0 alone: 25%.
But components 1, 2, 3 fill in the gaps!
Component 1: offsets 1, 5, 9, ..., 125.
Cache line 0: 1, 5, 9, 13, 17, 21, 25, 29 = 8 floats.
Cache line 1: 33, 37, ..., 61 = 8 floats.
Cache line 2: 65, 69, ..., 93 = 8 floats.
Cache line 3: 97, 101, ..., 125 = 8 floats.
Component 2: offsets 2, 6, ..., 126.
Cache line 0: 2, 6, 10, 14, 18, 22, 26, 30 = 8 floats.
...
Component 3: offsets 3, 7, ..., 127.
Cache line 0: 3, 7, 11, 15, 19, 23, 27, 31 = 8 floats.
Cache line 3: 99, 103, ..., 127 = 8 floats.
Now adding ALL 4 components:
Cache line 0: components 0,1,2,3 use all offsets 0-31 = 32 floats! Full utilization.
Cache line 1: all offsets 32-63 = 32 floats! Full.
Cache line 2: all offsets 64-95 = 32 floats! Full.
Cache line 3: all offsets 96-127 = 32 floats! Full.
Total useful: 128 floats = 512 bytes.
Total loaded: 4 cache lines = 512 bytes.
Efficiency: 100%!
This is amazing! By having each warp load 4 components that together fill an entire 4-cache-line region, we get 100% cache line utilization.
But this is only for alignments where the starting offset is aligned to 4 floats. What about general inner_idx?
For a warp starting at inner_idx = base_inner:
Component 0 loads at base_inner + 0, base_inner + 4, ..., base_inner + 124.
Component 1 loads at base_inner + 1, base_inner + 5, ..., base_inner + 125.
Component 2 loads at base_inner + 2, base_inner + 6, ..., base_inner + 126.
Component 3 loads at base_inner + 3, base_inner + 7, ..., base_inner + 127.
These 128 offsets span [base_inner, base_inner + 127]. This is always exactly 128 consecutive floats = 512 bytes = 4 cache lines! As long as the memory requests for the warp are coalesced together (which they are, since the warp accesses a contiguous 512-byte region), we get full utilization.
So with this 4-component approach, each warp fully utilizes 4 cache lines per r. This is much better than the scalar approach!
The key insight is that although each COMPONENT only uses 25% of each cache line, the 4 components TOGETHER use 100% because they interleave perfectly.
This means the 4-component approach has 100% global memory bandwidth utilization, plus it offers ILP (4 independent loads per thread).
This is better than the shared memory approach because:
1. No shared memory overhead
2. No __syncthreads()
3. Full cache line utilization
4. ILP
Let me design this kernel:
```cuda
__global__ void sum_reduce_kernel(const float* __restrict__ input,
float* __restrict__ output,
int outer_size,
int reduce_size,
int inner_size) {
int outer_idx = blockIdx.y;
int inner_base = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
if (outer_idx >= outer_size) return;
float sum0 = 0.0f, sum1 = 0.0f, sum2 = 0.0f, sum3 = 0.0f;
int64_t base = (int64_t)outer_idx * reduce_size * inner_size + inner_base;
#pragma unroll 4
for (int r = 0; r < reduce_size; ++r) {
const float* row_ptr = input + base + (int64_t)r * inner_size;
float val0 = row_ptr[0];
float val1 = (inner_base + 1 < inner_size) ? row_ptr[1] : 0.0f;
float val2 = (inner_base + 2 < inner_size) ? row_ptr[2] : 0.0f;
float val3 = (inner_base + 3 < inner_size) ? row_ptr[3] : 0.0f;
sum0 += val0;
sum1 += val1;
sum2 += val2;
sum3 += val3;
}
int output_idx = outer_idx * inner_size + inner_base;
if (inner_base < inner_size) output[output_idx] = sum0;
if (inner_base + 1 < inner_size) output[output_idx + 1] = sum1;
if (inner_base + 2 < inner_size) output[output_idx + 2] = sum2;
if (inner_base + 3 < inner_size) output[output_idx + 3] = sum3;
}
```
Wait, each thread processes 4 inner elements. So blockDim.x threads process 4*blockDim.x inner elements.
For inner_size=4095, with 256 threads per block, each block processes 1024 inner elements. Blocks needed = ceil(4095/1024) = 4.
Total blocks = 128 * 4 = 512 blocks. That's a bit low.
We can increase the number of blocks by having each thread process fewer elements but with more blocks, or by using a grid-stride loop.
Option: use 128 threads per block, each processing 4 elements = 512 elements per block. Blocks = ceil(4095/512) = 8 per outer. Total = 128 * 8 = 1024 blocks.
Or better: use 256 threads, each processing 4 elements, but also use grid-stride where each thread processes multiple groups.
Actually, let me compute how many blocks are ideal for H100. H100 has many SMs, maybe 132 or so. We want at least 2-4 blocks per SM for good occupancy. So 1024 blocks is plenty.
Let me go with 128 threads per block, each thread handles 4 elements. That's 1024 blocks total.
But wait, with 128 threads per block and only 32 threads per warp, we have 4 warps per block. The ILP of 4 loads per thread helps hide latency within a warp.
Actually, on H100, we might want more warps per block to hide latency better. Let's use 256 threads per block, which is 8 warps. But with each thread processing 4 elements, that's 1024 elements per block. For 4095 inner elements, we need 4 blocks_x per outer = 512 total blocks. Still decent.
To get more blocks, I can use 128 threads per block, giving 1024 total blocks. Or I can have each thread process only 2 elements with 256 threads = 512 elements per block = 8 blocks_x per outer = 1024 total blocks. Same total blocks but fewer loads per thread.
More loads per thread = better ILP and less thread launch overhead. But fewer blocks could reduce parallelism if the GPU is underutilized.
For batch=128, reduce_size=4096, inner_size=4095:
Total work = 128 * 4095 = 524,160 output elements.
With 4 elements per thread: 131,072 threads needed.
With 256 threads/block and 4 blocks per outer: 131,072 threads total. Perfect.
Actually 128 * 4 = 512 blocks, each with 256 threads = 131,072 threads. Each thread does 4096 * 4 = 16,384 operations but most are loads. 131,072 * 4 = 524,288 output elements, which is the exact amount of work.
Wait, output elements = 128 * 1 * 4095 = 524,160. But 131,072 threads * 4 = 524,288, which is slightly more. Some threads will process elements beyond inner_size.
This is fine.
Actually, let me think about the grid structure more carefully. If I use 256 threads per block and each handles 4 inner elements, then:
blockDim.x = 256
Number of blocks_x = ceil(inner_size / (256 * 4)) = ceil(4095 / 1024) = 4.
For block (x=0, y=outer_idx):
inner_base = (0 * 256 + tid) * 4 = tid * 4.
Threads 0-255 handle inner elements 0, 4, 8, ..., 1020.
Thread 255 handles 1020, 1021, 1022, 1023.
For block (x=1, y=outer_idx):
inner_base = (1 * 256 + tid) * 4 = 1024 + tid * 4.
Thread 0 handles 1024, 1025, 1026, 1027.
Thread 255 handles 2044, 2045, 2046, 2047.
For block (x=2, y=outer_idx):
inner_base = 2048 + tid * 4.
Thread 0 handles 2048-2051.
Thread 255 handles 3068-3071.
For block (x=3, y=outer_idx):
inner_base = 3072 + tid * 4.
Thread 0 handles 3072-3075.
Thread 255 handles 4092-4095. But inner_size=4095, so 4095 is out of bounds (max valid is 4094).
Wait, 3072 + 255*4 = 3072 + 1020 = 4092. So thread 255 handles 4092, 4093, 4094, 4095. Only 4092, 4093, 4094 are valid. 4095 is out of bounds. That's handled by the bounds check.
This all seems correct.
One issue: the loads within a warp for component 0, 1, 2, 3 at a given r. Are the 4 loads issued simultaneously and coalesced together? On H100, memory coalescing happens per warp per instruction. The 4 loads are separate instructions (4 LDs). For each LD, the coalescer looks at the addresses of the 32 threads.
For LD val0: addresses are base + r*inner_size + tid*4 + 0. These form a contiguous 128-float region with stride 4. As I showed, the coalescer will combine these into 4 128-byte transactions, each fully utilized when considering all 4 components together.
But from the hardware's perspective, each LD instruction is processed separately. For LD val0, the transaction includes all 32 threads. The memory subsystem will load 4 cache lines (512 bytes). Then for LD val1, it loads the same 4 cache lines. But since the cache lines were just loaded, they should be in L1/L2 cache! So LD val1, val2, val3 will hit L1 cache and be much faster.
This is crucial. The warp's val0, val1, val2, val3 loads for a given r access the same 512-byte region. After val0 loads the 4 cache lines, val1-3 will cache-hit and be very fast.
So the effective cost is one set of global memory loads per r, plus 3 fast cached loads. This is excellent!
In the scalar version, for a given r, each warp loads either 1 or 2 cache lines. Subsequent r values load different cache lines, so there's no reuse.
Let me now write the kernel. I need to be careful about:
1. Using int64_t for large index calculations to avoid overflow
2. Proper bounds checking
3. Good unrolling
Also, I want to try using float4 to load 4 floats at once when possible. But as we discussed, alignment might be an issue. For the 4-component approach, each thread reads 4 floats at addresses offset + 0, 1, 2, 3. These are consecutive! So we can load them as a float4!
The address is base + r * inner_size + inner_base.
inner_base = (blockIdx.x * blockDim.x + threadIdx.x) * 4.
Is this address 16-byte aligned?
base = outer_idx * reduce_size * inner_size, which may or may not be divisible by 4.
inner_base is divisible by 4 but not necessarily by 16.
For the float4 to load 4 consecutive floats at address `ptr`, `ptr` must be 16-byte aligned.
ptr = input + outer_idx * reduce_size * inner_size + r * inner_size + inner_base
= input + outer_idx * reduce_size * inner_size + r * inner_size + (blockIdx.x * blockDim.x + threadIdx.x) * 4
For r=0 and threadIdx.x=0:
ptr = input + outer_idx * reduce_size * inner_size + blockIdx.x * blockDim.x * 4
If input and outer_idx * reduce_size * inner_size and blockIdx.x * blockDim.x * 4 are such that ptr is 16-byte aligned, then float4 works.
blockIdx.x * blockDim.x * 4 = blockIdx.x * 256 * 4 = blockIdx.x * 1024. This is always divisible by 16.
outer_idx * reduce_size * inner_size = outer_idx * 4096 * 4095.
4096 * 4095 = 16773120. This is divisible by 16? 16773120 / 16 = 1048320. Yes! So outer_idx * 16773120 is always divisible by 16.
And input is 16-byte aligned (PyTorch ensures this for CUDA tensors).
So ptr is 16-byte aligned for all blocks, all threads, at r=0.
For r=1:
ptr = input + outer_idx * 4096 * 4095 + 1 * 4095 + inner_base
= (aligned) + 4095 + inner_base
4095 is NOT divisible by 4 (4095 = 4096 - 1). So ptr at r=1 is at offset 4095 + inner_base.
If inner_base is divisible by 4, then 4095 + inner_base ≡ 3 (mod 4), which is NOT divisible by 4, let alone 16.
So float4 at r=1 would be misaligned. We can't use float4 for arbitrary r unless inner_size % 4 == 0.
Since inner_size = 4095, we can't use float4 for r >= 1.
However, we CAN use float4 for r=0, and scalar loads for r >= 1. This gives a small boost for 1/4096 of the loads.
Alternatively, we can handle the r loop with r running from 0 to reduce_size - 1, and for each r, use scalar loads. The compiler might optimize this well.
Actually, for the 4-component scalar approach, the compiler might not be able to optimize 4 separate loads into one vector load even when aligned.
Let me just write the scalar version first. With the caching behavior I described (4 components sharing cache lines), the scalar version should perform very well.
But wait, there's an important point: will the compiler issue the 4 loads in close proximity so they can share cache lines? If the compiler serializes them too much, val1-3 might evict val0's cache lines before using them.
Since 512 bytes is small (fits in L1 cache), and H100 L1 is large, this shouldn't be a problem. But the compiler's instruction scheduling matters.
To help the compiler, I can write:
```cuda
float val0, val1, val2, val3;
#pragma unroll
for (int r = 0; r < reduce_size; ++r) {
const float* row_ptr = input + base + (int64_t)r * inner_size;
val0 = row_ptr[0];
val1 = row_ptr[1];
val2 = row_ptr[2];
val3 = row_ptr[3];
sum0 += val0;
sum1 += val1;
sum2 += val2;
sum3 += val3;
}
```
This should allow the 4 loads to be issued close together.
Actually, I want to let the compiler schedule loads and adds freely. I might reorder:
```cuda
#pragma unroll 4
for (int r = 0; r < reduce_size; ++r) {
const float* row_ptr = input + base + (int64_t)r * inner_size;
sum0 += row_ptr[0];
sum1 += row_ptr[1];
sum2 += row_ptr[2];
sum3 += row_ptr[3];
}
```
The compiler should optimize this.
Now, one more optimization: I can process multiple r values in the loop body to increase ILP further. For instance:
```cuda
#pragma unroll 2
for (int r = 0; r < reduce_size; r += 2) {
const float* row_ptr0 = input + base + (int64_t)r * inner_size;
const float* row_ptr1 = input + base + (int64_t)(r+1) * inner_size;
sum0 += row_ptr0[0] + row_ptr1[0];
sum1 += row_ptr0[1] + row_ptr1[1];
sum2 += row_ptr0[2] + row_ptr1[2];
sum3 += row_ptr0[3] + row_ptr1[3];
}
```
This doubles the ILP (8 loads in flight). Similarly, unroll by 4:
```cuda
#pragma unroll 4
for (int r = 0; r < reduce_size; r += 4) {
const float* row_ptr0 = input + base + (int64_t)(r+0) * inner_size;
const float* row_ptr1 = input + base + (int64_t)(r+1) * inner_size;
const float* row_ptr2 = input + base + (int64_t)(r+2) * inner_size;
const float* row_ptr3 = input + base + (int64_t)(r+3) * inner_size;
sum0 += row_ptr0[0] + row_ptr1[0] + row_ptr2[0] + row_ptr3[0];
...
}
```
This gives 16 loads in flight. On H100, each SM can have many loads in flight, so this should help hide latency.
But I need to be careful about register pressure. With 16 pointers + 4 sums + loop variables, we might exceed the register limit. Let me be careful.
Actually, `row_ptr0`, etc. might not stay as registers - the compiler will optimize them to address calculations. Let me check by compiling.
Also, add bounds check for `reduce_size` not being a multiple of 4.
Let me also consider using `__builtin_assume(reduce_size % 4 == 0)` if I know it's true, but for generic correctness I shouldn't.
For the test case, reduce_size=4096, which IS divisible by 4. But the kernel should work for any shape.
Actually, the unroll factor and ILP can be tuned. Let me try unroll by 4 on reduce_size since 4096/4=1024, which is a nice number. But CUDA compilation with unroll might need the loop bound to be known at compile time, which it isn't.
`#pragma unroll 4` should work even with dynamic bounds - it just unrolls 4 iterations and handles the remainder.
Let me write the kernel and compile it. I'll also make sure it works for the exact test dimensions.
Actually, I'll write it with a moderate unroll of 4 and see.
Wait, there's another consideration: register spilling. With 256 threads per block and many registers per thread, occupancy might drop. H100 has 64K registers per SM. If each thread uses 64 registers, then 256 threads use 16K registers. Up to 4 blocks per SM = 64K, which is the limit. If we use 128 registers per thread, 256 threads = 32K, so only 2 blocks per SM.
To maximize occupancy, let's use 128 threads per block. Then even with 128 registers, we get 128*128 = 16K per block, so 4 blocks per SM.
Or 256 threads with ~64 registers gives 4 blocks per SM.
Let me aim for 64 or fewer registers per thread. I'll monitor during compilation.
Actually, I don't get register counts from compile_kernel. Only static_check gives some diagnostics but not register counts. I'll need to submit to see performance.
For now, let me write a clean kernel and submit.
Let me also check: should I use `int` or `int64_t` for the indices?
For outer_size up to very large values, int64_t is safer. But `blockIdx` and `threadIdx` are `int`. For this problem, outer_size=128, reduce_size=4096, inner_size=4095, all fit in int. I'll use int for performance and only cast when needed for multiplication.
Actually, `outer_idx * reduce_size * inner_size` can overflow int if these are large. Let me use int64_t for the base offset.
Let me compile and test the 4-component kernel. I'll also handle the case where inner_size is not divisible by 4.
Actually, to make the code cleaner, I can write:
```cuda
__global__ void sum_reduce_kernel(const float* __restrict__ input,
float* __restrict__ output,
int outer_size,
int reduce_size,
int inner_size) {
int outer_idx = blockIdx.y;
int tid = threadIdx.x;
int block_inner = blockIdx.x * blockDim.x * 4;
int inner_idx = block_inner + tid * 4;
if (outer_idx >= outer_size) return;
float4 sums = {0.0f, 0.0f, 0.0f, 0.0f};
int64_t base = (int64_t)outer_idx * reduce_size * inner_size + inner_idx;
#pragma unroll 4
for (int r = 0; r < reduce_size; ++r) {
const float* row_ptr = input + base + (int64_t)r * inner_size;
sums.x += (inner_idx + 0 < inner_size) ? row_ptr[0] : 0.0f;
sums.y += (inner_idx + 1 < inner_size) ? row_ptr[1] : 0.0f;
sums.z += (inner_idx + 2 < inner_size) ? row_ptr[2] : 0.0f;
sums.w += (inner_idx + 3 < inner_size) ? row_ptr[3] : 0.0f;
}
int64_t out_base = (int64_t)outer_idx * inner_size + inner_idx;
if (inner_idx + 0 < inner_size) output[out_base + 0] = sums.x;
if (inner_idx + 1 < inner_size) output[out_base + 1] = sums.y;
if (inner_idx + 2 < inner_size) output[out_base + 2] = sums.z;
if (inner_idx + 3 < inner_size) output[out_base + 3] = sums.w;
}
```
Wait, the bounds check inside the loop might hurt performance. For the test case, inner_idx ranges from 0 to 4091 for the last block, and inner_size=4095. So for the last block, some threads have inner_idx + 3 = 4094, which is < 4095. Only a few threads in the last block might exceed.
Specifically:
Block (3, y): block_inner = 3 * 256 * 4 = 3072.
Thread 255: inner_idx = 3072 + 255*4 = 4092.
4092 + 3 = 4095, which is NOT < 4095. So thread 255's 4th element is out of bounds.
For block (3, y), only the last thread has one out-of-bound element. All other threads in all blocks are fully in bounds because:
Block (3, y): inner_idx for thread 0 = 3072 < 4095, thread 0-254 all have inner_idx+3 <= 4091+3 = 4094 < 4095.
Thread 255: 4092 < 4095 (valid), 4093 < 4095 (valid), 4094 < 4095 (valid), 4095 is invalid.
So very few threads have divergent bounds. We can optimize by moving the bounds check outside the loop:
```cuda
bool valid0 = inner_idx + 0 < inner_size;
bool valid1 = inner_idx + 1 < inner_size;
bool valid2 = inner_idx + 2 < inner_size;
bool valid3 = inner_idx + 3 < inner_size;
for (int r = 0; r < reduce_size; ++r) {
const float* row_ptr = input + base + (int64_t)r * inner_size;
if (valid0) sums.x += row_ptr[0];
if (valid1) sums.y += row_ptr[1];
if (valid2) sums.z += row_ptr[2];
if (valid3) sums.w += row_ptr[3];
}
```
This is slightly better as the valid flags are computed once.
But actually, for performance, it might be even better to just not worry about the few threads at the edge. Most threads process all 4 elements without overhead. The compiler can optimize this.
Let me check: for thread 255 in block 3, valid3 is false. The warp has 31 threads with valid3=true and 1 thread with valid3=false. On CUDA, branches within a warp are serialized: first execute the true threads, then the false threads. So there's a small penalty.
To avoid this, I can pad the input so that inner_size is always a multiple of 4, but I can't modify the input. Alternatively, I can ensure we never launch threads that go out of bounds by computing the grid size carefully.
But with `ceil(inner_size / 1024) = 4` blocks, we inevitably process 4096 elements but only need 4095. The simplest solution is to just handle the bounds.
Actually, we can handle it by not using a 4-wide vector at the boundary:
```cuda
// Main loop with 4-wide processing
int main_blocks = (inner_size / 4) * 4; // Largest multiple of 4 <= inner_size
// For elements before main_blocks
...
// For elements in main_blocks, process 4 at a time
...
// For remaining elements after main_blocks
...
```
But this complicates the grid. Let me just keep the simple version with bounds checks inside the loop. For most threads, all bounds are true and there's no divergence.
Actually wait, for the last block, about 1/4 of the threads might have some false bounds. Let me think:
Block 3 covers inner indices 3072-4095. inner_size=4095.
Thread j covers 3072 + j*4 to 3072 + j*4 + 3.
For this to all be < 4095: 3072 + j*4 + 3 < 4095, so j*4 < 1020, j < 255.
So for j=0 to 254, ALL 4 elements are valid.
For j=255, elements 4092, 4093, 4094 are valid, 4095 is invalid.
Only 1 thread out of 256 has any divergence, and only 1 out of 4 loads diverges. The cost is negligible.
I can even simplify to just handle 3 valid elements for thread 255:
```cuda
if (inner_idx >= inner_size) return; // Early exit for completely out-of-bounds threads
float sum0 = 0.0f, sum1 = 0.0f, sum2 = 0.0f, sum3 = 0.0f;
for (int r = 0; r < reduce_size; ++r) {
const float* row_ptr = input + base + (int64_t)r * inner_size;
sum0 += row_ptr[0];
if (inner_idx + 1 < inner_size) sum1 += row_ptr[1];
if (inner_idx + 2 < inner_size) sum2 += row_ptr[2];
if (inner_idx + 3 < inner_size) sum3 += row_ptr[3];
}
output[out_base] = sum0;
if (inner_idx + 1 < inner_size) output[out_base + 1] = sum1;
if (inner_idx + 2 < inner_size) output[out_base + 2] = sum2;
if (inner_idx + 3 < inner_size) output[out_base + 3] = sum3;
```
The `if (inner_idx >= inner_size) return;` handles threads where inner_idx is way out of bounds. In our grid, there shouldn't be any such threads, but it's good for safety.
Actually, for block 3, inner_idx ranges from 3072 to 4092. All are < 4095 = inner_size. So no thread is completely out of bounds.
But if inner_size were 3072, then block 3 would have inner_idx >= inner_size for all threads. The early exit handles this.
OK, let me finalize the kernel and compile. I'll also try the shared memory version and compare. Since we can't benchmark without submitting, I'll test correctness first, then submit to measure performance.
Let me also think about whether I should issue more ILP. Currently each thread accumulates 4 sums. I can also accumulate across multiple r values:
```cuda
float sum0_a = 0.0f, sum1_a = 0.0f, sum2_a = 0.0f, sum3_a = 0.0f;
float sum0_b = 0.0f, sum1_b = 0.0f, sum2_b = 0.0f, sum3_b = 0.0f;
#pragma unroll 2
for (int r = 0; r < reduce_size; r += 2) {
sum0_a += ...;
sum1_a += ...;
sum0_b += ...;
sum1_b += ...;
}
sum0 = sum0_a + sum0_b;
```
This gives even more ILP. But with 8 sums + more pointers, register pressure might increase too much. Let me keep it simpler first.
Actually, let me try both approaches: the simple 4-element version and the shared memory version. I'll start with the 4-element version since it's simpler and should have good performance.
Let me write the full code and compile.