From 357c4d9cc40f461fe7b95af6524bbf8b2ecd9f74 Mon Sep 17 00:00:00 2001
From: Jiewen Tan <jwtan@fb.com>
Date: Wed, 7 Jul 2021 13:15:40 -0700
Subject: [PATCH] Add a test case for findDanglingImpls (#61104)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/61104

This patch added a new test case for findDanglingImpls. The test case introduces a C++ extension which has a dangling impl such that findDanglingImpls can find it and output its information.

Test Plan:
python test/test_dispatch.py TestDispatch.test_find_dangling_impls_ext

Imported from OSS

Reviewed By: ezyang

Differential Revision: D29512520

fbshipit-source-id: 6883fb8f065f2c0ae0e7a1adf6fd298591497e2b
---
 .../dangling_impl_extension.cpp               | 11 +++++
 test/test_dispatch.py                         | 40 ++++++++++++++-----
 2 files changed, 42 insertions(+), 9 deletions(-)
 create mode 100644 test/cpp_extensions/dangling_impl_extension.cpp

diff --git a/test/cpp_extensions/dangling_impl_extension.cpp b/test/cpp_extensions/dangling_impl_extension.cpp
new file mode 100644
index 00000000000..ea43eab3313
--- /dev/null
+++ b/test/cpp_extensions/dangling_impl_extension.cpp
@@ -0,0 +1,11 @@
+#include <torch/extension.h>
+
+void foo() { }
+
+TORCH_LIBRARY_IMPL(__test, CPU, m) {
+  m.impl("foo", foo);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("bar", foo);
+}
diff --git a/test/test_dispatch.py b/test/test_dispatch.py
index d22eaf96721..62cd5d64a1a 100644
--- a/test/test_dispatch.py
+++ b/test/test_dispatch.py
@@ -4,7 +4,9 @@ from torch._python_dispatcher import PythonDispatcher
 
 from collections import namedtuple
 import itertools
+import os
 import re
+import torch.utils.cpp_extension
 
 # TODO: Expand the dispatcher API to be a generic API for interfacing with
 # the dispatcher from Python!
@@ -755,6 +757,35 @@ CompositeImplicitAutograd[alias] (inactive): fn1 :: (Tensor _0) -> (Tensor _0) [
 '''
         )
 
+    def test_find_dangling_impls(self):
+        dangling_impls = C._dispatch_find_dangling_impls()
+        self.assertEqual(
+            0,
+            len(dangling_impls),
+            msg=f"Expect zero dangling impls, but found: {dangling_impls}"
+        )
+
+    def test_find_dangling_impls_ext(self):
+        extension_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'cpp_extensions', 'dangling_impl_extension.cpp')
+        module = torch.utils.cpp_extension.load(
+            name="dangling_impl_extension",
+            sources=[
+                extension_path,
+            ],
+            extra_cflags=["-g"],
+            verbose=True,
+        )
+
+        impls = C._dispatch_find_dangling_impls()
+        self.assertEqual(1, len(impls))
+        self.assertEqual(
+            '''\
+name: __test::foo
+schema: (none)
+CPU: registered at {}:5 :: () -> () [ boxed unboxed ]
+'''.format(extension_path),
+            impls[0])
+
 class TestPythonDispatcher(TestCase):
     def test_basic(self):
         dispatcher = PythonDispatcher()
@@ -886,14 +917,5 @@ CompositeImplicitAutograd[alias] fn_CompositeImplicitAutograd
                 r"Registration to both CompositeImplicitAutograd and CompositeExplicitAutograd is not allowed"):
             dispatcher.register(["CompositeExplicitAutograd", "CompositeImplicitAutograd"])
 
-    # TODO(jwtan): Use a C++ extension, like msnpu, to introduce a dangling impl and examine the output.
-    def test_find_dangling_impls(self):
-        dangling_impls = C._dispatch_find_dangling_impls()
-        self.assertEqual(
-            0,
-            len(dangling_impls),
-            msg=f"Expect zero dangling impls, but found: {dangling_impls}"
-        )
-
 if __name__ == '__main__':
     run_tests()