Python中使用共享数组进行快速FFT的内存对齐

Question

Python中使用共享数组进行快速FFT的内存对齐

pythonnumpymultiprocessingfftmemory-alignment

5

我编写了一个图像处理应用程序，需要完成多个任务，并尽可能实时地完成它们。数据的获取和处理在不同的进程中运行（主要是为了性能原因）。数据本身相当大（2MPix 16位灰度图像）。

我可以按照这篇文章中所描述的方式在进程之间共享数组：如何在 Python 子进程之间传递大型 numpy 数组而无需保存到磁盘？（我使用 numpy-shared 包中的 shmarray 脚本）。我可以对这些数据执行提供的 Numpy FFT，但速度较慢。

调用 FFTW 可能会更快，但为了充分从中受益，我应该在内存对齐的数组上运行操作。

问题是：有没有一种方法可以创建和共享类似于 Numpy 的数组，同时保证它们在内存上对齐？

- bubla

1

你使用哪个Python FFTW绑定？这些绑定应该提供一种分配正确对齐内存块的方法。 - Sven Marnach

你说得没错，但棘手的部分是我需要在多个进程之间共享这些数组。我很确定这些函数不会创建“共享”数组。 - bubla

3个回答

3

继承Sven的回答，该函数将返回任何numpy数组的对齐副本（如果需要）：

import numpy as np
def aligned(a, alignment=16):
    if (a.ctypes.data % alignment) == 0:
        return a

    extra = alignment / a.itemsize
    buf = np.empty(a.size + extra, dtype=a.dtype)
    ofs = (-buf.ctypes.data % alignment) / a.itemsize
    aa = buf[ofs:ofs+a.size].reshape(a.shape)
    np.copyto(aa, a)
    assert (aa.ctypes.data % alignment) == 0
    return aa

- payne

0

我在2021年运行了Payne的答案，但出现了类型错误（Python 3.7，Numpy 1.18.5），因此我已经调整了代码：

def aligned(a, alignment = 16):
    if (a.ctypes.data % alignment) == 0:
        return a
    assert alignment % a.itemsize == 0
    extra = alignment // a.itemsize
    buf = np.empty(a.size + extra, dtype = a.dtype)
    ofs = (-buf.ctypes.data % alignment) // a.itemsize
    aa = buf[ofs:ofs + a.size].reshape(a.shape)
    np.copyto(aa, a)
    assert aa.ctypes.data % alignment == 0
    return aa

我改用整数除法来避免类型错误，并添加了额外的assert进行健全性检查。

- Dan Houghton

虽然这可能会证明有用，但您应该提供详细的解释。也许对先前答案进行编辑或评论会更合适？ - PJProudhon

很不幸，我是一个新成员，既不能评论（“您必须拥有50个声望才能评论”），也不能编辑（“编辑队列已满”）。 - Dan Houghton

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Sven Marnach · Accepted Answer

获取正确对齐内存的最简单标准技巧是多分配一些内存，如果对齐不正确，则跳过前几个字节。如果我没有记错，NumPy数组将始终是8字节对齐的，而FFTW需要16字节对齐才能发挥最佳性能。因此，您只需比所需的多分配8个字节，并在必要时跳过前8个字节。

编辑：这相当容易实现。NumPy数组的数据指针可以作为整数在ctypes.data属性中获得。使用移位块可以通过切片、查看为不同的数据类型和重新整形来实现——所有这些都不会复制数据，而是重复使用相同的buf。

要分配一个16字节对齐的1000x1000的64位浮点数数组，我们可以使用以下代码：

m = n = 1000
dtype = numpy.dtype(numpy.float64)
nbytes = m * n * dtype.itemsize
buf = numpy.empty(nbytes + 16, dtype=numpy.uint8)
start_index = -buf.ctypes.data % 16
a = buf[start_index:start_index + nbytes].view(dtype).reshape(m, n)

现在，a是一个具有所需属性的数组，通过检查a.ctypes.data % 16是否确实为0可以验证。