Repository: pytorch
Issue: Dispatch key reorganization
Labels: module: internals,triaged

Motivation:
* Projects like XLA occasionally need the ability to override meta operators in PyTorch. Right now, it's not possible to do so using a normal override at XLA, as meta operators get processed before we get to the XLA override. This leads to things like `generic_convolution` hack to make things work out.
* There is some confusion about what the priority of various "fallback" mechanisms are. In the current dispatcher, we have a per-operator default fallback kernel, as well as per-backend fallback. Which one takes precedence?
* Profiling and tracing are traditionally code generated as part of autograd, even though they are separate concerns. We should separate these concerns.

Right now, we define dispatch keys Variable, XLA, CPU; we attempt dispatches to concrete implementations in that order. The new pseudocode for dispatch looks like this:

```
# included_keys and excluded_keys are thread local
# getDispatchKey extracts dispatch keys from tensor arguments in args
# dispatch_set is an ordered set; first() retrieves the highest priority dispatch key in the set
def dispatch(op, args):
  dispatch_set = ((get_dispatch_keys(args) & included_keys) - excluded_keys
  return dispatchWithSet(op, args)

def dispatchWithSet(dispatch_set, op, args):
    dispatch_key = dispatch_set.first()
    if dispatch_key is None:
      raise RuntimeError("no dispatch key available")
    op_fn = op.table[dispatch_key]
    if op_fn is not None:
     # TODO: It would be convenient if dispatch_set is passed along to these
     # calls to, so redispatch can be written as dispatchWithSet(dispatch_set.remove(dispatch_key), op, args)
     return op_fn(args)
    fallback_fn = fallback_table[dispatch_key]
    return fallback_fn(dispatch_set, dispatch_key, op, args)

# reference implementations 
def fallthrough_fn(dispatch_set, dispatch_key, op, args):
  return dispatchWithSet(dispatch_set.remove(dispatch_key), op, args)

def error_fn(dispatch_set, dispatch_key, op, args):
  raise RuntimeError("no implementation")
```

The ordering of dispatch keys in the set specify the priority of various handlers. The proposed new dispatch order is as follows:

1. Profiling
2. Tracing
3. XLAPreAutograd, MSNPUPreAutograd...
4. CompoundOperator
5. Autograd
6. XLA, MSNPU...
7. CPU

For each backend, you can register both a per-op handler and a backend global handler; per-op handler takes precedence. There is a default for fallthrough, we suggest the default be raise an error (to avoid unintentional fallthrough when it is not permissible.)

Some things to note:

* Fallthrough performance is important. The reference implementation above processes fallthrough in a loop, but a better optimization is to maintain a set of fallthrough keys, and mask them out of the dispatch set immediately.
* Fallthrough does NOT add a dispatch key to the excluded set. If you need to redispatch and skip earlier layers which don't add themselves to the excluded key set (Compound, in particular), you MUST call `dispatchWithSet` and mask out the key you don't want to get redispatched to.

Explanation of the new keys

* CompoundOperator. Implementations of compound operators are registered at this layer. This is a replacement for the current “fallback” dispatch key. It occurs BEFORE variable, since CompoundOperators may call other CompoundOperators and Variables. Usually, CompoundOperators are terminal: there are never implementations for Variable/XLA/CPU, because, a CompoundOperator will not redispatch; it will usually call other functions. Every tensor has CompoundOperator set.
* XLAPreAutograd. Implementations of backend-specific, autograd implementations. Sometimes, a backend needs a different implementation of a derivative that was already defined in PyTorch, or needs to override a meta operator. In those cases, they need a dispatch key that occurs before MetaOperator/Variable. So we insert a pre-autograd hook for every extensible backend here. All XLA tensors would have both the XLA and XLAPreAutograd dispatch keys set on all tensors.
* Autograd: This is just what Variable used to be, but with Profiling and Tracing removed

Some other changes related to this proposal:

* Our `variable_factories.h` currently always turn off VariableDispatch no matter it's compound or not. In other words, we treat all factory functions as non-compound ops, which is confusing.  It’d be helpful to call compound factory ops without `AutoNonVariableTypeMode`. That way factory functions are no longer special, they’re the same as normal ops.

**CHANGELOG**

* Jan 16, 2020. Updated to not privilege fallthrough or error reporting: they're expressible purely as fallthrough functions, but we special case them.

**APPENDIX: Optimized fallthrough.** This implementation applies the optimization of masking out fallthrough immediately, so that no loop is necessary.

```
def dispatch(op, args):
    dispatch_set = ((get_dispatch_keys(args) & included_keys) - excluded_keys
    return dispatchWithSet(op, args)

def dispatchWithSet(dispatch_set, op, args):
    dispatch_set_without_fallthrough = dispatch_set - op.fallthrough_keys
    dispatch_key = dispatch_set_without_fallthrough.first()
    if dispatch_key is None:
        raise RuntimeError("no dispatch key available")
    op_fn = op.table[dispatch_key]
    if op_fn is not None:
         return op_fn(args)
    fallback_fn = fallback_table[dispatch_key]
    if fallback_fn == error_fn:
        raise RuntimeError("no implementation")
    else:
        return fallback_fn(dispatch_set, dispatch_key, op, args)
```