Add CUTLASS_DEBUG_TRACE_LEVEL flag for AITemplate compilation

Summary: Add CUTLASS_DEBUG_TRACE_LEVEL flag for AITemplate compilation in case debugging needs more information Differential Revision: D49242578
facebookincubator · Sep 13, 2023 · 24d439d · 24d439d
1 parent 6bbe03c
commit 24d439d
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/python/aitemplate/backend/cuda/target_def.py b/python/aitemplate/backend/cuda/target_def.py
@@ -158,6 +158,7 @@ def _build_nvcc_compiler_options(self) -> List[str]:
             environ.get_compiler_opt_level(),
             "-std=c++17",
             "--expt-relaxed-constexpr",
+            "-DCUTLASS_DEBUG_TRACE_LEVEL=" + environ.get_cutlass_debug_trace_level(),
         ]
         if environ.enable_ptxas_info():
             options.extend(
@@ -429,6 +430,8 @@ def _build_compile_options(self):
                     if not FBCUDA.optimize_for_compilation_time_
                     else "-O1",
                     "-std=c++17",
+                    "-DCUTLASS_DEBUG_TRACE_LEVEL="
+                    + environ.get_cutlass_debug_trace_level(),
                 ]
                 + (
                     ["-DOPTIMIZE_FOR_COMPILATION_TIME"]

diff --git a/python/aitemplate/utils/environ.py b/python/aitemplate/utils/environ.py
@@ -285,6 +285,13 @@ def get_cuda_nvcc_debug_level():
     return level
 
 
+def get_cutlass_debug_trace_level():
+    """
+    Return level of CUTLASS lib debug trace information. Default to no debug info.
+    """
+    level = os.getenv("CUTLASS_DEBUG_TRACE_LEVEL", "0")
+    return level
+
 def enable_cuda_source_navigation_fix():
     """
     When this flag is enabled, the FBCUDA Target will copy every *.cu file in build dirs into