将复数数组传递给 PyCUDA 内核

Question

我试图将一个二维复数数组传递给 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

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

Answer 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

将复数数组传递给 PyCUDA 内核

Passing complex number arrays into PyCUDA Kernel

python

cuda

pycuda