reasoning
Now I understand the GPU specs. Let me analyze the problem:
The model is a depthwise separable convolution block with:
1. Pointwise Conv1d (256 -> 512 channels)
2. Depthwise Conv1d (512 channels, kernel_size=9)
3. GELU activation
4. Project back Conv1d (512 -> 256 channels)
5. Residual addition
Input shape: [32, 256, 4096]
After pointwise: [32, 512, 4096]
After depthwise: [32, 512, 4096]
After GELU: [32, 512, 4096]
After project_back: [32, 256, 4096]
Let me write a fused kernel that handles all operations. I'll need to:
1. Fuse pointwise convolution with the first part
2. Apply depthwise convolution
3. Apply GELU
4. Fuse project_back with residual addition
Actually, let me think about this more carefully. The operations are:
1. Pointwise Conv1d: input [B, C, L] -> [B, E, L] where C=256, E=512
- Weight shape: [E, C, 1]
- For each position in L: output[b, e, l] = sum_c(input[b, c, l] * weight[e, c, 0]) + bias[e]
2. Depthwise Conv1d: [B, E, L] -> [B, E, L]
- Weight shape: [E, 1, K] where K=9
- For each channel e: output[b, e, l] = sum_k(input[b, e, l + k - K/2] * weight[e, 0, k]) + bias[e]
- This is essentially E independent 1D convolutions
3. GELU: element-wise activation
4. Project back: [B, E, L] -> [B, C, L]
- Weight shape: [C, E, 1]
- output[b, c, l] = sum_e(input[b, e, l] * weight[c, e, 0]) + bias[c]
5. Residual: output[b, c, l] = output[b, c, l] + residual[b, c, l]
Let me think about how to fuse these efficiently. One approach is to:
- Have each thread handle one output element (B, C, L)
- Compute the pointwise + depthwise + GELU + project_back in a fused manner
But this might require a lot of registers since we need to store intermediate values for each expanded channel.
Actually, let me think about a different approach. We can:
1. Compute the pointwise convolution and store it in shared memory or registers
2. Apply depthwise convolution on the stored values
3. Apply GELU
4. Compute project_back and add residual
Let me design a kernel where each thread block handles one (batch, output_channel) pair and processes the entire sequence length. This way, we can use shared memory for intermediate storage.
Actually, for better parallelism, let me consider having each thread block handle one (batch, position) pair and process all channels. This gives us B * L = 32 * 4096 = 131,072 thread blocks, which is a good number for the H100 with 132 SMs.
Within each block, we need to:
1. Compute pointwise: for each expanded channel, sum over input channels
- We need 256 threads minimum (one per input channel) or use reduction
2. Apply depthwise: for each expanded channel, convolve with kernel
- Need to load 9 values per expanded channel
3. Apply GELU: element-wise
4. Compute project_back: for each output channel, sum over expanded channels
- Need to reduce over 512 expanded channels
Let me think about memory access patterns. The input is [B, C, L] and we're reading contiguous L values for each (B, C) pair. If we have one block per (B, L), we'd read 256 contiguous values from global memory, which is efficient.
For the pointwise convolution:
- Each thread reads one input channel value and multiplies by weight
- Then we reduce to get each expanded channel output
Actually, let me try a simpler approach first: have each thread compute one output element (B, C, L). This might not be optimal but is straightforward.
For thread (b, c, l):
1. Compute pointwise output for all 512 expanded channels at position l
- For each e: sum over 256 input channels
- This is 256 * 512 = 131,072 operations per thread
- That's a lot! We need to parallelize this.
Let me reconsider. A better approach:
Have each warp handle one output element (B, C, L). Within the warp:
- Each thread handles 256/32 = 8 input channels for pointwise
- Each thread handles 512/32 = 16 expanded channels for the intermediate computation
- Then reduce across warp for final output
Actually, let me try a different strategy. Let me use shared memory to store intermediate results:
Block configuration: each block handles one (B, L) position
- 512 threads (one per expanded channel)
- Each thread:
1. Computes one expanded channel output from pointwise (reduce over 256 input channels)
2. Computes depthwise convolution (needs 9 values from neighbors)
3. Applies GELU
4. Reduces over expanded channels for each output channel contribution
But wait, step 4 requires reducing over expanded channels to get 256 output channels values. That's 512 threads reducing to 256 values.
Actually, let me think about this more carefully. Let me try a tiled approach:
Block: handles multiple positions in L dimension
- Each thread handles one (B, expanded_channel) pair for multiple positions
Hmm, this is getting complex. Let me start with a simpler but still efficient approach:
1. Kernel 1: Pointwise + Depthwise + GELU
- Each block handles one (B, position) for all expanded channels
- Load input values for this position and neighbors into shared memory
- Each thread computes one expanded channel output
- Store result in global memory
2. Kernel 2: Project back + Residual
- Each block handles one (B, position) for all output channels
- Load intermediate values for this position
- Each thread computes one output channel
- Add residual and store
But two kernels means extra global memory traffic. Let me try to fuse everything.
Actually, let me reconsider the problem size:
- B = 32, C = 256, E = 512, L = 4096, K = 9
Total operations:
- Pointwise: B * E * L * C = 32 * 512 * 4096 * 256 = 17.2 billion FMA
- Depthwise: B * E * L * K = 32 * 512 * 4096 * 9 = 604 million FMA
- Project back: B * C * L * E = 32 * 256 * 4096 * 512 = 17.2 billion FMA
Total: ~35 billion FMA operations
Memory:
- Input: B * C * L * 4 bytes = 32 * 256 * 4096 * 4 = 134 MB
- Weights: pointwise (512 * 256 * 4 = 0.5 MB), depthwise (512 * 9 * 4 = 18 KB), project_back (256 * 512 * 4 = 0.5 MB)
- Intermediate: B * E * L * 4 = 32 * 512 * 4096 * 4 = 268 MB
If we fuse everything, we can avoid storing the intermediate 268 MB.
Let me design a fused kernel where each block handles one (B, L_out) pair:
```cuda
// Thread block: one per (B, L_out)
// Threads: 256 (one per output channel C)
// Shared memory:
// - Input: (K + 1) positions * C channels = 10 * 256 * 4 = 10 KB
// - Intermediate expanded: E * 4 = 2048 bytes
// But wait, we need to compute E values for pointwise and depthwise
// Let me try a different approach:
// Thread block: one per (B, L_out)
// Threads: 512 (one per expanded channel E)
// Shared memory:
// - Input: (K + 1) positions * C channels for depthwise? No, that's for expanded channels
// - Actually for depthwise, we need E * K neighboring values in the expanded channel space
```
Hmm, this is tricky. Let me think about it differently.
For depthwise convolution, we need neighboring values in the L dimension of the expanded channel output. So we need to compute pointwise for multiple positions first.
Let me try:
- Thread block handles one (B, output_channel) for all L positions
- We compute the entire sequence for that output channel
This approach:
1. For each L position, compute pointwise output (sum over 256 input channels) for all 512 expanded channels
2. Apply depthwise convolution (need 9 values from neighboring positions, which we have since we're processing all L)
3. Apply GELU
4. Compute project_back output (sum over 512 expanded channels)
5. Add residual
But this requires a lot of storage for intermediate expanded channel values across L.
Actually, let me try a simpler approach using tiling:
Block: handles one tile of (B, L_tile_start) for all channels
- Process a tile of L positions (say 64 positions)
- Use shared memory for input and weights
- Compute all operations in a fused manner
Let me implement this:
```cuda
// Block: handles tile_size positions in L dimension for one batch
// Threads: 256 (or some other convenient number)
// Each thread computes multiple output elements
// For efficiency, let me try:
// Block: handles one (B, L) position
// Threads: 512
// Each thread handles one expanded channel
// Step 1: Pointwise
// Each thread computes one expanded channel output
// Needs reduction over 256 input channels
// If 512 threads, each thread handles 256/512 = 0.5 channels? That doesn't work.
// Alternative: 256 threads
// Each thread handles one output channel in the end
// For pointwise, each thread computes partial sums for multiple expanded channels
// For project_back, each thread reduces over expanded channels
```
Let me try a more pragmatic approach. I'll use a hybrid kernel:
1. Each thread block handles one (B, L) position
2. Use 256 threads
3. Each thread handles one output channel C
The computation:
1. For pointwise + depthwise, we need to compute all 512 expanded channel values at this L position and neighboring positions (for depthwise)
2. Then apply GELU
3. Then reduce over expanded channels for the project_back
4. Add residual
The challenge is step 1: we need to compute 512 * (K = 9) values for depthwise. That's 4608 values.
If we have 256 threads, each thread would compute 512/256 = 2 expanded channels, but for all K positions, that's 18 values per thread.
Let me try this:
- 256 threads per block
- Each thread computes contributions from expanded channels 0, 1, ..., E-1 for one output channel
- But computing one expanded channel requires:
1. Load input at this L position for all C channels (shared memory)
2. Compute pointwise output
3. Load neighboring values (from global memory or recomputed)
4. Compute depthwise output
5. Apply GELU
This is expensive. Let me try reordering:
1. Load all input values for this L position and neighbors (K positions) into shared memory
- Shared memory: K * C = 9 * 256 = 2304 floats = 9 KB
2. Compute pointwise output for all expanded channels at this L position
- Each thread computes E/256 = 2 expanded channel outputs
- Needs K * E = 9 * 512 = 4608 values for depthwise
- But we only have pointwise output at this L position, not neighbors
Hmm, the issue is that depthwise needs neighboring positions' pointwise outputs.
Let me try a different approach:
1. Thread block handles multiple positions in L dimension
2. Compute pointwise for all positions in the tile, store in shared memory
3. Compute depthwise for all positions in the tile
4. Compute project_back for all positions in the tile, add residual
Tile size: let's say 32 positions
- Shared memory for pointwise output: 32 * 512 = 16,384 floats = 64 KB
- Shared memory for input: 32 * 256 = 8192 floats = 32 KB
- Total: 96 KB, fits within 227 KB per block
Actually, we need input for depthwise padding too. Let's account for that:
- Input for tile + padding: (32 + 8) * 256 = 40 * 256 = 10,240 floats = 40 KB
- Pointwise output: 32 * 512 = 16,384 floats = 64 KB
- Total: 104 KB, still fits
Threads per block: 512 (one per expanded channel per position)
Wait, let me reconsider. We have:
- 32 positions in tile
- 512 expanded channels
- 256 output channels
If we have 512 threads and process one position at a time:
- Each thread handles one expanded channel
- Compute pointwise: reduce over 256 input channels
- This needs parallel reduction or each thread reads all 256 inputs
- Compute depthwise: need neighboring expanded outputs
- We can compute these since we have all positions in tile
- Compute project_back: reduce over 512 expanded channels
- Use warp reduction
Let me try this approach:
```
Block: handles one batch and a tile of L positions
Threads: 512
Shared memory:
- Input: (tile_size + padding) * C floats
- Expanded: tile_size * E floats (pointwise output, also depthwise input)
- Depthwise output: tile_size * E floats (reuse same memory after depthwise)
- Weights: load into registers or constant memory
Each thread:
1. Loop over positions in tile:
a. Load input for position into shared memory (cooperative)
b. Compute pointwise for one expanded channel (need to reduce over C)
c. Store in expanded shared memory
2. Sync
3. Loop over positions in tile:
a. Load expanded values for position and neighbors from shared memory
b. Compute depthwise
c. Apply GELU
d. Store back to shared memory (reuse space)
4. Sync
5. Loop over positions in tile:
a. Load expanded values (depthwise + GELU output) from shared memory
b. Compute project_back contribution to all output channels
c. Warp reduce to get final output
d. Load residual and add
e. Store to global memory
```
Actually, step 1b is problematic: computing pointwise for one expanded channel requires reducing over 256 input channels, but we have 512 threads. We can either:
- Have each thread read all 256 inputs (256 * 4 = 1024 bytes per thread, might be slow)
- Use parallel reduction (need 256 threads)
Let me try 256 threads instead:
```
Block: handles one batch and a tile of L positions
Threads: 256
Each thread: handles one output channel
Shared memory:
- Input: (tile_size + padding) * C floats
- Expanded output (pointwise + depthwise + GELU): tile_size * E floats
Computation:
1. Load input for all positions in tile into shared memory (cooperative)
2. For each position in tile:
a. Compute pointwise for all expanded channels
- This is E reductions over C channels
- We have 256 threads, can compute E values with each thread computing E/256 contributions
b. Store in expanded shared memory
3. Sync
4. For each position in tile:
a. Compute depthwise for all expanded channels
- Each thread handles E/256 expanded channels
- Load neighbors from shared memory (pointwise output)
- Compute depthwise, apply GELU
- Store back
5. Sync
6. For each position in tile:
a. Compute project_back for one output channel
- Reduce over E expanded channels
- Use warp shuffle for reduction
b. Add residual
c. Store to global memory
```
Let me think about the reduction patterns:
- Pointwise: E outputs, each needs reduction over C = 256 inputs
- We have 256 threads
- Each thread can compute one partial sum for E/C expanded channels
- Then reduce across threads
Actually, let me simplify. Let's just have each thread compute its part of the reduction directly without complex parallel reduction patterns.
For 256 threads and E = 512 expanded channels:
- Each thread handles 2 expanded channels for pointwise output
- Thread i handles expanded channels 2*i and 2*i+1
For pointwise:
- Thread i reads all C = 256 input values
- Computes partial sum for expanded channels 2*i and 2*i+1
- This requires reading 256 floats = 1 KB, which is reasonable
For depthwise:
- Thread i computes depthwise output for expanded channels 2*i and 2*i+1
- Needs to read K = 9 neighboring values from shared memory
- Total reads: 2 * 9 = 18 floats = 72 bytes
For project_back:
- Thread i computes output channel i
- Needs to reduce over E = 512 expanded channels
- Use warp shuffle reduction
Let me refine the kernel design:
```cuda
#define TILE_SIZE 64
#define C 256
#define E 512
#define K 9
#define PADDING (K / 2) // 4
// Block dimensions: 256 threads
// Grid dimensions: (B, (L + TILE_SIZE - 1) / TILE_SIZE)
__global__ void fused_conv_block(
const float* __restrict__ input, // [B, C, L]
const float* __restrict__ pw_weight, // [E, C]
const float* __restrict__ pw_bias, // [E]
const float* __restrict__ dw_weight, // [E, K]
const float* __restrict__ dw_bias, // [E]
const float* __restrict__ pb_weight, // [C, E]
const float* __restrict__ pb_bias, // [C]
float* __restrict__ output, // [B, C, L]
int B, int L
) {
int batch = blockIdx.x;
int tile_start = blockIdx.y * TILE_SIZE;
int tid = threadIdx.x;
// Shared memory
__shared__ float input_shared[(TILE_SIZE + 2*PADDING) * C];
__shared__ float expanded_shared[TILE_SIZE * E];
// Load input into shared memory
// Each thread loads multiple elements
int input_start = tile_start - PADDING;
for (int i = tid; i < (TILE_SIZE + 2*PADDING) * C; i += blockDim.x) {
int pos = i / C;
int ch = i % C;
int global_pos = input_start + pos;
float val = 0.0f;
if (global_pos >= 0 && global_pos < L) {
val = input[batch * C * L + ch * L + global_pos];
}
input_shared[i] = val;
}
__syncthreads();
// Compute pointwise for each position in tile
// Each thread handles 2 expanded channels
for (int pos = 0; pos < TILE_SIZE; pos++) {
int e_base = tid * 2;
for (int e_offset = 0; e_offset < 2; e_offset++) {
int e = e_base + e_offset;
float sum = pw_bias[e];
for (int c = 0; c < C; c++) {
float in_val = input_shared[(pos + PADDING) * C + c];
sum += in_val * pw_weight[e * C + c];
}
expanded_shared[pos * E + e] = sum;
}
}
__syncthreads();
// Compute depthwise + GELU for each position
for (int pos = 0; pos < TILE_SIZE; pos++) {
int e_base = tid * 2;
for (int e_offset = 0; e_offset < 2; e_offset++) {
int e = e_base + e_offset;
float sum = dw_bias[e];
for (int k = 0; k < K; k++) {
int neighbor_pos = pos + k - PADDING; // PADDING = K/2 = 4
float in_val;
if (neighbor_pos >= 0 && neighbor_pos < TILE_SIZE) {
in_val = expanded_shared[neighbor_pos * E + e];
} else {
// Need to handle boundary - this shouldn't happen if we computed expanded_shared correctly
// Actually, we need to compute expanded values for padding positions too!
}
sum += in_val * dw_weight[e * K + k];
}
// GELU approximation
float gelu = 0.5f * sum * (1.0f + tanhf(0.7978845608f * (sum + 0.044715f * sum * sum * sum)));
expanded_shared[pos * E + e] = gelu;
}
}
__syncthreads();
// Compute project_back + residual + store
// Each thread handles one output channel
int c = tid;
for (int pos = 0; pos < TILE_SIZE; pos++) {
int global_pos = tile_start + pos;
if (global_pos >= L) continue;
float sum = pb_bias[c];
for (int e = 0; e < E; e++) {
sum += expanded_shared[pos * E + e] * pb_weight[c * E + e];
}
// Add residual
sum += input[batch * C * L + c * L + global_pos];
output[batch * C * L + c * L + global_pos] = sum;
}
}
```
Wait, I realize there's an issue with depthwise: we need expanded values for positions outside the tile (for the convolution kernel). We have two options:
1. Compute expanded values for padding positions too
2. Handle boundary by computing on-the-fly or using different logic
Let me compute expanded values for the padding region as well:
Actually, I think we should include the padding positions in the expanded_shared computation. Let me adjust:
```cuda
// Shared memory
__shared__ float input_shared[(TILE_SIZE + 2*PADDING) * C]; // +padding on both sides
__shared__ float expanded_shared[(TILE_SIZE + 2*PADDING) * E]; // also need padding for depthwise
// Then compute expanded for all (TILE_SIZE + 2*PADDING) positions
// But then for depthwise, we only compute for TILE_SIZE positions (the middle ones)
```
This increases shared memory:
- Input shared: (64 + 8) * 256 = 18,432 floats = 72 KB
- Expanded shared: (64 + 8) * 512 = 36,864 floats = 144 KB
- Total: 216 KB, still fits in 227 KB
Let me rewrite the kernel with this approach.
Actually wait, I realize there's another issue. The pointwise convolution uses weights of shape [E, C], and we're reading these for every position. That's a lot of weight reads. Let me think about how to optimize this.
Option 1: Load weights into shared memory
- pw_weight: E * C = 512 * 256 = 131,072 floats = 512 KB
- Too large for shared memory!
Option 2: Use __ldg to cache weights in read-only cache
- This should help with repeated reads
Option 3: Have each thread compute its part of the weights and reuse across positions
- Since each thread handles 2 expanded channels, it only needs 2 * C = 512 weight values
- These can be loaded into registers
Let me try option 3:
```cuda
// Load weight values into registers
// Each thread needs weights for 2 expanded channels, each with C values
// 2 * C = 512 floats = 2048 bytes per thread
// Total registers: 256 threads * 512 floats = way too many!
// Actually, we can't load all weights into registers
// Let's use __ldg for read-only cache
```
Actually, I think using __ldg (or const __restrict__) should work well since the weights are read-only and the cache should handle repeated accesses.
Let me also think about the project_back reduction. We need to reduce over E = 512 expanded channels. With 256 threads:
- Each thread computes one output channel
- Needs to sum over 512 expanded channels
- That's 512 reads and additions per output position
We could use warp shuffles to parallelize this, but with 256 threads and only 256 outputs, each thread already handles one output. So the reduction isn't parallelized.
Let me think about an alternative: what if we have more threads?
If we have 512 threads:
- 2 threads per output channel
- Each thread sums over 256 expanded channels
- Then reduce the pair to get final output
Let me try this:
```cuda
// 512 threads per block
// Thread pair (2*tid, 2*tid+1) computes output channel tid
// For project_back:
// Thread 2*tid: sum over e = 0, 2, 4, ..., 510
// Thread 2*tid+1: sum over e = 1, 3, 5, ..., 511
// Then reduce pair using __shfl_down_sync
```
But wait, this makes the pointwise and depthwise computation different. Let me think...
Actually, let me simplify and just use 256 threads. The project_back reduction over 512 values per thread isn't ideal but might be acceptable.
Let me also think about memory access patterns:
- Input: stored as [B, C, L], so for a given (B, L), we access C contiguous values
- Weight: stored as [E, C] or [C, E], need to check PyTorch conventions
- Output: [B, C, L]
For the pointwise weight, PyTorch Conv1d with kernel_size=1 stores weight as [out_channels, in_channels, 1], which is [E, C, 1].
Let me write the complete kernel now.
Hmm actually, I realize that my analysis of the depthwise convolution is incomplete. Depthwise Conv1d with groups=expanded_channels means each expanded channel has its own kernel. The weight shape is [expanded_channels, 1, kernel_size] = [E, 1, K].
For each expanded channel e:
output[b, e, l] = sum_k(input[b, e, l + k - K/2] * weight[e, 0, k]) + bias[e]
So I need to load weight[e * K + k] for thread handling expanded channel e.
Let me also reconsider the problem constraints. Looking at the reference code:
```python
self.depthwise = nn.Conv1d(expanded_channels, expanded_channels, kernel_size,
padding=kernel_size // 2, groups=expanded_channels)
```
This uses padding = kernel_size // 2 = 4, which maintains the sequence length.
OK let me write a complete kernel. I'll start with a simpler version and optimize if needed.
Actually, let me reconsider the tile size. With TILE_SIZE = 64:
- Input shared: (64 + 8) * 256 * 4 = 72 KB
- Expanded shared: (64 + 8) * 512 * 4 = 144 KB
- Total: 216 KB
This is quite close to the 227 KB limit. Let me use TILE_SIZE = 32:
- Input shared: (32 + 8) * 256 * 4 = 40 KB
- Expanded shared: (32 + 8) * 512 * 4 = 80 KB
- Total: 120 KB
This is more comfortable and allows more blocks to run concurrently.
But with smaller tiles, we have more blocks and potentially more overhead. Let me try TILE_SIZE = 64 first and see if it fits.
Let me write the kernel code:
call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = \"\"\"\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n#define TILE_SIZE 32\n#define C 256\n#define E 512\n#define K 9\n#define PADDING 4\n\n__device__ __forceinline__ float gelu_activation(float x) {\n // GELU approximation: 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))\n const float SQRT_2_OVER_PI = 0.7978845608028654f;\n float x3 = x * x * x;\n return 0.5f * x * (1.0f + tanhf(SQRT_2_OVER_PI * (x + 0.044715f * x3)));\n}\n\n__global__ void fused_depthwise_block_kernel(\n const float* __restrict__ input,\n const float* __restrict__ pw_weight,\n const float* __restrict__ pw_bias,\n const float* __restrict__ dw_weight,\n const float* __restrict__ dw_bias,\n const float* __restrict__ pb_weight,\n const float* __restrict__ pb_bias,\n float* __restrict__ output,\n int L\n) {\n int batch = blockIdx.x;\n int tile_start = blockIdx.y * TILE_SIZE;\n int tid = threadIdx.x;\n \n // Shared memory for input (including padding for depthwise)\n __shared__ float input_shm[(TILE_SIZE + 2*PADDING) * C];\n // Shared memory for expanded channels (including padding)\n __shared__ float expanded_shm[(TILE_SIZE + 2*PADDING) * E];\n \n // Load input into shared memory\n int input_start = tile_start - PADDING;\n for (int i = tid; i < (TILE_SIZE + 2*PADDING) * C; i += blockDim.x) {\n int pos = i / C;\n int ch = i % C;\n int global_pos = input_start + pos;\n float val = 0.0f;\n if (global_pos >= 0 && global_pos < L) {\n val = input[batch * C * L + ch * L + global_pos];\n }\n input_shm[i] = val;\n }\n __syncthreads();\n \n // Phase 1: Pointwise convolution\n // Each thread computes 2 expanded channels\n int e_base = tid * 2;\n for (int pos = 0; pos < TILE_SIZE + 2*PADDING; pos++) {\n for (int e_off = 0; e_off < 2; e_off++) {\n int e = e_base + e_off;\n if (e >= E) continue;\n \n float sum = pw_bias[e];\n for (int c = 0; c < C; c++) {\n sum += input_shm[pos * C + c] * pw_weight[e * C + c];\n }\n expanded_shm[pos * E + e] = sum;\n }\n }\n __syncthreads();\n \n // Phase 2: Depthwise convolution + GELU\n // Only compute for positions within the tile (not padding)\n for (int pos = 0; pos < TILE_SIZE; pos++) {\n for (int e_off = 0; e_off < 2; e_off++) {\n int e = e_base + e_off;\n if (e >= E) continue;\n \n int shm_pos = pos + PADDING; // position in shared memory (with padding offset)\n float sum = dw_bias[e];\n for (int k = 0; k < K; k++) {\n int neighbor_pos = shm_pos + k - PADDING;\n sum += expanded_shm[neighbor_pos * E + e] * dw_weight[e * K + k];\n }\n // Apply GELU and store back\n expanded_shm[shm_pos * E + e] = gelu_activation(sum);\n }\n }\n __syncthreads();\n \n // Phase 3: Project back + residual + store\n // Each thread handles one output channel\n int c = tid;\n for (int pos = 0; pos < TILE_SIZE; pos++) {\n int global_pos = tile_start + pos;\n if (global_pos >= L) continue;\n \n int shm_pos = pos + PADDING;\n float sum = pb_bias[c];\n for (int e = 0; e < E; e++) {\n sum += expanded_shm[shm_pos * E + e] * pb_weight[c * E + e];\n }\n // Add residual from original input\n sum += input_shm[shm_pos * C + c];\n output[batch * C * L + c * L + global_pos] = sum;\n }\n}\n\ntorch::Tensor fused_depthwise_block_forward(\n torch::Tensor input,\n torch::Tensor pw_weight, torch::Tensor pw_bias,\n torch::Tensor dw_weight, torch::Tensor dw_bias,\n torch::Tensor pb_weight, torch::Tensor pb_bias\n) {\n int B = input.size(0);\n int L = input.size(2);\n \n auto output = torch::empty({B, C, L}, input.options());\n \n dim3 grid(B, (L + TILE_SIZE - 1) / TILE_SIZE);\n dim3 block(256);\n \n fused_depthwise_block_kernel<<<grid, block>>>(\n input.data_ptr<float>(),\n pw_weight.data_ptr<float>(),\n pw_bias.data_ptr<float>(),\n dw_weight.data_ptr<float>(),\n dw_bias.data_ptr<float>(),\n pb_weight.data_ptr<float>(),\n pb_bias.data_ptr<float>(),\n output.data_ptr<float>(),\n L\n );\n \n return output;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor fused_depthwise_block_forward(\n torch::Tensor input,\n torch::Tensor pw_weight, torch::Tensor pw_bias,\n torch::Tensor dw_weight, torch::Tensor dw_bias,\n torch::Tensor pb_weight, torch::Tensor pb_bias\n);\n\"\"\"\n\nmodule = load_inline(\n name=\"fused_depthwise_block\",\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=[\"fused_depthwise_block_forward\"],\n verbose=True\n)\n\nclass ModelNew(nn.Module):\n def __init__(self, channels, expanded_channels, kernel_size):\n super(ModelNew, self).__init__()\n # Store weights as nn.Parameter for proper initialization\n self.pointwise = nn.Conv1d(channels, expanded_channels, 1)\n self.depthwise = nn.Conv1d(expanded_channels, expanded_channels, kernel_size,\n padding=kernel_size // 2, groups=expanded_channels)\n self.activation = nn.GELU()\n self.project_back = nn.Conv1d(expanded_channels, channels, 1)\n \n self.fused_forward = module.fused_depthwise_block_forward\n \n def forward(self, x):\n # Extract weights - pointwise weight is [E, C, 1], squeeze to [E, C]\n pw_weight = self.pointwise.weight.squeeze(-1)\n pw_bias = self.pointwise.bias\n \n # Depthwise weight is [E, 1, K], squeeze to [E, K]\n dw_weight = self.depthwise.weight.squeeze(1)\n dw_bias = self.depthwise.bias\n \n # Project back weight is [C, E, 1], squeeze to [C, E]\n pb_weight = self.project_back.weight.squeeze(-1)\n pb_bias = self.project_back.bias\n \n return self.fused_forward(\n x, pw_weight, pw_bias,\n dw_weight, dw_bias,\n pb_weight, pb_bias\n )"
}