File size: 673 Bytes
9dd3461 |
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 |
#pragma once
#include <c10/util/Exception.h>
namespace c10 {
namespace detail {
C10_API int64_t
maybe_wrap_dim_slow(int64_t dim, int64_t dim_post_expr, bool wrap_scalar);
}
static inline int64_t maybe_wrap_dim(
int64_t dim,
int64_t dim_post_expr,
bool wrap_scalar = true) {
// Inline the fast paths
if (C10_LIKELY(-dim_post_expr <= dim && dim < dim_post_expr)) {
// Branch-less version of dim + (dim < 0 ? dim_post_expr : 0)
return dim + dim_post_expr * (dim < 0);
}
// Check edge-cases out-of-line (wrapping scalars and out-of-bounds errors)
return c10::detail::maybe_wrap_dim_slow(dim, dim_post_expr, wrap_scalar);
}
} // namespace c10
|