Use general ATen dispatch mechanism

[fengyuan14]

🚀 Feature

Use general ATen dispatch mechanism in some case, like at::lstm.

Motivation

When implementing a new out-of-source ATen backend extension for PyTorch, we find it is no chance to hook whole lstm cell based on existing at::lstm dispatch strategy. Our backend gets bad performance here.

In detail, backends of at::lstm are bypassed at runtime. VariableType::lstm calls TypeDefaultType::lstm directly, bypasses c10::dispatcher. Except for two backends, cudnn and mimopen, which are dispatched by hard-code “if-else”, others backends fall down to component operators, like at::matmul, at::sigmoid, at::tanh...

// TypeDefault::lstm -> at::native::lstm
at::native::lstm() {
  if (at::cudnn_is_acceptable(_input)) {
    ...
    return at::_cudnn_rnn();
  }
  if (at::use_miopen(_input, dropout_p)) {
    ...
    return at::miopen_rnn();
  }
  lstm_default_impl(); // fall down to component Ops, at::matmul, at::tanh...
}

We also find same issue in at::_embedding_bag_backward.

Pitch

at::convolution_overrideable, at::native_batch_norm and at::native_layer_norm are good examples in similar case.

[ezyang]

Yeah, basically we should do the same thing as convolution_overrideable in this case. Or move to the glorious new world described in #29548

[fengyuan14]

Do you mean you won't bypass at::lstm, and TypeDefault::lstm should be a fallback for redispatch. Please correct me.

[ezyang]

No, we will still bypass lstm. But if it's not one of our hardcoded backends (cpu/cuda), we'll call lstm_overrideable, which is NOT bypassed. This is how the convolution overrideable business works.

[fengyuan14]

Thanks, it is clear to me.