[PROTON-DEV] Add the instrumentation mode and clean up dependencies (#5742)

https://github.com/orgs/triton-lang/projects/6?pane=issue&itemId=95371680

Author: Jokeren

third_party/proton/proton/hook.py

@@ -1,13 +1,15 @@
 from .state import enter_state, exit_state
 from .scope import enter_scope, exit_scope
-from triton.compiler import CompiledKernel, LazyDict
 
 COMPUTE_METADATA_SCOPE_NAME = "__proton_launch_metadata"
 
 
 class TritonHook:
+    from triton.compiler import LazyDict
+
     flops_width = [8, 16, 32, 64]
-    metrics = [f"flops{width}" for width in flops_width] + ["bytes"] + ["flops"]
+    metrics = [f"flops{width}" for width in flops_width] + \
+        ["bytes"] + ["flops"]
 
     @staticmethod
     def enter(lazy_dict: LazyDict) -> None:
@@ -23,12 +25,14 @@ def exit(lazy_dict: LazyDict) -> None:
 
 
 def register_triton_hook() -> None:
+    from triton.compiler import CompiledKernel
     if CompiledKernel.launch_enter_hook is None:
         CompiledKernel.launch_enter_hook = TritonHook.enter
         CompiledKernel.launch_exit_hook = TritonHook.exit
 
 
 def unregister_triton_hook() -> None:
+    from triton.compiler import CompiledKernel
     if CompiledKernel.launch_enter_hook == TritonHook.enter:
         CompiledKernel.launch_enter_hook = None
         CompiledKernel.launch_exit_hook = None

third_party/proton/proton/profile.py

@@ -43,12 +43,25 @@ def _check_env(backend: str) -> None:
                 )
 
 
+def _check_mode(backend: str, mode: Optional[str]) -> None:
+    # TODO(Keren): Need a better mode registration mechanism
+    backend_modes = {
+        "cupti": [None, "pcsampling"],
+        "roctracer": [None],
+        "instrumentation": [None],
+    }
+
+    if mode not in backend_modes[backend]:
+        raise ValueError(f"Invalid mode {mode} for backend {backend}")
+
+
 def start(
     name: Optional[str] = None,
     *,
     context: Optional[str] = "shadow",
     data: Optional[str] = "tree",
     backend: Optional[str] = None,
+    mode: Optional[str] = None,
     hook: Optional[str] = None,
 ):
     """
@@ -65,15 +78,20 @@ def start(
     Args:
         name (str, optional): The name (with path) of the profiling session.
                               If not provided, the default name is "~/proton.hatchet".
-        backend (str, optional): The backend to use for profiling.
-                                 Available options are [None, "cupti", "cupti_pcsampling", "roctracer"].
-                                 Defaults to None, which automatically selects the backend matching the current active runtime.
         context (str, optional): The context to use for profiling.
                                  Available options are ["shadow", "python"].
                                  Defaults to "shadow".
         data (str, optional): The data structure to use for profiling.
                               Available options are ["tree"].
                               Defaults to "tree".
+        backend (str, optional): The backend to use for profiling.
+                                 Available options are [None, "cupti", "roctracer", "instrumentation"].
+                                 Defaults to None, which automatically selects the backend matching the current active runtime.
+        mode (str, optional): The "mode" to use for profiling, which is specific to the backend.
+                              Defaults to None.
+                              For "cupti", available options are [None, "pcsampling"].
+                              For "roctracer", available options are [None].
+                              For "instrumentation", available options are [None].
         hook (str, optional): The hook to use for profiling.
                               Available options are [None, "triton"].
                               Defaults to None.
@@ -91,6 +109,7 @@ def start(
         backend = _select_backend()
 
     _check_env(backend)
+    _check_mode(backend, mode)
 
     backend_path = _get_backend_default_path(backend)
 
@@ -167,6 +186,7 @@ def _profiling(
     context: Optional[str] = "shadow",
     data: Optional[str] = "tree",
     backend: Optional[str] = None,
+    mode: Optional[str] = None,
     hook: Optional[str] = None,
 ):
     """
