[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.

⚠️ 27 Awaiting Approval

As of commit a62495a with merge base 6c7b213 ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build Aarch64 Linux CUDA Wheels (gh)
• Build Aarch64 Linux Native Wheels (gh)
• Build Aarch64 Linux Wheels (gh)
• Build documentation (gh)
• Build Linux CUDA Wheels (gh)
• Build Linux Native Wheels (gh)
• Build Linux Wheels (gh)
• Build macOS Wheels (gh)
• Build Presets (gh)
• Build Windows CUDA Wheels (gh)
• Build Windows Wheels (gh)
• Cadence Build & Test (gh)
• Lint (gh)
• periodic (gh)
• pull (gh)
• Test ARM Backend (gh)
• Test CoreML Backend (gh)
• Test Cortex-M Backend (gh)
• Test CPU Backend (gh)
• Test Native Engines (gh)
• Test NXP Backend (gh)
• Test QNN Backend (gh)
• Test QNN Windows MSVC build (gh)
• Test Vulkan Backend (gh)
• Test Vulkan Backend (specialized runners) (gh)
• Test WebGPU Backend (gh)
• Test XNNPACK Backend (gh)

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?

[yenhao-huang]

Sorry, I pinged the wrong name. @JakeStevens could you rerun CI again?

[JakeStevens]

Thanks for responding to the first pass. A few more pieces of feedback:

1. Sum with padding fails in exported models — line 215. For nonempty indices and padding_idx=0, export allocates a static offset2bag of length num_indices. This kernel resizes it to zero, causing InvalidArgument before computing the result even when auxiliary outputs are unused. Match ATen’s conditional auxiliary shape and populate it when required.

2. Transposed weight inputs silently produce incorrect results Export preserves transposed runtime inputs, which the C++ runtime accepts, but this access assumes contiguous rows. Add a test for this as well.

3. Partial final offsets change max-mode results. With weight=[[2],[7]], indices [0,1], offsets [0,1], and include_last_offset=true, ATen max returns [[7]]; this kernel returns [[2]] and leaves offset2bag[1] unwritten. Match ATen’s final-bag behavior.

[JakeStevens]

see above

[yenhao-huang]

@JakeStevens, thanks for the review. Here's the fixes:

1. Sum with padding fails with static offset2bag outputs — ad0c42946b

Added fast-path selection logic to match ATen, ensuring the correct offset2bag output shape for sum with padding.

2. Incorrect results with transposed weights — 635ccc2970

Used tensor strides to read weights and write outputs instead of assuming contiguous storage.

3. Incorrect handling of include_last_offset=true — 18d0721ac3

Updated the final bag boundaries to match ATen's behavior across modes and paths.

---

Please let me know if there’s anything else.

[JakeStevens]

Thanks, the stride and partial-final-offset fixes look good. Two blockers left:

1. BF16 sum in exported models. Export sizes offset2bag with meta_embedding_bag, which doesn't count BF16 as fast path, unlike eager ATen. A BF16 sum with no padding gets a static offset2bag of length num_indices, and the kernel's resize to {0} fails. Since sum never reads it, keep and fill whatever length the caller passed (num_indices or 0) instead of predicting it. Please add a BF16 test for the static layout.
2. Frozen weights. With requires_grad=False, embedding_bag lowers to _embedding_bag_forward_only, which fails to_edge. Please add decomps[aten._embedding_bag_forward_only.default] = aten._embedding_bag.default to _default_decomposition_table() in exir/tracer.py (same pattern as _is_all_true), plus an exir test. I checked that this lowers to _embedding_bag.out and matches eager.

Nits: fill offset2bag when there are zero bags but indices are nonempty; check the weight dtype before the empty-out early return; consider looping over indices on the outside for performance.

[yenhao-huang]

@JakeStevens, thanks for the review. Here are the fixes:

1. BF16 sum with static offset2bag — c9d2c08cc5

Sum now keeps and fills whatever length the caller passed (num_indices or 0) instead of predicting it. Added a BF16 test for the static layout.

2. Frozen weights — 7d1ddae751

Added the _embedding_bag_forward_only → _embedding_bag decomposition plus an exir test.

---

Nits: zero bags now fill offset2bag with -1. The weight dtype is now checked before the early return.

Please let me know if there's anything else.

[JakeStevens]

Thanks, the BF16 static offset2bag and _embedding_bag_forward_only fixes look good. Two blockers left:

1. Long indices with Int offsets fail in exported models. The kernel requires the auxiliary outputs to have promoteTypes(indices, offsets) (op_embedding_bag.cpp:85-94). That matches eager ATen. But export plans these outputs with meta_embedding_bag, which allocates all three with offsets.new_empty(...), so they get the offsets dtype. An EmbeddingBag exported with int64 indices and int32 offsets fails with InvalidArgument on every run. index_at/set_index already handle either width, so please accept Int or Long for each output independently. Also add an exir test with int32 offsets, since test_dtypes allocates Long outputs by hand and misses this.
2. The tests fail on the ATen variant. op_embedding_bag_test is registered for ["aten", "portable"] with no is_aten guards. For Float/Half/BF16 sum, ATen resizes offset2bag to [0], so SumStaticOffset2BagAllDtypes fails there. WrongAuxiliaryDtypeRejected and DecreasingOffsetsRejected probably fail too, because ATen casts on copy and doesn't check offset monotonicity. Please wrap the ExecuTorch-only cases in ET_SKIP_IF(SupportedFeatures::get()->is_aten, ...), as op_embedding_test.cpp:222 does.

Nits:

• fills_offset2bag depends on offset2bag.numel() from the previous call. With dynamic shapes, a BF16 sum output planned at the upper bound gets shrunk to 0 on the first smaller input and never refills. Mirroring meta's fast-path predicate directly (Float/Half, unit strides, no padding) would remove the dependence on state.
• Please remove aten::_embedding_bag from UNSUPPORTED_PORTABLE_OPS in backends/test/suite/runner.py, so the existing test_embedding_bag.py suite runs against this kernel.
• Please extend test_embedding_bag_frozen_weight to cover BF16 and mean/max. Right now it only exercises Float sum, where meta already gives a size-0 offset2bag.
• The loop order is still bag → column → index, with index_at in the innermost loop. Looping over indices on the outside and accumulating whole rows would be much friendlier to the cache.

[JakeStevens]

lintrunner failing as well

[yenhao-huang]

@JakeStevens, thanks for the review. Here are the fixes:

1. Long indices with Int offsets — 1665007fb5

Each auxiliary output now accepts Int or Long independently. Added an exir test with int32 offsets (sum/mean/max), which fails with InvalidArgument before the fix.

2. Tests on the ATen variant — a62495a521

Wrapped SumStaticOffset2BagAllDtypes in ET_SKIP_IF(is_aten, ...). It's the only case that fails on ATen; the other two still throw there.

---

Lintrunner is fixed in 1665007fb5.

Please let me know if there's anything else.