-
Notifications
You must be signed in to change notification settings - Fork 0
/
cufft_import.py
72 lines (51 loc) · 2.46 KB
/
cufft_import.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
# -*- coding: utf-8 -*-
__all__ = [
"cufft_c2c",
"cufft_c2r",
"cufft_plan",
"cufft_r2c",
"cufft_setstream",
]
import os
from numpy.ctypeslib import ndpointer
from ctypes import (c_int,
c_void_p)
# Load the shared library
from shared_utils import load_lib
lib_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "lib"))
cufft_lib = load_lib(lib_path,"cufft")
# Define argtypes for all functions to import
argtype_defs = {
"cufft_addredundants" : [c_void_p, #Device pointer to non-redundant data
c_void_p, #Device pointer to redundant data
c_int, c_int, #Dimensions of the array
c_int, #Data type identifier
c_void_p], #Pointer to CUDA stream
"cufft_c2c" : [c_void_p, #Pointer to the cuFFT plan
c_void_p, #Device pointer to the input data
c_void_p, #Device pointer to the output data
c_int, #Direction of the fft (forward or inverse)
c_int], #Data type identifier
"cufft_c2r" : [c_void_p, #Pointer to the cuFFT plan
c_void_p, #Device pointer to the input data
c_void_p, #Device pointer to the output data
c_int], #Data type identifier
"cufft_r2c" : [c_void_p, #Pointer to the cuFFT plan
c_void_p, #Device pointer to the input data
c_void_p, #Device pointer to non-redundant data
c_int], #Data type identifier
"cufft_plan" : [ndpointer(), #Dimensions of the fft
c_int, #Type of fft
c_int], #Batch size of fft
"cufft_setstream" : [c_void_p, #Pointer to the cuFFT plan
c_void_p] #Pointer to CUDA stream
}
restype_defs = {
"cufft_plan": c_void_p
}
# Import functions from DLL
for func, argtypes in argtype_defs.items():
restype = restype_defs.get(func)
vars().update({func: cufft_lib[func]})
vars()[func].argtypes = argtypes
vars()[func].restype = restype