reasoning
reasoning (truncated 52759 chars): Let me analyze this task:
1. The reference implementation computes pairwise Euclidean distances from 3D coordinates
2. Input: coords of shape (B, N, 3) where B=8, N=512
3. Output: dist of shape (B, N, N) - pairwise distance matrix
The reference computes:
```python
diff = coords.unsqueeze(2) - coords.unsqueeze(1) # (B, N, N, 3)
dist = torch.sqrt((diff ** 2).sum(dim=-1) + 1e-8)
```
This is essentially computing for all pairs (i,j):
dist[b,i,j] = sqrt((coords[b,i,0] - coords[b,j,0])^2 + (coords[b,i,1] - coords[b,j,1])^2 + (coords[b,i,2] - coords[b,j,2])^2 + 1e-8)
For optimization:
- The output matrix is N×N = 512×512 = 262,144 elements per batch
- The matrix is symmetric (dist[i,j] = dist[j,i]), but we need to compute all values
- We can potentially optimize by computing only the upper/lower triangle, but the output needs to be full matrix
A custom CUDA kernel can:
1. Process each batch independently
2. Each thread can compute one element of the output matrix
3. Use shared memory to load coordinates and reduce global memory accesses
Let me write an optimized kernel:
```python
import torch
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>
__global__ void pairwise_distance_kernel(
const float* __restrict__ coords,
float* __restrict__ dist,
int batch_size,
int num_atoms
) {
// Each thread block handles one batch
int b = blockIdx.x;
if (b >= batch_size) return;
// Load coordinates into shared memory
extern __shared__ float shared_coords[];
int tid = threadIdx.x;
int num_threads = blockDim.x;
// Cooperative load of coordinates (N atoms * 3 values)
for (int i = tid; i < num_atoms * 3; i += num_threads) {
shared_coords[i] = coords[b * num_atoms * 3 + i];
}
__syncthreads();
// Each thread computes multiple output elements
// Output is num_atoms x num_atoms
int total_outputs = num_atoms * num_atoms;
for (int idx = tid; idx < total_outputs; idx += num_threads) {
int i = idx / num_atoms; // row
int j = idx % num_atoms; // column
float dx = shared_coords[i * 3 + 0] - shared_coords[j * 3 + 0];
float dy = shared_coords[i * 3 + 1] - shared_coords[j * 3 + 1];
float dz = shared_coords[i * 3 + 2] - shared_coords[j * 3 + 2];
float dist_sq = dx * dx + dy * dy + dz * dz;
dist[b * num_atoms * num_atoms + i * num_atoms + j] = sqrtf(dist_sq + 1e-8f);
}
}
torch::Tensor pairwise_distance_cuda(torch::Tensor coords) {
TORCH_CHECK(coords.dim() == 3, "coords must be 3D");
TORCH_CHECK(coords.size(2) == 3, "coords last dim must be 3");
int batch_size = coords.size(0);
int num_atoms = coords.size(1);
auto dist = torch::empty({batch_size, num_atoms, num_atoms}, coords.options());
int threads_per_block = 256;
int shared_mem_size = num_atoms * 3 * sizeof(float);
pairwise_distance_kernel<<<batch_size, threads_per_block, shared_mem_size>>>(
coords.data_ptr<float>(),
dist.data_ptr<float>(),
batch_size,
num_atoms
);
return dist;
}
"""
cpp_source = """
torch::Tensor pairwise_distance_cuda(torch::Tensor coords);
"""
# Load the extension
pairwise_distance_ext = load_inline(
name='pairwise_distance_ext',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['pairwise_distance_cuda'],
extra_cuda_cflags=['-O3'],
verbose=True
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(self, coords: torch.Tensor) -> torch.Tensor:
return pairwise_distance_ext.pairwise_distance_cuda(coords)
```
Wait, I need to be more careful. Let me reconsider:
1. The shared memory approach might hit the limit (512 * 3 * 4 = 6144 bytes, which should be fine)
2. But I should consider more threads for better parallelism
Actually, let me think about a more optimized approach. With 8 batches and 262K outputs per batch:
- We could use 1D grid covering all outputs
- Or use 2D thread blocks where each thread computes one output element
Let me try a cleaner implementation with better parallelism:
Actually, let me reconsider the approach. With 512 atoms:
- 512 * 512 = 262,144 outputs per batch
- Total = 8 * 262,144 = 2,097,152 outputs
A better approach might be:
1. Use 2D thread blocks (e.g., 16x16 = 256 threads)
2. Grid covers the output matrix
3. Each thread computes one distance value
4. Use shared memory for coordinate blocks (tiling)
But for simplicity and to ensure correctness first, let me start with a straightforward implementation:
```cuda
__global__ void pairwise_distance_kernel(
const float* __restrict__ coords,
float* __restrict__ dist,
int batch_size,
int num_atoms
) {
int b = blockIdx.z;
if (b >= batch_size) return;
int i = blockIdx.y * blockDim.y + threadIdx.y;
int j = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= num_atoms || j >= num_atoms) return;
const float* ci = coords + b * num_atoms * 3 + i * 3;
const float* cj = coords + b * num_atoms * 3 + j * 3;
float dx = ci[0] - cj[0];
float dy = ci[1] - cj[1];
float dz = ci[2] - cj[2];
float dist_sq = dx * dx + dy * dy + dz * dz;
dist[b * num_atoms * num_atoms + i * num_atoms + j] = sqrtf(dist_sq + 1e-8f);
}
```
This is simple but each thread makes 6 global memory accesses (or 2 coalesced accesses if we're lucky). We can do better with shared memory tiling.
Let me write a tiled version:
```cuda
#define TILE_SIZE 16
__global__ void pairwise_distance_tiled_kernel(
const float* __restrict__ coords,
float* __restrict__ dist,
int batch_size,
int num_atoms
) {
int b = blockIdx.z;
if (b >= batch_size) return;
__shared__ float tile_i[TILE_SIZE][3];
__shared__ float tile_j[TILE_SIZE][3];
int i = blockIdx.y * TILE_SIZE + threadIdx.y;
int j = blockIdx.x * TILE_SIZE + threadIdx.x;
int ty = threadIdx.y;
int tx = threadIdx.x;
// Load tile_i (coordinates for rows)
if (i < num_atoms) {
tile_i[ty][0] = coords[b * num_atoms * 3 + i * 3 + 0];
tile_i[ty][1] = coords[b * num_atoms * 3 + i * 3 + 1];
tile_i[ty][2] = coords[b * num_atoms * 3 + i * 3 + 2];
}
// Load tile_j (coordinates for columns)
if (j < num_atoms) {
tile_j[tx][0] = coords[b * num_atoms * 3 + j * 3 + 0];
tile_j[tx][1] = coords[b * num_atoms * 3 + j * 3 + 1];
tile_j[tx][2] = coords[b * num_atoms * 3 + j * 3 + 2];
}
__syncthreads();
if (i < num_atoms && j < num_atoms) {
float dx = tile_i[ty][0] - tile_j[tx][0];
float dy = tile_i[ty][1] - tile_j[tx][1];
float dz = tile_i[ty][2] - tile_j[tx][2];
float dist_sq = dx * dx + dy * dy + dz * dz;
dist[b * num_atoms * num_atoms + i * num_atoms + j] = sqrtf(dist_sq + 1e-8f);
}
}
```
Hmm, this doesn't really use tiling effectively since we're not iterating over tiles. The standard tiling optimization for matrix multiplication iterates over tiles in the reduction dimension, but here there's no reduction - each output element is independent.
Let me think again. The key observation is:
- Each output (i,j) needs coords[i] and coords[j]
- coords[i] is reused across all j values
- coords[j] is reused across all i values
So if we have a thread block computing a tile of outputs (say 16x16), we should:
1. Load 16 coordinates for rows into shared memory
2. Load 16 coordinates for columns into shared memory
3. Compute 16x16 = 256 distances
But this only gives us a small benefit since each coordinate is only used once per thread in the tile.
A better approach: have each thread compute multiple outputs along a row or column. This way, coordinates loaded into registers are reused.
Actually, let me t
...[ omitted 52759 characters ]...
t ci0 = s_coords[i * 3 + 0];
float ci1 = s_coords[i * 3 + 1];
float ci2 = s_coords[i * 3 + 2];
float dx = ci0 - cj0;
float dy = ci1 - cj1;
float dz = ci2 - cj2;
float dist_sq = dx * dx + dy * dy + dz * dz;
dist[b * num_atoms * num_atoms + i * num_atoms + j] = sqrtf(dist_sq + 1e-8f);
}
}
torch::Tensor pairwise_distance_cuda(torch::Tensor coords) {
TORCH_CHECK(coords.dim() == 3, "coords must be 3D");
TORCH_CHECK(coords.size(2) == 3, "coords last dim must be 3");
auto coords_contiguous = coords.contiguous();
int batch_size = coords_contiguous.size(0);
int num_atoms = coords_contiguous.size(1);
auto dist = torch::empty({batch_size, num_atoms, num_atoms}, coords_contiguous.options());
int shared_mem_size = num_atoms * 3 * sizeof(float);
pairwise_distance_kernel<<<batch_size, num_atoms, shared_mem_size>>>(
coords_contiguous.data_ptr<float>(),
dist.data_ptr<float>(),
num_atoms
);
return dist;
}
"""
cpp_source = "torch::Tensor pairwise_distance_cuda(torch::Tensor coords);"
pairwise_distance_ext = load_inline(
name='pairwise_distance_ext',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['pairwise_distance_cuda'],
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_distance_ext.pairwise_distance_cuda(coords)
```
I removed `batch_size` from the kernel arguments since it's not used inside the kernel (we use `blockIdx.x` directly).
Actually wait, I just realized that the kernel uses `blockIdx.x` for batch index, which means the grid has `batch_size` blocks in the x-dimension. The kernel launch is `<<<batch_size, num_atoms, shared_mem_size>>>`, which sets:
- gridDim.x = batch_size
- blockDim.x = num_atoms
This is correct.
Let me also add a CUDA error check for safety:
```cpp
cudaError_t err = cudaGetLastError();
if (err != cudaSuccess) {
TORCH_CHECK(false, "CUDA kernel launch failed: ", cudaGetErrorString(err));
}
```
Actually, PyTorch's torch::empty and other operations should handle errors. But it's good practice to check.
Let me also consider using `__launch_bounds__` to optimize register usage:
```cuda
__global__ void __launch_bounds__(512, 2) pairwise_distance_kernel(...)
```
This tells the compiler:
- Max threads per block: 512
- Min blocks per SM: 2
With min blocks per SM = 2, the compiler will try to keep register usage low enough to allow 2 blocks per SM. This can improve occupancy.
Let me add this to the kernel.
Actually, I'm not sure about the optimal settings. Let me just use `__launch_bounds__(512)` for now:
```cuda
__global__ void __launch_bounds__(512) pairwise_distance_kernel(...)
```
This tells the compiler the max threads per block is 512, allowing it to optimize accordingly.
Final code:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>
__global__ void __launch_bounds__(512) pairwise_distance_kernel(
const float* __restrict__ coords,
float* __restrict__ dist,
int num_atoms
) {
int b = blockIdx.x;
extern __shared__ float s_coords[];
// Cooperatively load all coordinates for this batch
for (int idx = threadIdx.x; idx < num_atoms * 3; idx += blockDim.x) {
s_coords[idx] = coords[b * num_atoms * 3 + idx];
}
__syncthreads();
int j = threadIdx.x;
if (j >= num_atoms) return;
// Load coords[j] into registers
float cj0 = s_coords[j * 3 + 0];
float cj1 = s_coords[j * 3 + 1];
float cj2 = s_coords[j * 3 + 2];
for (int i = 0; i < num_atoms; i++) {
float ci0 = s_coords[i * 3 + 0];
float ci1 = s_coords[i * 3 + 1];
float ci2 = s_coords[i * 3 + 2];
float dx = ci0 - cj0;
float dy = ci1 - cj1;
float dz = ci2 - cj2;
float dist_sq = dx * dx + dy * dy + dz * dz;
dist[b * num_atoms * num_atoms + i * num_atoms + j] = sqrtf(dist_sq + 1e-8f);
}
}
torch::Tensor pairwise_distance_cuda(torch::Tensor coords) {
TORCH_CHECK(coords.dim() == 3, "coords must be 3D");
TORCH_CHECK(coords.size(2) == 3, "coords last dim must be 3");
auto coords_contiguous = coords.contiguous();
int batch_size = coords_contiguous.size(0);
int num_atoms = coords_contiguous.size(1);
auto dist = torch::empty({batch_size, num_atoms, num_atoms}, coords_contiguous.options());
int shared_mem_size = num_atoms * 3 * sizeof(float);
pairwise_distance_kernel<<<batch_size, num_atoms, shared_mem_size>>>(
coords_contiguous.data_ptr<float>(),
dist.data_ptr<float>(),
num_atoms
);
return dist;
}
"""
cpp_source = "torch::Tensor pairwise_distance_cuda(torch::Tensor coords);"
pairwise_distance_ext = load_inline(
name='pairwise_distance_ext',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['pairwise_distance_cuda'],
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_distance_ext.pairwise_distance_cuda(coords)
```
Hmm, wait. I want to double-check the write coalescing. Let me trace through:
At iteration i of the loop:
- Thread 0 computes dist[b, i, 0]
- Thread 1 computes dist[b, i, 1]
- ...
- Thread 511 computes dist[b, i, 511]
The addresses are:
- dist + b * N * N + i * N + 0
- dist + b * N * N + i * N + 1
- ...
- dist + b * N * N + i * N + 511
These are consecutive addresses! Perfect for coalesced writes.
Great, the kernel should have good write bandwidth.
Now let me also check the shared memory access for s_coords[i * 3 + k] in the loop. At iteration i:
- All threads access s_coords[i * 3 + 0], s_coords[i * 3 + 1], s_coords[i * 3 + 2]
- Since i is the same for all threads, they all access the same addresses
- This is a broadcast pattern, which is efficient
For s_coords[j * 3 + k] (loading cj values):
- Thread 0 accesses s_coords[0], s_coords[1], s_coords[2]
- Thread 1 accesses s_coords[3], s_coords[4], s_coords[5]
- ...
The addresses accessed by consecutive threads are consecutive (0, 3, 6, 9, ... for k=0). With shared memory banks, this might cause conflicts.
Let me check: for s_coords[j * 3 + k]:
- j=0: address 0+k, bank (0+k) % 32
- j=1: address 3+k, bank (3+k) % 32
- j=2: address 6+k, bank (6+k) % 32
- ...
For j=0..31 (threads in the first warp) and k=0:
- Banks: 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
All unique! No bank conflicts for the first warp.
For j=32..63 (threads in the second warp) and k=0:
- j=32: address 96, bank 96 % 32 = 0
- j=33: address 99, bank 99 % 32 = 3
- ...
Same pattern, all unique banks within each warp. Good!
Actually, wait. The warps in a 512-thread block are:
- Warp 0: threads 0-31
- Warp 1: threads 32-63
- ...
- Warp 15: threads 480-511
Each warp has threads with consecutive threadIdx.x values. For the first warp (threads 0-31), j values are 0-31. For the second warp (threads 32-63), j values are 32-63. And so on.
Within each warp, the bank access pattern for s_coords[j * 3 + k] is the same as I calculated above - all unique banks. No bank conflicts.
Great, the kernel should be efficient in terms of shared memory access.
I think this implementation is solid. Let me submit it.
Actually, one more thing: I should check if the shared memory size is within limits. 6144 bytes is well within the typical 48KB or 96KB limit per block.
OK, I'm confident in this implementation. Let me submit it.