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