[release/2.11] [inductor] Backport num_stages=1 pipelining-disable autotune fallback (#3408)

Cherry-pick of pytorch#180892 (30e2dee) minus the
mix_order_reduction_allow_multi_stages default flip.

The mix_order_reduction_allow_multi_stages flip is an unrelated
mix-order-reduction perf feature that the pytorch#17189 conv OOM fix doesn't
need — the num_stages=1 fallback delivers that on its own. It was
dropped because it's tuned for NVIDIA's ~228 KB shared memory and on
gfx942's 64 KB LDS would mostly cause double-compile churn and new
overflow risk with no proven benefit, keeping this a conservative,
fix-only backport.

When every Triton config for a choice fails to build with an
OutOfResources/OutOfMemoryError and pipelining is active (num_stages >
1), CachingAutotuner now recompiles that config with num_stages=1 and
retries instead of raising. On gfx942 (64KB LDS) this recovers the
BLOCK_K=64/128 triton_convolution2d configs that otherwise overflow
shared memory, eliminating the "No valid triton configs ... out of
resource" autotuning errors seen in ROCm/frameworks-internal#17189 and
keeping those kernels available as real choices.

Excludes the mix_order_reduction_allow_multi_stages opt-in -> default-on
change from the original PR to keep this a fix-only backport.

---------

Co-authored-by: shunting314 <shunting@meta.com>

Author: iupaikov-amd

test/inductor/test_max_autotune_blackwell.py

@@ -280,6 +280,19 @@ def addmm(x, a, b):
                 "triton.enable_template_tma_store": tma_store,
                 "test_configs.autotune_choice_name_regex": "blackwell_ws_persistent_device_tma",
                 "test_configs.autotune_choice_desc_regex": epilogue_subtile_regex,
+                # If we dynamically disable pipelining,
+                # triton_blackwell_ws_persistent_device_tma template will
+                # be picked and then cause mis-aligned memory access.
+                # If we don't dynamically disable pipelining,
+                # this template is skipped and triton_tem_fused_addmm_repeat_3
+                # is used.
+                #
+                # Fundamentally we should fix the triton_blackwell_ws_persistent_device_tma template
+                # The flag below work around the problem.
+                #
+                # This only happens in fbcode https://www.internalfb.com/diff/D101855575.
+                # Can not repro it in OSS build.
+                "triton.dynamic_disable_pipelining": False,
             }
         ):
             c_actual, code = run_and_get_code(

torch/_inductor/codegen/triton.py

@@ -5298,6 +5298,7 @@ def inductor_meta_common(cls):
             "deterministic": config.deterministic,
             "force_filter_reduction_configs": config.test_configs.force_filter_reduction_configs,
             "mix_order_reduction_allow_multi_stages": config.triton.mix_order_reduction_allow_multi_stages,
+            "dynamic_disable_pipelining": config.triton.dynamic_disable_pipelining,
         }
 
         if config.write_are_deterministic_algorithms_enabled:

torch/_inductor/config.py

@@ -1844,6 +1844,8 @@ class triton:
         os.environ.get("TORCHINDUCTOR_TRITON_PROTON_PER_CTA_OCCUPANCY", "1") == "1"
     )
 
+    dynamic_disable_pipelining = True
+
 
 class aot_inductor:
     """

torch/_inductor/runtime/triton_heuristics.py

@@ -642,6 +642,20 @@ def _dynamic_scale_rblock(self):
 
             self._make_launchers()
 
+    def compile_by_disabling_pipelining(self, config):
+        # self.fn.fn is dropped by prepare_for_pickle() after the initial
+        # compile; reload it so triton.compile can access the source.
+        if self.fn.fn is None:
+            assert callable(self._reload_kernel)
+            self.fn = self._reload_kernel().fn
+        cfg = copy.deepcopy(config)
+        cfg.num_stages = 1
+        if "NUM_STAGES" in cfg.kwargs:
+            cfg.kwargs["NUM_STAGES"] = 1
+        result = self._precompile_config(cfg)
+        self.compile_results = [result]
+        return result.make_launcher()
+
     def _make_launchers(self):
         if len(self.launchers) == len(self.compile_results):
             return
@@ -660,8 +674,21 @@ def _make_launchers(self):
 
                 except (OutOfResources, PTXASError, torch.cuda.OutOfMemoryError) as e:
                     exc = e
-        if len(launchers) == 0:
-            raise RuntimeError(f"No valid triton configs. {type(exc).__name__}: {exc}")
+            if len(launchers) == 0:
+                result = self.compile_results[-1]
+                config = result.config
+                if (
+                    isinstance(exc, (OutOfResources, torch.cuda.OutOfMemoryError))
+                    and (
+                        config.num_stages > 1 or config.kwargs.get("NUM_STAGES", 1) > 1
+                    )
+                    and self.inductor_meta.get("dynamic_disable_pipelining", True)
+                ):
+                    self.launchers = [self.compile_by_disabling_pipelining(config)]
+                    return
+                raise RuntimeError(
+                    f"No valid triton configs. {type(exc).__name__}: {exc}"
+                )
         self.launchers = launchers
 
     def prepare_for_pickle(self) -> tuple[Any, Any, Any, Any, Any, Any]: