Multi-device Zero-Copy Path: Missing lowering for stream.cmd.flush/invalidate and zero-copy path

[Manewing]

Request description

Overview

In multi-device settings that include a heterogenous device setup including e.g. the CPU and some kind of accelerator, memory is often shared between the CPU and the accelerator. While the memory may be shared, explicit flush of written memory regions may be required in case non-coherent shared memory.

Currently there exist multiple gaps in the compiler to allow a true zero-copy lowering path for this use case.

This ticket proposes steps to close these gaps and aims to describe an implementation that achieves a true zero-copy lowering path:

The zero-copy path must preserve a synchronization operation representing the handoff between devices, even when no data copy is required. Allocation planning must make the transfer result alias the source storage. Transfer canonicalization and elision must preserve required cache maintenance and producer/consumer dependencies.

The current lowering for flow.tensor.transfer

Currently handling device assignment happens through the insertion of flow transfer operations on tensors:

%1 = flow.tensor.transfer %0 : tensor<16x16xi8> to #hal.device.affinity<@device_b>

these sub-sequently lower to async transfer in the stream dialect:

stream.async.transfer %src : ... -> to(#hal.device.affinity<@device_b>) ...

In the lowering from stream async to stream cmd this then lowers outgoing transfers to (

iree/compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleAllocation.cpp

Line 677 in 9acea6a

applyAsyncTransferOp(IREE::Stream::AffinityAttr executionAffinityAttr,

)

stream.cmd.copy %src[...], %dst[...], %length : ...
stream.cmd.flush to(#hal.device.affinity<@device_b>) %dst[...] : ...

incoming transfers on the other hand lower to: stream.cmd.invalidate and stream.cmd.copy

Missing lowering for stream.cmd.flush and stream.cmd.invalidate

The first gap in the lowering is that lowering of stream.cmd.flush is simply not implemented. The operations is simply erased in the current lowering. See

iree/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Patterns.cpp

Line 622 in 9acea6a

struct CmdFlushOpPattern : StreamConversionPattern<IREE::Stream::CmdFlushOp> {

Not only is the lowering of the operation missing also its implementation in the HAL command buffer is not present. The lowering to be implemented should follow what is implemented for the IREE::Stream::CmdCopyOp

iree/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Patterns.cpp

Line 674 in 9acea6a

struct CmdCopyOpPattern : StreamConversionPattern<IREE::Stream::CmdCopyOp> {

. This lowers to IREE::HAL::CommandBufferCopyBufferOp which is sub-sequently lowered to the VM call here

iree/compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/ConvertCommandBufferOps.cpp

Line 206 in 9acea6a

class CommandBufferCopyBufferOpConversion

It also requires the extension of the HAL with the correct command buffer handler. See handling for hal.command_buffer.copy_buffer it should be analogous.

Similar for stream.cmd.invalidate

Note that the HAL implementation already provides functions for buffer flush and invalidation:

• iree_hal_buffer_mapping_flush_range
• iree_hal_buffer_mapping_invalidate_range

Copy operation when lowering

The current lowering can produce a copy of the buffer which breaks a true zero-copy path:

stream.cmd.copy %src[...], %dst[...], %length : ...
stream.cmd.flush to(#hal.device.affinity<@device_b>) %dst[...] : ...

Transfers can be eliminated when resource lifetimes match and the affinity/topology checks permit it. This replaces the transfer result with its source.

iree/compiler/src/iree/compiler/Dialect/Stream/Transforms/ElideAsyncCopies.cpp

Line 994 in 9acea6a

IREE::Stream::AsyncTransferOp transferOp, ElisionAnalysis &analysis,

For non-coherent memory this elision is not valid.

For true zero-copy we need to avoid the introduction of stream.cmd.copy.

To achieve this we should allow to communicate that no explicit buffer allocation and copy is required for flow.tensor.transfer and stream.async.transfer. Either derived from the device affinity or from an explicit attribute set on the operation. Zero-copy requires that the backing allocation is accessible to both devices and satisfies their usage requirements. Device affinity or an operation attribute must convey or establish this compatibility, shared physical memory alone is insufficient.

This should control the lowering of stream.async.transfer to avoid the insertion of the stream.cmd.copy

Note that simply leaving out the copy is not possible, there also has to be change in allocation so that the consumers read from the correct buffer. This needs to be handled in

iree/compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleAllocation.cpp

Line 1642 in 9acea6a

allocateExecutionRegion(IREE::Stream::AsyncExecuteOp executeOp,

Flush and invalidate operations must have execution-time ordering semantics. A flush must execute after the writes it makes available, and completion must be included in the handoff’s signaled timepoint/fence. An invalidate must execute after the corresponding producer completion and before dependent reads. Implementations must not perform cache maintenance merely during command-buffer recording. Backends may satisfy these requirements through queued commands, ordered host work, or equivalent synchronization guarantees.

Invalidate/flush is ignored for single copy/fill

The current matching for transfer operations is ignoring flush/invalidate

iree/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Patterns.cpp

Line 1232 in 9acea6a

// }

This optimization must preserve required flush/invalidate operations or decline the optimization. A region containing only cache maintenance after copy removal must retain its execution and synchronization semantics.

Cases to be handled

• Coherent shared memory: shared backing allocation, no copy, preserved ordering.
• Non-coherent shared memory: shared backing allocation, no copy, correctly ordered flush/invalidate as required in each direction.
• Incompatible or discrete memory: retain allocation and copy.
• Multiple consumers and in-place updates: preserve tensor semantics and allocation lifetime.

What component(s) does this issue relate to?

No response

Additional context

No response

[purnima-nlp]

hey i have working on this already , may you assign it to me

[Manewing]

Hey @purnima-nlp

thanks for looking into this and also for reaching out on this! I opened this issue after I saw a discussion on Discord regarding this topic, now there already has been some work done by @l0r3m1psum as well [0]. So we will have to take a look how to align the work and avoid duplication.

Generally though very nice to see movement on this topic :) let's see how we can push this together!

[0] main...l0r3m1psum:iree:hal-command-buffer-flush

[purnima-nlp]

Hey @purnima-nlp

thanks for looking into this and also for reaching out on this! I opened this issue after I saw a discussion on Discord regarding this topic, now there already has been some work done by @l0r3m1psum as well [0]. So we will have to take a look how to align the work and avoid duplication.

Generally though very nice to see movement on this topic :) let's see how we can push this together!

[0] main...l0r3m1psum:iree:hal-command-buffer-flush

Yes I would love to contribute on this issue, let me see what's already implemented.
I will comment soon on this

[devtbi]

for context, two items should be mentioned here, 1) where the current topology was introduced: #20885 / #20856 2) runtime toplogy that was added here: #23573. We should look into having the device capabilities & toplogy /parts of it mirrored to compile-time, if that makes sense. Looks a bit inconsistent right now.

on flush: I think stream flush docs "Transfers a resource to an external target. The resource memory is made available to the target and can be made visible there using stream.cmd.invalidate." is a bit misleading. All flush_range implementations currently do is a 'host-CPU (mapped-memory) flush, such that host writes become available to the device owning the buffer', e.g. (vkFlushMappedMemoryRanges) - might be the reason why it's not lowered currently? just guessing.

Generally I agree with adding some attribute to transfer to establish the kind of transfer, instead of attempting to clean up after-the-fact. This can optionally be derived from the topology. Maybe starting with copy | alias (which will be "NATIVE" or "IMPORT" in terms of the runtime topology terminology, I don't think we need to differentiate import at compile-time). Not sure if coherence should be an attribute of transfer, and then resolved by resolve_memory_properties, or inferred from topology and the buffers.

[purnima-nlp]

for context, two items should be mentioned here, 1) where the current topology was introduced: #20885 / #20856 2) runtime toplogy that was added here: #23573. We should look into having the device capabilities & toplogy /parts of it mirrored to compile-time, if that makes sense. Looks a bit inconsistent right now.

on flush: I think stream flush docs "Transfers a resource to an external target. The resource memory is made available to the target and can be made visible there using stream.cmd.invalidate." is a bit misleading. All flush_range implementations currently do is a 'host-CPU (mapped-memory) flush, such that host writes become available to the device owning the buffer', e.g. (vkFlushMappedMemoryRanges) - might be the reason why it's not lowered currently? just guessing.

Generally I agree with adding some attribute to transfer to establish the kind of transfer, instead of attempting to clean up after-the-fact. This can optionally be derived from the topology. Maybe starting with copy | alias (which will be "NATIVE" or "IMPORT" in terms of the runtime topology terminology, I don't think we need to differentiate import at compile-time). Not sure if coherence should be an attribute of transfer, and then resolved by resolve_memory_properties, or inferred from topology and the buffers.

Thanks for your comment this is helpful for me. I have few things to ask you below !! please provide your valuable opinion

For my PR, I was intentionally keeping the scope limited to introducing HAL/runtime command-buffer plumbing for flush/invalidate and leaving Stream-to-HAL lowering disabled. The current inline implementation uses iree_hal_buffer_mapping_flush_range / invalidate_range, so as you pointed out that is specifically mapped host-memory cache maintenance rather than a generic device-to-device ownership transfer.

My current thinking was to keep these HAL operations as the runtime representation for that cache-maintenance step, and handle the actual zero-copy transfer decision separately. In particular, I agree that distinguishing something like copy vs alias earlier on the transfer seems cleaner than trying to infer it after allocation/scheduling.

Would you prefer that hal.command_buffer.flush_buffer / invalidate_buffer be explicitly scoped/documented as mapped-memory cache maintenance only, with stream.cmd.flush/invalidate remaining a higher-level concept for now? Or do you think these HAL operations should not exist independently and the design should instead start from transfer kind/topology first?

Also, for the transfer kind, would you expect copy | alias to be carried on flow.tensor.transfer / stream.async.transfer, with topology used to derive it where possible, while coherence is resolved separately from buffer/memory properties?