forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
python_variable.cpp
2398 lines (2193 loc) · 85.8 KB
/
python_variable.cpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
#include <ATen/NamedTensorUtils.h>
#include <c10/core/DeviceType.h>
#include <c10/core/impl/GPUTrace.h>
#include <c10/core/impl/HermeticPyObjectTLS.h>
#include <c10/core/impl/PythonDispatcherTLS.h>
#include <c10/util/irange.h>
#include <pybind11/pytypes.h>
#include <torch/csrc/Device.h>
#include <torch/csrc/DynamicTypes.h>
#include <torch/csrc/Exceptions.h>
#include <torch/csrc/PyInterpreter.h>
#include <torch/csrc/Size.h>
#include <torch/csrc/THP.h>
#include <torch/csrc/Types.h>
#include <torch/csrc/autograd/autograd.h>
#include <torch/csrc/autograd/edge.h>
#include <torch/csrc/autograd/function.h>
#include <torch/csrc/autograd/python_cpp_function.h>
#include <torch/csrc/autograd/python_hook.h>
#include <torch/csrc/autograd/python_variable_indexing.h>
#include <torch/csrc/autograd/utils/error_messages.h>
#include <torch/csrc/autograd/utils/wrap_outputs.h>
#include <torch/csrc/autograd/variable.h>
#include <torch/csrc/jit/frontend/tracer.h>
#include <torch/csrc/jit/python/pybind_utils.h>
#include <torch/csrc/tensor/python_tensor.h>
#include <torch/csrc/utils/pybind.h>
#include <torch/csrc/utils/pycfunction_helpers.h>
#include <torch/csrc/utils/pyobject_preservation.h>
#include <torch/csrc/utils/python_arg_parser.h>
#include <torch/csrc/utils/python_dispatch.h>
#include <torch/csrc/utils/python_strings.h>
#include <torch/csrc/utils/tensor_new.h>
#include <torch/csrc/utils/tensor_numpy.h>
#include <torch/csrc/utils/torch_dispatch_mode.h>
#include <ATen/ATen.h>
#include <c10/core/SymIntArrayRef.h>
#include <structmember.h>
#include <cstdint>
#include <memory>
#include <utility>
#include <vector>
using namespace at;
using namespace torch;
using namespace torch::autograd;
std::pair<py::object, py::dict> parseIValuesToPyArgsKwargs(
const c10::OperatorHandle& op,
const std::vector<c10::IValue>& arguments) {
TORCH_CHECK(
PyGILState_Check(),
"GIL must be held before you call parseIValuesToPyArgsKwargs");
const auto& schema = op.schema();
py::dict kwargs;
// About all the pointers:
//
// f(int x, int y = 0, *, int z = 0)
// ^- arguments.size()
// ^- kwarg_only_start
// ^- positional_default_start
// ^- 0
// Find the split point between kwarg-only and regular. Since most functions
// don't have kwarg-only arguments, it is more efficient to scan from the
// right (but ideally, this would just be precomputed in FunctionSchema
// itself). (NB: minus one in the loop is because we're testing if the
// *next* argument is kwarg-only before we advance the starting index)
int64_t kwarg_only_start = static_cast<int64_t>(arguments.size());
for (; kwarg_only_start > 0; kwarg_only_start--) {
const auto& arg = schema.arguments()[kwarg_only_start - 1];
if (!arg.kwarg_only()) {
break;
}
}
// Find the first positional argument that isn't defaulted
auto is_default = [&](size_t idx) -> bool {
const auto& arg = schema.arguments()[idx];
if (!arg.default_value().has_value()) {
return false;
}
const auto& default_ivalue = *arg.default_value();
const auto& ivalue = arguments[idx];
if (default_ivalue != ivalue) {
return false;
}
return true;
};
int64_t positional_default_start = kwarg_only_start;
for (; positional_default_start > 0; positional_default_start--) {
if (!is_default(positional_default_start - 1)) {
break;
}
}
auto args =
py::reinterpret_steal<py::object>(PyTuple_New(positional_default_start));
auto schemaAwareToPyObject = [&](size_t idx) -> py::object {
const auto& arg = schema.arguments()[idx];
auto match = [&](c10::TypeKind kind) {
const auto& t = arg.real_type();
if (t->kind() == kind)
return true;
if (auto opt_t = t->cast<c10::OptionalType>()) {
if (opt_t->getElementType()->kind() == kind)
return true;
}
return false;
};
if (arguments[idx].isNone()) {
return py::none();
} else if (match(c10::ScalarTypeType::Kind)) {
auto* obj =
getTHPDtype(static_cast<c10::ScalarType>(arguments[idx].toInt()));
return py::reinterpret_borrow<py::object>(
reinterpret_cast<PyObject*>(obj));
} else if (match(c10::LayoutType::Kind)) {
auto* obj =
getTHPLayout(static_cast<c10::Layout>(arguments[idx].toInt()));
return py::reinterpret_borrow<py::object>(
reinterpret_cast<PyObject*>(obj));
} else if (match(c10::MemoryFormatType::Kind)) {
return py::cast(static_cast<c10::MemoryFormat>(arguments[idx].toInt()));
} else {
return torch::jit::toPyObject(arguments[idx]);
}
};
// Populate positional arguments
for (const auto idx : c10::irange(positional_default_start)) {
PyTuple_SET_ITEM(
args.ptr(), idx, schemaAwareToPyObject(idx).release().ptr());
}
// Populate keyword arguments
for (const auto idx : c10::irange(kwarg_only_start, arguments.size())) {
// But don't populate default keyword arguments
if (is_default(idx))
continue;
const auto& arg = schema.arguments()[idx];
kwargs[py::cast(arg.name())] = schemaAwareToPyObject(idx);
}
return std::make_pair(std::move(args), std::move(kwargs));
}
void pushPyOutToStack(
const c10::OperatorHandle& op,
torch::jit::Stack* stack,
py::object out,
const char* msg) {
TORCH_CHECK(
PyGILState_Check(), "GIL must be held before you call pushPyOutToStack");
auto schema_returns = op.schema().returns();
const auto num_returns = schema_returns.size();
if (num_returns == 0) {
// Check that we got a None return from Python. Anything else is an error.
TORCH_CHECK(
out.is_none(),
"Expected ",
msg,
" for ",
op.operator_name(),
" to return None but it returned something else instead.");
} else if (num_returns == 1) {
torch::jit::push(
stack, torch::jit::toIValue(out.ptr(), schema_returns[0].real_type()));
} else {
auto outs = py::cast<py::sequence>(out);
for (const auto idx : c10::irange(outs.size())) {
torch::jit::push(
stack,
torch::jit::toIValue(
outs[idx].ptr(), schema_returns[idx].real_type()));
}
}
}
namespace {
c10::TensorImpl::SizesStridesPolicy parseSizesStridesPolicyArgument(
c10::string_view arg) {
if (arg == "strides") {
return c10::TensorImpl::SizesStridesPolicy::CustomStrides;
}
if (arg == "sizes") {
return c10::TensorImpl::SizesStridesPolicy::CustomSizes;
}
TORCH_CHECK_VALUE(
false,
"Unknown sizes_strides_policy: ",
arg,
"; expected 'strides' or 'sizes'");
}
} // anonymous namespace
PyObject* THPVariableClass = nullptr;
PyObject* ParameterClass = nullptr;
static PyObject* THPVariable_NewWithVar(
PyTypeObject* type,
Variable _var,
c10::impl::PyInterpreterStatus status,
bool allow_preexisting_pyobj = false);
// clang-tidy gets confused by static const
static const char* VOLATILE_WARNING =
"volatile was removed and now has no effect. Use "
"`with torch.no_grad():` instead.";
static bool check_has_torch_dispatch(PyObject* obj) {
PyTypeObject* tp = Py_TYPE(obj);
if (THPVariable_CheckTypeExact(tp)) {
return false;
}
py::object attr = PyObject_FastGetAttrString(obj, "__torch_dispatch__");
return (
attr.ptr() != nullptr &&
attr.ptr() != torch::disabled_torch_dispatch_impl());
}
// NOLINTNEXTLINE(*-c-arrays,cppcoreguidelines-avoid-non-const-global-variables)
static PyObject* device_to_py_class_[static_cast<size_t>(
c10::DeviceType::COMPILE_TIME_MAX_DEVICE_TYPES)];
void registerPythonTensorClass(
const std::string& device,
PyObject* python_tensor_class) {
c10::Device dev(device);
TORCH_CHECK(
dev.type() == kXLA, "Only the python class for XLA can be overriden");
if (device_to_py_class_[static_cast<size_t>(dev.type())] != nullptr) {
TORCH_WARN(
"Overriding a previously registered python class for ", dev.str());
}
device_to_py_class_[static_cast<size_t>(dev.type())] = python_tensor_class;
}
static PyObject* getPythonTensorClass(c10::Device d) {
return device_to_py_class_[static_cast<size_t>(d.type())];
}
void activateGPUTrace() {
c10::impl::GPUTrace::set_trace(getPyInterpreter());
}
// TODO: Make this take Variable by const reference
PyObject* THPVariable_Wrap(at::TensorBase var) {
if (!var.defined()) {
Py_RETURN_NONE;
}
if (c10::impl::HermeticPyObjectTLS::get_state()) {
return THPVariable_NewWithVar(
(PyTypeObject*)THPVariableClass,
std::move(var),
c10::impl::PyInterpreterStatus::DEFINITELY_UNINITIALIZED);
}
std::optional<PyObject*> mb_obj =
var.unsafeGetTensorImpl()->pyobj_slot()->check_pyobj(
getPyInterpreter(), /*ignore_hermetic_tls=*/false);
c10::impl::PyInterpreterStatus status{};
if (mb_obj.has_value()) {
auto obj = *mb_obj;
if (obj) {
if (var.unsafeGetTensorImpl()->pyobj_slot()->owns_pyobj()) {
// C++ owns the Python object; this implies there weren't any other
// owning references to the Python object. Since we're making the
// object "live" again on Python side, let's flip back the ownership
// (Python owns C++) as it would now be unsound to deallocate the C++
// object if all C++ references go to zero
var.unsafeGetTensorImpl()->pyobj_slot()->set_owns_pyobj(false);
reinterpret_cast<THPVariable*>(obj)->cdata =
MaybeOwned<Variable>::owned(std::move(var));
// NB: incref is not necessary, because we are "stealing" the previous
// ownership from the Variable to return it here for the wrap
return obj;
}
Py_INCREF(obj);
return obj;
}
// TODO: a better invariant is that if we tagged, we MUST have a valid
// PyObject. That's PyObject preservation
// (https://github.com/pytorch/pytorch/pull/56017). Prior to this PR
// being a thing, the PyObject field will get cleared when all references
// to the Python object are removed.
status = c10::impl::PyInterpreterStatus::TAGGED_BY_US;
} else {
// Assumption: if a Tensor has been shared across threads, this induces
// a refcount bump. Therefore, if the use count 1, we are the sole thread
// with access to this tensor and no race is possible.
if (var.use_count() <= 1) {
status = c10::impl::PyInterpreterStatus::DEFINITELY_UNINITIALIZED;
} else {
status = c10::impl::PyInterpreterStatus::MAYBE_UNINITIALIZED;
}
}
if (C10_LIKELY(var.device().type() != c10::kXLA)) {
return THPVariable_NewWithVar(
(PyTypeObject*)THPVariableClass, std::move(var), status);
}
if (auto clazz = getPythonTensorClass(var.device())) {
return THPVariable_NewWithVar((PyTypeObject*)clazz, std::move(var), status);
}
return THPVariable_NewWithVar(
(PyTypeObject*)THPVariableClass, std::move(var), status);
}
bool isResurrectable(THPVariable* self) {
// We want to divide this check into 2 cases.
// 1. C++ owns PyObject (in this case, self->cdata.unsafeIsBorrowed() is
// true). You might think that in this case, it is impossible for tp_clear to
// be called: surely the C++ reference to the PyObject is keeping it live? And
// you'd be right! In fact, when C++ owns the PyObject, we have an invariant
// that the refcount on the PyObject should be precisely one (because if you
// take out another reference to the PyObject, we're supposed to flip the
// ownership pointer back). In reality, you can violate this invariant
// temporarily with weak references, so we don't test for it in asserts.
// 2. PyObject owns C++ (in this case, self->cdata.unsafeIsBorrowed() is
// false). In this case, tp_clear can get called if the PyObject is referenced
// from a dead cycle, and nowhere else. But if resurrection did not occur,
// then the reference to C++ from the PyObject must be the ONLY reference to
// the C++ object.
if (self->cdata.unsafeIsBorrowed()) {
return false;
}
auto const& tensor = THPVariable_Unpack(self);
if (!tensor.defined() || tensor.use_count() <= 1) {
return false;
}
// Check if this is hermetic. If it is, no resurrection.
if (tensor.unsafeGetTensorImpl()->pyobj_slot()->check_pyobj(
getPyInterpreter(), /*ignore_hermetic_tls=*/false) !=
std::make_optional((PyObject*)self)) {
return false;
}
return true;
}
// returns true if successfully rezzed; if so, cancel the
// rest of deallocation
static bool THPVariable_tryResurrect(THPVariable* self) {
const auto& tensor = THPVariable_Unpack(self);
if (!isResurrectable(self)) {
return false;
}
// At this point, we are definitely going to resurrect the tensor. So, the
// tensor better be defined :)
TORCH_INTERNAL_ASSERT(tensor.defined());
// There are other C++ owners of the tensor. Flip ownership
// so that C++ owns this Python object, and cancel deallocation.
TORCH_INTERNAL_ASSERT(
!tensor.unsafeGetTensorImpl()->pyobj_slot()->owns_pyobj());
c10::TensorImpl* tensor_impl = tensor.unsafeGetTensorImpl();
auto maybe_pyobj = tensor_impl->pyobj_slot()->check_pyobj(
getPyInterpreter(),
/*ignore_hermetic_tls=*/false);
TORCH_INTERNAL_ASSERT(
maybe_pyobj.has_value(),
"Trying to preserve a Python tensor whose PyObjectSlot does not have a PyObject");
tensor_impl->pyobj_slot()->set_owns_pyobj(true);
// Resurrect the Python object. This is something CPython does
// internally occasionally, see
// https://github.com/python/cpython/blob/b98eba5bc2ffbe7a0ed49d540ebc4f756ae61985/Objects/object.c#L248-L259
// so we just copy the pattern here. Note that we don't have to worry
// about saving and restoring the refcount (as the quoted code does)
// because we actually DO need to reset the refcount to one here, we
// can't assume that some other code has taken care of it.
// NB: this will overreport _Py_RefTotal but based on inspection of object.c
// there is no way to avoid this
// When resurrecting, we MUST use _Py_NewReference and not Py_INCREF to
// ensure the PyObject is in a valid state
_Py_NewReference((PyObject*)self);
// Flip THPVariable to be non-owning
// (near use-after-free miss here: fresh MaybeOwned is created breaking
// reference on Tensor in struct BEFORE we overwrite the old one)
TORCH_INTERNAL_ASSERT(!c10::impl::HermeticPyObjectTLS::get_state());
self->cdata = MaybeOwned<Variable>::borrowed(tensor);
// NB: At this point, tensor *could* be dead (e.g., some other C++ thread
// decrefed it.) At this point, it is probably waiting on the GIL to
// deallocate the Python object and will kill self, BUT NOT YET.
return true;
}
static int THPVariable_subclass_clear(THPVariable* self) {
// Is it OK for an object to still be live after running
// tp_clear? Yes. When Python is breaking reference cycles, it can't assume
// that an object will dealloc after it's cleared. The source code explicitly
// handles this case:
// https://github.com/python/cpython/blob/4e661cd69164318c1f871faa476c68a04092ddc4/Modules/gcmodule.c#L1010-L1025
// Note that we don't need to actually resurrect here. There are 2 cases:
// 1. The PyObject is not part of a reference cycle. In this case, we don't
// need to do anything. The GC will move on to try and break the reference
// cycle on another object, which will eventually trigger tp_dealloc (and thus
// resurrection).
// 2. The PyObject is part of a reference cycle. This case should not actually
// be possible, due to the logic in our tp_traverse
// (THPVariable_subclass_traverse).
// In fact, resurrecting here breaks the invariant that "C++ owns Python only
// when PyObject's refcount would otherwise be 0". Most immediately, as we're
// merely breaking reference cycles here, there can be other references to the
// PyObject. *However*, if other objects in the refcycle resurrect, then we
// will be in a state where the PyObject has multiple Python references, yet
// C++ owns the PyObject.
// See https://github.com/pytorch/pytorch/pull/75933 for more discussion.
if (isResurrectable((THPVariable*)self)) {
return 0;
}
Py_CLEAR(self->backward_hooks);
Py_CLEAR(self->post_accumulate_grad_hooks);
const auto& tensor = THPVariable_Unpack(self);
if (tensor.defined()) {
// Two situations to consider:
// PyObject -owns-> Tensor
// unsafeIsBorrowed() is FALSE. We're obligated to look through
// Tensor to break references. Clearing cdata must induce the
// destruction of the C++ Tensor. If there were other references
// to C++ tensor, the Python object would have been resurrected
// by flipping the ownership.
// Tensor -owns-> PyObject
// unsafeIsBorrowed() is TRUE. We're deallocating the PyObject
// because Tensor asked us to (it's already destructing).
if (!self->cdata.unsafeIsBorrowed() &&
tensor.unsafeGetTensorImpl()->pyobj_slot()->check_pyobj(
getPyInterpreter(), /*ignore_hermetic_tls=*/false) ==
std::make_optional((PyObject*)self)) {
// TODO: empirically, on OS X this assert appears to be untrue
// In test_py_tensors_multi_async_call - ProcessGroupRpcTestWithSpawn
// distributed/rpc/test_process_group_agent.py
//
// libc++abi.dylib: terminating with uncaught exception of type
// c10::Error:
// !tensor.unsafeGetTensorImpl()->pyobj_slot()->owns_pyobj()INTERNAL
// ASSERT FAILED at "../torch/csrc/autograd/python_variable.cpp":171,
// please report a bug to PyTorch. Exception raised from
// THPVariable_subclass_clear at
// ../torch/csrc/autograd/python_variable.cpp:171 (most recent call
// first): frame #0: c10::Error::Error(c10::SourceLocation,
// std::__1::basic_string<char, std::__1::char_traits<char>,
// std::__1::allocator<char> >) + 98 (0x1158a0442 in libc10.dylib) frame
// #1: c10::detail::torchCheckFail(char const*, char const*, unsigned
// int, char const*) + 205 (0x11589ed3d in libc10.dylib) frame #2:
// c10::detail::torchInternalAssertFail(char const*, char const*,
// unsigned int, char const*, c10::detail::CompileTimeEmptyString) + 9
// (0x1141e3f89 in libtorch_python.dylib) frame #3:
// THPVariable_subclass_clear(THPVariable*) + 412 (0x1148a547c in
// libtorch_python.dylib) frame #4:
// THPVariable_subclass_dealloc(_object*) + 453 (0x1148a5035 in
// libtorch_python.dylib) frame #5: (anonymous
// namespace)::concrete_decref_fn(c10::impl::PyInterpreter const*,
// _object*) + 53 (0x1148a5ea5 in libtorch_python.dylib) frame #6:
// c10::TensorImpl::release_resources() + 182 (0x11588c4a6 in
// libc10.dylib) frame #7:
// c10::MaybeOwned<at::Tensor>::operator=(c10::MaybeOwned<at::Tensor>&&)
// + 91 (0x11488c11b in libtorch_python.dylib) frame #8:
// THPVariable_subclass_dealloc(_object*) + 607 (0x1148a50cf in
// libtorch_python.dylib) <omitting python frames> frame #47: start + 1
// (0x7fff6ffc7cc9 in libdyld.dylib) frame #48: 0x0 + 4 (0x4 in ???)
// TORCH_INTERNAL_ASSERT(!tensor.unsafeGetTensorImpl()->pyobj_slot()->owns_pyobj());
if (auto grad_acc =
torch::autograd::impl::try_get_grad_accumulator(tensor)) {
grad_acc->pre_hooks().clear();
grad_acc->tensor_pre_hooks().clear();
grad_acc->retains_grad_hooks().clear();
}
}
}
TORCH_INTERNAL_ASSERT(!isResurrectable((THPVariable*)self));
{
// MapAllocator can take significant time to release large tensors;
// release the GIL here to avoid impacting main thread perf.
pybind11::gil_scoped_release no_gil;
self->cdata = MaybeOwned<Variable>();
}
return 0;
}
int THPFake_traverse(THPVariable* self, visitproc visit, void* arg) {
TORCH_INTERNAL_ASSERT(
false, "TensorBase tp_traverse function was not overriden properly");
return 0;
}
int THPFake_clear(THPVariable* self) {
TORCH_INTERNAL_ASSERT(
false, "TensorBase tp_clear function was not overriden properly");
return 0;
}
PyObject* THPVariable_pynew(
PyTypeObject* type,
PyObject* args,
PyObject* kwargs);
static PyObject* THPVariable_fix_weakref(PyObject* self, PyObject* noargs) {
const auto& var = THPVariable_Unpack(self);
Py_DECREF(THPVariable_Wrap(var));
Py_RETURN_NONE;
}
// Maps the given python callable over a vector of items, returning a vector
// of the same type of items.
template <typename T>
static std::vector<T> map_py_func(
const py::function& func,
const std::vector<T>& items) {
std::vector<T> new_items;
new_items.reserve(items.size());
for (auto& item : items) {
new_items.emplace_back(py::cast<T>(func(item)));
}
return new_items;
}
template <>
std::vector<at::Tensor> map_py_func(
const py::function& func,
const std::vector<at::Tensor>& items) {
std::vector<at::Tensor> new_items;
new_items.reserve(items.size());
for (auto& item : items) {
auto output = func(item);
if (output.is(py::none())) {
// treat None value as an undefined tensor
new_items.emplace_back();
} else {
new_items.emplace_back(py::cast<at::Tensor>(output));
}
}
return new_items;
}
static PyObject* view_func_impl(
PyObject* _self,
PyObject* args,
PyObject* kwargs,
bool check_has_same_meta) {
HANDLE_TH_ERRORS
const auto& self = THPVariable_Unpack(_self);
static PythonArgParser parser({
"_view_func(Tensor new_base, PyObject* symint_visitor_fn=None, PyObject* tensor_visitor_fn=None)",
});
ParsedArgs<3> parsed_args{};
auto r = parser.parse(_self, args, kwargs, parsed_args);
auto new_base = r.tensor(0);
PyObject* symint_visitor_fn = r.pyobject(1);
PyObject* tensor_visitor_fn = r.pyobject(2);
// Ensure that self is indeed a backward differentiable view
// If not, we return an undefined Tensor (None) and let the user handle it.
auto diff_view_meta = torch::autograd::impl::get_view_autograd_meta(self);
at::Tensor out;
if (diff_view_meta && diff_view_meta->has_bw_view()) {
const auto& view_info = diff_view_meta->get_backward_view();
// Ensure that the newly provided base is similar to the original base
if (!check_has_same_meta ||
torch::autograd::utils::has_same_meta(new_base, view_info.base_)) {
// Do the actual view replay
if (view_info.has_view_fn()) {
auto& view_func = view_info.view_fn();
// Determine new SymInt / tensor state as needed.
std::optional<std::vector<c10::SymInt>> new_symints = std::nullopt;
if (symint_visitor_fn != Py_None) {
new_symints = map_py_func(
py::cast<py::function>(symint_visitor_fn),
view_func.get_symints());
}
std::optional<std::vector<at::Tensor>> new_tensors = std::nullopt;
if (tensor_visitor_fn != Py_None) {
new_tensors = map_py_func(
py::cast<py::function>(tensor_visitor_fn),
view_func.get_tensors());
}
// call view func
if (new_symints.has_value() || new_tensors.has_value()) {
out = (*view_func.clone_and_set(new_symints, new_tensors))(new_base);
} else {
out = view_func(new_base);
}
} else {
out = new_base.as_strided(
self.sizes(), self.strides(), self.storage_offset());
}
}
}
return THPVariable_Wrap(std::move(out));
END_HANDLE_TH_ERRORS
}
static PyObject* THPVariable_view_func(
PyObject* self_,
PyObject* args,
PyObject* kwargs) {
return view_func_impl(self_, args, kwargs, /*check_has_same_meta=*/true);
}
static PyObject* THPVariable_view_func_unsafe(
PyObject* self_,
PyObject* args,
PyObject* kwargs) {
return view_func_impl(self_, args, kwargs, /*check_has_same_meta=*/false);
}
static PyObject* rev_view_func_impl(PyObject* self_, PyObject* arg) {
HANDLE_TH_ERRORS
const auto& self = THPVariable_Unpack(self_);
TORCH_CHECK(
THPVariable_Check(arg),
"_rev_view_func expect a single argument that is a Tensor");
const auto& new_view = THPVariable_Unpack(arg);
// Ensure that self is indeed a backward differentiable view
// If not, we return an undefined Tensor (None) and let the user handle it.
auto diff_view_meta = torch::autograd::impl::get_view_autograd_meta(self);
at::Tensor out;
if (diff_view_meta && diff_view_meta->has_bw_view()) {
const auto& view_info = diff_view_meta->get_backward_view();
// Do the actual view replay
TORCH_CHECK(view_info.has_view_fn(), "No _rev_view_func() found");
out = view_info.rev_view_fn()(new_view);
}
return THPVariable_Wrap(std::move(out));
END_HANDLE_TH_ERRORS
}
static PyObject* THPVariable_rev_view_func_unsafe(
PyObject* self_,
PyObject* arg) {
return rev_view_func_impl(self_, arg);
}
// Instantiates a subclass of self with the same data.
static PyObject* THPVariable_as_subclass(
PyObject* _self,
PyObject* args,
PyObject* kwargs) {
HANDLE_TH_ERRORS
const auto& self = THPVariable_Unpack(_self);
static PythonArgParser parser({
"as_subclass(PyObject* cls)",
});
ParsedArgs<1> parsed_args{};
auto r = parser.parse(_self, args, kwargs, parsed_args);
PyObject* cls = r.pyobject(0);
TORCH_CHECK_TYPE(
PyType_Check(cls),
"cls must be a type (got ",
Py_TYPE(cls)->tp_name,
")");
// guard completely turns off torch dispatch modes, doesn't just pop off the
// stack
torch_dispatch_mode::StashTorchDispatchStackGuard td_g;
c10::impl::DisablePythonDispatcher dpd_g;
return THPVariable_NewWithVar(
(PyTypeObject*)cls,
self.alias(),
c10::impl::PyInterpreterStatus::DEFINITELY_UNINITIALIZED);
END_HANDLE_TH_ERRORS
}
static PyObject* THPVariable_make_subclass(
PyObject* _ignored,
PyObject* args,
PyObject* kwargs) {
HANDLE_TH_ERRORS
static PythonArgParser parser({
"_make_subclass(PyObject* cls, Tensor data, bool require_grad=False, *, c10::string_view? dispatch_sizes_strides_policy=None, bool dispatch_device=False, bool dispatch_layout=False, Device? device_for_backend_keys=None)",
});
ParsedArgs<7> parsed_args{};
auto r = parser.parse(args, kwargs, parsed_args);
PyObject* cls = r.pyobject(0);
TORCH_CHECK_TYPE(
PyType_Check(cls),
"cls must be a type (got ",
Py_TYPE(cls)->tp_name,
")");
// guard completely turns off torch dispatch modes, doesn't just pop off the
// stack
torch_dispatch_mode::StashTorchDispatchStackGuard td_g;
c10::impl::DisablePythonDispatcher dpd_g;
auto data =
r.tensor(1).detach(); // creates a fresh Tensor (DEFINITELY_UNINITIALIZED)
// We set `data`'s `allow_tensor_metadata_change` to true here, because we
// want to allow the following use case for backward compatibility:
//
// ```python
// rnn = torch.nn.RNN(100, 100, 2)
// # The following calls `torch._cudnn_rnn_flatten_weight(rnn._flat_weights,
// ...)`, # which changes storage of `rnn`'s weights in-place
// rnn.flatten_parameters()
// ```
data.unsafeGetTensorImpl()->set_allow_tensor_metadata_change(true);
data.set_requires_grad(r.toBool(2));
const auto sizes_strides_policy = r.stringViewOptional(3);
if (sizes_strides_policy.has_value()) {
data.unsafeGetTensorImpl()->set_python_custom_sizes_strides(
parseSizesStridesPolicyArgument(*sizes_strides_policy));
}
if (r.toBool(4)) {
data.unsafeGetTensorImpl()->set_python_custom_device(true);
}
if (r.toBool(5)) {
data.unsafeGetTensorImpl()->set_python_custom_layout(true);
}
if (!r.isNone(6)) {
data.unsafeGetTensorImpl()->_change_backend_component_keys(r.device(6));
}
return THPVariable_NewWithVar(
(PyTypeObject*)cls,
data,
c10::impl::PyInterpreterStatus::DEFINITELY_UNINITIALIZED);
END_HANDLE_TH_ERRORS
}
static PyObject* THPVariable_make_wrapper_subclass(
PyObject*,
PyObject* args,
PyObject* kwargs) {
HANDLE_TH_ERRORS
// NB: pin_memory doesn't actually do anything
// TODO: strides variant?
// cls: Python subclass type
// size, strides, storage_offset, memory_format, dtype: self-explanatory
// layout: memory layout, e.g. for types of Nested Tensors or other sparse
// tensors
// pin_memory, requires_grad: self-explanatory
// dispatch_sizes_strides_policy: string - which sizes/strides we should
// dispatch to a custom python implementation.
// dispatch_device: whether to dispatch to a custom python implementation
// for device
// dispatch_layout: whether to dispatch to a custom python implementation
// for layout
// _extra_dispatch_keys: additional dispatch keys to add to the tensor
// storage_size: if provided, skip storage size calculation and just use the
// value provided. One use case is for Nested Tensor, where the
// storage size cannot be calculated from the sizes/strides
// (because they contain a NestedInt).
static PythonArgParser parser({
"_make_wrapper_subclass(PyObject* cls, SymIntArrayRef size, SymIntArrayRef? strides=None, "
"SymInt? storage_offset=None, MemoryFormat? memory_format=None, ScalarType dtype=None, "
"Layout layout=torch.strided, Device device=None, bool pin_memory=False, bool requires_grad=False, "
"c10::string_view? dispatch_sizes_strides_policy=None, bool dispatch_device=False, bool dispatch_layout=False, "
"DispatchKeySet _extra_dispatch_keys=None, SymInt? storage_size=None)",
});
ParsedArgs<15> parsed_args{};
auto r = parser.parse(args, kwargs, parsed_args);
PyObject* cls = r.pyobject(0);
TORCH_CHECK_TYPE(
PyType_Check(cls),
"cls must be a type (got ",
Py_TYPE(cls)->tp_name,
")");
// This is an important safety check; without it, the default behavior will be
// to continue on to the underlying CPU/CUDA kernel advertised by the dispatch
// key, which will immediately segfault because the data pointer is null. By
// forcing users to define __torch_dispatch__ we ensure this does not happen
// TODO: This check is not complete; because the user can disable torch
// dispatch and then go again, triggering segfault. TBH I'm thinking I want
// to delete this function entirely
py::object attr = PyObject_FastGetAttrString(cls, "__torch_dispatch__");
TORCH_CHECK_TYPE(
attr.ptr() != nullptr &&
attr.ptr() != torch::disabled_torch_dispatch_impl(),
((PyTypeObject*)cls)->tp_name,
" must define __torch_dispatch__");
const auto options = TensorOptions()
.dtype(r.scalartype(5))
.device(r.device(7))
.layout(r.layoutOptional(6))
// NB: long standing issue, requires_grad is not
// respected here; you have to set it post facto, see
// https://github.com/pytorch/pytorch/issues/26428
// .requires_grad(r.toBool(7))
.pinned_memory(r.toBool(8));
// don't bother releasing GIL here, as we are not allocating any nontrivial
// data
Tensor tensor;
{
AutoDispatchBelowADInplaceOrView guard{}; // TODO: Remove.
tracer::impl::NoTracerDispatchMode tracer_guard{};
auto sym_sizes = r.symintlist(1);
auto sym_strides_own = r.symintlistOptional(2);
auto sym_strides =
static_cast<std::optional<c10::SymIntArrayRef>>(sym_strides_own);
auto sym_storage_offset = r.toSymIntOptional(3);
c10::SymInt size_bytes;
auto dtype_itemsize = static_cast<int64_t>(options.dtype().itemsize());
auto storage_size = r.toSymIntOptional(14);
if (storage_size.has_value()) {
size_bytes = storage_size.value();
} else if (sym_strides.has_value()) {
size_bytes = at::detail::computeStorageNbytes(
sym_sizes,
sym_strides.value(),
dtype_itemsize,
sym_storage_offset.value_or(0));
} else {
size_bytes = at::detail::computeStorageNbytesContiguous(
sym_sizes, dtype_itemsize, sym_storage_offset.value_or(0));
}
// We use storages **only** to track aliasing of subclasses during tracing.
// The actual data pointers are not valid.
Storage storage{
Storage::use_byte_size_t{},
size_bytes,
/*allocator=*/c10::GetAllocator(c10::kMeta),
/*resizable=*/true};
// TODO: constructor should probably accept data pointer
storage.set_data_ptr_noswap(at::DataPtr{nullptr, r.device(7)});
auto keys = c10::DispatchKeySet({options.computeDispatchKey()});
if (auto mb_extra_keys = r.toDispatchKeySetOptional(13)) {
keys = keys | *mb_extra_keys;
}
tensor = at::detail::make_tensor<TensorImpl>(
std::move(storage), keys, options.dtype());
TensorImpl* tensor_impl = tensor.unsafeGetTensorImpl();
if (sym_strides.has_value()) {
tensor_impl->set_sizes_and_strides(
sym_sizes, sym_strides.value(), sym_storage_offset);
} else {
TORCH_CHECK(
!sym_storage_offset.has_value(),
"setting storage offset without stride not supported");
tensor_impl->generic_set_sizes_contiguous(sym_sizes);
}
const auto sizes_strides_policy = r.stringViewOptional(10);
if (sizes_strides_policy.has_value()) {
tensor.unsafeGetTensorImpl()->set_python_custom_sizes_strides(
parseSizesStridesPolicyArgument(*sizes_strides_policy));
}
}
tensor.set_requires_grad(r.toBool(9));
if (r.toBool(11)) {
tensor.unsafeGetTensorImpl()->set_python_custom_device(true);
}
if (r.toBool(12)) {
tensor.unsafeGetTensorImpl()->set_python_custom_layout(true);
}
return THPVariable_NewWithVar(
(PyTypeObject*)cls,
tensor,
c10::impl::PyInterpreterStatus::DEFINITELY_UNINITIALIZED);
END_HANDLE_TH_ERRORS
}
using getter = PyObject* (*)(PyObject*, void*);
using setter = int (*)(PyObject*, PyObject*, void*);
PyObject* THPVariable_get_python_dispatch(THPVariable* self, void* unused) {
HANDLE_TH_ERRORS
const auto& var = THPVariable_Unpack(self);
return torch::autograd::utils::wrap(
var.unsafeGetTensorImpl()->is_python_dispatch());
END_HANDLE_TH_ERRORS
}
// CRTP base class to implement the python bindings for a Tensor property in
// PyTorch A class that implements a property is expected to have:
// - static constexpr const char* name;
// - This variable should hold the Python name of the property
// - static Tensor fn(const Tensor&);
// - This function calls the relevant ATen on the tensor
template <typename T>
struct GetterBase {
static PyObject* getter(THPVariable* self, void* /*unused*/) {
HANDLE_TH_ERRORS
if (check_has_torch_function((PyObject*)self)) {
return handle_torch_function_getter(self, T::name);
}
return THPVariable_Wrap(T::fn(THPVariable_Unpack(self)));
END_HANDLE_TH_ERRORS
}
};
struct PropertyT : GetterBase<PropertyT> {
static constexpr const char* name = "T";
static Tensor fn(const Tensor& t) {
return t.numpy_T();
}
};
struct PropertyH : GetterBase<PropertyH> {
static constexpr const char* name = "H";
static Tensor fn(const Tensor& t) {
return t.matrix_H();
}
};
struct PropertymT : GetterBase<PropertymT> {
static constexpr const char* name = "mT";
static Tensor fn(const Tensor& t) {
return t.mT();
}
};
struct PropertymH : GetterBase<PropertymH> {
static constexpr const char* name = "mH";
static Tensor fn(const Tensor& t) {
return t.mH();
}
};
struct PropertyData : GetterBase<PropertyData> {
static constexpr const char* name = "data";
static Tensor fn(const Tensor& t) {
return t.variable_data();
}
};
struct PropertyGrad : GetterBase<PropertyGrad> {
static constexpr const char* name = "grad";
static Tensor fn(const Tensor& t) {
return t.grad();
}
};
struct PropertyReal : GetterBase<PropertyReal> {
static constexpr const char* name = "real";
static Tensor fn(const Tensor& t) {
return at::real(t);
}
};
struct PropertyImag : GetterBase<PropertyImag> {
static constexpr const char* name = "imag";
static Tensor fn(const Tensor& t) {
return at::imag(t);
}
};
PyObject* THPVariable_get_cdata(THPVariable* self, void* unused) {
HANDLE_TH_ERRORS
if (check_has_torch_function((PyObject*)self)) {
return handle_torch_function_getter(self, "_cdata");
}
const auto& var = THPVariable_Unpack(self);
return PyLong_FromVoidPtr(var.unsafeGetTensorImpl());
END_HANDLE_TH_ERRORS
}
PyObject* THPVariable_get_version(THPVariable* self, void* unused) {
HANDLE_TH_ERRORS
if (check_has_torch_function((PyObject*)self)) {
return handle_torch_function_getter(self, "_version");
}
const auto& var = THPVariable_Unpack(self);