google / google/heir

Support profiling in HEIR

Open
#2,074 5 comments 0 reactions 1 assignee Claimed by @TomiAkeju View on GitHub
newcomer project
Dominant language
MLIR
Stars
906
Forks
171
Avg merge
4d 12h
Merged PRs (30d)
32

Description

From the 2025-07-08 HEIR Meeting ([notes](https://docs.google.com/document/d/1v3P8LhabG75UCM3xhNcp9yTes3P7DpOqsVDFKVCSdwg/edit?usp=sharing)) we discussed the idea of having HEIR emit timing information for various ops. Ideally we could emit a json file that can be converted to a ["folded stacks"](https://queue.acm.org/detail.cfm?id=2927301#:~:text=The%20folded%20stack%2Dtrace%20format,trace%2C%20followed%20by%20a%20semicolon.) representation, which can then be viewed as a [flamegraph](https://github.com/brendangregg/FlameGraph).

The starting motivation for this, however, is to provide a blended per-op latency cost that can be used to inform scheme selection passes like https://github.com/google/heir/issues/2019

I think the right approach would be instead to use the `Location` part of MLIR to keep track of where the lowered ops came from, and then build up a folded stack from the location information.

In the meeting we talked about trying something clever like use a region-holding op (or memory-effect `secret.separator`-style) to keep track of which ops were lowered from high-level ops. E.g.,

```
timing.time {
Linalg.matvec
}
```

would lower to

```
timing.time "linalg.matvec" {
timing.time "tensor_ext.rotate" {
...
}
timing.time "arith.addi" {
...
}
...
}
```

And the stack trace could include the lowerings of these ops to associate a total op runtime with the individual components.

_However_, then I realized that lower-level optimizations completely reshuffle the IR around (e.g., operation rebalancing, lazy-relin, even CSE) which would either invalidate the above nesting, or else the nested structure/memory-effect would prevent those optimizations from running since MLIR wouldn't be allowed to move operations across the boundaries or analyze the IR through the region. So using Locations seems more proper.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.