[inductor] Support multiple symbolic numel expr in CudaWrapperCodeGen (pytorch#102093)

Summary: Add a set to avoid generating extra `auto` when seeing the
symbolic numel expression for the second time.

Pull Request resolved: pytorch#102093
Approved by: https://github.com/jansel

Author: desertfire

test/inductor/test_cpp_wrapper.py

@@ -202,7 +202,7 @@ class BaseTest(NamedTuple):
         BaseTest("test_embedding_bag"),  # test default FallbackKernel
         BaseTest("test_index_put_deterministic_fallback"),
         BaseTest("test_linear1"),
-        # BaseTest("test_linear2"),
+        BaseTest("test_linear2"),
         BaseTest("test_mm_views"),
         BaseTest("test_multi_device"),
         BaseTest("test_profiler_mark_wrapper_call"),

torch/_inductor/codegen/triton.py

@@ -1665,7 +1665,8 @@ def dense_size_str(self):
                 sizes.append("1")
         return f"[{', '.join(sizes)}]"
 
-    def call_kernel(self, code, name: str):
+    def call_kernel(self, name: str):
+        wrapper = V.graph.wrapper_code
         _, call_args, _ = self.args.python_argdefs()
         # dynamo wraps unspec variable as 0d CPU tensor, need convert to scalar
         for i in range(len(call_args)):
@@ -1677,26 +1678,31 @@ def call_kernel(self, code, name: str):
             if isinstance(tree.numel, (sympy.Integer, sympy.Symbol)):
                 expr = tree.numel
             else:
+                expr = f"{name}_{tree.prefix}numel"
+                # TODO(voz): Tragic. This should at the very least be a util to slapp on declare and ending.
+                # The real fix here is to revisit our cross language calling convention.
+                if expr not in wrapper.kenel_numel_expr:
+                    wrapper.kenel_numel_expr.add(expr)
+                    wrapper.writeline(
+                        f"{wrapper.declare}{expr} = {pexpr(tree.numel)}{wrapper.ending}"
+                    )
+                else:
+                    wrapper.writeline(f"{expr} = {pexpr(tree.numel)}{wrapper.ending}")
                 # We can get symbolic expressions here, like s0*64
                 # It is fine to have them here, but we need to handle them correctly as their own type
                 # This is tricky to do, so we wrap in a custom type, distinct from scalars, but also from sympy*
                 # scalars as well.
                 # This is handled in `generate_args_decl` which has a correct comment of: TODO: only works for
                 # constant now, need type info. I agree, this needs type info, and while this is not true type info
                 # it suffices as a type hint for the purposes of producing the correct code for this type.
-                expr = SymbolicCallArg(f"{name}_{tree.prefix}numel")
-                # TODO(voz): Tragic. This should at the very least be a util to slapp on declare and ending.
-                # The real fix here is to revisit our cross language calling convention.
-                code.writeline(
-                    f"{code.declare}{expr} = {pexpr(tree.numel)}{code.ending}"
-                )
+                expr = SymbolicCallArg(expr)
 
             if tree.prefix != "r" or self.inside_reduction:
                 call_args.append(expr)
             if tree.prefix != "r":
                 grid.append(expr)
 
-        code.generate_kernel_call(
+        wrapper.generate_kernel_call(
             name,
             call_args,
             grid,
@@ -1985,7 +1991,7 @@ def codegen_node_schedule(self, node_schedule, numel, reduction_numel):
         src_code = kernel.codegen_kernel()
         kernel_name = self.define_kernel(src_code, node_schedule)
 
-        kernel.call_kernel(V.graph.wrapper_code, kernel_name)
+        kernel.call_kernel(kernel_name)
 
         if (
             V.graph.wrapper_code.supports_intermediate_hooks
@@ -2082,7 +2088,7 @@ def codegen_template(self, template_node, epilogue_nodes):
 
         src_code = render()
         kernel_name = self.define_kernel(src_code, [template_node, *epilogue_nodes])
-        kernel.call_kernel(V.graph.wrapper_code, kernel_name)
+        kernel.call_kernel(kernel_name)
         self.scheduler.free_buffers()
 
     def codegen_sync(self):

torch/_inductor/codegen/wrapper.py

@@ -24,7 +24,7 @@
     sympy_product,
 )
 from ..virtualized import V
-from .common import CodeGen, DeferredLine, IndentedBuffer, Kernel, PythonPrinter
+from .common import CodeGen, DeferredLine, IndentedBuffer, PythonPrinter
 
 
 pexpr = PythonPrinter().doprint
@@ -253,6 +253,7 @@ def __init__(self):
         self.wrapper_call = IndentedBuffer()
         self.src_to_kernel = {}
         self.kernel_to_hash = {}
+        self.kenel_numel_expr = set()
         self.lines = []
         self.declare = ""
         self.ending = ""
@@ -665,14 +666,6 @@ def generate_kernel_call(
         else:
             self.writeline(self.wrap_kernel_call(name, call_args))
 
-    def call_kernel(self, name: str, kernel: Kernel):
-        tmp = IndentedBuffer()
-        kernel.call_kernel(self, tmp, name)
-        for line in tmp.getvalue().split("\n"):
-            line = line.strip()
-            if line:
-                self.writeline(line)
-
     def writeline(self, line):
         self.lines.append(line)
 

torch/_inductor/select_algorithm.py

@@ -289,26 +289,29 @@ def initialize_range_tree(self, pid_cache):
         self.body.clear()
         self.indexing_code.clear()
 
-    def call_kernel(self, code, name: str):
+    def call_kernel(self, name: str):
+        wrapper = V.graph.wrapper_code
         _, call_args, _ = self.args.python_argdefs()
 
         for i in range(len(call_args)):
             if V.graph.is_unspec_arg(call_args[i]):
                 call_args[i] = call_args[i] + ".item()"
         call_args = ", ".join(call_args)
 
-        stream_name = code.write_get_cuda_stream(V.graph.scheduler.current_device.index)
+        stream_name = wrapper.write_get_cuda_stream(
+            V.graph.scheduler.current_device.index
+        )
 
-        V.graph.wrapper_code.add_import_once(f"import {self.grid_fn.__module__}")
-        meta = V.graph.wrapper_code.add_meta_once(self.meta)
+        wrapper.add_import_once(f"import {self.grid_fn.__module__}")
+        meta = wrapper.add_meta_once(self.meta)
 
         grid_call = [texpr(V.graph.sizevars.simplify(s)) for s in self.call_sizes] + [
             meta
         ]
         grid_call = (
             f"{self.grid_fn.__module__}.{self.grid_fn.__name__}({', '.join(grid_call)})"
         )
-        code.writeline(
+        wrapper.writeline(
             f"{name}.run({call_args}, grid={grid_call}, stream={stream_name})"
         )