[Inductor] handle AMD special launch options (#124146)

sijiac · pytorchmergebot · commit 0f6ce45bcbd7 · 2024-04-16T11:07:17.000Z
Summary: `matrix_instr_nonkdim` and `waves_per_eu` are AMD specific launch configs that can't be treated as fn input args Test Plan: HIP_VISIBLE_DEVICES=7 numactl --cpunodebind=1 --membind=1 buck2 run mode/{opt,amd-gpu} -c fbcode.triton_backend=amd -c fbcode.enable_gpu_sections=true -c fbcode.rocm_arch=mi300 //hammer/modules/sequential/encoders/tests:hstu_bench -- --torch-compile=True the E2E works well on the magic model Differential Revision: D56165438 Pull Request resolved: #124146 Approved by: https://github.com/aakhundov
diff --git a/torch/_inductor/triton_heuristics.py b/torch/_inductor/triton_heuristics.py
@@ -300,6 +300,13 @@ def _precompile_config(self, cfg: Config, warm_cache_only_with_cc: Optional[int]
         """Ahead of time compile a given autotuner config."""
         compile_meta = copy.deepcopy(self.triton_meta)
         for k, v in cfg.kwargs.items():
+            if torch.version.hip is not None:
+                if k == "matrix_instr_nonkdim":
+                    compile_meta["matrix_instr_nonkdim"] = v
+                    continue
+                if k == "waves_per_eu":
+                    compile_meta["waves_per_eu"] = v
+                    continue
             compile_meta["constants"][self.fn.arg_names.index(k)] = v
         compile_meta["num_warps"] = cfg.num_warps
         compile_meta["num_stages"] = cfg.num_stages
@@ -340,6 +347,13 @@ def _precompile_config(self, cfg: Config, warm_cache_only_with_cc: Optional[int]
                 "num_stages": compile_meta["num_stages"],
                 "debug": compile_meta["debug"],
             }
+            if torch.version.hip is not None:
+                if "waves_per_eu" in compile_meta:
+                    options["waves_per_eu"] = compile_meta["waves_per_eu"]
+                if "matrix_instr_nonkdim" in compile_meta:
+                    options["matrix_instr_nonkdim"] = compile_meta[
+                        "matrix_instr_nonkdim"
+                    ]
             compile_kwargs = {
                 "target": target,
                 "options": options,
