[PROTON][NFC] Clean up code (triton-lang#5109)

Including typos, unused function, README, tool usage, and function
simplification

Author: Jokeren

third_party/proton/README.md

@@ -141,14 +141,17 @@ By default, proton profiles are in the *json* format and can be read by *Hatchet
 pip install llnl-hatchet
 proton-viewer -m time/s <profile.hatchet>
 ```
+
 NOTE: `pip install hatchet` does not work because the API is slightly different.
 
 ### Visualizing sorted profile data
+
 In addition visualizing the profile data on terminal through Hatchet. A sorted list of the kernels by the first metric can be done using the --print-sorted flag with proton-viewer
 
 ```bash
 proton-viewer -m time/ns,time/% <profile.hatchet> --print-sorted
 ```
+
 prints the sorted kernels by the time/ns since it is the first listed.
 
 More options can be found by running the following command.
@@ -157,21 +160,27 @@ More options can be found by running the following command.
 proton-viewer -h
 ```
 
-### Advanced features
+## Advanced features
+
+### Instrumentation (experimental)
+
 In addition to profiling, Proton also incorporates MLIR/LLVM based compiler instrumentation passes to get Triton level analysis
 and optimization information. This feature is under active development and the list of available passes is expected to grow.
 
 #### Available passes
+
 print-mem-spaces: this pass prints the load and store address spaces (e.g. global, flat, shared) chosen by the compiler and attributes back to Triton source information.
 
 Example usage with the Proton matmul tutorial:
+
 ```bash
 $ proton --instrument=print-mem-spaces matmul.py
 0     matmul_kernel     matmul.py:180:20     SHARED     STORE
 1     matmul_kernel     matmul.py:181:20     SHARED     STORE
 2     matmul_kernel     matmul.py:180:20     SHARED     LOAD
 3     matmul_kernel     matmul.py:181:20     SHARED     LOAD
 ```
+
 Notes: The instrument functionality is currently only available from the command line. Additionally the instrument and profile command line arguments can not be use simulantously.
 
 ### Instruction sampling (experimental)

third_party/proton/csrc/include/Profiler/GPUProfiler.h

@@ -58,13 +58,15 @@ class GPUProfiler : public Profiler,
 
     ThreadState(ConcreteProfilerT &profiler) : profiler(profiler) {}
 
-    void record(size_t scopeId) {
+    size_t record() {
+      auto scopeId = Scope::getNewScopeId();
       if (profiler.isOpInProgress())
-        return;
+        return scopeId;
       std::set<Data *> dataSet = profiler.getDataSet();
       for (auto data : dataSet)
         data->addScope(scopeId);
       profiler.correlation.apiExternIds.insert(scopeId);
+      return scopeId;
     }
 
     void enterOp(size_t scopeId) {

third_party/proton/csrc/include/Session/Session.h

@@ -75,7 +75,7 @@ class SessionManager : public Singleton<SessionManager> {
 
   void finalizeAllSessions(OutputFormat outputFormat);
 
-  void activateSession(size_t sesssionId);
+  void activateSession(size_t sessionId);
 
   void deactivateSession(size_t sessionId);
 
@@ -97,7 +97,7 @@ class SessionManager : public Singleton<SessionManager> {
                                        const std::string &contextSourceName,
                                        const std::string &dataName);
 
-  void activateSessionImpl(size_t sesssionId);
+  void activateSessionImpl(size_t sessionId);
 
   void deActivateSessionImpl(size_t sessionId);
 
@@ -135,7 +135,7 @@ class SessionManager : public Singleton<SessionManager> {
   // path -> session id
   std::map<std::string, size_t> sessionPaths;
   // session id -> active
-  std::map<size_t, bool> activeSessions;
+  std::map<size_t, bool> sessionActive;
   // session id -> session
   std::map<size_t, std::unique_ptr<Session>> sessions;
   // scope -> active count

third_party/proton/csrc/lib/Profiler/Cupti/CuptiProfiler.cpp

@@ -323,8 +323,7 @@ void CuptiProfiler::CuptiProfilerPimpl::callbackFn(void *userData,
         static_cast<const CUpti_CallbackData *>(cbData);
     auto *pImpl = dynamic_cast<CuptiProfilerPimpl *>(profiler.pImpl.get());
     if (callbackData->callbackSite == CUPTI_API_ENTER) {
-      auto scopeId = Scope::getNewScopeId();
-      threadState.record(scopeId);
+      auto scopeId = threadState.record();
       threadState.enterOp(scopeId);
       size_t numInstances = 1;
       if (cbId == CUPTI_DRIVER_TRACE_CBID_cuGraphLaunch ||

third_party/proton/csrc/lib/Profiler/RocTracer/RoctracerProfiler.cpp

@@ -231,8 +231,7 @@ void RoctracerProfiler::RoctracerProfilerPimpl::apiCallback(
     const hip_api_data_t *data = (const hip_api_data_t *)(callbackData);
     if (data->phase == ACTIVITY_API_PHASE_ENTER) {
       // Valid context and outermost level of the kernel launch
-      auto scopeId = Scope::getNewScopeId();
-      threadState.record(scopeId);
+      auto scopeId = threadState.record();
       threadState.enterOp(scopeId);
       size_t numInstances = 1;
       if (cid == HIP_API_ID_hipGraphLaunch) {

third_party/proton/csrc/lib/Session/Session.cpp

@@ -91,20 +91,20 @@ void SessionManager::deactivateSession(size_t sessionId) {
 
 void SessionManager::activateSessionImpl(size_t sessionId) {
   throwIfSessionNotInitialized(sessions, sessionId);
-  if (activeSessions[sessionId])
+  if (sessionActive[sessionId])
     return;
-  activeSessions[sessionId] = true;
+  sessionActive[sessionId] = true;
   sessions[sessionId]->activate();
   registerInterface<ScopeInterface>(sessionId, scopeInterfaceCounts);
   registerInterface<OpInterface>(sessionId, opInterfaceCounts);
 }
 
 void SessionManager::deActivateSessionImpl(size_t sessionId) {
   throwIfSessionNotInitialized(sessions, sessionId);
-  if (!activeSessions[sessionId]) {
+  if (!sessionActive[sessionId]) {
     return;
   }
-  activeSessions[sessionId] = false;
+  sessionActive[sessionId] = false;
   sessions[sessionId]->deactivate();
   unregisterInterface<ScopeInterface>(sessionId, scopeInterfaceCounts);
   unregisterInterface<OpInterface>(sessionId, opInterfaceCounts);
@@ -204,7 +204,7 @@ void SessionManager::addMetrics(
     size_t scopeId, const std::map<std::string, MetricValueType> &metrics,
     bool aggregable) {
   std::shared_lock<std::shared_mutex> lock(mutex);
-  for (auto [sessionId, active] : activeSessions) {
+  for (auto [sessionId, active] : sessionActive) {
     if (active) {
       sessions[sessionId]->data->addMetrics(scopeId, metrics, aggregable);
     }

third_party/proton/proton/proton.py

@@ -1,7 +1,6 @@
 import argparse
 import sys
 import os
-from glob import glob
 import pathlib
 from .profile import start, finalize, _select_backend
 from .flags import set_command_line
@@ -91,7 +90,6 @@ def run_profiling(args, target_args):
 
 
 def run_instrumentation(args, target_args):
-    backend = args.backend if args.backend else _select_backend()
     do_setup_and_execute(target_args, args.instrument)
 
 

third_party/proton/proton/viewer.py

@@ -92,7 +92,7 @@ def get_min_time_bytes(df, device_info):
 }
 
 # FLOPS have a specific width to their metric
-default_flop_factor_dict = {f"flop/s": 1, f"gflop/s": 1e9, f"tflop/s": 1e12}
+default_flop_factor_dict = {"flop/s": 1, "gflop/s": 1e9, "tflop/s": 1e12}
 derivable_metrics.update(
     {key: FactorDict("flops", default_flop_factor_dict)
      for key in default_flop_factor_dict.keys()})
@@ -278,7 +278,7 @@ def main():
         type=str,
         default=None,
         help="""Exclude frames that match the given regular expression and their children.
-For example, the following command will exclude all paths that contain frames that contains "test":
+For example, the following command will exclude all paths starting from "test":
 ```
 proton-viewer -e ".*test.*" path/to/file.json
 ```