Skip to content

Optimized le.Tensor_out returns b <= a for reverse-argument broadcasts #24485

Open

Description

@SofianElmotiem

Describe the bug

With the optimized CPU kernels, le.Tensor_out returns b <= a instead of a <= b when the left operand is the one being broadcast and both inputs and the output share a dtype (for example two Bool masks, or a non-Bool output).

Repro (C++ kernel test)

TensorFactory<ScalarType::Bool> tf;
Tensor a = tf.make({4}, {false, true, false, true});
Tensor b = tf.make({2, 4}, {false, false, true, true, true, true, false, false});
Tensor out = tf.zeros({2, 4});
torch::executor::aten::le_outf(context_, a, b, out);
Result
Expected (a <= b, also what the portable kernel returns) {1, 0, 1, 1, 1, 1, 1, 0}
Optimized kernel {1, 1, 0, 1, 0, 1, 1, 1} (that is b <= a)

The same happens for Float and Int inputs with a Float/Int output, on all three reverse-argument broadcast paths: (4,) vs (2, 4), (2, 1) vs (2, 3), and (2, 1, 2) vs (2, 3, 2). Float inputs with a Bool output are not affected, because select_optimized_path returns kNone when the output dtype differs and the portable comparison is used.

Cause

select_optimized_path returns kBroadcast2dBy1dReverseArguments, kBroadcastLastDimReverseArguments or kBroadcastNdByNdReverseArguments when the left operand is the broadcast one, and handle_broadcast_elementwise then calls the lambda with the operands swapped. op_div.cpp and op_add_sub_impl.h account for this by swapping their lambda on those paths, but op_le.cpp always passes x.le(y):

https://github.com/pytorch/executorch/blob/aed328df53/kernels/optimized/cpu/op_le.cpp#L57-L63

The existing BroadcastLastDimReverseTest, BroadcastNdByNdReverseTest and similar tests in kernels/test/op_le_test.cpp use a Bool output with Int inputs, so they take the portable path and never exercise the optimized broadcast code.

I have a fix with regression tests and will open a PR.

Versions

main at aed328d, built from source with CMake (EXECUTORCH_BUILD_TESTS=ON, EXECUTORCH_BUILD_KERNELS_OPTIMIZED=ON), Ubuntu, x86_64, GCC 15.2.0.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions