Enthought Python 中的线程化FFT

Question

Numpy/SciPy中的快速傅里叶变换（FFT）不支持线程。Enthought Python自带Intel MKL数值库，可执行多线程FFT。要使用这些例程，该怎么做？

- bdforbes

Answer 1

以下代码在我的 Windows 7 Ultimate 64 位操作系统上，使用 Enthought 7.3-1（64 位）可以正常运行。我没有进行基准测试，但它确实同时使用了所有核心而不仅仅是一个核心。

from ctypes import *

class Mkl_Fft:
    c_double_p = POINTER(c_double)

    def __init__(self,num_threads=8):
        self.dfti = cdll.LoadLibrary("mk2_rt.dll")
        self.dfti.MKL_Set_Num_Threads(num_threads)
        self.Create = self.dfti.DftiCreateDescriptor_d_md
        self.Commit = self.dfti.DftiCommitDescriptor
        self.ComputeForward = self.dfti.DftiComputeForward

    def fft(self,a):
        Desc_Handle = c_void_p(0)
        dims = (c_int*2)(*a.shape)
        DFTI_COMPLEX = c_int(32)
        rank = 2

        self.Create(byref(Desc_Handle), DFTI_COMPLEX, rank, dims )
        self.Commit(Desc_Handle)
        self.ComputeForward(Desc_Handle, a.ctypes.data_as(self.c_double_p) )

使用方法：

import numpy as np
a = np.ones( (32,32), dtype = complex128 )
fft = Mkl_Fft()
fft.fft(a)

- bdforbes

Answer 2

我的原始答案的更简洁版本如下：

from ctypes import *

mkl = cdll.LoadLibrary("mk2_rt.dll")
c_double_p = POINTER(c_double)
DFTI_COMPLEX = c_int(32)
DFTI_DOUBLE = c_int(36)

def fft2(a):
    Desc_Handle = c_void_p(0)
    dims = (c_int*2)(*a.shape)

    mkl.DftiCreateDescriptor(byref(Desc_Handle), DFTI_DOUBLE, DFTI_COMPLEX, 2, dims )
    mkl.DftiCommitDescriptor(Desc_Handle)
    mkl.DftiComputeForward(Desc_Handle, a.ctypes.data_as(c_void_p) )
    mkl.DftiFreeDescriptor(byref(Desc_Handle))

    return a

def ifft2(a):
    Desc_Handle = c_void_p(0)
    dims = (c_int*2)(*a.shape)

    mkl.DftiCreateDescriptor(byref(Desc_Handle), DFTI_DOUBLE, DFTI_COMPLEX, 2, dims )
    mkl.DftiCommitDescriptor(Desc_Handle)
    mkl.DftiComputeBackward(Desc_Handle, a.ctypes.data_as(c_void_p) )
    mkl.DftiFreeDescriptor(byref(Desc_Handle))

    return a

- bdforbes

Answer 3

新的改进版本可以处理输入和输出数组中的任意步长。 默认情况下，该版本不再是原地操作，并且会创建一个新的数组。 它模仿了Numpy FFT例程，但具有不同的归一化。

''' Wrapper to MKL FFT routines '''

import numpy as _np
import ctypes as _ctypes

mkl = _ctypes.cdll.LoadLibrary("mk2_rt.dll")
_DFTI_COMPLEX = _ctypes.c_int(32)
_DFTI_DOUBLE = _ctypes.c_int(36)
_DFTI_PLACEMENT = _ctypes.c_int(11)
_DFTI_NOT_INPLACE = _ctypes.c_int(44)
_DFTI_INPUT_STRIDES = _ctypes.c_int(12)
_DFTI_OUTPUT_STRIDES = _ctypes.c_int(13)

