53_Gemm_Scaling_Hardtanh_GELU • atomic_operations_optimized_baseimport 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 = 512;
// Atomic add function for float
__device__ inline float atomicAddFloat(float* address, float val) {
return atomicAdd(address, val);
}
// Atomic add function for double
__device__ inline double atomicAddDouble(double* address, double val) {
unsigned long long int* address_as_ull = (unsigned long long int*)address;
unsigned long long int old = *address_as_ull, assumed;
do {
assumed = old;
old = atomicCAS(address_as_ull, assumed,
__double_as_longlong(val + __longlong_as_double(assumed)));
} while (assumed != old);
return __longlong_as_double(old);
}
// Kernel function with atomic operations optimization
template <typename scalar_t>
__global__ void optimized_fused_kernel(
scalar_t* __restrict__ x,
scalar_t scaling_factor,
scalar_t hardtanh_min,
scalar_t hardtanh_max,
int64_t numel) {
int idx = blockIdx.x * BLOCK_SIZE + threadIdx.x;
int stride = gridDim.x * BLOCK_SIZE;
for (; idx < numel; idx += stride) {
scalar_t val = x[idx];
val *= scaling_factor;
val = fminf(fmaxf(val, hardtanh_min), hardtanh_max);
// GELU approximation
const scalar_t c = (scalar_t)0.044715;
const scalar_t sqrt_2_over_pi = (scalar_t)0.7978845608028654;
scalar_t x_cube = val * val * val;
scalar_t tanh_arg = sqrt_2_over_pi * (val + c * x_cube);
scalar_t tanh_res = tanh(tanh_arg);
val = 0.5 * val * (1.0 + tanh_res);
// Use atomic operation to update the value
if constexpr (std::is_same<scalar_t, float>::value) {
atomicAddFloat(&x[idx], val - x[idx]);
} else if constexpr (std::is_same<scalar_t, double>::value) {
atomicAddDouble(&x[idx], val - x[idx]);
} else {
x[idx] = val; // Fallback for other types
}
}
}
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", ([&] {
optimized_fused_kernel<scalar_t><<<blocks, BLOCK_SIZE>>>(
x.data_ptr<scalar_t>(),
(scalar_t)scaling_factor,
(scalar_t)hardtanh_min,
(scalar_t)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, "Optimized module function forward (CUDA)");
}| Metric | Value | Unit | Variance | Samples |
|---|---|---|---|---|
| Executed Ipc Active | 0.554 | inst/cycle | 0.001 | 5 |
| Executed Ipc Elapsed | 0.186 | inst/cycle | 0.000 | 5 |
| Issue Slots Busy | 14.296 | % | 0.856 | 5 |
| Issued Ipc Active | 0.572 | inst/cycle | 0.001 | 5 |
| SM Busy | 14.296 | % | 0.856 | 5 |
| Memory Throughput | 79971739396.428 | byte/second | 14146422227604330496.000 | 5 |
| Mem Busy | 11.154 | % | 0.275 | 5 |
| Max Bandwidth | 7.626 | % | 0.125 | 5 |
| L1/TEX Hit Rate | 0.000 | % | 0.000 | 5 |
| L2 Hit Rate | 82.582 | % | 0.118 | 5 |
| Mem Pipes Busy | 2.376 | % | 0.012 | 5 |
| Warp Cycles Per Issued Instruction | 23.630 | cycle | 0.042 | 5 |
| Warp Cycles Per Executed Instruction | 24.394 | cycle | 0.045 | 5 |
| Avg. Active Threads Per Warp | 30.190 | 0.000 | 5 | |
| Avg. Not Predicated Off Threads Per Warp | 28.620 | 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 | 4.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 | 22.906 | % | 0.067 | 5 |
| Achieved Active Warps Per SM | 14.660 | warp | 0.027 | 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.8%) 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 | 560026.75 | μs |
| Device Time | 196.19 | μs |
| Self CPU Time | 13017.92 | μ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 | 547008.84 | μs |
| Device Time | 196.19 | μs |
| Self CPU Time | 134.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 |
| aten::empty_strided | ||
| CPU Time | 546353.78 | μs |
| Device Time | 0.00 | μs |
| Self CPU Time | 128.83 | μ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 | 545806.13 | μs |
| Device Time | 0.00 | μs |
| Self CPU Time | 545806.13 | μ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 | 325779.47 | μs |
| Device Time | 171380.58 | μs |
| Self CPU Time | 12765.08 | μ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 | 313014.39 | μs |
| Device Time | 171380.58 | μs |
| Self CPU Time | 181399.03 | μs |
| Self Device Time | 171380.58 | μ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 | 558448.26 | μs |
| Device Time | 17621.84 | μs |
| Self CPU Time | 558448.26 | μs |
| Self Device Time | 17621.84 | μ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 | 515311.53 | μs |
| Device Time | 923579.02 | μs |
| Self CPU Time | 21107.65 | μ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 | 494208.36 | μs |
| Device Time | 923579.02 | μs |
| Self CPU Time | 25170.56 | μs |
| Self Device Time | 923579.02 | μ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 | 923579.02 | μs |
| Self CPU Time | 0.00 | μs |
| Self Device Time | 923579.02 | μ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.