Skip to content

Commit f15367d

Browse files
wolfvwolfv
committed
Plan broadcast and permutation assignments
1 parent d456997 commit f15367d

3 files changed

Lines changed: 379 additions & 2 deletions

File tree

‎include/xtensor/core/xassign.hpp‎

Lines changed: 194 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,9 @@
1111
#define XTENSOR_ASSIGN_HPP
1212

1313
#include <algorithm>
14+
#include <array>
1415
#include <functional>
16+
#include <tuple>
1517
#include <type_traits>
1618
#include <utility>
1719

@@ -58,6 +60,12 @@ namespace xt
5860
template <class E1, class E2>
5961
void strided_assign(E1& e1, const E2& e2, std::true_type /*enable*/);
6062

63+
template <class E1, class E2>
64+
bool try_permutation_assign(E1&, const E2&)
65+
{
66+
return false;
67+
}
68+
6169
/************************
6270
* xexpression_assigner *
6371
************************/
@@ -334,6 +342,187 @@ namespace xt
334342
{
335343
};
336344

345+
template <class E, class = void>
346+
struct runtime_plan_expression : std::false_type
347+
{
348+
};
349+
350+
template <class E>
351+
struct runtime_plan_expression<
352+
E,
353+
std::void_t<decltype(std::declval<const E&>().storage().data())>>
354+
: std::bool_constant<
355+
data_interface_expression<E> && std::is_arithmetic<typename E::value_type>::value
356+
&& std::is_pointer<decltype(std::declval<const E&>().storage().data())>::value
357+
&& std::is_same<
358+
std::remove_cv_t<std::remove_pointer_t<decltype(std::declval<const E&>().storage().data())>>,
359+
typename E::value_type>::value>
360+
{
361+
};
362+
363+
template <class T>
364+
struct runtime_plan_expression<xscalar<T>, void> : std::bool_constant<std::is_arithmetic<T>::value>
365+
{
366+
};
367+
368+
template <class F, class... CT>
369+
struct runtime_plan_expression<xfunction<F, CT...>, void>
370+
: std::conjunction<runtime_plan_expression<std::decay_t<CT>>...>
371+
{
372+
};
373+
374+
template <class E, class = void>
375+
struct runtime_plan_output : std::false_type
376+
{
377+
};
378+
379+
template <class E>
380+
struct runtime_plan_output<
381+
E,
382+
std::void_t<
383+
decltype(std::declval<E&>().storage().data()),
384+
decltype(std::declval<E&>().data_offset())>>
385+
: std::bool_constant<
386+
E::contiguous_layout && E::static_layout == layout_type::row_major
387+
&& std::is_pointer<decltype(std::declval<E&>().storage().data())>::value
388+
&& std::is_same<
389+
std::remove_cv_t<std::remove_pointer_t<decltype(std::declval<E&>().storage().data())>>,
390+
typename E::value_type>::value>
391+
{
392+
};
393+
394+
template <class E>
395+
struct runtime_leaf_plan
396+
{
397+
using value_type = typename E::value_type;
398+
const value_type* data;
399+
std::array<std::ptrdiff_t, 3> strides;
400+
401+
value_type operator()(std::size_t i, std::size_t j, std::size_t k) const
402+
{
403+
const auto offset = static_cast<std::ptrdiff_t>(i) * strides[0]
404+
+ static_cast<std::ptrdiff_t>(j) * strides[1];
405+
if (strides[2] == 0)
406+
{
407+
return data[offset];
408+
}
409+
if (strides[2] == 1)
410+
{
411+
return data[offset + static_cast<std::ptrdiff_t>(k)];
412+
}
413+
return data[offset + static_cast<std::ptrdiff_t>(k) * strides[2]];
414+
}
415+
};
416+
417+
template <class T>
418+
struct runtime_scalar_plan
419+
{
420+
T value;
421+
422+
T operator()(std::size_t, std::size_t, std::size_t) const
423+
{
424+
return value;
425+
}
426+
};
427+
428+
template <class F, class P>
429+
struct runtime_function_plan
430+
{
431+
F functor;
432+
P arguments;
433+
434+
auto operator()(std::size_t i, std::size_t j, std::size_t k) const
435+
{
436+
return std::apply(
437+
[&](const auto&... argument) { return functor(argument(i, j, k)...); },
438+
arguments
439+
);
440+
}
441+
};
442+
443+
template <class E>
444+
struct runtime_plan_builder
445+
{
446+
static auto make(const E& expression)
447+
{
448+
std::array<std::ptrdiff_t, 3> strides = {0, 0, 0};
449+
const std::size_t offset = 3 - expression.dimension();
450+
for (std::size_t i = 0; i < expression.dimension(); ++i)
451+
{
452+
if (expression.shape()[i] != 1)
453+
{
454+
strides[offset + i] = expression.strides()[i];
455+
}
456+
}
457+
return runtime_leaf_plan<E>{expression.storage().data() + expression.data_offset(), strides};
458+
}
459+
};
460+
461+
template <class T>
462+
struct runtime_plan_builder<xscalar<T>>
463+
{
464+
static auto make(const xscalar<T>& expression)
465+
{
466+
return runtime_scalar_plan<T>{expression()};
467+
}
468+
};
469+
470+
template <class F, class... CT>
471+
struct runtime_plan_builder<xfunction<F, CT...>>
472+
{
473+
using expression_type = xfunction<F, CT...>;
474+
475+
template <std::size_t... I>
476+
static auto make(const expression_type& expression, std::index_sequence<I...>)
477+
{
478+
auto arguments = std::make_tuple(
479+
runtime_plan_builder<std::decay_t<decltype(std::get<I>(expression.arguments()))>>::make(
480+
std::get<I>(expression.arguments())
481+
)...
482+
);
483+
return runtime_function_plan<std::decay_t<F>, decltype(arguments)>{
484+
expression.functor(),
485+
std::move(arguments)
486+
};
487+
}
488+
489+
static auto make(const expression_type& expression)
490+
{
491+
return make(expression, std::index_sequence_for<CT...>{});
492+
}
493+
};
494+
495+
template <class E1, class E2>
496+
bool try_runtime_plan(E1& dst, const E2& expression, bool trivial)
497+
{
498+
if (trivial)
499+
{
500+
return false;
501+
}
502+
if constexpr (
503+
runtime_plan_output<E1>::value && runtime_plan_expression<E2>::value)
504+
{
505+
if (dst.dimension() == 3 && expression.dimension() <= 3)
506+
{
507+
auto plan = runtime_plan_builder<E2>::make(expression);
508+
auto* output = dst.storage().data() + dst.data_offset();
509+
std::size_t n = 0;
510+
for (std::size_t i = 0; i < dst.shape()[0]; ++i)
511+
{
512+
for (std::size_t j = 0; j < dst.shape()[1]; ++j)
513+
{
514+
for (std::size_t k = 0; k < dst.shape()[2]; ++k)
515+
{
516+
output[n++] = plan(i, j, k);
517+
}
518+
}
519+
}
520+
return true;
521+
}
522+
}
523+
return false;
524+
}
525+
337526
template <class E>
338527
struct supports_row_major_gather : std::bool_constant<E::static_layout != layout_type::column_major>
339528
{
@@ -474,6 +663,11 @@ namespace xt
474663
const E2& de2 = e2.derived_cast();
475664
using traits = xassign_traits<E1, E2>;
476665

666+
if (try_permutation_assign(de1, de2) || detail::try_runtime_plan(de1, de2, trivial))
667+
{
668+
return;
669+
}
670+
477671
bool linear_assign = traits::linear_assign(de1, de2, trivial);
478672
constexpr bool simd_assign = traits::simd_assign();
479673
constexpr bool simd_linear_assign = traits::simd_linear_assign();

‎include/xtensor/core/xoperation.hpp‎

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
#define XTENSOR_OPERATION_HPP
1212

1313
#include <algorithm>
14+
#include <array>
1415
#include <type_traits>
1516

1617
#include <xtl/xsequence.hpp>
@@ -148,6 +149,7 @@ namespace xt
148149
struct functor
149150
{
150151
using result_type = R;
152+
using cast_result_type = R;
151153

152154
template <class A1>
153155
constexpr result_type operator()(const A1& arg) const
@@ -167,6 +169,117 @@ namespace xt
167169
};
168170
};
169171

