3_ConvTranspose3d_Sum_LayerNorm_AvgPool_GELU • optimized_fused_3d_norm_baseimport torch
import torch.nn as nn
import torch.nn.functional as F
def module_fn(
x: torch.Tensor,
conv_transpose_weight: torch.Tensor,
conv_transpose_bias: torch.Tensor,
sum_weight: torch.Tensor,
norm_weight: torch.Tensor,
norm_bias: torch.Tensor,
stride: tuple,
padding: tuple,
output_padding: tuple,
pool_kernel_size: tuple,
norm_shape: tuple,
) -> torch.Tensor:
"""
Functional implementation of a sequence of operations:
1. 3D transposed convolution
2. Addition with a learnable weight
3. Layer normalization
4. 3D average pooling
5. GELU activation
Args:
x (torch.Tensor): Input tensor of shape (batch_size, in_channels, depth, height, width)
conv_transpose_weight (torch.Tensor): Weight tensor for transposed convolution
conv_transpose_bias (torch.Tensor): Bias tensor for transposed convolution
sum_weight (torch.Tensor): Learnable weight for addition
norm_weight (torch.Tensor): Weight tensor for layer normalization
norm_bias (torch.Tensor): Bias tensor for layer normalization
stride (tuple): Stride for transposed convolution, as (depth_stride, height_stride, width_stride)
padding (tuple): Padding for transposed convolution, as (depth_pad, height_pad, width_pad)
output_padding (tuple): Output padding for transposed convolution, as (depth_pad, height_pad, width_pad)
pool_kernel_size (tuple): Kernel size for average pooling, as (depth_kernel, height_kernel, width_kernel)
norm_shape (tuple): Shape for layer normalization
Returns:
torch.Tensor: Output tensor after applying all operations
"""
x = F.conv_transpose3d(
x,
conv_transpose_weight,
bias=conv_transpose_bias,
stride=stride,
padding=padding,
output_padding=output_padding,
)
x = x + sum_weight
x = F.layer_norm(x, norm_shape, norm_weight, norm_bias)
x = F.avg_pool3d(x, kernel_size=pool_kernel_size)
x = F.gelu(x)
return x
class Model(nn.Module):
"""
Model that performs a 3D transposed convolution, followed by a sum, layer normalization, average pooling, and GELU activation.
"""
def __init__(
self,
in_channels,
out_channels,
kernel_size,
stride,
padding,
output_padding,
sum_weight,
norm_shape,
pool_kernel_size,
):
super(Model, self).__init__()
conv = nn.ConvTranspose3d(
in_channels,
out_channels,
kernel_size,
stride=stride,
padding=padding,
output_padding=output_padding,
)
self.conv_transpose_weight = nn.Parameter(conv.weight)
self.conv_transpose_bias = nn.Parameter(conv.bias)
self.sum_weight = nn.Parameter(torch.tensor(sum_weight))
norm = nn.LayerNorm(norm_shape)
self.norm_weight = nn.Parameter(norm.weight + torch.randn(norm_shape) * 0.02)
self.norm_bias = nn.Parameter(norm.bias + torch.randn(norm_shape) * 0.02)
def forward(
self,
x,
stride,
padding,
output_padding,
pool_kernel_size,
norm_shape,
fn=module_fn,
):
return fn(
x,
self.conv_transpose_weight,
self.conv_transpose_bias,
self.sum_weight,
self.norm_weight,
self.norm_bias,
stride,
padding,
output_padding,
pool_kernel_size,
norm_shape,
)
batch_size = 128
in_channels = 32
out_channels = 64
depth, height, width = 16, 32, 32
kernel_size = (3, 3, 3)
stride = (2, 2, 2)
padding = (1, 1, 1)
output_padding = (1, 1, 1)
sum_weight = 1.0
norm_shape = (out_channels,)
pool_kernel_size = (2, 2, 2)
def get_inputs():
return [
torch.randn(batch_size, in_channels, depth, height, width),
stride,
padding,
output_padding,
pool_kernel_size,
norm_shape,
]
def get_init_inputs():
return [
in_channels,
out_channels,
kernel_size,
stride,
padding,
output_padding,
sum_weight,
norm_shape,
pool_kernel_size,
]
import torch
import torch.nn as nn
class Model(nn.Module):
"""
Model that performs a 3D transposed convolution, followed by a sum, layer normalization, average pooling, and GELU activation.
"""
def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, sum_weight, norm_shape, pool_kernel_size):
super(Model, self).__init__()
self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding)
self.sum_weight = nn.Parameter(torch.tensor(sum_weight))
self.norm = nn.LayerNorm(norm_shape)
self.norm.weight = nn.Parameter(self.norm.weight + torch.randn(norm_shape)*0.02)
self.norm.bias = nn.Parameter(self.norm.bias + torch.randn(norm_shape)*0.02)
self.avg_pool = nn.AvgPool3d(kernel_size=pool_kernel_size)
self.gelu = nn.GELU()
def forward(self, x):
x = self.conv_transpose(x)
x = x + self.sum_weight
x = self.norm(x)
x = self.avg_pool(x)
x = self.gelu(x)
return x
batch_size = 128
in_channels = 32
out_channels = 64
depth, height, width = 16, 32, 32
kernel_size = (3, 3, 3)
stride = (2, 2, 2)
padding = (1, 1, 1)
output_padding = (1, 1, 1)
sum_weight = 1.0
norm_shape = (out_channels,)
pool_kernel_size = (2, 2, 2)
def get_inputs():
return [torch.randn(batch_size, in_channels, depth, height, width)]
def get_init_inputs():
return [in_channels, out_channels, kernel_size, stride, padding, output_padding, sum_weight, norm_shape, pool_kernel_size]#include <torch/extension.h>
#include <ATen/ATen.h>
#include <cuda_runtime.h>
#include <vector>
#define BLOCK_SIZE 256
#define WARP_SIZE 32
template <typename T>
__global__ void fused_layernorm_kernel(
const T* __restrict__ input,
const T* __restrict__ gamma,
const T* __restrict__ beta,
T* __restrict__ output,
int n1,
int n2
) {
extern __shared__ float s_data[];
float* s_mean = s_data;
float* s_variance = &s_data[BLOCK_SIZE];
const int tid = threadIdx.x;
const int bid = blockIdx.x;
const int offset = bid * n2;
float local_sum = 0.0f;
float local_sq_sum = 0.0f;
#pragma unroll 4
for (int i = tid; i < n2; i += BLOCK_SIZE) {
float val = __ldg(&input[offset + i]);
local_sum += val;
local_sq_sum += val * val;
}
#pragma unroll
for (int offset = WARP_SIZE/2; offset > 0; offset >>= 1) {
local_sum += __shfl_down_sync(0xffffffff, local_sum, offset);
local_sq_sum += __shfl_down_sync(0xffffffff, local_sq_sum, offset);
}
if (tid % WARP_SIZE == 0) {
s_mean[tid / WARP_SIZE] = local_sum;
s_variance[tid / WARP_SIZE] = local_sq_sum;
}
__syncthreads();
if (tid < WARP_SIZE) {
local_sum = (tid < BLOCK_SIZE/WARP_SIZE) ? s_mean[tid] : 0.0f;
local_sq_sum = (tid < BLOCK_SIZE/WARP_SIZE) ? s_variance[tid] : 0.0f;
#pragma unroll
for (int offset = WARP_SIZE/2; offset > 0; offset >>= 1) {
local_sum += __shfl_down_sync(0xffffffff, local_sum, offset);
local_sq_sum += __shfl_down_sync(0xffffffff, local_sq_sum, offset);
}
if (tid == 0) {
s_mean[0] = local_sum;
s_variance[0] = local_sq_sum;
}
}
__syncthreads();
const float mean = s_mean[0] / n2;
const float variance = (s_variance[0] / n2) - (mean * mean);
const float inv_std = rsqrtf(variance + 1e-5f);
#pragma unroll 4
for (int i = tid; i < n2; i += BLOCK_SIZE) {
const int idx = offset + i;
const float normalized = (input[idx] - mean) * inv_std;
output[idx] = gamma[i] * normalized + beta[i];
}
}
torch::Tensor forward(
torch::Tensor x,
torch::Tensor conv_transpose_weight,
torch::Tensor conv_transpose_bias,
torch::Tensor sum_weight,
torch::Tensor norm_weight,
torch::Tensor norm_bias,
std::vector<int64_t> stride,
std::vector<int64_t> padding,
std::vector<int64_t> output_padding,
std::vector<int64_t> pool_kernel_size,
std::vector<int64_t> norm_shape
) {
x = x.contiguous();
conv_transpose_weight = conv_transpose_weight.contiguous();
sum_weight = sum_weight.contiguous();
norm_weight = norm_weight.contiguous();
norm_bias = norm_bias.contiguous();
at::IntArrayRef strideRef(stride);
at::IntArrayRef paddingRef(padding);
at::IntArrayRef outputPaddingRef(output_padding);
at::IntArrayRef poolKernelRef(pool_kernel_size);
auto out = at::conv_transpose3d(
x,
conv_transpose_weight,
conv_transpose_bias,
strideRef,
paddingRef,
outputPaddingRef,
1,
1
);
out.add_(sum_weight);
auto out_size = out.sizes();
int64_t n1 = 1;
for (int i = 0; i < out.dim() - norm_shape.size(); ++i) {
n1 *= out_size[i];
}
int64_t n2 = 1;
for (size_t i = 0; i < norm_shape.size(); ++i) {
n2 *= norm_shape[i];
}
auto output = torch::empty_like(out);
dim3 grid(n1);
dim3 block(BLOCK_SIZE);
fused_layernorm_kernel<<<grid, block, 2 * BLOCK_SIZE * sizeof(float)>>>(
out.data_ptr<float>(),
norm_weight.data_ptr<float>(),
norm_bias.data_ptr<float>(),
output.data_ptr<float>(),
n1,
n2
);
output = at::gelu(at::avg_pool3d(
output,
poolKernelRef,
poolKernelRef,
{0,0,0},
false,
true
));
return output;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("forward", &forward, "Optimized module_fn forward (CUDA)");
}| Metric | Value | Unit | Variance | Samples |
|---|---|---|---|---|
| Executed Ipc Active | 2.370 | inst/cycle | 0.000 | 5 |
| Executed Ipc Elapsed | 2.370 | inst/cycle | 0.000 | 5 |
| Issue Slots Busy | 60.016 | % | 0.000 | 5 |
| Issued Ipc Active | 2.400 | inst/cycle | 0.000 | 5 |
| SM Busy | 60.016 | % | 0.000 | 5 |
| Memory Throughput | 396487222709.812 | byte/second | 1727863615651133.500 | 5 |
| Mem Busy | 58.786 | % | 0.000 | 5 |
| Max Bandwidth | 53.220 | % | 0.000 | 5 |
| L1/TEX Hit Rate | 59.910 | % | 0.000 | 5 |
| L2 Hit Rate | 50.090 | % | 0.000 | 5 |
| Mem Pipes Busy | 53.220 | % | 0.000 | 5 |
| Warp Cycles Per Issued Instruction | 21.720 | cycle | 0.000 | 5 |
| Warp Cycles Per Executed Instruction | 21.990 | cycle | 0.000 | 5 |
| Avg. Active Threads Per Warp | 28.120 | 0.000 | 5 | |
| Avg. Not Predicated Off Threads Per Warp | 27.290 | 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 | 21.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 | 82.330 | % | 0.000 | 5 |
| Achieved Active Warps Per SM | 52.690 | warp | 0.000 | 5 |
| Rule | Description |
|---|---|
| INF HighPipeUtilization | ALU is the highest-utilized pipeline (24.4%) 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 (82.3%) 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::conv_transpose3d | ||
| CPU Time | 1353828.93 | μs |
| Device Time | 2644694.64 | μs |
| Self CPU Time | 586.31 | μ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 | 1353242.62 | μs |
| Device Time | 2644694.64 | μs |
| Self CPU Time | 808.54 | μ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 | 1352434.08 | μs |
| Device Time | 2644694.64 | μs |
| Self CPU Time | 1715.47 | μ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 |
| cudaLaunchKernel | ||
| CPU Time | 7677809.53 | μs |
| Device Time | 6551.49 | μs |
| Self CPU Time | 7677809.53 | μs |
| Self Device Time | 6551.49 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
| aten::add_ | ||
| CPU Time | 1976571.90 | μs |
| Device Time | 1994086.34 | μs |
| Self CPU Time | 4329.78 | μs |
| Self Device Time | 1994086.34 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
| void fused_layernorm_kernel<float>(float const*, float const*, float const*, float*, int, int) | ||
| CPU Time | 0.00 | μs |
| Device Time | 5541110.88 | μs |
| Self CPU Time | 0.00 | μs |
| Self Device Time | 5541110.88 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
| aten::avg_pool3d | ||
| CPU Time | 4930205.89 | μs |
| Device Time | 818949.24 | μs |
| Self CPU Time | 3263.88 | μs |
| Self Device Time | 818949.24 | μ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 | 2349463.18 | μs |
| Device Time | 3719.31 | μs |
| Self CPU Time | 2349463.18 | μs |
| Self Device Time | 3719.31 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
45298 warnings generated when compiling for host. Suppressed 45331 warnings (45284 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.