python - 将复数数组传递到 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