/usr/lib/python3/dist-packages/gpyfft/fft.py is in python3-gpyfft 0.7.0-1.
This file is owned by root:root, with mode 0o644.
The actual contents of the file can be viewed below.
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 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 | from __future__ import absolute_import, division, print_function
from .gpyfftlib import GpyFFT
import gpyfft.gpyfftlib as gfft
import pyopencl as cl
GFFT = GpyFFT(debug=False)
import pyopencl as cl
import numpy as np
# TODO:
class FFT(object):
def __init__(self, context, queue, in_array, out_array=None, axes = None,
fast_math = False,
real=False,
callbacks=None, #dict: 'pre', 'post'
):
# Callbacks: dict(pre=b'pre source (kernel named pre!)')
self.context = context
self.queue = queue
t_strides_in, t_distance_in, t_batchsize_in, t_shape, axes_transform = self.calculate_transform_strides(axes, in_array)
if out_array is not None:
t_inplace = False
t_strides_out, t_distance_out, t_batchsize_out, t_shape_out, foo = self.calculate_transform_strides(
axes, out_array)
#assert t_batchsize_out == t_batchsize_in and t_shape == t_shape_out, 'input and output size does not match' #TODO: fails for real-to-complex
else:
t_inplace = True
t_strides_out, t_distance_out = t_strides_in, t_distance_in
#assert np.issubclass(in_array.dtype, np.complexfloating) and \
# np.issubclass(in_array.dtype, np.complexfloating), \
#precision (+ fast_math!)
#complex64 <-> complex64
#complex128 <-> complex128
if in_array.dtype in (np.float32, np.complex64):
precision = gfft.CLFFT_SINGLE
elif in_array.dtype in (np.float64, np.complex128):
precision = gfft.CLFFT_DOUBLE
#TODO: add assertions that precision match
if in_array.dtype in (np.float32, np.float64):
layout_in = gfft.CLFFT_REAL
layout_out = gfft.CLFFT_HERMITIAN_INTERLEAVED
expected_out_shape = list(in_array.shape)
expected_out_shape[axes_transform[0]] = expected_out_shape[axes_transform[0]]//2 + 1
assert out_array.shape == tuple(expected_out_shape), \
'output array shape %s does not match expected shape: %s'%(out_array.shape,expected_out_shape)
elif in_array.dtype in (np.complex64, np.complex128):
if not real:
layout_in = gfft.CLFFT_COMPLEX_INTERLEAVED
layout_out = gfft.CLFFT_COMPLEX_INTERLEAVED
else:
# complex-to-real transform
layout_in = gfft.CLFFT_HERMITIAN_INTERLEAVED
layout_out = gfft.CLFFT_REAL
t_shape = t_shape_out
self.t_shape = t_shape
self.batchsize = t_batchsize_in
plan = GFFT.create_plan(context, t_shape)
plan.inplace = t_inplace
plan.strides_in = t_strides_in
plan.strides_out = t_strides_out
plan.distances = (t_distance_in, t_distance_out)
plan.batch_size = self.batchsize
plan.precision = precision
plan.layouts = (layout_in, layout_out)
if callbacks is not None:
if callbacks.has_key('pre'):
plan.set_callback(b'pre',
callbacks['pre'],
'pre')
if 'post' in callbacks:
plan.set_callback(b'post',
callbacks['post'],
'post')
if False:
print('axes', axes )
print('in_array.shape: ', in_array.shape)
print('in_array.strides/itemsize', tuple(s // in_array.dtype.itemsize for s in in_array.strides))
print('shape transform ', t_shape)
print('layout_in ', str(layout_in).split('.')[1])
print('t_strides ', t_strides_in)
print('distance_in ', t_distance_in)
print('batchsize ', t_batchsize_in)
print('layout_out ', str(layout_out).split('.')[1])
print('t_stride_out ', t_strides_out)
print('inplace ', t_inplace)
plan.bake(self.queue)
temp_size = plan.temp_array_size
if temp_size:
#print 'temp_size:', plan.temp_array_size
self.temp_buffer = cl.Buffer(self.context, cl.mem_flags.READ_WRITE, size = temp_size)
else:
self.temp_buffer = None
self.plan = plan
self.data = in_array
self.result = out_array
@classmethod
def calculate_transform_strides(cls, axes, array):
shape = np.array(array.shape)
strides = np.array(array.strides)
dtype = array.dtype
ddim = len(shape) #dimensionality data
#transform along all axes if transform axes are not given (None)
axes_transform = np.arange(ddim) if axes is None else np.array(axes)
tdim = len(axes_transform) #dimensionality transform
assert tdim <= ddim
# transform negative axis values (e.g. -1 for last axis) to positive
axes_transform[axes_transform<0] += ddim
# remaining, non-transformed axes
axes_notransform = np.lib.arraysetops.setdiff1d(range(ddim), axes_transform)
#sort non-transformed axes by strides
axes_notransform = axes_notransform[np.argsort(strides[axes_notransform])]
#print "axes_notransformed sorted", axes_notransform
# -> list of collapsable axes, [ [x,y], [z] ]
collapsable_axes_list = [] #result
collapsable_axes_candidates = axes_notransform[:1].tolist() #intermediate list of collapsable axes (magic code to get empty list if axes_notransform is empty)
for a in axes_notransform[1:]:
if strides[a] == strides[collapsable_axes_candidates[-1]] * shape[collapsable_axes_candidates[-1]]:
collapsable_axes_candidates.append(a) #add axes to intermediate list of collapsable axes
else: #does not fit into current intermediate list of collapsable axes
collapsable_axes_list.append(collapsable_axes_candidates) #store away intermediate list
collapsable_axes_candidates = [a] #start new intermediate list
collapsable_axes_list.append(collapsable_axes_candidates) #append last intermediate list to
assert len(collapsable_axes_list) == 1 #all non-transformed axes collapsed
axes_notransform = collapsable_axes_list[0] #all axes collapsable: take single group of collapsable axes
t_distances = strides[axes_notransform]//dtype.itemsize
if len(t_distances) == 0:
t_distance = 0
else:
t_distance = t_distances[0] #takes smalles stride (axes_notransform have been sorted by stride size)
batchsize = np.prod(shape[axes_notransform])
t_shape = shape[axes_transform]
t_strides = strides[axes_transform]//dtype.itemsize
return (tuple(t_strides), t_distance, batchsize, tuple(t_shape), tuple(axes_transform)) #, tuple(axes_notransform))
def enqueue(self, forward = True, wait_for_events = None):
return self.enqueue_arrays(forward=forward, data=self.data, result=self.result, wait_for_events=wait_for_events)
def enqueue_arrays(self, data = None, result = None, forward = True, wait_for_events = None):
"""enqueue transform"""
if data is None:
data = self.data
else:
assert data.shape == self.data.shape
assert data.strides == self.data.strides
assert data.dtype == self.data.dtype
if result is None:
result = self.result
else:
assert result.shape == self.result.shape
assert result.strides == self.result.strides
assert result.dtype == self.result.dtype
# get buffer for data
if data.offset != 0:
data = data._new_with_changes(data=data.base_data[data.offset:], offset=0)
data_buffer = data.base_data
if result is not None:
# get buffer for result
if result.offset != 0:
result = result._new_with_changes(data=result.base_data[result.offset:], offset=0)
result_buffer = result.base_data
events = self.plan.enqueue_transform((self.queue,), (data_buffer,), (result_buffer),
direction_forward = forward, temp_buffer = self.temp_buffer, wait_for_events = wait_for_events)
else:
events = self.plan.enqueue_transform((self.queue,), (data_buffer,),
direction_forward = forward, temp_buffer = self.temp_buffer, wait_for_events = wait_for_events)
return events
def update_arrays(self, input_array, output_array):
pass
|