Improve useless stack frame deletion algorithm more (#58)

pytorch · Jul 15, 2024 · e582e38 · e582e38
1 parent 588dd9d
commit e582e38
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 16 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -44,24 +44,33 @@ impl Default for ParseConfig {
 }
 
 fn maybe_remove_convert_frame_suffixes(frames: &mut Vec<FrameSummary>) {
-    let target_frames = [
-        ("torch/_dynamo/convert_frame.py", "catch_errors"),
-        ("torch/_dynamo/convert_frame.py", "_convert_frame"),
-        ("torch/_dynamo/convert_frame.py", "_convert_frame_assert"),
+    let all_target_frames = [
+        [
+            ("torch/_dynamo/convert_frame.py", "catch_errors"),
+            ("torch/_dynamo/convert_frame.py", "_convert_frame"),
+            ("torch/_dynamo/convert_frame.py", "_convert_frame_assert"),
+        ],
+        [
+            ("torch/_dynamo/convert_frame.py", "__call__"),
+            ("torch/_dynamo/convert_frame.py", "__call__"),
+            ("torch/_dynamo/convert_frame.py", "__call__"),
+        ],
     ];
 
     let len = frames.len();
-    if len >= target_frames.len() {
-        let suffix = &frames[len - target_frames.len()..];
-        if suffix
-            .iter()
-            .zip(target_frames.iter())
-            .all(|(frame, target)| {
-                simplify_filename(unintern_str(frame.filename).as_ref()) == target.0
-                    && frame.name == target.1
-            })
-        {
-            frames.truncate(len - target_frames.len());
+    for target_frames in all_target_frames {
+        if len >= target_frames.len() {
+            let suffix = &frames[len - target_frames.len()..];
+            if suffix
+                .iter()
+                .zip(target_frames.iter())
+                .all(|(frame, target)| {
+                    simplify_filename(unintern_str(frame.filename).as_ref()) == target.0
+                        && frame.name == target.1
+                })
+            {
+                frames.truncate(len - target_frames.len());
+            }
         }
     }
 }

diff --git a/src/types.rs b/src/types.rs
@@ -160,7 +160,7 @@ pub fn simplify_filename<'a>(filename: &'a str) -> &'a str {
     if parts.len() > 1 {
         return parts[1];
     }
-    let re = Regex::new(r"\d+e\d+-seed-nspid\d+_cgpid\d+-ns-\d+/").unwrap();
+    let re = Regex::new(r"[^/]+-seed-nspid[^/]+/").unwrap();
     if let Some(captures) = re.captures(filename) {
         if let Some(capture) = captures.get(0) {
             return &filename[capture.end()..];