Update on "[aotinductor] Avoid generating redundant kernel loading code"

desertfire · desertfire · commit b4b55bd4421e · 2023-10-05T08:04:37.000-07:00
Summary: 1) Stop forcing triton.unique_kernel_names to True for AOTInductor, because the unique kernel name can be read from metadata; 2) Only generate load_kernel once for each kernel since we don't have control flow in our generated code. This solves #105553. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler [ghstack-poisoned]
diff --git a/test/inductor/test_aot_inductor.py b/test/inductor/test_aot_inductor.py
@@ -806,7 +806,7 @@ def forward(self, x, y):
             with open(os.path.splitext(so_path)[0] + ".cpp") as cpp:
                 src_code = cpp.read()
                 FileCheck().check_count(
-                    "triton_poi_fused_sin_0 == nullptr) {",
+                    "triton_poi_fused_sin_0 = loadKernel(",
                     1,
                     exactly=True,
                 ).run(src_code)
diff --git a/torch/_inductor/codegen/wrapper.py b/torch/_inductor/codegen/wrapper.py
@@ -2014,17 +2014,13 @@ def generate_load_kernel_once(
         self, name: str, mangled_name: str, cubin_path: str, shared_mem: int
     ):
         if V.graph.aot_mode:
-            self.writeline(f"if (kernels.{name} == nullptr) {{")
             self.writeline(
-                f"""    kernels.{name} = loadKernel("{cubin_path}", "{mangled_name}", {shared_mem}, this->cubin_dir_);"""
+                f"""kernels.{name} = loadKernel("{cubin_path}", "{mangled_name}", {shared_mem}, this->cubin_dir_);"""
             )
-            self.writeline("}")
         else:
-            self.writeline(f"if ({name} == nullptr) {{")
             self.writeline(
-                f"""    {name} = loadKernel("{cubin_path}", "{mangled_name}", {shared_mem});"""
+                f"""{name} = loadKernel("{cubin_path}", "{mangled_name}", {shared_mem});"""
             )
-            self.writeline("}")
 
     def generate_args_decl(self, call_args):
         dynamic_symbols = V.graph.sizevars.free_symbols()
