Repository: pytorch
Issue: Operators that return dynamic-shape outputs that require_grad choke in AOTAutograd / unbacked SymInts are reallocated
Labels: triaged,oncall: pt2,module: dynamic shapes,module: aotdispatch,module: pt2-dispatcher

Custom autograd functions that use `impl_abstract` that give outputs with dynamic shapes are not supported in PT2.

Repro:
```python
import torch
from torch.library import Library, impl
from torch.testing._internal.optests import aot_autograd_check

m = Library("my_custom_library", "DEF")
m.define("custom_op(Tensor input, Tensor lengths) -> Tensor")

class CustomOp(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, lengths):
        ctx.save_for_backward(input, lengths)
        with torch._C._AutoDispatchBelowAutograd():
            return torch.ops.my_custom_library.custom_op(input, lengths)

    @staticmethod
    def backward(ctx, grad):
        input, lengths = ctx.saved_tensors
        idx = lengths.sum()
        grad_out = torch.zeros_like(input)
        grad_out[:idx] = grad
        return grad_out, None

@impl(m, "custom_op", "CPU")
def custom_op_impl(input, lengths):
    idx = lengths.sum()
    return input[:idx].clone()

@impl(m, "custom_op", "Autograd")
def custom_op_autograd(input, lengths):
    return CustomOp.apply(input, lengths)

@torch.library.impl_abstract("my_custom_library::custom_op")
def custom_op_meta(input, lengths):
    return input.new_empty(torch.library.get_ctx().new_dynamic_size())

def fn(x):
    return torch.ops.my_custom_library.custom_op(x, torch.tensor([1, 1, 1, 1], dtype=torch.int64))

aot_autograd_check(fn, (torch.randn(10, requires_grad=True),), {}, True)
```

