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

[pallas::mosaic_gpu] Turn the accumulator into a reference #23793

Closed

copybara-service wants to merge 0 commits into main from test_676801466

+0 −0

copybara-service bot commented Sep 20, 2024 •

edited

Loading

[pallas::mosaic_gpu] Turn the accumulator into a reference

Changes the accumulator into a reference
Creates a discharged flavor of the wgmma op
run_scoped lowering discharges the input jaxpr
dereferencing the accumulator ref is done by a new primitive that behaves as expected when discharged
the deref primitive implies flushing the wgmma pipeline.
run_scoped does not allow references to be leaked.

copybara-service bot force-pushed the test_676801466 branch 2 times, most recently from f8d783b to e979541 Compare

September 26, 2024 09:02

copybara-service bot closed this

copybara-service bot force-pushed the test_676801466 branch from e979541 to b6d668e Compare

September 26, 2024 09:18

copybara-service bot deleted the test_676801466 branch

September 26, 2024 09:18

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment