Vectorize optimized_portable_ops versions of portable ops?

[swolchok]

🚀 The feature, motivation and pitch

Similarly to #8932, we should be able to conditionally compile portable ops to do some vectorization. I imagine that this would look like either passing a second lambda to our util functions, or perhaps passing template lambdas that we then could use for both some scalar T and also Vectorized<T>. The second option would require us to get an std-workalike interface to Vectorized operations so that things like exp would work seemlessly, which probably would have a similar solution to pytorch/pytorch#144495 .

RFC

As a concrete example, op_add currently calls a util workhorse function with a lambda:

    utils::apply_bitensor_elementwise_fn<CTYPE_COMPUTE, op_name>(
        [val_alpha](const CTYPE_COMPUTE val_a, const CTYPE_COMPUTE val_b) {
          return val_a + val_alpha * val_b;
        },

We could imagine instead making the call look like this, with a template lambda, so that we could seamlessly use the lambda with Vectorized:

    utils::apply_bitensor_elementwise_fn<CTYPE_COMPUTE, op_name>(
        [val_alpha](const auto val_a, const auto val_b) {
          return val_a + val_alpha * val_b;
        },

A second, harder example is op_exp:

Tensor& exp_out(KernelRuntimeContext& ctx, const Tensor& in, Tensor& out) {
  return internal::unary_ufunc_realhbbf16_to_floathbf16(std::exp, ctx, in, out);
}

I think ideally we would find a solution to the above-mentioned PyTorch issue and then write this as

Tensor& exp_out(KernelRuntimeContext& ctx, const Tensor& in, Tensor& out) {
  return internal::unary_ufunc_realhbbf16_to_floathbf16_v2([](auto x) { return c10::math::exp(x); }, ctx, in, out);
}

using a template lambda that could be instantiated with either a scalar or Vectorized, as outlined above.

cc @larryliu0820 @manuelcandales

[swolchok]

Sketch of more detailed plan:

• Make sure we have committed code that uses ATen Vectorized (optimized/op_gelu does this already)
• Broaden "optimized portable ops" detection from just ET_USE_THREADPOOL -- something needs to tell us it's OK to use Vectorized. (ET_USE_VECTORIZED? what would define it? ET_BUILDING_OPTIMIZED_PORTABLE_OPS? probably we define a set of specific macros when we build the optimized_portable_ops target) (Make PyTorch headers available in optimized_portable_kernels, define ET_USE_PYTORCH_HEADERS #9384)
• Implement unary_ufunc functions using elementwise_util #9386
• Specialize elementwise_util ops for the non-mixed-dtype case. If we don't, vectorization doesn't make any sense. (RFC: Specialize for non-mixed-dtype in elementwise_util #9388)
• Proof-of-concept vectorization for elementwise_util ops using some sort of metaprogramming/SFINAE on their lambda (need to maintain backward compatibility with existing code)
• Roll out specialized, vectorized elementwise_util ops
• Establish c10::math namespace in upstream PyTorch
• Sync c10::math header into ET core so that we can use c10::math namespace in ExecuTorch portable ops, which are part of core
• Roll out vectorized unary_ufunc ops using c10::math namespace

[swolchok]

vectorization for elementwise ops

This is trickier than I remembered currently. We outline our loads in elementwise_util.h in the name of build time and code size; getting to a point where vectorization would be the next item on the list would take some time. Accordingly, I am going to skip elementwise_util.h for now and move on to unary_ufunc; will come back to elementwise afterwards.

[swolchok]

skip elementwise_util.h for now and move on to unary_ufunc

This isn't right. The unary_func_* ops in pattern.h seem to be somewhat redundant with elementwise_util.h.

[swolchok]

Plan above updated to reflect resolution of my confusion about whether to start with elementwise_util or unary_ufunc_* . (In short, unary_ufunc_* will call through to elementwise_util.)

[kimishpatel]

so that we could seamlessly use the lambda with Vectorized:

Do you have an example for this, for say add refactor that you described in the summary of the RFC

[kimishpatel]

template lambda that could be instantiated with Vectorized

Do you have an example of this

[kimishpatel]

something needs to tell us it's OK to use Vectorized.

Can we not have implementation of Vectorized that is just scalar? This is already the case in pytorch core and et, no?

[swolchok]

template lambda that could be instantiated with Vectorized

Do you have an example of this

[](auto x) { return c10::math::exp(x); }, where c10::math contains using std::exp; and also template <typename T> auto exp(at::Vectorized<T> v) { return v.exp(); }

[swolchok]

something needs to tell us it's OK to use Vectorized.

Can we not have implementation of Vectorized that is just scalar? This is already the case in pytorch core and et, no?

I don't particularly want to sign up to ensure that we generate code for that that is just as good as writing scalar code directly.

[swolchok]

I'm running into trouble with my intended implementation here. Apparently, SFINAE can't be used together with generic lambdas to detect whether they will actually compile when passed an argument of a particular type. See #9432 ; I expect to resolve this tomorrow after sleeping on the problem.