Error:
```bash
Traceback (most recent call last):
  File "/data/users/williamwen/pytorch/playground2.py", line 41, in <module>
    aot_autograd_check(fn, (torch.randn(10, requires_grad=True),), {}, True)
  File "/data/users/williamwen/pytorch/torch/testing/_internal/optests/aot_autograd.py", line 76, in aot_autograd_check
    _test_aot_autograd_forwards_backwards_helper(
  File "/data/users/williamwen/pytorch/torch/testing/_internal/optests/aot_autograd.py", line 137, in _test_aot_autograd_forwards_backwards_helper
    check(args, ignore_failure=False)
  File "/data/users/williamwen/pytorch/torch/testing/_internal/optests/aot_autograd.py", line 133, in check
    compiled_out, compiled_grad = call_forwards_backwards(compiled_f, args)
  File "/data/users/williamwen/pytorch/torch/testing/_internal/optests/aot_autograd.py", line 97, in call_forwards_backwards
    out = wrapper_set_seed(f, args)
  File "/data/users/williamwen/pytorch/torch/testing/_internal/common_methods_invocations.py", line 9245, in wrapper_set_seed
    return op(*args, **kwargs)
  File "/data/users/williamwen/pytorch/torch/_functorch/aot_autograd.py", line 3758, in returned_function
    compiled_fn = create_aot_dispatcher_function(
  File "/data/users/williamwen/pytorch/torch/_dynamo/utils.py", line 221, in time_wrapper
    r = func(*args, **kwargs)
  File "/data/users/williamwen/pytorch/torch/_functorch/aot_autograd.py", line 3462, in create_aot_dispatcher_function
    compiled_fn = compiler_fn(flat_fn, fake_flat_args, aot_config, fw_metadata=fw_metadata)
  File "/data/users/williamwen/pytorch/torch/_functorch/aot_autograd.py", line 2243, in aot_wrapper_dedupe
    return compiler_fn(flat_fn, leaf_flat_args, aot_config, fw_metadata=fw_metadata)
  File "/data/users/williamwen/pytorch/torch/_functorch/aot_autograd.py", line 2423, in aot_wrapper_synthetic_base
    return compiler_fn(flat_fn, flat_args, aot_config, fw_metadata=fw_metadata)
  File "/data/users/williamwen/pytorch/torch/_functorch/aot_autograd.py", line 2837, in aot_dispatch_autograd
    fx_g = aot_dispatch_autograd_graph(flat_fn, flat_args, aot_config, fw_metadata=fw_metadata)
  File "/data/users/williamwen/pytorch/torch/_functorch/aot_autograd.py", line 2814, in aot_dispatch_autograd_graph
    fx_g = create_functionalized_graph(
  File "/data/users/williamwen/pytorch/torch/_functorch/aot_autograd.py", line 1451, in create_functionalized_graph
    fx_g = make_fx(helper, decomposition_table=aot_config.decompositions)(*args)
  File "/data/users/williamwen/pytorch/torch/fx/experimental/proxy_tensor.py", line 858, in wrapped
    t = dispatch_trace(wrap_key(func, args, fx_tracer, pre_dispatch), tracer=fx_tracer, concrete_args=tuple(phs))
  File "/data/users/williamwen/pytorch/torch/_compile.py", line 24, in inner
    return torch._dynamo.disable(fn, recursive)(*args, **kwargs)
  File "/data/users/williamwen/pytorch/torch/_dynamo/eval_frame.py", line 401, in _fn
    return fn(*args, **kwargs)
  File "/data/users/williamwen/pytorch/torch/_dynamo/external_utils.py", line 17, in inner
    return fn(*args, **kwargs)
  File "/data/users/williamwen/pytorch/torch/fx/experimental/proxy_tensor.py", line 478, in dispatch_trace
    graph = tracer.trace(root, concrete_args)
  File "/data/users/williamwen/pytorch/torch/fx/_symbolic_trace.py", line 817, in trace
    (self.create_arg(fn(*args)),),
  File "/data/users/williamwen/pytorch/torch/fx/_symbolic_trace.py", line 684, in flatten_fn
    tree_out = root_fn(*tree_args)
  File "/data/users/williamwen/pytorch/torch/fx/experimental/proxy_tensor.py", line 514, in wrapped
    out = f(*tensors)
  File "/data/users/williamwen/pytorch/torch/_functorch/aot_autograd.py", line 1440, in joint_helper
    return functionalized_f_helper(primals, tangents)
  File "/data/users/williamwen/pytorch/torch/_functorch/aot_autograd.py", line 1396, in functionalized_f_helper
    f_outs = fn(*f_args)
  File "/data/users/williamwen/pytorch/torch/_functorch/aot_autograd.py", line 1365, in inner_fn_with_anomaly
    return inner_fn(*args)
  File "/data/users/williamwen/pytorch/torch/_functorch/aot_autograd.py", line 1327, in inner_fn
    out if out.shape == tangent.shape else out.view(tangent.shape)
  File "/data/users/williamwen/pytorch/torch/__init__.py", line 361, in __bool__
    return self.node.bool_()
  File "/data/users/williamwen/pytorch/torch/fx/experimental/symbolic_shapes.py", line 1052, in bool_
    return self.guard_bool("", 0)
  File "/data/users/williamwen/pytorch/torch/fx/experimental/symbolic_shapes.py", line 1022, in guard_bool
    r = self.shape_env.evaluate_expr(self.expr, self.hint, fx_node=self.fx_node)
  File "/data/users/williamwen/pytorch/torch/fx/experimental/recording.py", line 250, in wrapper
    return event.run(self)
  File "/data/users/williamwen/pytorch/torch/fx/experimental/recording.py", line 151, in run
    return self.f(*args, **kwargs)
  File "/data/users/williamwen/pytorch/torch/fx/experimental/symbolic_shapes.py", line 3986, in evaluate_expr
    concrete_val = self.size_hint(orig_expr)
  File "/data/users/williamwen/pytorch/torch/fx/experimental/symbolic_shapes.py", line 3780, in size_hint
    raise self._make_data_dependent_error(result_expr, expr)
torch.fx.experimental.symbolic_shapes.GuardOnDataDependentSymNode: It appears that you're trying to get a value out of symbolic int/float whose value is data-dependent (and thus we do not know the true value.)  The expression we were trying to evaluate is Eq(i1, i0) (unhinted: Eq(i1, i0)).  Scroll up to see where each of these data-dependent accesses originally occurred.
```

If we remove the dynamic shapes, then the script runs fine:
```python
import torch
from torch.library import Library, impl
from torch.testing._internal.optests import aot_autograd_check

m = Library("my_custom_library", "DEF")
m.define("custom_op(Tensor input, Tensor lengths) -> Tensor")

class CustomOp(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, lengths):
        ctx.save_for_backward(input, lengths)
        with torch._C._AutoDispatchBelowAutograd():
            return torch.ops.my_custom_library.custom_op(input, lengths)

    @staticmethod
    def backward(ctx, grad):
        input, lengths = ctx.saved_tensors
        # idx = lengths.sum()
        idx = 4
        grad_out = torch.zeros_like(input)
        grad_out[:idx] = grad
        return grad_out, None

@impl(m, "custom_op", "CPU")
def custom_op_impl(input, lengths):
    # idx = lengths.sum()
    idx = 4
    return input[:idx].clone()

@impl(m, "custom_op", "Autograd")
def custom_op_autograd(input, lengths):
    return CustomOp.apply(input, lengths)

@torch.library.impl_abstract("my_custom_library::custom_op")
def custom_op_meta(input, lengths):
    # return input.new_empty(torch.library.get_ctx().new_dynamic_size())
    return input.new_empty(4)

def fn(x):
    return torch.ops.my_custom_library.custom_op(x, torch.tensor([1, 1, 1, 1], dtype=torch.int64))

aot_autograd_check(fn, (torch.randn(10, requires_grad=True),), {}, True)
```

cc @ezyang @gchanan @zou3519 @kadeng @msaroufim @wconstab @bdhirsh @anijain2305