1 #pragma once
2
3 #include <c10/core/SymInt.h>
4 #include <c10/util/ArrayRef.h>
5 #include <c10/util/DimVector.h>
6 #include <c10/util/Exception.h>
7 #include <c10/util/irange.h>
8 #include <cstdint>
9 #include <optional>
10
11 namespace c10 {
12 using SymIntArrayRef = ArrayRef<SymInt>;
13
asIntArrayRefUnchecked(c10::SymIntArrayRef ar)14 inline at::IntArrayRef asIntArrayRefUnchecked(c10::SymIntArrayRef ar) {
15 return IntArrayRef(reinterpret_cast<const int64_t*>(ar.data()), ar.size());
16 }
17
18 // TODO: a SymIntArrayRef containing a heap allocated large negative integer
19 // can actually technically be converted to an IntArrayRef... but not with
20 // the non-owning API we have here. We can't reinterpet cast; we have to
21 // allocate another buffer and write the integers into it. If you need it,
22 // we can do it. But I don't think you need it.
23
asIntArrayRefSlowOpt(c10::SymIntArrayRef ar)24 inline std::optional<at::IntArrayRef> asIntArrayRefSlowOpt(
25 c10::SymIntArrayRef ar) {
26 for (const c10::SymInt& sci : ar) {
27 if (sci.is_heap_allocated()) {
28 return std::nullopt;
29 }
30 }
31
32 return {asIntArrayRefUnchecked(ar)};
33 }
34
asIntArrayRefSlow(c10::SymIntArrayRef ar,const char * file,int64_t line)35 inline at::IntArrayRef asIntArrayRefSlow(
36 c10::SymIntArrayRef ar,
37 const char* file,
38 int64_t line) {
39 for (const c10::SymInt& sci : ar) {
40 TORCH_CHECK(
41 !sci.is_heap_allocated(),
42 file,
43 ":",
44 line,
45 ": SymIntArrayRef expected to contain only concrete integers");
46 }
47 return asIntArrayRefUnchecked(ar);
48 }
49
50 // Even slower than asIntArrayRefSlow, as it forces an allocation for a
51 // destination int, BUT it is able to force specialization (it never errors)
asIntArrayRefSlowAlloc(c10::SymIntArrayRef ar,const char * file,int64_t line)52 inline c10::DimVector asIntArrayRefSlowAlloc(
53 c10::SymIntArrayRef ar,
54 const char* file,
55 int64_t line) {
56 c10::DimVector res(ar.size(), 0);
57 for (const auto i : c10::irange(ar.size())) {
58 res[i] = ar[i].guard_int(file, line);
59 }
60 return res;
61 }
62
63 #define C10_AS_INTARRAYREF_SLOW(a) c10::asIntArrayRefSlow(a, __FILE__, __LINE__)
64 #define C10_AS_INTARRAYREF_SLOW_ALLOC(a) \
65 c10::asIntArrayRefSlowAlloc(a, __FILE__, __LINE__)
66
67 // Prefer using a more semantic constructor, like
68 // fromIntArrayRefKnownNonNegative
fromIntArrayRefUnchecked(IntArrayRef array_ref)69 inline SymIntArrayRef fromIntArrayRefUnchecked(IntArrayRef array_ref) {
70 return SymIntArrayRef(
71 reinterpret_cast<const SymInt*>(array_ref.data()), array_ref.size());
72 }
73
fromIntArrayRefKnownNonNegative(IntArrayRef array_ref)74 inline SymIntArrayRef fromIntArrayRefKnownNonNegative(IntArrayRef array_ref) {
75 return fromIntArrayRefUnchecked(array_ref);
76 }
77
fromIntArrayRefSlow(IntArrayRef array_ref)78 inline SymIntArrayRef fromIntArrayRefSlow(IntArrayRef array_ref) {
79 for (long i : array_ref) {
80 TORCH_CHECK(
81 SymInt::check_range(i),
82 "IntArrayRef contains an int that cannot be represented as a SymInt: ",
83 i);
84 }
85 return SymIntArrayRef(
86 reinterpret_cast<const SymInt*>(array_ref.data()), array_ref.size());
87 }
88
89 } // namespace c10
90