Python与CUDA/ROCm的底层内存交互：实现GPU/加速器之间的高效数据传输

Python与CUDA/ROCm的底层内存交互：实现GPU/加速器之间的高效数据传输

大家好！今天我们来深入探讨一个关键而又复杂的课题：Python与CUDA/ROCm的底层内存交互。在深度学习和高性能计算领域，Python作为易用性极佳的脚本语言，经常被用于算法原型设计、模型训练流程编排以及后处理分析。而CUDA（NVIDIA）和ROCm（AMD）则分别是主流GPU加速平台的编程模型，它们能够充分发挥GPU的并行计算能力。如何高效地在Python和CUDA/ROCm之间传输数据，直接影响着整个计算流程的性能。

1. 问题：Python与GPU内存管理的固有差异

Python的内存管理由其自身的垃圾回收机制控制，而CUDA/ROCm则拥有独立的GPU内存空间和显式内存管理API。直接从Python对象访问GPU内存是不可能的，因为它们位于不同的地址空间。因此，我们需要借助桥梁来打通这两个世界，实现高效的数据传输。

2. 桥梁：NumPy和数组接口

NumPy是Python科学计算的核心库，提供了强大的多维数组对象。NumPy数组在内存中以连续的块存储数据，这使得它可以作为高效的中间表示，用于与CUDA/ROCm交互。NumPy数组实现了Python的数组接口（array interface），这是一个协议，允许不同的Python库共享数组数据，而无需复制。

3. CUDA/ROCm Python绑定库：PyCUDA/PyTorch/CuPy/Numba/HIP

为了实现Python与CUDA/ROCm的底层交互，我们需要使用专门的Python绑定库。这些库提供了访问CUDA/ROCm API的Python接口，并封装了数据传输和设备管理的复杂性。一些流行的库包括：

• PyCUDA: 提供了对CUDA C API的底层访问，灵活性高，但使用相对复杂。
• PyTorch: 一个流行的深度学习框架，内置了强大的CUDA支持，能够自动管理GPU内存和数据传输。
• CuPy: 一个与NumPy兼容的GPU数组库，提供了与NumPy类似的API，方便将NumPy代码移植到GPU上运行。
• Numba: 一个即时（JIT）编译器，可以将Python代码编译成高性能的CUDA/ROCm代码，无需手动编写CUDA/ROCm内核。
• HIP: 由AMD开发的异构计算接口，允许在AMD和NVIDIA GPU上运行相同的代码，降低了代码维护成本。

4. 数据传输方式：显式与隐式

数据传输可以分为显式和隐式两种方式：

• 显式数据传输: 开发者需要明确调用API函数来执行数据在CPU内存和GPU内存之间的拷贝。这种方式需要更多的代码，但可以更精确地控制数据传输的时机和方式。
• 隐式数据传输: 库会自动处理数据传输，例如，在PyTorch中，将NumPy数组转换为Tensor时，数据会自动传输到GPU。这种方式简化了代码，但可能会牺牲一些性能控制。

5. 数据传输API：内存拷贝函数

不同的库提供了不同的API来进行内存拷贝。常见的API包括：

• PyCUDA: 使用pycuda.driver.memcpy_htod() (host to device) 和 pycuda.driver.memcpy_dtoh() (device to host) 函数。
• PyTorch: 使用 .to() 方法将Tensor移动到不同的设备（CPU或GPU）。
• CuPy: 使用 cupy.asarray() 和 cupy.asnumpy() 函数在NumPy数组和CuPy数组之间转换。
• Numba: 通过cuda.to_device()和cuda.to_host()在Numba编译的函数中进行数据传输。

6. 代码示例：PyCUDA

以下是一个使用PyCUDA进行显式数据传输的例子：

import pycuda.driver as cuda
import pycuda.autoinit
from pycuda.compiler import SourceModule
import numpy as np

# 初始化CUDA
cuda.init()
dev = cuda.Device(0) # 选择GPU设备
ctx = dev.make_context()

# 创建一个NumPy数组
a_np = np.random.randn(400).astype(np.float32)

# 在GPU上分配内存
a_gpu = cuda.mem_alloc(a_np.nbytes)

# 将数据从CPU拷贝到GPU
cuda.memcpy_htod(a_gpu, a_np)

# 定义一个简单的CUDA内核
mod = SourceModule("""
  __global__ void multiply_by_two(float *a)
  {
    int idx = threadIdx.x + blockIdx.x*blockDim.x;
    a[idx] = a[idx] * 2;
  }
""")

# 获取内核函数
multiply_by_two = mod.get_function("multiply_by_two")

