如何利用GPU加速求解线性方程组？NumPy/PyTorch性能测试分析

Let’s break down why your GPU-based PyTorch solve is underperforming, and fix it to get the expected speedups.

First, Why Your Initial Test Failed #

Your initial benchmark has several critical issues that mask GPU performance:

• Tiny matrix size (default 5x5): GPUs have significant overhead for kernel launches and data transfer between CPU/GPU. For small matrices, this overhead completely outweighs any computational speedup. GPUs shine with large-scale problems or batched workloads.
• Inefficient GPU test setup: Your solve_torch_gpu function sets the default tensor type but reuses solve_torch_cpu, which generates new tensors on every loop iteration—this adds unnecessary data transfer overhead.
• Misconfigured Numba tests: solve_numpy_njit_a doesn’t actually execute the solving logic (it just JIT-compiles the function without running it), which is why its time is artificially low. numpy.linalg.solve already uses optimized BLAS libraries (like OpenBLAS or MKL), so Numba won’t provide meaningful speedups here for single solves.

Fixing the GPU Benchmark #

To properly leverage GPU acceleration for linear system solving, follow these steps:

1. Use Larger/Batched Matrices #

GPUs excel at parallelizing work across multiple problems. Instead of solving 1000 small 5x5 systems, solve a single large system or batch hundreds/thousands of systems at once.

2. Minimize Data Transfer Overhead #

Move your data to the GPU once, not on every iteration. Avoid generating random tensors inside the loop on the GPU.

3. Use Modern PyTorch APIs #

Prefer torch.linalg.solve (available in PyTorch 1.9+) over the deprecated torch.solve—it’s optimized for both CPU and GPU.

Here’s an improved benchmark code that addresses these issues:

import timeit
import torch
import numpy as np
from numba import njit

# Configurable parameters
DIM = 500  # Larger matrix size
NUM_SOLVES = 1000  # Number of solves to perform

def solve_numpy_cpu():
    a = np.random.rand(DIM, DIM)
    b = np.random.rand(DIM)
    for _ in range(NUM_SOLVES):
        np.linalg.solve(a, b)

def solve_torch_cpu():
    a = torch.rand(DIM, DIM)
    b = torch.rand(DIM)
    for _ in range(NUM_SOLVES):
        torch.linalg.solve(a, b)

def solve_torch_gpu():
    # Move data to GPU once
    a = torch.rand(DIM, DIM, device='cuda')
    b = torch.rand(DIM, device='cuda')
    # Warm up the GPU kernel (first run has overhead)
    torch.linalg.solve(a, b)
    # Benchmark the actual solves
    start = timeit.default_timer()
    for _ in range(NUM_SOLVES):
        torch.linalg.solve(a, b)
    torch.cuda.synchronize()  # Wait for GPU operations to finish before timing
    end = timeit.default_timer()
    return end - start

# Numba is unlikely to help here, but for completeness
@njit
def solve_numpy_njit():
    a = np.random.rand(DIM, DIM)
    b = np.random.rand(DIM)
    for _ in range(NUM_SOLVES):
        np.linalg.solve(a, b)

def main():
    print(f"Benchmarking with {DIM}x{DIM} matrices, {NUM_SOLVES} solves:")
    # NumPy CPU
    time = timeit.timeit(solve_numpy_cpu, number=1)
    print(f"{solve_numpy_cpu.__name__:<20s}: {time:.4f}s")
    # PyTorch CPU
    time = timeit.timeit(solve_torch_cpu, number=1)
    print(f"{solve_torch_cpu.__name__:<20s}: {time:.4f}s")
    # PyTorch GPU
    time = solve_torch_gpu()
    print(f"{solve_torch_gpu.__name__:<20s}: {time:.4f}s")
    # Numba
    time = timeit.timeit(solve_numpy_njit, number=1)
    print(f"{solve_numpy_njit.__name__:<20s}: {time:.4f}s")

if __name__ == "__main__":
    main()

Key Optimizations in the Improved Code #

• Larger matrix size: We use 500x500 matrices (adjust this to even larger sizes like 1000x1000 for more dramatic GPU gains).
• GPU warm-up: The first GPU solve incurs kernel launch overhead, so we run it once before benchmarking.
• CUDA synchronization: torch.cuda.synchronize() ensures we wait for all GPU operations to complete before stopping the timer—this avoids underestimating GPU time.
• Single data transfer: Tensors are created directly on the GPU, so no repeated CPU<->GPU data transfer during the loop.

Additional Tips for Maximum GPU Performance #

• Batch multiple linear systems: If you have many independent linear systems, solve them in a single batch using a 3D tensor for a (shape (num_batches, dim, dim)) and b (shape (num_batches, dim)). torch.linalg.solve natively supports batched inputs, which fully utilizes GPU parallelism.
  Example:

  batch_size = 1000
  a_batch = torch.rand(batch_size, DIM, DIM, device='cuda')
  b_batch = torch.rand(batch_size, DIM, device='cuda')
  solutions = torch.linalg.solve(a_batch, b_batch)  # Solves all 1000 systems at once
  

• Use mixed precision: If your problem allows it, use half-precision (torch.float16) or bfloat16 (torch.bfloat16) tensors to reduce memory usage and increase throughput. PyTorch’s torch.linalg.solve supports mixed precision on modern GPUs.
• Ensure GPU utilization: Use tools like nvidia-smi or PyTorch’s torch.cuda.utilization() to check if your GPU is being fully utilized. If utilization is low, your problem size is still too small—scale up the matrix size or batch count.

Why CPU PyTorch is Faster Than NumPy #

PyTorch’s CPU linear algebra operations use optimized libraries (like MKL or OpenBLAS) with multithreading enabled by default, which can outperform NumPy if NumPy isn’t configured to use the same optimized libraries. This is why you saw a 4x speedup with CPU PyTorch.

内容的提问来源于stack exchange，提问作者wehnsdaefflae