tree 4bca58ccf0364634adf715d26dab0373be7d4cc9
parent 63cdc7d2e3d3bf6f657bbeb6f38e234da128f7fc
author Benoit Jacob <jacob.benoit.1@gmail.com> 1734228796 -0500
committer GitHub <noreply@github.com> 1734228796 -0500
gpgsig -----BEGIN PGP SIGNATURE-----
 
 wsFcBAABCAAQBQJnXjs8CRC1aQ7uu5UhlAAAffcQABQB+pou9yvPWWjBQrfJrfAj
 ZAEZpi5M+c/PTXOrH2hc4tOAY4dW95Bk3wXzh3+tf0i9T+f3mfLuoQBFnJIBja/0
 C9WEdTXaWU/ZWmr/48Y24zj2J9dQURFSEMadFpnE4Aulxem3KbKOBywYDaPSvvzm
 Rau050N2IVwNNqj1d0qUVwQjUzV05HKrOjEp+LZBMpKP+z58sgH5RRRhoB+zlI/6
 VCvKmM0C9tVB/fEYhgy5jsGlCQ3lwOcknD20w6GbUFcsTteecltJP+spCn+CLgkq
 o6uvpRacFnDul5VDgywTqWmR0YuU08H41V0jb+KyWmmL75bB0duc123TlUuQkD84
 aorDl661ooKyFI8hB6zkBe9/6Sot7+6VGjptvK/rFNX6HQ75wTSY1be7qx3ZZwuz
 iuihl9CKSMgvkTSji2kXbiLOLon9OYt5zcNcBP2bnY4fqo6h6EHknjppzQN6uS6G
 zy4OigDHxpYelN29OR3GMI+P9R7cj6ccB4cpy2jt9yXcH34cLHFJ2vFLLnyGhori
 cj0v8n8qRQqkx/ztUqW+g0zVV3zvlhxf7aV1my7HSGl23lmttej3FEqoUL00n7II
 /PByFTDsz5fb64qiXE7AOx8G664GFKmYyQ0duPZTjv9W0ej+7ShSIEBvqzMxOQMv
 L7yPz5K3lbsUvfStHDqj
 =6+n6
 -----END PGP SIGNATURE-----
 

Move GPU ukernel selection to KernelConfig (#19440)

This moves the logic deciding whether an op should be a ukernel out of
the GPULowerToUKernels pass, into KernelConfig.

So KernelConfig decides whether the op should be a ukernel, and encodes
that into the resulting `lowering_config`, in a new parameter, that is a
new attribute, UKernelSpecAttr. That attribute is directly modeled after
the equivalent C++ data structure that we have had in LowerToUKernels
passes, `FnNameAndDefAttrs`, which it replaces. If the attribute is
present, it means that the op was selected for ukernel lowering, with
the fields telling the ukernel name and some function definition
attributes (to import any dependencies, such as the `rocm` module for
runtime support symbols).

All the details about supplying the ukernel bitcode in a
`hal.executable.object` are also moved there, becoming a side effect of
`KernelConfig`.

The GPULowerToUKernels becomes much simpler, since all the
decision-making was already done for it. It just looks at the
`LoweringConfigAttr` and if it's there, it performs the requested
lowering.

The motivation for this split is that we need to know in KernelConfig
whether it's going to be a ukernel, because ops that will get lowered to
a ukernel require a different configuration. The important example for
us is `multi_mma`, which in the ukernel case needs to avoid
reduction-dimension tiling to 1 so that the ukernel gets to see the
reduction loop.

A few simplifications arise already in the current argmax ukernel logic,
confirming that this was the right design choice: the old ukernel's
matching logic was checking that the distribution tile sizes matched
what the ukernel could handle; now that is turned upside down: the
ukernel matching happens as a helper within KernelConfig where we know
we are setting the appropriate tile sizes on purpose.

Another nice improvement is that this puts just enough distance between
ukernel selection (which creates the `hal.executable.object`) and
ukernel lowering, that we are able to insert
`HoistExecutableObjectsPass` in between, simplifying the ukernel
lowering as it doesn't need to worry anymore about preserving the
`hal.executable.object`.

---------

Signed-off-by: Benoit Jacob <jacob.benoit.1@gmail.com>