@@ -181,7 +201,7 @@ def _profiling(
 
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
-        session = start(name, context=context, data=data, backend=backend, hook=hook)
+        session = start(name, context=context, data=data, backend=backend, mode=mode, hook=hook)
         ret = func(*args, **kwargs)
         deactivate(session)
         return ret
@@ -196,6 +216,7 @@ def profile(
     context: Optional[str] = "shadow",
     data: Optional[str] = "tree",
     backend: Optional[str] = None,
+    mode: Optional[str] = None,
     hook: Optional[str] = None,
 ):
     """
@@ -218,9 +239,9 @@ def foo():
     if func is None:
         # It's being used with parentheses, so return a decorator
         def decorator(f):
-            return _profiling(f, name=name, context=context, data=data, backend=backend, hook=hook)
+            return _profiling(f, name=name, context=context, data=data, backend=backend, mode=mode, hook=hook)
 
         return decorator
     else:
         # It's being used without parentheses, so apply the decorator directly
-        return _profiling(func, name=name, context=context, data=data, backend=backend, hook=hook)
+        return _profiling(func, name=name, context=context, data=data, backend=backend, mode=mode, hook=hook)

third_party/proton/proton/proton.py

@@ -16,9 +16,10 @@ def parse_arguments():
 """, formatter_class=argparse.RawTextHelpFormatter)
     parser.add_argument("-n", "--name", type=str, help="Name of the profiling session")
     parser.add_argument("-b", "--backend", type=str, help="Profiling backend", default=None,
-                        choices=["cupti", "cupti_pcsampling", "roctracer"])
+                        choices=["cupti", "roctracer"])
     parser.add_argument("-c", "--context", type=str, help="Profiling context", default="shadow",
                         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"])
     parser.add_argument("-k", "--hook", type=str, help="Profiling hook", default=None, choices=[None, "triton"])
     parser.add_argument("-i", "--instrument", type=str, help="Instrumentation analysis type", default=None,

third_party/proton/test/test_profile.py

@@ -256,7 +256,7 @@ def foo(x, y, size: tl.constexpr):
             tl.store(y + offs, tl.load(x + offs))
 
     temp_file = tmp_path / "test_pcsampling.hatchet"
-    proton.start(str(temp_file.with_suffix("")), hook="triton", backend="cupti_pcsampling")
+    proton.start(str(temp_file.with_suffix("")), hook="triton", backend="cupti", mode="pcsampling")
     with proton.scope("init"):
         x = torch.ones((1024, ), device="cuda", dtype=torch.float32)
         y = torch.zeros_like(x)

third_party/proton/tutorials/dynamic_net.py

@@ -5,7 +5,7 @@
 import triton.profiler as proton
 import argparse
 
-mode = "torch"
+engine = "torch"
 
 
 class DynamicNet(torch.nn.Module):
@@ -53,7 +53,7 @@ def run():
 
         # Construct our model by instantiating the class defined above
         model = DynamicNet().to("cuda")
-        if mode == "torchinductor":
+        if engine == "torchinductor":
             model = torch.compile(model)
 
         # Construct our loss function and an Optimizer. Training this strange model with
@@ -83,16 +83,17 @@ def run():
 
 argparser = argparse.ArgumentParser()
 argparser.add_argument("--profile", action="store_true")
-argparser.add_argument("--mode", default="torch", choices=["torch", "torchinductor"])
+argparser.add_argument("--engine", default="torch", choices=["torch", "torchinductor"])
 argparser.add_argument("--context", default="shadow", choices=["shadow", "python"])
-argparser.add_argument("--backend", default=None, choices=["cupti", "roctracer", "cupti_pcsampling"])
+argparser.add_argument("--backend", default=None, choices=["cupti", "roctracer"])
+argparser.add_argument("--mode", default=None)
 
 args = argparser.parse_args()
 
-mode = args.mode
+engine = args.engine
 
 if args.profile:
-    func = proton.profile(run, name="dynamic_net", context=args.context, backend=args.backend)
+    func = proton.profile(run, name="dynamic_net", context=args.context, backend=args.backend, mode=args.mode)
 else:
     func = run
 

third_party/proton/tutorials/matmul.py

@@ -308,7 +308,7 @@ def perf(ms):
 if args.profile:
     if args.pcsampling:
         # proton-viewer -m num_samples/%,time/s ./matmul.hatchet
-        proton.start("matmul", hook="triton", backend="cupti_pcsampling")
+        proton.start("matmul", hook="triton", backend="cupti", mode="pcsampling")
     else:
         # proton-viewer -m tflop/s,time/s ./matmul.hatchet
         proton.start("matmul", hook="triton")