mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
66a2600b6a
Summary: This is a re-attempt to land the iwyu header changes, by taking the diff from [PR 100304](https://github.com/pytorch/pytorch/pull/100304), and adding the bare minimal changes to make the diff build corectly in the internal builds. X-link: https://github.com/facebookresearch/pytorch3d/pull/1541 X-link: https://github.com/fairinternal/pytorch3d/pull/44 - Re-work D45769819 to fix header inclusions in c10 Test Plan: ``` buck2 build --no-remote-cache mode/dev-nosan //caffe2/c10/... buck2 build --no-remote-cache mode/dev-nosan //deeplearning/fbgemm/fbgemm_gpu/... buck2 build mode/dev-nosan //vision/fair/pytorch3d/pytorch3d:_C ``` Reviewed By: malfet Differential Revision: D45920611 Pull Request resolved: https://github.com/pytorch/pytorch/pull/101846 Approved by: https://github.com/malfet, https://github.com/Skylion007
45 lines
1.2 KiB
C++
45 lines
1.2 KiB
C++
#pragma once
|
|
|
|
#include <c10/core/SymInt.h>
|
|
|
|
namespace c10 {
|
|
|
|
namespace detail {
|
|
// This template can only be specialized at int64_t and c10::SymInt;
|
|
// you'll get linker errors otherwise
|
|
template <typename T>
|
|
C10_API T maybe_wrap_dim_slow(T dim, T dim_post_expr, bool wrap_scalar);
|
|
} // namespace detail
|
|
|
|
template <typename T>
|
|
T _maybe_wrap_dim(T dim, T dim_post_expr, bool wrap_scalar = true) {
|
|
// Inline the fast paths
|
|
if (C10_LIKELY(dim_post_expr * -1 <= dim && dim < dim_post_expr)) {
|
|
// For SymInts, we want an explicit control flow to trigger a guard, so we
|
|
// may as well branch too.
|
|
if (dim < 0) {
|
|
return dim + dim_post_expr;
|
|
}
|
|
return dim;
|
|
}
|
|
// Check edge-cases out-of-line (wrapping scalars and out-of-bounds errors)
|
|
return c10::detail::maybe_wrap_dim_slow<T>(
|
|
std::move(dim), std::move(dim_post_expr), wrap_scalar);
|
|
}
|
|
|
|
inline int64_t maybe_wrap_dim(
|
|
int64_t dim,
|
|
int64_t dim_post_expr,
|
|
bool wrap_scalar = true) {
|
|
return _maybe_wrap_dim(dim, dim_post_expr, wrap_scalar);
|
|
}
|
|
|
|
inline c10::SymInt maybe_wrap_dim(
|
|
c10::SymInt dim,
|
|
c10::SymInt dim_post_expr,
|
|
bool wrap_scalar = true) {
|
|
return _maybe_wrap_dim(std::move(dim), std::move(dim_post_expr), wrap_scalar);
|
|
}
|
|
|
|
} // namespace c10
|