[JakeStevens]

Summary:

Add an optimized _embedding_bag.out with FP32 weights, int32/int64 indices and offsets, sum/mean/max reductions, padding, per-sample weights for sum, and all four functional outputs. The CPU provider now routes _embedding_bag.default to the linked _embedding_bag.out, which works with the portable kernel too.

Validation and output resizing move into a new shared kernels/portable/cpu/util/embedding_bag_util.cpp, which the portable kernel now also uses. The kernel takes default-dim-order FP32 weights and outputs, with the index outputs typed like offsets as export allocates them, and falls back to the portable kernel for anything else. Bags are split across the threadpool. Sum and mean accumulate rows with at::vec over whole vectors and a scalar loop for each row's remainder; max is a scalar loop that preserves ATen's first-winner indices and NaN behavior. The kernel matches ATen's contiguous FP32 handling of a final offset below indices.numel(). No temporary tensor allocation is required.

Profiling: warm FP32 sum, 4096 × 128 weight table, 256 bags × 32 indices. The 1-thread gain compares optimized with the single-threaded portable kernel; scaling compares the optimized kernel at 1 and 4 threads. These are kernel-call measurements, not end-to-end model latency.

Surface	1-thread gain vs portable	Optimized 1→4-thread scaling
x86 · AMD EPYC	7.54×	3.04×
Mac · Apple Silicon	4.14×	2.35×
iPhone 17	5.90×	1.58×
Pixel 10	4.90×	No benefit established

Pixel 10 had measurement issues, with both baseline and test arm fluctuating.

Authored with AI assistance using Codex and Claude Code.

Differential Revision: D123412521

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (4 Unrelated Failures)

As of commit 47dbaba with merge base a70d32f ():

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / unittest / linux / linux-job (gh) (trunk failure)
  examples/models/muse-glimmer/tests/test_cuda_sequence_export_inputs.py::SequenceExportInputsTest::test_small_logits_example_preserves_token_and_embedding_capacity
• Test CPU Backend / test-cpu / test-backend-linux (cpu_fp32, models) / linux-job (gh) (trunk failure)
  /__w/executorch/executorch/pytorch/executorch/backends/cpu/runtime/providers/xnnpack/XNNPACKProvider.cpp:1128:16: error: call of overloaded ‘clamp_bound(const ptn::Argument&, const float&, float&)’ is ambiguous
• Test CPU Backend / test-cpu / test-backend-linux (cpu_fp32, operators) / linux-job (gh) (trunk failure)
  /__w/executorch/executorch/pytorch/executorch/backends/cpu/runtime/providers/xnnpack/XNNPACKProvider.cpp:1128:16: error: call of overloaded ‘clamp_bound(const ptn::Argument&, const float&, float&)’ is ambiguous
• Test Native Engines / native-cpu (gh) (trunk failure)
  /__w/executorch/executorch/backends/cpu/runtime/providers/xnnpack/XNNPACKProvider.cpp:1128:16: error: call of overloaded 'clamp_bound(const ptn::Argument&, const float&, float&)' is ambiguous

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

[meta-codesync]

@JakeStevens has exported this pull request. If you are a Meta employee, you can view the originating Diff in D123412521.

[JakeStevens]

@claude review effort:max

[claude]

Claude finished @JakeStevens's task in 5m 38s —— View job

---

I found no blockers or correctness bugs; the remaining issues are minor (logic copied from the portable kernel, gaps in test coverage, and some small performance and docs items), so I recommend Approve.

Detailed review

PR Review: #24851

Summary

This PR adds an optimized FP32 _embedding_bag.out kernel that splits bags across the threadpool, and routes _embedding_bag through the CPU backend's reference provider. Three independent reviewers traced every fast-path branch against kernels/portable/cpu/op_embedding_bag.cpp. These all match the portable kernel:

• the offsets and indices checks
• when offset2bag gets filled
• the cut at a partial final offset
• the mean divisor
• bag_size and max_indices
• first-winner and NaN behaviour in max mode

There are no thread-safety problems: workers only write bag-disjoint ranges, because offsets are validated as nondecreasing. All findings below are minor. Nothing was built or run; everything comes from reading the code.

Code Quality