# 执行内核函数
block_size = 256
grid_size = (a_np.size + block_size - 1) // block_size
multiply_by_two(a_gpu, block=(block_size,1,1), grid=(grid_size,1))

# 在GPU上分配结果数组的内存
a_doubled_np = np.empty_like(a_np)

# 将结果从GPU拷贝到CPU
cuda.memcpy_dtoh(a_doubled_np, a_gpu)

# 打印结果
print("Original array:", a_np[:5])
print("Doubled array:", a_doubled_np[:5])

# 清理
a_gpu.free()
ctx.pop()

7. 代码示例：PyTorch

以下是一个使用PyTorch进行隐式数据传输的例子：

import torch
import numpy as np

# 创建一个NumPy数组
a_np = np.random.randn(400).astype(np.float32)

# 将NumPy数组转换为PyTorch Tensor
a_torch = torch.from_numpy(a_np)

# 将Tensor移动到GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
a_gpu = a_torch.to(device)

# 在GPU上执行操作
a_doubled_gpu = a_gpu * 2

# 将结果移动回CPU
a_doubled_cpu = a_doubled_gpu.cpu().numpy()

# 打印结果
print("Original array:", a_np[:5])
print("Doubled array:", a_doubled_cpu[:5])

8. 代码示例：CuPy

以下是一个使用CuPy进行数据传输的例子：

import cupy as cp
import numpy as np

# 创建一个NumPy数组
a_np = np.random.randn(400).astype(np.float32)

# 将NumPy数组转换为CuPy数组
a_gpu = cp.asarray(a_np)

# 在GPU上执行操作
a_doubled_gpu = a_gpu * 2

# 将结果移动回CPU
a_doubled_cpu = cp.asnumpy(a_doubled_gpu)

# 打印结果
print("Original array:", a_np[:5])
print("Doubled array:", a_doubled_cpu[:5])

9. 代码示例：Numba

以下是一个使用Numba进行隐式数据传输和GPU加速的例子：

from numba import cuda
import numpy as np

@cuda.jit
def add_gpu(x, y, out):
  idx = cuda.grid(1)
  out[idx] = x[idx] + y[idx]

n = 1024
x = np.arange(n, dtype=np.float32)
y = np.ones_like(x)

out = np.empty_like(x)

threads_per_block = 32
blocks_per_grid = (x.size + (threads_per_block - 1)) // threads_per_block

add_gpu[blocks_per_grid, threads_per_block](x, y, out)

print(out[:5])

10. 代码示例：HIP

import numpy as np
from hip import hipDeviceSynchronize, hipMemcpy, hipMalloc, hipFree
from hip.hiprtc import hiprtcCompileProgram, hiprtcCreateProgram, hiprtcGetCode, hiprtcDestroyProgram

# Kernel code (HIP code)
kernel_code = """
extern "C" __global__
void vecAdd(float *A, float *B, float *C, int N) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i < N) {
    C[i] = A[i] + B[i];
  }
}
"""

# Array size
N = 1024

# Allocate memory on the host
A_h = np.random.rand(N).astype(np.float32)
B_h = np.random.rand(N).astype(np.float32)
C_h = np.zeros_like(A_h)

# Allocate memory on the device
hipSuccess = 0
A_d = hipMalloc(A_h.nbytes)
B_d = hipMalloc(B_h.nbytes)
C_d = hipMalloc(C_h.nbytes)

# Copy data from host to device
hipMemcpy(A_d, A_h.ctypes.data, A_h.nbytes, 0)  # 0 corresponds to hipMemcpyHostToDevice
hipMemcpy(B_d, B_h.ctypes.data, B_h.nbytes, 0)

# Compile the kernel using hiprtc
prog = hiprtcCreateProgram(kernel_code, "vecAdd.cu", [], [])
result = hiprtcCompileProgram(prog, [])
code = hiprtcGetCode(prog)

# Load the module
module = hip.hipModuleLoadData(code)

# Get the function
func = hip.hipModuleGetFunction(module, "vecAdd")

# Configure the kernel
threads_per_block = 256
blocks_per_grid = (N + threads_per_block - 1) // threads_per_block

# Set up the arguments
args = [A_d, B_d, C_d, np.int32(N)]

# Launch the kernel
hip.hipLaunchKernel(func, blocks_per_grid, 1, 1, threads_per_block, 1, 1, 0, 0, args)
hipDeviceSynchronize()

# Copy data from device to host
hipMemcpy(C_h.ctypes.data, C_d, C_h.nbytes, 1) # 1 corresponds to hipMemcpyDeviceToHost

# Print the first 5 elements of the result
print("A + B = C:", C_h[:5])

