import torch
import torch.nn as nn
import torch.nn.functional as F
def module_fn(
x: torch.Tensor,
weight: torch.Tensor,
bias: torch.Tensor,
scaling_factor: float,
) -> torch.Tensor:
"""
Applies linear transformation, Swish activation, and scaling.
Args:
x (torch.Tensor): Input tensor of shape (batch_size, in_features)
weight (torch.Tensor): Weight matrix of shape (out_features, in_features)
bias (torch.Tensor): Bias vector of shape (out_features)
scaling_factor (float): Factor to scale the output by
Returns:
torch.Tensor: Output tensor of shape (batch_size, out_features)
"""
x = F.linear(x, weight, bias)
x = x * torch.sigmoid(x) # Swish activation
x = x * scaling_factor
return x
class Model(nn.Module):
"""
Simple model that performs a matrix multiplication, applies Swish activation, and scales the result.
"""
def __init__(self, in_features, out_features, scaling_factor):
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
def forward(self, x, fn=module_fn):
return fn(x, self.weight, self.bias, self.scaling_factor)
batch_size = 128
in_features = 1024
out_features = 512
scaling_factor = 2.0
def get_inputs():
return [torch.randn(batch_size, in_features)]
def get_init_inputs():
return [in_features, out_features, scaling_factor]
import torch
import torch.nn as nn
class Model(nn.Module):
"""
Simple model that performs a matrix multiplication, applies Swish activation, and scales the result.
"""
def __init__(self, in_features, out_features, scaling_factor):
super(Model, self).__init__()
self.matmul = nn.Linear(in_features, out_features)
self.scaling_factor = scaling_factor
def forward(self, x):
x = self.matmul(x)
x = x * torch.sigmoid(x) # Swish activation
x = x * self.scaling_factor
return x
batch_size = 128
in_features = 1024
out_features = 512
scaling_factor = 2.0
def get_inputs():
return [torch.randn(batch_size, in_features)]
def get_init_inputs():
return [in_features, out_features, scaling_factor]#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
template<bool UseVectorization>
__global__ void adaptive_swish_scaling_kernel(
const float* __restrict__ input,
float* output,
float scaling_factor,
int rows,
int cols) {
if constexpr(UseVectorization) {
// Vectorized version for large contiguous data
const int tid = threadIdx.x;
const int elements_per_thread = 4;
const int block_elements = blockDim.x * elements_per_thread;
int base_idx = blockIdx.x * block_elements + tid * elements_per_thread;
const int N = rows * cols;
extern __shared__ float shared_data[];
if (base_idx < N) {
float4* in_vec = (float4*)(&input[base_idx]);
float4* shared_vec = (float4*)(&shared_data[tid * elements_per_thread]);
if (base_idx + 3 < N) {
*shared_vec = *in_vec;
} else {
for (int i = 0; i < elements_per_thread && (base_idx + i) < N; i++) {
shared_data[tid * elements_per_thread + i] = input[base_idx + i];
}
}
}
__syncthreads();
if (base_idx < N) {
float4 result;
float* result_f = (float*)&result;
#pragma unroll
for (int i = 0; i < elements_per_thread && (base_idx + i) < N; i++) {
float x = shared_data[tid * elements_per_thread + i];
float sigmoid = 1.0f / (1.0f + expf(-x));
result_f[i] = x * sigmoid * scaling_factor;
}
if (base_idx + 3 < N) {
float4* out_vec = (float4*)(&output[base_idx]);
*out_vec = result;
} else {
for (int i = 0; i < elements_per_thread && (base_idx + i) < N; i++) {
output[base_idx + i] = result_f[i];
}
}
}
} else {
// 2D version for better cache locality with matrix operations
int row = blockIdx.y * blockDim.y + threadIdx.y;
int col = blockIdx.x * blockDim.x + threadIdx.x;
if (row < rows && col < cols) {
int idx = row * cols + col;
float x = input[idx];
float sigmoid = 1.0f / (1.0f + expf(-x));
output[idx] = x * sigmoid * scaling_factor;
}
}
}
torch::Tensor forward(
torch::Tensor x,
torch::Tensor weight,
torch::Tensor bias,
double scaling_factor) {
x = x.contiguous();
weight = weight.contiguous();
bias = bias.contiguous();
TORCH_CHECK(x.is_cuda(), "Input tensor 'x' must be a CUDA tensor.");
TORCH_CHECK(weight.is_cuda(), "Weight tensor must be a CUDA tensor.");
TORCH_CHECK(bias.is_cuda(), "Bias tensor must be a CUDA tensor.");
TORCH_CHECK(x.scalar_type() == at::kFloat, "Input tensor 'x' must be of type torch.float32.");
auto y = at::addmm(bias, x, weight.t());
auto output = at::empty_like(y);
const int rows = y.size(0);
const int cols = y.size(1);
const int total_elements = rows * cols;
// Choose kernel configuration based on input size and shape
if (cols >= 512 && rows == 1) {
// Use vectorized version for large 1D operations
const int threads = 256;
const int elements_per_thread = 4;
const int elements_per_block = threads * elements_per_thread;
const int blocks = (total_elements + elements_per_block - 1) / elements_per_block;
const size_t shared_mem_size = threads * elements_per_thread * sizeof(float);
adaptive_swish_scaling_kernel<true><<<blocks, threads, shared_mem_size>>>(
y.data_ptr<float>(),
output.data_ptr<float>(),
static_cast<float>(scaling_factor),
rows,
cols);
} else {
// Use 2D version for matrix operations
dim3 threads(32, 32);
dim3 blocks((cols + threads.x - 1) / threads.x, (rows + threads.y - 1) / threads.y);
adaptive_swish_scaling_kernel<false><<<blocks, threads, 0>>>(
y.data_ptr<float>(),
output.data_ptr<float>(),
static_cast<float>(scaling_factor),
rows,
cols);
}
cudaError_t err = cudaGetLastError();
TORCH_CHECK(err == cudaSuccess, "CUDA kernel failed : ", cudaGetErrorString(err));
return output;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("forward", &forward, "Adaptive CUDA forward function");
}| Metric | Value | Unit | Variance | Samples |
|---|---|---|---|---|
| Executed Ipc Active | 0.554 | inst/cycle | 0.002 | 5 |
| Executed Ipc Elapsed | 0.090 | inst/cycle | 0.000 | 5 |
| Issue Slots Busy | 15.574 | % | 1.223 | 5 |
| Issued Ipc Active | 0.626 | inst/cycle | 0.002 | 5 |
| SM Busy | 15.574 | % | 1.223 | 5 |
| Memory Throughput | 81229030763.742 | byte/second | 4227931271296517120.000 | 5 |
| Mem Busy | 11.302 | % | 0.097 | 5 |
| Max Bandwidth | 7.316 | % | 0.039 | 5 |
| L1/TEX Hit Rate | 0.000 | % | 0.000 | 5 |
| L2 Hit Rate | 82.594 | % | 0.055 | 5 |
| Mem Pipes Busy | 4.690 | % | 0.018 | 5 |
| Warp Cycles Per Issued Instruction | 43.758 | cycle | 5.720 | 5 |
| Warp Cycles Per Executed Instruction | 49.274 | cycle | 7.246 | 5 |
| Avg. Active Threads Per Warp | 32.000 | 0.000 | 5 | |
| Avg. Not Predicated Off Threads Per Warp | 30.930 | 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 | 4.000 | block | 0.000 | 5 |
| Block Limit Shared Mem | 8.000 | block | 0.000 | 5 |
| Block Limit Warps | 2.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 | 44.076 | % | 1.755 | 5 |
| Achieved Active Warps Per SM | 28.208 | warp | 0.718 | 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 (44.4%) 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 | 292302.63 | μs |
| Device Time | 197.37 | μs |
| Self CPU Time | 54.53 | μ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 | 292248.11 | μs |
| Device Time | 197.37 | μs |
| Self CPU Time | 97.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::empty_strided | ||
| CPU Time | 313876.61 | μs |
| Device Time | 0.00 | μs |
| Self CPU Time | 22355.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 |
| cudaDeviceGetStreamPriorityRange | ||
| CPU Time | 291097.07 | μs |
| Device Time | 0.00 | μs |
| Self CPU Time | 291097.07 | μ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::addmm | ||
| CPU Time | 530572.51 | μs |
| Device Time | 131394.49 | μs |
| Self CPU Time | 186759.39 | μs |
| Self Device Time | 131394.49 | μ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_gemm_f32f32_f32f32_f32_tn_n_tilesize32x32x8_stage3_warpsize1x2x1_ffma_aligna4_alignc4_execute_kernel__51_cublas | ||
| CPU Time | 0.00 | μs |
| Device Time | 118430.95 | μs |
| Self CPU Time | 0.00 | μs |
| Self Device Time | 118430.95 | μ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 | 66377.80 | μs |
| Device Time | 615412.36 | μs |
| Self CPU Time | 12624.33 | μ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 | 53754.99 | μs |
| Device Time | 615412.36 | μs |
| Self CPU Time | 17971.74 | μs |
| Self Device Time | 615412.36 | μ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 | 615412.36 | μs |
| Self CPU Time | 0.00 | μs |
| Self Device Time | 615412.36 | μ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 45324 warnings (45277 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.