[FlagPrism][Mthreads] Integrate profiler and debugger

cmake/FlagTreeOptions.cmake

@@ -98,27 +98,46 @@ endmacro()
 # FlagPrism: configure the external profiler/debugger after base options exist.
 macro(flagtree_configure_flagprism)
   set(_flagprism_default OFF)
-  if(FLAGTREE_BACKEND STREQUAL "ascend" OR FLAGTREE_BACKEND STREQUAL "iluvatar")
+  # FlagPrism: retain the original supported-backend condition for reference.
+  # if(FLAGTREE_BACKEND STREQUAL "ascend" OR FLAGTREE_BACKEND STREQUAL "iluvatar")
+  # FlagPrism: enable the external tools for the supported mthreads backend.
+  if(FLAGTREE_BACKEND STREQUAL "ascend" OR
+     FLAGTREE_BACKEND STREQUAL "iluvatar" OR
+     FLAGTREE_BACKEND STREQUAL "mthreads")
     set(_flagprism_default ON)
   endif()
   option(TRITON_BUILD_FLAGPRISM
          "Build the FlagPrism debugger and profiler"
          ${_flagprism_default})
 
   if(TRITON_BUILD_FLAGPRISM)
-    if(NOT (FLAGTREE_BACKEND STREQUAL "ascend" OR FLAGTREE_BACKEND STREQUAL "iluvatar"))
+    # FlagPrism: retain the original supported-backend guard for reference.
+    # if(NOT (FLAGTREE_BACKEND STREQUAL "ascend" OR FLAGTREE_BACKEND STREQUAL "iluvatar"))
+    # FlagPrism: accept mthreads as a supported integration backend.
+    if(NOT (FLAGTREE_BACKEND STREQUAL "ascend" OR
+            FLAGTREE_BACKEND STREQUAL "iluvatar" OR
+            FLAGTREE_BACKEND STREQUAL "mthreads"))
       message(FATAL_ERROR
         "TRITON_BUILD_FLAGPRISM is only supported when "
-        "FLAGTREE_BACKEND is ascend or iluvatar.")
+        # FlagPrism: original diagnostic was "FLAGTREE_BACKEND is ascend or iluvatar."
+        "FLAGTREE_BACKEND is ascend, iluvatar, or mthreads.")
     endif()
     if(TRITON_BUILD_PROTON)
       message(FATAL_ERROR
         "TRITON_BUILD_FLAGPRISM and TRITON_BUILD_PROTON cannot both be enabled. "
         "Select exactly one profiler implementation.")
     endif()
 
-    set(FLAGPRISM_CMAKE_FILE
-        "${CMAKE_CURRENT_SOURCE_DIR}/third_party/FlagPrism/cmake/FlagPrism.cmake")
+    # FlagPrism: retain the original bundled source path for reference.
+    # set(FLAGPRISM_CMAKE_FILE
+    #     "${CMAKE_CURRENT_SOURCE_DIR}/third_party/FlagPrism/cmake/FlagPrism.cmake")
+    # FlagPrism: permit a separate local checkout during backend development.
+    if(NOT FLAGPRISM_SOURCE_DIR)
+      set(FLAGPRISM_SOURCE_DIR
+          "${CMAKE_CURRENT_SOURCE_DIR}/third_party/FlagPrism")
+    endif()
+    get_filename_component(FLAGPRISM_SOURCE_DIR "${FLAGPRISM_SOURCE_DIR}" ABSOLUTE)
+    set(FLAGPRISM_CMAKE_FILE "${FLAGPRISM_SOURCE_DIR}/cmake/FlagPrism.cmake")
     if(EXISTS "${FLAGPRISM_CMAKE_FILE}")
       include("${FLAGPRISM_CMAKE_FILE}")
       add_compile_definitions(__FLAGPRISM__=1)

python/setup_tools/setup_helper.py

@@ -266,14 +266,20 @@ def __init__(self, project_root, dependency_cmake_args):
         self.project_root = Path(project_root)
         backend = configs.flagtree_backend or ""
         supported_backends = {"ascend", "iluvatar"}