# Free the memory on the device
hipFree(A_d)
hipFree(B_d)
hipFree(C_d)

hiprtcDestroyProgram(prog)

11. 性能优化：减少数据传输量

数据传输是GPU加速的瓶颈之一。为了提高性能，我们需要尽可能减少数据传输量。一些常用的技巧包括：

• 在GPU上进行尽可能多的计算: 避免频繁地在CPU和GPU之间传输数据。
• 使用pinned memory (page-locked memory): pinned memory可以提高数据传输速度，因为它避免了CPU缓存的干扰。 PyCUDA 和 Numba 都提供了对 pinned memory 的支持.
• 使用异步数据传输: 异步数据传输允许CPU和GPU同时执行计算和数据传输，从而提高并行性。CUDA streams可以实现异步数据传输。
• 数据类型选择: 选择合适的数据类型，避免使用过大的数据类型，从而减少数据传输量。 例如，如果精度要求不高，可以使用 float32 代替 float64。
• 数据压缩: 对于某些类型的数据，可以使用压缩算法来减少数据传输量。

12. pinned memory 的使用示例 (PyCUDA)

import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np

# 创建 pinned memory
size = 1024 * 1024  # 1MB
a_np = np.zeros((size,), dtype=np.float32)
a_pinned = a_np.ctypes.data_as(cuda.POINTER(cuda.c_float))

# 在 GPU 上分配内存
a_gpu = cuda.mem_alloc(a_np.nbytes)

# 使用 pinned memory 进行数据传输
cuda.memcpy_htod(a_gpu, a_pinned) # 快速传输
# ... GPU kernel 操作 ...
cuda.memcpy_dtoh(a_np, a_gpu) #快速传回

a_gpu.free()

13. CUDA Streams的使用示例 (PyCUDA)

import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
from pycuda.compiler import SourceModule

# 初始化 CUDA
cuda.init()
dev = cuda.Device(0)
ctx = dev.make_context()

# 创建 streams
stream1 = cuda.Stream()
stream2 = cuda.Stream()

# 创建 numpy 数组
a_np = np.random.randn(400).astype(np.float32)
b_np = np.random.randn(400).astype(np.float32)

# 在 GPU 上分配内存
a_gpu = cuda.mem_alloc(a_np.nbytes)
b_gpu = cuda.mem_alloc(b_np.nbytes)
c_gpu = cuda.mem_alloc(a_np.nbytes) # 用于存储结果

# 定义 CUDA kernel (向量加法)
mod = SourceModule("""
  __global__ void vector_add(float *a, float *b, float *c, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
      c[idx] = a[idx] + b[idx];
    }
  }
""")
vector_add = mod.get_function("vector_add")

# 异步拷贝数据到 GPU (使用 stream1 和 stream2)
cuda.memcpy_htod_async(a_gpu, a_np, stream1)
cuda.memcpy_htod_async(b_gpu, b_np, stream2)

# 配置 kernel
block_size = 256
grid_size = (a_np.size + block_size - 1) // block_size

# 异步执行 kernel (等待 stream1 和 stream2 完成)
vector_add(a_gpu, b_gpu, c_gpu, np.int32(a_np.size), block=(block_size, 1, 1), grid=(grid_size, 1), stream=stream1)
vector_add(a_gpu, b_gpu, c_gpu, np.int32(a_np.size), block=(block_size, 1, 1), grid=(grid_size, 1), stream=stream2)

# 创建结果数组
c_np = np.empty_like(a_np)

# 异步拷贝结果到 CPU (等待 stream1 和 stream2 完成)
cuda.memcpy_dtoh_async(c_np, c_gpu, stream1)
cuda.memcpy_dtoh_async(c_np, c_gpu, stream2)

# 同步 streams 以确保所有操作完成
stream1.synchronize()
stream2.synchronize()

# 打印结果 (前 5 个元素)
print("Result (first 5 elements):", c_np[:5])

# 清理
a_gpu.free()
b_gpu.free()
c_gpu.free()
ctx.pop()

14. 结论：选择合适的工具和策略

Python与CUDA/ROCm的底层内存交互是一个复杂的主题，涉及到多个库和API。理解不同库的特点和适用场景，选择合适的数据传输方式，并采取性能优化措施，是实现GPU/加速器之间高效数据传输的关键。在实际应用中，需要根据具体的应用场景和性能需求，选择合适的工具和策略。

15. 总结：掌握底层交互，提升计算性能

Python作为高级语言，需要借助绑定库才能与CUDA/ROCm进行底层内存交互。 通过显式或隐式数据传输，结合性能优化技巧，可以显著提升GPU加速的计算性能。

更多IT精英技术系列讲座，到智猿学院