add test and move external_matmul to config.py

maxyanghu · pytorchmergebot · commit c83a47f5314a · 2024-09-28T21:23:32.000Z
diff --git a/test/inductor/test_external_callables.py b/test/inductor/test_external_callables.py
@@ -0,0 +1,69 @@
+# Owner(s): ["module: inductor"]
+import torch
+
+from torch._inductor import config
+
+from torch._inductor.test_case import run_tests, TestCase
+
+
+class MatMulModule(torch.nn.Module):
+    def __init__(self):
+        super()
+        self.matrix = torch.nn.Parameter(torch.eye(128, 128) * 2, requires_grad=True)
+
+    def forward(self, x):
+        return torch.matmul(x, self.matrix)
+
+
+# torch.add performs better than torch.mm and got choosed during tuning
+def matmul_cpu(a: torch.Tensor, b: torch.Tensor, out: torch.Tensor) -> None:
+    torch.add(a, b, out=out)
+
+
+def matmul_cuda(a: torch.Tensor, b: torch.Tensor, out: torch.Tensor) -> None:
+    torch.add(a, b, out=out)
+
+
+class TestInductorExternalCallable(TestCase):
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        cls._saved_config = config.save_config()
+
+    def tearDown(self):
+        super().tearDown()
+        config.load_config(self._saved_config)
+
+    def test_matmul_cpu(self):
+        # 2I + 2I == (2I)(2I)
+        x = torch.eye(128, 128) * 2
+        opt_fn = torch.compile(
+            MatMulModule(),
+            options={"max_autotune": True, "external_matmul": [matmul_cpu]},
+        )
+        opt_fn_golden = torch.compile(MatMulModule(), options={"max_autotune": True})
+        torch.testing.assert_close(
+            opt_fn(x),
+            opt_fn_golden(x),
+            msg=f"torch.compile(..., external_matmul = {matmul_cpu}) failed",
+        )
+
+    def test_matmul_cuda(self):
+        device = torch.device("cuda")
+        x = (torch.eye(128, 128) * 2).to(device=device)
+        opt_fn = torch.compile(
+            MatMulModule().to(device),
+            options={"max_autotune": True, "external_matmul": [matmul_cuda]},
+        )
+        opt_fn_golden = torch.compile(
+            MatMulModule().to(device), options={"max_autotune": True}
+        )
+        torch.testing.assert_close(
+            opt_fn(x),
+            opt_fn_golden(x),
+            msg=f"torch.compile(..., external_matmul = {matmul_cuda}) failed",
+        )
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/torch/_inductor/config.py b/torch/_inductor/config.py
@@ -1260,6 +1260,9 @@ class trace:
     "compile_threads",
 ]
 
+# External callable for matmul tuning candidates
+external_matmul: List[Callable[[torch.Tensor, torch.Tensor, torch.Tensor], None]] = []
+
 if TYPE_CHECKING:
     from torch.utils._config_typing import *  # noqa: F401, F403
 
diff --git a/torch/_inductor/external_callable.py b/torch/_inductor/external_callable.py
diff --git a/torch/_inductor/kernel/mm.py b/torch/_inductor/kernel/mm.py
@@ -245,7 +245,7 @@ def tuned_mm(mat1, mat2, *, layout=None):
         log.warning("No choices for GEMM, using ATen backend as fallback")
         return aten_mm.bind((mat1, mat2), aten_layout).output_node()
 
-    from ..external_callable import external_matmul
+    from ..config import external_matmul
 
     for k in external_matmul:
         choices.append(ExternKernelChoice(k).bind((mat1, mat2), layout))
