[Inductor] Support calling user defined kernels with different type of arguments

oulgen · oulgen · commit 0b730644f04c · 2023-10-24T09:29:06.000-07:00
[ghstack-poisoned]
diff --git a/test/dynamo/test_functions.py b/test/dynamo/test_functions.py
@@ -1831,6 +1831,36 @@ def call_triton_add(
         self.assertTrue("aten.copy" not in code)
         self.assertTrue("aten.clone" not in code)
 
+    @requires_cuda()
+    @requires_triton()
+    @common_utils.parametrize("grad", [False, True])
+    def test_triton_kernel_multi_kernel(self, grad):
+        @torch.compile
+        def call_triton(
+            x: torch.Tensor,
+            y: torch.Tensor,
+            xi: torch.Tensor,
+            yi: torch.Tensor,
+        ):
+            output = torch.zeros_like(x, requires_grad=grad)
+            outputi = torch.zeros_like(xi)
+            n_elements = output.numel()
+
+            grid = (x.numel(),)
+            add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=16)
+            add_kernel[grid](xi, yi, outputi, n_elements, BLOCK_SIZE=16)
+
+            return (output, outputi)
+
+        t1 = torch.rand(5, device="cuda", requires_grad=grad)
+        t2 = torch.rand(5, device="cuda", requires_grad=grad)
+        t1i = torch.randint(0, 5, (5,), device="cuda")
+        t2i = torch.randint(0, 5, (5,), device="cuda")
+
+        (result, resulti) = call_triton(t1, t2, t1i, t2i)
+        self.assertEqual(t1 + t2, result)
+        self.assertEqual(t1i + t2i, resulti)
+
     @requires_cuda()
     @requires_triton()
     @common_utils.parametrize("grad", [False, True])
diff --git a/torch/_inductor/codegen/wrapper.py b/torch/_inductor/codegen/wrapper.py
@@ -307,6 +307,7 @@ def __init__(self):
         self.supports_intermediate_hooks = True
         self.expr_printer = pexpr
         self.cached_thread_locals = set()
+        self.user_defined_kernel_count = 0
 
         self.write_header()
         self.write_prefix()
@@ -752,11 +753,15 @@ def define_kernel(
         metadata_comment = f"{metadata}\n" if metadata else ""
         self.header.splice(f"\n\n{metadata_comment}{name} = {kernel}")
 
-    def define_user_defined_triton_kernel(self, kernel, kwargs):
-        name = kernel.__name__
+    def get_unique_kernel_name(self, name: str) -> str:
+        new_name = f"{name}_{self.user_defined_kernel_count}"
+        self.user_defined_kernel_count += 1
+        return new_name
 
+    def define_user_defined_triton_kernel(self, name, kernel, kwargs):
+        original_name = kernel.__name__
         compile_wrapper = IndentedBuffer()
-        compile_wrapper.writeline(f"async_compile.triton({name!r}, '''")
+        compile_wrapper.writeline(f"async_compile.triton({original_name!r}, '''")
 
         compile_wrapper.splice(
             """
diff --git a/torch/_inductor/ir.py b/torch/_inductor/ir.py
@@ -3745,14 +3745,15 @@ def codegen(self, wrapper):
         from torch._higher_order_ops.triton_kernel_wrap import kernel_side_table
 
         kernel = kernel_side_table.get_kernel(self.kernel_idx)
+        new_name = wrapper.get_unique_kernel_name(kernel.__name__)
 
         self.codegen_comment(wrapper)
         wrapper.generate_user_defined_triton_kernel(
-            kernel.__name__,
+            new_name,
             self.grid,
             self.codegen_kwargs(),
         )
-        wrapper.define_user_defined_triton_kernel(kernel, self.kwargs)
+        wrapper.define_user_defined_triton_kernel(new_name, kernel, self.kwargs)
 
     def should_allocate(self):
         return False
@@ -3765,6 +3766,9 @@ def has_side_effects(self):
     def get_unbacked_symbol_defs(self):
         return {}
 
+    def get_mutation_names(self):
+        return [t.get_name() for t in self.inputs]
+
     def __init__(self, *, kernel_idx, grid, kernel_args):
         inputs = []
         kwargs = dict()
@@ -3776,6 +3780,7 @@ def __init__(self, *, kernel_idx, grid, kernel_args):
                 t = InputsKernel.unwrap_storage_for_input(self.realize_input(v))
                 inputs.append(t)
                 kwargs[k] = t
+                V.graph.mark_buffer_mutated(t.get_name())
             else:
                 constant_args.append(v)
                 kwargs[k] = v
