有什么方法可以将我的CPU并行代码转换为cuda吗？

问题描述 投票：0回答：1

我编写了一段在 CPU 内核上运行良好的代码。但对我来说还不够快。我想在 CUDA 核心上运行它，并且我已经尝试为其 montecarlo 部分编写内核等。但无论我做什么，我都无法正确运行它。

所以我有两个问题：

  1. 在 CUDA 上运行它是个好主意吗？会有什么不同吗？
  2. 你能告诉我一个可以将其更改为 CUDA 版本的方法吗？
import os
import numpy as np
import matplotlib.pyplot as plt
from tqdm.notebook import trange
from numba import jit
from numba import misc as msc
import random
import itertools
from joblib import Parallel, delayed  # Import joblib for parallel processing
from scipy.special import binom, comb
from math import factorial
import pickle

# # Ensure 'output' folder exists in the current working directory
# script_dir = os.path.dirname(os.path.realpath(__file__))
# output_dir = os.path.join(script_dir, 'output')
# if not os.path.exists(output_dir):
#     os.makedirs(output_dir)

# Running simulation
n = 2  # Number of spins in the many-body system
omegas = np.arange(0.1, 12, 0.2)
ents = []

# Parameters
dt = 0.002
kappa = 1.0
N = 2  # Number of trajectories
nt = 100000  # Number of timesteps which makes total time T = nt * dt = 500
J = n // 2

def precompute_binomial_coefficients(N):
    binomial_coeffs = np.zeros((N + 1, N + 1))
    for n in range(N + 1):
        for k in range(n + 1):
            binomial_coeffs[n, k] = comb(n, k)
    return binomial_coeffs
binomial_coeffs = precompute_binomial_coefficients(N)


# Incase it's needed
@jit(nopython=True)
def dicke_basis(n, m):
    """Generate Dicke basis states for a specified number of spins and value of m."""
    basis_states = []
    num_up_spins = int(m + (J))
    for positions in itertools.combinations(range(n), num_up_spins):
        state = np.zeros(n, dtype=np.int8)
        for pos in positions:
            state[pos] = 1
        basis_states.append(state)
    return basis_states

@jit(nopython=True)
def s_z_operator(n):
    return np.diag(np.arange((J), (-J) - 1, -1))

@jit(nopython=True)
def s_minus_operator(n):
    s_minus = np.zeros((int(n + 1), int(n + 1)), dtype=np.complex128)
    for m in range(1, int(n + 1)):
        S_z = (n / 2) - (m - 1)
        s_minus[m, m - 1] = np.sqrt((J) * ((J) + 1) - S_z * (S_z - 1))
    return s_minus

@jit(nopython=True)
def s_plus_operator(n):
    return s_minus_operator(n).T.conj()

@jit(nopython=True)
def s_x_operator(n):
    return 0.5 * (s_plus_operator(n) + s_minus_operator(n))

@jit(nopython=True)
def s_y_operator(n):
    return -1j * (s_plus_operator(n) - s_minus_operator(n)) * 0.5




@jit(nopython=True)
def normstate(state):
    return np.real(np.dot(state.conj().T, state))
s_z = s_z_operator(n)
s_m = s_minus_operator(n)
s_x = s_x_operator(n)
s_y = s_y_operator(n)
s_p = s_plus_operator(n)

