pytorch / pytorch/executorch

Vectorize optimized_portable_ops versions of portable ops?

Open
#9,241 16 comments 0 reactions 1 assignee View on GitHub

@swolchok is already working on this.

Since Mar 13, 2025.

actionable module: kernels
Dominant language
Python
Stars
5k
Forks
1.2k
Avg merge
2d 10h
Merged PRs (30d)
581

Description

🚀 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 https://github.com/pytorch/pytorch/issues/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

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.