set AliasAnalysisKind of embedding_bag and interaction to PURE_FUNCTI…

…ON (#163)
intel · Aug 17, 2021 · 4201cef · 4201cef
1 parent 8d52321
commit 4201cef
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 9 deletions.
diff --git a/torch_ipex/csrc/cpu/embeddingbag.cpp b/torch_ipex/csrc/cpu/embeddingbag.cpp
@@ -494,9 +494,13 @@ at::Tensor AtenIpexJITDev::dil_qembeddingbag(
 }  // namespace torch_ipex
 
 namespace {
-static auto dispatch =
-    torch::RegisterOperators()
-        .op("torch_ipex::embedding_bag", &torch_ipex::AtenIpexTypeExt::embedding_bag);
+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::AtenIpexTypeExt::embedding_bag);
+}
 }
 
 namespace torch_ipex {

diff --git a/torch_ipex/csrc/cpu/interaction.cpp b/torch_ipex/csrc/cpu/interaction.cpp
@@ -422,12 +422,16 @@ at::Tensor AtenIpexJITDev::dil_qinteraction(const std::vector<at::Tensor> input,
 } // namespace torch_ipex
 
 namespace {
-static auto dispatch =
-    torch::RegisterOperators()
-        .op("torch_ipex::interaction_forward",
-            &torch_ipex::AtenIpexTypeExt::interaction_forward)
-        .op("torch_ipex::interaction_backward",
-            &torch_ipex::AtenIpexTypeExt::interaction_backward);
+TORCH_LIBRARY_FRAGMENT(torch_ipex, m) {
+  m.def(
+      torch::schema("torch_ipex::interaction_forward(Tensor[] input) -> Tensor",
+                    c10::AliasAnalysisKind::PURE_FUNCTION),
+      torch_ipex::AtenIpexTypeExt::interaction_forward);
+  m.def(torch::schema("torch_ipex::interaction_backward(Tensor grad_out, "
+                      "Tensor[] input) -> Tensor[]",
+                      c10::AliasAnalysisKind::PURE_FUNCTION),
+        torch_ipex::AtenIpexTypeExt::interaction_backward);
+}
 }
 
 namespace torch_ipex {