53_Gemm_Scaling_Hardtanh_GELU
• shared_mem_optimized_kernel_base
import torch
import torch.nn as nn
import torch.nn.functional as F
def module_fn(
x: torch.Tensor,
scaling_factor: float,
hardtanh_min: float,
hardtanh_max: float,
weight: torch.Tensor,
bias: torch.Tensor,
) -> torch.Tensor:
"""
Applies GEMM, scaling, hardtanh and GELU activation.
Args:
x (torch.Tensor): Input tensor of shape (batch_size, in_features)
scaling_factor (float): Factor to scale the GEMM output
hardtanh_min (float): Minimum value for hardtanh
hardtanh_max (float): Maximum value for hardtanh
weight (torch.Tensor): Weight matrix of shape (out_features, in_features)
bias (torch.Tensor): Bias vector of shape (out_features)
Returns:
torch.Tensor: Output tensor after applying GEMM, scaling, hardtanh and GELU,
with shape (batch_size, out_features)
"""
x = F.linear(x, weight, bias)
x = x * scaling_factor
x = F.hardtanh(x, min_val=hardtanh_min, max_val=hardtanh_max)
x = F.gelu(x)
return x
class Model(nn.Module):
"""
Model that performs a GEMM, scaling, hardtanh, and GELU activation.
"""
def __init__(
self, in_features, out_features, scaling_factor, hardtanh_min, hardtanh_max
):
super(Model, self).__init__()
gemm = nn.Linear(in_features, out_features)
self.weight = nn.Parameter(gemm.weight)
self.bias = nn.Parameter(gemm.bias)
self.scaling_factor = scaling_factor
self.hardtanh_min = hardtanh_min
self.hardtanh_max = hardtanh_max
def forward(self, x, fn=module_fn):
return fn(
x,
self.scaling_factor,
self.hardtanh_min,
self.hardtanh_max,
self.weight,
self.bias,
)
batch_size = 128
in_features = 1024
out_features = 512
scaling_factor = 0.5
hardtanh_min = -2
hardtanh_max = 2
def get_inputs():
return [torch.randn(batch_size, in_features)]
def get_init_inputs():
return [in_features, out_features, scaling_factor, hardtanh_min, hardtanh_max]
import torch
import torch.nn as nn
class Model(nn.Module):
"""
Model that performs a GEMM, scaling, hardtanh, and GELU activation.
"""
def __init__(self, in_features, out_features, scaling_factor, hardtanh_min, hardtanh_max):
super(Model, self).__init__()
self.gemm = nn.Linear(in_features, out_features)
self.scaling_factor = scaling_factor
self.hardtanh = nn.Hardtanh(min_val=hardtanh_min, max_val=hardtanh_max)
self.gelu = nn.GELU()
def forward(self, x):
x = self.gemm(x)
x = x * self.scaling_factor
x = self.hardtanh(x)
x = self.gelu(x)
return x
batch_size = 128
in_features = 1024
out_features = 512
scaling_factor = 0.5
hardtanh_min = -2
hardtanh_max = 2
def get_inputs():
return [torch.randn(batch_size, in_features)]
def get_init_inputs():
return [in_features, out_features, scaling_factor, hardtanh_min, hardtanh_max]
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
constexpr int BLOCK_SIZE = 256;
__device__ inline float scale_value(float val, float factor) {
return val * factor;
}
__device__ inline float hard_tanh(float val, float min_val, float max_val) {
return fminf(fmaxf(val, min_val), max_val);
}
__device__ inline float gelu_activation(float val) {
constexpr float c = 0.044715;
constexpr float sqrt_2_over_pi = 0.7978845608028654;
float x_cube = val * val * val;
float tanh_arg = sqrt_2_over_pi * (val + c * x_cube);
float tanh_res = tanh(tanh_arg);
return 0.5 * val * (1.0 + tanh_res);
}
__global__ void fused_activation_with_shared_memory(
float* __restrict__ x,
const float scaling_factor,
const float hardtanh_min,
const float hardtanh_max,
const int64_t numel) {
__shared__ float shared_data[BLOCK_SIZE];
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int stride = gridDim.x * blockDim.x;
for (int i = idx; i < numel; i += stride) {
// Load data into shared memory
float val = x[i];
// Use shared memory for subsequent operations
shared_data[threadIdx.x] = scale_value(val, scaling_factor);
shared_data[threadIdx.x] = hard_tanh(
shared_data[threadIdx.x], hardtanh_min, hardtanh_max);
val = gelu_activation(shared_data[threadIdx.x]);
// Write back the computed value
x[i] = val;
}
}
void fused_activation_cuda(
torch::Tensor& x,
double scaling_factor,
double hardtanh_min,
double hardtanh_max) {
const auto numel = x.numel();
const dim3 blocks((numel + BLOCK_SIZE - 1) / BLOCK_SIZE);
AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), "fused_activation_cuda", ([&] {
fused_activation_with_shared_memory<<<blocks, BLOCK_SIZE>>>(
x.data_ptr<float>(),
static_cast<float>(scaling_factor),
static_cast<float>(hardtanh_min),
static_cast<float>(hardtanh_max),
numel);
}));
}
torch::Tensor module_fn_forward(
torch::Tensor x,
double scaling_factor,
double hardtanh_min,
double hardtanh_max,
torch::Tensor weight,
torch::Tensor bias) {
x = x.contiguous().cuda();
weight = weight.contiguous().cuda();
bias = bias.contiguous().cuda();
auto xw = torch::matmul(x, weight.t()) + bias;
fused_activation_cuda(xw, scaling_factor, hardtanh_min, hardtanh_max);
return xw;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("forward", &module_fn_forward, "Module function forward (CUDA)");
}
Metric | Value | Unit | Variance | Samples |
---|---|---|---|---|
Executed Ipc Active | 0.504 | inst/cycle | 0.000 | 5 |
Executed Ipc Elapsed | 0.192 | inst/cycle | 0.000 | 5 |
Issue Slots Busy | 13.518 | % | 0.019 | 5 |
Issued Ipc Active | 0.542 | inst/cycle | 0.000 | 5 |
SM Busy | 13.518 | % | 0.019 | 5 |
Memory Throughput | 74711585836.802 | byte/second | 2680222953560535040.000 | 5 |
Mem Busy | 10.350 | % | 0.037 | 5 |
Max Bandwidth | 6.718 | % | 0.013 | 5 |
L1/TEX Hit Rate | 50.000 | % | 0.000 | 5 |
L2 Hit Rate | 82.732 | % | 0.045 | 5 |
Mem Pipes Busy | 2.218 | % | 0.002 | 5 |
Warp Cycles Per Issued Instruction | 25.482 | cycle | 0.099 | 5 |
Warp Cycles Per Executed Instruction | 27.324 | cycle | 0.114 | 5 |
Avg. Active Threads Per Warp | 30.550 | 0.000 | 5 | |
Avg. Not Predicated Off Threads Per Warp | 29.160 | 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 | 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 | 21.910 | % | 0.011 | 5 |
Achieved Active Warps Per SM | 14.022 | warp | 0.005 | 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 (22.0%) 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 | 637639.96 | μs |
Device Time | 184.28 | μs |
Self CPU Time | 14149.32 | μ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 | 623490.64 | μs |
Device Time | 184.28 | μs |
Self CPU Time | 108.81 | μ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 | 622873.71 | μs |
Device Time | 0.00 | μs |
Self CPU Time | 116.96 | μ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 | 620626.95 | μs |
Device Time | 0.00 | μs |
Self CPU Time | 620626.95 | μ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::matmul | ||
CPU Time | 331408.72 | μs |
Device Time | 169624.04 | μs |
Self CPU Time | 12403.87 | μ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::mm | ||
CPU Time | 319004.85 | μs |
Device Time | 169624.04 | μs |
Self CPU Time | 178649.25 | μs |
Self Device Time | 169624.04 | μ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 | 545675.72 | μs |
Device Time | 17437.62 | μs |
Self CPU Time | 545675.72 | μs |
Self Device Time | 17437.62 | μ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 | 496378.41 | μs |
Device Time | 914002.06 | μs |
Self CPU Time | 21134.17 | μ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::fill_ | ||
CPU Time | 475245.87 | μs |
Device Time | 914002.06 | μs |
Self CPU Time | 24958.92 | μs |
Self Device Time | 914002.06 | μs |
CPU Memory Usage | 0 | B |
Device Memory Usage | 0 | B |
Self CPU Memory Usage | 0 | B |
Self Device Memory Usage | 0 | B |
void at::native::vectorized_elementwise_kernel<4, at::native::FillFunctor<int>, at::detail::Array<char*, 1> >(int, at::native::FillFunctor<int>, at::detail::Array<char*, 1>) | ||
CPU Time | 0.00 | μs |
Device Time | 914002.06 | μs |
Self CPU Time | 0.00 | μs |
Self Device Time | 914002.06 | μ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 45326 warnings (45279 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.