Skip to content

Commit

Permalink
Introduce cstride_slice<N>, constant stride slice
Browse files Browse the repository at this point in the history 
It's very rare to have non-constant stride, and using a constant unlocks
various optimization, including potential improvement for #2095
  • Loading branch information
@serge-sans-paille
serge-sans-paille committed Apr 15, 2023
1 parent a4b6e22 commit 5edc917
Show file tree
Hide file tree
Showing 8 changed files with 531 additions and 81 deletions.
17 changes: 12 additions & 5 deletions pythran/backend.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1088,12 +1088,19 @@ def visit_Slice(self, node):
else 'pythonic::builtins::None')
args.append(arg)

if node.step is None or (isnum(node.step) and node.step.value == 1):
if self.all_positive(node.lower) and self.all_positive(node.upper):
builder = "pythonic::types::fast_contiguous_slice({},{})"
nstep = node.step
if nstep is None or (isnum(nstep) and nstep.value > 0):
if nstep is None or nstep.value == 1:
if self.all_positive(node.lower) and self.all_positive(node.upper):
builder = "pythonic::types::fast_contiguous_slice({0},{1})"
else:
builder = "pythonic::types::contiguous_slice({0},{1})"
step = 1
else:
builder = "pythonic::types::contiguous_slice({},{})"
return builder.format(args[0], args[1])
builder = "pythonic::types::cstride_slice<{2}>({0},{1})"
step = nstep.value

return builder.format(args[0], args[1], step)
else:
return "pythonic::types::slice({},{},{})".format(*args)

Expand Down
73 changes: 49 additions & 24 deletions pythran/pythonic/include/types/numpy_gexpr.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,11 @@ namespace types
static constexpr size_t value = 0;
};

template <long stride>
struct count_long<cstride_normalized_slice<stride>> {
static constexpr size_t value = 0;
};

template <class T, class... Types>
struct count_long<T, Types...> {
static constexpr size_t value =
Expand Down Expand Up @@ -511,6 +516,13 @@ namespace types
: gexpr_shape<pshape<Tys..., std::integral_constant<long, 1>>,
pshape<oTys...>, S...> {
};
template <class... Tys, class... oTys, class... S, long stride>
struct gexpr_shape<pshape<Tys...>,
pshape<std::integral_constant<long, 1>, oTys...>,
cstride_normalized_slice<stride>, S...>
: gexpr_shape<pshape<Tys..., std::integral_constant<long, 1>>,
pshape<oTys...>, S...> {
};
template <class... Tys, class... oTys, class... S>
struct gexpr_shape<pshape<Tys...>,
pshape<std::integral_constant<long, 1>, oTys...>,
Expand Down Expand Up @@ -546,11 +558,9 @@ namespace types
static_assert(
utils::all_of<std::is_same<S, normalize_t<S>>::value...>::value,
"all slices are normalized");
static_assert(
utils::all_of<(std::is_same<S, long>::value ||
std::is_same<S, contiguous_normalized_slice>::value ||
std::is_same<S, normalized_slice>::value)...>::value,
"all slices are valid");
static_assert(utils::all_of<(std::is_same<S, long>::value ||
is_normalized_slice<S>::value)...>::value,
"all slices are valid");
static_assert(std::decay<Arg>::type::value >= sizeof...(S),
"slicing respects array shape");

Expand All @@ -572,6 +582,11 @@ namespace types
static constexpr size_t value =
std::remove_reference<Arg>::type::value - count_long<S...>::value;

using last_arg_stride_t =
decltype(std::declval<Arg>().template strides<sizeof...(S) - 1>());
using last_slice_t =
typename std::tuple_element<sizeof...(S) - 1, std::tuple<S...>>::type;

// It is not possible to vectorize everything. We only vectorize if the
// last dimension is contiguous, which happens if
// 1. Arg is an ndarray (this is too strict)
Expand All @@ -580,18 +595,15 @@ namespace types
static const bool is_vectorizable =
std::remove_reference<Arg>::type::is_vectorizable &&
(sizeof...(S) < std::remove_reference<Arg>::type::value ||
std::is_same<contiguous_normalized_slice,
typename std::tuple_element<
sizeof...(S) - 1, std::tuple<S...>>::type>::value);
std::is_same<contiguous_normalized_slice, last_slice_t>::value);
static const bool is_flat =
std::remove_reference<Arg>::type::is_flat && value == 1 &&
utils::all_of<std::is_same<contiguous_normalized_slice, S>::value...>::value;
utils::all_of<
std::is_same<contiguous_normalized_slice, S>::value...>::value;
static const bool is_strided =
std::remove_reference<Arg>::type::is_strided ||
(((sizeof...(S) - count_long<S...>::value) == value) &&
!std::is_same<contiguous_normalized_slice,
typename std::tuple_element<
sizeof...(S) - 1, std::tuple<S...>>::type>::value);
!std::is_same<contiguous_normalized_slice, last_slice_t>::value);

