Extending roofline

[dangdana]

Is there interest from the JAX community or maintainers in extending coverage and examples for jax.experimental.roofline? @zacmustin

I’d like to contribute while preserving its minimal API. Before implementing anything, I’d appreciate feedback on which of the following would fit the library’s intended scope.

1. Documentation: Add a few worked examples connecting roofline accounting to How to Scale Your Model, especially Chapter 3 on sharding. @jacobaustin123
2. Example: Demonstrate how to convert the existing RooflineResult into compute, HBM, and communication time bounds using hardware specifications, without changing the API (perhaps part of 1 above)
3. Bug: scan counts its body once, regardless of iteration count.
   Fix: Multiply work and traffic by iteration count, treating peak memory separately.
4. Bug: Conditional branch work is omitted silently.
   Fix: Warn when conditional work is omitted, following the existing approach for unsupported primitives.
5. Bug: reduce_max returns zero work and traffic.
   Fix: Add a rule following the existing reduce_sum convention.

Bug were observed in JAX 0.11.2. I can provide small reproducers.

Would these contributions be welcome, and are there other priorities or existing efforts I should coordinate with?

[zacmustin]

Thanks for the question! Although I did flesh roofline out a bit (mostly by adding ops), I believe @epiqueras is the true code owner, so I defer to him for decisions on design / scope of the tool. AFAIK there's no ongoing work to this tool; my guess is that contributions are welcome.

[dangdana]

To start small, I tried addressing number 5 above by adding more reductions to the cost model in #41825 .

Also tagging @matthiaskramm , who git claims originally added the sum reduction.