forked from ccf-ai-infra/GPUCodeForces
finish lppool1d #6
This commit is contained in:
parent
f876a28ada
commit
dec6fce985
|
|
@ -0,0 +1,121 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
import math
|
||||
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
|
||||
torch::Tensor lppool1d_cuda_forward(
|
||||
const torch::Tensor& input,
|
||||
int norm_power,
|
||||
int kernel_size,
|
||||
int stride
|
||||
);
|
||||
"""
|
||||
|
||||
# CUDA 源代码
|
||||
cuda_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cmath>
|
||||
|
||||
template <typename T>
|
||||
__global__ void lppool1d_kernel(
|
||||
T* output,
|
||||
const T* input,
|
||||
const T norm_power,
|
||||
const int kernel_size,
|
||||
const int stride,
|
||||
const int N, const int C, const int L_in, const int L_out,
|
||||
const int64_t n_elements_out)
|
||||
{
|
||||
int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx >= n_elements_out) return;
|
||||
|
||||
// --- 1D 输出索引映射回 3D 输出坐标 (n, c, l_out) ---
|
||||
const int l_out = idx % L_out;
|
||||
const int c = (idx / L_out) % C;
|
||||
const int n = idx / (L_out * C);
|
||||
|
||||
// --- 计算输入窗口的范围 ---
|
||||
const int start = l_out * stride;
|
||||
const int end = fminf(start + kernel_size, L_in);
|
||||
|
||||
// --- 在寄存器中循环累加 ---
|
||||
T sum_val = 0.0f;
|
||||
const T* input_channel_ptr = input + (n * C + c) * L_in;
|
||||
|
||||
for (int k = start; k < end; ++k) {
|
||||
sum_val += powf(input_channel_ptr[k], norm_power);
|
||||
}
|
||||
|
||||
// --- 最终计算并写入结果 ---
|
||||
output[idx] = powf(sum_val, 1.0f / norm_power);
|
||||
}
|
||||
|
||||
torch::Tensor lppool1d_cuda_forward(
|
||||
const torch::Tensor& input,
|
||||
const int norm_power,
|
||||
const int kernel_size,
|
||||
const int stride)
|
||||
{
|
||||
TORCH_CHECK(input.is_cuda(), "Input tensor must be a CUDA tensor");
|
||||
TORCH_CHECK(input.dim() == 3, "Input tensor must be 3D");
|
||||
TORCH_CHECK(input.is_contiguous(), "Input tensor must be contiguous");
|
||||
|
||||
const int N = input.size(0);
|
||||
const int C = input.size(1);
|
||||
const int L_in = input.size(2);
|
||||
|
||||
const int L_out = floor(((float)L_in - kernel_size) / stride) + 1;
|
||||
|
||||
auto output = torch::empty({N, C, L_out}, input.options());
|
||||
const int64_t n_elements_out = output.numel();
|
||||
|
||||
if (n_elements_out == 0) {
|
||||
return output;
|
||||
}
|
||||
|
||||
const int block_size = 256;
|
||||
const int num_blocks = (n_elements_out + block_size - 1) / block_size;
|
||||
|
||||
AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "lppool1d_kernel", ([&] {
|
||||
lppool1d_kernel<scalar_t><<<num_blocks, block_size>>>(
|
||||
output.data_ptr<scalar_t>(),
|
||||
input.data_ptr<scalar_t>(),
|
||||
static_cast<scalar_t>(norm_power),
|
||||
kernel_size, stride,
|
||||
N, C, L_in, L_out, n_elements_out
|
||||
);
|
||||
}));
|
||||
|
||||
return output;
|
||||
}
|
||||
"""
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
"""
|
||||
使用自定义 CUDA 内核进行优化的 LPPool1d 模型。
|
||||
"""
|
||||
def __init__(self, norm_power, kernel_size, stride):
|
||||
super(ModelNew, self).__init__()
|
||||
self.norm_power = norm_power
|
||||
self.kernel_size = kernel_size
|
||||
self.stride = stride
|
||||
|
||||
self.op = load_inline(
|
||||
name='lppool1d_op',
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=cuda_source,
|
||||
functions=['lppool1d_cuda_forward'],
|
||||
verbose=False
|
||||
)
|
||||
|
||||
def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
|
||||
return self.op.lppool1d_cuda_forward(
|
||||
input_tensor.contiguous(), # 确保输入是连续的
|
||||
self.norm_power,
|
||||
self.kernel_size,
|
||||
self.stride
|
||||
)
|
||||
|
|
@ -0,0 +1,28 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
BATCH_SIZE = 512
|
||||
CHANNELS = 256
|
||||
L_IN = 1024
|
||||
NORM_POWER = 2
|
||||
KERNEL_SIZE = 3
|
||||
STRIDE = 1
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, norm_power, kernel_size, stride):
|
||||
super(Model, self).__init__()
|
||||
self.pool = nn.LPPool1d(norm_type=norm_power, kernel_size=kernel_size, stride=stride)
|
||||
|
||||
def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
|
||||
return self.pool(input_tensor)
|
||||
|
||||
def get_inputs():
|
||||
"""
|
||||
生成用于测试的输入张量。
|
||||
"""
|
||||
# LPPool 对负值敏感，使用正值以保证与 powf 的行为一致
|
||||
input_tensor = torch.rand(BATCH_SIZE, CHANNELS, L_IN, dtype=torch.float32) + 0.1
|
||||
return [input_tensor.contiguous()]
|
||||
|
||||
def get_init_inputs():
|
||||
return [NORM_POWER, KERNEL_SIZE, STRIDE]
|
||||
|
|
@ -0,0 +1,80 @@
|
|||
Write a custom CUDA kernel to optimize `torch.nn.LPPool1d`.
|
||||
|
||||
The operation performs 1D LP-norm pooling over an input signal. For each sliding window, it computes `(sum(x^p))^(1/p)`, where `p` is the `norm_power`.
|
||||
|
||||
**Problem Analysis:**
|
||||
The standard PyTorch implementation of pooling layers often relies on an `unfold` (or `im2col`) operation to extract the sliding windows. This approach has significant performance drawbacks:
|
||||
1. **Memory Explosion**: The `unfold` operation creates a massive intermediate tensor containing all extracted windows. For a 1D signal, this can increase memory usage by a factor of `kernel_size`, becoming a major memory bandwidth bottleneck.
|
||||
2. **Multiple Kernel Launches**: After unfolding, a sequence of separate element-wise and reduction kernels are launched (`pow`, `sum`, `pow`), each requiring a full pass over the data and incurring kernel launch latency.
|
||||
|
||||
**Optimization Strategy: Fused Output-Oriented Kernel**
|
||||
|
||||
The optimization strategy is to create a single, fused CUDA kernel that computes the pooling result directly, avoiding the `unfold` operation entirely.
|
||||
|
||||
1. **Output-Oriented Parallelism**: The kernel is launched with one thread for each element of the **output tensor**. Each thread is uniquely responsible for computing one final output value.
|
||||
|
||||
2. **Direct Window Computation**: Each thread first calculates its position `(n, c, l_out)` in the output tensor. From this, it computes the corresponding window's start and end indices in the input tensor based on `stride` and `kernel_size`.
|
||||
|
||||
3. **In-Register Reduction**: The thread then loops over the elements of its assigned input window. The entire LP-norm calculation (`pow(x, p)`, summation) is performed within the thread's private registers. This is extremely fast and completely avoids writing any intermediate data to global memory.
|
||||
|
||||
4. **Full Fusion**: After the loop, the final `pow(1/p)` operation is applied, and the thread writes the single, final result to its designated position in the output tensor. This approach fuses the `unfold`, `pow`, `sum`, and final `pow` operations into a single memory pass, dramatically improving performance by minimizing memory traffic and kernel launch overhead. It also benefits from good data locality, as adjacent threads access overlapping regions of the input tensor, leading to efficient cache utilization.
|
||||
|
||||
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
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
||||
def forward(self, a, b):
|
||||
return a + b
|
||||
|
||||
|
||||
def get_inputs():
|
||||
# randomly generate input tensors based on the model architecture
|
||||
a = torch.randn(1, 128).cuda()
|
||||
b = torch.randn(1, 128).cuda()
|
||||
return [a, b]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
# randomly generate tensors required for initialization based on the model architecture
|
||||
return []
|
||||
```
|
||||
|
||||
The example new arch with custom CUDA kernels looks like this:
|
||||
```python
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
BATCH_SIZE = 512
|
||||
CHANNELS = 256
|
||||
L_IN = 1024
|
||||
NORM_POWER = 2
|
||||
KERNEL_SIZE = 3
|
||||
STRIDE = 1
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, norm_power, kernel_size, stride):
|
||||
super(Model, self).__init__()
|
||||
self.pool = nn.LPPool1d(norm_type=norm_power, kernel_size=kernel_size, stride=stride)
|
||||
|
||||
def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
|
||||
return self.pool(input_tensor)
|
||||
|
||||
def get_inputs():
|
||||
"""
|
||||
生成用于测试的输入张量。
|
||||
"""
|
||||
# LPPool 对负值敏感，使用正值以保证与 powf 的行为一致
|
||||
input_tensor = torch.rand(BATCH_SIZE, CHANNELS, L_IN, dtype=torch.float32) + 0.1
|
||||
return [input_tensor.contiguous()]
|
||||
|
||||
def get_init_inputs():
|
||||
return [NORM_POWER, KERNEL_SIZE, STRIDE]
|
||||
```
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from lppool1d_torch import Model,get_inputs,get_init_inputs
|
||||
from lppool1d_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