[python_utils] add try import utils (#387)

Author: xuzhao9

tritonbench/operators/flash_attention/operator.py

@@ -57,23 +57,19 @@
 )
 
 from tritonbench.utils.env_utils import get_nvidia_gpu_model, is_cuda, is_hip
-
 from tritonbench.utils.path_utils import add_ld_library_path
+from tritonbench.utils.python_utils import try_import
 from tritonbench.utils.triton_op import is_fbcode
 
 
 # [Optional] flash_attn v2
-try:
+with try_import("HAS_FLASH_V2"):
     from flash_attn.flash_attn_interface import (
         flash_attn_qkvpacked_func as flash_attn_func,
     )
 
     from .test_fmha_utils import make_packed_qkv
 
-    HAS_FLASH_V2 = True
-except (ImportError, IOError, AttributeError):
-    HAS_FLASH_V2 = False
-
 HAS_CUDA_124 = (
     torch.cuda.is_available() and torch.version.cuda and torch.version.cuda >= "12.4"
 )
@@ -83,59 +79,37 @@
 # only enabling the variants known to be working on B200 (trunk).
 if not IS_B200:
     # [Optional] flash_attn v3
-    try:
-        torch_lib_path = os.path.join(os.path.dirname(__file__), "lib")
-        with add_ld_library_path(torch_lib_path):
-            from flash_attn_interface import flash_attn_func as flash_attn_v3
-        HAS_FLASH_V3 = True
-    except (ImportError, IOError, AttributeError):
+    with try_import("HAS_FLASH_V3"):
         try:
-            from fa3.hopper.flash_attn_interface import flash_attn_func as flash_attn_v3
-
-            HAS_FLASH_V3 = True
+            torch_lib_path = os.path.join(os.path.dirname(__file__), "lib")
+            with add_ld_library_path(torch_lib_path):
+                from flash_attn_interface import flash_attn_func as flash_attn_v3
         except (ImportError, IOError, AttributeError):
-            HAS_FLASH_V3 = False
+            from fa3.hopper.flash_attn_interface import flash_attn_func as flash_attn_v3
 
-    try:
+    with try_import("HAS_TILELANG"):
         import tilelang
 
         from .tilelang_mha import tilelang_mha
 
-        HAS_TILELANG = True
-    except (ImportError, IOError, AttributeError, TypeError):
-        HAS_TILELANG = False
-
     # [Optional] ThunderKittens backend
-    try:
+    with try_import("HAS_TK"):
         from .tk import tk_attn
 
-        HAS_TK = True
-    except (ImportError, IOError, AttributeError):
-        HAS_TK = False
-
     # [Optional] JAX Pallas backend
-    try:
+    with try_import("HAS_PALLAS"):
         import jax
-
         from tritonbench.utils.jax_utils import torch_to_jax_tensor
 
         from .pallas import mha as pallas_mha
 
-        HAS_PALLAS = True
-    except (ImportError, IOError, AttributeError):
-        HAS_PALLAS = False
-
 # [Optional] xformers backend
-try:
+with try_import("HAS_XFORMERS"):
     import xformers  # @manual=//fair/xformers:xformers
     import xformers.ops.fmha as xformers_fmha  # @manual=//fair/xformers:xformers
 
     from .test_fmha_utils import permute_qkv
 
-    HAS_XFORMERS = True
-except (ImportError, IOError, AttributeError, TypeError):
-    HAS_XFORMERS = False
-
 from typing import Any, Generator, List
 
 from tritonbench.utils.input import input_filter

tritonbench/utils/python_utils.py

@@ -0,0 +1,13 @@
+import inspect
+from contextlib import contextmanager
+
+
+@contextmanager
+def try_import(cond_name: str):
+    frame = inspect.currentframe().f_back.f_back
+    _caller_globals = frame.f_globals
+    try:
+        yield
+        _caller_globals[cond_name] = True
+    except (ImportError, ModuleNotFoundError) as e:
+        _caller_globals[cond_name] = False