Enthough Python中的线程FFT

Numpy/SciPy中的快速傅立叶变换（FFT）没有线程。Enthough Python随“英特尔MKL数字库”一起提供，该库支持线程化FFT。如何访问这些例程？

以下代码适用于我在Windows 7 Ultimate 64位上使用的Enthound 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)

---

我的原始答案的更清晰版本如下：

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

---

新的和改进的版本，可处理输入和输出阵列中的任意跨步。 默认情况下，此阵列现在不在适当位置，并创建一个新阵列。 它模仿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