[SPEC] Fix spec __init__.py, _filecheck.py, _internal_testing.py in triton dir

Author: zhzhcookie

python/triton/__init__.py

@@ -4,11 +4,6 @@
 # ---------------------------------------
 # Note: import order is significant here.
 
-# flagtree backend path specialization
-from .flagtree_spec import spec_path
-
-spec_path(__path__)
-
 # submodules
 from .runtime import (
     autotune,

python/triton/_filecheck.py

@@ -17,23 +17,13 @@
 # ===-----------------------------------------------------------------------===#
 
 # Stub target for testing the frontend.
-stub_target = GPUTarget("cuda", 100, 32)
+# flagtree backend path specialization
+from triton.flagtree_spec import spec
+if not spec("spec_get_stub_target"):
+    stub_target = GPUTarget("cuda", 100, 32)
 
 triton_dir = os.path.dirname(__file__)
-_filecheck_local = os.path.join(triton_dir, "FileCheck")
-_filecheck_system = shutil.which("FileCheck")
-_filecheck_path = _filecheck_local if os.path.isfile(_filecheck_local) else _filecheck_system
-
-_MISSING_FILECHECK_MSG = ("FileCheck binary not found.  Install it with your package manager\n"
-                          "  (e.g. apt-get install llvm-15-tools) or place it next to this module:\n"
-                          f"  {_filecheck_local}")
-
-
-def _get_filecheck_path():
-    """Return the path to the FileCheck binary, or raise FileNotFoundError."""
-    if _filecheck_path is None:
-        raise FileNotFoundError(_MISSING_FILECHECK_MSG)
-    return _filecheck_path
+filecheck_path = os.path.join(triton_dir, "FileCheck")
 
 
 class MatchError(ValueError):

python/triton/_internal_testing.py

@@ -43,7 +43,11 @@ def is_ampere_or_newer():
 
 
 def is_blackwell():
-    return is_cuda() and torch.cuda.get_device_capability()[0] == 10
+    return is_cuda() and torch.cuda.get_device_capability()[0] in [10, 11]  # Triton 3.7
+
+
+def is_blackwell_ultra():
+    return is_cuda() and torch.cuda.get_device_capability()[0:2] == (10, 3)  # Triton 3.7
 
 
 def is_hopper_or_newer():
@@ -106,6 +110,17 @@ def is_xpu():
     return False if target is None else target.backend == "xpu"
 
 
+# flagtree: mthreads
+def is_musa():
+    target = get_current_target()
+    return False if target is None else target.backend == "musa"
+
+
+# flagtree: mthreads
+def is_musa_ph1():
+    return is_musa() and torch.musa.get_device_capability() == (3, 1)
+
+
 def get_arch():
     target = get_current_target()
     return "" if target is None else str(target.arch)
@@ -156,6 +171,8 @@ def to_triton(x: np.ndarray, device, dst_type=None) -> Union[TensorWrapper, torc
         if dst_type and 'float8' in dst_type:
             return reinterpret(torch.tensor(x, device=device), getattr(tl, dst_type))
         if t == 'float32' and dst_type == 'bfloat16':
+            if is_musa():  # flagtree mthreads
+                return torch.tensor(x, device='cpu').bfloat16().to(device)
             return torch.tensor(x, device=device).bfloat16()
         return torch.tensor(x, device=device)
 

python/triton/flagtree_spec.py

@@ -16,18 +16,12 @@ def _triton_root() -> str | None:
 def _get_spec_module():
     global _spec_module
     from ._flagtree_backend import FLAGTREE_BACKEND
-    if not FLAGTREE_BACKEND:
-        return None
-    triton_root = _triton_root()
-    if triton_root is None:
-        return None
-    spec_dir = os.path.join(triton_root, "backends", FLAGTREE_BACKEND, "spec")
-    if not os.path.isdir(spec_dir):
-        return None
     if _spec_module is not None:
         return _spec_module
+    if not FLAGTREE_BACKEND:
+        return None
     try:
-        _spec_module = importlib.import_module(f"triton.backends.{FLAGTREE_BACKEND}.spec")
+        _spec_module = importlib.import_module(f"triton.backends.{FLAGTREE_BACKEND}.spec.triton")
     except ImportError:
         return None
     return _spec_module

…_party/mthreads/backend/spec/__init__.py …mthreads/backend/spec/triton/__init__.pythird_party/mthreads/backend/spec/__init__.py renamed to third_party/mthreads/backend/spec/triton/__init__.py third_party/mthreads/backend/spec/__init__.py renamed to third_party/mthreads/backend/spec/triton/__init__.py

@@ -82,3 +82,8 @@ def _experimental_descriptor_store(desc_pointer, value, offsets, _semantic=None)
     _ext._experimental_descriptor_store = _experimental_descriptor_store
 
     bind_language_extension_symbols_to_tl(_ext)
+
+
+def spec_get_stub_target() -> GPUTarget:
+    arch = os.environ.get("TRITON_OVERRIDE_ARCH") or os.environ.get("TRITON_MUSA_ARCH") or "ph1"
+    return GPUTarget("musa", arch, 32)