Add mechanism to error out when registering a batching rule for a CompositeImplicitAutograd operation

[zou3519]

The problem

There are two types of batching rules:

1. decompositions for CompositeImplicitAutograd operations
2. everything else

Registering a type-2 batching rule for a CompositeImplicitAutograd operator "op" is a trap. When we do something like vmap(grad(op))(x), then op gets decomposed, and the batching rule never gets hit. It's possible that op's decomposition decomposes into other ops that don't have a batching rule or straight up doesn't work with vmap due to doing .data_ptr accesses.

We should change it so that registering a batching rule for a CompositeImplicitAutograd operator throws an error or there is some test that checks this.

The solution

@albanD tells me we have a way to tell if an operator is CompositeImplicitAutograd print(torch._C._dispatch_has_kernel_for_dispatch_key("aten::reshape", "CompositeImplicitAutograd"))

I'm not completely sure this is sufficient, because of the order in which kernels get registered:

• Is it guaranteed that when we're registering a batching rule for aten::reshape that there is a CompositeImplicitAutograd kernel for aten::reshape that we can check and then error out on?

Thoughts? cc @albanD @ezyang @samdow @Chillee

[ezyang]

The easiest way to check for this will be to load all the registrations, and then look for dead registrations (registrations that have both Batched and CompositeImplicitAutograd). This is how we test for dangling impls.

[zou3519]

The easiest way to check for this will be to load all the registrations, and then look for dead registrations (registrations that have both Batched and CompositeImplicitAutograd).

This doesn't quite work, because we manually register CompositeImplicitAutograd decompositions to the Batched key (as well as regular batching rules). So some operators have both CompositeImplicitAutograd and Batched and that is expected.

In the end state we should just have Batched always decompose on CompositeImplicitAutograd, but... the problem is many CompositeImplicitAutograd decompositions don't actually work with vmap, so we currently treat those as a black box and run them through the vmap fallback

[ezyang]

Are you looking at the computed table? Don't look at the computed table, look at the actual registrations. (Or is that what you mean? Can you give a code pointer?) How can I tell if a "legit" version is ok versus a not legit version.

[zou3519]

How can I tell if a "legit" version is ok versus a not legit version.

Lol I think the problem is that we are unable to tell, so we need to add some structure.

For an operator that is CompositeImplicitAutograd, the correct registration is either no registration or one of these: https://github.com/pytorch/pytorch/blob/0a7d8b40b6f956a14b7ea02e04f596e914414c47/aten/src/ATen/functorch/BatchRulesDecompositions.cpp#L29-L32
(OP_DECOMPOSE(arctan2) expands to is m.impl(arctan2, at::native::arctan2)).

We would like it to be an error if, for an operator that is CompositeImplicitAutograd, a developer tries to do anything else. For arctan2, this would be like registering a binary pointwise batching rule to it: https://github.com/pytorch/pytorch/blob/0a7d8b40b6f956a14b7ea02e04f596e914414c47/aten/src/ATen/functorch/BatchRulesBinaryOps.cpp#L369 . This macro ends up expanding to m.impl(clamp_min, /*some function that is not at::native::clamp_min */)

[ezyang]

You can cheaply add an alias key to distinguish between OP_DECOMPOSE regs, and other regs.

[zou3519]

You can cheaply add an alias key to distinguish between OP_DECOMPOSE regs, and other regs.

This makes a lot of sense to me. To unpack that for others:

First step is to get the batching rules in BatchRulesDecompositions to register to an alias dispatch key

• we'd add a new alias dispatch key called CompositeImplicitBatched, to match CompositeImplicitAutograd
• Change the registrations in BatchRulesDecompositions to register everything in there to the CompositeImplicitBatched key.
• Map registrations on CompositeImplicitBatched to the FuncTorchBatched key so that vmap actually gets support for them (example: https://github.com/pytorch/pytorch/blob/eb99c1efce20aedf67e696ba5aa61fecb5651838/aten/src/ATen/core/dispatch/OperatorEntry.cpp#L351-L356)

Second step is to set up a test (similar to the dangling impls test). If an operator has a kernel for both FuncTorchBatched AND CompositeImplicitAutograd, then we raise an error:

• https://github.com/pytorch/pytorch/blob/b2679dc61cf792e922ba56b1ccb75982e6c20553/test/test_dispatch.py#L762-L766

[bdhirsh]

I'm probably just forgetting something, but - if an op has a CompositeImplicitAutograd decomposition, then should vmap always run that decomposition? Since as you mentioned above, we definitely don't want to register a custom batching rule. If that's true, then is there any reason not to add FuncTorchBatched to the alias set for CompositeImplicitAutograd?

At the very least the new alias key gives you the flexibility to opt / not opt a decomp into vmap.

[zou3519]

if an op has a CompositeImplicitAutograd decomposition, then should vmap always run that decomposition?

We (@Chillee) tried this before for a period of time and then undid the change. The problem is that not all decompositions were composite compliant, so the default of not running the decomposition is a bit better than a default of always decomposing.

Now that most (if not all) of our OpInfo-tested operations are composite compliant, we can revisit the problem. In general I want more confidence that composite compliance testing covers all PyTorch operations before we change the default to "always decompose" (I don't have enough confidence in OpInfos coverage of PyTorch operations yet)