测试代码如下：

import torch

import time

import numpy as np

import sys

def test_lstsq_performance():

if torch.cuda.is_available():

device = torch.device("cuda")

device_name = torch.cuda.get_device_name(0)

print(f"使用设备: {device} - {device_name}")

else:

device = torch.device("cpu")

print(f"使用设备: {device}")

# 测试不同的问题规模

test_sizes = [

(100, 2, 2), # 小规模：100个2x2线性方程组

(1000, 2, 2), # 中规模：1000个2x2线性方程组

(10000, 2, 2), # 大规模：1万个2x2线性方程组

(50000, 2, 2), # 超大规模：5万个2x2线性方程组

(100000, 2, 2), # 极大规模：10万个2x2线性方程组（接近实际问题规模）

]

print("\n开始性能测试...")

print("="*80)

print(f"{'问题规模':<15} {'耗时(秒)':<12} {'状态':<10} {'备注'}")

print("="*80)

for n_equations, m, n in test_sizes:

try:

print(f"{n_equations}个{m}x{n}方程组", end="")

# 创建随机测试数据

A = torch.randn(n_equations, m, n, device=device, dtype=torch.float32)

b = torch.randn(n_equations, m, 1, device=device, dtype=torch.float32)

# 预热GPU

if device.type == 'cuda':

torch.cuda.synchronize()

# 开始计时

start_time = time.time()

# 设置超时检测

max_wait_time = 30.0 # 最多等待30秒

# 执行求解

try:

if device.type == 'cuda':

torch.cuda.synchronize()

x, residuals = torch.linalg.lstsq(A, b)[:2]

if device.type == 'cuda':

torch.cuda.synchronize()

end_time = time.time()

elapsed = end_time - start_time

if elapsed > max_wait_time:

print(f" {elapsed:.3f} 超时 耗时过长")

else:

print(f" {elapsed:.3f} 正常 ")

except Exception as e:

end_time = time.time()

elapsed = end_time - start_time

print(f" {elapsed:.3f} 错误 {str(e)[:30]}...")

except KeyboardInterrupt:

print(" 中断 手动停止 用户中断")

break

except Exception as e:

print(f" 失败 异常 {str(e)[:30]}...")

print("="*80)

def main():

print("DCU平台torch.linalg.lstsq性能测试")

print(f"PyTorch版本: {torch.__version__}")

print(f"CUDA可用: {torch.cuda.is_available()}")

if torch.cuda.is_available():

print(f"CUDA版本: {torch.version.cuda}")

print(f"设备数量: {torch.cuda.device_count()}")

# 基本性能测试

test_lstsq_performance()

if __name__ == "__main__":