[Feature Request] Single compilation unit kernels and/or improved error messages

[owenhilyard]

What is your request?

This is a 2-part request, but bundled since they both address the same UX issue.

Part one is to make the "custom op not found" error message direct users to documentation on how to include kernels. This is important since this is different for the Python and Mojo APIs, with Python doing it as part of the InferenceSession constructor in all of the examples but Mojo only being able to use the custom_ops_paths kwarg when loading a graph (which Python can as well, but it is deprecated there in favor of custom_extensions, further confusing things).

Part two is to make it so that a user can use a "single source" model for custom ops, ideally allowing for anything in the Mojo module tree to be pulled in. I know this likely means bundling MLIR inside of the binary in the data section (at least for ELF), but at present it appears that there is no way to have a self-contained binary that uses MAX. Ideally, this would be extended to Modular's kernels as well, to help with deploying MAX-based programs if the user is willing to write them in Mojo and to help with the total size of a MAX program/container.

Pinging @joshpeterson as requested.

What is your motivation for this change?

My initial assumption was that, since the Mojo compiler can see that I am using MAX, and I have the custom op I need right there and I import it, that I would be able to make use of it without needing to package a separate .mojopkg which has to live alongside my pure mojo binary. This also means iterating on a custom op is a 2-step process. First, compile the ops, then mojo run the program that is using the ops. That this didn't work was very confusing, since both SYCL and CUDA can do this. This was not helped by the error message, which did not direct me towards a solution for my problem, nor the divergence of the Mojo and Python APIs for MAX, since I would expect that the Mojo and Python APIs would be similar, with Mojo potentially needing to be more statically typed.

This also makes me concerned about the potential "minimum package size" for Mojo programs using MAX, since if they need to ship Modular's entire library of kernels due to a lack of tree shaking that may become very large over time and our current inability to tree shake .mojopkg libraries ourselves.

Any other details?

No response

[joshpeterson]

Thanks for submitting this request! We have some upcoming work planned around custom op ergonomics, so we will consider these suggestions.

[owenhilyard]

Thanks for submitting this request! We have some upcoming work planned around custom op ergonomics, so we will consider these suggestions.

I'll be happy to see it. Pure mojo is rough right now but I can see the potential.