AOT autograd - Views

Karthik_Nagasubraman: <p>AOT autograd functionalizes the views correctly only if at least one of the input tensors is having required_grad == True.</p>
<p>In <a href="https://github.com/pytorch/pytorch/blob/master/torch/_functorch/aot_autograd.py" class="inline-onebox" rel="noopener nofollow ugc">pytorch/aot_autograd.py at master · pytorch/pytorch · GitHub</a><br>
if needs_autograd:<br>
compiler_fn = aot_dispatch_autograd → This functionalizes the views correctly<br>
else:<br>
compiler_fn = aot_dispatch_base → This is resulting in incorrect view handling</p>
<p>import torch<br>
from functorch.compile import aot_function</p>
<p>def fn(a, b):<br>
e = a[:]<br>
e =  a.as_strided([2, 2], [4, 2], 0) #[:,0:-1:2] <span class="hashtag">#a</span>.view([8])<br>
e.add_(1.0)<br>
a.mul_(2.0)<br>
e[2:-1] = 3.0<br>
return  e[1:-1:2]</p>
<h1>
<a name="test-that-it-works-1" class="anchor" href="#test-that-it-works-1"></a>Test that it works</h1>
<p>a = torch.randn(2, 4, requires_grad=False)<br>
b = torch.randn(2, 4, requires_grad=True)</p>
<h1>
<a name="the-compiler_fn-is-called-after-the-forward-and-backward-graphs-are-extracted-2" class="anchor" href="#the-compiler_fn-is-called-after-the-forward-and-backward-graphs-are-extracted-2"></a>The compiler_fn is called after the forward and backward graphs are extracted.</h1>
<h1>
<a name="here-we-just-print-the-code-in-the-compiler_fn-return-of-this-function-is-a-callable-3" class="anchor" href="#here-we-just-print-the-code-in-the-compiler_fn-return-of-this-function-is-a-callable-3"></a>Here, we just print the code in the compiler_fn. Return of this function is a callable.</h1>
<p>def compiler_fn(fx_module: torch.fx.GraphModule, _):<br>
print(fx_module.code)<br>
<span class="hashtag">#print</span>(fx_module.graph)<br>
return fx_module</p>
<h1>
<a name="pass-on-the-compiler_fn-to-the-aot_function-api-4" class="anchor" href="#pass-on-the-compiler_fn-to-the-aot_function-api-4"></a>Pass on the compiler_fn to the aot_function API</h1>
<p>aot_print_fn = aot_function(fn, fw_compiler=compiler_fn, bw_compiler=compiler_fn)</p>
<h1>
<a name="run-the-aot_print_fn-once-to-trigger-the-compilation-and-print-the-graphs-5" class="anchor" href="#run-the-aot_print_fn-once-to-trigger-the-compilation-and-print-the-graphs-5"></a>Run the aot_print_fn once to trigger the compilation and print the graphs</h1>
<p>res = aot_print_fn(a, b)</p>
<p><strong>Output:</strong><br>
def forward(self, primals_1, primals_2):<br>
clone = torch.ops.aten.clone.default(primals_1);  primals_1 = None<br>
as_strided = torch.ops.aten.as_strided.default(clone, [2, 2], [4, 2], 0)<br>
add = torch.ops.aten.add.Tensor(as_strided, 1.0);  as_strided = None<br>
as_strided_scatter = torch.ops.aten.as_strided_scatter.default(clone, add, [2, 2], [4, 2], 0);  clone = add = None<br>
mul = torch.ops.aten.mul.Tensor(as_strided_scatter, 2.0);  as_strided_scatter = None<br>
_tensor_constant0 = self._tensor_constant0<br>
lift_fresh_copy = torch.ops.aten.lift_fresh_copy.default(_tensor_constant0);  _tensor_constant0 = None<br>
as_strided_2 = torch.ops.aten.as_strided.default(mul, [2, 2], [4, 2], 0)<br>
slice_3 = torch.ops.aten.slice.Tensor(as_strided_2, 0, 2, -1);  as_strided_2 = None<br>
fill = torch.ops.aten.fill.Tensor(slice_3, lift_fresh_copy);  slice_3 = lift_fresh_copy = None<br>
as_strided_3 = torch.ops.aten.as_strided.default(mul, [2, 2], [4, 2], 0)<br>
slice_scatter = torch.ops.aten.slice_scatter.default(as_strided_3, fill, 0, 2, -1);  as_strided_3 = fill = None<br>
as_strided_scatter_1 = torch.ops.aten.as_strided_scatter.default(mul, slice_scatter, [2, 2], [4, 2], 0);  mul = slice_scatter = None<br>
return [as_strided_scatter_1, 0, 2, 8, 2, 4]</p>
<p>Now if we chnage b = torch.randn(2, 4, requires_grad=False) (True → False)</p>
<p>output becomes:<br>
def forward(self, arg0_1, arg1_1):<br>
slice_1 = torch.ops.aten.slice.Tensor(arg0_1, 0, 0, 9223372036854775807)<br>
as_strided = torch.ops.aten.as_strided.default(arg0_1, [2, 2], [4, 2], 0)<br>
add_ = torch.ops.aten.add_.Tensor(as_strided, 1.0);  as_strided = None<br>
mul_ = torch.ops.aten.mul_.Tensor(arg0_1, 2.0);  arg0_1 = None<br>
<em>tensor_constant0 = self.<em>tensor_constant0<br>
lift_fresh_copy = torch.ops.aten.lift_fresh_copy.default(<em>tensor_constant0);  <em>tensor_constant0 = None<br>
slice_2 = torch.ops.aten.slice.Tensor(add</em>, 0, 2, -1)<br>
fill</em> = torch.ops.aten.fill</em>.Tensor(slice_2, lift_fresh_copy);  slice_2 = lift_fresh_copy = None<br>
slice_3 = torch.ops.aten.slice.Tensor(add</em>, 0, 1, -1, 2);  add_ = None<br>
return [slice_3]</p>
<p>In the second case, the generate code is functionally incorrect</p>
<p>Question:<br>
Why not make aot_dispatch_autograd the default/only  option as it seem to functionalize views correctly?</p>

bdhirsh: <p>There’s actually a WIP PR to change this! We want functionalization to always run in AOT Autograd: <a href="https://github.com/pytorch/pytorch/pull/92857" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/pull/92857</a></p>