Implement the portable kernel for _embedding_bag.out

exir/program/test/test_program.py

@@ -345,6 +345,37 @@ def forward(self, x):
 
         to_edge(export(M(), (torch.randn(4),))).to_executorch()
 
+    def test_embedding_bag_frozen_weight(self):
+        model = torch.nn.EmbeddingBag(10, 3, mode="sum").requires_grad_(False)
+        inputs = (torch.tensor([1, 2, 4, 5, 4, 3]), torch.tensor([0, 2, 4]))
+        executorch_module = _load_for_executorch_from_buffer(
+            to_edge(export(model, inputs)).to_executorch().buffer
+        )
+        self.assertTrue(
+            torch.allclose(
+                executorch_module.run_method("forward", inputs)[0], model(*inputs)
+            )
+        )
+
+    def test_embedding_bag_int32_offsets(self):
+        # Export allocates the auxiliary outputs with the offsets dtype.
+        inputs = (
+            torch.tensor([1, 2, 4, 5, 4, 3]),
+            torch.tensor([0, 2, 4], dtype=torch.int32),
+        )
+        for mode in ("sum", "mean", "max"):
+            with self.subTest(mode=mode):
+                model = torch.nn.EmbeddingBag(10, 3, mode=mode)
+                executorch_module = _load_for_executorch_from_buffer(
+                    to_edge(export(model, inputs)).to_executorch().buffer
+                )
+                self.assertTrue(
+                    torch.allclose(
+                        executorch_module.run_method("forward", inputs)[0],
+                        model(*inputs),
+                    )
+                )
+
     def test_antialiased_resize(self):
         class Resize(torch.nn.Module):
             def forward(self, x):

exir/tracer.py

@@ -669,6 +669,11 @@ def _default_decomposition_table(
     decomps[torch.ops.aten.empty_permuted.default] = _empty_permuted
     # torch._check_tensor_all exports to this op, which has no decomposition in core.
     decomps[torch.ops.aten._is_all_true.default] = torch.ops.aten.all.default
+    # embedding_bag with frozen weights exports to this op; it only skips
+    # saving tensors for the backward pass.
+    decomps[torch.ops.aten._embedding_bag_forward_only.default] = (
+        torch.ops.aten._embedding_bag.default
+    )
     # pyre-fixme[7]: Expected `Dict[OpOverload, typing.Callable[..., executorch.exir....
 
     never_decompose = []

kernels/aten/functions.yaml

@@ -8,6 +8,8 @@
 
 - op: _conj_physical.out
 
+- op: _embedding_bag.out
+
 - op: _fake_quantize_per_tensor_affine_cachemask_tensor_qparams.out
 
 - op: _fft_c2r.out