172+
template <class E, class = void>
173+
struct cast_expression : std::false_type
174+
{
175+
};
176+
177+
template <class F, class CT>
178+
struct cast_expression<xfunction<F, CT>, std::void_t<typename F::cast_result_type>> : std::true_type
179+
{
180+
using result_type = typename F::cast_result_type;
181+
};
182+
183+
template <class E>
184+
struct permutation_view : std::false_type
185+
{
186+
};
187+
188+
template <class CT, class S, layout_type L, class FST>
189+
struct permutation_view<xstrided_view<CT, S, L, FST>>
190+
: std::bool_constant<
191+
L == layout_type::dynamic && is_specialization_of<inner_storage_getter, FST>::value>
192+
{
193+
};
194+
195+
template <std::size_t A, std::size_t B, std::size_t C, class E1, class V, class R, class S>
196+
inline void run_permutation_assign(E1& dst, const V& src, S scale)
197+
{
198+
const std::size_t shape_a = src.shape()[A];
199+
const std::size_t shape_b = src.shape()[B];
200+
const std::size_t shape_c = src.shape()[C];
201+
const std::size_t src_stride_a = static_cast<std::size_t>(src.strides()[A]);
202+
const std::size_t src_stride_b = static_cast<std::size_t>(src.strides()[B]);
203+
const std::size_t src_stride_c = static_cast<std::size_t>(src.strides()[C]);
204+
const std::size_t dst_stride_a = static_cast<std::size_t>(dst.strides()[A]);
205+
const std::size_t dst_stride_b = static_cast<std::size_t>(dst.strides()[B]);
206+
const std::size_t dst_stride_c = static_cast<std::size_t>(dst.strides()[C]);
207+
auto* src_data = src.data() + src.data_offset();
208+
auto* dst_data = dst.data() + dst.data_offset();
209+
if (
210+
shape_c == 3 && src_stride_c == 1 && src_stride_b == 3 && src_stride_a == shape_b * 3
211+
&& dst_stride_b == 1 && dst_stride_a == shape_b && dst_stride_c == shape_a * shape_b)
212+
{
213+
const std::size_t plane_size = shape_a * shape_b;
214+
auto* dst_0 = dst_data;
215+
auto* dst_1 = dst_data + plane_size;
216+
auto* dst_2 = dst_data + 2 * plane_size;
217+
for (std::size_t i = 0; i < plane_size; ++i)
218+
{
219+
dst_0[i] = static_cast<R>(src_data[3 * i]) / static_cast<R>(scale);
220+
dst_1[i] = static_cast<R>(src_data[3 * i + 1]) / static_cast<R>(scale);
221+
dst_2[i] = static_cast<R>(src_data[3 * i + 2]) / static_cast<R>(scale);
222+
}
223+
return;
224+
}
225+
constexpr std::size_t tile_size = 32;
226+
for (std::size_t a0 = 0; a0 < shape_a; a0 += tile_size)
227+
{
228+
for (std::size_t b0 = 0; b0 < shape_b; b0 += tile_size)
229+
{
230+
const std::size_t a_end = std::min(a0 + tile_size, shape_a);
231+
const std::size_t b_end = std::min(b0 + tile_size, shape_b);
232+
for (std::size_t a = a0; a < a_end; ++a)
233+
{
234+
for (std::size_t b = b0; b < b_end; ++b)
235+
{
236+
for (std::size_t c = 0; c < shape_c; ++c)
237+
{
238+
const auto src_offset = a * src_stride_a + b * src_stride_b + c * src_stride_c;
239+
const auto dst_offset = a * dst_stride_a + b * dst_stride_b + c * dst_stride_c;
240+
dst_data[dst_offset] = static_cast<R>(src_data[src_offset]) / static_cast<R>(scale);
241+
}
242+
}
243+
}
244+
}
245+
}
246+
}
247+
248+
template <class E1, class V, class R, class S>
249+
inline bool try_3d_permutation_assign(E1& dst, const V& src, S scale)
250+
{
251+
if (dst.dimension() != 3 || src.dimension() != 3 || dst.shape() != src.shape()
252+
|| !dst.is_contiguous() || dst.layout() != layout_type::row_major)
253+
{
254+
return false;
255+
}
256+
257+
std::array<std::size_t, 3> axes = {0, 1, 2};
258+
std::sort(axes.begin(), axes.end(), [&](auto lhs, auto rhs) { return src.strides()[lhs] > src.strides()[rhs]; });
259+
if (src.strides()[axes[2]] != 1
260+
|| static_cast<std::size_t>(src.strides()[axes[1]]) != src.shape()[axes[2]]
261+
|| static_cast<std::size_t>(src.strides()[axes[0]])
262+
!= src.shape()[axes[2]] * src.shape()[axes[1]])
263+
{
264+
return false;
265+
}
266+
267+
#define XTENSOR_RUN_PERMUTATION(A, B, C) \
268+
if (axes == std::array<std::size_t, 3>{A, B, C}) \
269+
{ \
270+
run_permutation_assign<A, B, C, E1, V, R>(dst, src, scale); \
271+
return true; \
272+
}
273+
XTENSOR_RUN_PERMUTATION(0, 1, 2)
274+
XTENSOR_RUN_PERMUTATION(0, 2, 1)
275+
XTENSOR_RUN_PERMUTATION(1, 0, 2)
276+
XTENSOR_RUN_PERMUTATION(1, 2, 0)
277+
XTENSOR_RUN_PERMUTATION(2, 0, 1)
278+
XTENSOR_RUN_PERMUTATION(2, 1, 0)
279+
#undef XTENSOR_RUN_PERMUTATION
280+
return false;
281+
}
282+
170283
template <class Tag, class F, class... E>
171284
struct select_xfunction_expression;
172285

