import 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>
#include <cuda_runtime.h>
// Maximum channels supported for constant memory usage
#define MAX_CHANNELS 1024
// Constant memory for fast broadcast of the channel sum values
__constant__ float d_sum_tensor_const[MAX_CHANNELS];
// Inline device function: LeakyReLU
__device__ inline float leaky_relu(float x, float alpha = 0.2f) {
return (x > 0.0f) ? x : alpha * x;
}
// Inline device function: Clamp between -1 and 1
__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: GELU 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);
}
// Optimized CUDA kernel using a grid-stride loop and constant memory for sum_tensor
__global__ void optimized_kernel(
const float* __restrict__ input,
float* __restrict__ output,
int64_t num_elements,
int64_t width,
int64_t height,
int64_t depth,
int64_t channels) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int stride = blockDim.x * gridDim.x;
// Use grid-stride loop to cover all elements
for (int i = idx; i < num_elements; i += stride) {
// Compute multi-dimensional indices from flattened index
int64_t w = i % width;
int64_t h = (i / width) % height;
int64_t d = (i / (width * height)) % depth;
int64_t c = (i / (width * height * depth)) % channels;
float x = input[i];
float y = leaky_relu(x);
// Use value from constant memory for the corresponding channel
y += d_sum_tensor_const[c];
y = clamp_val(y);
y = gelu(y);
output[i] = y;
}
}
// Kernel launcher: copies the sum_tensor to constant memory and launches the optimized kernel
void optimized_kernel_launcher(torch::Tensor &x, torch::Tensor &sum_tensor) {
int64_t batch_size = x.size(0);
int64_t channels = x.size(1);
int64_t depth = x.size(2);
int64_t height = x.size(3);
int64_t width = x.size(4);
int64_t num_elements = x.numel();
TORCH_CHECK(channels <= MAX_CHANNELS, "Channels exceed maximum constant memory limit");
int bytes = channels * sizeof(float);
// Copy sum_tensor (assumed to have shape compatible with channels) into constant memory
cudaMemcpyToSymbol(d_sum_tensor_const, sum_tensor.data_ptr<float>(), bytes, 0, cudaMemcpyDeviceToDevice);
// Choose an optimal block size. The grid-stride loop handles arbitrary sizes
const int threads = 256;
const int blocks = (num_elements + threads - 1) / threads;
optimized_kernel<<<blocks, threads>>>(
x.data_ptr<float>(),
x.data_ptr<float>(),
num_elements,
width,
height,
depth,
channels);
cudaDeviceSynchronize();
}
// Forward function: performs convolution followed by the optimized element-wise operations
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 contiguous memory layout for efficient memory access
auto output = x_conv.contiguous();
// Launch the optimized kernel that uses constant memory and grid-stride loop
optimized_kernel_launcher(output, sum_tensor);
return output;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("forward", &forward, "Optimized Conv3d LeakyReLU Sum Clamp GELU forward using grid-stride loop and constant memory (CUDA)");
}
| Metric | Value | Unit | Variance | Samples |
|---|---|---|---|---|
| Executed Ipc Active | 3.298 | inst/cycle | 0.000 | 5 |
| Executed Ipc Elapsed | 3.246 | inst/cycle | 0.000 | 5 |
| Issue Slots Busy | 82.484 | % | 0.006 | 5 |
| Issued Ipc Active | 3.298 | inst/cycle | 0.000 | 5 |
| SM Busy | 82.484 | % | 0.006 | 5 |
| Memory Throughput | 1209544106619.502 | byte/second | 1793617844545018112.000 | 5 |
| Mem Busy | 20.456 | % | 0.001 | 5 |
| Max Bandwidth | 36.086 | % | 0.002 | 5 |
| L1/TEX Hit Rate | 49.858 | % | 0.000 | 5 |
| L2 Hit Rate | 50.326 | % | 0.012 | 5 |
| Mem Pipes Busy | 15.214 | % | 0.000 | 5 |
| Warp Cycles Per Issued Instruction | 16.250 | cycle | 0.000 | 5 |
| Warp Cycles Per Executed Instruction | 16.258 | cycle | 0.000 | 5 |
| Avg. Active Threads Per Warp | 29.460 | 0.000 | 5 | |
| Avg. Not Predicated Off Threads Per Warp | 27.090 | 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 | 10.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 | 84.530 | % | 0.013 | 5 |
| Achieved Active Warps Per SM | 54.098 | warp | 0.005 | 5 |
| Rule | Description |
|---|---|
| INF HighPipeUtilization | ALU is the highest-utilized pipeline (49.9%) 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. |
| 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 (84.6%) 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 | 498346.22 | μs |
| Device Time | 2516.95 | μs |
| Self CPU Time | 69.73 | μ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 | 498276.48 | μs |
| Device Time | 2516.95 | μs |
| Self CPU Time | 133.83 | μ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 | 495320.43 | μs |
| Device Time | 0.00 | μs |
| Self CPU Time | 135.02 | μ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 |
| cudaDeviceGetStreamPriorityRange | ||
| CPU Time | 495947.06 | μs |
| Device Time | 0.00 | μs |
| Self CPU Time | 495947.06 | μ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 | 389931.21 | μs |
| Device Time | 5137742.48 | μs |
| Self CPU Time | 12626.08 | μ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 | 377305.12 | μs |
| Device Time | 5137742.48 | μs |
| Self CPU Time | 18710.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::_convolution | ||
| CPU Time | 358594.60 | μs |
| Device Time | 5137742.48 | μs |
| Self CPU Time | 34454.22 | μ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 | 246957.01 | μs |
| Device Time | 4453195.13 | μs |
| Self CPU Time | 178278.11 | μs |
| Self Device Time | 4453195.13 | μ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 | 4453192.45 | μs |
| Self CPU Time | 0.00 | μs |
| Self Device Time | 4453192.45 | μ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 | 6397224.33 | μs |
| Device Time | 120424.87 | μs |
| Self CPU Time | 6397224.33 | μs |
| Self Device Time | 120424.87 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
45288 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.