forked from ccf-ai-infra/GPUCodeForces
Merge pull request 'finish logish #18' (#248) from hli28146/GPUCodeForces:h18 into main
This commit is contained in:
commit
6d00a4f2cf
|
|
@ -0,0 +1,96 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
|
||||
torch::Tensor logish_cuda_forward(const torch::Tensor& input);
|
||||
"""
|
||||
|
||||
cuda_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <math.h>
|
||||
|
||||
struct __align__(16) Float4 {
|
||||
float x, y, z, w;
|
||||
};
|
||||
|
||||
__device__ __forceinline__ float logish_op(float x) {
|
||||
// f(x) = x * log(1 + sigmoid(x))
|
||||
// sigmoid(x) = 1 / (1 + exp(-x))
|
||||
float s = 1.0f / (1.0f + expf(-x));
|
||||
return x * logf(1.0f + s);
|
||||
}
|
||||
|
||||
__global__ void logish_kernel(
|
||||
const float* __restrict__ input,
|
||||
float* __restrict__ output,
|
||||
const int n_elements)
|
||||
{
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int stride = blockDim.x * gridDim.x;
|
||||
|
||||
int vec_loops = n_elements / 4;
|
||||
const Float4* vec_input = reinterpret_cast<const Float4*>(input);
|
||||
Float4* vec_output = reinterpret_cast<Float4*>(output);
|
||||
|
||||
for (int i = idx; i < vec_loops; i += stride) {
|
||||
Float4 in_val = vec_input[i];
|
||||
Float4 out_val;
|
||||
|
||||
out_val.x = logish_op(in_val.x);
|
||||
out_val.y = logish_op(in_val.y);
|
||||
out_val.z = logish_op(in_val.z);
|
||||
out_val.w = logish_op(in_val.w);
|
||||
|
||||
vec_output[i] = out_val;
|
||||
}
|
||||
|
||||
// 处理尾部剩余的元素
|
||||
int tail_start = vec_loops * 4;
|
||||
for (int i = tail_start + idx; i < n_elements; i += stride) {
|
||||
output[i] = logish_op(input[i]);
|
||||
}
|
||||
}
|
||||
|
||||
// C++ Wrapper
|
||||
torch::Tensor logish_cuda_forward(const torch::Tensor& input) {
|
||||
TORCH_CHECK(input.is_cuda(), "Input tensor must be a CUDA tensor");
|
||||
TORCH_CHECK(input.is_contiguous(), "Input tensor must be contiguous");
|
||||
|
||||
auto output = torch::empty_like(input);
|
||||
const int n_elements = input.numel();
|
||||
|
||||
const int block_size = 256;
|
||||
|
||||
int grid_size = (n_elements + block_size * 4 - 1) / (block_size * 4);
|
||||
|
||||
// 限制 Grid 大小以防止过度占用
|
||||
if (grid_size > 65535) grid_size = 65535;
|
||||
|
||||
logish_kernel<<<grid_size, block_size>>>(
|
||||
input.data_ptr<float>(),
|
||||
output.data_ptr<float>(),
|
||||
n_elements
|
||||
);
|
||||
|
||||
return output;
|
||||
}
|
||||
"""
|
||||
|
||||
class ModelNew(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super(ModelNew, self).__init__()
|
||||
# 即时编译 CUDA 代码
|
||||
self.op = load_inline(
|
||||
name='logish_op',
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=cuda_source,
|
||||
functions=['logish_cuda_forward'],
|
||||
verbose=False,
|
||||
extra_cuda_cflags=['-O3']
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.op.logish_cuda_forward(x)
|
||||
|
|
@ -0,0 +1,31 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
BATCH_SIZE = 4096
|
||||
DIM = 4096
|
||||
SHAPE = (BATCH_SIZE, DIM)
|
||||
|
||||
class Logish(nn.Module):
|
||||
"""
|
||||
公式: f(x) = x * log(1 + sigmoid(x))
|
||||
"""
|
||||
def __init__(self):
|
||||
super(Logish, self).__init__()
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return x * torch.log(1 + torch.sigmoid(x))
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self):
|
||||
super(Model, self).__init__()
|
||||
self.logish = Logish()
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.logish(x)
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(SHAPE, dtype=torch.float32)
|
||||
return [x.contiguous()]
|
||||
|
||||
def get_init_inputs():
|
||||
return []
|
||||
|
|
@ -0,0 +1,57 @@
|
|||
Write a custom CUDA kernel to optimize the Logish activation function.
|
||||
|
||||
The mathematical definition is:
|
||||
f(x) = x * log(1 + sigmoid(x))
|
||||
|
||||
Problem Analysis:
|
||||
The standard PyTorch implementation is memory-bound because it involves a chain of element-wise operations: sigmoid, addition, log, and multiplication.
|
||||
1. sigmoid(x) creates an intermediate tensor.
|
||||
2. 1 + temp adds overhead.
|
||||
3. log(temp) creates another intermediate tensor.
|
||||
4. x * temp creates the final output.
|
||||
This results in multiple read/write passes over the GPU global memory, creating a bandwidth bottleneck.
|
||||
|
||||
Optimization Strategy: Fused Element-wise Kernel with Vectorized Access
|
||||
|
||||
1. Operator Fusion: Create a single CUDA kernel that performs the entire mathematical calculation in registers for each element. This reduces global memory access to just one read and one write per element.
|
||||
|
||||
2. Vectorized Memory Access: Since this is a memory-bound operation, maximizing bandwidth is critical. We will use float4 data types to load and store 128 bits (4 floats) per instruction. This reduces the number of memory instructions and improves bus utilization.
|
||||
|
||||
3. Grid-Stride Loop: Implement the kernel using a grid-stride loop pattern to handle input tensors of arbitrary size efficiently, regardless of the grid dimension.
|
||||
|
||||
4. Numerical Implementation: Use fast hardware intrinsics where appropriate (e.g., expf, logf) to ensure the computation throughput matches the optimized memory bandwidth. The calculation will be performed as: s = 1 / (1 + exp(-x)); result = x * log(1 + s).
|
||||
|
||||
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:
|
||||
|
||||
```python
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
BATCH_SIZE = 4096
|
||||
DIM = 4096
|
||||
SHAPE = (BATCH_SIZE, DIM)
|
||||
|
||||
class Logish(nn.Module):
|
||||
"""
|
||||
公式: f(x) = x * log(1 + sigmoid(x))
|
||||
"""
|
||||
def __init__(self):
|
||||
super(Logish, self).__init__()
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return x * torch.log(1 + torch.sigmoid(x))
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self):
|
||||
super(Model, self).__init__()
|
||||
self.logish = Logish()
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.logish(x)
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(SHAPE, dtype=torch.float32)
|
||||
return [x.contiguous()]
|
||||
|
||||
def get_init_inputs():
|
||||
return []
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from logish_torch import Model,get_inputs,get_init_inputs
|
||||
from logish_cuda import ModelNew
|
||||
|
||||
def run_benchmark():
|
||||
# 检查 CUDA 是否可用
|
||||
if not torch.cuda.is_available():
|
||||
print("CUDA 不可用，请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
|
||||
return
|
||||
else:
|
||||
device = torch.device("cuda")
|
||||
|
||||
# 初始化模型
|
||||
init_inputs = get_init_inputs()
|
||||
init_inputs = [
|
||||
x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in init_inputs
|
||||
]
|
||||
inputs = get_inputs()
|
||||
inputs = [
|
||||
x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in inputs
|
||||
]
|
||||
|
||||
torch_model = Model(*init_inputs).cuda()
|
||||
cuda_model = ModelNew(*init_inputs).cuda()
|
||||
|
||||
torch_model.eval()
|
||||
cuda_model.eval()
|
||||
|
||||
print("-------------------- 精度对齐验证 --------------------")
|
||||
with torch.no_grad():
|
||||
output_torch = torch_model( *inputs)
|
||||
output_cuda = cuda_model(*inputs)
|
||||
|
||||
precision_flag = torch.allclose(output_torch, output_cuda,rtol=1e-03)
|
||||
if precision_flag:
|
||||
print("✅ 精度对齐：两个模型的输出结果非常接近。")
|
||||
else:
|
||||
print("❌ 精度不一致！")
|
||||
|
||||
print("\n-------------------- 性能加速比测试 --------------------")
|
||||
num_iterations = 100
|
||||
|
||||
# PyTorch 模型计时
|
||||
torch.cuda.synchronize()
|
||||
start_time = time.time()
|
||||
for _ in range(num_iterations):
|
||||
_ = torch_model(*inputs)
|
||||
torch.cuda.synchronize()
|
||||
torch_time = (time.time() - start_time) / num_iterations
|
||||
|
||||
# 自定义 CUDA 内核计时
|
||||
torch.cuda.synchronize()
|
||||
start_time = time.time()
|
||||
for _ in range(num_iterations):
|
||||
_ = cuda_model(*inputs)
|
||||
torch.cuda.synchronize()
|
||||
cuda_time = (time.time() - start_time) / num_iterations
|
||||
|
||||
print(f"PyTorch torch.relu 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA 内核 平均执行时间: {cuda_time:.6f} 秒")
|
||||
speedup = 0
|
||||
if cuda_time > 0:
|
||||
speedup = torch_time / cuda_time
|
||||
print(f"加速比 (Speedup): {speedup:.2f}x")
|
||||
else:
|
||||
print("CUDA 内核执行时间为0，无法计算加速比。")
|
||||
return precision_flag,speedup
|
||||
if __name__ == "__main__":
|
||||
precision_flag,speedup = run_benchmark()
|
||||
Loading…
Reference in New Issue