import torch
import torch.nn as nn
import torch.nn.functional as F
def module_fn(
x: torch.Tensor,
stride: int,
padding: int,
conv_weight: torch.Tensor,
conv_bias: torch.Tensor,
) -> torch.Tensor:
"""
Applies 3D convolution followed by Mish and Tanh activations.
Args:
x (torch.Tensor): Input tensor of shape (batch_size, in_channels, D, H, W)
stride (int): Stride of the convolution
padding (int): Padding of the convolution
conv_weight (torch.Tensor): Convolution weight tensor of shape
(out_channels, in_channels, kernel_size, kernel_size, kernel_size)
conv_bias (torch.Tensor): Bias tensor for convolution of shape (out_channels)
Returns:
torch.Tensor: Output tensor after applying convolution, Mish and Tanh activations
"""
x = F.conv3d(x, conv_weight, bias=conv_bias, stride=stride, padding=padding)
x = F.mish(x)
x = torch.tanh(x)
return x
class Model(nn.Module):
"""
Model that performs a 3D convolution, applies Mish activation, and then applies Tanh activation.
"""
def __init__(self, in_channels, out_channels, kernel_size, stride, padding):
super(Model, self).__init__()
conv = nn.Conv3d(
in_channels, out_channels, kernel_size, stride=stride, padding=padding
)
self.conv_weight = nn.Parameter(conv.weight)
self.conv_bias = nn.Parameter(
conv.bias
+ torch.randn(
conv.bias.shape, device=conv.bias.device, dtype=conv.bias.dtype
)
* 0.02
)
def forward(self, x, stride, padding, fn=module_fn):
return fn(x, stride, padding, self.conv_weight, self.conv_bias)
batch_size = 16
in_channels = 3
out_channels = 16
D, H, W = 16, 32, 32
kernel_size = 3
stride = 1
padding = 0
def get_inputs():
return [torch.randn(batch_size, in_channels, D, H, W), stride, padding]
def get_init_inputs():
return [in_channels, out_channels, kernel_size, stride, padding]
import torch
import torch.nn as nn
class Model(nn.Module):
"""
Model that performs a 3D convolution, applies Mish activation, and then applies Tanh activation.
"""
def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0):
super(Model, self).__init__()
self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
self.conv.bias = nn.Parameter(self.conv.bias + torch.randn(self.conv.bias.shape, device=self.conv.bias.device, dtype=self.conv.bias.dtype) * 0.02)
def forward(self, x):
"""
Args:
x (torch.Tensor): Input tensor of shape (batch_size, in_channels, D, H, W).
Returns:
torch.Tensor: Output tensor of shape (batch_size, out_channels, D', H', W').
"""
x = self.conv(x)
x = torch.nn.functional.mish(x)
x = torch.tanh(x)
return x
batch_size = 16
in_channels = 3
out_channels = 16
D, H, W = 16, 32, 32
kernel_size = 3
def get_inputs():
return [torch.randn(batch_size, in_channels, D, H, W)]
def get_init_inputs():
return [in_channels, out_channels, kernel_size]#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>
// Warp size for NVIDIA GPUs
#define WARP_SIZE 32
__device__ __forceinline__ float fused_mish_tanh_activation(float x) {
float softplus = logf(1.0f + expf(x));
float mish = x * tanhf(softplus);
return tanhf(mish);
}
__device__ __forceinline__ float4 load_float4(const float* addr) {
float4 val;
val = *reinterpret_cast<const float4*>(addr);
return val;
}
__device__ __forceinline__ void store_float4(float* addr, float4 val) {
*reinterpret_cast<float4*>(addr) = val;
}
__global__ void warp_optimized_mish_tanh_kernel(
float* __restrict__ output,
const float* __restrict__ input,
const int total_elements
) {
// Get warp and lane information
const int warp_id = threadIdx.x / WARP_SIZE;
const int lane_id = threadIdx.x % WARP_SIZE;
const int warp_per_block = blockDim.x / WARP_SIZE;
const int global_warp_id = blockIdx.x * warp_per_block + warp_id;
// Calculate base index for this warp
int base_idx = global_warp_id * (WARP_SIZE * 4); // Process 4 elements per thread
// Process 4 elements per thread when possible
if (base_idx + (WARP_SIZE * 4) <= total_elements) {
float4 in_val4 = load_float4(input + base_idx + lane_id * 4);
// Process each element
float4 out_val4;
out_val4.x = fused_mish_tanh_activation(in_val4.x);
out_val4.y = fused_mish_tanh_activation(in_val4.y);
out_val4.z = fused_mish_tanh_activation(in_val4.z);
out_val4.w = fused_mish_tanh_activation(in_val4.w);
// Store result
store_float4(output + base_idx + lane_id * 4, out_val4);
}
else {
// Handle remaining elements
for (int i = base_idx + lane_id; i < total_elements; i += WARP_SIZE) {
output[i] = fused_mish_tanh_activation(input[i]);
}
}
}
torch::Tensor module_fn_forward(
torch::Tensor x,
int64_t stride,
int64_t padding,
torch::Tensor conv_weight,
torch::Tensor conv_bias
) {
TORCH_CHECK(x.is_cuda(), "Input tensor x must be a CUDA tensor");
TORCH_CHECK(conv_weight.is_cuda(), "Convolution weight must be a CUDA tensor");
TORCH_CHECK(conv_bias.is_cuda(), "Convolution bias must be a CUDA tensor");
auto x_conv = at::conv3d(
x,
conv_weight,
conv_bias,
{stride, stride, stride},
{padding, padding, padding}
);
auto output = torch::empty_like(x_conv);
const int total_elements = x_conv.numel();
// Configure kernel launch parameters
const int threads_per_block = 256; // Multiple of WARP_SIZE
const int num_blocks = (total_elements + threads_per_block * 4 - 1) / (threads_per_block * 4);
warp_optimized_mish_tanh_kernel<<<num_blocks, threads_per_block>>>(
output.data_ptr<float>(),
x_conv.data_ptr<float>(),
total_elements
);
return output;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("forward", &module_fn_forward, "Warp optimized convolution with Mish and Tanh activations (CUDA)");
}| Metric | Value | Unit | Variance | Samples |
|---|---|---|---|---|
| Executed Ipc Active | 2.826 | inst/cycle | 0.000 | 5 |
| Executed Ipc Elapsed | 2.284 | inst/cycle | 0.000 | 5 |
| Issue Slots Busy | 71.250 | % | 0.075 | 5 |
| Issued Ipc Active | 2.850 | inst/cycle | 0.000 | 5 |
| SM Busy | 71.250 | % | 0.075 | 5 |
| Memory Throughput | 1091418463681.024 | byte/second | 7086087547311440896.000 | 5 |
| Mem Busy | 28.090 | % | 0.005 | 5 |
| Max Bandwidth | 32.686 | % | 0.003 | 5 |
| L1/TEX Hit Rate | 0.000 | % | 0.000 | 5 |
| L2 Hit Rate | 50.684 | % | 0.012 | 5 |
| Mem Pipes Busy | 6.704 | % | 0.000 | 5 |
| Warp Cycles Per Issued Instruction | 18.402 | cycle | 0.014 | 5 |
| Warp Cycles Per Executed Instruction | 18.562 | cycle | 0.014 | 5 |
| Avg. Active Threads Per Warp | 20.890 | 0.000 | 5 | |
| Avg. Not Predicated Off Threads Per Warp | 20.340 | 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 | 82.204 | % | 0.007 | 5 |
| Achieved Active Warps Per SM | 52.614 | warp | 0.003 | 5 |
| Rule | Description |
|---|---|
| INF HighPipeUtilization | ALU is the highest-utilized pipeline (30.7%) 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 ThreadDivergence | Instructions are executed in warps, which are groups of 32 threads. Optimal instruction throughput is achieved if all 32 threads of a warp execute the same instruction. The chosen launch configuration, early thread completion, and divergent flow control can significantly lower the number of active threads in a warp per cycle. This kernel achieves an average of 20.9 threads being active per cycle. This is further reduced to 20.3 threads per warp due to predication. The compiler may use predication to avoid an actual branch. Instead, all instructions are scheduled, but a per-thread condition code or predicate controls which threads execute the instructions. Try to avoid different execution paths within a warp when possible. In addition, ensure your kernel makes use of Independent Thread Scheduling, which allows a warp to reconverge after a data-dependent conditional block by explicitly calling __syncwarp(). |
| 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.2%) 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::conv3d | ||
| CPU Time | 1152920.08 | μs |
| Device Time | 1165156.11 | μs |
| Self CPU Time | 21817.89 | μ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 | 1131102.19 | μs |
| Device Time | 1165156.11 | μs |
| Self CPU Time | 27806.98 | μ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 | 1103295.21 | μs |
| Device Time | 1165156.11 | μs |
| Self CPU Time | 54578.64 | μ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 | 926882.77 | μs |
| Device Time | 1015011.63 | μs |
| Self CPU Time | 225778.87 | μs |
| Self Device Time | 1015011.63 | μ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 | 1015008.75 | μs |
| Self CPU Time | 0.00 | μs |
| Self Device Time | 1015008.75 | μ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 | 950518.98 | μs |
| Device Time | 40672.45 | μs |
| Self CPU Time | 950518.98 | μs |
| Self Device Time | 40672.45 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
45287 warnings generated when compiling for host. Suppressed 45325 warnings (45278 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.