• Matching rules copied from portable (kernels/optimized/cpu/op_embedding_bag.cpp:48-105, :127-151, :180-183)
  • Problem: check_embedding_bag_args, the offset/index checks, the fills_offset2bag rule (:95) and the reduce_end cut re-implement logic already in the portable kernel. The fill rule is written in a different but equivalent form (padding_idx >= 0 rather than the portable fast-path flag), and the error messages differ.
  • Risk: if the portable rules change later, the two drift apart, and output shapes and errors will depend on whether the fast path or the fallback runs.
  • Fix: move these rules into a shared helper under kernels/portable/cpu/util/, as other optimized ops do.
  • Reported by all three reviewers.
• Outputs are resized before inputs are validated (op_embedding_bag.cpp:303-317 vs :127-151)
  • Problem: on the fast path, offsets and indices are checked inside embedding_bag_kernel, after resize_embedding_bag_outputs has already run. The portable kernel validates first.
  • Trigger: bad input such as indices=[5] with a 2-row weight, or offsets=[0,5] with 4 indices. A third case is Short indices: is_fast_path never checks the indices dtype, so the call fails in ET_SWITCH_TWO_TYPES only after the resize.
  • Effect: the call still fails, but the caller's dynamic-bound outputs are left reshaped, which the fallback path never does.
  • Fix: run the offset/index checks before resizing, and add the indices dtype (Long or Int) to is_fast_path.
  • Reported by all three reviewers.
• Stale README (backends/cpu/runtime/providers/reference/README.md:7): it still says "linked portable kernel set". The PR changed the matching comment in ReferenceProvider.cpp:69 to "linked kernel".

Testing

• Weighted sum is never tested at vector widths in the C++ op tests
  • The per_sample_weights path through map2 with Vec(scale) (op_embedding_bag.cpp:199-206) is only tested at width 2 (test_weighted_sum_mixed_index_types).
  • ReductionsAcrossVectorWidths never passes weights.
  • Full-vector weighted sums are only reached by the Python test_embedding_bag_all_outputs (width 17).
  • Fix: add a weighted case to test_reductions, run at widths ≥ 8.
• ManyBagsAllModes covers little of the multi-task logic
  • Coverage:
    • With about 162 reads per bag, the grain is 202, so 512 bags make at most 3 tasks.
    • It only runs with include_last_offset=false, padding_idx=-1 and no weights.
    • So the boundary logic on the last bag (end = num_indices, and reduce_end = offsets[bag+1] when the final offset is below numel) never runs at a task boundary.
  • Single-thread runners: nothing forces more than one threadpool thread, so a runner with a single-threaded pool only exercises the serial path.
  • Fix: add a case with include_last_offset=true, a partial final offset, padding and weights, and enough bags to split.
• The int32 error paths are never run. expect_invalid (op_embedding_bag_test.cpp:140) always builds Long indices and Long offsets, so the range checks are never run on bad input for int32_t indices or int32_t offsets.
• Python test relaxes the dtype check on all four outputs (backends/cpu/test/test_runtime.py, the assert_close(..., check_dtype=False) call). Only the three index outputs need it. As written, a regression that gives the float out the wrong dtype would still pass.

Performance

• Unweighted sum/mean does an unneeded multiply (op_embedding_bag.cpp:200-206). With scales == nullptr, the lambda still computes x + y * Vec(1.0f) for every element on the benchmarked path.
  • Fix: split out a no-scale branch (x + y), and use fmadd for the weighted one.
• Serial work before parallel_for limits scaling (op_embedding_bag.cpp:161-162).
  • Problem: out (and max_indices) are zero-filled over the whole tensor on the calling thread, and then every bag's row is rewritten in the workers anyway. Together with the serial index-validation pass, this caps the 1→4-thread scaling reported in the description (1.58× on iPhone).
  • Fix: zero out_row and winners per bag inside the worker.
• Loop-invariant check in the inner loop (op_embedding_bag.cpp:187): offset2bag.numel() != 0 is re-checked for every index. Compute it once outside the loop.

Recommendation

Approve

Nothing blocks merging. The two worth addressing before landing are the validate-before-resize ordering and the weighted-sum and multi-task test gaps. The rest are optional cleanups.

[JakeStevens]

@claude review effort=max

[claude]

Claude encountered an error after 3m 16s —— View job

---

I'll analyze this and get back to you.