1#include <ATen/core/Tensor.h>
2#include <ATen/core/dispatch/Dispatcher.h>
3#include <ATen/core/op_registration/op_registration.h>
4#include <ATen/native/UnaryOps.h>
5#include <ATen/native/Resize.h>
6#include <c10/util/irange.h>
7#include <torch/library.h>
8
9#ifndef AT_PER_OPERATOR_HEADERS
10#include <ATen/Functions.h>
11#else
12#include <ATen/ops/clone.h>
13
14#include <utility>
15#endif
16
17namespace at {
18namespace native {
19// This fallback should only be used for operations that are self inverse and have a corresponding tensor
20// bit (internally implemented using DispatchKey) to maintain the state on tensor using tensor bit.
21// Currently there are two tensor bits that trigger this fallback: conjugate bit and negative bit.
22// Conjugate bit is set on a tensor when `.conj()` is called and neg bit is set on a tensor when `.conj().imag` is called.
23
24// NOTE: To use this fallback, `clone` and `copy_` should fully understand and be able to correctly handle the semantic of your math bit.
25struct MathOpFallback {
26 MathOpFallback(DispatchKey key_, string op_name_) : key(key_), op_name(std::move(op_name_)) {}
27 virtual bool is_bit_set(const Tensor&) = 0;
28 void fallback_impl(const c10::OperatorHandle& op, DispatchKeySet dispatch_keys, torch::jit::Stack* stack) {
29 /*
30 Situations to handle:
31 1. Out-of-place operation. Easy: materialize all inputs and
32 call it a day.
33 2. Inplace operation. Desugar x.add_(2) into x.conj_().add_(2).conj_().
34 Materialize other inputs as in (1).
35 3. out= operation. Desugar add(x, 2, out=y) into y.copy_(add(x, 2))
36 Materialize other inputs as in (1).
37
38 It is important to be able to tell if we READ from an argument and if we
39 WRITE to an argument. Conservative approach is to assume that we always
40 READ from an argument, but in out= operations you can skip
41 conjugating inputs on entry that never get used. In the current schema we
42 can't easily tell if the operation is in in-place or out= operation.
43
44 Note:
45 1. Mutable tensorlists containing tensors whose math bit set to true are disallowed.
46 2. Mutable tensors with math bit set to true are unconditionally cloned to ensure
47 correct behavior in the case when the mutable tensor shares memory with non mutable arguments.
48
49 If we were to in-place resolve the math bit for mutable inputs, then the non-mutable inputs sharing partial or full memory
50 with these mutable inputs would read into wrong values in the following cases:
51 1. Non mutable inputs have their math bit set to false.
52 2. Math bit for mutable input(s) is resolved before the non mutable inputs (with bit set to true and sharing memory
53 with one or more mutable arg(s)) are cloned.
54 At the end, the final value of the mutable arguments from the stack are copied into the original input mutable tensor inputs.
55 */
56 const auto& arguments = op.schema().arguments();
57 const auto num_arguments = arguments.size();
58 const auto stack_start = stack->size() - num_arguments;
59
60 c10::optional<bool> is_write;
61 for (const auto i : c10::irange(num_arguments)) {
62 // Three possible states:
63 // 1. alias_info has no value --> out-of-place operation
64 // 2. alias_info does have a value, alias_info->is_write=True --> in-place or out= operation
65 // 3. alias_info does have a value, alias_info->is_write=False --> view operation
66 const AliasInfo* alias_info = arguments[i].alias_info();
67 if (alias_info != nullptr) {
68 if (is_write.has_value()) {
69 TORCH_CHECK(*is_write == alias_info->isWrite(),
70 "Unsupported operator for ", op_name, " fallback: ", op.schema().name(),
71 op_name, " fallback doesn't work for operators with a mix "
72 "mutable and non-mutable inputs that alias with outputs, "
73 "this must be implemented manually. "
74 "If you got this error on a core op, please report a bug to PyTorch.");
75 } else {
76 is_write = alias_info->isWrite();
77 }
78 }
79 }
80
81 if (is_write.has_value() && !*is_write) {
82 // We assume that view operators automatically handle the math bit
83 // correctly by propagating the dispatch key in key_set.
84 // This is not necessarily always right, so you should test these cases.
85 op.redispatchBoxed(dispatch_keys & c10::DispatchKeySet(DispatchKeySet::FULL_AFTER, key), stack);
86 return;
87 }
88
89 // Mutable inputs with math bit set to True and their clones
90 std::vector<std::pair<Tensor, Tensor>> mutable_inputs_with_their_clones;
91 for (const auto i : c10::irange(num_arguments)) {
92 auto& ivalue = (*stack)[stack_start + i];
93 if (!(ivalue.isTensor() || ivalue.isTensorList())) {
94 continue;
95 }
96 const auto& argument = arguments[i];
97 bool mut_arg = false;
98 if (argument.alias_info()) {
99 // Was already tested by is_write loop above
100 TORCH_INTERNAL_ASSERT_DEBUG_ONLY(argument.alias_info()->isWrite());
101 mut_arg = true;
102 }
103 if (ivalue.isTensor()) {
104 if (!is_bit_set(ivalue.toTensor())) {
105 continue;
106 }
107 auto tensor = std::move(ivalue).toTensor();
108 auto resolved_tensor = at::clone(tensor);
109 if (mut_arg) {
110 TORCH_CHECK(mutable_inputs_with_their_clones.empty(), op_name, " fallback does not support operators with more than one mutable tensors with ",
111 op_name, "bit set to true.");
112 mutable_inputs_with_their_clones.emplace_back(std::move(tensor), resolved_tensor);
113 }
114 (*stack)[stack_start + i] = std::move(resolved_tensor);
115 } else if (ivalue.isTensorList()) {
116 auto tensors = std::move(ivalue).toTensorList();
117 for(const auto j : c10::irange(tensors.size())) {
118 const auto& tensor = tensors[j];
119 if (!is_bit_set(tensor)) {
120 continue;
121 }
122 TORCH_CHECK(!mut_arg, " fallback doesn't currently support mutable TensorLists with ",
123 op_name, " inputs. Please materialize all the ", op_name, " input tensor(s) in the mutable TensorList inputs before calling ",
124 op.schema().name());
125 tensors[j] = at::clone(tensor);
126 }
127 (*stack)[stack_start + i] = std::move(tensors);
128 }
129 }
130
131 op.redispatchBoxed(dispatch_keys & c10::DispatchKeySet(DispatchKeySet::FULL_AFTER, key), stack);
132
133 TORCH_INTERNAL_ASSERT(mutable_inputs_with_their_clones.size() <= 1);
134
135 for (std::pair<Tensor, Tensor> mut_tensors: mutable_inputs_with_their_clones) {
136 auto& mutable_input = mut_tensors.first;
137 auto& cloned_mutable_input = mut_tensors.second;
138 auto& ivalue = (*stack)[stack_start];
139 auto returned_output = std::move(ivalue).toTensor();
140
141 // sanity check to ensure that the tensor in stack aliases the cloned_mutable_input
142 TORCH_INTERNAL_ASSERT(cloned_mutable_input.is_same(returned_output));
143
144 // necessary for out= arg
145 at::native::resize_output(mutable_input, returned_output.sizes());
146
147 mutable_input.copy_(returned_output);
148 (*stack)[stack_start] = std::move(mutable_input);
149 }
150 }
151
152 virtual ~MathOpFallback() = default;
153
154 DispatchKey key;
155 string op_name;
156};
157}
158}// namespace at
159