reasoning
Let me understand the task. I need to replace a PyTorch model that implements a differentiable Smith-Waterman-style local sequence alignment scoring with a custom CUDA kernel.
The model:
1. Takes a score_matrix of shape (B, M, N)
2. Computes a DP matrix H where H[i][j] = max(0, H[i-1][j-1]+S[i][j], H[i-1][j]+gap, H[i][j-1]+gap)
3. Uses logsumexp instead of max for differentiability (soft-max relaxation)
4. Returns the max value over all cells for each batch
The key insight is that this is a DP algorithm with dependencies - each cell depends on the cell to its left, above, and diagonally above-left. This means we can't easily parallelize across cells, but we can parallelize across the batch dimension.
Let me think about how to parallelize this on GPU:
1. **Batch parallelism**: We can process each batch element independently in parallel. This is the most straightforward approach.
2. **Diagonal wavefront**: For a single sequence, cells on the same anti-diagonal (i+j = constant) are independent of each other. We could process anti-diagonals in sequence, with each anti-diagonal processed in parallel. But this is more complex.
Given B=16, M=32, N=32, the simplest approach is to use batch parallelism - one thread block per batch element.
Let me write the CUDA kernel:
```cpp
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <float.h>
__global__ void smith_waterman_kernel(
const float* score_matrix,
float* H,
float* output,
float gap_open,
float temperature,
int B, int M, int N
) {
int b = blockIdx.x;
if (b >= B) return;
// Each block handles one batch element
// We need shared memory for H matrix
extern __shared__ float shared_H[];
// Initialize H[0][j] = 0 and H[i][0] = 0
for (int i = threadIdx.x; i <= M; i += blockDim.x) {
shared_H[i * (N + 1) + 0] = 0.0f;
}
for (int j = threadIdx.x; j <= N; j += blockDim.x) {
shared_H[0 * (N + 1) + j] = 0.0f;
}
__syncthreads();
// DP computation
for (int i = 1; i <= M; i++) {
for (int j = 1; j <= N; j++) {
// Compute logsumexp([0, match, delete, insert])
float match = shared_H[(i-1) * (N+1) + (j-1)] + score_matrix[b * M * N + (i-1) * N + (j-1)];
float delete = shared_H[(i-1) * (N+1) + j] + gap_open;
float insert = shared_H[i * (N+1) + (j-1)] + gap_open;
// logsumexp of [0, match, delete, insert]
float vals[4] = {0.0f, match, delete, insert};
float max_val = vals[0];
for (int k = 1; k < 4; k++) {
if (vals[k] > max_val) max_val = vals[k];
}
float sum_exp = 0.0f;
for (int k = 0; k < 4; k++) {
sum_exp += expf(vals[k] - max_val);
}
shared_H[i * (N+1) + j] = temperature * logf(sum_exp) + max_val;
__syncthreads();
}
}
// Find max value
float max_val = -INFINITY;
for (int idx = threadIdx.x; idx < (M+1)*(N+1); idx += blockDim.x) {
if (shared_H[idx] > max_val) max_val = shared_H[idx];
}
// Reduce to find max
__shared__ float shared_max[256];
shared_max[threadIdx.x] = max_val;
__syncthreads();
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
if (threadIdx.x < s) {
if (shared_max[threadIdx.x + s] > shared_max[threadIdx.x]) {
shared_max[threadIdx.x] = shared_max[threadIdx.x + s];
}
}
__syncthreads();
}
if (threadIdx.x == 0) {
output[b] = shared_max[0];
}
}
```
Wait, I have an issue. The DP computation has dependencies within a row - each cell depends on the cell to its left. So we can't parallelize within a row easily. We need to process cells sequentially.
Let me reconsider. For the DP part, we process each cell sequentially but have multiple threads cooperatively. Actually, for this problem size (M=32, N=32), a single thread can process the entire DP matrix for one batch element. With 16 batch elements, we have 16 independent computations.
Let me simplify:
```cpp
__global__ void smith_waterman_kernel(
const float* score_matrix,
float* output,
float gap_open,
float temperature,
int B, int M, int N
) {
int b = blockIdx.x;
if (b >= B) return;
// Shared memory for H matrix
extern __shared__ float H[];
// One thread per block processes the DP
if (threadIdx.x == 0) {
// Initialize first row and column to 0
for (int j = 0; j <= N; j++) {
H[j] = 0.0f;
}
for (int i = 0; i <= M; i++) {
H[i * (N + 1)] = 0.0f;
}
// DP computation
for (int i = 1; i <= M; i++) {
for (int j = 1; j <= N; j++) {
float match = H[(i-1) * (N+1) + (j-1)] + score_matrix[b * M * N + (i-1) * N + (j-1)];
float delete = H[(i-1) * (N+1) + j] + gap_open;
float insert = H[i * (N+1) + (j-1)] + gap_open;
// logsumexp of [0, match, delete, insert]
float vals[4] = {0.0f, match, delete, insert};
float max_val = vals[0];
for (int k = 1; k < 4; k++) {
if (vals[k] > max_val) max_val = vals[k];
}
float sum_exp = 0.0f;
for (int k = 0; k < 4; k++) {
sum_exp += expf(vals[k] - max_val);
}
H[i * (N+1) + j] = temperature * logf(sum_exp) + max_val;
}
}
// Find max
float max_val = -INFINITY;
for (int idx = 0; idx < (M+1)*(N+1); idx++) {
if (H[idx] > max_val) max_val = H[idx];
}
output[b] = max_val;
}
}
```
This is very inefficient - only 16 threads doing any work on a GPU with thousands of cores. Let me think about better parallelization.
Actually, let's think about the anti-diagonal approach. Cells on the same anti-diagonal (i+j = constant) are independent because:
- H[i][j] depends on H[i-1][j-1] (diagonal), H[i-1][j] (above), H[i][j-1] (left)
- None of these are on the same anti-diagonal as H[i][j]
So we can process anti-diagonals in sequence, with all cells on an anti-diagonal processed in parallel.
For a 32x32 matrix, we have 64 anti-diagonals. Each anti-diagonal has at most 32 cells. This gives us some parallelism, but not huge.
Actually, the original reference code iterates:
```python
for i in range(1, M + 1):
for j in range(1, N + 1):
```
This is row-major order. Each cell H[i][j] depends on:
- H[i-1][j-1]: computed in the previous row
- H[i-1][j]: computed in the previous row
- H[i][j-1]: computed just before in the same row
So within a row, we have sequential dependencies. The anti-diagonal approach is correct for parallelization.
Let me implement the anti-diagonal approach:
```cpp
__global__ void smith_waterman_kernel(
const float* score_matrix,
float* output,
float gap_open,
float temperature,
int B, int M, int N
) {
int b = blockIdx.x;
if (b >= B) return;
// Shared memory for H matrix
extern __shared__ float H[];
// Initialize H to 0
for (int idx = threadIdx.x; idx < (M+1)*(N+1); idx += blockDim.x) {
H[idx] = 0.0f;
}
__syncthreads();
// Process anti-diagonals
for (int d = 2; d <= M + N; d++) {
// anti-diagonal d has cells where i+j = d
// i ranges from max(1, d-N) to min(M, d-1)
int i_start = max(1, d - N);
int i_end = min(M, d - 1);
for (int i = i_start + threadIdx.x; i <= i_end; i += blockDim.x) {
int j = d - i;
float match = H[(i-1) * (N+1) + (j-1)] + score_matrix[b * M * N + (i-1) * N + (j-1)];
float delete = H[(i-1) * (N+1) + j] + gap_open;
float insert = H[i * (N+1) + (j-1)] + gap_open;
float vals[4] = {0.0f, match, delete, insert};
float max_val = vals[0];
for (int k = 1; k < 4; k++) {
if (vals[k] > max_val) max_val = vals[k];
}
float sum_exp = 0.0f;
for (int k = 0; k < 4; k++) {
sum_exp += expf(vals[k] - max_val);
}
H[i * (N+1) + j] = temperature * logf(sum_exp) + max_val;
}
__syncthreads();
}
// Find max using parallel reduction
float local_max = -INFINITY;
for (int idx = threadIdx.x; idx < (M+1)*(N+1); idx += blockDim.x) {
if (H[idx] > local_max) local_max = H[idx];
}
__shared__ float shared_max[256];
shared_max[threadIdx.x] = local_max;
__syncthreads();
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
if (threadIdx.x < s) {
if (shared_max[threadIdx.x + s] > shared_max[threadIdx.x]) {
shared_max[threadIdx.x] = shared_max[threadIdx.x + s];
}
}
__syncthreads();
}
if (threadIdx.x == 0) {
output[b] = shared_max[0];
}
}
```
Actually, I think for this problem size, we might be better off just having one thread per batch element, since M=32 and N=32 are small. The overhead of synchronization might outweigh the benefits of parallelism within a batch.
Let me first write a simple version and then profile to see where the bottlenecks are.
Actually wait, let me re-read the problem. The temperature is used in the logsumexp:
```python
H[:, i, j] = self.temperature * torch.logsumexp(stack / self.temperature, dim=-1)
```
So we compute logsumexp(stack / temperature) * temperature. This is numerically more stable than just logsumexp(stack).
Let me make sure I implement this correctly:
- vals = [0, match, delete, insert] / temperature
- result = temperature * logsumexp(vals)
Let me also think about numerical stability. In the logsumexp, we need to handle -inf correctly.
Actually, looking at the reference code more carefully:
```python
neg_inf = torch.tensor(float("-inf"), device=score_matrix.device, dtype=score_matrix.dtype)
# ... but neg_inf is never used
```
Hmm, neg_inf is defined but not used. The zero is always 0.0 (not -inf), so the logsumexp always includes 0.0. This means the result is always >= 0 (since max includes 0.0).
Let me write the kernel:
```cpp
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <float.h>
__device__ float logsumexp4(float a, float b, float c, float d, float temp) {
// logsumexp([a, b, c, d] / temp) * temp
float vals[4] = {a/temp, b/temp, c/temp, d/temp};
float max_val = vals[0];
for (int k = 1; k < 4; k++) {
if (vals[k] > max_val) max_val = vals[k];
}
float sum_exp = 0.0f;
for (int k = 0; k < 4; k++) {
sum_exp += expf(vals[k] - max_val);
}
return temp * (logf(sum_exp) + max_val);
}
__global__ void smith_waterman_kernel(
const float* __restrict__ score_matrix,
float* __restrict__ output,
float gap_open,
float temperature,
int B, int M, int N
) {
int b = blockIdx.x;
if (b >= B) return;
// Shared memory for H matrix
extern __shared__ float H[];
// Initialize H to 0
for (int idx = threadIdx.x; idx < (M+1)*(N+1); idx += blockDim.x) {
H[idx] = 0.0f;
}
__syncthreads();
// DP computation - anti-diagonal wavefront
for (int d = 2; d <= M + N; d++) {
int i_start = max(1, d - N);
int i_end = min(M, d - 1);
for (int i = i_start + threadIdx.x; i <= i_end; i += blockDim.x) {
int j = d - i;
float match = H[(i-1) * (N+1) + (j-1)] + score_matrix[b * M * N + (i-1) * N + (j-1)];
float delete = H[(i-1) * (N+1) + j] + gap_open;
float insert = H[i * (N+1) + (j-1)] + gap_open;
H[i * (N+1) + j] = logsumexp4(0.0f, match, delete, insert, temperature);
}
__syncthreads();
}
// Find max using parallel reduction
float local_max = -INFINITY;
for (int idx = threadIdx.x; idx < (M+1)*(N+1); idx += blockDim.x) {
if (H[idx] > local_max) local_max = H[idx];
}
__shared__ float shared_max[256];
shared_max[threadIdx.x] = local_max;
__syncthreads();
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
if (threadIdx.x < s) {
if (shared_max[threadIdx.x + s] > shared_max[threadIdx.x]) {
shared_max[threadIdx.x] = shared_max[threadIdx.x + s];
}
}
__syncthreads();
}
if (threadIdx.x == 0) {
output[b] = shared_max[0];
}
}
torch::Tensor smith_waterman_cuda(
torch::Tensor score_matrix,
float gap_open,
float temperature
) {
int B = score_matrix.size(0);
int M = score_matrix.size(1);
int N = score_matrix.size(2);
auto output = torch::empty({B}, score_matrix.options());
int block_size = 256;
int shared_mem = (M + 1) * (N + 1) * sizeof(float);
smith_waterman_kernel<<<B, block_size, shared_mem>>>(
score_matrix.data_ptr<float>(),
output.data_ptr<float>(),
gap_open,
temperature,
B, M, N
);
return output;
}
```
Wait, I have a typo - `shared_max` vs `shared_max`. Let me fix that.
Also, I should check if the shared memory size is within limits. For M=32, N=32, we need (33*33) = 1089 floats = 4356 bytes, which is well within the 48KB shared memory limit.
Let me also think about optimizations:
1. Use `__restrict__` pointers
2. Use shared memory for H
3. Optimize the logsumexp computation
Actually, for better performance, I might want to consider:
1. Storing H in a way that reduces bank conflicts
2. Using warp-level primitives for the reduction
3. Loop unrolling
But let me first get a working version and then profile it.
Let me write the complete Python file:
```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>
__device__ float logsumexp4(float a, float b, float c, float d, float temp) {
float vals[4] = {a/temp, b/temp, c/temp, d/temp};
float max_val = vals[0];
for (int k = 1; k < 4; k++) {
if (vals[k] > max_val) max_val = vals[k];
}
float sum_exp = 0.0f;
for (int k = 0; k < 4; k++) {
sum_exp += expf(vals[k] - max_val);
}
return temp * (logf(sum_exp) + max_val);
}
__global__ void smith_waterman_kernel(
const float* __restrict__ score_matrix,
float* __restrict__ output,
float gap_open,
float temperature,
int B, int M, int N
) {
int b = blockIdx.x;
if (b >= B) return;
extern __shared__ float H[];
// Initialize H to 0
for (int idx = threadIdx.x; idx < (M+1)*(N+1); idx += blockDim.x) {
H[idx] = 0.0f;
}
__syncthreads();
// DP computation - anti-diagonal wavefront
for (int d = 2; d <= M + N; d++) {
int i_start = max(1, d - N);
int i_end = min(M, d - 1);
for (int i = i_start + threadIdx.x; i <= i_end; i += blockDim.x) {
int j = d - i;
float match = H[(i-1) * (N+1) + (j-1)] + score_matrix[b * M * N + (i-1) * N + (j-1)];
float delete = H[(i-1) * (N+1) + j] + gap_open;
float insert = H[i * (N+1) + (j-1)] + gap_open;
H[i * (N+1) + j] = logsumexp4(0.0f, match, delete, insert, temperature);
}
__syncthreads();
}
// Find max using parallel reduction
float local_max = -INFINITY;
for (int idx = threadIdx.x; idx < (M+1)*(N+1); idx += blockDim.x) {
if (H[idx] > local_max) local_max = H[idx];
}
__shared__ float shared_max[256];
shared_max[threadIdx.x] = local_max;
__syncthreads();
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
if (threadIdx.x < s) {
if (shared_max[threadIdx.x + s] > shared_max[threadIdx.x]) {
shared_max[threadIdx.x] = shared_max[threadIdx.x + s];
}
}
__syncthreads();
}
if (threadIdx.x == 0) {
output[b] = shared_max[0];
}
}
torch::Tensor smith_waterman_cuda(
torch::Tensor score_matrix,
float gap_open,
float temperature
) {
int B = score_matrix.size(0);
int M = score_matrix.size(1);
int N = score_matrix.size(2);
auto output = torch::empty({B}, score_matrix.options());
int block_size = 256;
int shared_mem = (M + 1) * (N + 1) * sizeof(float);
smith_waterman_kernel<<<B, block_size, shared_mem>>>(
score_matrix.data_ptr<float>(),
output.data_ptr<float>(),
gap_open,
temperature,
B, M, N
);
return output;
}
"""
cpp_source = "torch::Tensor smith_waterman_cuda(torch::Tensor score_matrix, float gap_open, float temperature);"
module = load_inline(
name="smith_waterman",
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=["smith_waterman_cuda"],
verbose=True,
extra_cuda_cflags=["-O3"],
)
class ModelNew(nn.Module):
def __init__(self, gap_open=-1.0, temperature=1.0):
super().__init__()
self.gap_open = gap_open
self.temperature = temperature
def forward(self, score_matrix: torch.Tensor) -> torch.Tensor:
return module.smith_waterman_cuda(score_matrix, self.gap_open, self.temperature)
```
Wait, I see a bug in my kernel - I'm using `shared_max` but the array is named `shared_max`. Let me fix that.
Also, I need to be careful with the thread block size. If blockDim.x is 256, but shared_max is declared as `float shared_max[256]`, we need to make sure threadIdx.x is always less than 256. This should be fine since we launch with block_size = 256.
But wait, there's another issue - the reference uses `self.gap_open` which is added to delete and insert, not multiplied. Let me double-check:
```python
delete = H[:, i - 1, j] + self.gap_open
insert = H[:, i, j - 1] + self.gap_open
```
Yes, it's added. And `self.gap_open = -1.0` by default, so this is a penalty.
Let me compile and test this.
Actually, wait. I need to re-read the code more carefully. In the reference:
```python
stack = torch.stack([zero, match, delete, insert], dim=-1)
H[:, i, j] = self.temperature * torch.logsumexp(stack / self.temperature, dim=-1)
```
So the stack is [zero, match, delete, insert] (in that order), then divided by temperature, then logsumexp, then multiplied by temperature.
My logsumexp4 function does: logsumexp([a/temp, b/temp, c/temp, d/temp]) * temp
This is equivalent to: temp * log(sum(exp((val - max)/temp))) + temp * max/temp = temp * log(sum(...)) + max
Wait, I need to be more careful. Let me trace through:
logsumexp([x1, x2, x3, x4]) = log(exp(x1) + exp(x2) + exp(x3) + exp(x4))
= log(exp(x1 - max) * exp(max) + ...)
= log(exp(max) * (exp(x1-max) + ...))
= max + log(exp(x1-max) + ...)
So logsumexp([a/temp, b/temp, c/temp, d/temp]) = max + log(sum(exp(val - max)))
And multiplying by temp: temp * max + temp * log(sum(...))
But I compute: temp * (log(sum(...)) + max) which is the same thing.
OK so my implementation should be correct. Let me just double-check the formula:
```python
H[:, i, j] = self.temperature * torch.logsumexp(stack / self.temperature, dim=-1)
```
stack = [0, match, delete, insert]
So H[i,j] = temperature * logsumexp([0/temp, match/temp, delete/temp, insert/temp])
In my kernel:
```cpp
H[i * (N+1) + j] = logsumexp4(0.0f, match, delete, insert, temperature);
```
And logsumexp4 does:
```cpp
float vals[4] = {a/temp, b/temp, c/temp, d/temp};
...
return temp * (logf(sum_exp) + max_val);
```
This should be correct.
Let me also check if I handle the boundary correctly. The anti-diagonal starts at d=2 (for i=1, j=1) and goes to d=M+N (for i=M, j=N).
For d=2: i ranges from max(1, 2-N) to min(M, 1), which is [1, 1]. So only i=1, j=1.
For d=M+N: i ranges from max(1, M+N-N) to min(M, M+N-1), which is [M, M]. So only i=M, j=N.
This looks correct.
Let me compile and test.
Actually, I realized there's still a typo in my kernel code - I have `shared_max` and `shared_max` mixed up. Let me fix that and submit.
Actually, let me reconsider the algorithm. For this problem size (M=32, N=32), each anti-diagonal has at most 32 elements. With a block size of 256, most threads are idle during the DP computation. This is not efficient.
Let me think about a different approach. What if each thread processes a row, and we use wavefront parallelism?
Actually, let's try a simpler approach first: one thread per batch element. This gives us 16 threads (one per batch), which is still underutilizing the GPU but might be faster than the Python reference due to CUDA overhead being lower than Python loop overhead.
Actually, the Python reference has a double nested loop that runs M*N = 1024 iterations per batch, times B=16 = 16384 iterations total. Each iteration does a stack, logsumexp, etc. This is very slow in Python.
Even a single-threaded CUDA kernel per batch element would be much faster because:
1. No Python interpreter overhead
2. Native floating point operations
3. Better memory access patterns
Let me try the simplest approach first and see if it works.
Actually, the anti-diagonal approach with wavefront parallelism should work well. The issue is that with small M and N (32x32), the parallelism within a batch is limited. But we have 16 batches, so we get 16 blocks running in parallel. That's already better than the sequential Python code.
Let me also consider using warp-level primitives for the reduction. But first, let me get the basic version working.
Here's the corrected kernel:
```cpp
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
__device__ float logsumexp4(float a, float b, float c, float d, float temp) {
float vals[4] = {a/temp, b/temp, c/temp, d/temp};
float max_val = vals[0];
for (int k = 1; k < 4; k++) {
if (vals[k] > max_val) max_val = vals[k];
}
float sum_exp = 0.0f;
for (int k = 0; k < 4; k++) {
sum_exp += expf(vals[k] - max_val);
}
return temp * (logf(sum_exp) + max_val);
}
__global__ void smith_waterman_kernel(
const float* __restrict__ score_matrix,
float* __restrict__ output,
float gap_open,
float temperature,
int B, int M, int N
) {
int b = blockIdx.x;
if (b >= B) return;
extern __shared__ float H[];
// Initialize H to 0
for (int idx = threadIdx.x; idx < (M+1)*(N+1); idx += blockDim.x) {
H[idx] = 0.0f;
}
__syncthreads();
// DP computation - anti-diagonal wavefront
for (int d = 2; d <= M + N; d++) {
int i_start = max(1, d - N);
int i_end = min(M, d - 1);
for (int i = i_start + threadIdx.x; i <= i_end; i += blockDim.x) {
int j = d - i;
float match = H[(i-1) * (N+1) + (j-1)] + score_matrix[b * M * N + (i-1) * N + (j-1)];
float delete = H[(i-1) * (N+1) + j] + gap_open;
float insert = H[i * (N+1) + (j-1)] + gap_open;
H[i * (N+1) + j] = logsumexp4(0.0f, match, delete, insert, temperature);
}
__syncthreads();
}
// Find max using parallel reduction
float local_max = -INFINITY;
for (int idx = threadIdx.x; idx < (M+1)*(N+1); idx += blockDim.x) {
if (H[idx] > local_max) local_max = H[idx];
}
__shared__ float s_max[256];
s_max[threadIdx.x] = local_max;
__syncthreads();
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
if (threadIdx.x < s) {
if (s_max[threadIdx.x + s] > s_max[threadIdx.x]) {
s_max[threadIdx.x] = s_max[threadIdx.x + s];
}
}
__syncthreads();
}
if (threadIdx.x == 0) {
output[b] = s_max[0];
}
}
```
OK let me now write the complete Python file and test it.