fix emb (#6)

Author: zhuhaozhe

intel_pytorch_extension_py/ops/embeddingbag.py

@@ -1,15 +1,18 @@
 import torch
-from torch import nn
-from torch.autograd import Function
 import _torch_ipex as core
+import warnings
 
-# # extension for BF16 fast path only
+torch_embedding_bag = torch.embedding_bag
 
-
-def embeddingbag(weights, indices, offsets, scale_grad_by_freq, mode, sparse, per_sample_weights, include_last_offset):
-    ret = torch.ops.torch_ipex.embedding_bag(weights, indices, offsets, scale_grad_by_freq, mode, sparse, per_sample_weights, include_last_offset)
-    if len(ret)==1:
-        ret += [torch.Tensor(), torch.Tensor(), torch.Tensor()]
+def embeddingbag(weights, indices, offsets, scale_grad_by_freq, mode, sparse, per_sample_weights, include_last_offset, padding_idx):
+    if core.embedding_bag_fast_path_sum(weights, per_sample_weights, mode, padding_idx):
+        ret = torch.ops.torch_ipex.embedding_bag(weights, indices, offsets, sparse, include_last_offset)
+        # torch.embedding_bag expected 4 Tensor returned
+        # here we only return 1 tensor since the other three tensors are not needed in our fast path
+        ret = [ret, torch.Tensor(), torch.Tensor(), torch.Tensor()]
+    else:
+        warnings.warn('Fallback to torch.embedding bag')
+        ret = torch_embedding_bag(weights, indices, offsets, scale_grad_by_freq, mode, sparse, per_sample_weights, include_last_offset, padding_idx)
     return ret
 
 torch.embedding_bag = embeddingbag

tests/cpu/test_emb.py

@@ -1,45 +1,48 @@
 import torch
 import torch.nn as nn
-import intel_pytorch_extension as ipex
 import unittest
 import copy
 from common_utils import TestCase
 
 class TestEMB(TestCase):
-    def test_emb(self):
+    def _test_emb(self, mode):
         #E = nn.EmbeddingBag(10, 5, mode="sum", sparse=True)
-        cpu_emb = nn.EmbeddingBag(10, 3, mode='sum', sparse=True)
-        dpcpp_emb = copy.deepcopy(cpu_emb)
-        bf16_emb = copy.deepcopy(cpu_emb).bfloat16()
+        aten_emb = nn.EmbeddingBag(10, 3, mode=mode, sparse=True)
+        ipex_emb = copy.deepcopy(aten_emb)
+        bf16_emb = copy.deepcopy(aten_emb).bfloat16()
         # a batch of 2 samples of 4 indices each
-        cpu_input = torch.LongTensor([1,2,4,5,4,3,2,9])
-        dpcpp_input = cpu_input.clone().detach()
-
-        cpu_offsets = torch.LongTensor([0,1,2,3,4,5,6,7])
-        dpcpp_offsets = cpu_offsets.clone().detach()
-
-        cpu_out = cpu_emb(cpu_input, cpu_offsets)
-
-        #torch.embedding_bag = ipex.embeddingbag
-        dpcpp_out = dpcpp_emb(dpcpp_input, dpcpp_offsets)
-        bf16_out = bf16_emb(dpcpp_input, dpcpp_offsets)
-
-        self.assertEqual(cpu_out, dpcpp_out.to('cpu'))
-        self.assertEqual(cpu_out, bf16_out.to('cpu').float(), 0.01)
-
-        cpu_out.mean().backward()
-        dpcpp_out.mean().backward()
-        bf16_out.float().mean().backward()
-
-        self.assertEqual(cpu_emb.weight.grad.data._nnz(), dpcpp_emb.weight.grad.data._nnz())
-        self.assertEqual(cpu_emb.weight.grad.data.sparse_dim(), dpcpp_emb.weight.grad.data.sparse_dim())
-        self.assertEqual(cpu_emb.weight.grad.data.dense_dim(), dpcpp_emb.weight.grad.data.dense_dim())
-        self.assertEqual(cpu_emb.weight.grad.data.is_coalesced(), dpcpp_emb.weight.grad.data.is_coalesced())
-        self.assertEqual(cpu_emb.weight.grad.data._indices(), dpcpp_emb.weight.grad.data._indices().to('cpu'))
-        self.assertEqual(cpu_emb.weight.grad.data._values(), dpcpp_emb.weight.grad.data._values().to('cpu'))
-
-        self.assertEqual(cpu_emb.weight.grad.data._values(), dpcpp_emb.weight.grad.data._values().to('cpu'), 0.01)
-        self.assertEqual(bf16_emb.weight.grad.data._values().dtype, torch.bfloat16)
+        input = torch.LongTensor([1,2,4,5,4,3,2,9])
+        offsets = torch.LongTensor([0,1,2,3,4,5,6,7])
+        # aten path
+        aten_out = aten_emb(input, offsets)
+        aten_out.mean().backward()
+
+        # ipex fast path (both fp32/bf16)
+        import intel_pytorch_extension
+        ipex_out = ipex_emb(input, offsets)
+        ipex_out.mean().backward()
+        if mode == 'sum':
+            bf16_out = bf16_emb(input, offsets)
+            bf16_out.mean().backward()
+            self.assertEqual(aten_out, bf16_out.float(), 0.01)
+            self.assertEqual(bf16_emb.weight.grad.data._values().dtype, torch.bfloat16)
+        del(intel_pytorch_extension)
+
+        self.assertEqual(aten_out, ipex_out)
+
+        self.assertEqual(aten_emb.weight.grad.data._nnz(), ipex_emb.weight.grad.data._nnz())
+        self.assertEqual(aten_emb.weight.grad.data.sparse_dim(), ipex_emb.weight.grad.data.sparse_dim())
+        self.assertEqual(aten_emb.weight.grad.data.dense_dim(), ipex_emb.weight.grad.data.dense_dim())
+        self.assertEqual(aten_emb.weight.grad.data.is_coalesced(), ipex_emb.weight.grad.data.is_coalesced())
+        self.assertEqual(aten_emb.weight.grad.data._indices(), ipex_emb.weight.grad.data._indices())
+        self.assertEqual(aten_emb.weight.grad.data._values(), ipex_emb.weight.grad.data._values())
+        self.assertEqual(aten_emb.weight.grad.data._values(), ipex_emb.weight.grad.data._values(), 0.01)
+
+    def test_emb_fast_path(self):
+        self._test_emb(mode='mean')
+
+    def test_emb_fallback_path(self):
+        self._test_emb(mode='sum')
 
 if __name__ == '__main__':
     test = unittest.main()

torch_ipex/csrc/cpu/ExtendOPs.h

@@ -39,11 +39,19 @@ class AtenIpexTypeExt {
   static std::vector<at::Tensor> interaction_backward(const at::Tensor & grad_out, 
                                                                      const std::vector<at::Tensor> & input);
 
-  static std::vector<at::Tensor> embedding_bag(
-      const at::Tensor &weight, const at::Tensor &indices,
-      const at::Tensor &offsets, bool scale_grad_by_freq, int64_t mode,
-      bool sparse, const c10::optional<at::Tensor> &per_sample_weights,
+  static at::Tensor embedding_bag(
+      const at::Tensor &weight,
+      const at::Tensor &indices,
+      const at::Tensor &offsets,
+      bool sparse,
       bool include_last_offset);
+
+  static bool embedding_bag_fast_path_sum(
+      const at::Tensor weight, 
+      const c10::optional<at::Tensor> per_sample_weights, 
+      int64_t mode, 
+      const c10::optional<int64_t> padding_idx);
+
 };
 
 }  // namespace torch_ipex