Call for agenda items: Oct. 18 Community Meeting

[theadactyl]

Hi folks,

What would you be most interested in discussing and/or presenting at the upcoming community meeting on Oct. 18? We are looking for the following:

• Community presentations (10-15 minutes each)
• Technical discussion topics (please be prepared to share a proposal in advance)
• Request for more information from project maintainers on XLA, StableHLO or community process.

If you have anything in mind, please respond below OR email me at thealamkin at google dot com with your ideas by Tuesday, Oct. 11.

If you would like to add this event to your calendar, please do so by subscribing to the OpenXLA Community Calendar (click on the +GoogleCalendar icon in the bottom right corner).

[ezhulenev]

I have a proof of concept for integrating Cuda Graphs + XLA runtime. Without going into too much technical details I should be able to cover the high level idea in ~15 minutes.

[nouiz]

If you "recapture and update" a captured graph, it will work.
But if you do that at every call you lose most (all?) of the speed up.

We did a cudaGraph PoC some time ago for TF/XLA.
To WAR the input ptr being dynamic issue, we captured and cached the graphs when the input ptr change.

This was working for some model when the input ptr is only a small set that always come again.
But it wasn't always working as for some models the input ptr is always a new set of ptr.
I suppose with JAX the input ptr will be more stable. But we didn't validate this.

I'll be happy to hear if you see this input ptr stability to be an issue or not with JAX.

[ezhulenev]

Yes, my hope was that pointers will not change often (didn't verify it in any real workloads yet, will not be surprised if it's not true). For this to work XLA might require more end-to-end changes and rethinking how it allocates memory.

[ezhulenev]

We don't have to capture XLA module as a whole into a graph, we can start with regions that are more likely to have stable pointers:

1. Capture only operations that have arguments inside temp allocation, presumably it should be relatively stable from run to run, or at least easier to make them stable by holding allocated temp memory
2. Capture graphs only inside loops with known trip count > N (to compensate for cache misses)

[nouiz]

That looks interesting direction.
Another potential approach: pytorch copy the model input to some constant ptr to make them stable.

If we start to hold the tmp memory ptr, we should take care to not explose the memory usage. JAX can build many small XLA module that aren't used frequently. But we can probably find a good heuristics.

Another approach would be to allocate the tmp ptr inside the graph with cudaMallocAsync. This is being optimized by the driver. The memory allocation is automatically kept unless it is needed elsewhere. This would bypass the increase of GPU memory usage issue.

[ndeep27]

@ezhulenev How can we enable the CUDA Graph Capture with XLA on a TF Saved Model?

[GleasonK]

I have started working on a StableHLO compatibility tool, to support evolving the dialect with forward and backward compatibility. I'd like to share the proposed compatibility guarantees and talk about how the tool works from a high level. Interested in community feedback / use cases. (~10min)