using value_type =
typename std::decay<decltype(numpy_iexpr_helper<value>::get(
Expand All @@ -614,7 +626,21 @@ namespace types

shape_t _shape;
dtype *buffer;
array<long, value> _strides;

template <long stride>
static constexpr types::pshape<std::integral_constant<long, stride>>
last_stride(cstride_normalized_slice<stride>);
static constexpr types::pshape<std::integral_constant<long, 1>>
last_stride(contiguous_normalized_slice);
static constexpr types::array<long, 1> last_stride(...);

sutils::concat_t<types::array<long, value - 1>,
typename std::conditional<
sizeof...(S) == std::decay<Arg>::type::value,
decltype(last_stride(std::declval<last_slice_t>())),
types::array<long, 1>>::type>
_strides; // strides

template <size_t I>
auto shape() const -> decltype(std::get<I>(_shape))
{
Expand All @@ -636,17 +662,11 @@ namespace types
numpy_gexpr(numpy_gexpr<Argp, S...> const &other);

template <size_t J, class Slice>
typename std::enable_if<
std::is_same<Slice, normalized_slice>::value ||
std::is_same<Slice, contiguous_normalized_slice>::value,
void>::type
typename std::enable_if<is_normalized_slice<Slice>::value, void>::type
init_shape(Slice const &s, utils::int_<1>, utils::int_<J>);

template <size_t I, size_t J, class Slice>
typename std::enable_if<
std::is_same<Slice, normalized_slice>::value ||
std::is_same<Slice, contiguous_normalized_slice>::value,
void>::type
typename std::enable_if<is_normalized_slice<Slice>::value, void>::type
init_shape(Slice const &s, utils::int_<I>, utils::int_<J>);

template <size_t J>
Expand Down Expand Up @@ -877,9 +897,14 @@ namespace types

explicit operator bool() const;


dtype* data() { return buffer;}
const dtype* data() const { return buffer;}
dtype *data()
{
return buffer;
}
const dtype *data() const
{
return buffer;
}
long flat_size() const;
long size() const;
ndarray<dtype, shape_t> copy() const
Expand Down
129 changes: 126 additions & 3 deletions pythran/pythonic/include/types/slice.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,18 +49,27 @@ namespace types
};

struct slice;
template <long stride>
struct cstride_slice;
template <long stride>
struct cstride_normalized_slice;
struct contiguous_slice;
struct fast_contiguous_slice;
struct contiguous_normalized_slice;

struct normalized_slice {
long lower, upper, step;
normalized_slice();
normalized_slice() = default;
normalized_slice(long lower, long upper, long step = 1);

normalized_slice operator*(normalized_slice const &other) const;
normalized_slice operator*(contiguous_normalized_slice const &other) const;
template <long stride>
normalized_slice
operator*(cstride_normalized_slice<stride> const &other) const;
normalized_slice operator*(slice const &other) const;
template <long stride>
normalized_slice operator*(cstride_slice<stride> const &other) const;
normalized_slice operator*(contiguous_slice const &other) const;
normalized_slice operator*(fast_contiguous_slice const &other) const;

Expand All @@ -77,6 +86,9 @@ namespace types

slice operator*(slice const &other) const;

template <long stride>
slice operator*(cstride_slice<stride> const &other) const;

slice operator*(contiguous_slice const &other) const;

slice operator*(fast_contiguous_slice const &other) const;
Expand All @@ -97,6 +109,79 @@ namespace types
long get(long i) const;
};

template <long stride>
struct cstride_normalized_slice {
long lower, upper;
static constexpr long step = stride;
cstride_normalized_slice();
cstride_normalized_slice(long lower, long upper, long = 0);

normalized_slice operator*(normalized_slice const &other) const;
cstride_normalized_slice
operator*(contiguous_normalized_slice const &other) const;
template <long other_stride>
typename std::conditional<(stride < 256 && other_stride < 256),
cstride_normalized_slice<stride * other_stride>,
normalized_slice>::type
operator*(cstride_normalized_slice<other_stride> const &other) const;

normalized_slice operator*(slice const &other) const;
template <long other_stride>
typename std::conditional<(stride < 256 && other_stride < 256),
cstride_normalized_slice<stride * other_stride>,
normalized_slice>::type
operator*(cstride_slice<other_stride> const &other) const;
cstride_normalized_slice operator*(contiguous_slice const &other) const;
cstride_normalized_slice
operator*(fast_contiguous_slice const &other) const;

long size() const;
inline long get(long i) const;
};

template <long stride>
constexpr long cstride_normalized_slice<stride>::step;

template <long stride>
struct cstride_slice {
using normalized_type = cstride_normalized_slice<stride>;

bound<long> lower, upper;
static constexpr long step = stride;

cstride_slice(none<long> lower, none<long> upper);
cstride_slice();

slice operator*(slice const &other) const;

template <long other_stride>
typename std::conditional<(stride < 256 && other_stride < 256),
cstride_slice<stride * other_stride>, slice>::type
operator*(cstride_slice<other_stride> const &other) const;

cstride_slice operator*(contiguous_slice const &other) const;

cstride_slice operator*(fast_contiguous_slice const &other) const;

/*
Normalize change a[:-1] to a[:len(a)-1] to have positif index.
It also check for value bigger than len(a) to fit the size of the
container
*/
cstride_normalized_slice<stride> normalize(long max_size) const;

/*
* An assert is raised when we can't compute the size without more
* informations.
*/
long size() const;

long get(long i) const;
};

template <long stride>
constexpr long cstride_slice<stride>::step;

struct contiguous_normalized_slice {
long lower, upper;
static constexpr long step = 1;
Expand All @@ -105,11 +190,17 @@ namespace types

contiguous_normalized_slice
operator*(contiguous_normalized_slice const &other) const;
template <long stride>
cstride_normalized_slice<stride>
operator*(cstride_normalized_slice<stride> const &other) const;
contiguous_normalized_slice operator*(contiguous_slice const &other) const;
contiguous_normalized_slice
operator*(fast_contiguous_slice const &other) const;
normalized_slice operator*(normalized_slice const &other) const;
normalized_slice operator*(slice const &other) const;
template <long stride>
cstride_normalized_slice<stride>
operator*(cstride_slice<stride> const &other) const;

long size() const;

Expand Down Expand Up @@ -167,6 +258,11 @@ namespace types
using type = normalized_slice;
};

template <long stride>
struct normalized<cstride_slice<stride>> {
using type = cstride_normalized_slice<stride>;
};

template <>
struct normalized<contiguous_slice> {
using type = contiguous_normalized_slice;
Expand All @@ -189,6 +285,23 @@ namespace types
template <>
struct is_slice<slice> : std::true_type {
};
template <long stride>
struct is_slice<cstride_slice<stride>> : std::true_type {
};

template <class S>
struct is_normalized_slice : std::false_type {
};
template <>
struct is_normalized_slice<contiguous_normalized_slice> : std::true_type {
};
template <>
struct is_normalized_slice<normalized_slice> : std::true_type {
};
template <long stride>
struct is_normalized_slice<cstride_normalized_slice<stride>>
: std::true_type {
};

template <class S>
using normalize_t = typename normalized<S>::type;
Expand Down Expand Up @@ -253,7 +366,7 @@ namespace types
template <class S>
typename std::enable_if<is_slice<S>::value, std::ostream &>::type
operator<<(std::ostream &os, S const &s);
}
} // namespace types
namespace builtins
{
template <class T>
Expand All @@ -271,7 +384,7 @@ namespace builtins
{
return s.step;
}
}
} // namespace builtins
PYTHONIC_NS_END

#ifdef ENABLE_PYTHON_MODULE
Expand All @@ -295,6 +408,16 @@ struct to_python<types::contiguous_normalized_slice> {
static PyObject *convert(types::contiguous_normalized_slice const &n);
};

template <long stride>
struct to_python<types::cstride_slice<stride>> {
static PyObject *convert(types::cstride_slice<stride> const &n);
};

template <long stride>
struct to_python<types::cstride_normalized_slice<stride>> {
static PyObject *convert(types::cstride_normalized_slice<stride> const &n);
};

template <>
struct to_python<types::slice> {
static PyObject *convert(types::slice const &n);
Expand Down
Loading

0 comments on commit 5edc917

Please sign in to comment.