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>
// This kernel uses shared memory to load a tile of the input matrix from global memory.
// Each block loads its tile into shared memory, computes swish activation, and writes the result back.
__global__ void shared_memory_swish_scaling_kernel(
const float* __restrict__ input,
float* __restrict__ output,
float scaling_factor,
int rows,
int cols) {
extern __shared__ float tile[];
// 2D thread indices within the block
int tx = threadIdx.x;
int ty = threadIdx.y;
// Global indices
int col = blockIdx.x * blockDim.x + tx;
int row = blockIdx.y * blockDim.y + ty;
// Index within the shared memory tile
int index = ty * blockDim.x + tx;
// Load data from global memory to shared memory (if within bounds)
if (row < rows && col < cols) {
tile[index] = input[row * cols + col];
}
// Synchronize to ensure the tile is fully loaded
__syncthreads();
// Process the data in shared memory and write back to global memory
if (row < rows && col < cols) {
float x = tile[index];
float sigmoid = 1.0f / (1.0f + expf(-x));
float res = x * sigmoid * scaling_factor;
output[row * cols + col] = res;
}
}
// The forward function performs the linear transformation (addmm) then launches the CUDA kernel
// The kernel uses a 2D grid and allocates shared memory to cache a tile of the matrix.
torch::Tensor forward(
torch::Tensor x,
torch::Tensor weight,
torch::Tensor bias,
double scaling_factor) {
// Ensure tensors are contiguous
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.");
TORCH_CHECK(weight.scalar_type() == at::kFloat, "Weight tensor must be of type torch.float32.");
TORCH_CHECK(bias.scalar_type() == at::kFloat, "Bias tensor must be of type torch.float32.");
// Compute linear transformation: y = x @ weight.T + bias
auto y = at::addmm(bias, x, weight.t());
auto output = at::empty_like(y);
int rows = y.size(0);
int cols = y.size(1);
// Configure a 2D grid of threads
dim3 threads(32, 32);
dim3 blocks((cols + threads.x - 1) / threads.x, (rows + threads.y - 1) / threads.y);
// Allocate shared memory size: blockDim.x * blockDim.y * sizeof(float)
size_t shared_mem_size = threads.x * threads.y * sizeof(float);
shared_memory_swish_scaling_kernel<<<blocks, threads, shared_mem_size>>>(
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, "Shared Memory CUDA forward function");
}
| Metric | Value | Unit | Variance | Samples |
|---|---|---|---|---|
| Executed Ipc Active | 0.613 | inst/cycle | 0.000 | 3 |
| Executed Ipc Elapsed | 0.110 | inst/cycle | 0.000 | 3 |
| Issue Slots Busy | 18.263 | % | 0.259 | 3 |
| Issued Ipc Active | 0.730 | inst/cycle | 0.000 | 3 |
| SM Busy | 18.263 | % | 0.259 | 3 |
| Memory Throughput | 78370589910.880 | byte/second | 863580889414865408.000 | 3 |
| Mem Busy | 10.833 | % | 0.012 | 3 |
| Max Bandwidth | 7.017 | % | 0.005 | 3 |
| L1/TEX Hit Rate | 0.000 | % | 0.000 | 3 |
| L2 Hit Rate | 82.713 | % | 0.040 | 3 |
| Mem Pipes Busy | 4.920 | % | 0.002 | 3 |
| Warp Cycles Per Issued Instruction | 43.693 | cycle | 3.157 | 3 |
| Warp Cycles Per Executed Instruction | 51.887 | cycle | 4.441 | 3 |
| Avg. Active Threads Per Warp | 32.000 | 0.000 | 3 | |
| Avg. Not Predicated Off Threads Per Warp | 31.180 | 0.000 | 3 | |
| Max Active Clusters | 0.000 | cluster | 0.000 | 3 |
| Max Cluster Size | 8.000 | block | 0.000 | 3 |
| Overall GPU Occupancy | 0.000 | % | 0.000 | 3 |
| Cluster Occupancy | 0.000 | % | 0.000 | 3 |
| Block Limit SM | 32.000 | block | 0.000 | 3 |
| Block Limit Registers | 4.000 | block | 0.000 | 3 |
| Block Limit Shared Mem | 3.000 | block | 0.000 | 3 |
| Block Limit Warps | 2.000 | block | 0.000 | 3 |
| Theoretical Active Warps per SM | 64.000 | warp | 0.000 | 3 |
| Theoretical Occupancy | 100.000 | % | 0.000 | 3 |
| Achieved Occupancy | 48.330 | % | 0.003 | 3 |
| Achieved Active Warps Per SM | 30.930 | warp | 0.001 | 3 |
| 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 (48.3%) 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 | 450528.05 | μs |
| Device Time | 221.54 | μs |
| Self CPU Time | 69.72 | μ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 | 450458.34 | μs |
| Device Time | 221.54 | μs |
| Self CPU Time | 130.02 | μ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 | 472585.07 | μs |
| Device Time | 0.00 | μs |
| Self CPU Time | 23024.62 | μ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 | 448975.55 | μs |
| Device Time | 0.00 | μs |
| Self CPU Time | 448975.55 | μ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 | 595594.85 | μs |
| Device Time | 150032.62 | μs |
| Self CPU Time | 222453.52 | μs |
| Self Device Time | 150032.62 | μ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 | 135252.56 | μs |
| Self CPU Time | 0.00 | μs |
| Self Device Time | 135252.56 | μ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 | 79012.19 | μs |
| Device Time | 702593.48 | μs |
| Self CPU Time | 14375.59 | μ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 | 64641.44 | μs |
| Device Time | 702593.48 | μs |
| Self CPU Time | 22196.69 | μs |
| Self Device Time | 702593.48 | μ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 | 702593.48 | μs |
| Self CPU Time | 0.00 | μs |
| Self Device Time | 702593.48 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
45286 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.