Сложное деление на мнимую часть

Я ищу быструю реализацию (x / y).imag, куда x, y представляют собой сложные 2D-массивы (тензоры PyTorch уже на GPU). Мой подход заключается в переносе вычислений в ядро ​​CUDA через cupy, сопрягая C и Python. Требования:

• Быстрее, чем (x / y).imag
• Использование такой же или меньшей памяти, чем (x / y).imag
• Выход такой же, как (x / y).imag в пределах точности с плавающей запятой
• #includes должен поддерживаться CuPy (который в значительной степени ограничен стандартными библиотеками CUDA afaik)

Моя попытка:

extern "C" __global__
void cdiv_imag(float x[240][240000][2], float y[240][240000][2],
               float out[240][240000]) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int j = blockIdx.y * blockDim.y + threadIdx.y;
    if (i >= 240 || j >= 240000)
        return;

    float A = x[i][j][0];
    float B = x[i][j][1];
    float C = y[i][j][0];
    float D = y[i][j][1];

    out[i][j] = (B*C - A*D) / (C*C + D*D);
}

Это использует половину памяти, так как x / y дает временный массив действительных и воображаемых, но все же в конечном итоге медленнее, согласно моим тестам на GTX 1070:

0.02320581099999992 sec (avg of 1000 runs)
0.01924530900000002 sec (avg of 1000 runs)

Полный код; Win 10 x64, Python 3.7.9, CuPy 8.3.0, PyTorch 1.8.0.

import torch
import numpy as np
import cupy as cp
from collections import namedtuple
from timeit import timeit
 
def timer(fn, number=1000):
    print(timeit(fn, number=number) / number) 
 
Stream = namedtuple('Stream', ['ptr'])
 
#%%##################################################################
kernel = cp.RawKernel(r'''
extern "C" __global__
void cdiv_imag(float x[240][240000][2], float y[240][240000][2],
               float out[240][240000]) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int j = blockIdx.y * blockDim.y + threadIdx.y;
    if (i >= 240 || j >= 240000)
        return;
 
    float A = x[i][j][0];
    float B = x[i][j][1]; 
    float C = y[i][j][0];
    float D = y[i][j][1];
 
    out[i][j] = (B*C - A*D) / (C*C + D*D);
}
 
''', 'cdiv_imag')
#%%##################################################################
shape = (240, 240000)
x = torch.randn(shape, dtype=torch.complex64, device="cuda")
y = torch.fliplr(x.clone().detach())
 
threadsperblock = (32, 32)
blockspergrid_x = int(np.ceil(y.shape[0] / threadsperblock[0]))
blockspergrid_y = int(np.ceil(y.shape[1] / threadsperblock[1]))
blockspergrid = (blockspergrid_x, blockspergrid_y)
 
#%%##################################################################
def f1(x, y):
    out = torch.zeros(shape, dtype=torch.float32, device="cuda")
    xv, yv = torch.view_as_real(x), torch.view_as_real(y)
    kernel(
        grid=blockspergrid, block=threadsperblock,
        args=[xv.data_ptr(), yv.data_ptr(), out.data_ptr()],
        stream=Stream(ptr=torch.cuda.current_stream().cuda_stream))
    return out    
 
def f2(x, y):
    return (x / y).imag
 
# ensure outputs agree within float precision
adiff = torch.abs(f1(x, y) - f2(x, y))
assert adiff.mean() < 1e-8
 
# dummy run to factor out caching & warm up GPU
for _ in range(2000):
  f1(x, y)
  f2(x, y)
 
#%%##################################################################
timer(lambda: f1(x, y))
timer(lambda: f2(x, y))