Get cutlass_library import working under fbcode (#125257)

Differential Revision: D56764089

Pull Request resolved: #125257
Approved by: https://github.com/chenyang78

torch/_inductor/codegen/cuda/cutlass_utils.py

@@ -10,6 +10,7 @@
 import sympy
 
 import torch
+from ... import config
 from ...config import cuda as inductor_cuda_config
 from ...ir import Layout
 
@@ -48,6 +49,9 @@ def _gen_cutlass_file(
 
 @functools.lru_cache(None)
 def try_import_cutlass() -> bool:
+    if config.is_fbcode():
+        return True
+
     # 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.

torch/_inductor/config.py

@@ -9,6 +9,9 @@ def is_fbcode():
     return not hasattr(torch.version, "git_version")
 
 
+if is_fbcode():
+    from triton.fb import build_paths
+
 # add some debug printouts
 debug = False
 
@@ -753,7 +756,9 @@ class cuda:
     # The default path only works under PyTorch local development environment.
     cutlass_dir = os.environ.get(
         "TORCHINDUCTOR_CUTLASS_DIR",
-        os.path.abspath(
+        build_paths.cutlass()
+        if is_fbcode()
+        else os.path.abspath(
             os.path.join(os.path.dirname(torch.__file__), "../third_party/cutlass/")
         ),
     )