Guard {set,rebase}_history on grad_fn check (#19623)
Summary:
We would previously have statements like
```
set_history(flatten_tensor_args( result ), grad_fn);
```
Internally, {set,rebase}_history would check grad_fn and short circuit if it is nullptr. However, this means that we are executing the expression `flatten_tensor_args( result )` and immediately throwing away the results. This was causing unnecessary allocations + overhead.
My JIT overhead benchmark script (with custom benchmark method):
```
import torch, time
torch.jit.script
def add(x, y):
return x + y
a = torch.rand([])
b = torch.rand([])
niter = 1000000
with torch.no_grad():
s = time.time()
add.__getattr__('forward').benchmark(niter, a, b)
e = time.time() - s
print('overhead per call (us)', e / niter * 1e6)
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/19623
Differential Revision: D15053399
Pulled By: jamesr66a
fbshipit-source-id: 8777e1a2b5c5a5bbd3a035b7247c8154c5fc4aa6
diff --git a/tools/autograd/gen_variable_type.py b/tools/autograd/gen_variable_type.py
index 87f1318..f7a68d8 100644
--- a/tools/autograd/gen_variable_type.py
+++ b/tools/autograd/gen_variable_type.py
@@ -196,7 +196,9 @@
""")
SET_HISTORY = CodeTemplate("""\
-${fn}_history(${differentiable_outputs}, grad_fn);
+if (grad_fn) {
+ ${fn}_history(${differentiable_outputs}, grad_fn);
+}
""")
CONDITIONAL = CodeTemplate("""\
diff --git a/torch/csrc/autograd/functions/utils.h b/torch/csrc/autograd/functions/utils.h
index c632526..b4ee46d 100644
--- a/torch/csrc/autograd/functions/utils.h
+++ b/torch/csrc/autograd/functions/utils.h
@@ -51,14 +51,13 @@
inline void set_history(
at::Tensor& variable,
const std::shared_ptr<Function>& grad_fn) {
- if (grad_fn) {
- if (variable.defined()) {
- auto output_nr =
- grad_fn->add_input_metadata(variable);
- as_variable_ref(variable).set_gradient_edge({grad_fn, output_nr});
- } else {
- grad_fn->add_input_metadata(Function::undefined_input());
- }
+ AT_ASSERT(grad_fn);
+ if (variable.defined()) {
+ auto output_nr =
+ grad_fn->add_input_metadata(variable);
+ as_variable_ref(variable).set_gradient_edge({grad_fn, output_nr});
+ } else {
+ grad_fn->add_input_metadata(Function::undefined_input());
}
}
