Skip to content
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

Jump to bottom

[Mosaic GPU] Pass in TMA descriptors through kernel parameters #22175

Merged
merged 1 commit into from
Jul 2, 2024
Merged

[Mosaic GPU] Pass in TMA descriptors through kernel parameters #22175

merged 1 commit into from
Jul 2, 2024

Commits on Jul 2, 2024

  1. [Mosaic GPU] Pass in TMA descriptors through kernel parameters 

    As we've established (sigh) we can't pass in TMA descriptors through global memory.
The current workaround was to use constant memory instead, but this raises a number of
potential concurrency issues. So, instead, we use the freshly added support for grid_constant
parameters in upstream LLVM to pass the descriptors as kernel arguments. This seems to work
fine and should in fact have lower overheads than both previous methods.

PiperOrigin-RevId: 648744363
    @apaszke
    apaszke authored and jax authors committed Jul 2, 2024
    Configuration menu
    Copy the full SHA
    265a54d View commit details
    Browse the repository at this point in the history