[PROTON] Add a flag to disable proton in order to use other profilers without commenting out code (#8293)

Author: Jokeren

python/triton/knobs.py

@@ -515,6 +515,7 @@ class amd_knobs(base_knobs):
 
 
 class proton_knobs(base_knobs):
+    disable: env_bool = env_bool("TRITON_PROTON_DISABLE", False)
     cupti_lib_dir: env_str = env_str(
         "TRITON_CUPTI_LIB_PATH",
         str(pathlib.Path(__file__).parent.absolute() / "backends" / "nvidia" / "lib" / "cupti"))

third_party/proton/proton/flags.py

@@ -17,7 +17,7 @@
 
 @dataclass
 class ProfilerFlags:
-    # Whether to enable profiling. Default is False.
+    # Whether profiling is enabled. Default is False.
     profiling_on: bool = False
     # Whether instrumentation is enabled. Default is False.
     instrumentation_on: bool = False

third_party/proton/proton/profile.py

@@ -94,8 +94,8 @@ def start(
     Returns:
         session (int): The session ID of the profiling session.
     """
-    if flags.command_line:
-        # Ignore the start() call if the script is run from the command line.
+    if flags.command_line or triton.knobs.proton.disable:
+        # Ignore the start() call if the script is run from the command line or profiling is disabled.
         return
 
     flags.profiling_on = True

third_party/proton/proton/proton.py

@@ -19,7 +19,7 @@ def parse_arguments():
                         choices=["shadow", "python"])
     parser.add_argument("-m", "--mode", type=str, help="Profiling mode", default=None)
     parser.add_argument("-d", "--data", type=str, help="Profiling data", default="tree", choices=["tree", "trace"])
-    parser.add_argument("-k", "--hook", type=str, help="Profiling hook", default=None, choices=[None, "launch"])
+    parser.add_argument("-k", "--hook", type=str, help="Profiling hook", default=None, choices=[None, "triton"])
     parser.add_argument('target_args', nargs=argparse.REMAINDER, help='Subcommand and its arguments')
     args = parser.parse_args()
     return args, args.target_args

third_party/proton/test/test_api.py

@@ -4,6 +4,7 @@
 Profile correctness tests involving GPU kernels should be placed in `test_profile.py`.
 """
 
+import pytest
 import json
 import triton.profiler as proton
 import pathlib
@@ -385,3 +386,17 @@ def test_throw(tmp_path: pathlib.Path):
     finally:
         proton.finalize()
     assert "Session has not been initialized: " + str(session_id + 1) in deactivate_error
+
+
+@pytest.mark.parametrize("disable", [True, False])
+def test_profile_disable(disable, fresh_knobs, tmp_path: pathlib.Path):
+    fresh_knobs.proton.disable = disable
+    temp_file = tmp_path / "test_profile_disable.hatchet"
+    proton.start(str(temp_file.with_suffix("")))
+    proton.enter_scope("test0")
+    proton.exit_scope()
+    proton.finalize()
+    if disable:
+        assert not temp_file.exists()
+    else:
+        assert temp_file.exists()