def setitem_computation(dest, source):
"""
Returns a compiled computation that broadcasts ``source`` to ``dest``,
where ``dest`` is a GPU array, and ``source`` is either a GPU array or a scalar.
"""
if len(source.shape) == 0:
trf = transformations.broadcast_param(dest)
return PureParallel.from_trf(trf, guiding_array=trf.output)
else:
source_dt = Type.from_value(source).with_dtype(dest.dtype)
trf = transformations.copy(source_dt, dest)
comp = PureParallel.from_trf(trf, guiding_array=trf.output)
cast_trf = transformations.cast(source, dest.dtype)
comp.parameter.input.connect(cast_trf, cast_trf.output, src_input=cast_trf.input)
return comp
def test_trf_with_guiding_output(thr):
"""
Test the creation of ``PureParallel`` out of a transformation,
with an output parameter as a guiding array.
"""
N = 1000
coeff = 3
dtype = numpy.float32
arr_t = Type(dtype, shape=N)
trf = mul_param(arr_t, dtype)
p = PureParallel.from_trf(trf, trf.output)
# The new PureParallel has to preserve the parameter list of the original transformation.
assert list(p.signature.parameters.values()) == list(
trf.signature.parameters.values())
a = get_test_array_like(p.parameter.input)
a_dev = thr.to_device(a)
res_dev = thr.empty_like(p.parameter.output)
pc = p.compile(thr)
pc(res_dev, a_dev, coeff)
assert diff_is_negligible(res_dev.get(), a * 3)
def _build_plan(
self, plan_factory, device_params,
result_a, result_b, result_cv, messages, key, noises_a, noises_b):
plan = plan_factory()
mul_key = MatrixMulVector(noises_a)
fill_b_cv = Transformation([
Parameter('result_b', Annotation(result_b, 'o')),
Parameter('result_cv', Annotation(result_cv, 'o')),
Parameter('messages', Annotation(messages, 'i')),
Parameter('noises_a_times_key', Annotation(noises_b, 'i')),
Parameter('noises_b', Annotation(noises_b, 'i'))],
"""
${result_b.store_same}(
${noises_b.load_same}
+ ${messages.load_same}
+ ${noises_a_times_key.load_same});
${result_cv.store_same}(${noise**2});
""",
connectors=['noises_a_times_key'],
render_kwds=dict(noise=self._noise))
mul_key.parameter.output.connect(
fill_b_cv, fill_b_cv.noises_a_times_key,
b=fill_b_cv.result_b, cv=fill_b_cv.result_cv, messages=fill_b_cv.messages,
noises_b=fill_b_cv.noises_b)
plan.computation_call(mul_key, result_b, result_cv, messages, noises_b, noises_a, key)
plan.computation_call(
PureParallel.from_trf(transformations.copy(noises_a)),
result_a, noises_a)
return plan
def _build_plan(self, plan_factory, device_params, output, input_, inverse):
if helpers.product([input_.shape[i] for i in self._axes]) == 1:
return self._build_trivial_plan(plan_factory, output, input_)
plan = plan_factory()
axes = tuple(sorted(self._axes))
shape = list(input_.shape)
if all(shape[axis] % 2 == 0 for axis in axes):
# If all shift axes have even length, it is possible to perform the shift inplace
# (by swapping pairs of elements).
# Note that the inplace fftshift is its own inverse.
shape[axes[0]] //= 2
plan.kernel_call(
TEMPLATE.get_def('fftshift_inplace'), [output, input_],
kernel_name="kernel_fftshift_inplace",
global_size=shape,
render_kwds=dict(axes=axes))
else:
# Resort to an out-of-place shift to a temporary array and then copy.
temp = plan.temp_array_like(output)
plan.kernel_call(
TEMPLATE.get_def('fftshift_outplace'), [temp, input_, inverse],
kernel_name="kernel_fftshift_outplace",
global_size=shape,
render_kwds=dict(axes=axes))
copy_trf = copy(input_, out_arr_t=output)
copy_comp = PureParallel.from_trf(copy_trf, copy_trf.input)
plan.computation_call(copy_comp, output, temp)
return plan
def test_tgsw_polynomial_decomp_trf(thread):
shape = (2, 3)
params = NuFHEParameters()
tgsw_params = params.tgsw_params
decomp_length = tgsw_params.decomp_length
mask_size = tgsw_params.tlwe_params.mask_size
polynomial_degree = tgsw_params.tlwe_params.polynomial_degree
sample = get_test_array(shape + (mask_size + 1, polynomial_degree), Torus32, (0, 1000))
result = numpy.empty(shape + (mask_size + 1, decomp_length, polynomial_degree), dtype=Int32)
sample_dev = thread.to_device(sample)
result_dev = thread.empty_like(result)
trf = get_tgsw_polynomial_decomp_trf(tgsw_params, shape)
test = PureParallel.from_trf(trf, guiding_array='result').compile(thread)
ref = tgsw_polynomial_decomp_trf_reference(tgsw_params, shape)
test(result_dev, sample_dev)
result_test = result_dev.get()
ref(result, sample)
assert (result == result_test).all()
def test_from_trf(thr, guiding_array):
"""
Test the creation of ``PureParallel`` out of a transformation
with various values of the guiding array.
"""
N = 1000
coeff = 3
dtype = numpy.float32
arr_t = Type(dtype, shape=N)
trf = mul_param(arr_t, dtype)
if guiding_array == 'input':
arr = trf.input
elif guiding_array == 'output':
arr = trf.output
elif guiding_array == 'none':
arr = None
p = PureParallel.from_trf(trf, guiding_array=arr)
# The new PureParallel has to preserve the parameter list of the original transformation.
assert list(p.signature.parameters.values()) == list(trf.signature.parameters.values())
a = get_test_array_like(p.parameter.input)
a_dev = thr.to_device(a)
res_dev = thr.empty_like(p.parameter.output)
pc = p.compile(thr)
pc(res_dev, a_dev, coeff)
assert diff_is_negligible(res_dev.get(), a * 3)
def test_from_trf(thr, guiding_array):
"""
Test the creation of ``PureParallel`` out of a transformation
with various values of the guiding array.
"""
N = 1000
coeff = 3
dtype = numpy.float32
arr_t = Type(dtype, shape=N)
trf = mul_param(arr_t, dtype)
if guiding_array == 'input':
arr = trf.input
elif guiding_array == 'output':
arr = trf.output
elif guiding_array == 'none':
arr = None
p = PureParallel.from_trf(trf, guiding_array=arr)
# The new PureParallel has to preserve the parameter list of the original transformation.
assert list(p.signature.parameters.values()) == list(
trf.signature.parameters.values())
a = get_test_array_like(p.parameter.input)
a_dev = thr.to_device(a)
res_dev = thr.empty_like(p.parameter.output)
pc = p.compile(thr)
pc(res_dev, a_dev, coeff)
assert diff_is_negligible(res_dev.get(), a * 3)
def _build_plan(self, plan_factory, device_params, output, input_,
inverse):
if helpers.product([input_.shape[i] for i in self._axes]) == 1:
return self._build_trivial_plan(plan_factory, output, input_)
plan = plan_factory()
axes = tuple(sorted(self._axes))
shape = list(input_.shape)
if all(shape[axis] % 2 == 0 for axis in axes):
# If all shift axes have even length, it is possible to perform the shift inplace
# (by swapping pairs of elements).
# Note that the inplace fftshift is its own inverse.
shape[axes[0]] //= 2
plan.kernel_call(TEMPLATE.get_def('fftshift_inplace'),
[output, input_],
kernel_name="kernel_fftshift_inplace",
global_size=shape,
render_kwds=dict(axes=axes))
else:
# Resort to an out-of-place shift to a temporary array and then copy.
temp = plan.temp_array_like(output)
plan.kernel_call(TEMPLATE.get_def('fftshift_outplace'),
[temp, input_, inverse],
kernel_name="kernel_fftshift_outplace",
global_size=shape,
render_kwds=dict(axes=axes))
copy_trf = copy(input_, out_arr_t=output)
copy_comp = PureParallel.from_trf(copy_trf, copy_trf.input)
plan.computation_call(copy_comp, output, temp)
return plan
def test_trf_with_guiding_output(thr):
"""
Test the creation of ``PureParallel`` out of a transformation,
with an output parameter as a guiding array.
"""
N = 1000
coeff = 3
dtype = numpy.float32
arr_t = Type(dtype, shape=N)
trf = mul_param(arr_t, dtype)
p = PureParallel.from_trf(trf, trf.output)
# The new PureParallel has to preserve the parameter list of the original transformation.
assert list(p.signature.parameters.values()) == list(trf.signature.parameters.values())
a = get_test_array_like(p.parameter.input)
a_dev = thr.to_device(a)
res_dev = thr.empty_like(p.parameter.output)
pc = p.compile(thr)
pc(res_dev, a_dev, coeff)
assert diff_is_negligible(res_dev.get(), a * 3)
def __init__(self, arr):
copy_trf = copy(arr, out_arr_t=arr)
self._copy_comp = PureParallel.from_trf(copy_trf, copy_trf.input)
Computation.__init__(self, [
Parameter('outer_output', Annotation(arr, 'o')),
Parameter('outer_input', Annotation(arr, 'i'))])
def setitem_computation(dest, source, is_array):
"""
Returns a compiled computation that broadcasts ``source`` to ``dest``,
where ``dest`` is a GPU array, and ``source`` is either a GPU array or a scalar.
"""
if is_array:
source_dt = Type.from_value(source).with_dtype(dest.dtype)
trf = transformations.copy(source_dt, dest)
comp = PureParallel.from_trf(trf, guiding_array=trf.output)
cast_trf = transformations.cast(source, dest.dtype)
comp.parameter.input.connect(cast_trf,
cast_trf.output,
src_input=cast_trf.input)
return comp
else:
trf = transformations.broadcast_param(dest)
return PureParallel.from_trf(trf, guiding_array=trf.output)
def __init__(self, arr):
copy_trf = copy(arr, out_arr_t=arr)
self._copy_comp = PureParallel.from_trf(copy_trf, copy_trf.input)
Computation.__init__(self, [
Parameter('outer_output', Annotation(arr, 'o')),
Parameter('outer_input', Annotation(arr, 'i'))
])
def _build_trivial_plan(self, plan_factory, output, input_):
# Trivial problem. Need to add a dummy kernel
# because we still have to run transformations.
plan = plan_factory()
copy_trf = copy(input_, out_arr_t=output)
copy_comp = PureParallel.from_trf(copy_trf, copy_trf.input)
plan.computation_call(copy_comp, output, input_)
return plan
def _build_trivial_plan(self, plan_factory, output, input_):
# Trivial problem. Need to add a dummy kernel
# because we still have to run transformations.
plan = plan_factory()
copy_trf = copy(input_, out_arr_t=output)
copy_comp = PureParallel.from_trf(copy_trf, copy_trf.input)
plan.computation_call(copy_comp, output, input_)
return plan
def _build_plan(self, plan_factory, device_params, array, shift):
plan = plan_factory()
temp = plan.temp_array_like(array)
plan.computation_call(roll_computation(array, self._axis), temp, array, shift)
tr = transformations.copy(temp, out_arr_t=array)
copy_comp = PureParallel.from_trf(tr, guiding_array=tr.output)
plan.computation_call(copy_comp, array, temp)
return plan
def _build_plan(self, plan_factory, device_params, array, shift):
plan = plan_factory()
temp = plan.temp_array_like(array)
plan.computation_call(roll_computation(array, self._axis), temp, array,
shift)
tr = transformations.copy(temp, out_arr_t=array)
copy_comp = PureParallel.from_trf(tr, guiding_array=tr.output)
plan.computation_call(copy_comp, array, temp)
return plan
def test_tlwe_transformed_add_mul_to_trf(thread):
shape = (2, 3)
params = NuFHEParameters(transform_type='NTT')
perf_params = PerformanceParameters(params).for_device(
thread.device_params)
tgsw_params = params.tgsw_params
decomp_length = tgsw_params.decomp_length
mask_size = tgsw_params.tlwe_params.mask_size
polynomial_degree = tgsw_params.tlwe_params.polynomial_degree
transform_type = tgsw_params.tlwe_params.transform_type
transform = get_transform(transform_type)
tlength = transform.transformed_length(polynomial_degree)
tdtype = transform.transformed_dtype()
result_shape = shape + (mask_size + 1, tlength)
sample_shape = shape + (mask_size + 1, decomp_length, tlength)
bk_len = 10
bootstrap_key_shape = (bk_len, mask_size + 1, decomp_length, mask_size + 1,
tlength)
bk_row_idx = 2
result = numpy.empty(result_shape, tdtype)
sample = get_test_array(sample_shape, 'ff_number')
bootstrap_key = get_test_array(bootstrap_key_shape, 'ff_number')
result_dev = thread.empty_like(result)
sample_dev = thread.to_device(sample)
bootstrap_key_dev = thread.to_device(bootstrap_key)
trf = get_tlwe_transformed_add_mul_to_trf(tgsw_params, shape, bk_len,
perf_params)
test = PureParallel.from_trf(trf, guiding_array='result').compile(thread)
ref = tlwe_transformed_add_mul_to_trf_reference(tgsw_params, shape, bk_len,
perf_params)
test(result_dev, sample_dev, bootstrap_key_dev, bk_row_idx)
result_test = result_dev.get()
ref(result, sample, bootstrap_key, bk_row_idx)
if numpy.issubdtype(tdtype, numpy.integer):
assert (result == result_test).all()
else:
assert numpy.allclose(result, result_test)
def test_array_views(thr):
a = get_test_array((6, 8, 10), numpy.int32)
a_dev = thr.to_device(a)
b_dev = thr.empty_like(a)
in_view = a_dev[2:4, ::2, ::-1]
out_view = b_dev[4:, 1:5, :]
move = PureParallel.from_trf(
transformations.copy(in_view, out_arr_t=out_view),
guiding_array='output').compile(thr)
move(out_view, in_view)
b_res = b_dev.get()[4:, 1:5, :]
b_ref = a[2:4, ::2, ::-1]
assert diff_is_negligible(b_res, b_ref)
def _build_plan(self, plan_factory, device_params, result, phase):
plan = plan_factory()
tr = Transformation([
Parameter('result', Annotation(result, 'o')),
Parameter('phase', Annotation(phase, 'i')),
],
"""
<%
interv = 2**32 // mspace_size
half_interv = interv // 2
%>
${phase.ctype} phase = ${phase.load_same};
${result.store_same}(((unsigned int)phase + ${half_interv}) / ${interv});
""",
render_kwds=dict(mspace_size=self._mspace_size,
uint64=dtypes.ctype(
numpy.uint64)),
connectors=['result', 'phase'])
plan.computation_call(
PureParallel.from_trf(tr, guiding_array='result'), result, phase)
return plan