Missing broadcast capability in optimized kernels.

[cptspacemanspiff]

Missing broadcast causes abort.

So, I ran into a bug where the optimized kernel library, op_add, op_sub, and op_div are missing broadcast support. I specifically ran into it when trying to add a 2x8x12x12 tensor with a 2x1x12x12 tensor. (I only ran into it when I was doing batch size of greater than 1, which is my assumption why this has not been an issue?)

The cause of this seems to be an update a while back that presumably added new optimization paths for use with op_mul:

enum class ElementwiseOptimizedPath {
  kNone,
  kTreatAs1d,
  kBroadcast2dBy1d,
  kBroadcast2dBy1dReverseArguments,
  kBroadcastNdByNd,
  kBroadcastNdByNdReverseArguments,
  kBroadcastLastDim,
  kBroadcastLastDimReverseArguments,
};

However, the first 4 are the only ones known in op_add, op_sub, and op_div, and the logic falls to an error if it encounters an unexpected optimization path.

Playing with it, a quick and dirty fix is to add a check for if we are using one of the un-implemented broadcasts, then fallback to kNone

  auto selected_optimized_path = select_optimized_path(a, b, out);
  if (selected_optimized_path == ElementwiseOptimizedPath::kBroadcastNdByNd ||
      selected_optimized_path == ElementwiseOptimizedPath::kBroadcastNdByNdReverseArguments ||
      selected_optimized_path == ElementwiseOptimizedPath::kBroadcastLastDim || 
      selected_optimized_path == ElementwiseOptimizedPath::kBroadcastLastDimReverseArguments) {
      selected_optimized_path = ElementwiseOptimizedPath::kNone;
  }

That is obviously non-ideal, but was my temp fix for myself.

Versions

Executorch main branch.

cc @larryliu0820 @manuelcandales

[manuelcandales]

cc @swolchok

[swolchok]

Looks like @kimishpatel did the update in question

[kimishpatel]

Let me look into this. Enums added should have been implemented, but clearly something is off

[kimishpatel]

trying to repro you using inputs you mentioned

[kimishpatel]

@cptspacemanspiff there is a test that seems similar https://github.com/pytorch/executorch/blob/main/kernels/test/op_mul_test.cpp#L176? What is the issue you ran into? Can you provide a repro?

[cptspacemanspiff]

To be clear, multiply works, add/sub/div are broken, due to the updates to multiply.

I don't have a (easy) repro at the moment, is there a similar test for add, sub, and div?

The issue was that the changes were implemented for multiply but not the others, (to be honest I only ran into add, but looking at the code the others have the same problem.)

Here is the code where it errors out (from op_add.cpp)

if (selected_optimized_path == ElementwiseOptimizedPath::kTreatAs1d) {
    // Resize for dynamic shape
    auto error = resize_tensor(out, a.sizes());
    ET_KERNEL_CHECK_MSG(
        ctx,
        error == Error::Ok,
        InvalidArgument,
        out,
        "Failed to resize output tensor.");

    ET_SWITCH_REALB_TYPES(a_type, ctx, "add.out", CTYPE, [&]() {
      CTYPE alpha_val;
      ET_KERNEL_CHECK(
          ctx, utils::extract_scalar(alpha, &alpha_val), InvalidArgument, );

      using Vec = executorch::vec::Vectorized<CTYPE>;
      executorch::vec::map2<CTYPE>(
          [alpha_val](Vec x, Vec y) { return x + Vec(alpha_val) * y; },
          out.mutable_data_ptr<CTYPE>(),
          a.const_data_ptr<CTYPE>(),
          b.const_data_ptr<CTYPE>(),
          out.numel());
    });
  } else if (selected_optimized_path != ElementwiseOptimizedPath::kNone) {
    const Tensor* lhs;
    const Tensor* rhs;
    if (selected_optimized_path ==
        ElementwiseOptimizedPath::kBroadcast2dBy1dReverseArguments) {
      lhs = &b;
      rhs = &a;
    } else {
      // Catch failure to update logic when adding new broadcasting
      // possibility.
      ET_DCHECK(
          selected_optimized_path ==
          ElementwiseOptimizedPath::kBroadcast2dBy1d);  ///////////////errors out in this check/////////////////////////
      lhs = &a;
      rhs = &b;
    }

The optimization path is ElementwiseOptimizedPath::kBroadcastNdByNd, it falls through and when it finds an unexpected case (because it was not updated, it errors out.

[kimishpatel]

falls

do you have stack trace

[kimishpatel]

falls

do you have stack trace

nevermind, I see your comment

[kimishpatel]

I think you are right. This fell off my radar to update in other elementwise ops in optimized op. Let me fix that

[kimishpatel]

@cptspacemanspiff WIP stack #8205. sub and div are next