reasoning
Let me analyze the problem:
1. We need to compute pairwise Euclidean distance matrix for (B, N, 3) coordinates
2. Reference: `diff = coords.unsqueeze(2) - coords.unsqueeze(1)` then sqrt of sum of squares + 1e-8
3. Output: (B, N, N) tensor
This is essentially computing:
`dist[b, i, j] = sqrt(sum_k (coords[b, i, k] - coords[b, j, k])^2 + 1e-8)`
This is a classic GEMM-like operation that can be optimized:
`||x_i - x_j||^2 = ||x_i||^2 + ||x_j||^2 - 2*x_i·x_j`
So we can compute:
1. `norms = (coords ** 2).sum(dim=-1)` -> (B, N)
2. `dist_sq = norms.unsqueeze(2) + norms.unsqueeze(1) - 2 * coords @ coords.transpose(-1, -2)`
3. `dist = sqrt(dist_sq + 1e-8)`
With N=512, B=8, a custom kernel can be very efficient. Let's think about the best CUDA kernel approach:
A direct element-wise kernel computing (B, N, N) distances could use shared memory for tiles of coords. This avoids forming the large (B, N, N, 3) intermediate tensor.
For each batch b, we need to compute N×N distances from 3D coordinates. That's 8 * 512 * 512 = ~2M outputs.
For a fused kernel, we can tile over the N dimension. Each thread block computes a tile of distances. We load tiles of coords into shared memory.
But actually, we can also use the GEMM formulation with cuBLAS for the matrix multiplication. Let me think...
`dist_sq = ||x_i||^2 + ||x_j||^2 - 2 * X @ X^T`
This is exactly how standard pairwise distance is computed efficiently. The `-2 * X @ X^T` is a batch GEMM operation.
For N=512, a batched GEMM is extremely fast on H100. But wait, we need to also add the norms and take sqrt. We could:
1. Compute norms with a custom kernel or PyTorch
2. Do a batch GEMM for X @ X^T
3. Do the rest with a fused kernel
But we need everything in a custom kernel. Let me think about writing a single kernel that does all of this.
Actually, using a single custom CUDA kernel that loads data cooperatively into shared memory and computes pairwise distances is likely very efficient for this problem size. The computation is memory bandwidth bound.
Let me design a kernel:
- Grid: (B, N/TileX, N/TileY) where each block computes a tile of the output
- Actually, since the output is symmetric (distances from i to j = j to i), but computing both is simpler and the reference computes both.
For N=512, if we use 32x32 tiles:
- Blocks per batch: (512/32) * (512/32) = 16 * 16 = 256
- Total blocks: 8 * 256 = 2048
Each block loads a tile of coords_i (TileX, 3) and coords_j (TileY, 3) into shared memory, then each thread computes the squared distance between one pair (i, j).
Wait, but each thread in a block needs different i and j. Let's use:
- BlockDim: (32, 32, 1) - 1024 threads per block
- Each thread computes one output element (i, j)
- We can tile over the 3 dimension since it's small
Actually, with 3 dimensions, we don't even need to tile over K. We can just load all 3 values per point.
Let me design it more carefully:
- Tile size: 32x32
- Block size: 32x32 threads
- Each thread block processes one tile of the (N, N) output matrix
- For computing dist[i, j], we need coords[i, :] and coords[j, :] (3 floats each)
- We can either load from global memory directly, or use shared memory for the tile
Since 3 is very small, we might as well just load coords for the tile into shared memory:
- Shared memory for block: 32 rows * 3 floats for i-tile, 32 rows * 3 floats for j-tile
- Total: 32 * 3 + 32 * 3 = 192 floats ≈ 768 bytes - very small!
Actually wait, each element in the tile is computed by one thread, and each thread needs 3 values from row i and 3 from row j. If we put coords for the whole tile in shared memory, all threads can access them.
But actually, a simpler approach: each thread just loads the 3 coords for its row i and column j directly from global memory. Since there are 32 threads in each dimension, this results in coalesced reads (each warp reads 32 consecutive rows for the same coordinate value if we access properly).
However, with H100 and L2 cache, direct load might already be well-cached since nearby thread blocks access overlapping data.
Let me think about the best shared memory tiling strategy:
Option 1: 2D tiling with shared memory
- Block (Bx, By) = (16, 16), each thread computes one distance
- Block computes a 16x16 tile
- Shared memory: load 16 points for i-tile and 16 points for j-tile
- For K=3, this is negligible shared memory
- But a 16x16 block uses 256 threads. With 24 warps per SM (2048 threads / 32), and 512 threads per block, we get 4 blocks per SM.
Option 2: Larger tiles
- Block (32, 32) = 1024 threads
- This uses most of the thread budget per SM (2048 on H100), so 2 blocks per SM
For this problem where computation is very light (just 3 subtractions, 3 multiplications, 2 additions, 1 sqrt), we want high occupancy and good memory coalescing.
Actually, since the computation is memory-bound, what matters most is:
1. Minimizing global memory reads
2. Maximizing coalescing
3. Good occupancy
Each output element requires reading 6 floats (3 from row i, 3 from row j). The output is N*N = 512*512 = 262K floats per batch. Total reads: 8 * 262K * 6 = ~12.5M floats = ~50MB.
Memory bandwidth is 3.35 TB/s, so this should take about 15 microseconds in theory. But with overhead and the fact that some reads are reused, shared memory can help.
Actually, the reference PyTorch implementation:
1. `coords.unsqueeze(2) - coords.unsqueeze(1)` creates a (B, N, N, 3) tensor of 8*512*512*3 = ~6M elements = 24MB
2. Then squares and sums, then sqrt
This creates large intermediates and does multiple passes over memory.
A custom fused kernel that goes from coords directly to dist without intermediates would be much more memory efficient.
Let me write a kernel using shared memory tiling:
```cuda
template <int BLOCK_SIZE>
__global__ void pairwise_distance_kernel(const float* coords, float* dist, int B, int N) {
int b = blockIdx.z;
int i0 = blockIdx.y * BLOCK_SIZE;
int j0 = blockIdx.x * BLOCK_SIZE;
int ti = threadIdx.y;
int tj = threadIdx.x;
int i = i0 + ti;
int j = j0 + tj;
__shared__ float coords_i[BLOCK_SIZE][3];
__shared__ float coords_j[BLOCK_SIZE][3];
// Load coords for this tile into shared memory
if (i < N && tj < 3) {
coords_i[ti][tj] = coords[b * N * 3 + i * 3 + tj];
}
if (j < N && ti < 3) {
coords_j[tj][ti] = coords[b * N * 3 + j * 3 + ti];
}
// Wait, I'm mixing up the indexing. Let me be more careful.
```
Wait, the problem with having both ti and tj participate in loading is that we need each thread to potentially load two things. Let me use a different approach:
Since K=3 is so small, instead of having threads cooperatively load, we can just have each thread read the coords it needs directly. With the right thread layout, reads will be coalesced.
Let's use:
- Block dimension: (32, 32)
- grid: (ceil(N/32), ceil(N/32), B)
- Thread (tx, ty) computes output at (i = by*32 + ty, j = bx*32 + tx)
For coalescing, adjacent threads in x should read adjacent memory. The coords are stored as (B, N, 3) so coords[b, n, :] are contiguous.
When thread (tx, ty) reads coords[b, i, :], threads with same ty but different tx read the same i, so this is a broadcast (not coalesced but cached).
When thread (tx, ty) reads coords[b, j, :], threads with same ty and adjacent tx read adjacent j values, so this is coalesced!
Actually, the memory for coords[b, n, k] where n varies and k fixed is NOT contiguous because coords[b, n, :] are contiguous. So coords[b, n, 0] and coords[b, n+1, 0] are 3 floats apart, not 1. But that's still fine for coalescing since adjacent threads access memory with stride 3, which is still coalesced on modern GPUs (L1 cache handles this, or if not L1, the memory controller handles it via sector cache lines).
Actually on modern GPUs, coalescing works for any stride as long as threads in a warp access addresses in the same 128-byte cache line. With 32 threads accessing floats at stride 3, the addresses are base, base+12, base+24, ... The total span is 31*12 = 372 bytes, which spans about 3 cache lines. This is not perfectly coalesced but still reasonably efficient.
Wait, let me reconsider. If we use block (32, 32), then each warp (say warp 0 with threads (0-31, 0)) all have ty=0, so they compute j = bx*32 + tx with tx from 0 to 31, and fixed i = by*32.
For reading coords[b, i, :], all threads in the warp read the same 3 floats. This is broadcast, handled well by L1 cache.
For reading coords[b, j, :], threads read j = bx*32 + 0 to bx*32 + 31, each reading 3 contiguous floats. This has a stride of 3 floats = 12 bytes between threads. With 32 threads, the addresses are 0, 12, 24, ..., 372 bytes apart. This spans about 3-4 cache lines (128 bytes each). So roughly 3-4 sector requests per warp. This is acceptable.
For better coalescing, we could transpose the coordinate data so coords are (B, 3, N). Then reading coords[b, :, j] with j varying would be perfectly coalesced. But we can't change the input format.
Alternatively, we can transpose in shared memory.
Actually, for simplicity and because N=512 is small, let me try a simple approach first, then optimize.
Actually, I just realized: for this specific problem, we can also compute it via the GEMM identity. Let me think about whether that's faster.
Identity:
`||x_i - x_j||^2 = ||x_i||^2 + ||x_j||^2 - 2 x_i^T x_j`
Steps:
1. Compute `norms = sum(coords^2, dim=2)` -> (B, N)
2. Compute `dots = coords @ coords.transpose(1, 2)` -> (B, N, N) via batch matmul
3. Compute `dist_sq = norms.unsqueeze(2) + norms.unsqueeze(1) - 2 * dots`
4. `dist = sqrt(dist_sq + 1e-8)`
Step 2 is a batch GEMM of (B, N, 3) @ (B, 3, N) -> (B, N, N).
For N=512 and B=8, the effective matrix sizes are:
- M = 512, K = 3, N = 512 per batch
- Total: 8 * 512 * 3 * 512 ≈ 6.3M FLOPs
This is very small for a GEMM. The overhead of calling cuBLAS might be significant. A custom kernel might be faster for this tiny GEMM.
But actually, the batch size is 8, and the problem is memory bound anyway. I think a direct custom kernel is the way to go.
Let me write a custom kernel with shared memory tiling for better performance.
```cuda
#define BLOCK 32
__global__ void pairwise_distance_kernel(
const float* __restrict__ coords,
float* __restrict__ dist,
int B, int N) {
int b = blockIdx.z;
int i0 = blockIdx.y * BLOCK;
int j0 = blockIdx.x * BLOCK;
int ti = threadIdx.y;
int tj = threadIdx.x;
int i = i0 + ti;
int j = j0 + tj;
__shared__ float smem_i[BLOCK][3];
__shared__ float smem_j[BLOCK][3];
// Each thread loads the coords for its row i into shared memory
// We can have threads with tj < 3 do the loading for smem_i
if (tj < 3 && i < N) {
smem_i[ti][tj] = coords[(b * N + i) * 3 + tj];
}
// Each thread loads coords for its column j into shared memory
// We can have threads with ti < 3 do the loading for smem_j
// But wait, we need all threads to participate efficiently
if (ti < 3 && j < N) {
smem_j[tj][ti] = coords[(b * N + j) * 3 + ti];
}
__syncthreads();
if (i < N && j < N) {
float dx = smem_i[ti][0] - smem_j[tj][0];
float dy = smem_i[ti][1] - smem_j[tj][1];
float dz = smem_i[ti][2] - smem_j[tj][2];
float d2 = dx*dx + dy*dy + dz*dz + 1e-8f;
dist[(b * N + i) * N + j] = sqrtf(d2);
}
}
```
Wait, there's an issue with the shared memory loading. If ti ranges from 0 to 31 and tj ranges from 0 to 31, then:
- For smem_i: threads with tj < 3 load coords[b, i, 0], coords[b, i, 1], coords[b, i, 2]
- But all threads with same ti but different tj where tj < 3 all try to load different elements for the SAME row i.
- 3 threads per row load 3 values. That's 32*3 = 96 loads, but they span 32 different rows. Threads (ti, 0), (ti, 1), (ti, 2) load row ti. This is fine.
- Coalescing: for fixed tj, threads with ti varying from 0 to 31 load coords[b, i0+ti, tj]. These are contiguous in memory! Perfect coalescing.
- For smem_j: threads with ti < 3 and tj varying load coords[b, j0+tj, ti].
- For fixed ti, threads with tj varying from 0 to 31 load coords[b, j0+tj, ti]. These are strided by 3 floats. Not perfectly coalesced but okay.
Actually, I realize we can be smarter about this. Since the j loading isn't perfectly coalesced, what if we reorganize? We can have threads load j values in a way that threads in a warp load contiguous memory.
For the j data loading: if we have a warp where ti varies, each thread in the warp loads a different tj. Actually no, ti and tj are both in the thread index. Let me think about warps.
A warp of 32 threads could be (ti=0, tj=0..31), or (ti=0..3, tj=0..7), etc., depending on how we map threadIdx to (ti, tj).
CUDA maps threadIdx.x to the fastest varying dimension in warps. So threads with consecutive threadIdx.x values are in the same warp.
If we use `int ti = threadIdx.y; int tj = threadIdx.x;`:
- threadIdx ranges from 0 to 1023 (32x32)
- threadIdx.x = tj, threadIdx.y = ti
- Warps contain threads with consecutive tj and fixed ti (or nearby ti since 32 tj values = 1 warp)
Actually, for a 2D block of (32, 32):
- Threads with threadIdx in [0, 31] have ti=0, tj=0..31
- Threads with threadIdx in [32, 63] have ti=1, tj=0..31
- etc.
So warp 0: ti=0, tj=0..31
Warp 1: ti=1, tj=0..31
...
For loading smem_i:
- In warp 0 (ti=0), threads with tj=0,1,2 load coords[b, i0+0, 0..2]. Threads with tj>=3 do nothing.
- The active threads (tj=0,1,2) load contiguous memory: coords[b, i0, 0], coords[b, i0, 1], coords[b, i0, 2]. This is a single thread loading 3 contiguous values? No, 3 different threads each load one value. These are contiguous, so it's coalesced. Other threads in the warp are inactive.
For loading smem_j:
- In warp 0 (ti=0), threads with tj=0..31 load coords[b, j0+tj, 0]. These are strided by 3. Not perfectly coalesced.
- Actually wait, ti=0 is the condition for loading smem_j, and all 32 threads in warp 0 have ti=0. So all 32 threads load coords[b, j0+tj, 0] for tj=0..31. These have stride 3 floats = 12 bytes. This spans multiple cache lines but is still reasonably efficient.
For better loading of j coords, we could have 3 warps load the j coords. Actually, what if we have each thread read 3 floats directly into registers instead of using shared memory?
Simpler approach without shared memory:
```cuda
__global__ void pairwise_distance_kernel(
const float* __restrict__ coords,
float* __restrict__ dist,
int B, int N) {
int b = blockIdx.z;
int i = blockIdx.y * BLOCK + threadIdx.y;
int j = blockIdx.x * BLOCK + threadIdx.x;
if (i < N && j < N) {
float xi0 = coords[(b * N + i) * 3 + 0];
float xi1 = coords[(b * N + i) * 3 + 1];
float xi2 = coords[(b * N + i) * 3 + 2];
float xj0 = coords[(b * N + j) * 3 + 0];
float xj1 = coords[(b * N + j) * 3 + 1];
float xj2 = coords[(b * N + j) * 3 + 2];
float dx = xi0 - xj0;
float dy = xi1 - xj1;
float dz = xi2 - xj2;
float d2 = dx*dx + dy*dy + dz*dz + 1e-8f;
dist[(b * N + i) * N + j] = sqrtf(d2);
}
}
```
This is much simpler. Would it be fast enough?
Memory access pattern:
- For i coords: all threads in a warp with same threadIdx.y read the same i coords (broadcast, cached well by L1).
- For j coords: threads with consecutive threadIdx.x read consecutive j values with stride 3. As discussed, not perfectly coalesced but okay.
Actually, since L2 cache is large and N=512 is small, the coords data for one batch is 512*3*4 = 6144 bytes = 6KB. This easily fits in L1 cache! So broadcast reads and strided reads will mostly hit cache.
For B=8, total coords is 48KB. Also tiny.
The main benefit of shared memory would be to improve the reuse within a block. Each block of 32x32 reads 32 rows and 32 columns. The rows are reused 32 times each (one row used by 32 threads), and columns also reused 32 times each. With L1 cache, this reuse might already happen naturally.
Let me try the direct approach first - it's simpler and might be fast enough since the working set is tiny.
But wait, can I do better with vectorized loads and larger tiles?
Another idea: since N=512 and the output is symmetric (dist[i,j] = dist[j,i]), we only need to compute the upper or lower triangle. But actually, computing the full matrix with a simple kernel is already fast, and the reference computes the full matrix.
Actually, thinking more about the symmetric property: if we only compute the upper triangle plus diagonal, we save half the computation and writes. But N=512, so 512*512/2 = 131K elements per batch. For B=8, that's about 1M elements. Each element requires 6 float reads and 1 float write. Still very small.
But implementing the symmetric optimization with proper coalescing is tricky. Let me stick to the full matrix approach for simplicity.
What about vectorized loads? Instead of loading 1 float at a time, we can load float3 or float4. But float3 is 12 bytes, which is not power-of-2 aligned unless we pad. The PyTorch tensor is (B, N, 3), so coords[b, n, :] are contiguous and aligned to at least the tensor's alignment. Since each row is 12 bytes, alignment is at least 4 bytes but possibly more. On H100, float3 loads might not be vectorized. float4 would require padding.
Actually, we can load 3 floats as a single 12-byte load. Modern GPUs can do this. But for simplicity, let's use 3 separate loads.
Let me consider another optimization: compute multiple outputs per thread to amortize the i-coordinate load.
If each thread computes a 2x2 or 4x4 tile of outputs:
- Thread loads its i coords once
- Then computes for 4 different j values
- This amortizes the i load and gets better instruction-level parallelism
- But the j loads become less coalesced if we're not careful
Actually, let me think about this more carefully. With block (16, 16):
- 256 threads per block
- Each thread computes 1 output
- More blocks = better occupancy
With block (32, 32):
- 1024 threads per block
- Fewer blocks but still 2048 total
For a memory-bound kernel on H100, occupancy matters less than memory throughput. Let me try 32x32 blocks with the direct approach and see if it's fast enough.
Wait, I should also consider the register approach. Maybe I can use the GEMM-based method with a custom kernel that fuses everything.
Let me also consider: what if we treat this as 3 separate outer products?
For each dimension k in {0, 1, 2}:
diff_k = coords[:, :, k].unsqueeze(2) - coords[:, :, k].unsqueeze(1)
sum += diff_k ** 2
This could be done in a single kernel that iterates over k=0,1,2.
Actually, the simplest high-performance kernel is:
```cuda
template <int BLOCK>
__global__ void pairwise_distance_kernel(
const float* __restrict__ coords,
float* __restrict__ dist,
int B, int N) {
int b = blockIdx.z;
int i0 = blockIdx.y * BLOCK;
int j0 = blockIdx.x * BLOCK;
int ti = threadIdx.y;
int tj = threadIdx.x;
int i = i0 + ti;
int j = j0 + tj;
__shared__ float s_i[BLOCK][4]; // pad to avoid bank conflicts
__shared__ float s_j[BLOCK][4];
// Load coords for tile
if (i < N && tj < 3) {
s_i[ti][tj] = coords[(b * N + i) * 3 + tj];
}
if (j < N && ti < 3) {
s_j[tj][ti] = coords[(b * N + j) * 3 + ti];
}
__syncthreads();
if (i < N && j < N) {
float dx = s_i[ti][0] - s_j[tj][0];
float dy = s_i[ti][1] - s_j[tj][1];
float dz = s_i[ti][2] - s_j[tj][2];
dist[(b * N + i) * N + j] = sqrtf(dx*dx + dy*dy + dz*dz + 1e-8f);
}
}
```
Wait, I need to be more careful about the shared memory loading. Let me verify:
For s_i: all threads in the block participate. Threads with same ti and tj=0,1,2 load three values for row i. With BLOCK=32, each warp handles one or two ti values.
For warp with ti=0, tj=0..31:
- tj=0 loads coords[b, i0+0, 0]
- tj=1 loads coords[b, i0+0, 1]
- tj=2 loads coords[b, i0+0, 2]
- tj>=3 do nothing
These three loads are coalesced (contiguous addresses).
For warp with ti=1, tj=0..31:
- Same pattern but for row i0+1.
So loading s_i is perfectly coalesced.
For s_j: threads with ti=0,1,2 and all tj load coords.
Warp with ti=0, tj=0..31:
- All threads load coords[b, j0+tj, 0] - strided by 3 floats.
Warp with ti=1, tj=0..31:
- All threads load coords[b, j0+tj, 1] - strided by 3 floats.
Warp with ti=2, tj=0..31:
- All threads load coords[b, j0+tj, 2] - strided by 3 floats.
This is 3 warps with strided loads. Not ideal but okay.
Total shared memory: 32 * 4 * 4 bytes * 2 = 1024 bytes. Very small.
Let me code this up properly. But first, I realize I should also check: can we use float2/float4 to read 2 or 4 floats at once for better bandwidth?
coords[b, n, 0:2] are contiguous (8 bytes), and coords[b, n, 2] is the remaining 4 bytes. We could read coords[b, n, 0:2] as float2 and coords[b, n, 2] as float. This might help.
Let me reconsider the shared memory layout. With `s_i[BLOCK][4]`, a row has 4 floats = 16 bytes. Since shared memory banks are 4 bytes wide on H100, this is 4 banks per row. For BLOCK=32, adjacent rows access banks 0,4,8,12 and 1,5,9,13, etc. depending on access pattern. Since each thread accesses s_i[ti][0..2] and s_j[tj][0..2], threads in a warp access s_i with same ti (no conflict, broadcast) and s_j with different tj (each accesses a different row, so different banks if the rows map to different banks).
Actually, with [BLOCK][4] where BLOCK=32:
- Row 0: banks 0,1,2,3
- Row 1: banks 4,5,6,7
- ...
- Row 7: banks 28,29,30,31
- Row 8: banks 0,1,2,3 (wrap around)
For warp computing 32 outputs in a row (fixed i, varying j):
- s_i[ti][0..2]: all threads read same bank (broadcast, no conflict)
- s_j[tj][0..2]: threads read rows tj=0..31
- Rows 0..7 map to banks 0..31 (4 banks per row), no conflict within these
- Wait, each thread reads 3 floats from row tj. So 32 threads access 32 different rows. Since there are 32 banks and each row uses 4 consecutive banks, rows 0 and 8 both use banks 0-3. So threads tj=0 and tj=8 have a 2-way bank conflict. Similarly for other rows.
Hmm, bank conflicts. With BLOCK=32 and [BLOCK][4], rows 0 and 8 conflict, 1 and 9 conflict, etc.
To avoid bank conflicts, we could pad more: `s_j[BLOCK][8]` or use a different access pattern. But actually, is this a problem? On newer GPUs, shared memory bank conflicts are less of an issue because each bank might have multiple ports. But still, let's try to avoid them.
If we use `s_j[BLOCK][4]` but transpose the layout so that s_j[3][BLOCK]... Actually, if we store as `s_j[4][BLOCK]`, then for a fixed channel c and varying tj, threads access s_j[c][tj] which are all in different banks (stride BLOCK*4 bytes between elements in the same column, but within a row they're contiguous, wait...).
Let me think again. `float s_j[4][BLOCK];`
- s_j[0][0], s_j[0][1], ..., s_j[0][BLOCK-1] are contiguous in memory
- s_j[1][0] comes after s_j[0][BLOCK-1]
- Each float is 4 bytes.
Thread tj reads s_j[0][tj], s_j[1][tj], s_j[2][tj].
All threads in a warp read s_j[0][0..31]. These are contiguous, so banks 0..31 (assuming 32 banks of 4 bytes each). No conflicts!
Then they read s_j[1][0..31] - also contiguous, banks 0..31. No conflicts.
Then s_j[2][0..31] - same. No conflicts!
So if we store as `s_j[4][BLOCK]` instead of `s_j[BLOCK][4]`, we get no bank conflicts when all threads in a warp read the same channel with varying tj.
For s_i: `float s_i[4][BLOCK];` - threads with fixed ti read s_i[0..2][ti]. This is just 3 sequential reads for each thread, no conflicts since threads read different channels/rows.
Let's use this layout!
```cuda
__shared__ float s_i[4][BLOCK];
__shared__ float s_j[4][BLOCK];
```
Loading:
```cuda
if (i < N && tj < 3) {
s_i[tj][ti] = coords[(b * N + i) * 3 + tj];
}
if (j < N && ti < 3) {
s_j[ti][tj] = coords[(b * N + j) * 3 + ti];
}
```
For s_i loading: warp with ti=0, tj varies. Threads with tj=0,1,2 load s_i[0][0], s_i[1][0], s_i[2][0]. These are spaced BLOCK*4 bytes apart. If BLOCK=32, that's 128 bytes apart. So they fall in banks 0, 0, 0 (since 128 bytes = 32 floats = 32 banks * 4 bytes, so 128 bytes aligns with bank 0 again). Wait, 128 / 4 = 32, so s_i[1][0] is 32 floats from s_i[0][0]. With 32 banks of 4 bytes, 32 floats span exactly all 32 banks, so s_i[1][0] is in bank 0 again. So threads tj=0,1,2 all access bank 0! 3-way bank conflict.
Oh no. But threads tj>=3 don't participate. So we have 3 threads accessing bank 0. That's a 3-way bank conflict.
To fix this, we can have each thread load all 3 values of its row using vectorized or sequential loads, rather than having different threads load different channels.
Actually, since each thread already knows its own i, the simplest thing is:
```cuda
// Each thread loads its own row's coords
if (i < N) {
float x = coords[(b * N + i) * 3 + 0];
float y = coords[(b * N + i) * 3 + 1];
float z = coords[(b * N + i) * 3 + 2];
s_i[0][ti] = x;
s_i[1][ti] = y;
s_i[2][ti] = z;
}
if (j < N) {
float x = coords[(b * N + j) * 3 + 0];
float y = coords[(b * N + j) * 3 + 1];
float z = coords[(b * N + j) * 3 + 2];
s_j[0][tj] = x;
s_j[1][tj] = y;
s_j[2][tj] = z;
}
```
This way, all threads in a warp load sequentially from memory:
- For s_i: warp with ti=0, tj=0..31, all threads have same i. Wait, all threads in the warp have the same ti but different tj? No, in our layout:
- Warp 0: ti=0, tj=0..31. All 32 threads load coords[b, i0+0, :] - the same 3 floats! But wait, each thread loads 3 floats for its i, but all threads have the same i (ti=0). So all 32 threads try to read the same 3 floats and write to s_i[0..2][0]. This is terrible - 32 threads broadcasting the same memory address.
Hmm, let me reconsider the block layout.
What if we use 1D blocks? Or a different 2D layout?
Let me try a completely different approach. Use 1D thread blocks where each thread computes one output:
```cuda
__global__ void pairwise_distance_kernel(
const float* __restrict__ coords,
float* __restrict__ dist,
int B, int N) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = B * N * N;
if (idx >= total) return;
int b = idx / (N * N);
int rem = idx % (N * N);
int i = rem / N;
int j = rem % N;
float xi = coords[(b * N + i) * 3 + 0];
float yi = coords[(b * N + i) * 3 + 1];
float zi = coords[(b * N + i) * 3 + 2];
float xj = coords[(b * N + j) * 3 + 0];
float yj = coords[(b * N + j) * 3 + 1];
float zj = coords[(b * N + j) * 3 + 2];
float dx = xi - xj;
float dy = yi - yj;
float dz = zi - zj;
dist[idx] = sqrtf(dx*dx + dy*dy + dz*dz + 1e-8f);
}
```
With this, memory reads for i are scattered (each thread has different i), but they're strided by N*3 between threads in a warp. And j reads are also scattered.
Actually, the thread-to-output mapping is (b, i, j) -> linear idx. For adjacent threads in a block:
- Thread 0: (b, i, j)
- Thread 1: (b, i, j+1)
- ...
So adjacent threads have same i and adjacent j. Great!
For i coords: all threads in a warp with same i read the same 3 floats -> broadcast, cached.
For j coords: adjacent threads read adjacent j values with stride 3 -> strided but okay.
And we only need to do 3 broadcasts per warp for the i coords, which is very efficient with L1 cache.
But the problem is that j increases until it wraps to the next i. When j wraps from N-1 to 0, i increases by 1. At that boundary, the i coords change. But this happens at most once per 512 threads, which is 16 warps. So most warps have uniform i.
Actually wait, if blockDim.x = 256, then a warp is 32 threads. For N=512:
- Warp 0: threads 0-31, (b, i=0, j=0..31) - same i
- Warp 1: threads 32-63, (b, i=0, j=32..63) - same i
- ...
- Warp 15: threads 480-511, (b, i=0, j=480..511) - same i
- Warp 16: threads 512-543, (b, i=1, j=0..31) - same i (i=1)
So within each warp, all threads have the same i! This is perfect.
And within each warp, j values are consecutive. Great coalescing for j coords (with stride 3).
This seems like a very clean approach. Let's compute the launch config:
Total outputs = 8 * 512 * 512 = 2,097,152
If block size = 256 threads, we need 8192 blocks.
If block size = 512 threads, we need 4096 blocks.
If block size = 1024 threads, we need 2048 blocks.
With 2048 threads per block, we'd have exactly 1024 blocks... no wait, 2,097,152 / 1024 = 2048 blocks.
For H100, having 2048 blocks of 1024 threads each gives us 2 blocks per SM (2048 threads max per SM). With 132 SMs on H100, that's 264 blocks active at a time, wait no, total threads launched is 2M, but they run in waves. With 2048 blocks and 132 SMs, if each SM can run 2 blocks, that's 264 blocks concurrently. So it takes 2048/264 ≈ 8 waves.
With 256 threads per block, we have 8192 blocks, and each SM can run 8 blocks (2048/256 = 8). So 132*8 = 1056 blocks concurrently. 8192/1056 ≈ 8 waves.
Waves are similar. But smaller blocks might have better occupancy characteristics? Actually occupancy is about concurrent warps. With 1024 threads per block (32 warps), and max 64 warps per SM... wait, H100 has max 2048 threads per SM, which is 64 warps. So a 1024-thread block uses 32 warps, meaning 2 blocks per SM max. That's fine.
With 256 threads (8 warps), we can have 8 blocks per SM for 64 warps total. Same total concurrent warps.
I think either is fine. Let me go with 256 threads per block for more flexibility, or try 512. Actually, let me just try 1024 threads for fewer blocks and less scheduling overhead.
No wait, with 1D blocks and 1024 threads, the total is 2048 blocks. For H100 with 132 SMs, that's about 15-16 blocks per SM, which is fine.
Actually, I just realized: with this 1D approach, the i coords are constant within a warp but change between warps. Since i changes every 16 warps (for N=512, 512/32=16 warps per i), the same i coords will be reused by 16 warps. If these warps run concurrently or near-concurrently on the same SM, the i coords will be in L1 cache. This is good.
But the j coords change every thread. Each j is used by B*N different i values. With B=8, N=512, each j is used 4096 times. So there's a lot of reuse potential if we tile properly. But with a 1D approach, we're not explicitly reusing j coords across different i values.
Actually, let's go back to 2D tiling. With a 2D tile, each block loads a tile of i rows and j columns, and all threads in the block reuse those loaded values. This is the classic optimization for matrix operations.
With shared memory 2D tiling and proper loading, we get much better data reuse. Let me design this carefully.
For a BLOCK_SIZE x BLOCK_SIZE tile, each block loads BLOCK_SIZE rows and BLOCK_SIZE columns of coords. Each loaded coord is reused BLOCK_SIZE times.
For BLOCK_SIZE=32, each block loads 32*3 + 32*3 = 192 floats. Then computes 32*32 = 1024 outputs. Each output requires 6 float accesses (if not cached), so naive would be 6144 float reads. With shared memory tile, we only read 192 floats from global memory. That's a 32x reduction in global memory reads!
Wait, but each output doesn't just need one float. Each output needs coords for one i and one j. In a tile of 32x32, there are 32 distinct i's and 32 distinct j's. Each i is used by 32 outputs, each j is used by 32 outputs. The total "float accesses" are 1024 outputs * 6 floats = 6144. But there are only 32*3 + 32*3 = 192 unique floats in the tile. So each unique float is reused 32 times on average.
If we load these 192 floats into shared memory, all subsequent accesses hit shared memory instead of global memory. This could be much faster.
Let me design the 2D shared memory kernel more carefully.
```cuda
#define BLOCK_SIZE 32
__global__ void pairwise_distance_kernel(
const float* __restrict__ coords,
float* __restrict__ dist,
int B, int N) {
int b = blockIdx.z;
int i_base = blockIdx.y * BLOCK_SIZE;
int j_base = blockIdx.x * BLOCK_SIZE;
int ti = threadIdx.y;
int tj = threadIdx.x;
int i = i_base + ti;
int j = j_base + tj;
__shared__ float s_coords_i[BLOCK_SIZE][3];
__shared__ float s_coords_j[BLOCK_SIZE][3];
// Load coords_i[ti][0..2] for all ti
// Each thread loads one row's 3 values
if (i < N) {
s_coords_i[ti][0] = coords[(b * N + i) * 3 + 0];
s_coords_i[ti][1] = coords[(b * N + i) * 3 + 1];
s_coords_i[ti][2] = coords[(b * N + i) * 3 + 2];
}
if (j < N) {
s_coords_j[tj][0] = coords[(b * N + j) * 3 + 0];
s_coords_j[tj][1] = coords[(b * N + j) * 3 + 1];
s_coords_j[tj][2] = coords[(b * N + j) * 3 + 2];
}
__syncthreads();
if (i < N && j < N) {
float dx = s_coords_i[ti][0] - s_coords_j[tj][0];
float dy = s_coords_i[ti][1] - s_coords_j[tj][1];
float dz = s_coords_i[ti][2] - s_coords_j[tj][2];
dist[(b * N + i) * N + j] = sqrtf(dx*dx + dy*dy + dz*dz + 1e-8f);
}
}
```
Wait, there's still the issue with loading s_coords_j. In a warp with ti=0 and tj=0..31, all 32 threads load s_coords_j[tj][0..2] (each thread loads 3 values for its tj). The memory accesses for these loads are coords[b, j_base+tj, 0..2] which have stride 3 between adjacent threads. This is the same strided access as before.
But on H100 with its large L2 cache and the fact that we're only reading 32*3=96 unique floats per block, the strided access isn't that bad. And once in shared memory, accesses are fast.
For s_coords_i, in a warp with ti=0 and tj=0..31, all threads load coords[b, i_base+0, 0..2]. Wait, each thread loads ALL 3 values for row i, but all threads in the warp have the same ti=0 (if we use the warp layout I described earlier).
Wait, I got confused. Let me reconsider thread mapping.
If blockDim = (32, 32) with threadIdx.x = tj and threadIdx.y = ti:
- threadIdx linear = ti * 32 + tj
- Warp 0: threads 0-31, which are ti=0, tj=0..31
- Warp 1: threads 32-63, which are ti=1, tj=0..31
- etc.
For loading s_coords_i:
```cuda
if (i < N) {
s_coords_i[ti][0] = coords[(b * N + i) * 3 + 0];
s_coords_i[ti][1] = coords[(b * N + i) * 3 + 1];
s_coords_i[ti][2] = coords[(b * N + i) * 3 + 2];
}
```
Warp 0 (ti=0, tj=0..31): all 32 threads have i = i_base + 0. So all 32 threads try to load the same 3 floats and write to s_coords_i[0][0..2]. This is a terrible broadcast pattern - 32 threads loading the exact same address and writing to the exact same shared memory addresses.
Actually, I only need ONE thread per row to load the data. Let me change:
```cuda
if (i < N && tj == 0) {
s_coords_i[ti][0] = coords[(b * N + i) * 3 + 0];
s_coords_i[ti][1] = coords[(b * N + i) * 3 + 1];
s_coords_i[ti][2] = coords[(b * N + i) * 3 + 2];
}
if (j < N && ti == 0) {
s_coords_j[tj][0] = coords[(b * N + j) * 3 + 0];
s_coords_j[tj][1] = coords[(b * N + j) * 3 + 1];
s_coords_j[tj][2] = coords[(b * N + j) * 3 + 2];
}
```
Warp 0 (ti=0, tj=0..31): only thread with tj=0 loads row i. So 1 thread loads 3 contiguous floats. Not coalesced (only 1 thread active), but it's fine since it's just 3 floats.
For s_coords_j: only threads with ti=0 load. Warp 0: all 32 threads have ti=0! So all 32 threads load their respective j values. Thread tj loads coords[b, j_base+tj, 0..2]. This is strided by 3 floats between threads. But all 32 threads are active, so it's a "coalesced strided" pattern. The memory controller will handle it reasonably well.
This is much better! Only 32 threads per block actively load from global memory.
But we could also have all threads participate in a more coalesced way. What if we do:
```cuda
// Load s_coords_i: have ti threads load, with tj determining which component
if (i < N && tj < 3) {
s_coords_i[ti][tj] = coords[(b * N + i) * 3 + tj];
}
```
Warp 0 (ti=0, tj=0..31): threads with tj=0,1,2 load. These 3 threads load 3 contiguous floats. Coalesced! Threads with tj>=3 do nothing.
Warp 1 (ti=1, tj=0..31): same pattern for row 1.
This is very efficient loading! Only 3 active threads per warp, but each warp handles a different row, and the 3 loads are coalesced.
For s_coords_j, we can do something similar:
```cuda
if (j < N && ti < 3) {
s_coords_j[tj][ti] = coords[(b * N + j) * 3 + ti];
}
```
Warp 0 (ti=0, tj=0..31): all 32 threads have ti=0, so all 32 threads load coords[b, j_base+tj, 0]. This is strided by 3 floats.
Warp 1 (ti=1, tj=0..31): all 32 threads load coords[b, j_base+tj, 1]. Also strided.
Warp 2 (ti=2, tj=0..31): all 32 threads load coords[b, j_base+tj, 2]. Also strided.
3 warps, each with 32 strided loads.
Alternatively, for s_coords_j we could use the same pattern:
```cuda
if (j < N && ti < 3) {
// Transpose the storage
s_coords_j[ti][tj] = coords[(b * N + j) * 3 + ti];
}
```
With s_coords_j stored as [3][BLOCK_SIZE]:
- Warp 0: ti varies? No, in our mapping, warp 0 has ti=0. Let me reconsider.
Actually, if we store as s_coords_j[3][BLOCK_SIZE], and we want thread (ti, tj) to compute output using j=tj, then the computation would be:
```cuda
float dx = s_coords_i[ti][0] - s_coords_j[0][tj];
```
For this, we need s_coords_j[0][0..31] to be accessed by a warp. If s_coords_j is [3][32], then s_coords_j[0][0..31] is one row, all in different banks. Perfect!
Loading:
```cuda
if (j < N && ti < 3) {
s_coords_j[ti][tj] = coords[(b * N + j) * 3 + ti];
}
```
Warp 0 (ti=0, tj=0..31): all 32 threads load coords[b, j_base+tj, 0] and store to s_coords_j[0][tj]. Strided global read.
Warp 1 (ti=1, tj=0..31): all 32 threads load coords[b, j_base+tj, 1] and store to s_coords_j[1][tj].
Warp 2 (ti=2, tj=0..31): same for channel 2.
This stores to s_coords_j in a bank-conflict-free way because s_coords_j[0][0..31] are contiguous in shared memory, mapping to banks 0..31.
And for s_coords_i:
```cuda
if (i < N && tj < 3) {
s_coords_i[tj][ti] = coords[(b * N + i) * 3 + tj];
}
```
Store as s_coords_i[3][BLOCK_SIZE]. Warp 0 (ti=0, tj=0..31): threads 0,1,2 load and store to s_coords_i[0][0], s_coords_i[1][0], s_coords_i[2][0]. These are 32 floats apart. With 32 banks, they all map to bank 0! 3-way bank conflict.
Hmm. To avoid this, we could store as s_coords_i[BLOCK_SIZE][3]:
```cuda
if (i < N && tj < 3) {
s_coords_i[ti][tj] = coords[(b * N + i) * 3 + tj];
}
```
Warp 0: threads 0,1,2 store to s_coords_i[0][0], s_coords_i[0][1], s_coords_i[0][2]. These are contiguous: banks 0,1,2. No conflict!
Warp 1: threads store to s_coords_i[1][0..2]. Banks 3,4,5? Wait, with 32 banks of 4 bytes, 3 floats = 12 bytes = 3 banks. Row 0 uses banks 0,1,2. Row 1 uses banks 3,4,5. No conflict within the same warp. Different warps write to different rows, also no conflict.
But when reading s_coords_i[ti][0..2] in the computation, a warp with fixed ti=0 and varying tj all reads s_coords_i[0][0], but only 3 threads actually read (the ones doing the dx, dy, dz). Actually all threads read s_coords_i[ti][0], s_coords_i[ti][1], s_coords_i[ti][2] sequentially. Since all threads read the same row, and the row spans 3 banks, all threads read banks 0,1,2 sequentially. No bank conflicts because it's sequential, not simultaneous.
So `s_coords_i[BLOCK_SIZE][3]` is fine for reading.
For `s_coords_j`, if we store as `s_coords_j[BLOCK_SIZE][3]`, reading s_coords_j[tj][0..2] for varying tj means each thread reads a different row. Row tj spans banks (tj*3)%32, (tj*3+1)%32, (tj*3+2)%32. For tj=0..31:
- tj=0: banks 0,1,2
- tj=1: banks 3,4,5
- tj=2: banks 6,7,8
- ...
- tj=10: banks 30,31,0
- tj=11: banks 1,2,3
At tj=10 and tj=11, we have overlap. Specifically, for a given channel c:
- All threads read s_coords_j[tj][0] for tj=0..31.
- This accesses banks (tj*3)%32.
- tj=0: bank 0
- tj=1: bank 3
- ...
- The sequence is 0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 1, 4, 7, 10, 13, 16, 19, 22, 25, 28, 31, 2, 5, 8, 11, 14, 17, 20, 23, 26, 29
Are any two values in this sequence the same? No! Because 3 and 32 are coprime, multiplication by 3 modulo 32 is a permutation. So all 32 threads access 32 different banks when reading channel 0. Similarly for channels 1 and 2.
So `s_coords_j[BLOCK_SIZE][3]` has NO bank conflicts for reading!
This is great. Let me verify:
- `s_coords_i[BLOCK_SIZE][3]`: load with no conflicts, read with no conflicts.
- `s_coords_j[BLOCK_SIZE][3]`: load with strided global reads but no shared conflicts, read with no conflicts.
So the kernel is:
```cuda
#define BLOCK 32
__global__ void pairwise_distance_kernel(
const float* __restrict__ coords,
float* __restrict__ dist,
int B, int N) {
int b = blockIdx.z;
int i0 = blockIdx.y * BLOCK;
int j0 = blockIdx.x * BLOCK;
int ti = threadIdx.y;
int tj = threadIdx.x;
int i = i0 + ti;
int j = j0 + tj;
__shared__ float s_i[BLOCK][3];
__shared__ float s_j[BLOCK][3];
if (i < N && tj < 3) {
s_i[ti][tj] = coords[(b * N + i) * 3 + tj];
}
if (j < N && ti < 3) {
s_j[tj][ti] = coords[(b * N + j) * 3 + ti];
}
__syncthreads();
if (i < N && j < N) {
float dx = s_i[ti][0] - s_j[tj][0];
float dy = s_i[ti][1] - s_j[tj][1];
float dz = s_i[ti][2] - s_j[tj][2];
float d2 = dx*dx + dy*dy + dz*dz + 1e-8f;
dist[(b * N + i) * N + j] = sqrtf(d2);
}
}
```
Wait, I want to double-check the j loading. For warp 0 (ti=0, tj=0..31), all threads have ti=0. The condition `ti < 3` is true for all. So all 32 threads load. Thread tj loads:
- s_j[tj][0] = coords[(b * N + j0 + tj) * 3 + 0]
These are strided by 3 floats in global memory. But as discussed, this is fine.
For warp 1 (ti=1, tj=0..31):
- s_j[tj][1] = coords[(b * N + j0 + tj) * 3 + 1]
And warp 2: channel 2.
Each warp does 32 strided loads. The three warps combined cover all the data.
Is there a way to make the j loading more coalesced? If we swap the roles so that adjacent threads in a warp load adjacent memory... Actually, adjacent threads in tj DO load adjacent j values (j0+tj), but they access different channels (0 for warp 0, 1 for warp 1, 2 for warp 2). Within a warp, all threads access the same channel for different j values. These are strided by 3 floats.
To make it coalesced, we could have threads in a warp load all 3 channels for a subset of j values. But then we're changing the parallelism.
Actually, an alternative: use float4 loads or float2 loads. For coords[b, j, 0:1], we can load as float2. But coords[b, j, 0] and coords[b, j, 1] are contiguous, but coords[b, j, 2] is separate.
For j loading with a warp: all threads load coords[b, j0+tj, 0]. If each thread loads a float2 (coords[b, j0+tj, 0:1]), that's 8 bytes per thread, coalesced. Then another load for coords[b, j0+tj, 2].
Or even better, what if we transpose the coords in global memory first? But we can't change input format.
Actually, let me think about whether the strided load is really a problem. For 32 threads accessing floats with stride 3:
- Thread 0: address A
- Thread 1: address A + 12
- Thread 2: address A + 24
- ...
- Thread 31: address A + 372
The address range is 376 bytes. Each 128-byte cache line covers addresses A to A+127. So:
- Cache line 1: covers threads 0-10 (addresses A to A+110). Thread 11 address is A+132, which is in the next cache line.
- Cache line 2: covers threads 11-21 (addresses A+132 to A+252). Thread 22 is A+264.
- Cache line 3: covers threads 22-31 (addresses A+264 to A+372).
So each warp needs about 3 cache lines = 384 bytes. A single perfectly coalesced 32-float load would need 1-2 cache lines (128 bytes). So we're reading about 2-3x more data than necessary. But since the data is cached and reused (each j is used by many blocks), this is acceptable.
For the i loading: each warp loads 3 contiguous floats. This is 1 cache line, perfectly efficient.
So the shared memory approach should be quite fast. Let's go with it.
One more optimization: what if we compute multiple elements per thread? E.g., each thread computes a 2x2 tile of outputs. This would double the computation per thread, potentially improving instruction throughput, and also means we only load each i coord once per 2 outputs instead of once per output (but shared memory already handles this). It also reduces the number of blocks, which might not be beneficial.
Actually, with N=512 and BLOCK=32, we have 16x16=256 blocks per batch, 2048 total. With 132 SMs, each SM gets about 15 blocks. That's plenty.
Let me also consider: should we use BLOCK=64? Then we have 8x8=64 blocks per batch, 512 total. Each block has 64*64=4096 threads. But max threads per block is 1024, so we'd need each thread to compute 4 outputs. This gets more complex. Let's stick with BLOCK=32.
Wait, another idea: use wmma or tensor cores? For N=512, K=3, doing a batch "GEMM" via tensor cores seems like overkill and unlikely to be faster for such a tiny K.
Let me proceed with the kernel and see how it performs.
Actually, I want to also consider using `__ldg` (load through read-only cache) for the coords loads. Since coords is read-only within the kernel, `__ldg` can help.
```cuda
if (i < N && tj < 3) {
s_i[ti][tj] = __ldg(&coords[(b * N + i) * 3 + tj]);
}
```
This is a good idea. Let me include it.
Also, we can use `__restrict__` on the pointers, which we already do.
And we should mark the kernel with `__launch_bounds__` to help the compiler with register allocation. But for a simple kernel, it might not matter.
Let me also consider: can we exploit symmetry? `dist[i][j] = dist[j][i]`. If we compute only the upper triangle (i <= j), we'd roughly halve the computation and writes. But we'd need a more complex thread mapping and store pattern. For N=512, the total output is 2M floats. Halving to 1M might help, but the added complexity might hurt. Let me first try the full computation and see if it's fast enough.
One more thing: we should use `float` (fp32) precision. The problem asks for fp32 precision.
Let me write the Python code now.
```python
import torch
from torch.utils.cpp_extension import load_inline
cuda_source = '''
#include <torch/extension.h>
#include <cuda_runtime.h>
#define BLOCK 32
__global__ void pairwise_distance_kernel(
const float* __restrict__ coords,
float* __restrict__ dist,
int B, int N) {
int b = blockIdx.z;
int i0 = blockIdx.y * BLOCK;
int j0 = blockIdx.x * BLOCK;
int ti = threadIdx.y;
int tj = threadIdx.x;
int i = i0 + ti;
int j = j0 + tj;
__shared__ float s_i[BLOCK][3];
__shared__ float s_j[BLOCK][3];
if (i < N && tj < 3) {
s_i[ti][tj] = __ldg(&coords[(b * N + i) * 3 + tj]);
}
if (j < N && ti < 3) {
s_j[tj][ti] = __ldg(&coords[(b * N + j) * 3 + ti]);
}
__syncthreads();
if (i < N && j < N) {
float dx = s_i[ti][0] - s_j[tj][0];
float dy = s_i[ti][1] - s_j[tj][1];
float dz = s_i[ti][2] - s_j[tj][2];
float d2 = dx*dx + dy*dy + dz*dz + 1e-8f;
dist[(b * N + i) * N + j] = sqrtf(d2);
}
}
torch::Tensor pairwise_distance_forward(torch::Tensor coords) {
int B = coords.size(0);
int N = coords.size(1);
auto dist = torch::empty({B, N, N}, coords.options());
dim3 block(BLOCK, BLOCK);
dim3 grid((N + BLOCK - 1) / BLOCK, (N + BLOCK - 1) / BLOCK, B);
pairwise_distance_kernel<<<grid, block>>>(
coords.data_ptr<float>(),
dist.data_ptr<float>(),
B, N);
return dist;
}
'''
cpp_source = '''
torch::Tensor pairwise_distance_forward(torch::Tensor coords);
'''
pairwise_dist = load_inline(
name='pairwise_dist',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['pairwise_distance_forward'],
extra_cuda_cflags=['-O3'],
verbose=False,
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(self, coords):
return pairwise_dist.pairwise_distance_forward(coords)
```
Wait, I should double-check the grid launch. Block is (32, 32) meaning threadIdx.x=32, threadIdx.y=32. Grid is ((N+31)/32, (N+31)/32, B).
For cuda, dim3 takes (x, y, z). When we pass to <<<grid, block>>>:
- grid.x = (N+31)/32
- grid.y = (N+31)/32
- grid.z = B
And in the kernel:
- blockIdx.x ranges 0..grid.x-1
- blockIdx.y ranges 0..grid.y-1
- blockIdx.z ranges 0..grid.z-1
But I used:
```cuda
int b = blockIdx.z;
int i0 = blockIdx.y * BLOCK;
int j0 = blockIdx.x * BLOCK;
```
Wait, this means j is the "x" dimension and i is the "y" dimension. That's fine.
But in the grid setup:
```cuda
dim3 grid((N + BLOCK - 1) / BLOCK, (N + BLOCK - 1) / BLOCK, B);
```
This gives grid.x = grid.y = ceil(N/32), grid.z = B.
So blockIdx.x ranges for j, blockIdx.y ranges for i. This is correct.
However, CUDA prefers the x dimension to be the fastest-varying one, and we want memory accesses for j to be coalesced (adjacent threads in x have adjacent j). The output dist[(b*N+i)*N+j] has j as the last dimension, so adjacent j values are adjacent in memory. Our thread mapping has tj = threadIdx.x, and we compute adjacent j for adjacent tj. So writes to dist are coalesced. Good!
But wait, in my kernel, j0 = blockIdx.x * BLOCK, and i0 = blockIdx.y * BLOCK. Thread (ti, tj) computes element (i, j). So as threadIdx.x increases, j increases, meaning we write to adjacent memory locations in the inner dimension. This is perfect coalescing.
Let me now also think about fp32 accumulation. The operations are:
- dx = s_i[ti][0] - s_j[tj][0] (float)
- dx*dx (float)
- sum = dx*dx + dy*dy + dz*dz (float)
- sum + 1e-8f (float)
- sqrtf(sum) (float)
All in fp32, which matches the reference.
One more thing: the reference uses `1e-8` which is a Python float (double precision), but in CUDA I used `1e-8f` which is fp32. This should be fine and match to within tolerance.
Actually wait, `1e-8` in Python is exactly 1e-8, but `1e-8f` in CUDA might be slightly different due to float representation. Let me check: 1e-8 as a float32 is approximately 9.999999747378752e-09. The difference from 1e-8 is about 2.5e-17, negligible for fp32. The tolerance is 1e-4, so this is fine.
Let me write the complete code and compile it. I'll make sure to include the necessary imports and setup.
Actually, let me reconsider using `torch.utils.cpp_extension.load_inline`. I'll need to make sure the C++ function signature matches what pybind11 expects. Since `load_inline` auto-generates the pybind11 module, I just need to provide the function declaration in `cpp_sources` and the definition in `cuda_sources`.
One subtlety: `torch::Tensor` should be the return type and parameter type. And I need to make sure to use the right pybind11 template for overloading. Since we're not overloading, this should be fine.
Let me write the code:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
CUDA_SOURCE = r'''
#include <torch/extension.h>
#include <cuda_runtime.h>
#define BLOCK 32
__global__ void pairwise_distance_kernel(
const float* __restrict__ coords,
float* __restrict__ dist,
int B, int N) {
int b = blockIdx.z;
int i0 = blockIdx.y * BLOCK;
int j0 = blockIdx.x * BLOCK;
int ti = threadIdx.y;
int tj = threadIdx.x;
int i = i0 + ti;
int j = j0 + tj;
__shared__ float s_i[BLOCK][3];
__shared__ float s_j[BLOCK][3];
if (i < N && tj < 3) {
s_i[ti][tj] = __ldg(&coords[(b * N + i) * 3 + tj]);
}
if (j < N && ti < 3) {
s_j[tj][ti] = __ldg(&coords[(b * N + j) * 3 + ti]);
}
__syncthreads();
if (i < N && j < N) {
float dx = s_i[ti][0] - s_j[tj][0];
float dy = s_i[ti][1] - s_j[tj][1];
float dz = s_i[ti][2] - s_j[tj][2];
float d2 = dx * dx + dy * dy + dz * dz + 1e-8f;
dist[(b * N + i) * N + j] = sqrtf(d2);
}
}
torch::Tensor pairwise_distance_forward(torch::Tensor coords) {
int B = coords.size(0);
int N = coords.size(1);
auto dist = torch::empty({B, N, N}, coords.options());
dim3 block(BLOCK, BLOCK);
dim3 grid((N + BLOCK - 1) / BLOCK, (N + BLOCK - 1) / BLOCK, B);
pairwise_distance_kernel<<<grid, block>>>(
coords.data_ptr<float>(),
dist.data_ptr<float>(),
B, N);
return dist;
}
'''
CPP_SOURCE = r'''
torch::Tensor pairwise_distance_forward(torch::Tensor coords);
'''
pairwise_dist = load_inline(
name='pairwise_dist',
cpp_sources=CPP_SOURCE,
cuda_sources=CUDA_SOURCE,
functions=['pairwise_distance_forward'],
extra_cuda_cflags=['-O3'],
verbose=False,
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(self, coords: torch.Tensor) -> torch.Tensor:
return pairwise_dist.pairwise_distance_forward(coords)
```
I should be careful about `coords.options()` - this will preserve the device and dtype. Since the input is fp32, the output will be fp32.
Wait, I need to make sure that `coords` is contiguous and on CUDA. Since this is called from PyTorch, the tensor should already be on CUDA and contiguous (assuming standard inputs). But let me add `.contiguous()` just in case, though it adds overhead.
Actually, looking at the reference implementation, `diff = coords.unsqueeze(2) - coords.unsqueeze(1)` doesn't require coords to be contiguous in any particular way. But for performance, contiguous is better. Since we use `__ldg` and direct indexing, we do need the tensor to be contiguous in memory. For a standard (B, N, 3) tensor created by `torch.randn`, it is contiguous.
I think `.contiguous()` inside the Python wrapper is safe and cheap if already contiguous.
Let me also consider if we need to handle `coords` not being contiguous. The reference works with any stride. Our kernel assumes a specific layout. I should probably call `.contiguous()` to be safe.
Actually, let me reconsider one more thing. Should I use `BLOCK=64` with each thread computing multiple outputs? With N=512, BLOCK=64 would give 8x8=64 blocks per batch. Each block needs 64*3*2=384 floats of shared memory. Still tiny. Each thread would need to compute some multiple of outputs.
With BLOCK=64 and block size (32, 32), each thread computes 2x2 outputs:
- i = i0 + ti*2 + local_i
- j = j0 + tj*2 + local_j
This is more complex. Let's stick with BLOCK=32 for simplicity.
Actually, there's another consideration. For N=512 which is a multiple of 32, there are no edge cases - all tiles are fully populated. This makes the kernel very efficient because there are no warp divergence issues from boundary checks.
Let me compile and test.