Why pytorch cuda tensor need two dispatch key, CUDA and AutogradCUDA?

username: ryanc
cooked: <p>I write a test code</p>
<pre><code class="lang-auto">import torch

device = torch.device("cuda")
x = torch.randn(2, 2).to(device)
y = torch.randn(2).to(device)

x.requires_grad_(True)
z = torch.add(x, y)

z.sum().backward()
</code></pre>
<p>Tensor on cuda has two dispatch key, CUDA and AutogradCUDA. I read pytorch dispatcher source code and find dispatch key AutogradCUDA will be excluded in VariableTypeEverything.cpp</p>
<pre><code class="lang-auto">at::Tensor add_Tensor(c10::DispatchKeySet ks, const at::Tensor &amp; self, const at::Tensor &amp; other, const at::Scalar &amp; alpha) {
  auto _tmp = ([&amp;]() {
    at::AutoDispatchBelowADInplaceOrView guard;
    return at::redispatch::add(ks &amp; c10::after_autograd_keyset, self_, other_, alpha);
  })();
}
</code></pre>
<p>Finally, pytorch dispatch to wrapper_CUDA_add_Tensor function which is previously registered into pytorch with CUDA dispatch key. Then do the real computation and get the results</p>
<p><strong>My question is</strong> why we need AutogradCUDA when create cuda tensor? What will pytorch do for dispatch key AutogradCUDA？I guess there are some relation between pytorch dispatcher and pytorch autograd system, but I can’t figure out the connection.</p>

username: ptrblck
cooked: <p>I would assume the Autograd dispatching is triggered by calling <code>x.requires_grad_(True)</code> which let’s Autograd track operations on <code>x</code> and not by the usage of the CUDA backend.</p>

username: ryanc
cooked: <p><strong>Thank you very much for your reply.</strong></p>
<p>I printed the call chain, the key call sequence is as follows</p>
<ol>
<li>torch::autograd::THPVariable_add at pytorch/torch/csrc/autograd/generated/python_torch_functions_2.cpp</li>
<li>torch::autograd::VariableType::(anonymous namespace)::add_Tensor at pytorch/torch/csrc/autograd/generated/VariableType_2.cpp</li>
</ol>
<p>VariableType_2.cpp:add_Tensor function code shows as following</p>
<pre><code class="lang-auto">at::Tensor add_Tensor(c10::DispatchKeySet ks, const at::Tensor &amp; self, const at::Tensor &amp; other, const at::Scalar &amp; alpha) {
  auto&amp; self_ = unpack(self, "self", 0);
  auto&amp; other_ = unpack(other, "other", 1);
  auto _any_requires_grad = compute_requires_grad( self, other );
  
  (void)_any_requires_grad;
  auto _any_has_forward_grad_result = (isFwGradDefined(self) || isFwGradDefined(other));
  (void)_any_has_forward_grad_result;
  std::shared_ptr&lt;AddBackward0&gt; grad_fn;
  if (_any_requires_grad) {
    grad_fn = std::shared_ptr&lt;AddBackward0&gt;(new AddBackward0(), deleteNode);
    grad_fn-&gt;set_next_edges(collect_next_edges( self, other ));
    grad_fn-&gt;other_scalar_type = other.scalar_type();
    grad_fn-&gt;alpha = alpha;
    grad_fn-&gt;self_scalar_type = self.scalar_type();
  }

  auto _tmp = ([&amp;]() {
    at::AutoDispatchBelowADInplaceOrView guard;
    return at::redispatch::add(ks &amp; c10::after_autograd_keyset, self_, other_, alpha);
  })();
  auto result = std::move(_tmp);
  if (grad_fn) {
      set_history(flatten_tensor_args( result ), grad_fn);
  }
  return result;
}
</code></pre>
<p>And add_Tensor has registered to dispatch key autograd at pytorch/torch/csrc/autograd/generated/VariableType_2.cpp</p>
<pre><code class="lang-auto">TORCH_LIBRARY_IMPL(aten, Autograd, m) {
  m.impl("add.Tensor", TORCH_FN(VariableType::add_Tensor));
}
</code></pre>
<p>So I guess that pytorch first dispatched to add_Tensor function according to the dispatch key AutogradCUDA. In add_Tensor, it builds the nodes and edges required for the computation graph for pytorch autograd.</p>
<p>But I still haven’t figured out how pytorch dispatch tensor with dispatch key AutogradCUDA to dispatch key Autograd specific function. After all, add_Tensor is registered to  dispatch key Autograd, but the tensor which dispatch key set contains AutogradCUDA is still dispatched to add_Tensor function.</p>