more robust imports in heuristics.py and heuristics_provider.py (#2596)

2025-08-28 22:32:55 -04:00
parent 496654bf2c
commit b2dd65dc86
2 changed files with 9 additions and 2 deletions
@@ -41,7 +41,13 @@ import logging
 import ctypes
 import functools

-from library import DataType, LayoutType
+try:
+  import builtins
+  if hasattr(builtins, "CUTLASS_IGNORE_PACKAGE") and CUTLASS_IGNORE_PACKAGE == True:
+    raise ImportError("Disabling attempt to import cutlass_library")
+  from cutlass_library.library import DataType, LayoutType
+except ImportError:
+  from library import DataType, LayoutType

 class MatmulHeuristics: