[Refactor] Refactor env into a more flexible version (#740)

* Fix environment variable name for compilation print setting in `env.py`

* Remove deprecated test file for warp specialized pass configuration and refactor environment variable access in `env.py` to utilize a centralized `EnvVar` class for better management and clarity.

* lint fix

* Refactor cache check to use `env.is_cache_enabled()` for consistency in `tuner.py`

Author: LeiWang1999

testing/python/components/test_tilelang_env.py

@@ -0,0 +1,17 @@
+import tilelang
+import os
+
+
+def test_env_var():
+    # test default value
+    assert tilelang.env.TILELANG_PRINT_ON_COMPILATION == "1"
+    # test forced value
+    os.environ["TILELANG_PRINT_ON_COMPILATION"] = "0"
+    assert tilelang.env.TILELANG_PRINT_ON_COMPILATION == "0"
+    # test forced value with class method
+    tilelang.env.TILELANG_PRINT_ON_COMPILATION = "1"
+    assert tilelang.env.TILELANG_PRINT_ON_COMPILATION == "1"
+
+
+if __name__ == "__main__":
+    test_env_var()

testing/python/pass_config/test_tilelang_pass_config_disable_warp_specialized.py renamed to testing/python/components/test_tilelang_pass_config_disable_warp_specialized.py

@@ -53,8 +53,8 @@ def _init_logger():
 
 logger = logging.getLogger(__name__)
 
-from .env import SKIP_LOADING_TILELANG_SO
 from .env import enable_cache, disable_cache, is_cache_enabled  # noqa: F401
+from .env import env as env  # noqa: F401
 
 import tvm
 import tvm.base
@@ -76,12 +76,12 @@ def _load_tile_lang_lib():
 
 
 # only load once here
-if SKIP_LOADING_TILELANG_SO == "0":
+if env.SKIP_LOADING_TILELANG_SO == "0":
     _LIB, _LIB_PATH = _load_tile_lang_lib()
 
 from .jit import jit, JITKernel, compile  # noqa: F401
 from .profiler import Profiler  # noqa: F401
-from .cache import cached  # noqa: F401
+from .cache import clear_cache  # noqa: F401
 
 from .utils import (
     TensorSupplyType,  # noqa: F401

tilelang/__init__.py

@@ -25,13 +25,7 @@
 import traceback
 from pathlib import Path
 
-from tilelang.env import (
-    TILELANG_CACHE_DIR,
-    TILELANG_AUTO_TUNING_CPU_UTILITIES,
-    TILELANG_AUTO_TUNING_CPU_COUNTS,
-    TILELANG_AUTO_TUNING_MAX_CPU_COUNT,
-    is_cache_enabled,
-)
+from tilelang import env
 from tilelang.autotuner.param import CompileArgs, ProfileArgs, AutotuneResult
 from tilelang.autotuner.capture import get_autotune_inputs
 from tilelang.jit.param import _P, _RProg
@@ -111,7 +105,7 @@ class AutoTuner:
     _kernel_parameters: Optional[Tuple[str, ...]] = None
     _lock = threading.Lock()  # For thread safety
     _memory_cache = {}  # In-memory cache dictionary
-    cache_dir: Path = Path(TILELANG_CACHE_DIR) / "autotuner"
+    cache_dir: Path = Path(env.TILELANG_CACHE_DIR) / "autotuner"
 
     def __init__(self, fn: Callable, configs):
         self.fn = fn
@@ -285,7 +279,7 @@ def run(self, warmup: int = 25, rep: int = 100, timeout: int = 30):
         key = self.generate_cache_key(parameters)
 
         with self._lock:
-            if is_cache_enabled():
+            if env.is_cache_enabled():
                 # First check in-memory cache
                 if key in self._memory_cache:
                     logger.warning("Found kernel in memory cache. For better performance," \
@@ -437,9 +431,9 @@ def shape_equal(a, b):
                 return autotuner_result
         # get the cpu count
         available_cpu_count = get_available_cpu_count()
-        cpu_utilizations = float(TILELANG_AUTO_TUNING_CPU_UTILITIES)
-        cpu_counts = int(TILELANG_AUTO_TUNING_CPU_COUNTS)
-        max_cpu_count = int(TILELANG_AUTO_TUNING_MAX_CPU_COUNT)
+        cpu_utilizations = float(env.TILELANG_AUTO_TUNING_CPU_UTILITIES)
+        cpu_counts = int(env.TILELANG_AUTO_TUNING_CPU_COUNTS)
+        max_cpu_count = int(env.TILELANG_AUTO_TUNING_MAX_CPU_COUNT)
         if cpu_counts > 0:
             num_workers = min(cpu_counts, available_cpu_count)
             logger.info(
@@ -543,7 +537,7 @@ def device_wrapper(func, device, **config_arg):
             logger.warning("DLPack backend does not support cache saving to disk.")
         else:
             with self._lock:
-                if is_cache_enabled():
+                if env.is_cache_enabled():
                     self._save_result_to_disk(key, autotuner_result)
 
         self._memory_cache[key] = autotuner_result

tilelang/autotuner/tuner.py

@@ -4,8 +4,8 @@
 from tvm.target import Target
 from tvm.tir import PrimFunc
 from tilelang.jit import JITKernel
+from tilelang import env
 from .kernel_cache import KernelCache
-from tilelang.env import TILELANG_CLEAR_CACHE
 
 # Create singleton instance of KernelCache
 _kernel_cache_instance = KernelCache()
@@ -44,5 +44,5 @@ def clear_cache():
     _kernel_cache_instance.clear_cache()
 
 
-if TILELANG_CLEAR_CACHE.lower() in ("1", "true", "yes", "on"):
+if env.TILELANG_CLEAR_CACHE.lower() in ("1", "true", "yes", "on"):
     clear_cache()

tilelang/cache/__init__.py

@@ -14,7 +14,7 @@
 from tvm.tir import PrimFunc
 
 from tilelang.engine.param import KernelParam
-from tilelang.env import TILELANG_CACHE_DIR, TILELANG_TMP_DIR, is_cache_enabled
+from tilelang import env
 from tilelang.jit import JITKernel
 from tilelang.version import __version__
 
@@ -61,8 +61,8 @@ def __new__(cls):
 
     @staticmethod
     def _create_dirs():
-        os.makedirs(TILELANG_CACHE_DIR, exist_ok=True)
-        os.makedirs(TILELANG_TMP_DIR, exist_ok=True)
+        os.makedirs(env.TILELANG_CACHE_DIR, exist_ok=True)
+        os.makedirs(env.TILELANG_TMP_DIR, exist_ok=True)
 
     def _generate_key(
         self,
@@ -132,7 +132,7 @@ def cached(
         Returns:
             JITKernel: The compiled kernel, either freshly compiled or from cache
         """
-        if not is_cache_enabled():
+        if not env.is_cache_enabled():
             return JITKernel(
                 func,
                 out_idx=out_idx,
@@ -190,7 +190,7 @@ def cached(
             self.logger.warning("DLPack backend does not support cache saving to disk.")
         else:
             with self._lock:
-                if is_cache_enabled():
+                if env.is_cache_enabled():
                     self._save_kernel_to_disk(key, kernel, func, verbose)
 
         # Store in memory cache after compilation
@@ -215,7 +215,7 @@ def _get_cache_path(self, key: str) -> str:
         Returns:
             str: Absolute path to the cache directory for this kernel.
         """
-        return os.path.join(TILELANG_CACHE_DIR, key)
+        return os.path.join(env.TILELANG_CACHE_DIR, key)
 
     @staticmethod
     def _load_binary(path: str):
@@ -226,7 +226,7 @@ def _load_binary(path: str):
     @staticmethod
     def _safe_write_file(path: str, mode: str, operation: Callable):
         # Random a temporary file within the same FS as the cache directory
-        temp_path = os.path.join(TILELANG_TMP_DIR, f"{os.getpid()}_{uuid.uuid4()}")
+        temp_path = os.path.join(env.TILELANG_TMP_DIR, f"{os.getpid()}_{uuid.uuid4()}")
         with open(temp_path, mode) as temp_file:
             operation(temp_file)
 
@@ -396,7 +396,7 @@ def _clear_disk_cache(self):
         """
         try:
             # Delete the entire cache directory
-            shutil.rmtree(TILELANG_CACHE_DIR)
+            shutil.rmtree(env.TILELANG_CACHE_DIR)
 
             # Re-create the cache directory
             KernelCache._create_dirs()

tilelang/cache/kernel_cache.py

@@ -6,7 +6,7 @@
 import os
 import subprocess
 import warnings
-from ..env import CUDA_HOME
+from tilelang.env import CUDA_HOME
 
 import tvm.ffi
 from tvm.target import Target