reasoning
Let me analyze the problem:
1. There's a depthwise 2D convolution with a kernel of shape `(kernel_size, 1)` - so it's actually a 1D convolution applied along the height dimension (with width dimension being 1).
2. The input is `(batch_size, in_channels, height, width)`.
3. The kernel size is `(kernel_size, 1)`, meaning it convolves only along the height dimension, treating width as 1 (no horizontal convolution).
4. This is essentially a depthwise separable convolution along the height dimension only.
Wait, let me re-read: `kernel_size=(kernel_size, 1)` means the kernel has `kernel_size` rows and 1 column. So for each output pixel at (h, w), it takes a vertical strip of the input with height `kernel_size * dilation` and width 1, centered at w.
Actually, since the kernel width is 1, this convolves along the height dimension independently for each width position.
I need to implement this with a custom CUDA kernel to beat PyTorch.
Key considerations:
- The default PyTorch implementation may not be super optimized for this asymmetric kernel case.
- I should write a specialized kernel for this specific case: `(K, 1)` kernel, depthwise.
- Since kernel width is 1, each output element only depends on `kernel_size` input elements vertically aligned.
Let me think about memory access patterns:
- For each output position (b, c, h, w), we read input at (b, c, h', w) for h' in [h * stride + k * dilation - padding] for k in 0..kernel_size-1.
- Since width is shared and we're reading vertically, memory access is strided by width. This is not great for coalescing unless we handle it carefully.
Actually for a (K, 1) kernel, each output channel depends only on the same input channel. Within a channel, for a given (h, w), we read K elements along the height at the same width.
For better memory coalescing, threads in a warp should access contiguous memory. So if I assign threads to contiguous w positions, then for a given (b, c, h), threads access (b, c, h', w), (b, c, h', w+1), etc. which is contiguous! That's great for coalescing.
So the kernel should:
- Output dimensions: H_out = floor((H + 2*padding - dilation*(K-1) - 1) / stride) + 1
- W_out = floor((W + 2*padding - dilation*(0) - 1) / stride) + 1 = W when padding is 0 and stride is 1. Wait, for kernel width 1: W_out = floor((W + 2*padding - 0 - 1)/stride) + 1.
Actually for (K, 1) kernel: W_out = floor((W + 2*padding - 1*(1-1) - 1)/stride) + 1 = floor((W + 2*padding - 1)/stride) + 1.
If padding=0, stride=1: W_out = W.
If padding=1, stride=1: W_out = W + 2.
Let me assume general cases but optimize for the common test case.
For each output pixel, the computation is:
output[b, c, h, w] = sum_k weight[c, k, 0] * input[b, c, h * stride + k * dilation - padding, w] + bias[c]
Wait, PyTorch uses NCHW format, and for Conv2d:
output[b, c, h, w] = sum_{i=0}^{K_h-1} sum_{j=0}^{K_w-1} input[b, c, h*stride + i*dilation - padding, w*stride + j*dilation - padding] * weight[c, i, j]
Since K_w = 1, j only takes 0:
output[b, c, h, w] = sum_{i=0}^{K_h-1} input[b, c, h*stride + i*dilation - padding, w*stride - padding] * weight[c, i, 0]
Wait! The width also gets strided and padded! So:
output[b, c, h, w] = sum_{i=0}^{K_h-1} input[b, c, h*stride_h + i*dilation_h - padding_h, w*stride_w - padding_w] * weight[c, i, 0]
Hmm, actually for PyTorch Conv2d, stride can be a tuple (stride_h, stride_w) but here it's just an int, so stride_h = stride_w = stride.
padding is also the same for both.
So: output[b, c, h, w] = sum_{i=0}^{K-1} input[b, c, h*stride + i*dilation - padding, w*stride - padding] * weight[c, i, 0]
Wait that's strange - the width position is `w*stride - padding` which doesn't depend on i, but it also means output width depends on stride and padding.
For the test case: stride=1, padding=0, dilation=1:
W_out = W, H_out = H - K + 1 = 512 - 3 + 1 = 510.
And for each output: output[b, c, h, w] = sum_{i=0}^{2} input[b, c, h + i, w] * weight[c, i, 0].
So it's indeed a 1D conv along height only, independently for each width position.
Memory access: for a fixed (b, c, h), reading input[b, c, h+h', w] for all w - this is perfectly coalesced if threads handle consecutive w values.
Let me think about the kernel design:
Option 1: Each thread handles one output element.
- Great simple approach.
- Need to coalesce reads. If threads in a block handle consecutive w values, then reading input at the same (b, c, h+h', w) for consecutive w is coalesced.
Option 2: Use shared memory to tile and reuse.
- Could load a tile of input into shared memory.
For K=3 and large W (512), each thread reads 3 elements. Shared memory might not help much since each input element is only used once (well, K times across K different output rows). Actually for 1D conv along height, input element at row h is used in K output rows: h, h-1, h-2, ... (if they exist). So there's reuse vertically but not horizontally.
If a thread block handles a tile of (H_tile, W_tile) output elements, then it needs input of size (H_tile + K - 1, W_tile). The W_tile dimension can be loaded coalesced. But H_tile elements per thread block... if threads are organized along W, then only output rows within the block are sharing input rows.
Let me think about a sliding window approach in shared memory:
- A block of threads handles some output rows and some range of width.
- Load input tile into shared memory.
- Compute output.
For simplicity, let me start with a direct kernel where each thread computes one output element, but block is organized to maximize coalescing.
Grid: blocks handle ranges of (b, c, h, w).
Maybe: blockDim = (256, 1, 1). Each block handles 256 consecutive width positions.
Grid: (ceil(W/256), H_out, B*C) or similar.
Wait, the total output elements per (b,c,h) is W. With W=512, we could have 2 blocks per row.
If blockDim.x = 256, grid:
- x: ceil(W / 256)
- y: H_out
- z: B * C
Each thread computes one output at (b, c, h, w) where:
b = blockIdx.z / C
c = blockIdx.z % C
h = blockIdx.y
w = blockIdx.x * blockDim.x + threadIdx.x
Reads: for i in 0..K-1:
in_h = h * stride + i * dilation - padding
in_w = w * stride - padding # Wait, shouldn't there be j for width? For kernel width 1, j=0 so term is 0*dilation_w - padding_w.
Actually, let me verify PyTorch conv2d behavior. The formula is:
out_h = (H + 2*pad_h - dilation_h*(K_h-1) - 1) // stride_h + 1
out_w = (W + 2*pad_w - dilation_w*(K_w-1) - 1) // stride_w + 1
And the actual index is:
in_h = out_h_idx * stride_h - pad_h + k_h * dilation_h
in_w = out_w_idx * stride_w - pad_w + k_w * dilation_w
Since K_w = 1, k_w is always 0. So:
in_w = out_w_idx * stride_w - pad_w
For test case where stride=1, pad=0:
in_w = out_w_idx
in_h = out_h_idx + k_h * dilation
So yes, each output at (h, w) reads from input at (h + k_h*dilation, w) for k_h in 0..K_h-1.
With W=512 and blockDim.x=256, threads in a warp read consecutive w values, which is coalesced.
Let me write this kernel. But I should also consider using shared memory. For K=3, each thread reads 3 floats. Total memory traffic is 3 * 4 bytes = 12 bytes per thread per output. With 256 threads per block, that's 3KB of input per block.
If I use shared memory, I could potentially reduce global memory reads, but since each input element is read by K different output positions (vertically adjacent), the cache should handle this well. Given the H100 has huge cache bandwidth, direct global memory access with coalescing might be sufficient.
Actually, let me think about register usage. Each thread needs to accumulate a sum and read K weights (can be kept in registers). For K=3:
- Each thread loads weight[c, 0, 0], weight[c, 1, 0], weight[c, 2, 0] (shared among threads in the same c).
- Each thread reads K input values and does K multiply-adds.
I can optimize by having threads in a block collaboratively load weights into shared memory. But weights are only K*C elements, which is tiny. Each thread in the same channel c needs the same 3 weights. If a block spans multiple channels or batch elements, each might need different weights.
Wait, if I organize so blockIdx.z = b*C + c, then all threads in a block process the same (b, c) pair. They all need the same K weights. I can load these K weights into shared memory once per block (or even just have thread 0 load them into registers/shmem, and then broadcast).
Actually, maybe the compiler handles this well if I just load the weights directly. K is likely small, so unrolling the loop helps.
Let me also consider: is there a case where K is large? The test case has K=3, but the model should work for any kernel_size.
I should probably template the kernel or use a compile-time constant if possible. With `load_inline`, I can use `#define KERNEL_SIZE %d` in the string.
Let me write a flexible kernel but optimize for small K by unrolling.
Here's my plan:
1. In `__init__`, store parameters.
2. In `forward`, compute output shape.
3. Call a CUDA kernel that handles the depthwise conv.
CUDA Kernel:
```cuda
__global__ void depthwise_conv_k1_kernel(
const float* __restrict__ input,
const float* __restrict__ weight,
const float* __restrict__ bias,
float* __restrict__ output,
int B, int C, int H, int W,
int K, int stride, int padding, int dilation,
int H_out, int W_out
) {
int w = blockIdx.x * blockDim.x + threadIdx.x;
int h = blockIdx.y;
int bc = blockIdx.z;
int b = bc / C;
int c = bc % C;
if (w >= W_out) return;
// Load weights for this channel into registers
// Use shared memory for weights if K is large, otherwise registers
float sum = 0.0f;
#pragma unroll
for (int k = 0; k < K; k++) {
int in_h = h * stride + k * dilation - padding;
int in_w = w * stride - padding; // since kernel width is 1
if (in_h >= 0 && in_h < H && in_w >= 0 && in_w < W) {
float val = input[((b * C + c) * H + in_h) * W + in_w];
float wgt = weight[(c * K + k)]; // weight is [C, K, 1]
sum += val * wgt;
}
}
if (bias != nullptr) {
sum += bias[c];
}
output[((b * C + c) * H_out + h) * W_out + w] = sum;
}
```
Wait, I need to be careful about the `in_w` bounds check. For PyTorch conv2d with stride=1, padding=0, kernel_width=1: W_out = W. And in_w = w, so it's always in bounds. But for general padding/stride, I should keep the check.
But for performance, the bounds check hurts. Since this is a specialized kernel for (K, 1), and for the common case where padding and stride keep everything in bounds, maybe I can skip the check if we know the input is valid. Actually, no - for general padding, some output positions might reference out-of-bound input positions, which PyTorch treats as 0.
Actually wait, in PyTorch, if padding causes input indices to be out of bounds, the effective padding is such that no out-of-bounds access happens for the computed output size. Let me think again.
The output size is computed as:
H_out = (H + 2*pad - dilation*(K-1) - 1) / stride + 1
For any output index h in [0, H_out), the input index in_h = h * stride - pad + k * dilation.
The minimum in_h is -pad (when h=0, k=0).
The maximum in_h is (H_out-1)*stride - pad + (K-1)*dilation.
For these to be valid [0, H), we need:
0 <= -pad + something, and max < H.
Actually, PyTorch's conv2d with padding just implicitly pads with zeros. So for negative indices or indices >= H, the value is 0. So I do need the bounds check.
But for the test case with padding=0, there's no out-of-bounds for width (since kernel width is 1 and stride=1), but there might be for height.
For height with padding=0, K=3, dilation=1, stride=1:
in_h = h + k. For h=0, k can be 0,1,2. The max in_h for h=H_out-1 is H_out-1+K-1 = H-1+K-1 = H+K-2 which is >= H. Wait.
Actually H_out = H - K + 1 for padding=0, stride=1, dilation=1.
For h in [0, H_out-1]:
min in_h = 0 (when h=0, k=0)
max in_h = (H - K) + (K - 1) = H - 1.
So all input indices are in bounds! For the height dimension with padding=0, everything is valid.
In general, with padding=p:
H_out = H + 2p - K + 1 (for stride=1, dilation=1).
For h=0: in_h could be -p + k. Some might be negative.
For h=H_out-1: in_h = H_out-1-p+k = H+2p-K+1-1-p+k = H+p-K+k. For k=K-1, in_h = H+p-1 >= H if p>0.
So for p>0, some accesses are out of bounds and should be treated as 0.
I need to handle this but optimize the inner loop. One way is to precompute valid range and unroll.
Actually, for simplicity and since K is typically small, I'll keep the bounds check but use it carefully. The compiler might optimize it well.
Wait, I just realized: since all threads in a warp have the same h and k, they all have the same in_h condition result. So the branch is perfectly uniform across a warp. That's great for performance.
Similarly for the width check, if W_out is such that some threads in a block are out of bounds, those threads just return early. The remaining threads execute uniformly.
Now, about memory access: for a given (b, c, h, k), all threads in a warp read input at different w positions, so they read consecutive memory addresses. This is perfectly coalesced.
Let me consider if we can do better. H100 has 3.35 TB/s memory bandwidth. For an output element, we read K floats from input + K multiplies + 1 store. If K=3:
- Read 3 floats from global memory = 12 bytes
- Read 3 floats from weight = 12 bytes (but these are likely cached/L1 since same for all threads in block)
- Write 1 float = 4 bytes
Total: ~16 bytes per output element.
Compute: 3 FMAs = 3 operations.
Arithmetic intensity: 3 / 16 = 0.1875 FLOP/byte.
At 3.35 TB/s, this can sustain 0.1875 * 3.35e12 = 628 GFLOPS. FP32 peak is 67 TFLOPS. So this is definitely memory bandwidth bound.
The peak memory bandwidth is 3.35 TB/s. To get good performance, I need:
1. Good coalescing (done)
2. Maybe use vectorized loads or half-precision if possible
3. Use shared memory to reduce redundant weight loads
Actually, the weight loads are not redundant in terms of global memory bandwidth because they're broadcast (all threads in warp read same address, which gets cached, but counts toward bandwidth once per warp). The input reads are the main cost.
Can I use half-precision?) No, the problem says fp32 precision and correctness.
Wait, could I use tensor cores or something? For K=3, that's way too small.
What about using shared memory to cache input and enable reuse? For a tile of output rows, we can load a strip of input that covers all needed height positions, reducing global memory reads if rows overlap. But since each input row is only used by K output rows, and if a block handles output rows h_start to h_end, the input needed is h_start to h_end+K-1.
If I have a 2D thread block handling, say, a 16x64 tile of output, I can load a 19x64 tile of input into shared memory. But then each input element is still only read once from global memory - we don't actually save global memory bandwidth unless there are multiple output positions per (b, c, h) that reuse the same input, which only happens across different h outputs.
Actually, the benefit of shared memory here is mainly to improve latency hiding and cache utilization. But the H100 has excellent L1/L2 cache, so direct global access might already be close to peak.
Let me also consider if unrolling the K loop helps. With a compile-time K, I get better register allocation and instruction scheduling. But since K is a constructor argument, it varies.
I can template the kernel on K, or use a runtime parameter and hope the compiler unrolls. For the test case K=3, loop unrolling should help.
Actually, I can use a function where K is passed as a template parameter or generated dynamically. With `load_inline`, I can generate multiple kernels for common K values and dispatch, or just use the runtime parameter version.
Let me write a runtime parameter version first and see if it's fast enough. If not, I can generate specialized versions.
Another optimization: since `in_w` is the same across all k iterations (because kernel width is 1), I can compute the base address once and just add H*W strides.
Let me optimize the indexing:
```cuda
int base_in_idx = ((b * C + c) * H) * W + in_w;
int base_in_offset = k * dilation * W;
```
Actually:
input[((b * C + c) * H + in_h) * W + in_w] = input[base + in_h * W + in_w]
where base = (b * C + c) * H * W.
For consecutive k, in_h increases by dilation, so index increases by dilation * W.
Let me write a cleaner kernel:
```cuda
__global__ void depthwise_conv_k1_kernel(
const float* __restrict__ input,
const float* __restrict__ weight,
const float* __restrict__ bias,
float* __restrict__ output,
int B, int C, int H, int W,
int K, int stride, int pad_h, int pad_w, int dilation,
int H_out, int W_out
) {
int w = blockIdx.x * blockDim.x + threadIdx.x;
int h = blockIdx.y;
int bc = blockIdx.z;
int b = bc / C;
int c = bc % C;
if (w >= W_out) return;
int in_w = w * stride - pad_w;
int in_h_start = h * stride - pad_h;
float sum = 0.0f;
// Load weights for this channel
const float* wptr = weight + c * K;
#pragma unroll 4
for (int k = 0; k < K; k++) {
int in_h = in_h_start + k * dilation;
if (in_h >= 0 && in_h < H && in_w >= 0 && in_w < W) {
float val = input[((b * C + c) * H + in_h) * W + in_w];
sum += val * wptr[k];
}
}
if (bias != nullptr) {
sum += bias[c];
}
output[((b * C + c) * H_out + h) * W_out + w] = sum;
}
```
Wait, I need to be more careful. The weight shape for `nn.Conv2d(in_channels, in_channels, kernel_size=(kernel_size, 1), ...)` is `(out_channels, in_channels_per_group, K_h, K_w)`. Since it's depthwise (groups=in_channels), out_channels = in_channels, and in_channels_per_group = 1. So weight shape is `(C, 1, K, 1)`.
PyTorch stores conv2d weights as `(out_channels, in_channels // groups, K_h, K_w)`. For depthwise, that's `(C, 1, K, 1)`. In memory this is contiguous as `[c][0][k][0]`. So the offset is `c * 1 * K * 1 + 0 * K * 1 + k * 1 + 0 = c * K + k`.
So `weight[c * K + k]` is correct.
Now, about output shape. I can compute it in Python or CUDA. I'll compute it in Python and pass it in.
Let me think about block size. 256 threads per block is standard. For W_out = 512, that's 2 blocks per row. For H_out = 510, that's 510 rows per (b, c). Total blocks = 2 * 510 * 64 * 8 = 522240 blocks. That's a lot! With ~132 SMs on H100, we have lots of blocks.
Actually 522240 / 132 = ~3956 blocks per SM. Each block has 256 threads. 3956 blocks with 256 threads is a lot, but since each block is small (computes ~256 output elements with K=3 loads each), the total work per block is tiny.
The overhead of launching so many blocks might be significant. Let me increase work per thread.
Instead of each thread computing 1 output element, each thread can compute multiple output elements. For example:
- Each block has 128 threads.
- Each thread computes 4 consecutive output elements in W.
- Block handles 512 output elements in W.
Or organized differently:
- blockDim.x = 256
- Each thread computes 1 output element
This should be fine on H100.
But actually, let me check if there's a better organization. For height dimension, since vertical neighbors share input rows (with offset), maybe processing multiple height outputs per block using shared memory is better.
Consider a thread block that handles a 2D tile: say 16 rows x 64 columns of output.
- Needs input of (16 + K - 1) rows x 64 cols = 18x64 for K=3.
- 18*64 = 1152 floats = 4.5KB. Fits nicely in shared memory.
- Block could be 16x4 or 8x8 or whatever.
But the issue is that threads need to load input cooperatively, then compute. The shared memory load would be coalesced. The computation would then access shared memory, which is fast.
However, with direct global access, the data should be in L1 cache since consecutive warps processing consecutive rows will have significant row overlap. Actually, on H100, L2 cache should capture a lot of this. The real question is whether the direct approach is already close to peak bandwidth.
Peak bandwidth is 3.35 TB/s. Let's calculate theoretical throughput for direct:
- Per output: K global reads (input) + 1 global write (output) = 4 bytes * (K + 1).
- For K=3: 16 bytes per output.
- At 3.35 TB/s: 3.35e12 / 16 = 209e9 outputs/second.
- Total outputs: B * C * H_out * W_out = 64 * 8 * 510 * 512 = 133.6e6 outputs.
- Theoretical minimum time: 133.6e6 / 209e9 = 0.000639 seconds = 639 microseconds.
Hmm, that's the absolute minimum at peak bandwidth. If PyTorch takes, say, 1-2 ms, we'd need to get close to this. Shared memory won't reduce the fundamental bytes needed.
Wait, actually the weight loads also consume some bandwidth, but they're broadcast and heavily cached. If we don't count them (because they hit L1), then the main traffic is input reads + output writes = 16 bytes/output for K=3.
But PyTorch's Conv2d might be doing more work because it's a general-purpose implementation. It might not recognize this as a pure 1D conv and could do extra indexing, padding, etc.
Actually, PyTorch might dispatch to cuDNN. cuDNN is usually very well optimized. But for this unusual case of asymmetric kernel (K, 1), it might not have a highly tuned kernel and could use a more general GEMM-based approach, which might involve im2col expansion. That would increase memory traffic significantly!
If cuDNN uses im2col for this, it would:
1. Expand input into a column matrix of size (K*1) x (H_out*W_out) per channel.
2. Multiply by weight vector of size K*1.
The im2col step would read each input element K times (since each input is used in K output positions vertically). That's K times more memory traffic for input! So total input traffic would be K * H * W * B * C * 4 bytes, vs our approach which reads each input element only once (approximately, per output). Wait, no - each input element is used by K output positions in our direct computation too.
Actually, in the direct computation, for K=3, each input element (except edges) is read 3 times total across all outputs. In im2col, it's also read 3 times but stored to intermediate memory, then read again. So im2col has an extra write and read pass.
If cuDNN is using im2col for this small K, our direct kernel could be significantly faster by avoiding the intermediate memory.
Let me check by writing the kernel and comparing. I'll write the most straightforward direct kernel first.
Another thought: could I use a 1D convolution PyTorch op instead? Like `F.conv1d`? The input could be reshaped to (B, C, H) for each w, but that would require transposing. Probably not faster.
Or `F.conv2d` with a grouped convolution might already be well-optimized by cuDNN. I need to check.
Let me also vectorize memory access. If W is a multiple of 4 or 8, I can use float4 loads. For W=512, that's divisible by 4. This could improve bandwidth utilization.
Vectorized kernel:
- Process 4 consecutive width positions per thread using float4.
- Input pointer is cast to float4*.
- Output pointer is cast to float4*.
But this is a bit messy with general W. I'll try the simple version first.
Actually, let me think about another optimization: since `in_w` doesn't change with k, I can prefetch or access memory more efficiently.
Let me write the initial kernel:
```python
kernel_code = '''
import torch
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void depthwise_conv_k1_kernel(
const float* __restrict__ input,
const float* __restrict__ weight,
const float* __restrict__ bias,
float* __restrict__ output,
int B, int C, int H, int W,
int K, int stride, int pad_h, int pad_w, int dilation,
int H_out, int W_out
) {
int w = blockIdx.x * blockDim.x + threadIdx.x;
int h = blockIdx.y;
int bc = blockIdx.z;
int b = bc / C;
int c = bc % C;
if (w >= W_out) return;
int in_w = w * stride - pad_w;
int in_h_start = h * stride - pad_h;
float sum = 0.0f;
const float* wptr = weight + c * K;
#pragma unroll 4
for (int k = 0; k < K; k++) {
int in_h = in_h_start + k * dilation;
if (in_h >= 0 && in_h < H && in_w >= 0 && in_w < W) {
float val = input[(((b * C) + c) * H + in_h) * W + in_w];
sum += val * wptr[k];
}
}
if (bias != nullptr) {
sum += bias[c];
}
output[(((b * C) + c) * H_out + h) * W_out + w] = sum;
}
torch::Tensor depthwise_conv_k1_forward(
torch::Tensor input,
torch::Tensor weight,
c10::optional<torch::Tensor> bias,
int stride, int pad_h, int pad_w, int dilation
) {
int B = input.size(0);
int C = input.size(1);
int H = input.size(2);
int W = input.size(3);
int K = weight.size(2); // weight is [C, 1, K, 1]
int H_out = (H + 2 * pad_h - dilation * (K - 1) - 1) / stride + 1;
int W_out = (W + 2 * pad_w - 1 * (1 - 1) - 1) / stride + 1; // K_w = 1
auto output = torch::empty({B, C, H_out, W_out}, input.options());
const int threads = 256;
dim3 block(threads);
dim3 grid((W_out + threads - 1) / threads, H_out, B * C);
const float* bias_ptr = bias.has_value() ? bias.value().data_ptr<float>() : nullptr;
depthwise_conv_k1_kernel<<<grid, block>>>(
input.data_ptr<float>(),
weight.data_ptr<float>(),
bias_ptr,
output.data_ptr<float>(),
B, C, H, W,
K, stride, pad_h, pad_w, dilation,
H_out, W_out
);
return output;
}
"""
cpp_source = """
torch::Tensor depthwise_conv_k1_forward(
torch::Tensor input,
torch::Tensor weight,
c10::optional<torch::Tensor> bias,
int stride, int pad_h, int pad_w, int dilation
);
"""
depthwise_conv = load_inline(
name="depthwise_conv_k1",
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=["depthwise_conv_k1_forward"],
extra_cuda_cflags=["--use_fast_math", "-O3"],
verbose=False,
)
class ModelNew(nn.Module):
...
'''
```
Wait, I need to be more careful with the C++ signature and pybind11. The `load_inline` auto-generates the pybind11 module. I should not include `PYBIND11_MODULE`. Let me verify the pattern.
The `cpp_sources` argument should contain declarations, and `functions` lists the function names. `load_inline` generates the binding.
Actually, looking at PyTorch docs, `cpp_sources` can contain helper code, and the functions in `functions=[]` are automatically bound.
Let me write a cleaner version.
```python
cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void depthwise_conv_k1_kernel(
const float* __restrict__ input,
const float* __restrict__ weight,
const float* __restrict__ bias,
float* __restrict__ output,
int B, int C, int H, int W,
int K, int stride, int pad_h, int pad_w, int dilation,
int H_out, int W_out
) {
int w = blockIdx.x * blockDim.x + threadIdx.x;
int h = blockIdx.y;
int bc = blockIdx.z;
int b = bc / C;
int c = bc % C;
if (w >= W_out) return;
int in_w = w * stride - pad_w;
int in_h_start = h * stride - pad_h;
float sum = 0.0f;
const float* wptr = weight + c * K;
#pragma unroll 4
for (int k = 0; k < K; k++) {
int in_h = in_h_start + k * dilation;
if (in_h >= 0 && in_h < H && in_w >= 0 && in_w < W) {
float val = input[(((b * C) + c) * H + in_h) * W + in_w];
sum += val * wptr[k];
}
}
if (bias != nullptr) {
sum += bias[c];
}
output[(((b * C) + c) * H_out + h) * W_out + w] = sum;
}
torch::Tensor depthwise_conv_k1_forward(
torch::Tensor input,
torch::Tensor weight,
c10::optional<torch::Tensor> bias,
int stride, int pad_h, int pad_w, int dilation
) {
int B = input.size(0);
int C = input.size(1);
int H = input.size(2);
int W = input.size(3);
int K = weight.size(2);
int H_out = (H + 2 * pad_h - dilation * (K - 1) - 1) / stride + 1;
int W_out = (W + 2 * pad_w - 1) / stride + 1;
auto output = torch::empty({B, C, H_out, W_out}, input.options());
const int threads = 256;
dim3 block(threads);
dim3 grid((W_out + threads - 1) / threads, H_out, B * C);
const float* bias_ptr = bias.has_value() ? bias.value().data_ptr<float>() : nullptr;
depthwise_conv_k1_kernel<<<grid, block>>>(
input.data_ptr<float>(),
weight.data_ptr<float>(),
bias_ptr,
output.data_ptr<float>(),
B, C, H, W,
K, stride, pad_h, pad_w, dilation,
H_out, W_out
);
return output;
}
"""
cpp_source = "torch::Tensor depthwise_conv_k1_forward(torch::Tensor input, torch::Tensor weight, c10::optional<torch::Tensor> bias, int stride, int pad_h, int pad_w, int dilation);"
```
Actually, the C++ source needs to be valid. `c10::optional` should work. Let me also make sure to include the right headers.
Wait, after reviewing PyTorch's `load_inline` examples, I think the simplest approach is:
```python
cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void depthwise_conv_k1_kernel(...) { ... }
torch::Tensor depthwise_conv_k1_forward(
torch::Tensor input, torch::Tensor weight,
c10::optional<torch::Tensor> bias,
int64_t stride, int64_t pad_h, int64_t pad_w, int64_t dilation
) {
...
}
"""
```
And then:
```python
module = load_inline(
name="depthwise_conv_k1",
cpp_sources="", # or declaration
cuda_sources=cuda_source,
functions=["depthwise_conv_k1_forward"],
extra_cuda_cflags=["-O3"],
)
```
Actually, for `load_inline`, either `cpp_sources` or `cuda_sources` can contain the implementation. Usually people put the CUDA code in `cuda_sources` and C++ wrapper in `cpp_sources`, or put everything in `cuda_sources`.
Let me just put the implementation in `cuda_sources` and pass an empty string or forward declaration in `cpp_sources`.
Wait, I think for `load_inline` with `cuda_sources`, the CUDA source is compiled with nvcc, and the `cpp_sources` is compiled with the C++ compiler, then linked. So I can put all the code in `cuda_sources` and leave `cpp_sources` as just the forward declarations.
Let me check the exact syntax from memory...:
```python
module = torch.utils.cpp_extension.load_inline(
name="my_custom",
cpp_sources="torch::Tensor my_func(torch::Tensor x);",
cuda_sources="""#include <torch/extension.h>
__global__ void kernel(...) {}
torch::Tensor my_func(torch::Tensor x) { ... }
""",
functions=["my_func"],
)
```
Yes, something like this. Let me write the complete code.
One more thing: PyTorch conv2d weight shape. For `nn.Conv2d(in_channels, in_channels, kernel_size=(kernel_size, 1), groups=in_channels)`, the weight is initialized as:
`self.weight = Parameter(torch.empty((out_channels, in_channels // groups, *kernel_size)))`
= `(in_channels, 1, kernel_size, 1)`.
So `weight.size(0) = C`, `weight.size(1) = 1`, `weight.size(2) = K`, `weight.size(3) = 1`.
I'll make sure to handle weight reshaping/access correctly.
For the `ModelNew` class:
```python
class ModelNew(nn.Module):
def __init__(self, in_channels, kernel_size, stride=1, padding=0, dilation=1, bias=False):
super().__init__()
self.in_channels = in_channels
self.kernel_size = kernel_size
self.stride = stride
self.padding = padding
self.dilation = dilation
self.has_bias = bias
self.weight = nn.Parameter(torch.empty(in_channels, 1, kernel_size, 1))
if bias:
self.bias = nn.Parameter(torch.empty(in_channels))
else:
self.register_parameter('bias', None)
# init
nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
if self.bias is not None:
fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
bound = 1 / math.sqrt(fan_in)
nn.init.uniform_(self.bias, -bound, bound)
def forward(self, x):
return depthwise_conv.depthwise_conv_k1_forward(
x, self.weight, self.bias,
self.stride, self.padding, self.padding, self.dilation
)
```
Wait, if I define my own weight and bias, I need to copy from the reference model during initialization? No, the test framework will instantiate `ModelNew` with the constructor args and run forward on random inputs. It doesn't copy weights from `Model`. But wait, the reference uses `nn.Conv2d` which initializes weights with a specific scheme. My `ModelNew` should probably use the same initialization for correctness if weights matter, but since inputs are random and we test output equivalence for the same model, not for a pre-trained model, I think any consistent initialization is fine.
Actually, wait - in the evaluation framework, they probably instantiate both `Model` and `ModelNew` and might copy parameters. Let me check the task description again... "Replace the forward computation of the PyTorch model below with a custom CUDA kernel."
Typically in these problems, they instantiate `ModelNew` and run it on inputs, comparing to `Model` run on the same inputs. But if `ModelNew` uses different initialization, the outputs will differ unless the test code specifically copies state_dict.
Given the test code says:
```python
def get_inputs():
x = torch.rand(batch_size, in_channels, height, width)
return [x]
def get_init_inputs():
return [in_channels, kernel_size, stride, padding, dilation]
```
It seems like they instantiate `ModelNew(*get_init_inputs())` and `Model(*get_init_inputs())` separately. If both initialize differently, outputs won't match.
So I should make sure my initialization matches PyTorch's default. PyTorch Conv2d uses `kaiming_uniform_` for weights and a uniform for bias. Let me use the same.
Actually, looking at PyTorch's Conv2d source:
```python
init.kaiming_uniform_(self.weight, a=math.sqrt(5))
if self.bias is not None:
fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
bound = 1 / math.sqrt(fan_in)
init.uniform_(self.bias, -bound, bound)
```
I'll copy this exactly.
But wait, there's another issue: `nn.Conv2d` stores weight as `(out_channels, in_channels // groups, *kernel_size)`. For my custom weight, I'll store it the same way so that if the test framework copies parameters, shapes match.
Actually, does the test framework copy parameters? Let me think... The most likely scenario is that they use `get_init_inputs()` to create both models, then use state_dict transfer or something. Or they might just test that the models produce the same output for the same random input, regardless of parameter values (since both are freshly initialized).
If both are freshly initialized with different random seeds, the outputs will differ. So either:
1. The framework copies the state_dict from `Model` to `ModelNew`.
2. The framework seeds the RNG before creating both.
3. The framework compares architecture/behavior, not actual outputs (unlikely given the correctness tools).
Given that the task says "Produce numerically equivalent outputs at fp32 precision", they must be ensuring either initialization is the same or parameters are copied.
To be safe, I'll use `nn.Conv2d` in my `__init__` and just call its forward method OR store its parameters. Actually, the simplest approach:
```python
class ModelNew(nn.Module):
def __init__(self, *args, **kwargs):
super().__init__()
self.conv2d = nn.Conv2d(*args, **kwargs)
def forward(self, x):
return my_custom_kernel(x, self.conv2d.weight, self.conv2d.bias, ...)
```
This way, initialization is exactly the same, and if the framework copies state_dict, it will work. This is the most robust approach.
Yes! I'll use this approach. In `__init__`, create a `nn.Conv2d` to get the exact same parameter storage and initialization, then in `forward`, call my custom kernel with the stored weight and bias.
This is much safer.
Now let me write the first version and test it.
```python
import torch
import torch.nn as nn
import math
class ModelNew(nn.Module):
def __init__(self, in_channels, kernel_size, stride=1, padding=0, dilation=1, bias=False):
super(ModelNew, self).__init__()
self.conv2d = nn.Conv2d(in_channels, in_channels, kernel_size=(kernel_size, 1), stride=stride, padding=padding, dilation=dilation, groups=in_channels, bias=bias)
def forward(self, x):
return depthwise_conv_k1_forward(
x, self.conv2d.weight, self.conv2d.bias,
self.conv2d.stride[0], self.conv2d.padding[0], self.conv2d.padding[1] if len(self.conv2d.padding) > 1 else self.conv2d.padding[0], self.conv2d.dilation[0]
)
```
Wait, for PyTorch Conv2d:
- `stride` is a tuple `(stride_h, stride_w)` even if initialized with int.
- `padding` is a tuple `(pad_h, pad_w)` even if initialized with int.
- `dilation` is a tuple `(dilation_h, dilation_w)`.
For our kernel width = 1, stride_w = stride_h = stride, pad_w = pad_h = padding, dilation_w = dilation_h = dilation.
But the formula for W_out:
W_out = (W + 2*pad_w - dilation_w*(K_w - 1) - 1) / stride_w + 1
With K_w = 1: W_out = (W + 2*pad_w - 1) / stride_w + 1.
Actually that's wrong: K_w = 1, so K_w - 1 = 0:
W_out = (W + 2*pad_w - 0 - 1) / stride_w + 1 = (W + 2*pad_w - 1) / stride_w + 1.
Wait, that's still not right. For kernel size 1, stride 1, padding 0:
W_out = W.
Using the formula: (W - 1) / 1 + 1 = W. Yes, that's correct.
For kernel size 1, stride 1, padding 1:
W_out = (W + 2 - 1) / 1 + 1 = W + 2. Correct.
So the formula is right.
But I need to be careful: the `Conv2d` module might store padding as a tuple. Let me access it correctly.
Actually, looking at PyTorch nn.Conv2d, `self.padding` is stored as a tuple of ints.
Now, for the CUDA code, `c10::optional<torch::Tensor>` is the right type for optional bias. Let me make sure the pybind11 binding works with this.
PyTorch's extension system should handle optional tensors. But in the function signature for pybind11, I might need to use `torch::optional<torch::Tensor>` or just accept a pointer.
Actually, to avoid issues with optional tensor binding, I can just pass the bias tensor directly, or pass it as a regular tensor and check if it's defined. But if there's no bias, I'd pass an undefined tensor.
Wait, in PyTorch C++ API, an undefined tensor can be created with `torch::Tensor()`. And checking with `.defined()`.
Let me simplify: in Python, if bias is None, I create an empty tensor or undefined tensor. But passing through pybind11...
Actually, `c10::optional<torch::Tensor>` should work fine with `torch.utils.cpp_extension` because PyTorch's bindings use this extensively.
Let me just test with the simplest approach.
Here's my complete kernel 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_runtime.h>
__global__ void depthwise_conv_k1_kernel(
const float* __restrict__ input,
const float* __restrict__ weight,
const float* __restrict__ bias,
float* __restrict__ output,
int B, int C, int H, int W,
int K, int stride, int pad_h, int pad_w, int dilation,
int H_out, int W_out
) {
int w = blockIdx.x * blockDim.x + threadIdx.x;
int h = blockIdx.y;
int bc = blockIdx.z;
int b = bc / C;
int c = bc % C;
if (w >= W_out) return;
int in_w = w * stride - pad_w;
int in_h_start = h * stride - pad_h;
float sum = 0.0f;
const float* wptr = weight + c * K;
#pragma unroll 4
for (int k = 0; k < K; k++) {
int in_h = in_h_start + k * dilation;
if (in_h >= 0 && in_h < H && in_w >= 0 && in_w < W) {
float val = input[(((b * C) + c) * H + in_h) * W + in_w];
sum += val * wptr[k];
}
}
if (bias != nullptr) {
sum += bias[c];
}
output[(((b * C) + c) * H_out + h) * W_out + w] = sum;
}
torch::Tensor depthwise_conv_k1_forward(
torch::Tensor input,
torch::Tensor weight,
c10::optional<torch::Tensor> bias,
int64_t stride, int64_t pad_h, int64_t pad_w, int64_t dilation
) {
int B = input.size(0);
int C = input.size(1);
int H = input.size(2);
int W = input.size(3);
int K = weight.size(2);
int H_out = (H + 2 * pad_h - dilation * (K - 1) - 1) / stride + 1;
int W_out = (W + 2 * pad_w - 1) / stride + 1;
auto output = torch::empty({B, C, H_out, W_out}, input.options());
const int threads = 256;
dim3 block(threads);
dim3 grid((W_out + threads - 1) / threads, H_out, B * C);
const float* bias_ptr = bias.has_value() ? bias.value().data_ptr<float>() : nullptr;
depthwise_conv_k1_kernel<<<grid, block>>>(
input.data_ptr<float>(),
weight.data_ptr<float>(),
bias_ptr,
output.data_ptr<float>(),
B, C, H, W,
K, stride, pad_h, pad_w, dilation,
H_out, W_out
);
return output;
}
"""
cpp_source = """
torch::Tensor depthwise_conv_k1_forward(
torch::Tensor input,
torch::Tensor weight,
c10::optional<torch::Tensor> bias,
int64_t stride, int64_t pad_h, int64_t pad_w, int64_t dilation
);
"""
depthwise_conv = load_inline(
name="depthwise_conv_k1",
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=["depthwise_conv_k1_forward"],
extra_cuda_cflags=["-O3", "--use_fast_math"],
verbose=False,
)
class ModelNew(nn.Module):
def __init__(self, in_channels, kernel_size, stride=1, padding=0, dilation=1, bias=False):
super(ModelNew, self).__init__()
self.conv2d = nn.Conv2d(in_channels, in_channels, kernel_size=(kernel_size, 1), stride=stride, padding=padding, dilation=dilation, groups=in_channels, bias=bias)
def forward(self, x):
return depthwise_conv.depthwise_conv_k1_forward(
x,
self.conv2d.weight,
self.conv2d.bias,
self.conv2d.stride[0],
self.conv2d.padding[0],
self.conv2d.padding[1] if len(self.conv2d.padding) > 1 else self.conv2d.padding[0],
self.conv2d.dilation[0]
)
```
Let me fix the padding access. For nn.Conv2d, `self.padding` is a tuple. If initialized with `padding=0`, it's `(0, 0)`. If initialized with `padding=(1, 0)`, it's `(1, 0)`. But since we only pass a single `padding` arg, it will be `(padding, padding)` or possibly `(padding, 0)` depending on PyTorch's handling.
Actually, with `kernel_size=(kernel_size, 1)`, padding might be interpreted differently. Let me check: if `padding=0` is passed, it becomes `(0, 0)` by default. But the user could also pass `padding=(pad_h, pad_w)`.
For safety, let me handle both cases:
```python
pad = self.conv2d.padding
if isinstance(pad, int):
pad_h = pad_w = pad
else:
pad_h, pad_w = pad[0], pad[1] if len(pad) > 1 else pad[0]
```
Same for stride and dilation.
Let me compile and test this.