for i in range(len(omegas)):

 

  #Full density matrix
  @jit(nopython=True)
  def full_density_matrix(P_ne):
      """Construct the full density matrix from probability amplitudes P_ne."""
      rho_N = np.outer(P_ne, P_ne.conj())
      return rho_N

  

  # Calculating reduced density matrix
  @jit(nopython=True)
  def reduced_density_matrix(rho_N, k, binomial_coeffs):
      """Calculate the reduced density matrix of k spins."""
      N = rho_N.shape[0] - 1
      L = N - k  # Number of spins in the remaining partition
      rho_k = np.zeros((k + 1, k + 1), dtype=np.complex128)

      for ne in range(N + 1):
          for ne_prime in range(N + 1):
              for le in range(min(ne, k) + 1):
                  coef = (
                      binomial_coeffs[L, le]
                      * np.sqrt(
                          binomial_coeffs[N - L, ne - le]
                          * binomial_coeffs[N - L, ne_prime - le]
                          / (binomial_coeffs[N, ne] * binomial_coeffs[N, ne_prime])
                      )
                  )
                  if 0 <= (ne - le) <= k and 0 <= (ne_prime - le) <= k:
                      rho_k[ne - le, ne_prime - le] += coef * rho_N[ne, ne_prime]

      return rho_k

  # Neumann entanglement entropy
  @jit(nopython=True)
  def entanglement_entropy(rho_k):
      """Calculate the entanglement entropy from the reduced density matrix rho_k."""
      eigenvalues = np.linalg.eigvalsh(rho_k)
      entropy = -np.sum(eigenvalues * np.log(eigenvalues + 1e-12))  # Avoid log(0)
      return entropy

  #Monte carlo process
  @jit(nopython=True)
  def montecarlo(state_init, nt, dt, binomial_coeffs):
      """Perform a Monte Carlo simulation for a single trajectory in Dicke basis."""
      state = state_init.copy()  # Use a copy to avoid modifying the original
      trajectory = np.zeros(nt + 1, dtype=np.float64)  # Preallocate trajectory array
      trajectory[0] = normstate(state)  # Store initial state norm
      num_jumps = 0
      entropies = np.zeros(nt + 1)
      eye = np.eye(n + 1, dtype=np.complex128)
      spm = np.dot(s_p, s_m)
      for step in range(nt):
          new_state = np.dot((eye - 1j * H_eff * dt), state)  # Evolve state
          norm_new_state = normstate(new_state)
          r = np.random.rand()
          p = dt * normstate(np.dot(op_coll, state))

          if r < p:
              # Calculate J_+ J_- expectation value for normalization
              norm_factor = np.dot(state.conj().T, np.dot(spm, state))

              # Apply the jump and normalize by sqrt(norm_factor)
              jump_state = np.dot(op_coll, state)
              state = jump_state / np.sqrt(norm_factor)
              num_jumps += 1
          else:
              state = new_state / np.sqrt(norm_new_state)
          trajectory[step + 1] = np.abs(state[0]) ** 2
          rho_N = full_density_matrix(state)
          k = J
          rho_k = reduced_density_matrix(rho_N, k, binomial_coeffs)  # Pass binomial_coeffs here
          entropies[step + 1] = entanglement_entropy(rho_k)
      return trajectory, num_jumps, entropies


  def run_trajectories(N, nt, dt, binomial_coeffs):
      """Run N trajectories in parallel and average the results."""
      # Use joblib to run Monte Carlo simulations in parallel
      results = Parallel(n_jobs=-1)(
          delayed(montecarlo)(state_init, nt, dt, binomial_coeffs) for _ in trange(N)
      )

      # Unzip the results
      trajectories, jumps, entropies = zip(*results)

      trajectories = np.array(trajectories)  # Convert list of trajectories to numpy array
      jumps = np.array(jumps)  # Convert list of jumps to numpy array
      entropies = np.array(entropies)  # Convert list of entropies to numpy array
      # Average the trajectories manually
      # averaged = np.zeros(nt + 1, dtype=np.float64)
      # for j in range(nt + 1):
      #     averaged[j] = np.sum(trajectories[:, j]) / N

      return entropies




  omega_0 = omegas[i]
  
  H = omega_0 * s_x
  H_eff = H - 1j * (kappa / (2 * (n / 2))) * np.dot(s_p, s_m)
  op_coll = (kappa / J) * s_m.astype(np.complex128)  # Ensure op_coll is complex128
  state_init = np.zeros(n + 1, dtype=np.complex128)
  state_init[0] = 1.0  # Initial state in Dicke basis
  entropies = run_trajectories(N, nt, dt, binomial_coeffs)

  avgent = np.mean(entropies, axis=0)
  ents.append(avgent)
  plt.plot(ents[i])
  plt.show()
  plt.cla()
  # Save omegas and entropies in the 'output' folder

