iree: Compilation of JAX model fails in Stream Execution passes
Compile the following pass (preferably on cuda)
https://gist.githubusercontent.com/silvasean/28f8dc8bf8afbe4c150d33d494e65339/raw
Compilation causes failure when performing ScheduleExecution and ScheduleConcurrency.
About this issue
- Original URL
- State: closed
- Created a year ago
- Comments: 17 (5 by maintainers)
Commits related to this issue
- Allow cloned ops to escape stream execution regions. Fixes #13459. — committed to iree-org/iree by benvanik a year ago
- Allow cloned ops to escape stream execution regions. Fixes #13459. — committed to iree-org/iree by benvanik a year ago
- Allow cloned ops to escape stream execution regions. (#13625) Fixes #13459. — committed to iree-org/iree by benvanik a year ago
- Allow cloned ops to escape stream execution regions. (#13625) Fixes #13459. — committed to NatashaKnk/iree by benvanik a year ago
Found the root cause - special case around cloned streamable ops that have uses outside of streams (here a splat that is both used by dispatches and returned from the main function).
For the moment I think these are all coming from many people running approximately the same model from the same modeling codebase, so the IR will be very similar to the original reproducer.
I think I found both the issue and sadness, despair, and the need for a better utility for erasing sets of ops. Should be able to have a fix in a bit.
I’ll triage this now.