Reapply "[PROTON] Improve hook manager tests and fix a session_id=0 problem (#7745)" (#4960)

This reverts commit
f93ddb1.

Author: whitneywhtsang

third_party/proton/proton/hooks/hook.py

@@ -103,10 +103,10 @@ def register(hook: Hook, session: int) -> None:
 
     @staticmethod
     def unregister(session: Optional[int] = None) -> None:
-        if session and session not in HookManager.session_hooks:
+        if session is not None and session not in HookManager.session_hooks:
             return
 
-        if not session:
+        if session is None:
             for hook in HookManager.active_hooks:
                 hook.deactivate()
             HookManager.active_hooks.clear()

third_party/proton/proton/hooks/instrumentation.py

@@ -123,6 +123,8 @@ def _get_backend_name() -> str:
         return "nvidia"
     elif backend == "hip":
         return "amd"
+    elif backend == "xpu":
+        return "intel"
     else:
         raise RuntimeError(f"Unsupported backend: {backend}")
 

third_party/proton/test/test_api.py

@@ -2,6 +2,8 @@
 import triton.profiler as proton
 import pathlib
 from triton.profiler.hooks.hook import HookManager
+from triton.profiler.hooks.launch import LaunchHook
+from triton.profiler.hooks.instrumentation import InstrumentationHook
 
 
 def test_profile_single_session(tmp_path: pathlib.Path):
@@ -117,6 +119,34 @@ def test_hook(tmp_path: pathlib.Path):
     assert temp_file.exists()
 
 
+def test_hook_manager(tmp_path: pathlib.Path):
+    # Launch hook is a singleton
+    HookManager.register(LaunchHook(), 0)
+    HookManager.register(LaunchHook(), 0)
+    assert len(HookManager.active_hooks) == 1
+    assert isinstance(HookManager.active_hooks[0], LaunchHook)
+    assert HookManager.session_hooks[0][HookManager.active_hooks[0]] is True
+
+    # Only unregister one session
+    HookManager.register(LaunchHook(), 1)
+    HookManager.unregister(0)
+    assert len(HookManager.active_hooks) == 1
+    HookManager.unregister(1)
+    assert len(HookManager.active_hooks) == 0
+
+    # Heterogenous hooks
+    HookManager.register(InstrumentationHook(""), 2)
+    HookManager.register(LaunchHook(), 2)
+    assert len(HookManager.active_hooks) == 2
+    # Launch hook has a higher priority
+    assert isinstance(HookManager.active_hooks[0], LaunchHook)
+    assert isinstance(HookManager.active_hooks[1], InstrumentationHook)
+    assert HookManager.session_hooks[2][HookManager.active_hooks[0]] is True
+    assert HookManager.session_hooks[2][HookManager.active_hooks[1]] is True
+    HookManager.unregister()
+    assert len(HookManager.active_hooks) == 0
+
+
 def test_scope_metrics(tmp_path: pathlib.Path):
     temp_file = tmp_path / "test_scope_metrics.hatchet"
     session_id = proton.start(str(temp_file.with_suffix("")))