[pytorch] reintroduce static dispatch

[ljk53]

Stack from ghstack:

• [pytorch] reintroduce static dispatch #51554 [pytorch] reintroduce static dispatch

The new static dispatch and c10 registration can work together.

It generates static dispatch code for selected backends (if set) and
fallback to regular dispatch for the rest. This way, it can be
used to reduce dispatcher's overhead for perf sensitive use cases
without compromising the functionality.

If the static_dispatch_backends flag is not set, then the behavior is
the same as before.

Added back the E2E mobile static dispatch CI for testing purpose.

This PR doesn't try to optimize mobile build size yet. We can introduce
separate build flags to disable the fallback logic, with which the linker
can strip out unused op invocation code.

Static dispatch for manually registrated ops / custom ops / autograd
kernels are not handled by this PR. We can work on these special cases
progressively.

• Sample code (with static dispatch backend = CPU):

// aten::add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
Tensor Tensor::add(const Tensor & other, Scalar alpha) const {
    DispatchKeySet _dk_set = c10::detail::multi_dispatch_key_set(other, const_cast<Tensor&>(*this));
    DispatchKey _dk = c10::impl::dispatchTypeId(_dk_set, DispatchKeySet::FULL);
    switch (_dk) {
    case DispatchKey::BackendSelect:
        // fallthrough
    case DispatchKey::CPU:
        return at::cpu::add(const_cast<Tensor&>(*this), other, alpha);
    default:
        // fallback to regular dispatch
        // TORCH_CHECK(false, "Unsupported static dispatch", _dk);
        break;
    }
    static auto op = c10::Dispatcher::singleton()
        .findSchemaOrThrow("aten::add", "Tensor")
        .typed<Tensor (const Tensor &, const Tensor &, Scalar)>();
    return op.call(const_cast<Tensor&>(*this), other, alpha);
}

• If the op has BackendSelect kernel, then it should fallback to c10 dispatch:

// aten::arange(Scalar end, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
Tensor arange(Scalar end, const TensorOptions & options) {
    DispatchKey _dk = options.computeDispatchKey();

    switch (_dk) {
    case DispatchKey::CPU:
        return at::math::arange(end, options);
    default:
        // fallback to regular dispatch
        // TORCH_CHECK(false, "Unsupported static dispatch", _dk);
        break;
    }
    static auto op = c10::Dispatcher::singleton()
        .findSchemaOrThrow("aten::arange", "")
        .typed<Tensor (Scalar, c10::optional<ScalarType>, c10::optional<Layout>, c10::optional<Device>, c10::optional<bool>)>();
    return op.call(end, optTypeMetaToScalarType(options.dtype_opt()), options.layout_opt(), options.device_opt(), options.pinned_memory_opt());
}

• If the op only has math kernel and there is no tensor argument / tensor option
  to infer the dispatch key from, then always dispatch to math kernel (only if
  static_dispatch_backends is set).

// aten::_nnpack_available() -> bool
bool _nnpack_available() {
    return at::math::_nnpack_available();
    static auto op = c10::Dispatcher::singleton()
        .findSchemaOrThrow("aten::_nnpack_available", "")
        .typed<bool ()>();
    return op.call();
}

• If the op doesn't have CPU backend, then nothing changes:

// aten::quantized_batch_norm(Tensor input, Tensor? weight, Tensor? bias, Tensor mean, Tensor var, float eps, float output_scale, int output_zero_point) -> Tensor
Tensor quantized_batch_norm(const Tensor & input, const c10::optional<Tensor> & weight, const c10::optional<Tensor> & bias, const Tensor & mean, const Tensor & var, double eps, double output_scale, int64_t output_zero_point) {
    static auto op = c10::Dispatcher::singleton()
        .findSchemaOrThrow("aten::quantized_batch_norm", "")
        .typed<Tensor (const Tensor &, const c10::optional<Tensor> &, const c10::optional<Tensor> &, const Tensor &, const Tensor &, double, double, int64_t)>();
    return op.call(input, weight, bias, mean, var, eps, output_scale, output_zero_point);
}

Differential Revision: D26197326

[facebook-github-bot]

💊 CI failures summary and remediations

As of commit 8825104 (more details on the Dr. CI page):

---

• 2/2 failures possibly* introduced in this PR
  • 2/2 non-CircleCI failure(s)

---

Extra GitHub checks: 1 failed

• Failed: GitHub Actions - flake8-py3

---

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions to the (internal) Dr. CI Users group.

[ezyang]

nit: Today, functionally it doesn't make a difference, but it would be better to put this translation rule inside solve itself, so that we are still uniformly doing backward inference. Because the rule here is very simple it can be done with forward and backward, so it's mostly a uniformity thing.

[ezyang]

(The trouble with forward inference is when you start stuffing the context with tons and tons of possible conversions "just because they might help"; backward lets you be a lot more directed about things. Though it's not that bad of an idea; see Datalog for example :)

[ezyang]

You know what I changed my mind, unpacking of optional to tensor should be done as forward inference.

[ezyang]

I hope that c10::detail::multi_dispatch_key_set is just as good as this ;)

[ezyang]

When there's a structured delegate, the dispatch table is automatically generated based on the out variant (https://github.com/pytorch/rfcs/blob/rfc-0005/RFC-0005-structured-kernel-definitions.md#structured-keyword-proposal ; there is no dispatch table for upsample_nearest1d because it delegates its dispatch to upsample_nearest1d_out). We're still on the hook for generating wrapper functions for all the variants.

I'm not sure if that helps you globally here, still reading.