[inductor] set sanitize_overflow=False for triton kernels (pytorch#139502)

davidberard98 · rahulsingh-intel · commit d273954b6475 · 2024-11-05T12:15:25.000+02:00
In upstream triton, triton-lang/triton#4589 introduces overflow checks. However, overflow checks likely add some overhead, and have some correctness bugs at the moment (e.g. triton-lang/triton#5033). Let's set `sanitize_overflow=False` but keep `debug=True` so that we can keep using device_assert but without the additional asserts added by `sanitize_overflow`. Pull Request resolved: pytorch#139502 Approved by: https://github.com/bertmaher
diff --git a/torch/_inductor/runtime/triton_heuristics.py b/torch/_inductor/runtime/triton_heuristics.py
@@ -442,6 +442,7 @@ def _precompile_config(self, cfg: Config, warm_cache_only: bool):
                 "num_warps": compile_meta["num_warps"],
                 "num_stages": compile_meta["num_stages"],
                 "debug": compile_meta["debug"],
+                "sanitize_overflow": False,  # turn off additional asserts added for overflow checks
             }
             if self.device_props.type == "hip":
                 if "waves_per_eu" in compile_meta:

Original file line number	Diff line number	Diff line change
`@@ -442,6 +442,7 @@ def _precompile_config(self, cfg: Config, warm_cache_only: bool):`
`442`	`442`	`"num_warps": compile_meta["num_warps"],`
`443`	`443`	`"num_stages": compile_meta["num_stages"],`
`444`	`444`	`"debug": compile_meta["debug"],`
	`445`	`+ "sanitize_overflow": False, # turn off additional asserts added for overflow checks`
`445`	`446`	`}`
`446`	`447`	`if self.device_props.type == "hip":`
`447`	`448`	`if "waves_per_eu" in compile_meta:`