reasoning
Now let me understand the problem:
- Input: tensor of shape (R, B, H) = (4, 12, 32)
- We need to "gather" along the last dimension - essentially reshape from (R, B, H) to (B, R*H) = (12, 128)
- Then apply a linear projection to get (B, K) = (12, 56)
The gather operation is essentially:
- `x` has shape (R, B, H)
- `gathered = torch.cat([x[i] for i in range(R)], dim=-1)` results in (B, R*H)
Wait, let me check this more carefully. `x[i]` for each i gives a tensor of shape (B, H). We concatenate along dim=-1, so we get (B, R*H).
Actually, let me trace through:
- `x` shape: (R, B, H) = (4, 12, 32)
- `x[i]` for i in range(4) gives tensors of shape (B, H) = (12, 32)
- `torch.cat(parts, dim=-1)` concatenates 4 tensors of (12, 32) along last dim → (12, 128)
So the gather is really just a reshape/view operation from (R, B, H) to (B, R*H). The key insight is that the elements are contiguous in the right order already (since we're taking slices along the first dimension and concatenating along the last).
Actually, let me think about this more carefully. If x has shape (R, B, H), then x[i] for i=0,1,2,3 are:
- x[0] → elements at indices [0, :, :] → (B, H)
- x[1] → elements at indices [1, :, :] → (B, H)
- etc.
When we concatenate these along dim=-1, we get a tensor of shape (B, R*H) where:
- gathered[j, k] for k < H comes from x[0, j, k]
- gathered[j, k] for H <= k < 2*H comes from x[1, j, k-H]
- gathered[j, k] for 2*H <= k < 3*H comes from x[2, j, k-2*H]
- gathered[j, k] for 3*H <= k < 4*H comes from x[3, j, k-3*H]
So gathered[j, r*H + h] = x[r, j, h] for r in range(R), h in range(H), j in range(B).
Now, if x is stored in row-major order (which is the default in PyTorch), the elements are:
- x[0, 0, 0], x[0, 0, 1], ..., x[0, 0, H-1], x[0, 1, 0], ..., x[0, B-1, H-1], x[1, 0, 0], ...
And gathered in row-major order is:
- gathered[0, 0], ..., gathered[0, R*H-1], gathered[1, 0], ...
Where gathered[j, r*H + h] = x[r, j, h].
So the memory access pattern is:
- gathered[j, 0], ..., gathered[j, H-1] = x[0, j, 0], ..., x[0, j, H-1] (contiguous!)
- gathered[j, H], ..., gathered[j, 2*H-1] = x[1, j, 0], ..., x[1, j, H-1] (contiguous!)
- etc.
But wait, x[0, j, 0], ..., x[0, j, H-1] are contiguous in memory (row-major), and then x[1, j, 0], ..., x[1, j, H-1] come after all of x[0, :, :] in memory. So they're NOT contiguous after each other!
Actually, the full order of x in memory (row-major) is:
x[0,0,0], x[0,0,1], ..., x[0,0,H-1], x[0,1,0], ..., x[0,B-1,H-1], x[1,0,0], x[1,0,1], ..., x[R-1,B-1,H-1]
So the full sequence is:
- All of x[0, :, :] first
- Then all of x[1, :, :]
- etc.
And we want gathered to be:
gathered[0,0], ..., gathered[0,R*H-1], gathered[1,0], ..., gathered[B-1,R*H-1]
Where gathered[j, r*H + h] = x[r, j, h].
So:
- gathered[0, 0], ..., gathered[0, H-1] = x[0, 0, 0], ..., x[0, 0, H-1] ✓ contiguous
- gathered[0, H], ..., gathered[0, 2*H-1] = x[1, 0, 0], ..., x[1, 0, H-1]
- But in x's memory, after x[0, 0, H-1] comes x[0, 1, 0], not x[1, 0, 0]!
So the gather is NOT a simple reshape. We need to actually permute the data.
OK so the fused kernel needs to:
1. Read x in shape (R, B, H) and produce gathered in shape (B, R*H)
2. Apply linear: output[b, k] = sum_{r, h} gathered[b, r*H+h] * W[r*H+h, k] + bias[k]
We can fuse this into a single kernel:
output[b, k] = sum_{r, h} x[r, b, h] * W[r*H+h, k] + bias[k]
Let me write a fused kernel for this.
For efficiency, I'll have each thread compute one output element (b, k). Each thread will:
1. Loop over r in range(R)
2. Loop over h in range(H)
3. Read x[r, b, h] and W[r*H+h, k], multiply and accumulate
4. Add bias[k]
With B=12, K=56, we have 12*56=672 output elements. We can use a 1D grid with 672 threads, or a 2D grid with (B, K) threads.
Actually, since the problem is small, we might want to use a more optimized approach:
- Each thread block handles multiple output elements
- Use shared memory to cache parts of the input or weights
Let me think about the optimal approach...
Given the sizes (B=12, R=4, H=32, K=56):
- Input size: 4*12*32 = 1536 elements
- Weight size: 128*56 = 7168 elements
- Output size: 12*56 = 672 elements
The weights are reused across all B=12 rows. So we can:
1. Load the weights into shared memory (128*56 = 7168 fp32 = 28KB, might not fit well in shared mem)
2. Or, load portions of weights as needed
Actually, 28KB for weights fits in shared memory (228KB available per SM). But we'd need to load it once per thread block.
Alternative approach: Use tensor cores or optimized matrix multiply. But for such a small problem, the overhead might not be worth it.
Let me try a simple fused kernel first and see if it's faster than the PyTorch version.
Actually, the simplest optimization might just be to avoid the intermediate `gathered` tensor altogether. Instead of:
1. x (R, B, H) → gathered (B, R*H) - memory allocation + copy
2. gathered (B, R*H) × W (R*H, K) → output (B, K) - matmul
We do:
1. x (R, B, H) with W (R*H, K) → output (B, K) - fused kernel, no intermediate tensor
Let me write this kernel:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
# Define the CUDA kernel
cuda_source = """
__global__ void fused_gather_linear_kernel(
const float* __restrict__ x,
const float* __restrict__ weight,
const float* __restrict__ bias,
float* __restrict__ output,
int R, int B, int H, int K) {
int b = blockIdx.x * blockDim.x + threadIdx.x;
int k = blockIdx.y * blockDim.y + threadIdx.y;
if (b < B && k < K) {
float sum = bias[k];
for (int r = 0; r < R; r++) {
for (int h = 0; h < H; h++) {
int x_idx = r * B * H + b * H + h;
int w_idx = r * H + h; // Actually weight is (K, R*H), so we need to be careful
// Wait, let me check PyTorch Linear's weight shape
}
}
output[b * K + k] = sum;
}
}
"""
# Wait, I need to check the weight shape for nn.Linear
# nn.Linear(in_features, out_features) has weight of shape (out_features, in_features)
# So weight is (K, R*H) = (56, 128)
# And the computation is output = x @ weight.T + bias
# Or equivalently: output[b, k] = sum_{j} gathered[b, j] * weight[k, j] + bias[k]
```
Let me fix the weight indexing:
- weight shape is (K, R*H) = (56, 128)
- weight[k, r*H + h] is the weight for input at position r*H+h and output k
So:
```
output[b, k] = bias[k] + sum_{r, h} x[r, b, h] * weight[k, r*H + h]
```
Let me write the correct kernel:
```cuda
__global__ void fused_gather_linear_kernel(
const float* __restrict__ x,
const float* __restrict__ weight,
const float* __restrict__ bias,
float* __restrict__ output,
int R, int B, int H, int K) {
int b = blockIdx.x * blockDim.x + threadIdx.x;
int k = blockIdx.y * blockDim.y + threadIdx.y;
if (b < B && k < K) {
float sum = bias[k];
for (int r = 0; r < R; r++) {
for (int h = 0; h < H; h++) {
int x_idx = r * B * H + b * H + h;
int w_idx = k * (R * H) + r * H + h;
sum += x[x_idx] * weight[w_idx];
}
}
output[b * K + k] = sum;
}
}
```
For thread configuration, with B=12 and K=56:
- We can use a 2D grid: (B, K) = (12, 56)
- Block size: say (16, 16) or (32, 16), with enough threads to cover
- Grid size: ceil(12/16) x ceil(56/16) = 1 x 4 = 4 blocks
Actually, for such a small problem, maybe a 1D grid where each thread handles multiple outputs would be better. Or we can use vector loads.
Let me first try the straightforward approach and see if it works.
Actually, I realize I should think about memory access patterns:
- x is read with indices (r, b, h) → memory access is coalesced if threads in a warp access consecutive h values for the same (r, b)
- weight is read with indices (k, r*H+h) → memory access is coalesced if threads in a warp access consecutive k values for the same (r, h)
So for coalesced access:
- Group threads by (b, r, h) or by (k, ...) depending on which gives better coalescing
If threads are organized as (b, k), then:
- Each thread loops over (r, h)
- For the same (b, k), different threads in a warp would have different k values
- So they would read weight[k, r*H+h] with consecutive k values → coalesced!
- But they would read x[r, b, h] with the same b, different k values → all threads read the same x values → not coalesced, but cached
This seems reasonable. Let me try it.
Actually, let me also consider using vector loads for x. Since each thread reads R*H = 128 elements from x (contiguous in memory for each r), we could potentially use float4 loads.
Wait, but x is not contiguous in the way we need it. Let me think again...
For a given (b), we need x[:, b, :] which is:
- x[0, b, 0], x[0, b, 1], ..., x[0, b, H-1] (H contiguous elements)
- x[1, b, 0], x[1, b, 1], ..., x[1, b, H-1] (another H contiguous elements)
- ...
- x[R-1, b, 0], ..., x[R-1, b, H-1]
So for each b, we have R separate contiguous chunks of H elements each. Not fully contiguous, but each chunk is H elements long.
OK let me just write a straightforward kernel and see if it works:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
R = 4
B = 12
H = 32
K = 56
cuda_source = """
__global__ void fused_gather_linear_kernel(
const float* __restrict__ x,
const float* __restrict__ weight,
const float* __restrict__ bias,
float* __restrict__ output,
int R, int B, int H, int K) {
int b = blockIdx.x * blockDim.x + threadIdx.x;
int k = blockIdx.y * blockDim.y + threadIdx.y;
if (b < B && k < K) {
float sum = bias[k];
for (int r = 0; r < R; r++) {
for (int h = 0; h < H; h++) {
int x_idx = r * B * H + b * H + h;
int w_idx = k * (R * H) + r * H + h;
sum += x[x_idx] * weight[w_idx];
}
}
output[b * K + k] = sum;
}
}
void fused_gather_linear(
at::Tensor x,
at::Tensor weight,
at::Tensor bias,
at::Tensor output,
int R, int B, int H, int K) {
dim3 block(16, 16);
dim3 grid((B + 15) / 16, (K + 15) / 16);
fused_gather_linear_kernel<<<grid, block>>>(
x.data_ptr<float>(),
weight.data_ptr<float>(),
bias.data_ptr<float>(),
output.data_ptr<float>(),
R, B, H, K);
}
"""
cpp_source = """
void fused_gather_linear(
at::Tensor x,
at::Tensor weight,
at::Tensor bias,
at::Tensor output,
int R, int B, int H, int K);
"""
module = load_inline(
name='fused_gather_linear',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['fused_gather_linear'],
extra_cuda_cflags=['-O3'],
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
self.proj = nn.Linear(R * H, K, bias=True)
def forward(self, x: torch.Tensor) -> torch.Tensor:
r, b, h = x.shape
output = torch.empty(b, K, device=x.device, dtype=x.dtype)
module.fused_gather_linear(
x, self.proj.weight, self.proj.bias, output,
R, B, H, K)
return output
```
Wait, I need to be careful about the constants R, B, H, K. They're defined at the module level, so they should be accessible in the class. But let me pass them from the input shape to be more general.
Also, I realize I should check that my understanding of the weight shape is correct. Let me trace through the reference:
- `self.proj = nn.Linear(R * H, K, bias=True)` creates a Linear layer with in_features=128, out_features=56
- `self.proj.weight` has shape (K, R*H) = (56, 128) in PyTorch
- `self.proj.bias` has shape (K,) = (56,)
- The forward is `self.proj(gathered)` where gathered has shape (B, R*H) = (12, 128)
- This computes `gathered @ weight.T + bias` → output shape (B, K) = (12, 56)
So output[b, k] = sum_{j} gathered[b, j] * weight[k, j] + bias[k]
Where gathered[b, r*H + h] = x[r, b, h]
So output[b, k] = sum_{r, h} x[r, b, h] * weight[k, r*H + h] + bias[k]
This matches what I wrote. Let me compile and test it.