def fft2(a, out=None):
    ''' 
Forward two-dimensional double-precision complex-complex FFT.
Uses the Intel MKL libraries distributed with Enthought Python.
Normalisation is different from Numpy!
By default, allocates new memory like 'a' for output data.
Returns the array containing output data.
'''

    assert a.dtype == _np.complex128
    assert len(a.shape) == 2

    inplace = False

    if out is a:
        inplace = True

    elif out is not None:
        assert out.dtype == _np.complex128
        assert a.shape == out.shape
        assert not _np.may_share_memory(a, out)

    else:
        out = _np.empty_like(a)

    Desc_Handle = _ctypes.c_void_p(0)
    dims = (_ctypes.c_int*2)(*a.shape)

    mkl.DftiCreateDescriptor(_ctypes.byref(Desc_Handle), _DFTI_DOUBLE, _DFTI_COMPLEX, _ctypes.c_int(2), dims )

    #Set input strides if necessary
    if not a.flags['C_CONTIGUOUS']:
        in_strides = (_ctypes.c_int*3)(0, a.strides[0]/16, a.strides[1]/16)
        mkl.DftiSetValue(Desc_Handle, _DFTI_INPUT_STRIDES, _ctypes.byref(in_strides))

    if inplace:
        #Inplace FFT
        mkl.DftiCommitDescriptor(Desc_Handle)
        mkl.DftiComputeForward(Desc_Handle, a.ctypes.data_as(_ctypes.c_void_p) )

    else:
        #Not-inplace FFT
        mkl.DftiSetValue(Desc_Handle, _DFTI_PLACEMENT, _DFTI_NOT_INPLACE)

        #Set output strides if necessary
        if not out.flags['C_CONTIGUOUS']:
            out_strides = (_ctypes.c_int*3)(0, out.strides[0]/16, out.strides[1]/16)
            mkl.DftiSetValue(Desc_Handle, _DFTI_OUTPUT_STRIDES, _ctypes.byref(out_strides))

        mkl.DftiCommitDescriptor(Desc_Handle)
        mkl.DftiComputeForward(Desc_Handle, a.ctypes.data_as(_ctypes.c_void_p), out.ctypes.data_as(_ctypes.c_void_p) )

    mkl.DftiFreeDescriptor(_ctypes.byref(Desc_Handle))

    return out

def ifft2(a, out=None):
    ''' 
Backward two-dimensional double-precision complex-complex FFT.
Uses the Intel MKL libraries distributed with Enthought Python.
Normalisation is different from Numpy!
By default, allocates new memory like 'a' for output data.
Returns the array containing output data.
'''

    assert a.dtype == _np.complex128
    assert len(a.shape) == 2

    inplace = False

    if out is a:
        inplace = True

    elif out is not None:
        assert out.dtype == _np.complex128
        assert a.shape == out.shape
        assert not _np.may_share_memory(a, out)

    else:
        out = _np.empty_like(a)

    Desc_Handle = _ctypes.c_void_p(0)
    dims = (_ctypes.c_int*2)(*a.shape)

    mkl.DftiCreateDescriptor(_ctypes.byref(Desc_Handle), _DFTI_DOUBLE, _DFTI_COMPLEX, _ctypes.c_int(2), dims )

    #Set input strides if necessary
    if not a.flags['C_CONTIGUOUS']:
        in_strides = (_ctypes.c_int*3)(0, a.strides[0]/16, a.strides[1]/16)
        mkl.DftiSetValue(Desc_Handle, _DFTI_INPUT_STRIDES, _ctypes.byref(in_strides))

    if inplace:
        #Inplace FFT
        mkl.DftiCommitDescriptor(Desc_Handle)
        mkl.DftiComputeBackward(Desc_Handle, a.ctypes.data_as(_ctypes.c_void_p) )

    else:
        #Not-inplace FFT
        mkl.DftiSetValue(Desc_Handle, _DFTI_PLACEMENT, _DFTI_NOT_INPLACE)

        #Set output strides if necessary
        if not out.flags['C_CONTIGUOUS']:
            out_strides = (_ctypes.c_int*3)(0, out.strides[0]/16, out.strides[1]/16)
            mkl.DftiSetValue(Desc_Handle, _DFTI_OUTPUT_STRIDES, _ctypes.byref(out_strides))

        mkl.DftiCommitDescriptor(Desc_Handle)
        mkl.DftiComputeBackward(Desc_Handle, a.ctypes.data_as(_ctypes.c_void_p), out.ctypes.data_as(_ctypes.c_void_p) )

    mkl.DftiFreeDescriptor(_ctypes.byref(Desc_Handle))

    return out

- bdforbes