Add external registration API for matmul candidates

maxyanghu · pytorchmergebot · commit 41145f7ad1c0 · 2024-09-28T21:23:32.000Z
diff --git a/torch/_inductor/external_callable.py b/torch/_inductor/external_callable.py
@@ -0,0 +1,11 @@
+import inspect
+
+external_matmul = []
+
+def register_external_matmul_call(func):
+    if not callable(func):
+        raise Exception(f'{func} is not a callable')
+    params = inspect.signature(func).parameters
+    if len(params) != 3:
+        raise Exception(f'required 3 params but got {len(params)}')
+    external_matmul.append(func)
diff --git a/torch/_inductor/kernel/mm.py b/torch/_inductor/kernel/mm.py
@@ -245,6 +245,10 @@ 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
+    for k in external_matmul:
+        choices.append(ExternKernelChoice(k).bind((mat1, mat2), layout))
+
     try:
         return autotune_select_algorithm(name, choices, [mat1, mat2], layout)
     except NoValidChoicesError:
