Skip to content

Improve tlparse runtime by caching format_stack results #173

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yushangdi merged 1 commit into from
Mar 12, 2026
Merged

Improve tlparse runtime by caching format_stack results #173

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
182 changes: 115 additions & 67 deletions src/parsers.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
use crate::templates::TEMPLATE_QUERY_PARAM_SCRIPT;
use crate::{types::*, ParseConfig};
use fxhash::FxHashMap;
use html_escape::encode_text;
use std::cell::RefCell;
use std::collections::HashSet;
Expand All @@ -20,9 +21,23 @@ fn format_json_pretty(payload: &str) -> Result<String, anyhow::Error> {
}
}

use std::sync::OnceLock;
use syntect::highlighting::ThemeSet;
use syntect::parsing::SyntaxSet;

struct SyntectResources {
syntax_set: SyntaxSet,
theme_set: ThemeSet,
}

fn syntect_resources() -> &'static SyntectResources {
static RESOURCES: OnceLock<SyntectResources> = OnceLock::new();
RESOURCES.get_or_init(|| SyntectResources {
syntax_set: SyntaxSet::load_defaults_newlines(),
theme_set: ThemeSet::load_defaults(),
})
}

// Re-export types from types.rs for external use
pub use crate::types::{CompileId, EmptyMetadata, Envelope, GraphRuntime, Metadata, OpRuntime};

Expand Down Expand Up @@ -311,14 +326,13 @@ impl StructuredLogParser for InductorOutputCodeParser {
}

fn generate_html_output(payload: &str) -> Result<String, anyhow::Error> {
let syntax_set = SyntaxSet::load_defaults_newlines();
let theme_set = ThemeSet::load_defaults();
let syntax = syntax_set.find_syntax_by_extension("py").unwrap();
let res = syntect_resources();
let syntax = res.syntax_set.find_syntax_by_extension("py").unwrap();
let html = syntect::html::highlighted_html_for_string(
&payload,
&syntax_set,
&res.syntax_set,
&syntax,
&theme_set.themes["InspiredGitHub"],
&res.theme_set.themes["InspiredGitHub"],
);
Ok(html?)
}
Expand Down Expand Up @@ -385,6 +399,23 @@ fn format_stack(stack: &StackSummary, caption: &str, open: bool) -> String {
trie.fmt(None, caption, open).unwrap()
}

// HashMap requires cloning the key for tuple lookups. If this becomes a perf
// bottleneck, switch to a 2-level cache (FxHashMap<StackSummary, FxHashMap<String, String>>)
// which wouldn't require cloning the key.
fn format_stack_cached(
cache: &mut FxHashMap<(StackSummary, String), String>,
stack: &StackSummary,
caption: &str,
) -> String {
let key = (stack.clone(), caption.to_string());
Copy link
Copy Markdown
Contributor

@aorenste aorenste Mar 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You're not including open in the cache key? Should probably include it - but it looks like it's always false? So you could just remove it from the function.

Copy link
Copy Markdown
Contributor Author

@yushangdi yushangdi Mar 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

good catch! removing it

Copy link
Copy Markdown
Contributor

@aorenste aorenste Mar 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't like these clones, but unfortunately HashMap has a limitation that requires it for a tuple key. If we find that this is a perf bottleneck later we could switch to a 2-level cache (i.e. FxHashMap<StackSummary, FxHashMap<String, String>>) which wouldn't require us to clone the key - but that's proabably not necessary.

if let Some(cached) = cache.get(&key) {
return cached.clone();
}
let result = format_stack(stack, caption, false);
cache.insert(key, result.clone());
Copy link
Copy Markdown
Contributor

@aorenste aorenste Mar 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We could avoid this clone by returning the copy of the String the cache contains as a &str - but again... probably overkill.

result
}