+        # FlagPrism: register mthreads without replacing the original backend set.
+        supported_backends.add("mthreads")
         default = "ON" if backend in supported_backends else "OFF"
         self.enabled = self._check_env_flag("TRITON_BUILD_FLAGPRISM", default)
         self.build_config = None
         self._dependency_cmake_args = dependency_cmake_args
 
         if self.enabled and backend not in supported_backends:
+            # FlagPrism: retain the original diagnostic for reference.
+            # raise RuntimeError("TRITON_BUILD_FLAGPRISM is only supported when "
+            #                    "FLAGTREE_BACKEND=ascend or iluvatar.")
+            # FlagPrism: report the newly supported mthreads backend.
             raise RuntimeError("TRITON_BUILD_FLAGPRISM is only supported when "
-                               "FLAGTREE_BACKEND=ascend or iluvatar.")
+                               "FLAGTREE_BACKEND=ascend, iluvatar, or mthreads.")
         if not self.enabled:
             return
         if self._check_env_flag("TRITON_BUILD_PROTON"):
@@ -282,7 +288,12 @@ def __init__(self, project_root, dependency_cmake_args):
 
         # FlagPrism replaces Proton for the supported backend builds.
         os.environ["TRITON_BUILD_PROTON"] = "OFF"
-        source_root = self.project_root / "third_party" / "FlagPrism"
+        # FlagPrism: retain the original bundled checkout path for reference.
+        # source_root = self.project_root / "third_party" / "FlagPrism"
+        # FlagPrism: allow testing against a separate local FlagPrism checkout.
+        source_override = os.environ.get("FLAGPRISM_SOURCE_DIR", "").strip()
+        source_root = (Path(source_override).resolve() if source_override else self.project_root / "third_party" /
+                       "FlagPrism")
         # Keep FlagPrism as an external checkout. A local directory or symlink
         # is authoritative; only bootstrap the registered dependency when it
         # is absent.
@@ -294,7 +305,13 @@ def __init__(self, project_root, dependency_cmake_args):
             raise RuntimeError("FlagPrism sources are missing. Run the Python package build "
                                "to download third-party dependencies.")
         policy = runpy.run_path(str(helper_path), run_name="_flagprism_build")
-        self.build_config = policy["create_build_config"](self.project_root)
+        # FlagPrism: retain the original build-policy call for reference.
+        # self.build_config = policy["create_build_config"](self.project_root)
+        # FlagPrism: keep CMake and setuptools on the same external source tree.
+        if source_override:
+            self.build_config = policy["create_build_config"](self.project_root, source_root)
+        else:
+            self.build_config = policy["create_build_config"](self.project_root)
 
         legacy_link = self.project_root / "python" / "triton" / "profiler"
         if legacy_link.is_symlink():

python/test/unit/test_flagprism.py

@@ -165,7 +165,10 @@ def test_build_helper_uses_unified_switch(build_helper, monkeypatch, tmp_path, v
         ("cambricon", False),
         ("aipu", False),
         ("xpu", False),
-        ("mthreads", False),
+        # FlagPrism: retain the former unsupported expectation for reference.
+        # ("mthreads", False),
+        # FlagPrism: mthreads now enables the profiler/debugger by default.
+        ("mthreads", True),
     ),
 )
 def test_flagprism_is_enabled_by_default_for_supported_backends(flagprism_setup_factory, monkeypatch, backend, enabled):
