[HAL] Set HAL allocation usage bits (EXPORT/MAPPING/etc) based on affinities.

[benvanik]

In order to export buffers the IREE_HAL_BUFFER_USAGE_SHARING_EXPORT bit must be set upon allocation. For CPU interop the IREE_HAL_BUFFER_USAGE_MAPPING_PERSISTENT bit is required as well. The compiler needs to ensure that any allocations it requests have these bits when they are used on multiple devices that either need export or mapping. TargetDevice will need to provide a query for the required support but doing so in a decentralized way is still an open design question. The initial version will always return that the mapping bit is required when any device other than itself is requesting usage of the buffer. Future refinement can have devices agree on compatibility based on supported import/export or transparent usage paths.

[benvanik]

Interesting approach in #20953! I think we can refine this a bit to avoid the need to inject the target registry (which we really want to avoid doing when possible) and avoid the combinatorial explosion of setSharedUsageBits conditions. Good ideas in the description! I'll add that we also want to ensure what we have works with targets that exist out of tree and can't do that if they need to be matched up in code. The device names aren't always stable or unique, either, and may change based on the devices themselves (whether a Vulkan extension is available, whether the device is discrete/unified, etc).

Ideally we'd have a way to have links and device target attrs declare capabilities and then solve for the set without having to reach into TargetDevice at all. It's not so much that "metal" and "local" always have certain rules given that they are particular devices but that the declared link between @metal_device and @cpu_device has certain rules. We'll have to iterate a bit to figure out what the exact set of properties is and how they relate to the bits set on the enums. Ideally they are things that we can make generic but certainly some will be trickier. If there are device-specific properties that are not related to a link we can put those on the #hal.device.target attribute - maybe things like "can this device map memory into the host address space persistently or scoped" or "can this device operate directly on host pointers without import/export" live there. When automatically constructing topologies (vs having the user specify them, which we still want to allow) we can use this information to set the link parameters.

Where it's done is another good question. StreamToHAL is conceptually the easiest point and we could use the DeviceAnalysis like you have here. Having a helper on DeviceAnalysis to get the topology link(s) between devices would make it simple to get them in the patterns and use them to derive the bits. The downside is that debugging dialect conversion sucks and we want the patterns to be simple (well, simple-ish - they aren't simple already :).

There are a few tricks we use elsewhere that we could use here to keep the patterns simple and prevent conversion from needing global device/topology information. Today the buffer usage and memory type bits are attributes on the allocation ops but they don't have to be: we can make them SSA values. HAL_EncodingTypeOp is an example of this which lets us have nicer IR than just making everything arith.constant 1235893475: %encoding_type = hal.encoding_type<dense_row_major> : i32
Here we could have hal.buffer_usage<...> and hal.memory_type<...>, make the allocation ops take them as integer SSA values, and StreamToHAL change to emitting those.

The thing this buys us is the ability to then get those SSA values from another op. e.g., %buffer_usage, %memory_type = hal.allocator.resolve_memory_properties for(#hal.device.optimal<[@device_a, @device_b]>). The StreamToHAL patterns could always blindly emit that (even if only for one device e.g. hal.allocator.resolve_memory_properties for(#hal.device.affinity<@device_a>)) and not need to involve DeviceAnalysis during pattern application (a win on debuggability).

A subsequent pass can find those ops and replace them with either statically available information or runtime queries allowing us to vary the memory types based on what's available at runtime (so a program compiled for CUDA could work on Windows which doesn't support UVA as well as Linux which does, or with certain Vulkan extensions, etc etc). The hope would be that in the common case we'd have enough topology information to do it statically and everything folds back to constant bits - but it allows us the option of emitting hal.device.query ops on various edges that get resolved at runtime. For example, if @device_a supports transparent access and @device_b may support it, we'd emit the runtime query for @device_b (which gets hoisted/memoized to initialization-time) and emit IR to do the bitfield logic.

It's worth sketching out here some examples of what we can query and use from different targets to describe their capabilities and how we could use those to derive the bits. Once we convince ourselves we know how we'd do it then we can better figure out how to expose it in the IR/code. If we could enumerate for the kind of targets we currently have (with an eye towards ones we may) that may help us narrow down the commonalities. I don't think we need to be exhaustive to start implementing, but just want to make sure we don't hardcode ourselves into a corner :)

[ziereis]

Thanks a lot for this comment — it cleared up many things :)

One specific question:

To statically resolve hal.allocator.resolve_memory_properties, we would also have to statically resolve the hal.allocator.select operation first, right? From my understanding, we can’t know the buffer usage bits without also knowing which device performs the allocation. For resolving the allocation device statically — should this strictly be derived from the topology, or do we want something similar to the runtime flag used for choosing the allocation device? Something like --iree-lead-allocation-device=device_b?

Since this would involve quite a bit of work, I assume splitting it up into multiple tasks would be helpful.

I assume the following things could be done in isolation:

• (1) Refactor buffer_usage and memory_type bits so they are available as values and are no longer attributes on the allocation ops.
• (2) Add the hal.allocator.resolve_memory_properties operation and a pass that tries to resolve it statically using the topology (this should fail for now if no topology is available).
• (3) Improve the pass to lower the operation to a runtime query if the information from the topology is insufficient.
• (4) Add properties to the hal.device.target attribute that can be used to construct a default topology.

From these, I assume that work on (1) and (2) could start right now and would already enable useful functionality. I would also start with this if you think this sounds good.

[benvanik]

Yep, sounds great! You can follow the existing enum-like ops for 1. For 2/3 a ResolveTopologyQueriesPass could do that op to start and others in the future - if it runs before MemoizeDeviceQueriesPass then if it happens to fall back to runtime checks those will get memoized by that (like you say, future work!).

[ziereis]

Just as an update, i started working on the task for the memory_properties op and the ResolveTopologyQueriesPass. A draft PR for this is on my fork ziereis#2, will open this here as soon as the refactor is merged.