72_ConvTranspose3d_BatchNorm_AvgPool_AvgPool
• fully_unrolled_avgpool_base_base
import torch
import torch.nn as nn
import torch.nn.functional as F
def module_fn(
x: torch.Tensor,
stride: int,
padding: int,
conv_transpose: torch.Tensor,
conv_transpose_bias: torch.Tensor,
bn_weight: torch.Tensor,
bn_bias: torch.Tensor,
bn_running_mean: torch.Tensor,
bn_running_var: torch.Tensor,
bn_eps: torch.Tensor,
bn_momentum: torch.Tensor,
) -> torch.Tensor:
"""
Applies a 3D transposed convolution, batch normalization and two average pooling layers.
Args:
x (torch.Tensor): Input tensor of shape (batch_size, in_channels, depth, height, width)
stride (int): Stride of the transposed convolution
padding (int): Padding of the transposed convolution
conv_transpose (torch.Tensor): Transposed convolution weight tensor
conv_transpose_bias (torch.Tensor): Bias tensor for transposed convolution
bn_weight (torch.Tensor): Batch norm weight parameter
bn_bias (torch.Tensor): Batch norm bias parameter
bn_running_mean (torch.Tensor): Batch norm running mean
bn_running_var (torch.Tensor): Batch norm running variance
bn_eps (torch.Tensor): Small constant for numerical stability
bn_momentum (torch.Tensor): Momentum for running stats
Returns:
torch.Tensor: Output tensor after applying transposed conv, batch norm and avg pooling
"""
x = F.conv_transpose3d(
x, conv_transpose, bias=conv_transpose_bias, stride=stride, padding=padding
)
x = F.batch_norm(
x,
bn_running_mean,
bn_running_var,
bn_weight,
bn_bias,
training=True,
momentum=bn_momentum,
eps=bn_eps,
)
x = F.avg_pool3d(x, kernel_size=2)
x = F.avg_pool3d(x, kernel_size=2)
return x
class Model(nn.Module):
"""
A model that performs a 3D transposed convolution, followed by batch normalization,
two average pooling layers.
"""
def __init__(
self, in_channels, out_channels, kernel_size, stride, padding, bias_shape
):
super(Model, self).__init__()
conv = nn.ConvTranspose3d(in_channels, out_channels, kernel_size)
bn = nn.BatchNorm3d(out_channels)
self.conv_transpose_parameter = nn.Parameter(conv.weight)
self.conv_transpose_bias = nn.Parameter(conv.bias)
self.bn_weight = nn.Parameter(bn.weight + torch.randn(bn.weight.shape) * 0.02)
self.bn_bias = nn.Parameter(bn.bias + torch.randn(bn.bias.shape) * 0.02)
self.register_buffer(
"bn_running_mean",
bn.running_mean + torch.randn(bn.running_mean.shape) * 0.02,
)
self.register_buffer(
"bn_running_var",
bn.running_var + torch.randn(bn.running_var.shape).abs() * 0.02,
)
self.register_buffer("bn_eps", torch.tensor(1e-5))
self.register_buffer("bn_momentum", torch.tensor(0.1))
def forward(self, x, stride, padding, fn=module_fn):
return fn(
x,
stride,
padding,
self.conv_transpose_parameter,
self.conv_transpose_bias,
self.bn_weight,
self.bn_bias,
self.bn_running_mean,
self.bn_running_var,
self.bn_eps,
self.bn_momentum,
)
batch_size = 128
in_channels = 3
out_channels = 16
depth, height, width = 32, 32, 32
kernel_size = 3
stride = 2
padding = 1
bias_shape = (out_channels, 1, 1, 1)
def get_inputs():
return [torch.randn(batch_size, in_channels, depth, height, width), stride, padding]
def get_init_inputs():
return [in_channels, out_channels, kernel_size, stride, padding, bias_shape]
import torch
import torch.nn as nn
class Model(nn.Module):
"""
A model that performs a 3D transposed convolution, followed by batch normalization,
two average pooling layers.
"""
def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias_shape):
super(Model, self).__init__()
self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
self.batch_norm = nn.BatchNorm3d(out_channels)
# Add noise to batch norm parameters to match functional implementation
self.batch_norm.weight = nn.Parameter(self.batch_norm.weight + torch.randn(self.batch_norm.weight.shape) * 0.02)
self.batch_norm.bias = nn.Parameter(self.batch_norm.bias + torch.randn(self.batch_norm.bias.shape) * 0.02)
self.batch_norm.running_mean = self.batch_norm.running_mean + torch.randn(self.batch_norm.running_mean.shape) * 0.02
self.batch_norm.running_var = self.batch_norm.running_var + torch.randn(self.batch_norm.running_var.shape).abs() * 0.02
self.avg_pool1 = nn.AvgPool3d(kernel_size=2)
self.avg_pool2 = nn.AvgPool3d(kernel_size=2)
def forward(self, x):
x = self.conv_transpose(x)
x = self.batch_norm(x)
x = self.avg_pool1(x)
x = self.avg_pool2(x)
return x
batch_size = 128
in_channels = 3
out_channels = 16
depth, height, width = 32, 32, 32
kernel_size = 3
stride = 2
padding = 1
bias_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, stride, padding, bias_shape]
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <vector>
#define BLOCK_SIZE 256
__global__ void fully_unrolled_avgpool_kernel(const float* __restrict__ input,
float* __restrict__ output,
int N, int C, int D, int H, int W,
int pooled_D, int pooled_H, int pooled_W) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = N * C * pooled_D * pooled_H * pooled_W;
if (idx < total) {
// Decode index
int pw = idx % pooled_W;
int ph = (idx / pooled_W) % pooled_H;
int pd = (idx / (pooled_W * pooled_H)) % pooled_D;
int c = (idx / (pooled_W * pooled_H * pooled_D)) % C;
int n = idx / (pooled_W * pooled_H * pooled_D * C);
// Calculate base index for input
int base_d = pd * 4;
int base_h = ph * 4;
int base_w = pw * 4;
int batch_offset = n * C * D * H * W;
int channel_offset = c * D * H * W;
// Fully unrolled 4x4x4 pooling
float sum =
// d=0, h=0
input[batch_offset + channel_offset + (base_d+0)*H*W + (base_h+0)*W + (base_w+0)] +
input[batch_offset + channel_offset + (base_d+0)*H*W + (base_h+0)*W + (base_w+1)] +
input[batch_offset + channel_offset + (base_d+0)*H*W + (base_h+0)*W + (base_w+2)] +
input[batch_offset + channel_offset + (base_d+0)*H*W + (base_h+0)*W + (base_w+3)] +
// d=0, h=1
input[batch_offset + channel_offset + (base_d+0)*H*W + (base_h+1)*W + (base_w+0)] +
input[batch_offset + channel_offset + (base_d+0)*H*W + (base_h+1)*W + (base_w+1)] +
input[batch_offset + channel_offset + (base_d+0)*H*W + (base_h+1)*W + (base_w+2)] +
input[batch_offset + channel_offset + (base_d+0)*H*W + (base_h+1)*W + (base_w+3)] +
// d=0, h=2
input[batch_offset + channel_offset + (base_d+0)*H*W + (base_h+2)*W + (base_w+0)] +
input[batch_offset + channel_offset + (base_d+0)*H*W + (base_h+2)*W + (base_w+1)] +
input[batch_offset + channel_offset + (base_d+0)*H*W + (base_h+2)*W + (base_w+2)] +
input[batch_offset + channel_offset + (base_d+0)*H*W + (base_h+2)*W + (base_w+3)] +
// d=0, h=3
input[batch_offset + channel_offset + (base_d+0)*H*W + (base_h+3)*W + (base_w+0)] +
input[batch_offset + channel_offset + (base_d+0)*H*W + (base_h+3)*W + (base_w+1)] +
input[batch_offset + channel_offset + (base_d+0)*H*W + (base_h+3)*W + (base_w+2)] +
input[batch_offset + channel_offset + (base_d+0)*H*W + (base_h+3)*W + (base_w+3)] +
// d=1
input[batch_offset + channel_offset + (base_d+1)*H*W + (base_h+0)*W + (base_w+0)] +
input[batch_offset + channel_offset + (base_d+1)*H*W + (base_h+0)*W + (base_w+1)] +
input[batch_offset + channel_offset + (base_d+1)*H*W + (base_h+0)*W + (base_w+2)] +
input[batch_offset + channel_offset + (base_d+1)*H*W + (base_h+0)*W + (base_w+3)] +
input[batch_offset + channel_offset + (base_d+1)*H*W + (base_h+1)*W + (base_w+0)] +
input[batch_offset + channel_offset + (base_d+1)*H*W + (base_h+1)*W + (base_w+1)] +
input[batch_offset + channel_offset + (base_d+1)*H*W + (base_h+1)*W + (base_w+2)] +
input[batch_offset + channel_offset + (base_d+1)*H*W + (base_h+1)*W + (base_w+3)] +
input[batch_offset + channel_offset + (base_d+1)*H*W + (base_h+2)*W + (base_w+0)] +
input[batch_offset + channel_offset + (base_d+1)*H*W + (base_h+2)*W + (base_w+1)] +
input[batch_offset + channel_offset + (base_d+1)*H*W + (base_h+2)*W + (base_w+2)] +
input[batch_offset + channel_offset + (base_d+1)*H*W + (base_h+2)*W + (base_w+3)] +
input[batch_offset + channel_offset + (base_d+1)*H*W + (base_h+3)*W + (base_w+0)] +
input[batch_offset + channel_offset + (base_d+1)*H*W + (base_h+3)*W + (base_w+1)] +
input[batch_offset + channel_offset + (base_d+1)*H*W + (base_h+3)*W + (base_w+2)] +
input[batch_offset + channel_offset + (base_d+1)*H*W + (base_h+3)*W + (base_w+3)] +
// d=2, d=3 (similar pattern continues)
input[batch_offset + channel_offset + (base_d+2)*H*W + (base_h+0)*W + (base_w+0)] +
input[batch_offset + channel_offset + (base_d+2)*H*W + (base_h+0)*W + (base_w+1)] +
input[batch_offset + channel_offset + (base_d+2)*H*W + (base_h+0)*W + (base_w+2)] +
input[batch_offset + channel_offset + (base_d+2)*H*W + (base_h+0)*W + (base_w+3)] +
input[batch_offset + channel_offset + (base_d+2)*H*W + (base_h+1)*W + (base_w+0)] +
input[batch_offset + channel_offset + (base_d+2)*H*W + (base_h+1)*W + (base_w+1)] +
input[batch_offset + channel_offset + (base_d+2)*H*W + (base_h+1)*W + (base_w+2)] +
input[batch_offset + channel_offset + (base_d+2)*H*W + (base_h+1)*W + (base_w+3)] +
input[batch_offset + channel_offset + (base_d+2)*H*W + (base_h+2)*W + (base_w+0)] +
input[batch_offset + channel_offset + (base_d+2)*H*W + (base_h+2)*W + (base_w+1)] +
input[batch_offset + channel_offset + (base_d+2)*H*W + (base_h+2)*W + (base_w+2)] +
input[batch_offset + channel_offset + (base_d+2)*H*W + (base_h+2)*W + (base_w+3)] +
input[batch_offset + channel_offset + (base_d+2)*H*W + (base_h+3)*W + (base_w+0)] +
input[batch_offset + channel_offset + (base_d+2)*H*W + (base_h+3)*W + (base_w+1)] +
input[batch_offset + channel_offset + (base_d+2)*H*W + (base_h+3)*W + (base_w+2)] +
input[batch_offset + channel_offset + (base_d+2)*H*W + (base_h+3)*W + (base_w+3)] +
// d=3
input[batch_offset + channel_offset + (base_d+3)*H*W + (base_h+0)*W + (base_w+0)] +
input[batch_offset + channel_offset + (base_d+3)*H*W + (base_h+0)*W + (base_w+1)] +
input[batch_offset + channel_offset + (base_d+3)*H*W + (base_h+0)*W + (base_w+2)] +
input[batch_offset + channel_offset + (base_d+3)*H*W + (base_h+0)*W + (base_w+3)] +
input[batch_offset + channel_offset + (base_d+3)*H*W + (base_h+1)*W + (base_w+0)] +
input[batch_offset + channel_offset + (base_d+3)*H*W + (base_h+1)*W + (base_w+1)] +
input[batch_offset + channel_offset + (base_d+3)*H*W + (base_h+1)*W + (base_w+2)] +
input[batch_offset + channel_offset + (base_d+3)*H*W + (base_h+1)*W + (base_w+3)] +
input[batch_offset + channel_offset + (base_d+3)*H*W + (base_h+2)*W + (base_w+0)] +
input[batch_offset + channel_offset + (base_d+3)*H*W + (base_h+2)*W + (base_w+1)] +
input[batch_offset + channel_offset + (base_d+3)*H*W + (base_h+2)*W + (base_w+2)] +
input[batch_offset + channel_offset + (base_d+3)*H*W + (base_h+2)*W + (base_w+3)] +
input[batch_offset + channel_offset + (base_d+3)*H*W + (base_h+3)*W + (base_w+0)] +
input[batch_offset + channel_offset + (base_d+3)*H*W + (base_h+3)*W + (base_w+1)] +
input[batch_offset + channel_offset + (base_d+3)*H*W + (base_h+3)*W + (base_w+2)] +
input[batch_offset + channel_offset + (base_d+3)*H*W + (base_h+3)*W + (base_w+3)];
output[idx] = sum / 64.0f;
}
}
at::Tensor module_fn_forward(
at::Tensor x,
int64_t stride,
int64_t padding,
at::Tensor conv_transpose,
at::Tensor conv_transpose_bias,
at::Tensor bn_weight,
at::Tensor bn_bias,
at::Tensor bn_running_mean,
at::Tensor bn_running_var,
at::Tensor bn_eps,
at::Tensor bn_momentum
) {
TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
TORCH_CHECK(conv_transpose.is_cuda(), "conv_transpose must be a CUDA tensor");
TORCH_CHECK(conv_transpose_bias.is_cuda(), "conv_transpose_bias must be a CUDA tensor");
TORCH_CHECK(bn_weight.is_cuda(), "bn_weight must be a CUDA tensor");
TORCH_CHECK(bn_bias.is_cuda(), "bn_bias must be a CUDA tensor");
TORCH_CHECK(bn_running_mean.is_cuda(), "bn_running_mean must be a CUDA tensor");
TORCH_CHECK(bn_running_var.is_cuda(), "bn_running_var must be a CUDA tensor");
TORCH_CHECK(bn_eps.is_cuda(), "bn_eps must be a CUDA scalar tensor");
TORCH_CHECK(bn_momentum.is_cuda(), "bn_momentum must be a CUDA scalar tensor");
const double eps_val = bn_eps.item<double>();
const double momentum_val = bn_momentum.item<double>();
std::vector<int64_t> stride_3d = {stride, stride, stride};
std::vector<int64_t> pad_3d = {padding, padding, padding};
auto y = at::conv_transpose3d(x, conv_transpose, conv_transpose_bias, stride_3d, pad_3d);
y = at::batch_norm(y, bn_weight, bn_bias, bn_running_mean, bn_running_var,
true, momentum_val, eps_val, true);
auto sizes = y.sizes();
int N_val = sizes[0];
int C_val = sizes[1];
int D_val = sizes[2];
int H_val = sizes[3];
int W_val = sizes[4];
int pooled_D = D_val / 4;
int pooled_H = H_val / 4;
int pooled_W = W_val / 4;
auto output = at::empty({N_val, C_val, pooled_D, pooled_H, pooled_W}, y.options());
int total_elements = N_val * C_val * pooled_D * pooled_H * pooled_W;
int threads = BLOCK_SIZE;
int blocks = (total_elements + threads - 1) / threads;
fully_unrolled_avgpool_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
y.data_ptr<float>(),
output.data_ptr<float>(),
N_val, C_val, D_val, H_val, W_val,
pooled_D, pooled_H, pooled_W
);
AT_CUDA_CHECK(cudaGetLastError());
return output;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("forward", &module_fn_forward, "Fully unrolled 3D conv transpose + batch norm + avg pool (CUDA)");
}
Metric | Value | Unit | Variance | Samples |
---|---|---|---|---|
Executed Ipc Active | 0.610 | inst/cycle | 0.000 | 5 |
Executed Ipc Elapsed | 0.600 | inst/cycle | 0.000 | 5 |
Issue Slots Busy | 15.236 | % | 0.001 | 5 |
Issued Ipc Active | 0.610 | inst/cycle | 0.000 | 5 |
SM Busy | 15.236 | % | 0.001 | 5 |
Memory Throughput | 2972328140435.056 | byte/second | 8366229821247029248.000 | 5 |
Mem Busy | 56.786 | % | 0.010 | 5 |
Max Bandwidth | 88.666 | % | 0.008 | 5 |
L1/TEX Hit Rate | 75.250 | % | 0.000 | 5 |
L2 Hit Rate | 9.882 | % | 0.001 | 5 |
Mem Pipes Busy | 10.718 | % | 0.000 | 5 |
Warp Cycles Per Issued Instruction | 96.890 | cycle | 0.073 | 5 |
Warp Cycles Per Executed Instruction | 96.906 | cycle | 0.072 | 5 |
Avg. Active Threads Per Warp | 32.000 | 0.000 | 5 | |
Avg. Not Predicated Off Threads Per Warp | 29.180 | 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 | 8.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 | 92.130 | % | 0.001 | 5 |
Achieved Active Warps Per SM | 58.964 | warp | 0.000 | 5 |
Rule | Description |
---|---|
WRN HighPipeUtilization | All compute pipelines are under-utilized. Either this kernel is very small or it doesn't issue enough warps per scheduler. Check the Launch Statistics and Scheduler Statistics sections for further details. |
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. |
INF Occupancy | This kernel's theoretical occupancy is not impacted by any block limit. |
Operation / Metric | Value | Unit |
---|---|---|
aten::to | ||
CPU Time | 223385.72 | μs |
Device Time | 5070.79 | μs |
Self CPU Time | 70.56 | μ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 | 223315.16 | μs |
Device Time | 5070.79 | μs |
Self CPU Time | 135.45 | μ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 |
cudaStreamSynchronize | ||
CPU Time | 9592391.37 | μs |
Device Time | 0.00 | μs |
Self CPU Time | 9592391.37 | μ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::item | ||
CPU Time | 9604977.37 | μs |
Device Time | 1637.24 | μs |
Self CPU Time | 1866.26 | μ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::_local_scalar_dense | ||
CPU Time | 9603111.12 | μs |
Device Time | 1637.24 | μs |
Self CPU Time | 5048.63 | μs |
Self Device Time | 1637.24 | μs |
CPU Memory Usage | 0 | B |
Device Memory Usage | 0 | B |
Self CPU Memory Usage | 0 | B |
Self Device Memory Usage | 0 | B |
aten::conv_transpose3d | ||
CPU Time | 219241.00 | μs |
Device Time | 3403389.92 | μs |
Self CPU Time | 1210.80 | μ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::batch_norm | ||
CPU Time | 33154.55 | μs |
Device Time | 6025108.03 | μs |
Self CPU Time | 1093.44 | μ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::_batch_norm_impl_index | ||
CPU Time | 32061.11 | μs |
Device Time | 6025108.03 | μs |
Self CPU Time | 1206.03 | μ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_batch_norm | ||
CPU Time | 30855.09 | μs |
Device Time | 6025108.03 | μs |
Self CPU Time | 11940.54 | μs |
Self Device Time | 6025108.03 | μs |
CPU Memory Usage | 0 | B |
Device Memory Usage | 0 | B |
Self CPU Memory Usage | 0 | B |
Self Device Memory Usage | 0 | B |
void cudnn::bn_fw_tr_1C11_kernel_NCHW<float, float, int, 512, true, 1, true>(cudnnTensorStruct, float const*, cudnnTensorStruct, float*, float const*, float const*, float, float, float*, float*, float*, float*, float, float) | ||
CPU Time | 0.00 | μs |
Device Time | 6025108.03 | μs |
Self CPU Time | 0.00 | μs |
Self Device Time | 6025108.03 | μs |
CPU Memory Usage | 0 | B |
Device Memory Usage | 0 | B |
Self CPU Memory Usage | 0 | B |
Self Device Memory Usage | 0 | B |
45311 warnings generated when compiling for host. Suppressed 45346 warnings (45299 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.