90_Conv3d_LeakyReLU_Sum_Clamp_GELU • multidim_indexed_kernel_baseimport torch
import torch.nn as nn
import torch.nn.functional as F
def module_fn(
x: torch.Tensor,
conv_weight: torch.Tensor,
conv_bias: torch.Tensor,
sum_tensor: torch.Tensor,
) -> torch.Tensor:
"""
Applies 3D convolution, LeakyReLU, tensor addition, clamping and GELU activation.
Args:
x (torch.Tensor): Input tensor of shape (batch_size, in_channels, depth, height, width)
conv_weight (torch.Tensor): 3D convolution weight tensor of shape
(out_channels, in_channels, kernel_size, kernel_size, kernel_size)
conv_bias (torch.Tensor): Bias tensor for 3D convolution of shape (out_channels)
sum_tensor (torch.Tensor): Tensor to add of shape (out_channels, 1, 1, 1)
Returns:
torch.Tensor: Output tensor after applying convolution, LeakyReLU, addition,
clamping and GELU activation
"""
x = F.conv3d(x, conv_weight, bias=conv_bias)
x = F.leaky_relu(x, negative_slope=0.2)
x = x + sum_tensor
x = torch.clamp(x, min=-1.0, max=1.0)
x = F.gelu(x)
return x
class Model(nn.Module):
"""
Model that performs a 3D convolution, applies LeakyReLU, sums with a tensor, clamps, and applies GELU activation.
"""
def __init__(self, in_channels, out_channels, kernel_size, sum_tensor_shape):
super(Model, self).__init__()
conv = nn.Conv3d(in_channels, out_channels, kernel_size)
self.conv_weight = conv.weight
self.conv_bias = conv.bias
self.sum_tensor = nn.Parameter(torch.randn(sum_tensor_shape) * 0.02)
def forward(self, x, fn=module_fn):
return fn(x, self.conv_weight, self.conv_bias, self.sum_tensor)
batch_size = 128
in_channels = 3
out_channels = 16
depth, height, width = 16, 32, 32
kernel_size = 3
sum_tensor_shape = (out_channels, 1, 1, 1)
def get_inputs():
return [torch.randn(batch_size, in_channels, depth, height, width)]
def get_init_inputs():
return [in_channels, out_channels, kernel_size, sum_tensor_shape]
import torch
import torch.nn as nn
class Model(nn.Module):
"""
Model that performs a 3D convolution, applies LeakyReLU, sums with a tensor, clamps, and applies GELU activation.
"""
def __init__(self, in_channels, out_channels, kernel_size, sum_tensor_shape):
super(Model, self).__init__()
self.conv = nn.Conv3d(in_channels, out_channels, kernel_size)
self.sum_tensor = nn.Parameter(torch.randn(sum_tensor_shape)*0.02)
def forward(self, x):
x = self.conv(x)
x = torch.nn.functional.leaky_relu(x, negative_slope=0.2)
x = x + self.sum_tensor
x = torch.clamp(x, min=-1.0, max=1.0)
x = torch.nn.functional.gelu(x)
return x
batch_size = 128
in_channels = 3
out_channels = 16
depth, height, width = 16, 32, 32
kernel_size = 3
sum_tensor_shape = (out_channels, 1, 1, 1)
def get_inputs():
return [torch.randn(batch_size, in_channels, depth, height, width)]
def get_init_inputs():
return [in_channels, out_channels, kernel_size, sum_tensor_shape]#include <torch/extension.h>
#include <ATen/ATen.h>
#include <vector>
#include <cmath>
// Inline device function for LeakyReLU
__device__ inline float leaky_relu(float x, float alpha = 0.2f) {
return (x > 0.0f) ? x : alpha * x;
}
// Inline device function for clamping
__device__ inline float clamp_val(float x, float min_val = -1.0f, float max_val = 1.0f) {
return fmaxf(fminf(x, max_val), min_val);
}
// Inline device function for GELU activation using tanh approximation
__device__ inline float gelu(float x) {
float tanh_out = tanhf(0.7978845608f * (x + 0.044715f * x * x * x));
return x * 0.5f * (1.0f + tanh_out);
}
// CUDA kernel using 3D grid mapping for spatial and (batch,channel,depth) dimensions
__global__ void multi_dim_kernel(
const float* __restrict__ input,
const float* __restrict__ sum_tensor,
float* __restrict__ output,
int B,
int C,
int D,
int H,
int W) {
// Map spatial dimensions (width, height) using blockIdx.x and blockIdx.y
int w = blockIdx.x * blockDim.x + threadIdx.x;
int h = blockIdx.y * blockDim.y + threadIdx.y;
// Use gridDim.z to cover combined (B, C, D):
// Let each grid.z index represent a unique combination of (b, c, d)
int z_idx = blockIdx.z; // runs from 0 to (B * C * D - 1)
int d = z_idx % D; // depth index
int bc = z_idx / D; // combined batch and channel index
int b = bc / C; // batch index
int c = bc % C; // channel index
// Check bounds for spatial dimensions
if (w < W && h < H) {
// Compute linear index for contiguous tensor in [B, C, D, H, W] layout
int idx = (((b * C + c) * D + d) * H + h) * W + w;
float x = input[idx];
float y = leaky_relu(x);
y += sum_tensor[c];
y = clamp_val(y);
y = gelu(y);
output[idx] = y;
}
}
// Kernel launcher that sets up a 3D grid mapping the tensor dimensions
void multi_dim_kernel_launcher(torch::Tensor &x, torch::Tensor &sum_tensor) {
int B = x.size(0);
int C = x.size(1);
int D = x.size(2);
int H = x.size(3);
int W = x.size(4);
// Choose block dimensions for the spatial axes
const int BLOCK_W = 16;
const int BLOCK_H = 16;
dim3 block(BLOCK_W, BLOCK_H, 1);
// Grid dimensions:
// - grid.x covers width
// - grid.y covers height
// - grid.z covers the combined (B, C, D) dimensions
int grid_x = (W + BLOCK_W - 1) / BLOCK_W;
int grid_y = (H + BLOCK_H - 1) / BLOCK_H;
int grid_z = B * C * D;
dim3 grid(grid_x, grid_y, grid_z);
multi_dim_kernel<<<grid, block>>>(
x.data_ptr<float>(),
sum_tensor.data_ptr<float>(),
x.data_ptr<float>(),
B, C, D, H, W);
cudaDeviceSynchronize();
}
// Forward function: performs 3D convolution then applies our elementwise CUDA kernel
torch::Tensor forward(
torch::Tensor x,
torch::Tensor conv_weight,
torch::Tensor conv_bias,
torch::Tensor sum_tensor) {
TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
TORCH_CHECK(conv_weight.is_cuda(), "conv_weight must be a CUDA tensor");
TORCH_CHECK(conv_bias.is_cuda(), "conv_bias must be a CUDA tensor");
TORCH_CHECK(sum_tensor.is_cuda(), "sum_tensor must be a CUDA tensor");
TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
// Perform 3D convolution
auto x_conv = at::conv3d(x, conv_weight, conv_bias);
// Ensure the tensor is contiguous
auto output = x_conv.contiguous();
// Launch the CUDA kernel with multi-dimensional grid and block indexing
multi_dim_kernel_launcher(output, sum_tensor);
return output;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("forward", &forward, "Multi-dim indexed Conv3d LeakyReLU Sum Clamp GELU forward (CUDA)");
}
| Metric | Value | Unit | Variance | Samples |
|---|---|---|---|---|
| Executed Ipc Active | 3.140 | inst/cycle | 0.000 | 5 |
| Executed Ipc Elapsed | 3.090 | inst/cycle | 0.000 | 5 |
| Issue Slots Busy | 78.502 | % | 0.006 | 5 |
| Issued Ipc Active | 3.140 | inst/cycle | 0.000 | 5 |
| SM Busy | 78.502 | % | 0.006 | 5 |
| Memory Throughput | 1214948389363.868 | byte/second | 3520433275614237696.000 | 5 |
| Mem Busy | 29.408 | % | 0.003 | 5 |
| Max Bandwidth | 38.926 | % | 0.004 | 5 |
| L1/TEX Hit Rate | 53.216 | % | 0.000 | 5 |
| L2 Hit Rate | 68.916 | % | 0.010 | 5 |
| Mem Pipes Busy | 49.738 | % | 0.007 | 5 |
| Warp Cycles Per Issued Instruction | 15.736 | cycle | 0.000 | 5 |
| Warp Cycles Per Executed Instruction | 15.738 | cycle | 0.001 | 5 |
| Avg. Active Threads Per Warp | 27.670 | 0.000 | 5 | |
| Avg. Not Predicated Off Threads Per Warp | 24.960 | 0.000 | 5 | |
| Max Active Clusters | 0.000 | cluster | 0.000 | 5 |
| Max Cluster Size | 8.000 | block | 0.000 | 5 |
| Overall GPU Occupancy | 0.000 | % | 0.000 | 5 |
| Cluster Occupancy | 0.000 | % | 0.000 | 5 |
| Block Limit SM | 32.000 | block | 0.000 | 5 |
| Block Limit Registers | 16.000 | block | 0.000 | 5 |
| Block Limit Shared Mem | 32.000 | block | 0.000 | 5 |
| Block Limit Warps | 8.000 | block | 0.000 | 5 |
| Theoretical Active Warps per SM | 64.000 | warp | 0.000 | 5 |
| Theoretical Occupancy | 100.000 | % | 0.000 | 5 |
| Achieved Occupancy | 78.120 | % | 0.000 | 5 |
| Achieved Active Warps Per SM | 49.996 | warp | 0.000 | 5 |
| Rule | Description |
|---|---|
| INF HighPipeUtilization | ALU is the highest-utilized pipeline (54.3%) based on active cycles, taking into account the rates of its different instructions. It executes integer and logic operations. It is well-utilized, but should not be a bottleneck. |
| INF CPIStall | Check the Warp Stall Sampling (All Cycles) table for the top stall locations in your source based on sampling data. The Kernel Profiling Guide (https://docs.nvidia.com/nsight-compute/ProfilingGuide/index.html#metrics-reference) provides more details on each stall reason. |
| WRN Occupancy | This kernel's theoretical occupancy is not impacted by any block limit. The difference between calculated theoretical (100.0%) and measured achieved occupancy (78.1%) can be the result of warp scheduling overheads or workload imbalances during the kernel execution. Load imbalances can occur between warps within a block as well as across blocks of the same kernel. See the CUDA Best Practices Guide (https://docs.nvidia.com/cuda/cuda-c-best-practices-guide/index.html#occupancy) for more details on optimizing occupancy. |
| Operation / Metric | Value | Unit |
|---|---|---|
| aten::to | ||
| CPU Time | 470000.28 | μs |
| Device Time | 2732.40 | μs |
| Self CPU Time | 54.30 | μs |
| Self Device Time | 0.00 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
| aten::_to_copy | ||
| CPU Time | 469945.98 | μs |
| Device Time | 2732.40 | μs |
| Self CPU Time | 118.99 | μs |
| Self Device Time | 0.00 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
| aten::empty_strided | ||
| CPU Time | 466748.89 | μs |
| Device Time | 0.00 | μs |
| Self CPU Time | 132.53 | μs |
| Self Device Time | 0.00 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
| aten::conv3d | ||
| CPU Time | 444392.00 | μs |
| Device Time | 4829856.74 | μs |
| Self CPU Time | 14329.86 | μs |
| Self Device Time | 0.00 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
| aten::convolution | ||
| CPU Time | 430062.14 | μs |
| Device Time | 4829856.74 | μs |
| Self CPU Time | 20919.79 | μs |
| Self Device Time | 0.00 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
| aten::_convolution | ||
| CPU Time | 409142.36 | μs |
| Device Time | 4829856.74 | μs |
| Self CPU Time | 40295.99 | μs |
| Self Device Time | 0.00 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
| aten::cudnn_convolution | ||
| CPU Time | 281167.91 | μs |
| Device Time | 4187134.79 | μs |
| Self CPU Time | 203679.04 | μs |
| Self Device Time | 4187134.79 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
| sm80_xmma_fprop_implicit_gemm_indexed_f32f32_f32f32_f32_nchwkcrs_nchw_tilesize32x32x8_stage3_warpsize1x2x1_g1_ffma_aligna4_alignc4_execute_kernel__5x_cudnn | ||
| CPU Time | 0.00 | μs |
| Device Time | 4187132.17 | μs |
| Self CPU Time | 0.00 | μs |
| Self Device Time | 4187132.17 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
| cudaDeviceSynchronize | ||
| CPU Time | 5983533.02 | μs |
| Device Time | 106911.04 | μs |
| Self CPU Time | 5983533.02 | μs |
| Self Device Time | 106911.04 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
45289 warnings generated when compiling for host. Suppressed 45324 warnings (45277 in non-user code, 47 NOLINT). Use -header-filter=.* to display errors from all non-system headers. Use -system-headers to display errors from system headers as well.