dedup Externchoice

maxyanghu · pytorchmergebot · commit cca9969aa5fe · 2024-09-28T21:23:32.000Z
diff --git a/test/inductor/test_external_callables.py b/test/inductor/test_external_callables.py
@@ -1,10 +1,12 @@
 # Owner(s): ["module: inductor"]
 import torch
+import unittest
 
 from torch._inductor import config
-
 from torch._inductor.test_case import run_tests, TestCase
 
+from torch.testing._internal.common_cuda import TEST_CUDA
+
 
 class MatMulModule(torch.nn.Module):
     def __init__(self):
@@ -48,6 +50,7 @@ def test_matmul_cpu(self):
             msg=f"torch.compile(..., external_matmul = {matmul_cpu}) failed",
         )
 
+    @unittest.skipIf(not TEST_CUDA, "CUDA not found")   
     def test_matmul_cuda(self):
         device = torch.device("cuda")
         x = (torch.eye(128, 128) * 2).to(device=device)
diff --git a/torch/_inductor/kernel/mm.py b/torch/_inductor/kernel/mm.py
@@ -122,6 +122,11 @@
 """,
 )
 
+# prevent duplication registration of extern functions
+@functools.lru_cache(None)
+def lazy_register_extern_choice(fn):
+   return ExternKernelChoice(fn)
+
 aten_mm = ExternKernelChoice(torch.mm, "at::mm_out")
 
 aten_addmm = ExternKernelChoice(
@@ -248,7 +253,7 @@ def tuned_mm(mat1, mat2, *, layout=None):
     from ..config import external_matmul
 
     for k in external_matmul:
-        choices.append(ExternKernelChoice(k).bind((mat1, mat2), layout))
+        choices.append(lazy_register_extern_choice(k).bind((mat1, mat2), layout))
 
     try:
         return autotune_select_algorithm(name, choices, [mat1, mat2], layout)