pub struct CompilationMetricsParser<'t> {
pub tt: &'t TinyTemplate<'t>,
pub stack_index: &'t RefCell<StackIndex>,
Expand Down Expand Up @@ -413,6 +444,7 @@ impl StructuredLogParser for CompilationMetricsParser<'_> {
_payload: &str,
) -> anyhow::Result<ParserResults> {
let filename = format!("{}.html", self.name());
let mut stack_cache: FxHashMap<(StackSummary, String), String> = FxHashMap::default();
if let Metadata::CompilationMetrics(m) = metrics {
let id = compile_id
.clone()
Expand Down Expand Up @@ -446,90 +478,106 @@ impl StructuredLogParser for CompilationMetricsParser<'_> {
} else {
"".to_string()
};
let specializations = self
let specializations: Vec<_> = self
.symbolic_shape_specialization_index
.borrow_mut()
.remove(&cid)
.unwrap_or(Vec::new())
.unwrap_or_default()
.drain(..)
.map(|spec| SymbolicShapeSpecializationContext {
symbol: spec.symbol.unwrap_or("".to_string()),
sources: spec.sources.unwrap_or(Vec::new()),
value: spec.value.unwrap_or("".to_string()),
user_stack_html: format_stack(
&spec.user_stack.unwrap_or(Vec::new()),
"User Stack",
false,
),
stack_html: format_stack(
&spec.stack.unwrap_or(Vec::new()),
"Framework Stack",
false,
),
.map(|spec| {
let user_stack = spec.user_stack.unwrap_or_default();
let stack = spec.stack.unwrap_or_default();
SymbolicShapeSpecializationContext {
symbol: spec.symbol.unwrap_or("".to_string()),
sources: spec.sources.unwrap_or_default(),
value: spec.value.unwrap_or("".to_string()),
user_stack_html: format_stack_cached(
&mut stack_cache,
&user_stack,
"User Stack",
),
stack_html: format_stack_cached(
&mut stack_cache,
&stack,
"Framework Stack",
),
}
})
.collect();
let guards_added_fast = self
let guards_added_fast: Vec<_> = self
.guard_added_fast_index
.borrow_mut()
.remove(&cid)
.unwrap_or(Vec::new())
.unwrap_or_default()
.drain(..)
.map(|guard| GuardAddedFastContext {
expr: guard.expr.unwrap_or("".to_string()),
user_stack_html: format_stack(
&guard.user_stack.unwrap_or(Vec::new()),
"User Stack",
false,
),
stack_html: format_stack(
&guard.stack.unwrap_or(Vec::new()),
"Framework Stack",
false,
),
.map(|guard| {
let user_stack = guard.user_stack.unwrap_or_default();
let stack = guard.stack.unwrap_or_default();
GuardAddedFastContext {
expr: guard.expr.unwrap_or("".to_string()),
user_stack_html: format_stack_cached(
&mut stack_cache,
&user_stack,
"User Stack",
),
stack_html: format_stack_cached(
&mut stack_cache,
&stack,
"Framework Stack",
),
}
})
.collect();
let create_symbols = self
let create_symbols: Vec<_> = self
.create_symbol_index
.borrow_mut()
.remove(&cid)
.unwrap_or(Vec::new())
.unwrap_or_default()
.drain(..)
.map(|sym| CreateSymbolContext {
symbol: sym.symbol.unwrap_or("".to_string()),
val: sym.val.unwrap_or("".to_string()),
vr: sym.vr.unwrap_or("".to_string()),
source: sym.source.unwrap_or("".to_string()),
user_stack_html: format_stack(
&sym.user_stack.unwrap_or(Vec::new()),
"User Stack",
false,
),
stack_html: format_stack(
&sym.stack.unwrap_or(Vec::new()),
"Framework Stack",
false,
),
.map(|sym| {
let user_stack = sym.user_stack.unwrap_or_default();
let stack = sym.stack.unwrap_or_default();
CreateSymbolContext {
symbol: sym.symbol.unwrap_or("".to_string()),
val: sym.val.unwrap_or("".to_string()),
vr: sym.vr.unwrap_or("".to_string()),
source: sym.source.unwrap_or("".to_string()),
user_stack_html: format_stack_cached(
&mut stack_cache,
&user_stack,
"User Stack",
),
stack_html: format_stack_cached(
&mut stack_cache,
&stack,
"Framework Stack",
),
}
})
.collect();
let unbacked_symbols = self
let unbacked_symbols: Vec<_> = self
.unbacked_symbol_index
.borrow_mut()
.remove(&cid)
.unwrap_or(Vec::new())
.unwrap_or_default()
.drain(..)
.map(|sym| UnbackedSymbolContext {
symbol: sym.symbol.unwrap_or("".to_string()),
vr: sym.vr.unwrap_or("".to_string()),
user_stack_html: format_stack(
&sym.user_stack.unwrap_or(Vec::new()),
"User Stack",
false,
),
stack_html: format_stack(
&sym.stack.unwrap_or(Vec::new()),
"Framework Stack",
false,
),
.map(|sym| {
let user_stack = sym.user_stack.unwrap_or_default();
let stack = sym.stack.unwrap_or_default();
UnbackedSymbolContext {
symbol: sym.symbol.unwrap_or("".to_string()),
vr: sym.vr.unwrap_or("".to_string()),
user_stack_html: format_stack_cached(
&mut stack_cache,
&user_stack,
"User Stack",
),
stack_html: format_stack_cached(
&mut stack_cache,
&stack,
"Framework Stack",
),
}
})
.collect();
let remove_prefix = |x: &String| -> String {
Expand Down
4 changes: 2 additions & 2 deletions src/types.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -340,8 +340,8 @@ pub fn simplify_filename<'a>(filename: &'a str) -> &'a str {
if parts.len() > 1 {
return parts[1];
}
let re = Regex::new(r"[^/]+-seed-nspid[^/]+/").unwrap();
if let Some(captures) = re.captures(filename) {
static RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"[^/]+-seed-nspid[^/]+/").unwrap());
if let Some(captures) = RE.captures(filename) {
if let Some(capture) = captures.get(0) {
return &filename[capture.end()..];
}
Expand Down
Loading