将复数数组传递到 PyCUDA 内核答案

【问题标题】：Passing complex number arrays into PyCUDA Kernel将复数数组传递到 PyCUDA 内核
【发布时间】：2019-12-19 04:53:10
【问题描述】：

我正在尝试将一个二维复数数组传递给 PyCUDA 内核，但得到了意想不到的结果。

这是我的测试代码：

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

mod = SourceModule("""
  #include <pycuda-complex.hpp>
  #include <stdio.h>

  typedef pycuda::complex<float> cmplx;
  __global__ void myFunc(cmplx *A)
  {
    // A      : input, array shape (), of type complex64

    int ROWS = 3;
    int COLS = 2;

    printf("\\nKernel >>");
    for(int row = 0; row < ROWS; row++)
    {
        printf("\\n");
        for(int col = 0; col < COLS; col++)
        {
            printf("[row %d, col %d]: %f + %fi",row, col, A[row,col].real(), A[row,col].imag());
            printf("\\t");
        }
    }

    printf("\\n\\n");  
  }
  """)


A = np.zeros((3,2),dtype=complex)
A[0,0] = 1.23 + 3.5j
A[1,0] = 3.4 + 1.0j

A_gpu = gpuarray.to_gpu(A.astype(np.complex64))

print("Host >>")
print(A_gpu)

func = mod.get_function("myFunc")
func(A_gpu,
     block=(1,1,1), grid=(1, 1, 1)
    )

结果如下：

Host >>
[[1.23+3.5j 0.  +0.j ]
 [3.4 +1.j  0.  +0.j ]
 [0.  +0.j  0.  +0.j ]]

Kernel >>
[row 0, col 0]: 1.230000 + 3.500000i    [row 0, col 1]: 0.000000 + 0.000000i    
[row 1, col 0]: 1.230000 + 3.500000i    [row 1, col 1]: 0.000000 + 0.000000i    
[row 2, col 0]: 1.230000 + 3.500000i    [row 2, col 1]: 0.000000 + 0.000000i

谁能解释为什么内核中的数组看起来不像我传递给它的那个？

【问题讨论】：

标签： python cuda pycuda

【解决方案1】：

您的内核代码中的索引已损坏（请参阅here 了解原因）。

虽然A[row,col] 在 C++ 中在技术上是有效的语法，但它并不像在 Python 中那样暗示多维数组切片。事实上，A[row,col] 的计算结果为 A[row]，因此 print 语句的输出与您的预期不符的原因应该很明显了。

Numpy 数组是在内存中连续存储的，您必须使用自己的索引方案来访问数组。默认情况下，numpy 将row major ordering 用于多维数组。这个：

mod = SourceModule("""
  #include <pycuda-complex.hpp>
  #include <stdio.h>

  typedef pycuda::complex<float> cmplx;
  __global__ void myFunc(cmplx *A)
  {
    // A      : input, array shape (), of type complex64

    int ROWS = 3;
    int COLS = 2;

    printf("\\nKernel >>");
    for(int row = 0; row < ROWS; row++)
    {
        printf("\\n");
        for(int col = 0; col < COLS; col++)
        {
            printf("[row %d, col %d]: %f + %fi",row, col, A[row*COLS+col].real(), A[row*COLS+col].imag());
            printf("\\t");
        }
    }

    printf("\\n\\n");  
  }
""")

将按预期工作：

%run complex_print.py
Host >>
[[ 1.23000002+3.5j  0.00000000+0.j ]
 [ 3.40000010+1.j   0.00000000+0.j ]
 [ 0.00000000+0.j   0.00000000+0.j ]]

Kernel >>
[row 0, col 0]: 1.230000 + 3.500000i    [row 0, col 1]: 0.000000 + 0.000000i    
[row 1, col 0]: 3.400000 + 1.000000i    [row 1, col 1]: 0.000000 + 0.000000i    
[row 2, col 0]: 0.000000 + 0.000000i    [row 2, col 1]: 0.000000 + 0.000000i

【讨论】：