@@ -205,7 +208,11 @@ def test_non_ascend_explicit_flagprism_is_rejected_before_side_effects(flagprism
     monkeypatch.setenv("TRITON_BUILD_FLAGPRISM", "ON")
     monkeypatch.setenv("TRITON_BUILD_PROTON", "ON")
 
-    with pytest.raises(RuntimeError, match="ascend or iluvatar"):
+    # FlagPrism: retain the former diagnostic assertion for reference.
+    # with pytest.raises(RuntimeError, match="ascend or iluvatar"):
+    #     create("enflame")
+    # FlagPrism: include mthreads in the supported-backend diagnostic.
+    with pytest.raises(RuntimeError, match="ascend, iluvatar, or mthreads"):
         create("enflame")
 
     assert not downloads

third_party/mthreads/backend/driver.py

@@ -6,6 +6,7 @@
 from pathlib import Path
 
 from triton import knobs
+from flagtree import _flagprism  # FlagPrism
 from triton.backends.mthreads._musa_arch import musa_capability_from_arch, musa_warp_size_from_arch
 from triton.backends.compiler import GPUTarget
 from triton.backends.driver import DriverBase
@@ -820,6 +821,10 @@ def cst_key(i):
 
         constants = {cst_key(key): value for key, value in constants.items()}
         signature = {cst_key(key): value for key, value in src.signature.items()}
+        # FlagPrism: preserve user arguments and debugger launch metadata.
+        self.user_arg_count = len(signature)
+        self.metadata = metadata
+        self._debug_enabled = bool(getattr(metadata, "debug_launch_hidden_arg", False))
 
         ordered_sig_keys = sorted(signature.keys())
         self._signature_types = [signature[key] for key in ordered_sig_keys]
@@ -835,6 +840,9 @@ def cst_key(i):
 
         expanded_signature_types, expanded_index = _expand_signature_tree(self._signature_types, self._tensordesc_meta)
         expanded_signature = {idx: ty for idx, ty in enumerate(expanded_signature_types)}
+        # FlagPrism: reserve the final ABI slot for the debugger control pointer.
+        if self._debug_enabled:
+            expanded_signature[len(expanded_signature)] = "*i8"
         expanded_constants = {}
         for key, value in constants.items():
             path = _normalize_arg_path(key)
@@ -970,36 +978,95 @@ def _set_and_keep_global_scratch(self, buf):
             self._global_scratch_keepalive = self._global_scratch_keepalive[-4096:]
         self._set_global_scratch(buf.data_ptr())
 
+    # FlagPrism: retain the original MUSA launch path for reference.
+    # def __call__(self, *args, **kwargs):
+    #     if not self._needs_runtime_expansion:
+    #         gridX, gridY, gridZ = args[0], args[1], args[2]
+    #         stream = args[3]
+    #         gs_buf = self._alloc_global_scratch(gridX, gridY, gridZ, stream)
+    #         self._set_and_keep_global_scratch(gs_buf)
+    #         self.launch(*args, **kwargs)
+    #         return
+    #
+    #     # launch(gridX, gridY, gridZ, stream, function, kernel_metadata,
+    #     # launch_metadata, launch_enter_hook, launch_exit_hook, *kernel_args)
+    #     launch_prefix = args[:9]
+    #     kernel_args = args[9:]
+    #     if len(kernel_args) != len(self._signature_types):
+    #         raise RuntimeError("launcher argument count mismatch while expanding tensor descriptors")
+    #
+    #     expanded_kernel_args = []
+    #     launch_keepalive = []
+    #     tensordesc_state = [0]
+    #     for arg, ty in zip(kernel_args, self._signature_types):
+    #         self._expand_runtime_arg(arg, ty, expanded_kernel_args, launch_keepalive, tensordesc_state)
+    #
+    #     self._tensordesc_keepalive.extend(launch_keepalive)
+    #     if len(self._tensordesc_keepalive) > 4096:
+    #         self._tensordesc_keepalive = self._tensordesc_keepalive[-4096:]
+    #     gridX, gridY, gridZ = launch_prefix[0], launch_prefix[1], launch_prefix[2]
+    #     stream = launch_prefix[3]
+    #     gs_buf = self._alloc_global_scratch(gridX, gridY, gridZ, stream)
+    #     self._set_and_keep_global_scratch(gs_buf)
+    #     self.launch(*launch_prefix, *expanded_kernel_args, **kwargs)
+    #
+
+    # FlagPrism: append the debugger hidden argument through one launch wrapper.
     def __call__(self, *args, **kwargs):
-        if not self._needs_runtime_expansion:
-            gridX, gridY, gridZ = args[0], args[1], args[2]
-            stream = args[3]
+
+        def launch(hidden_args=()):
+            if not self._needs_runtime_expansion:
+                gridX, gridY, gridZ = args[0], args[1], args[2]
+                stream = args[3]
+                gs_buf = self._alloc_global_scratch(gridX, gridY, gridZ, stream)
+                self._set_and_keep_global_scratch(gs_buf)
+                return self.launch(*args, *hidden_args, **kwargs)
+
+            # launch(gridX, gridY, gridZ, stream, function, kernel_metadata,
+            # launch_metadata, launch_enter_hook, launch_exit_hook, *kernel_args)
+            launch_prefix = args[:9]
+            kernel_args = args[9:]
+            if len(kernel_args) != len(self._signature_types):
+                raise RuntimeError("launcher argument count mismatch while expanding "
+                                   "tensor descriptors")
+
+            expanded_kernel_args = []
+            launch_keepalive = []
+            tensordesc_state = [0]
+            for arg, ty in zip(kernel_args, self._signature_types):
+                self._expand_runtime_arg(
+                    arg,
+                    ty,
+                    expanded_kernel_args,
+                    launch_keepalive,
+                    tensordesc_state,
+                )
+
+            self._tensordesc_keepalive.extend(launch_keepalive)
+            if len(self._tensordesc_keepalive) > 4096:
+                self._tensordesc_keepalive = self._tensordesc_keepalive[-4096:]
+            gridX, gridY, gridZ = launch_prefix[0:3]
+            stream = launch_prefix[3]
             gs_buf = self._alloc_global_scratch(gridX, gridY, gridZ, stream)
             self._set_and_keep_global_scratch(gs_buf)
-            self.launch(*args, **kwargs)
-            return
-
-        # launch(gridX, gridY, gridZ, stream, function, kernel_metadata,
-        # launch_metadata, launch_enter_hook, launch_exit_hook, *kernel_args)
-        launch_prefix = args[:9]
-        kernel_args = args[9:]
-        if len(kernel_args) != len(self._signature_types):
-            raise RuntimeError("launcher argument count mismatch while expanding tensor descriptors")
-
-        expanded_kernel_args = []
-        launch_keepalive = []
-        tensordesc_state = [0]
-        for arg, ty in zip(kernel_args, self._signature_types):
-            self._expand_runtime_arg(arg, ty, expanded_kernel_args, launch_keepalive, tensordesc_state)
-
-        self._tensordesc_keepalive.extend(launch_keepalive)
-        if len(self._tensordesc_keepalive) > 4096:
-            self._tensordesc_keepalive = self._tensordesc_keepalive[-4096:]
-        gridX, gridY, gridZ = launch_prefix[0], launch_prefix[1], launch_prefix[2]
-        stream = launch_prefix[3]
-        gs_buf = self._alloc_global_scratch(gridX, gridY, gridZ, stream)
-        self._set_and_keep_global_scratch(gs_buf)
-        self.launch(*launch_prefix, *expanded_kernel_args, **kwargs)
+            return self.launch(*launch_prefix, *expanded_kernel_args, *hidden_args, **kwargs)
+
+        if not self._debug_enabled:
+            return launch()
+
+        user_args = args[-self.user_arg_count:] if self.user_arg_count else ()
+        launch_metadata = args[6] if len(args) > 6 else None
+        # FlagPrism: wrap MUSA launch only for debugger-instrumented kernels.
+        launch_grid = (args[0], args[1], args[2])
+        with _flagprism.debugger_launch_context(
+                "musa",
+                self.metadata,
+                launch_grid,
+                args[3],
+                launch_metadata,
+                user_args,
+        ) as hidden_args:
+            return launch(hidden_args)
 
 
 class MusaDriver(DriverBase):

third_party/mthreads/bin/CMakeLists.txt

@@ -5,6 +5,9 @@ add_executable(triton-opt triton-opt.cpp)
 target_compile_options(triton-opt PRIVATE -fno-rtti -fno-exceptions)
 target_link_libraries(triton-opt PRIVATE
   ${triton_libs}
+  # FlagPrism: link the external Proton test-pass registration used by the
+  # unified RegisterTritonDialects path.
+  $<$<BOOL:${TRITON_BUILD_FLAGPRISM}>:TritonTestProton>
   # MLIR core
   MLIROptLib
   MLIRPass
@@ -21,6 +24,9 @@ target_compile_options(triton-reduce PRIVATE -fno-rtti -fno-exceptions)
 
 target_link_libraries(triton-reduce PRIVATE
   ${triton_libs}
+  # FlagPrism: link the external Proton test-pass registration used by the
+  # unified RegisterTritonDialects path.
+  $<$<BOOL:${TRITON_BUILD_FLAGPRISM}>:TritonTestProton>
   # MLIR core
   MLIRReduceLib
   MLIRPass
@@ -52,6 +58,9 @@ add_executable(triton-tensor-layout triton-tensor-layout.cpp)
 target_compile_options(triton-tensor-layout PRIVATE -fno-rtti -fno-exceptions)
 target_link_libraries(triton-tensor-layout PRIVATE
   ${triton_libs}
+  # FlagPrism: link the external Proton test-pass registration used by the
+  # unified RegisterTritonDialects path.
+  $<$<BOOL:${TRITON_BUILD_FLAGPRISM}>:TritonTestProton>
   MLIRRegisterAllDialects
   MLIRRegisterAllPasses
   MLIRTransforms

third_party/mthreads/bin/RegisterTritonDialects.h

@@ -14,12 +14,19 @@
 #include "MTGPUToLLVM/Passes.h"
 #include "TritonMUSAGPUToLLVM/Passes.h"
 #include "TritonMUSAGPUTransforms/Passes.h"
+// FlagPrism: select the external component's dialect registration, matching
+// the shared FlagTree tool registration path.
+#ifdef __FLAGPRISM__
+// FlagPrism: use the external component's replacement Proton registration.
+#include "Integration/Registration.h" // FlagPrism: use the exported include root.
+#else
 #include "proton/Dialect/include/Conversion/ProtonGPUToLLVM/Passes.h"
 #include "proton/Dialect/include/Conversion/ProtonGPUToLLVM/ProtonNvidiaGPUToLLVM/Passes.h"
 #include "proton/Dialect/include/Conversion/ProtonToProtonGPU/Passes.h"
 #include "proton/Dialect/include/Dialect/Proton/IR/Dialect.h"
 #include "proton/Dialect/include/Dialect/ProtonGPU/IR/Dialect.h"
 #include "proton/Dialect/include/Dialect/ProtonGPU/Transforms/Passes.h"
+#endif
 #include "triton/Dialect/Gluon/Transforms/Passes.h"
 #include "triton/Dialect/NVGPU/IR/Dialect.h"
 #include "triton/Dialect/NVWS/IR/Dialect.h"
@@ -91,6 +98,13 @@ inline void registerTritonDialects(mlir::DialectRegistry &registry) {
   mlir::triton::registerMTGPUToLLVMPasses();
   mlir::registerTritonMUSAGPUPasses();
 
+  // FlagPrism: use the same external pass and dialect registration entry
+  // points as the shared FlagTree tools.
+#ifdef __FLAGPRISM__
+  mlir::triton::proton::registerFlagTreeProtonTestPasses();
+  mlir::triton::proton::registerFlagTreeProtonPassesAndDialects(registry);
+#endif
+
   // Plugin passes
   if (std::string filename =
           mlir::triton::tools::getStrEnv("TRITON_PASS_PLUGIN_PATH");

third_party/mthreads/spec/triton/compiler/code_generator.py

@@ -12,6 +12,7 @@
 from typing import Any, Callable, Dict, Optional, Tuple, Type, Union, Iterable, List
 
 from .. import knobs, language
+from flagtree import _flagprism  # FlagPrism
 from .._C.libtriton import ir, gluon_ir
 from ..language import constexpr, str_to_ty, tensor, tuple as tl_tuple
 from ..language.core import _unwrap_if_constexpr, base_value, base_type
@@ -726,6 +727,8 @@ def _sanitize_value(value):
                 values = _sanitize_value(self.visit(node.value))
         else:
             values = _sanitize_value(self.visit(node.value))
+        # FlagPrism: emit normalized statement data before symbol binding.
+        _flagprism.emit_statement_event("assignment", self, node, target, values)
         self.assignTarget(target, values)
 
     def visit_AugAssign(self, node):
@@ -1591,7 +1594,11 @@ def visit_Attribute(self, node):
 
     def visit_Expr(self, node):
         node.value._is_unused = True
-        ast.NodeVisitor.generic_visit(self, node)
+        # FlagPrism: retain the original traversal behavior for reference.
+        # ast.NodeVisitor.generic_visit(self, node)
+        value = self.visit(node.value)
+        # FlagPrism: retain the operation created by a void expression.
+        _flagprism.emit_statement_event("expression", self, node, None, value)
 
     def visit_NoneType(self, node):
         return None