Move cutlass_dir to general config

Signed-off-by: Lukas Sommer <lukas.sommer@codeplay.com>

Author: sommerlukas

torch/_inductor/codecache.py

@@ -2944,7 +2944,7 @@ def _cutlass_include_paths() -> list[str]:
 
         cutlass_path = parutil.get_dir_path("cutlass-3-headers")
     else:
-        cutlass_path = config.cuda.cutlass_dir
+        cutlass_path = config.cutlass_dir
     return [
         # Use realpath to get canonical absolute paths, in order not to mess up cache keys
         os.path.realpath(os.path.join(cutlass_path, "include")),

torch/_inductor/codegen/cuda/cutlass_utils.py

@@ -53,13 +53,13 @@ def try_import_cutlass() -> bool:
             log.warning("CUTLASS version < 3.7 is not recommended.")
 
         log.debug(
-            "Found cutlass_library in python search path, overriding config.cuda.cutlass_dir"
+            "Found cutlass_library in python search path, overriding config.cutlass_dir"
         )
         cutlass_library_dir = os.path.dirname(cutlass_library.__file__)
         assert os.path.isdir(cutlass_library_dir), (
             f"{cutlass_library_dir} is not a directory"
         )
-        config.cuda.cutlass_dir = os.path.abspath(
+        config.cutlass_dir = os.path.abspath(
             os.path.join(
                 cutlass_library_dir,
                 "source",
@@ -68,15 +68,15 @@ def try_import_cutlass() -> bool:
         return True
     except ModuleNotFoundError:
         log.debug(
-            "cutlass_library not found in sys.path, trying to import from config.cuda.cutlass_dir"
+            "cutlass_library not found in sys.path, trying to import from config.cutlass_dir"
         )
 
     # Copy CUTLASS python scripts to a temp dir and add the temp dir to Python search path.
     # This is a temporary hack to avoid CUTLASS module naming conflicts.
     # TODO(ipiszy): remove this hack when CUTLASS solves Python scripts packaging structure issues.
 
     cutlass_py_full_path = os.path.abspath(
-        os.path.join(config.cuda.cutlass_dir, "python/cutlass_library")
+        os.path.join(config.cutlass_dir, "python/cutlass_library")
     )
     tmp_cutlass_py_full_path = os.path.abspath(
         os.path.join(cache_dir(), "torch_cutlass_library")

torch/_inductor/config.py

@@ -871,6 +871,16 @@ def decide_compile_threads() -> int:
 # Adds NVTX annotations aroung training phases
 annotate_training: bool = os.environ.get("TORCHINDUCTOR_ANNOTATE_TRAINING", "0") == "1"
 
+# Path to the CUTLASS repo root directory.
+# The default path only works under PyTorch local development environment.
+cutlass_dir = os.environ.get(
+    "TORCHINDUCTOR_CUTLASS_DIR",
+    os.path.abspath(
+        os.path.join(os.path.dirname(torch.__file__), "../third_party/cutlass/")
+    ),
+)
+
+
 
 # config specific to codegen/cpp.py
 class cpp:
@@ -1280,15 +1290,6 @@ class cuda:
     # Whether to use fast math.
     use_fast_math = False
 
-    # Path to the CUTLASS repo root directory.
-    # The default path only works under PyTorch local development environment.
-    cutlass_dir = os.environ.get(
-        "TORCHINDUCTOR_CUTLASS_DIR",
-        os.path.abspath(
-            os.path.join(os.path.dirname(torch.__file__), "../third_party/cutlass/")
-        ),
-    )
-
     # Configures the maximum number of CUTLASS configs to profile in max_autotune.
     # By default it's None, so that all CUTLASS configs are tuned.
     # This is mainly used to reduce test time in CI.
@@ -1534,7 +1535,7 @@ class trace:
     # trace functions are not relevant to config caching
     "trace",
     # uses absolute path
-    "cuda.cutlass_dir",
+    "cutlass_dir",
     # not relevant
     "worker_start_method",
     "compile_threads",

torch/_inductor/fuzzer.py

@@ -469,7 +469,7 @@ def keys(self) -> KeysView[ComboType]:
         "aot_inductor.presets": DEFAULT,  # Typing
         "cuda.arch": DEFAULT,  # Out of Scope
         "cuda.version": DEFAULT,  # Out of Scope
-        "cuda.cutlass_dir": DEFAULT,  # Out of Scope
+        "cutlass_dir": DEFAULT,  # Out of Scope
         "cuda.cuda_cxx": DEFAULT,  # Out of Scope
         "rocm.arch": DEFAULT,  # Out of Scope
         "rocm.ck_supported_arch": DEFAULT,  # Out of Scope

torch/_inductor/utils.py

@@ -1405,7 +1405,7 @@ def use_cutlass_template(layout: Layout, m: int, n: int, k: int) -> bool:
         if not try_import_cutlass():
             log.warning(
                 "Failed to import CUTLASS lib. Please check whether "
-                "_inductor.config.cuda.cutlass_dir is set correctly. "
+                "_inductor.config.cutlass_dir is set correctly. "
                 "Skipping CUTLASS backend for now."
             )
             return False