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
|
#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
|