#entropy_file = os.path.join(output_dir, f"ents-{n}.npy")
np.save(f"ents-{n}.npy", ents)  # Use np.save to store numpy arrays
print(f"Saved results for n={n} in the output folder.")
python parallel-processing cuda numba
1个回答
0
投票

但对我来说还不够快

在尝试在 GPU 上运行此程序之前，您应该真正分析 CPU 版本。事实证明，大部分时间都花在编译代码上...事实上，JIT 函数处于循环中，因此每次都会重新定义和重新编译它们。编译是一个昂贵的过程，所以你应该避免这样做。为了性能，这些函数应该定义并编译一次（在某些情况下，您可以专门化这些函数，但是当编译时间大于运行时间时，这没有任何意义）。

事实证明，某些函数访问全局变量。请注意，这些变量在 Numba 中被视为编译时常量（理论上它们不能改变）。 Numba 可以用代码中的常量替换该值，甚至可以根据它们预先计算值（有时可能会非常慢并生成大量汇编代码）。您应该避免使用大型全局数组并且永远不要写入它们。事实上，通常最好将这样的数组传递给函数，因为全局变量在软件工程方面往往很糟糕，特别是当它们可以变异时（就像“幽灵般的远距离动作”;））。

将 Joblib 与 Numba 一起使用并不是一个好主意，因为它可能会创建需要为每个进程编译 Numba 函数的进程（这是非常低效的）。 Numba 支持多线程，这要归功于基于 

parallel=True
和 
prange
的循环，它们通常比使用 Joblib 更有效，因此调用即时 Numba 函数（并且在优化方面也更灵活）。

我强烈建议您先改进CPU版本，并在使用GPU之前检查所有核心至少已满负荷使用。

在 CUDA 上运行它是个好主意吗？

为了让代码在 GPU 上运行得更快，您需要大量的并行性；明显多于 CPU。问题是 GPU 不仅仅是具有更多内核的 CPU，而且是一种非常不同的硬件，有很多限制。该代码还需要SIMT 友好。基本上，您需要拆分工作，以便数千个 CUDA 线程可以同时运行（理想情况下甚至可以是数十万个线程）。此外，为了提高性能，应该编写代码，以便 CUDA 线程以“合并方式”访问内存。代码还应该避免“扭曲发散”。还有其他严格的限制：线程不需要使用太多寄存器（否则会降低占用率，通常也会降低性能），GPU 内核只能分配相当有限的内存量（例如 16 MiB，始终预先分配，并且可能在所有内核执行之前进行修改）。因此，将相对较大的代码（如这个代码）有效地移植到 GPU 通常相当困难。 这里 n 和

N

太小了；假设基于 

nt
的循环可以安全地并行执行（即没有循环携带的依赖项），只有 
nt
足够大才能提供足够的并行性。我想知道 
state
是否是一个依赖项...如果循环需要串行执行，那么除非我错过任何其他并行性来源，否则就没有机会在 GPU 上有效执行循环。
即使假设有足够的并行性，我认为将此代码有效地移植到 GPU 对于初学者来说也是一项相对庞大的工作！


你能告诉我一个可以将其更改为 CUDA 版本的方法吗？




Numba 中的标准方法是编写 Numba CUDA 内核。 
文档
对此进行了解释，因此请阅读它。


或者，您可以编写代码，以便在 Numpy 中对所有操作进行矢量化，然后使用 CuPy 轻松地将工作转移到 GPU 上。话虽这么说，这通常不如编写 GPU 内核那么高效（人们首先使用 GPU 是为了让事情变得更快）。我不确定这（即仅使用 Numpy 向量化函数或简单的映射函数）可以在这里轻松完成，甚至有效......
    
	



    

    


  

  

    
最新问题


  





  

    © www.soinside.com 2019 - 2024. All rights reserved.