Add option to ban/disallow autotuning (#184)

Author: oulgen

helion/exc.py

@@ -341,3 +341,7 @@ class TensorOperationsInHostCall(TensorOperationInWrapper):
 
 class WrongDevice(BaseWarning):
     message = "Operation {0} returned a tensor on {1} device, but the kernel is on {2} device. "
+
+
+class AutotuningDisallowedInEnvironment(BaseWarning):
+    message = "Autotuning is disabled {0}, please provide a config to @helion.kernel via the config= argument."

helion/runtime/kernel.py

@@ -410,6 +410,8 @@ def autotune(
 
                 config = FiniteSearch(self, args, self.configs).autotune()
         else:
+            self.settings.check_autotuning_disabled()
+
             from ..autotuner import DifferentialEvolutionSearch
 
             config = DifferentialEvolutionSearch(

helion/runtime/settings.py

@@ -11,12 +11,13 @@
 from typing import cast
 
 import torch
+from torch._environment import is_fbcode
+
+from helion import exc
 
 if TYPE_CHECKING:
     from contextlib import AbstractContextManager
 
-    from helion import exc
-
     class _TLS(Protocol):
         default_settings: Settings | None
 
@@ -110,6 +111,20 @@ def shallow_copy(x: object) -> object:
 
         return {k: shallow_copy(v) for k, v in dataclasses.asdict(self).items()}
 
+    def check_autotuning_disabled(self) -> None:
+        msg = None
+        if os.environ.get("HELION_DISALLOW_AUTOTUNING", "0") == "1":
+            msg = "by HELION_DISALLOW_AUTOTUNING=1"
+        if is_fbcode():
+            from aiplatform.runtime_environment.runtime_environment_pybind import (  # pyre-fixme[21]
+                RuntimeEnvironment,
+            )
+
+            if RuntimeEnvironment().get_mast_job_name() is not None:
+                msg = "because autotuning is not allowed in MAST environment"
+        if msg:
+            raise exc.AutotuningDisallowedInEnvironment(msg)
+
     @staticmethod
     def default() -> Settings:
         """

test/test_autotuner.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
+import os
 from pathlib import Path
 import random
 import tempfile
 import unittest
 from unittest.mock import patch
 
 from expecttest import TestCase
+import pytest
 import torch
 
 import helion
@@ -200,6 +202,27 @@ def add(a, b):
         result = add(*args)
         torch.testing.assert_close(result, sum(args))
 
+    def test_autotuner_disabled(self):
+        @helion.kernel
+        def add(a, b):
+            out = torch.empty_like(a)
+            for tile in hl.tile(out.size()):
+                out[tile] = a[tile] + b[tile]
+            return out
+
+        args = (
+            torch.randn([8, 512, 512], device=DEVICE),
+            torch.randn([8, 512, 512], device=DEVICE),
+        )
+        with (
+            patch.dict(os.environ, {"HELION_DISALLOW_AUTOTUNING": "1"}),
+            pytest.raises(
+                expected_exception=helion.exc.AutotuningDisallowedInEnvironment,
+                match="Autotuning is disabled by HELION_DISALLOW_AUTOTUNING=1, please provide a config to @helion.kernel via the config= argument.",
+            ),
+        ):
+            add(*args)
+
 
 if __name__ == "__main__":
     unittest.main()