44_ConvTranspose2d_Multiply_GlobalAvgPool_GlobalAvgPool_Mean • fused_global_avg_baseimport torch
import torch.nn as nn
import torch.nn.functional as F
def module_fn(
x: torch.Tensor,
stride: int,
padding: int,
output_padding: int,
conv_transpose: torch.Tensor,
conv_transpose_bias: torch.Tensor,
multiplier: float,
) -> torch.Tensor:
"""
Applies transposed convolution, scalar multiplication, and multiple global average pooling operations.
Args:
x (torch.Tensor): Input tensor of shape (batch_size, in_channels, height, width)
stride (int): Stride of the transposed convolution
padding (int): Padding of the transposed convolution
output_padding (int): Additional size added to output shape
conv_transpose (torch.Tensor): Transposed convolution weight tensor
conv_transpose_bias (torch.Tensor): Bias tensor for transposed convolution
multiplier (float): Scalar multiplier value
Returns:
torch.Tensor: Scalar output after applying operations
"""
x = F.conv_transpose2d(
x,
conv_transpose,
bias=conv_transpose_bias,
stride=stride,
padding=padding,
output_padding=output_padding,
)
x = x * multiplier
x = torch.mean(x, dim=[2, 3], keepdim=True)
x = torch.mean(x, dim=[2, 3], keepdim=True)
x = torch.mean(x)
return x
class Model(nn.Module):
"""
Model that performs a transposed convolution, multiplies by a scalar, applies global average pooling,
another global average pooling, and then calculates the mean.
"""
def __init__(
self,
in_channels,
out_channels,
kernel_size,
stride,
padding,
output_padding,
multiplier,
):
super(Model, self).__init__()
conv = nn.ConvTranspose2d(
in_channels,
out_channels,
kernel_size,
stride=stride,
padding=padding,
output_padding=output_padding,
)
self.conv_transpose_parameter = nn.Parameter(conv.weight)
self.conv_transpose_bias = nn.Parameter(
conv.bias
+ torch.randn(
conv.bias.shape, device=conv.bias.device, dtype=conv.bias.dtype
)
* 0.02
)
self.multiplier = multiplier
def forward(self, x, stride, padding, output_padding, fn=module_fn):
return fn(
x,
stride,
padding,
output_padding,
self.conv_transpose_parameter,
self.conv_transpose_bias,
self.multiplier,
)
batch_size = 128
in_channels = 3
out_channels = 16
height, width = 32, 32
kernel_size = 3
stride = 2
padding = 1
output_padding = 1
multiplier = 0.5
def get_inputs():
return [
torch.randn(batch_size, in_channels, height, width),
stride,
padding,
output_padding,
]
def get_init_inputs():
return [
in_channels,
out_channels,
kernel_size,
stride,
padding,
output_padding,
multiplier,
]
import torch
import torch.nn as nn
class Model(nn.Module):
"""
Model that performs a transposed convolution, multiplies by a scalar, applies global average pooling,
another global average pooling, and then calculates the mean.
"""
def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, multiplier):
super(Model, self).__init__()
self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding)
self.conv_transpose.bias = nn.Parameter(self.conv_transpose.bias + torch.randn(self.conv_transpose.bias.shape, device=self.conv_transpose.bias.device, dtype=self.conv_transpose.bias.dtype) * 0.02)
self.multiplier = multiplier
def forward(self, x):
x = self.conv_transpose(x)
x = x * self.multiplier
x = torch.mean(x, dim=[2, 3], keepdim=True) # First global average pooling
x = torch.mean(x, dim=[2, 3], keepdim=True) # Second global average pooling
x = torch.mean(x)
return x
batch_size = 128
in_channels = 3
out_channels = 16
height, width = 32, 32
kernel_size = 3
stride = 2
padding = 1
output_padding = 1
multiplier = 0.5
def get_inputs():
return [torch.randn(batch_size, in_channels, height, width)]
def get_init_inputs():
return [in_channels, out_channels, kernel_size, stride, padding, output_padding, multiplier]#include <torch/extension.h>
#include <ATen/ATen.h>
#include <cuda_runtime.h>
// This kernel fuses the elementwise multiplication by a scalar, spatial reduction (global average pooling)
// and accumulation across channels (final mean). Each block handles one (batch, channel) slice.
template <int BLOCK_SIZE>
__global__ void fused_global_avg_kernel(
const float* __restrict__ input,
float* __restrict__ global_accum, // pointer to a single float to accumulate all channel averages
const int N,
const int C,
const int H,
const int W,
const float multiplier
) {
extern __shared__ float sdata[];
int tid = threadIdx.x;
int bid = blockIdx.x; // each block corresponds to one (batch, channel) combination
int num_elements = H * W;
// Determine the batch and channel for this block
int batch_idx = bid / C;
int channel_idx = bid % C;
// Pointer to the (batch_idx, channel_idx) slice of the input
const float* channel_input = input + (batch_idx * C * num_elements) + (channel_idx * num_elements);
// Each thread accumulates partial sum for its assigned spatial elements
float sum = 0.0f;
for (int i = tid; i < num_elements; i += BLOCK_SIZE) {
sum += channel_input[i];
}
// Store each thread's partial sum in shared memory
sdata[tid] = sum;
__syncthreads();
// Reduction in shared memory
if (BLOCK_SIZE >= 512) {
if (tid < 256) { sdata[tid] += sdata[tid + 256]; }
__syncthreads();
}
if (BLOCK_SIZE >= 256) {
if (tid < 128) { sdata[tid] += sdata[tid + 128]; }
__syncthreads();
}
if (BLOCK_SIZE >= 128) {
if (tid < 64) { sdata[tid] += sdata[tid + 64]; }
__syncthreads();
}
// Warp-level reduction (no __syncthreads needed within a warp)
if (tid < 32) {
volatile float* vsmem = sdata;
vsmem[tid] += vsmem[tid + 32];
vsmem[tid] += vsmem[tid + 16];
vsmem[tid] += vsmem[tid + 8];
vsmem[tid] += vsmem[tid + 4];
vsmem[tid] += vsmem[tid + 2];
vsmem[tid] += vsmem[tid + 1];
}
// Thread 0 computes the channel average and adds it atomically to the global accumulator
if (tid == 0) {
// Instead of applying multiplier elementwise before reduction, we do it here on the summed value.
// Then compute channel average by dividing by the number of spatial elements.
float channel_avg = (sdata[0] * multiplier) / num_elements;
atomicAdd(global_accum, channel_avg);
}
}
// Module function that applies a transposed convolution, fuses elementwise multiplication and global average pooling
at::Tensor fused_module_fn(
at::Tensor x,
int64_t stride,
int64_t padding,
int64_t output_padding,
at::Tensor conv_transpose,
at::Tensor conv_transpose_bias,
double multiplier
) {
// Perform transposed convolution (implemented by PyTorch)
at::Tensor y = at::conv_transpose2d(
x,
conv_transpose,
conv_transpose_bias,
{stride, stride},
{padding, padding},
{output_padding, output_padding},
1,
{1, 1}
);
// Instead of doing an out-of-place elementwise multiplication, we fuse the multiplier in our reduction kernel
// Prepare a tensor on the same device to accumulate the global sum
auto options = torch::TensorOptions().device(y.device()).dtype(y.dtype());
at::Tensor global_accum_tensor = torch::zeros({1}, options);
// Kernel configuration
constexpr int BLOCK_SIZE = 256;
int N = y.size(0);
int C = y.size(1);
int H = y.size(2);
int W = y.size(3);
int num_channels = N * C; // one block per (batch, channel) combination
int shared_mem_size = BLOCK_SIZE * sizeof(float);
fused_global_avg_kernel<BLOCK_SIZE><<<num_channels, BLOCK_SIZE, shared_mem_size>>>(
y.data_ptr<float>(),
global_accum_tensor.data_ptr<float>(),
N, C, H, W,
static_cast<float>(multiplier)
);
// Ensure kernel completes
cudaDeviceSynchronize();
// The kernel has accumulated the per-channel averages into global_accum_tensor[0].
// Final result (global average) is the mean of these channel averages.
at::Tensor result = global_accum_tensor / static_cast<float>(num_channels);
return result;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("forward", &fused_module_fn, "Fused Global Average Pooling Module");
}
| Metric | Value | Unit | Variance | Samples |
|---|---|---|---|---|
| Executed Ipc Active | 0.932 | inst/cycle | 0.000 | 5 |
| Executed Ipc Elapsed | 0.744 | inst/cycle | 0.000 | 5 |
| Issue Slots Busy | 23.426 | % | 0.004 | 5 |
| Issued Ipc Active | 0.938 | inst/cycle | 0.000 | 5 |
| SM Busy | 23.426 | % | 0.004 | 5 |
| Memory Throughput | 2382781464804.190 | byte/second | 156675733414586679296.000 | 5 |
| Mem Busy | 40.168 | % | 0.041 | 5 |
| Max Bandwidth | 71.252 | % | 0.120 | 5 |
| L1/TEX Hit Rate | 0.000 | % | 0.000 | 5 |
| L2 Hit Rate | 2.922 | % | 0.000 | 5 |
| Mem Pipes Busy | 14.274 | % | 0.005 | 5 |
| Warp Cycles Per Issued Instruction | 55.542 | cycle | 1.742 | 5 |
| Warp Cycles Per Executed Instruction | 55.714 | cycle | 1.753 | 5 |
| Avg. Active Threads Per Warp | 31.750 | 0.000 | 5 | |
| Avg. Not Predicated Off Threads Per Warp | 28.720 | 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 | 16.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 | 80.558 | % | 0.022 | 5 |
| Achieved Active Warps Per SM | 51.558 | warp | 0.009 | 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. |
| 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 (80.5%) 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_transpose2d | ||
| CPU Time | 2096935.26 | μs |
| Device Time | 4572573.43 | μs |
| Self CPU Time | 52637.09 | μ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 | 2044298.17 | μs |
| Device Time | 4572573.43 | μs |
| Self CPU Time | 69142.18 | μ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 | 1975155.99 | μs |
| Device Time | 4572573.43 | μs |
| Self CPU Time | 136921.57 | μ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_transpose | ||
| CPU Time | 1571813.25 | μs |
| Device Time | 3719724.20 | μs |
| Self CPU Time | 651932.66 | μs |
| Self Device Time | 3719722.92 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
| aten::zero_ | ||
| CPU Time | 430699.76 | μs |
| Device Time | 2250329.67 | μs |
| Self CPU Time | 99312.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 |
| cudaDeviceSynchronize | ||
| CPU Time | 4919269.79 | μs |
| Device Time | 63911.16 | μs |
| Self CPU Time | 4919269.79 | μs |
| Self Device Time | 63911.16 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
45293 warnings generated when compiling for host. Suppressed 45327 warnings (45280 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.