forked from ccf-ai-infra/GPUCodeForces
Merge pull request 'finish NLLLoss# 32' (#59) from uucoco/GPUCodeForces:NLLLoss into main
This commit is contained in:
commit
1ca8c7d24c
|
|
@ -0,0 +1,280 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
import math
|
||||
|
||||
# -------------------------------------------------------------
|
||||
# 常量定义 (同上)
|
||||
# -------------------------------------------------------------
|
||||
N, C, H, W = 8, 10, 16, 16
|
||||
# 确保 weight 在 CUDA 上
|
||||
WEIGHT = torch.rand(C, dtype=torch.float32).cuda()
|
||||
IGNORE_INDEX = -100
|
||||
REDUCTION = 'mean'
|
||||
BLOCK_SIZE = 256
|
||||
|
||||
|
||||
# -------------------------------------------------------------
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
|
||||
def __init__(self, weight=None, size_average=None, ignore_index=-100,
|
||||
reduce=None, reduction='mean'):
|
||||
super().__init__()
|
||||
self.reduction_str = reduction
|
||||
self.ignore_index = ignore_index
|
||||
|
||||
if weight is not None:
|
||||
self.register_buffer('weight', weight.contiguous())
|
||||
else:
|
||||
self.weight = None
|
||||
|
||||
self.block_size = BLOCK_SIZE
|
||||
self._compile_cuda_kernel()
|
||||
|
||||
def _compile_cuda_kernel(self):
|
||||
|
||||
cpp_header = f"""
|
||||
#include <torch/extension.h>
|
||||
|
||||
// C++ 接口
|
||||
torch::Tensor nll_loss_forward_cuda(
|
||||
torch::Tensor input,
|
||||
torch::Tensor target,
|
||||
c10::optional<torch::Tensor> weight,
|
||||
int64_t ignore_index,
|
||||
std::string reduction
|
||||
);
|
||||
"""
|
||||
|
||||
cuda_source = f"""
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#define BLOCK_SIZE {self.block_size}
|
||||
|
||||
/* * 辅助函数：计算单个元素的 loss 和 weight
|
||||
* (在两个核函数之间共享)
|
||||
*/
|
||||
__device__ inline void compute_nll_loss_item(
|
||||
const float* __restrict__ input_data,
|
||||
const int64_t* __restrict__ target_data,
|
||||
const float* __restrict__ weight_data,
|
||||
int i, int N_spatial, int C, int64_t ignore_index,
|
||||
float& loss_val, // 输出
|
||||
float& weight_val // 输出
|
||||
) {{
|
||||
const int64_t target_idx = target_data[i];
|
||||
|
||||
weight_val = 1.0f;
|
||||
loss_val = 0.0f;
|
||||
|
||||
if (target_idx == ignore_index) {{
|
||||
weight_val = 0.0f;
|
||||
}} else if (target_idx < 0 || target_idx >= C) {{
|
||||
weight_val = 0.0f;
|
||||
}} else {{
|
||||
if (weight_data != nullptr) {{
|
||||
weight_val = weight_data[target_idx];
|
||||
}}
|
||||
|
||||
const int n = i / N_spatial;
|
||||
const int s = i % N_spatial;
|
||||
// const int s = i - n * N_spatial; // 优化的 modulo
|
||||
|
||||
const int64_t input_idx =
|
||||
(int64_t)n * C * N_spatial +
|
||||
(int64_t)target_idx * N_spatial +
|
||||
(int64_t)s;
|
||||
|
||||
loss_val = -weight_val * input_data[input_idx];
|
||||
}}
|
||||
}}
|
||||
|
||||
|
||||
__global__ void nll_loss_kernel_no_reduce(
|
||||
const float* __restrict__ input_data,
|
||||
const int64_t* __restrict__ target_data,
|
||||
const float* __restrict__ weight_data,
|
||||
float* __restrict__ loss_out_data,
|
||||
int N_spatial, int C, int N_total, int64_t ignore_index
|
||||
) {{
|
||||
for (int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
i < N_total;
|
||||
i += gridDim.x * blockDim.x)
|
||||
{{
|
||||
float loss_val, weight_val;
|
||||
compute_nll_loss_item(
|
||||
input_data, target_data, weight_data,
|
||||
i, N_spatial, C, ignore_index,
|
||||
loss_val, weight_val
|
||||
);
|
||||
// 写入未归约的损失
|
||||
loss_out_data[i] = loss_val;
|
||||
}}
|
||||
}}
|
||||
|
||||
|
||||
__global__ void nll_loss_kernel_reduce(
|
||||
const float* __restrict__ input_data,
|
||||
const int64_t* __restrict__ target_data,
|
||||
const float* __restrict__ weight_data,
|
||||
float* __restrict__ partial_loss_out, // Block-level
|
||||
float* __restrict__ partial_weight_out, // Block-level
|
||||
int N_spatial, int C, int N_total, int64_t ignore_index
|
||||
) {{
|
||||
// 共享内存用于 Block 内部归约
|
||||
__shared__ float s_loss[BLOCK_SIZE];
|
||||
__shared__ float s_weight[BLOCK_SIZE];
|
||||
|
||||
const int tid = threadIdx.x;
|
||||
float thread_loss_sum = 0.0f;
|
||||
float thread_weight_sum = 0.0f;
|
||||
|
||||
// 1. Grid-Stride Loop: 计算和累加
|
||||
for (int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
i < N_total;
|
||||
i += gridDim.x * blockDim.x)
|
||||
{{
|
||||
float loss_val, weight_val;
|
||||
compute_nll_loss_item(
|
||||
input_data, target_data, weight_data,
|
||||
i, N_spatial, C, ignore_index,
|
||||
loss_val, weight_val
|
||||
);
|
||||
thread_loss_sum += loss_val;
|
||||
thread_weight_sum += weight_val;
|
||||
}}
|
||||
|
||||
s_loss[tid] = thread_loss_sum;
|
||||
s_weight[tid] = thread_weight_sum;
|
||||
__syncthreads();
|
||||
|
||||
// 2. 共享内存归约 (Sum)
|
||||
for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) {{
|
||||
if (tid < offset) {{
|
||||
s_loss[tid] += s_loss[tid + offset];
|
||||
s_weight[tid] += s_weight[tid + offset];
|
||||
}}
|
||||
__syncthreads();
|
||||
}}
|
||||
|
||||
// 3. 线程 0 写入 Block 的总和
|
||||
if (tid == 0) {{
|
||||
partial_loss_out[blockIdx.x] = s_loss[0];
|
||||
partial_weight_out[blockIdx.x] = s_weight[0];
|
||||
}}
|
||||
}}
|
||||
|
||||
// C++ 封装函数
|
||||
torch::Tensor nll_loss_forward_cuda(
|
||||
torch::Tensor input,
|
||||
torch::Tensor target,
|
||||
c10::optional<torch::Tensor> weight_opt,
|
||||
int64_t ignore_index,
|
||||
std::string reduction
|
||||
) {{
|
||||
// 检查
|
||||
TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
|
||||
TORCH_CHECK(target.is_cuda(), "target must be a CUDA tensor");
|
||||
TORCH_CHECK(input.dim() >= 2, "input must be >= 2D");
|
||||
TORCH_CHECK(input.dim() == target.dim() + 1, "input/target dim mismatch");
|
||||
TORCH_CHECK(target.scalar_type() == torch::kLong, "target must be torch.long");
|
||||
|
||||
const int64_t N = input.size(0);
|
||||
const int64_t C = input.size(1);
|
||||
|
||||
int64_t N_spatial = 1;
|
||||
if (input.dim() > 2) {{
|
||||
for (int d = 2; d < input.dim(); ++d) {{
|
||||
TORCH_CHECK(input.size(d) == target.size(d-1), "spatial dim mismatch");
|
||||
N_spatial *= input.size(d);
|
||||
}}
|
||||
}}
|
||||
const int64_t N_total = N * N_spatial;
|
||||
|
||||
if (N_total == 0) {{
|
||||
return torch::tensor(0.0, input.options());
|
||||
}}
|
||||
|
||||
auto input_flat = input.contiguous().view({{N, C, N_spatial}});
|
||||
auto target_flat = target.contiguous().view({{N_total}});
|
||||
|
||||
const float* weight_ptr = nullptr;
|
||||
if (weight_opt.has_value()) {{
|
||||
auto& weight = weight_opt.value();
|
||||
TORCH_CHECK(weight.is_cuda() && weight.is_contiguous() &&
|
||||
weight.dim() == 1 && weight.size(0) == C, "weight size mismatch");
|
||||
weight_ptr = weight.data_ptr<float>();
|
||||
}}
|
||||
|
||||
// 4.核函数路由
|
||||
dim3 block_dim(BLOCK_SIZE);
|
||||
dim3 grid_dim((N_total + BLOCK_SIZE - 1) / BLOCK_SIZE);
|
||||
|
||||
if (reduction == "none") {{
|
||||
auto output_loss = torch::empty_like(target_flat, input.options());
|
||||
|
||||
nll_loss_kernel_no_reduce<<<grid_dim, block_dim>>>(
|
||||
input_flat.data_ptr<float>(),
|
||||
target_flat.data_ptr<long>(),
|
||||
weight_ptr,
|
||||
output_loss.data_ptr<float>(),
|
||||
N_spatial, C, N_total, ignore_index
|
||||
);
|
||||
return output_loss.view(target.sizes());
|
||||
|
||||
}} else {{ // "mean" or "sum"
|
||||
|
||||
// 创建小的部分和张量
|
||||
auto partial_loss_out = torch::empty({{grid_dim.x}}, input.options());
|
||||
auto partial_weight_out = torch::empty({{grid_dim.x}}, input.options());
|
||||
|
||||
nll_loss_kernel_reduce<<<grid_dim, block_dim>>>(
|
||||
input_flat.data_ptr<float>(),
|
||||
target_flat.data_ptr<long>(),
|
||||
weight_ptr,
|
||||
partial_loss_out.data_ptr<float>(),
|
||||
partial_weight_out.data_ptr<float>(),
|
||||
N_spatial, C, N_total, ignore_index
|
||||
);
|
||||
|
||||
// 5. 在 C++ 中对*小的*部分和张量进行归约
|
||||
torch::Tensor total_loss = partial_loss_out.sum();
|
||||
|
||||
if (reduction == "sum") {{
|
||||
return total_loss;
|
||||
}}
|
||||
|
||||
// reduction == "mean"
|
||||
double total_weight = partial_weight_out.sum().item<double>();
|
||||
if (total_weight == 0.0) {{
|
||||
return torch::tensor(0.0, input.options());
|
||||
}}
|
||||
return total_loss / total_weight;
|
||||
}}
|
||||
}}
|
||||
"""
|
||||
|
||||
nvcc_flags = ['-O3', '--use_fast_math']
|
||||
|
||||
# JIT (Just-In-Time) 编译
|
||||
self.loss_op = load_inline(
|
||||
name="nll_loss_op_v2_reduce", # 更改名称
|
||||
cpp_sources=cpp_header,
|
||||
cuda_sources=cuda_source,
|
||||
functions=["nll_loss_forward_cuda"],
|
||||
extra_cuda_cflags=nvcc_flags,
|
||||
verbose=False
|
||||
)
|
||||
|
||||
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
# C++ wrapper 现在处理 .contiguous()
|
||||
return self.loss_op.nll_loss_forward_cuda(
|
||||
input,
|
||||
target,
|
||||
self.weight,
|
||||
self.ignore_index,
|
||||
self.reduction_str
|
||||
)
|
||||
|
|
@ -0,0 +1,94 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
# -------------------------------------------------------------
|
||||
# 常量定义
|
||||
# -------------------------------------------------------------
|
||||
N, C, H, W = 8, 10, 16, 16 # (N, C, H, W)
|
||||
|
||||
# 损失函数参数
|
||||
WEIGHT = torch.rand(C, dtype=torch.float32) # (C,)
|
||||
IGNORE_INDEX = -100
|
||||
REDUCTION = 'mean'
|
||||
|
||||
|
||||
# -------------------------------------------------------------
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
nn.NLLLoss 的纯 PyTorch 基准实现
|
||||
(K-dim, 2D-example)
|
||||
"""
|
||||
|
||||
def __init__(self, weight=None, size_average=None, ignore_index=-100,
|
||||
reduce=None, reduction='mean'):
|
||||
super().__init__()
|
||||
|
||||
# 处理已弃用的 size_average 和 reduce
|
||||
if size_average is not None or reduce is not None:
|
||||
# (省略... 遵循 torch.nn.modules.loss)
|
||||
pass
|
||||
|
||||
self.reduction = reduction
|
||||
self.ignore_index = ignore_index
|
||||
|
||||
# 确保 weight 在正确的设备上
|
||||
if weight is not None:
|
||||
self.register_buffer('weight', weight)
|
||||
else:
|
||||
self.weight = None
|
||||
|
||||
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
input_flat = input.view(N, C, -1)
|
||||
|
||||
target_flat = target.view(N, -1)
|
||||
|
||||
|
||||
loss_unreduced = input_flat.gather(dim=1, index=target_flat.unsqueeze(1))
|
||||
loss_unreduced = -loss_unreduced.squeeze(1) # (N, H*W)
|
||||
|
||||
|
||||
if self.weight is not None:
|
||||
|
||||
weights_applied = self.weight[target_flat]
|
||||
loss_unreduced = loss_unreduced * weights_applied
|
||||
else:
|
||||
|
||||
weights_applied = torch.ones_like(target_flat, dtype=input.dtype)
|
||||
|
||||
|
||||
mask = (target_flat != self.ignore_index)
|
||||
loss_unreduced = loss_unreduced * mask
|
||||
weights_applied = weights_applied * mask
|
||||
|
||||
|
||||
if self.reduction == 'mean':
|
||||
|
||||
total_weight = weights_applied.sum()
|
||||
if total_weight == 0:
|
||||
return torch.tensor(0.0, device=input.device, dtype=input.dtype)
|
||||
return loss_unreduced.sum() / total_weight
|
||||
|
||||
elif self.reduction == 'sum':
|
||||
return loss_unreduced.sum()
|
||||
|
||||
else:
|
||||
return loss_unreduced.view_as(target)
|
||||
|
||||
|
||||
def get_inputs():
|
||||
|
||||
input_log_probs = F.log_softmax(torch.randn(N, C, H, W, dtype=torch.float32), dim=1)
|
||||
target = torch.empty(N, H, W, dtype=torch.long).random_(0, C)
|
||||
|
||||
|
||||
target.view(-1)[::10] = IGNORE_INDEX
|
||||
|
||||
return [input_log_probs, target]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
|
||||
return [WEIGHT, None, IGNORE_INDEX, None, REDUCTION]
|
||||
|
|
@ -0,0 +1,147 @@
|
|||
You write custom CUDA kernels to replace the PyTorch operators in the given EvoNorm architecture to get speedups.
|
||||
You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining normalization+affine_transform+nonlinear_gating), or algorithmic changes (such as optimized memory access patterns). You are only limited by your imagination.
|
||||
|
||||
Technologies Used :
|
||||
|
||||
PyTorch: Deep learning framework
|
||||
|
||||
CUDA: GPU acceleration for parallel computing
|
||||
|
||||
C++/CUDA C++: High-performance kernel programming
|
||||
|
||||
Inline C++/CUDA Extension (torch.utils.cpp_extension.load_inline): Just-In-Time (JIT) compilation of custom operators
|
||||
|
||||
Negative Log-Likelihood Loss (NLLLoss): Classification loss function for probability distributions
|
||||
|
||||
Dual-Kernel Strategy: Separate kernels for "none" reduction vs "mean"/"sum" reduction
|
||||
|
||||
Grid-Stride Loops: Efficiently processes data of arbitrary size using fixed thread blocks
|
||||
|
||||
Shared Memory Reduction: Uses __shared__ arrays for block-level parallel reduction
|
||||
|
||||
Tree Reduction Pattern: Binary tree reduction within thread blocks using __syncthreads()
|
||||
|
||||
Device Function: compute_nll_loss_item helper function shared between kernels
|
||||
|
||||
Conditional Weight Handling: Supports optional class weights with null pointer checking
|
||||
|
||||
Ignore Index Support: Filters out specified target indices from loss calculation
|
||||
|
||||
Multi-Dimensional Tensor Support: Handles 2D+ inputs with spatial dimensions
|
||||
|
||||
Tensor Flattening: Converts multi-dimensional tensors to flat views for kernel processing
|
||||
|
||||
Two-Stage Reduction: Block-level partial reduction followed by host-side final reduction
|
||||
|
||||
Boundary Checking: Validates target indices and handles out-of-range values
|
||||
|
||||
Memory Coalescing: Ensures contiguous tensor layout for optimal memory access
|
||||
|
||||
Fast Math Operations: Uses --use_fast_math compiler flag
|
||||
|
||||
Comprehensive Input Validation: Checks tensor dimensions, types, and device placement
|
||||
|
||||
Zero-Size Tensor Handling: Returns zero loss for empty inputs
|
||||
|
||||
Numerical Stability: Handles zero total weight case for mean reduction
|
||||
|
||||
Flexible Reduction Modes: Supports "none", "mean", and "sum" reduction strategies
|
||||
|
||||
Optional Tensor Handling: Uses c10::optional for optional weight parameter
|
||||
|
||||
|
||||
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
# -------------------------------------------------------------
|
||||
# 常量定义
|
||||
# -------------------------------------------------------------
|
||||
N, C, H, W = 8, 10, 16, 16 # (N, C, H, W)
|
||||
|
||||
# 损失函数参数
|
||||
WEIGHT = torch.rand(C, dtype=torch.float32) # (C,)
|
||||
IGNORE_INDEX = -100
|
||||
REDUCTION = 'mean'
|
||||
|
||||
|
||||
# -------------------------------------------------------------
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
nn.NLLLoss 的纯 PyTorch 基准实现
|
||||
(K-dim, 2D-example)
|
||||
"""
|
||||
|
||||
def __init__(self, weight=None, size_average=None, ignore_index=-100,
|
||||
reduce=None, reduction='mean'):
|
||||
super().__init__()
|
||||
|
||||
# 处理已弃用的 size_average 和 reduce
|
||||
if size_average is not None or reduce is not None:
|
||||
# (省略... 遵循 torch.nn.modules.loss)
|
||||
pass
|
||||
|
||||
self.reduction = reduction
|
||||
self.ignore_index = ignore_index
|
||||
|
||||
# 确保 weight 在正确的设备上
|
||||
if weight is not None:
|
||||
self.register_buffer('weight', weight)
|
||||
else:
|
||||
self.weight = None
|
||||
|
||||
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
input_flat = input.view(N, C, -1)
|
||||
|
||||
target_flat = target.view(N, -1)
|
||||
|
||||
|
||||
loss_unreduced = input_flat.gather(dim=1, index=target_flat.unsqueeze(1))
|
||||
loss_unreduced = -loss_unreduced.squeeze(1) # (N, H*W)
|
||||
|
||||
|
||||
if self.weight is not None:
|
||||
|
||||
weights_applied = self.weight[target_flat]
|
||||
loss_unreduced = loss_unreduced * weights_applied
|
||||
else:
|
||||
|
||||
weights_applied = torch.ones_like(target_flat, dtype=input.dtype)
|
||||
|
||||
|
||||
mask = (target_flat != self.ignore_index)
|
||||
loss_unreduced = loss_unreduced * mask
|
||||
weights_applied = weights_applied * mask
|
||||
|
||||
|
||||
if self.reduction == 'mean':
|
||||
|
||||
total_weight = weights_applied.sum()
|
||||
if total_weight == 0:
|
||||
return torch.tensor(0.0, device=input.device, dtype=input.dtype)
|
||||
return loss_unreduced.sum() / total_weight
|
||||
|
||||
elif self.reduction == 'sum':
|
||||
return loss_unreduced.sum()
|
||||
|
||||
else:
|
||||
return loss_unreduced.view_as(target)
|
||||
|
||||
|
||||
def get_inputs():
|
||||
|
||||
input_log_probs = F.log_softmax(torch.randn(N, C, H, W, dtype=torch.float32), dim=1)
|
||||
target = torch.empty(N, H, W, dtype=torch.long).random_(0, C)
|
||||
|
||||
|
||||
target.view(-1)[::10] = IGNORE_INDEX
|
||||
|
||||
return [input_log_probs, target]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
|
||||
return [WEIGHT, None, IGNORE_INDEX, None, REDUCTION]
|
||||
|
|
@ -0,0 +1,77 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from NLLLoss_torch import Model, get_inputs, get_init_inputs
|
||||
from NLLLoss_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