1 | #pragma once |
---|---|
2 | |
3 | #include <torch/csrc/Exceptions.h> |
4 | #include <torch/csrc/jit/frontend/tracer.h> |
5 | #include <torch/csrc/python_headers.h> |
6 | #include <torch/csrc/utils/object_ptr.h> |
7 | #include <torch/csrc/utils/tensor_numpy.h> |
8 | #include <cstdint> |
9 | #include <limits> |
10 | #include <stdexcept> |
11 | |
12 | // largest integer that can be represented consecutively in a double |
13 | const int64_t DOUBLE_INT_MAX = 9007199254740992; |
14 | |
15 | inline PyObject* THPUtils_packInt32(int32_t value) { |
16 | return PyLong_FromLong(value); |
17 | } |
18 | |
19 | inline PyObject* THPUtils_packInt64(int64_t value) { |
20 | return PyLong_FromLongLong(value); |
21 | } |
22 | |
23 | inline PyObject* THPUtils_packUInt32(uint32_t value) { |
24 | return PyLong_FromUnsignedLong(value); |
25 | } |
26 | |
27 | inline PyObject* THPUtils_packUInt64(uint64_t value) { |
28 | return PyLong_FromUnsignedLongLong(value); |
29 | } |
30 | |
31 | inline PyObject* THPUtils_packDoubleAsInt(double value) { |
32 | return PyLong_FromDouble(value); |
33 | } |
34 | |
35 | inline bool THPUtils_checkLong(PyObject* obj) { |
36 | #ifdef USE_NUMPY |
37 | if (torch::utils::is_numpy_int(obj)) { |
38 | return true; |
39 | } |
40 | #endif |
41 | |
42 | return PyLong_Check(obj) && !PyBool_Check(obj); |
43 | } |
44 | |
45 | inline int32_t THPUtils_unpackInt(PyObject* obj) { |
46 | // NOLINTNEXTLINE(cppcoreguidelines-init-variables) |
47 | int overflow; |
48 | long value = PyLong_AsLongAndOverflow(obj, &overflow); |
49 | if (value == -1 && PyErr_Occurred()) { |
50 | throw python_error(); |
51 | } |
52 | if (overflow != 0) { |
53 | throw std::runtime_error("Overflow when unpacking long"); |
54 | } |
55 | if (value > std::numeric_limits<int32_t>::max() || |
56 | value < std::numeric_limits<int32_t>::min()) { |
57 | throw std::runtime_error("Overflow when unpacking long"); |
58 | } |
59 | return (int32_t)value; |
60 | } |
61 | |
62 | inline int64_t THPUtils_unpackLong(PyObject* obj) { |
63 | // NOLINTNEXTLINE(cppcoreguidelines-init-variables) |
64 | int overflow; |
65 | long long value = PyLong_AsLongLongAndOverflow(obj, &overflow); |
66 | if (value == -1 && PyErr_Occurred()) { |
67 | throw python_error(); |
68 | } |
69 | if (overflow != 0) { |
70 | throw std::runtime_error("Overflow when unpacking long"); |
71 | } |
72 | return (int64_t)value; |
73 | } |
74 | |
75 | inline uint32_t THPUtils_unpackUInt32(PyObject* obj) { |
76 | unsigned long value = PyLong_AsUnsignedLong(obj); |
77 | if (PyErr_Occurred()) { |
78 | throw python_error(); |
79 | } |
80 | if (value > std::numeric_limits<uint32_t>::max()) { |
81 | throw std::runtime_error("Overflow when unpacking unsigned long"); |
82 | } |
83 | return (uint32_t)value; |
84 | } |
85 | |
86 | inline uint64_t THPUtils_unpackUInt64(PyObject* obj) { |
87 | unsigned long long value = PyLong_AsUnsignedLongLong(obj); |
88 | if (PyErr_Occurred()) { |
89 | throw python_error(); |
90 | } |
91 | return (uint64_t)value; |
92 | } |
93 | |
94 | inline bool THPUtils_checkIndex(PyObject* obj) { |
95 | if (PyBool_Check(obj)) { |
96 | return false; |
97 | } |
98 | if (THPUtils_checkLong(obj)) { |
99 | return true; |
100 | } |
101 | torch::jit::tracer::NoWarn no_warn_guard; |
102 | auto index = THPObjectPtr(PyNumber_Index(obj)); |
103 | if (!index) { |
104 | PyErr_Clear(); |
105 | return false; |
106 | } |
107 | return true; |
108 | } |
109 | |
110 | inline int64_t THPUtils_unpackIndex(PyObject* obj) { |
111 | if (!THPUtils_checkLong(obj)) { |
112 | auto index = THPObjectPtr(PyNumber_Index(obj)); |
113 | if (index == nullptr) { |
114 | throw python_error(); |
115 | } |
116 | // NB: This needs to be called before `index` goes out of scope and the |
117 | // underlying object's refcount is decremented |
118 | return THPUtils_unpackLong(index.get()); |
119 | } |
120 | return THPUtils_unpackLong(obj); |
121 | } |
122 | |
123 | inline bool THPUtils_unpackBool(PyObject* obj) { |
124 | if (obj == Py_True) { |
125 | return true; |
126 | } else if (obj == Py_False) { |
127 | return false; |
128 | } else { |
129 | throw std::runtime_error("couldn't convert python object to boolean"); |
130 | } |
131 | } |
132 | |
133 | inline bool THPUtils_checkDouble(PyObject* obj) { |
134 | #ifdef USE_NUMPY |
135 | if (torch::utils::is_numpy_scalar(obj)) { |
136 | return true; |
137 | } |
138 | #endif |
139 | return PyFloat_Check(obj) || PyLong_Check(obj); |
140 | } |
141 | |
142 | inline double THPUtils_unpackDouble(PyObject* obj) { |
143 | if (PyFloat_Check(obj)) { |
144 | return PyFloat_AS_DOUBLE(obj); |
145 | } |
146 | double value = PyFloat_AsDouble(obj); |
147 | if (value == -1 && PyErr_Occurred()) { |
148 | throw python_error(); |
149 | } |
150 | return value; |
151 | } |
152 | |
153 | inline c10::complex<double> THPUtils_unpackComplexDouble(PyObject* obj) { |
154 | Py_complex value = PyComplex_AsCComplex(obj); |
155 | if (value.real == -1.0 && PyErr_Occurred()) { |
156 | throw python_error(); |
157 | } |
158 | |
159 | return c10::complex<double>(value.real, value.imag); |
160 | } |
161 | |
162 | inline bool THPUtils_unpackNumberAsBool(PyObject* obj) { |
163 | if (PyFloat_Check(obj)) { |
164 | return (bool)PyFloat_AS_DOUBLE(obj); |
165 | } |
166 | |
167 | if (PyComplex_Check(obj)) { |
168 | double real_val = PyComplex_RealAsDouble(obj); |
169 | double imag_val = PyComplex_ImagAsDouble(obj); |
170 | return !(real_val == 0 && imag_val == 0); |
171 | } |
172 | |
173 | // NOLINTNEXTLINE(cppcoreguidelines-init-variables) |
174 | int overflow; |
175 | long long value = PyLong_AsLongLongAndOverflow(obj, &overflow); |
176 | if (value == -1 && PyErr_Occurred()) { |
177 | throw python_error(); |
178 | } |
179 | // No need to check overflow, because when overflow occured, it should |
180 | // return true in order to keep the same behavior of numpy. |
181 | return (bool)value; |
182 | } |
183 |