import numpy
from reikna.cluda import Snippet
import reikna.helpers as helpers
from reikna.cluda import dtypes
from reikna.cluda import OutOfResourcesError
from reikna.core import Computation, Parameter, Annotation, Type
from reikna.algorithms import Transpose
TEMPLATE = helpers.template_for(__file__)
class Predicate:
"""
A predicate used in :py:class:`~reikna.algorithms.Reduce`.
:param operation: a :py:class:`~reikna.cluda.Snippet` object with two parameters
which will take the names of two arguments to join.
:param empty: a numpy scalar with the empty value of the argument
(the one which, being joined by another argument, does not change it).
"""
def __init__(self, operation, empty):
self.operation = operation
self.empty = empty
def __process_modules__(self, process):
return Predicate(process(self.operation), self.empty)
def predicate_sum(dtype):
"""
"""
This module contains :py:class:`~reikna.cluda.Module` factories
which are used to compensate for the lack of complex number operations in OpenCL,
and the lack of C++ synthax which would allow one to write them.
"""
from warnings import warn
import numpy
from reikna.helpers import template_for
from reikna.cluda import dtypes
from reikna.cluda import Module
TEMPLATE = template_for(__file__)
def check_information_loss(out_dtype, expected_dtype):
if dtypes.is_complex(expected_dtype) and not dtypes.is_complex(out_dtype):
warn("Imaginary part ignored during the downcast from " +
str(expected_dtype) + " to " + str(out_dtype),
numpy.ComplexWarning)
def derive_out_dtype(out_dtype, *in_dtypes):
expected_dtype = dtypes.result_type(*in_dtypes)
if out_dtype is None:
out_dtype = expected_dtype
else:
check_information_loss(out_dtype, expected_dtype)
import reikna.helpers as helpers
from reikna.core import Computation, Parameter, Annotation, Type
import reikna.cluda.dtypes as dtypes
from reikna.cluda import OutOfResourcesError
from reikna.cluda import functions
TEMPLATE = helpers.template_for(__file__)
class MatrixMul(Computation):
"""
Bases: :py:class:`~reikna.core.Computation`
Multiplies two matrices using last two dimensions and batching over remaining dimensions.
For batching to work, the products of remaining dimensions should be equal
(then the multiplication will be performed piecewise), or one of them should equal 1
(then the multiplication will be batched over the remaining dimensions of the other matrix).
:param a_arr: an array-like defining the first argument.
:param b_arr: an array-like defining the second argument.
:param out_arr: an array-like definign the output; if not given, both shape and dtype
will be derived from ``a_arr`` and ``b_arr``.
:param block_width_override: if provided, it will used as a block size of
the multiplication kernel.
:param transposed_a: if ``True``, the first matrix will be transposed
before the multiplication.
:param transposed_b: if ``True``, the second matrix will be transposed
before the multiplication.
.. py:method:: compiled_signature(output:o, matrix_a:i, matrix_b:i)
def _build_plan(self, plan_factory, device_params, output, magspec, input_,
rect_x, rect_y, rect_r):
"""Assemble the Computation."""
"""Assert that the input/output size is quadratic and the same"""
assert output.shape[0] == output.shape[1] == input_.shape[
0] == input_.shape[1]
plan = plan_factory()
""" Initialize the temporary arrays.
mirrored contains the mirrored variant of the input image, thus has to
be four times the size.
"""
mirrored = np.empty(shape=[2 * x for x in input_.shape],
dtype=np.complex128)
temp = plan.temp_array_like(input_)
temp2 = plan.temp_array_like(input_)
mirror = plan.temp_array_like(mirrored)
mirror_temp1 = plan.temp_array_like(mirrored)
mirror_temp2 = plan.temp_array_like(mirrored)
mirror_sin = plan.temp_array_like(mirrored)
mirror_cos = plan.temp_array_like(mirrored)
r2s_c = plan.persistent_array(r2s(output.shape[0]).astype(np.float64))
"""Initialize computations for the fourier transforms, always on both
axis."""
fft = FFT(temp, axes=(0, 1))
fft_mirror = FFT(mirror, axes=(0, 1))
fftshift = FFTShift(temp, axes=(0, 1))
template = template_for("kernels") # Read template from kernels.mako
plan.computation_call(fft, temp, input_, 0)
plan.computation_call(fftshift, temp, temp)
plan.kernel_call(template.get_def('magspec'), [magspec, temp],
global_size=magspec.shape)
plan.kernel_call(template.get_def('cropcomp'),
[temp2, temp, rect_x, rect_y, rect_r],
global_size=output.shape)
plan.computation_call(fftshift, temp2, temp2)
plan.computation_call(fft, temp, temp2, 1)
plan.kernel_call(template.get_def('angle'), [output, temp],
global_size=output.shape)
plan.kernel_call(template.get_def('zeroarr'), [mirror],
global_size=mirror.shape)
plan.kernel_call(template.get_def('mirror4'), [mirror, output],
global_size=output.shape)
print(mirror.dtype, output.dtype)
plan.kernel_call(template.get_def('sine'), [mirror_sin, mirror],
global_size=mirror.shape)
plan.kernel_call(template.get_def('cosine'), [mirror_cos, mirror],
global_size=mirror.shape)
plan.computation_call(fft_mirror, mirror_temp1, mirror_sin, 0)
plan.computation_call(fft_mirror, mirror_temp2, mirror_cos, 0)
plan.kernel_call(template.get_def('mulmat'),
[mirror_temp1, r2s_c, mirror_temp1],
global_size=mirror.shape)
plan.kernel_call(template.get_def('mulmat'),
[mirror_temp2, r2s_c, mirror_temp2],
global_size=mirror.shape)
plan.computation_call(fft_mirror, mirror_temp1, mirror_temp1, 1)
plan.computation_call(fft_mirror, mirror_temp2, mirror_temp2, 1)
plan.kernel_call(template.get_def('mulmat'),
[mirror_temp1, mirror_cos, mirror_temp1],
global_size=mirror.shape)
plan.kernel_call(template.get_def('mulmat'),
[mirror_temp2, mirror_sin, mirror_temp2],
global_size=mirror.shape)
plan.kernel_call(template.get_def('submat'),
[mirror_temp1, mirror_temp1, mirror_temp2],
global_size=mirror.shape)
plan.computation_call(fft_mirror, mirror_temp1, mirror_temp1, 0)
plan.kernel_call(template.get_def('divmat'),
[mirror_temp1, mirror_temp1, r2s_c],
global_size=mirror.shape)
plan.computation_call(fft_mirror, mirror_temp1, mirror_temp1, 1)
plan.kernel_call(template.get_def('phicalc'),
[output, mirror_temp1, output],
global_size=output.shape)
return plan
