frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Inline asIntArrayRef (#86350)

Browse Source 
I was benchmarking and this is worth maybe 5% on at::empty, but it's basically
free so we should do it.

Signed-off-by: Edward Z. Yang <ezyang@fb.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/86350
Approved by: https://github.com/albanD
This commit is contained in:
Edward Z. Yang
2022-10-05 21:42:02 -07:00
committed by PyTorch MergeBot
parent cebf08afb2
commit 4d7728890b
2 changed files with 24 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
c10/core/SymIntArrayRef.cpp
View File

@ -1,29 +1,3 @@
#include <c10/core/SymIntArrayRef.h>
#include <c10/util/Optional.h>
#include <iostream>
namespace c10 {
at::IntArrayRef asIntArrayRefSlow(c10::SymIntArrayRef ar) {
auto r = asIntArrayRefSlowOpt(ar);
TORCH_CHECK(
r.has_value(),
"SymIntArrayRef expected to contain only concrete integers");
return *r;
}
c10::optional<at::IntArrayRef> asIntArrayRefSlowOpt(c10::SymIntArrayRef ar) {
for (c10::SymInt sci : ar) {
if (sci.is_symbolic()) {
return c10::nullopt;
}
}
return {asIntArrayRefUnchecked(ar)};
}
at::IntArrayRef asIntArrayRefUnchecked(c10::SymIntArrayRef ar) {
return IntArrayRef(reinterpret_cast<const int64_t*>(ar.data()), ar.size());
}
} // namespace c10
namespace c10 {} // namespace c10

27
c10/core/SymIntArrayRef.h
View File

@ -13,10 +13,29 @@
namespace c10 {
using SymIntArrayRef = ArrayRef<SymInt>;
TORCH_API at::IntArrayRef asIntArrayRefSlow(c10::SymIntArrayRef ar);
TORCH_API at::IntArrayRef asIntArrayRefUnchecked(c10::SymIntArrayRef ar);
TORCH_API c10::optional<at::IntArrayRef> asIntArrayRefSlowOpt(
c10::SymIntArrayRef ar);
inline at::IntArrayRef asIntArrayRefUnchecked(c10::SymIntArrayRef ar) {
return IntArrayRef(reinterpret_cast<const int64_t*>(ar.data()), ar.size());
}
inline c10::optional<at::IntArrayRef> asIntArrayRefSlowOpt(
c10::SymIntArrayRef ar) {
for (c10::SymInt sci : ar) {
if (sci.is_symbolic()) {
return c10::nullopt;
}
}
return {asIntArrayRefUnchecked(ar)};
}
inline at::IntArrayRef asIntArrayRefSlow(c10::SymIntArrayRef ar) {
for (c10::SymInt sci : ar) {
TORCH_CHECK(
!sci.is_symbolic(),
"SymIntArrayRef expected to contain only concrete integers");
}
return asIntArrayRefUnchecked(ar);
}
// Prefer using a more semantic constructor, like
// fromIntArrayRefKnownNonNegative