[PROTON] Fix scope related issues and add more tests

third_party/proton/README.md

@@ -10,7 +10,7 @@ The following command installs the latest version of Proton.
 
 ```bash
 git clone https://github.com/openai/triton
-cd triton/third_party/proton
+cd triton/python
 TRITON_BUILD_PROTON=1 pip install .
 ```
 

third_party/proton/csrc/lib/Data/TreeData.cpp

@@ -31,9 +31,9 @@ void TreeData::stopOp(const Scope &scope) {}
 void TreeData::addMetric(size_t scopeId, std::shared_ptr<Metric> metric) {
   std::unique_lock<std::shared_mutex> lock(mutex);
   auto scopeIdIt = scopeIdToContextId.find(scopeId);
-  if (scopeIdIt == scopeIdToContextId.end()) {
-    throw std::runtime_error("scopeId not found");
-  }
+  // The profile data is deactived, ignore the metric
+  if (scopeIdIt == scopeIdToContextId.end())
+    return;
   auto contextId = scopeIdIt->second;
   auto &node = tree->getNode(contextId);
   if (node.metrics.find(metric->getKind()) == node.metrics.end())
@@ -48,10 +48,9 @@ void TreeData::addMetrics(
   auto scopeIdIt = scopeIdToContextId.find(scopeId);
   auto contextId = Tree::TreeNode::DummyId;
   if (scopeIdIt == scopeIdToContextId.end()) {
-    if (contextSource == nullptr) {
-      // TODO: This is a hack, we should attribute it to an unknown context
-      throw std::runtime_error("scopeId not found");
-    }
+    if (contextSource == nullptr)
+      throw std::runtime_error("ContextSource is not set");
+    // Attribute the metric to the last context
     std::vector<Context> contexts = contextSource->getContexts();
     contextId = tree->addNode(contexts);
   } else {

third_party/proton/proton/scope.py

@@ -6,7 +6,6 @@
 from triton._C.libproton import proton as libproton
 
 _local = threading.local()
-_local.scopes = []
 
 
 class scope:
@@ -71,6 +70,8 @@ def enter_scope(name: str, *, triton_op: bool = False, metrics: Optional[dict[st
     if not get_profiling_on():
         return -1
     id = libproton.record_scope()
+    if not hasattr(_local, "scopes"):
+        _local.scopes = []
     _local.scopes.append((id, name))
     if triton_op:
         libproton.enter_op(id, name)

third_party/proton/proton/viewer.py

@@ -102,7 +102,7 @@ def derive_metrics(gf, metrics, raw_metrics):
     return derived_metrics + original_metrics
 
 
-def parse(metrics, filename, include, exclude, threshold):
+def parse(metrics, filename, include, exclude, threshold, depth):
     with open(filename, "r") as f:
         gf, raw_metrics = get_raw_metrics(f)
         assert len(raw_metrics) > 0, "No metrics found in the input file"
@@ -120,12 +120,12 @@ def parse(metrics, filename, include, exclude, threshold):
             # TODO: generalize to support multiple metrics
             query = ["*", {metrics[0]: f">= {threshold}"}]
             gf = gf.filter(query, squash=True)
-        print(gf.tree(metric_column=metrics, expand_name=True))
+        print(gf.tree(metric_column=metrics, expand_name=True, depth=depth))
 
 
 def show_metrics(file_name):
     with open(file_name, "r") as f:
-        gf, raw_metrics = get_raw_metrics(f)
+        _, raw_metrics = get_raw_metrics(f)
         print("Available metrics:")
         if raw_metrics:
             print("\n".join(raw_metrics))
@@ -183,6 +183,14 @@ def main():
         "Exclude frames(kernels) whose metrics are below the given threshold. This filter only applies on the first metric.",
     )
 
+    argparser.add_argument(
+        "-d",
+        "--depth",
+        type=int,
+        default=100,
+        help="The depth of the tree to display",
+    )
+
     args, target_args = argparser.parse_known_args()
     assert len(target_args) == 1, "Must specify a file to read"
 
@@ -191,9 +199,10 @@ def main():
     include = args.include
     exclude = args.exclude
     threshold = args.threshold
+    depth = args.depth
     if include and exclude:
         raise ValueError("Cannot specify both include and exclude")
     if args.list:
         show_metrics(file_name)
     elif metrics:
-        parse(metrics, file_name, include, exclude, threshold)
+        parse(metrics, file_name, include, exclude, threshold, depth)

third_party/proton/test/test_api.py

@@ -1,3 +1,4 @@
+import json
 import triton.profiler as proton
 import tempfile
 import pathlib
@@ -78,17 +79,28 @@ def test_hook():
 
 def test_scope_metrics():
     with tempfile.NamedTemporaryFile(delete=True, suffix=".hatchet") as f:
-        proton.start(f.name.split(".")[0])
-        with proton.scope("test", {"a": 1.0}):
+        session_id = proton.start(f.name.split(".")[0])
+        # Test different scope creation methods
+        with proton.scope("test0", {"a": 1.0}):
             pass
 
-        @proton.scope("test", {"a": 1.0})
+        @proton.scope("test1", {"a": 1.0})
         def foo():
             pass
 
         foo()
 
-        proton.enter_scope("test", metrics={"a": 1.0})
+        # After deactivation, the metrics should be ignored
+        proton.deactivate(session_id)
+        proton.enter_scope("test2", metrics={"a": 1.0})
+        proton.exit_scope()
+
+        # Metrics should be recorded again after reactivation
+        proton.activate(session_id)
+        proton.enter_scope("test3", metrics={"a": 1.0})
         proton.exit_scope()
+
         proton.finalize()
         assert pathlib.Path(f.name).exists()
+        data = json.load(f)
+        assert len(data[0]["children"]) == 3

third_party/proton/test/test_profile.py

@@ -38,13 +38,9 @@ def foo(x, y):
         proton.start(f.name.split(".")[0])
         with proton.scope("test0"):
             with proton.scope("test1"):
-                foo[
-                    1,
-                ](x, y)
+                foo[(1, )](x, y)
         with proton.scope("test2"):
-            foo[
-                1,
-            ](x, y)
+            foo[(1, )](x, y)
         proton.finalize()
         data = json.load(f)
         assert len(data[0]["children"]) == 2
@@ -65,16 +61,51 @@ def foo(x, y):
     with tempfile.NamedTemporaryFile(delete=True, suffix=".hatchet") as f:
         proton.start(f.name.split(".")[0])
         with proton.scope("test0", {"foo": 1.0}):
-            foo[
-                1,
-            ](x, y)
+            foo[(1, )](x, y)
         proton.finalize()
         data = json.load(f)
         assert len(data[0]["children"]) == 1
         assert data[0]["children"][0]["frame"]["name"] == "test0"
         assert data[0]["children"][0]["metrics"]["foo"] == 1.0
 
 
+def test_metrics_ignore():
+
+    @triton.jit
+    def foo(x, y):
+        tl.store(y, tl.load(x))
+
+    x = torch.tensor([2], device="cuda")
+    y = torch.zeros_like(x)
+    with tempfile.NamedTemporaryFile(delete=True, suffix=".hatchet") as f:
+        session_id = proton.start(f.name.split(".")[0])
+        proton.deactivate(session_id)
+        with proton.scope("test0", {"foo": 1.0}):
+            foo[(1, )](x, y)
+        proton.activate(session_id)
+        proton.finalize()
+        data = json.load(f)
+        assert len(data[0]["children"]) == 0
+
+
+def test_scope_backward():
+    with tempfile.NamedTemporaryFile(delete=True, suffix=".hatchet") as f:
+        proton.start(f.name.split(".")[0])
+        with proton.scope("ones1"):
+            a = torch.ones((100, 100), device="cuda", requires_grad=True)
+        with proton.scope("plus"):
+            a2 = a * a * a
+        with proton.scope("ones2"):
+            loss = torch.ones_like(a2)
+
+        # Backward triggers two kernels in a single scope
+        with proton.scope("backward"):
+            a2.backward(loss)
+        proton.finalize()
+        data = json.load(f)
+        assert len(data[0]["children"]) == 4
+
+
 def test_hook():
 
     def metadata_fn(grid: tuple, metadata: NamedTuple, args: dict):
@@ -95,9 +126,7 @@ def foo(x, size: tl.constexpr, y):
     with tempfile.NamedTemporaryFile(delete=True, suffix=".hatchet") as f:
         proton.start(f.name.split(".")[0], hook="triton")
         with proton.scope("test0"):
-            foo[
-                1,
-            ](x, 1, y, num_warps=4)
+            foo[(1, )](x, 1, y, num_warps=4)
         proton.finalize()
         data = json.load(f)
         assert len(data[0]["children"]) == 1