Implement the portable kernel for _embedding_bag.out

[yenhao-huang]

Summary

Fixes #11323

ExecuTorch has standalone embedding and reduction kernels, but no portable kernel for _embedding_bag.out. This adds a kernel that looks up embedding rows and reduces multiple bags in one operator call, supporting sum, mean, and max modes.

The implementation supports per-sample weights for sum mode, include_last_offset, padding_idx, and Int or Long indices and offsets. It also produces the auxiliary outputs in the shapes and forms expected by the ATen CPU operator.

Algorithm

Determine the number of bags and embedding dimension.
Resize and initialize the output and auxiliary tensors.

For each bag:
    Read its start and end positions from offsets.

    For each index in the bag:
        Look up the embedding row.

        If mode is max:
            Take the elementwise maximum and record its source row.
        Otherwise:
            Add the row to the output.
            For weighted sum, multiply it by per_sample_weights[i].
        Increment the count.

    If mode is mean and count > 0:
        Divide the output row by count.
    Write the bag's auxiliary results.

File changes

• kernels/portable/cpu/op_embedding_bag.cpp: implements the portable kernel.
• kernels/aten/functions.yaml and kernels/portable/functions.yaml: declare and register _embedding_bag.out.
• kernels/test/op_embedding_bag_test.cpp: adds cases for sum, mean, max, weighted sum, padding, final offset, empty bags, and invalid per-sample weights.
• kernels/test/CMakeLists.txt, kernels/test/targets.bzl, and shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl: include the kernel and tests in the build targets.

Test plan

Ran the targeted portable kernel tests on Linux:

GTEST_FILTER='OpEmbeddingBagOutTest.*' \
  ctest --test-dir cmake-out \
  -R '^portable_kernels_test$' \
  --output-on-failure

Result: portable_kernels_test passed; CTest reported 1/1 test target passed, 0 failed.

cc @larryliu0820 @manuelcandales @JakeStevens @nil-is-all

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/23871

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 9810885 with merge base 0486667 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[JakeStevens]

@claude please review

[claude]

Claude finished @JakeStevens's task in 1m 19s —— View job

---

Summary: One thing blocks this PR: the kernel never bounds-checks indices against weight.size(0), so a bad index reads memory outside weight. There are also a few smaller validation gaps and some missing test coverage.

Detailed review

PR Review: #23871

Summary

This PR adds a portable _embedding_bag.out kernel that supports sum, mean and max modes, per_sample_weights, padding_idx and include_last_offset. The main problem is that index values are never bounds-checked. There are also smaller gaps in offset validation and in test coverage.

Security

• Embedding rows are read without a bounds check (kernels/portable/cpu/op_embedding_bag.cpp:242-247). row = index_at(indices, i) goes straight into weight_data + row * dim. If an index is negative (other than padding_idx) or >= weight.size(0), the kernel reads outside the weight buffer. A malformed model or bad input is enough to trigger this. op_embedding.cpp:42-59 already rejects both cases with ET_KERNEL_CHECK_MSG(..., InvalidArgument, ...), and this kernel should do the same.
  • A related case: when padding_idx == -1 (meaning "no padding"), an index of -1 matches padding_idx and is silently skipped. ATen raises an error for this index instead. Run the range check before the padding_idx comparison.
  • The check can't go in check_embedding_bag_args, because that runs before any data is read. Put it in the per-index loop, or in a first pass over indices.
  • Fix this →%20before%20use%20(mirroring%20op_embedding.cpp%20ET_KERNEL_CHECK_MSG%20pattern)%2C%20performing%20the%20range%20check%20before%20the%20padding_idx%20comparison.%20Add%20a%20failure%20test%20in%20kernels%2Ftest%2Fop_embedding_bag_test.cpp%20for%20out-of-range%20and%20negative%20indices.&repo=pytorch/executorch)

Code Quality

• Offsets aren't validated (op_embedding_bag.cpp:196-205, :227-238). The loops clamp to [0, num_indices), so they never touch memory outside the buffers. However:
  • Offsets that don't start at 0, decrease, or run past the end are silently clipped and give wrong results.
  • ATen raises an error when offsets[0] != 0.
  • Checking once up front, as check_embedding_bag_args already does for other arguments, would remove the i < num_indices and i < 0 guards from both loops.
• Mixed index dtypes are accepted. indices and offsets are each checked against {Int, Long} separately, so mixing the two passes validation. I believe ATen requires them to match (checkSameType), but I didn't confirm that, so treat it as low confidence. If ATen does require it, add a tensors_have_same_dtype(indices, offsets) check to keep the portable and ATen behaviour the same.
• Half/BF16 lose precision (op_embedding_bag.cpp:257-269). With ET_SWITCH_FLOATHBF16_TYPES, the sums are accumulated directly in CTYPE. For Half/BF16, summing or averaging large bags in reduced precision can drift noticeably from ATen. Accumulating in float (or an accumulator type) and casting when writing out would fix this.

Testing

• Most supported inputs aren't tested. Every test uses Float weights and Long indices and offsets. The PR says Int indices and offsets plus Half/BF16 weights are supported, but none of them are exercised. Make the tests templated over dtype, as other portable op tests do with ET_FORALL_*, and add at least one Int index case.
• No failure tests for:
  • out-of-range or negative indices (needed once the bounds check above exists);
  • an invalid mode;
  • per_sample_weights with the wrong length;
  • a wrong dtype for offset2bag, bag_size or max_indices.

Recommendation

Request Changes: the missing index bounds check is a memory-safety bug. The other items can be fixed in the same pass.

[yenhao-huang]

@JakeSteven thanks for triggering the review. All points are addressed.

S1. Indices read without a bounds check

• Fix: offsets and indices are validated in check_offsets_and_indices, which runs before any weight read or output write — op_embedding_bag.cpp#L145-L154, called at L197-L201. It runs before the padding_idx comparison (L281), so -1 is rejected when padding_idx == -1.

C1. Offsets aren't validated

• Fix: offsets must start at 0, be nondecreasing and stay within [0, num_indices] — op_embedding_bag.cpp#L130-L143. The i < num_indices / i < 0 guards are gone from the loops.

C2. Mixed index dtypes are accepted

• Fix: kept, since ATen also accepts mixed Int/Long and promotes them (source). The auxiliary outputs now follow the promoted dtype (promoteTypes) instead of always Long, which matches ATen for Int/Int — op_embedding_bag.cpp#L84-L93

C3. Half/BF16 lose precision

• Fix: Half/BF16 accumulate in float and cast once on write — op_embedding_bag.cpp#L264-L301. A bag of 4096 ones previously summed to 2048 (Half) / 256 (BF16); it now gives 4096.

T1. Most supported inputs aren't tested

• Test: ET_FORALL_FLOATHBF16_TYPES × all four Int/Long indices/offsets combinations — op_embedding_bag_test.cpp#L58-L101, #L503-L511

T2. No failure tests

• Test: indices #L451-L475, #L519-L527; mode #L540-L542; per_sample_weights #L543-L554; auxiliary dtype #L556-L583

All 27 OpEmbeddingBagOutTest cases pass locally. Could we rerun it?