@@ -211,6 +324,31 @@ namespace xt
211324
enable_if_t<has_xexpression<std::decay_t<E>...>::value, xfunction_type<F, E...>>::type;
212325
}
213326

327+
template <class E1, class CT1, class CT2>
328+
inline bool try_permutation_assign(E1& dst, const xfunction<detail::divides, CT1, CT2>& expression)
329+
{
330+
using cast_expression_type = std::decay_t<decltype(std::get<0>(expression.arguments()))>;
331+
if constexpr (detail::cast_expression<cast_expression_type>::value && is_xscalar<std::decay_t<CT2>>::value)
332+
{
333+
const auto& cast_expression = std::get<0>(expression.arguments());
334+
const auto& view = std::get<0>(cast_expression.arguments());
335+
using view_type = std::decay_t<decltype(view)>;
336+
using cast_result_type = typename detail::cast_expression<cast_expression_type>::result_type;
337+
if constexpr (
338+
detail::permutation_view<view_type>::value
339+
&& std::is_same<cast_result_type, typename E1::value_type>::value
340+
&& std::is_floating_point<cast_result_type>::value)
341+
{
342+
return detail::try_3d_permutation_assign<E1, view_type, cast_result_type>(
343+
dst,
344+
view,
345+
std::get<1>(expression.arguments())()
346+
);
347+
}
348+
}
349+
return false;
350+
}
351+
214352
#undef UNARY_OPERATOR_FUNCTOR
215353
#undef BINARY_OPERATOR_FUNCTOR
216354

0 commit comments

Comments
 (0)