add meta backend for EmbeddingBag (#1525) (#1568)

* add meta backend for EmbeddingBag

* add UT

* modify UT

* fix UT

Author: jiayisunx

csrc/cpu/aten/EmbeddingBag.cpp

@@ -123,17 +123,6 @@ at::Tensor dil_qembeddingbag(
 } // namespace cpu
 } // namespace torch_ipex
 
-namespace {
-TORCH_LIBRARY_FRAGMENT(torch_ipex, m) {
-  m.def(
-      torch::schema(
-          "torch_ipex::embedding_bag(Tensor weight, Tensor indices, Tensor "
-          "offsets, bool sparse, bool include_last_offset) -> Tensor",
-          c10::AliasAnalysisKind::PURE_FUNCTION),
-      torch_ipex::embedding_bag);
-}
-} // namespace
-
 namespace torch_ipex {
 namespace autocast {
 
@@ -156,10 +145,6 @@ at::Tensor embedding_bag(
   return op.call(casted_weight, indices, offsets, sparse, include_last_offset);
 }
 
-TORCH_LIBRARY_IMPL(torch_ipex, AutocastCPU, m) {
-  m.impl("embedding_bag", torch_ipex::autocast::embedding_bag);
-}
-
 } // namespace autocast
 } // namespace torch_ipex
 
@@ -179,4 +164,36 @@ at::Tensor embedding_bag(
       weight, indices, offsets, sparse, include_last_offset);
 }
 
+at::Tensor embedding_bag_meta(
+    const at::Tensor& weight,
+    const at::Tensor& indices,
+    const at::Tensor& offsets,
+    bool sparse,
+    bool include_last_offset) {
+  auto num_bags = offsets.sym_size(0);
+  if (indices.dim() == 2) {
+    num_bags = indices.sym_size(0);
+  }
+  c10::SymDimVector output_size(2);
+  output_size[0] = num_bags;
+  output_size[1] = weight.sym_size(1);
+  auto output = at::empty_symint(output_size, weight.options());
+  return output;
+}
+
 } // namespace torch_ipex
+
+namespace {
+TORCH_LIBRARY_FRAGMENT(torch_ipex, m) {
+  m.def(
+      "embedding_bag(Tensor weight, Tensor indices, Tensor "
+      "offsets, bool sparse, bool include_last_offset) -> Tensor");
+  m.impl("embedding_bag", c10::DispatchKey::CPU, torch_ipex::embedding_bag);
+  m.impl(
+      "embedding_bag", c10::DispatchKey::Meta, torch_ipex::embedding_bag_meta);
+  m.impl(
+      "embedding_bag",
+      c10::DispatchKey::AutocastCPU,
+      torch_ipex::autocast::embedding_bag);
+}
+} // namespace

csrc/cpu/aten/EmbeddingBag.h

@@ -11,6 +11,13 @@ at::Tensor embedding_bag(
     bool sparse,
     bool include_last_offset);
 
+at::Tensor embedding_bag_meta(
+    const at::Tensor& weight,
+    const at::Tensor& indices,
+    const at::Tensor& offsets,
+    bool sparse,
+    bool include_last_offset);
+
 } // namespace torch_ipex
 
 namespace torch_ipex {

tests/cpu/test_emb.py

@@ -9,6 +9,14 @@
 ipex_emb_fn = ipex.nn.functional._embeddingbag._embeddingbag
 aten_emb_fn = ipex.nn.functional._embeddingbag.torch_embedding_bag
 
+class Embeddingbag(torch.nn.Module):
+    def __init__(self):
+        super(Embeddingbag, self).__init__()
+        self.embeddingbag = nn.EmbeddingBag(10, 3, mode='sum', sparse=True)
+
+    def forward(self, input, offsets):
+        return self.embeddingbag(input, offsets)
+
 class TestEMB(TestCase):
 
     def _test_emb(
@@ -96,5 +104,19 @@ def test_emb_jit_scriptable(self):
         out = script_emb(input, offsets)
         self.assertEqual(out, ref_out)
 
+    def test_emb_torch_compile(self):
+        emb = Embeddingbag().eval()
+        input = torch.LongTensor([1,2,4,5,4,3,2,9])
+        offsets = torch.LongTensor([0,1,2,3,4,5,6,7])
+        # TODO: add dynamic tests when 'ipex' backend supports it.
+        for dtype, backend, dynamic in itertools.product([torch.float32, torch.bfloat16], ['ipex', 'inductor'], [False]):
+            torch._dynamo.reset()
+            emb_torchcompile = torch.compile(emb, backend=backend, dynamic=dynamic)
+            with torch.cpu.amp.autocast(enabled=(dtype==torch.bfloat16)), torch.no_grad():
+                y0 = emb(input, offsets)
+                y1 = emb_torchcompile(input, offsets)
+            self.assertEqual(y0, y1)
+            self.assertEqual(y1.dtype, dtype)
+
 if __name__ == '__main__':
     test = unittest.main()