mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
622947afa8
It's a C++17 feature that usually makes code a bit more compact, and should have no side-effects otherwise. Pull Request resolved: https://github.com/pytorch/pytorch/pull/115938 Approved by: https://github.com/Skylion007
35 lines
836 B
C++
35 lines
836 B
C++
#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
|
|
#include <ATen/native/IndexingUtils.h>
|
|
|
|
namespace at::native {
|
|
|
|
bool canUse32BitIndexMath(const TensorBase& t, int64_t max_elem) {
|
|
auto elements = t.sym_numel();
|
|
if (elements >= max_elem) {
|
|
return false;
|
|
}
|
|
if (elements == 0) {
|
|
return max_elem > 0;
|
|
}
|
|
|
|
c10::SymInt offset = 0;
|
|
auto linearId = elements - 1;
|
|
|
|
// NOTE: Assumes all strides are positive, which is true for now
|
|
// NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
|
|
for (int i = t.dim() - 1; i >= 0; --i) {
|
|
auto curDimIndex = linearId % t.sym_size(i);
|
|
auto curDimOffset = curDimIndex * t.sym_stride(i);
|
|
offset += curDimOffset;
|
|
linearId /= t.sym_size(i);
|
|
}
|
|
|
|
if (offset >= max_elem) {
|
|
return false;
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
} // namespace at::native
|