Improve and demonstrate hackability at different levels of IR

[benvanik]

It'd be useful to have demos of hacking at different layers of the IR; for example, linalg-on-tensors + embedded SPIR-V IR, or tensors + linalg-on-buffers, etc. We should have these organized in a single place, tested to ensure they work, and cover the major use cases that are tricky to represent (dynamic shapes, etc).

OTOH some interesting examples:

• HLO + linalg-on-buffers
• HLO + VMLA (to keep us honest and not only work with linalg :)
• linalg-on-tensors + linalg-on-buffers (fallback for something not possible in linalg-on-tensors?)
• HLO + SPIR-V ("I want to hand-write this one kernel")

Most of this can be thought of as just allowing/better-handling mixed-dialect input at the top of various dialects (flow/hal/etc) and IR sugar (by way of pseudo ops, custom parsers/printers, etc) instead of anything specific to the dialects listed. This means that whatever is used to allow tensors to be passed between ops acting on tensors and buffers would allow any dialects using those two types to work together.

The particular gotcha here is that the existence of hal.interfaces is critical to how the IREE flow->hal lowering works and will always be required for working the hal dialect but is NOT required when still in flow w/ tensors. Whether explicitly provided (like if hand-coding a hal.executable in IR) or synthesized from IR sugar ops we just need to ensure it exists at the appropriate time (adjacent to the MaterializeInterfacesPass).

One way to refine this is to look through a print-ir-after-all dump and find in the IR at which point you want to be able to inject into other higher level dialects (like low-level linalg-on-buffers mixed with high-level HLO). Look at the nearby ops - which may be a mix of different dialects - and watch when they come in. The requirement to author at that point then becomes matching the expectations of the following transformations, which usually means replicating those ops (such as hal.interface being added while most ops are still in the flow dialect because flow->hal requires the interface). Whether that's done by hand or by sugar is up to the authors.

An example of where this is demonstrated (partially) today:
https://github.com/google/iree/blob/fc4abc18fda19458a473468df424b3f45f6f9d27/iree/compiler/Dialect/HAL/Conversion/FlowToHAL/test/stream_ops.mlir#L77-L117
That shows mixing hal.executables with flow ops -- note that the hal.executable.target would be your body that goes to the backend untouched, while flow->hal takes care of everything else

Note that a full flow stream isn't needed - you could just have the flow.dispatch referencing that executable:
https://github.com/google/iree/blob/85c5ef11e4dc4469591f3c9a8e7e2d2fab3fcca4/iree/compiler/Dialect/Flow/Transforms/test/form_streams.mlir#L15-L30

The major thing that needs to be provided (hal.interface) is added by the MaterializeInterfacesPass, test here:
https://github.com/google/iree/blob/a3e50b541cec6c6f4553747bac98b564c4ecf132/iree/compiler/Dialect/HAL/Transforms/test/materialize_interfaces.mlir#L48-L72
After MaterializeInterfacesPass the contents of the executable are no longer modified - so if you can provide an interface to use instead and make MaterializeInterfacesPass skip its work then you can have whatever you want inside. Makes sense: transforms after MaterializeInterfacesPass assume the interfaces are there, and if you want to not have MaterializeInterfacesPass mess with your IR you'd need to do its job for it.

This means today it should be possible to demonstrate HLO tensors + linalg-on-buffers, using the above stream_ops.mlir test as a reference (put the linalg-on-buffers code in the hal.executable.target and make sure the interface matches up). For examples of those target modules, look to each backend's tests: https://github.com/google/iree/blob/bd14ff42fddcafdf1efb09086c0ba598e436ca71/iree/compiler/Conversion/LinalgToSPIRV/test/vector_to_gpu.mlir#L3-L12
https://github.com/google/iree/blob/bf6a5dced5e16656f4bd05dee7bafba8a6a7ec49/iree/compiler/Conversion/LinalgToLLVM/test/convert_to_llvm.mlir#L4-L14

[benvanik]

(there may be some verification that prevents some of these, but they may be things we can work around - like flow.dispatch being able to accept anything with a DispatchableTrait that flow.executable and hal.executable both have so you could dispatch either and still verify, etc)

[benvanik]

(in theory that could allow flow.dispatches of gpu.kernels)

[benvanik]

Closing as obsolete.