[HAL] Add command buffer flush and invalidate operations (#24945) This PR adds HAL and runtime plumbing for ordered command-buffer buffer flush and invalidate operations. It introduces: * `hal.command_buffer.flush_buffer` * `hal.command_buffer.invalidate_buffer` * HAL-to-VM lowering and VM imports * runtime command-buffer API and vtable hooks * VM module/shim plumbing * inline command-buffer support * deferred command-buffer recording/replay support * HAL IR and HAL-to-VM test coverage The operations carry a buffer reference, offset, and length and are intended to represent execution-ordered cache maintenance for shared non-coherent memory. Stream-to-HAL lowering is intentionally left unchanged in this PR. Some backends, especially Vulkan, require queue/submission-level handling to preserve execution-time ordering because a HAL command buffer is currently replayed into a single native `VkCommandBuffer`. Enabling Stream lowering before that backend support exists could cause unsupported runtime paths. This PR therefore focuses on establishing the HAL/runtime infrastructure first. Backend-specific execution support and Stream lowering can follow separately. Related to #24930. AI Tool Disclosure Assisted by codex for code base exploration Signed-off-by: purnima <purnima.shrivastava@sandlogic.com>
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/ConvertCommandBufferOps.cpp b/compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/ConvertCommandBufferOps.cpp index bc85f19..2d24473 100644 --- a/compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/ConvertCommandBufferOps.cpp +++ b/compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/ConvertCommandBufferOps.cpp
@@ -100,6 +100,81 @@ mutable IREE::VM::ImportOp importOp; }; +class CommandBufferFlushBufferOpConversion + : public OpConversionPattern<IREE::HAL::CommandBufferFlushBufferOp> { +public: + CommandBufferFlushBufferOpConversion(MLIRContext *context, + SymbolTable &importSymbols, + TypeConverter &typeConverter, + StringRef importName) + : OpConversionPattern(typeConverter, context) { + importOp = importSymbols.lookup<IREE::VM::ImportOp>(importName); + assert(importOp); + } + + LogicalResult + matchAndRewrite(IREE::HAL::CommandBufferFlushBufferOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto importType = importOp.getFunctionType(); + auto [targetBufferSlot, targetBuffer] = + splitBufferSlot(op.getLoc(), adaptor.getTargetBuffer(), rewriter); + SmallVector<Value, 8> callOperands = { + adaptor.getCommandBuffer(), + targetBufferSlot, + targetBuffer, + castToImportType(adaptor.getTargetOffset(), rewriter.getI64Type(), + rewriter), + castToImportType(adaptor.getLength(), rewriter.getI64Type(), rewriter), + }; + auto callOp = rewriter.replaceOpWithNewOp<IREE::VM::CallOp>( + op, SymbolRefAttr::get(importOp), importType.getResults(), + callOperands); + copyImportAttrs(importOp, callOp); + return success(); + } + +private: + mutable IREE::VM::ImportOp importOp; +}; + +class CommandBufferInvalidateBufferOpConversion + : public OpConversionPattern<IREE::HAL::CommandBufferInvalidateBufferOp> { +public: + CommandBufferInvalidateBufferOpConversion(MLIRContext *context, + SymbolTable &importSymbols, + TypeConverter &typeConverter, + StringRef importName) + : OpConversionPattern(typeConverter, context) { + importOp = importSymbols.lookup<IREE::VM::ImportOp>(importName); + assert(importOp); + } + + LogicalResult + matchAndRewrite(IREE::HAL::CommandBufferInvalidateBufferOp op, + OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto importType = importOp.getFunctionType(); + auto [targetBufferSlot, targetBuffer] = + splitBufferSlot(op.getLoc(), adaptor.getTargetBuffer(), rewriter); + SmallVector<Value, 8> callOperands = { + adaptor.getCommandBuffer(), + targetBufferSlot, + targetBuffer, + castToImportType(adaptor.getTargetOffset(), rewriter.getI64Type(), + rewriter), + castToImportType(adaptor.getLength(), rewriter.getI64Type(), rewriter), + }; + auto callOp = rewriter.replaceOpWithNewOp<IREE::VM::CallOp>( + op, SymbolRefAttr::get(importOp), importType.getResults(), + callOperands); + copyImportAttrs(importOp, callOp); + return success(); + } + +private: + mutable IREE::VM::ImportOp importOp; +}; + class CommandBufferFillBufferOpConversion : public OpConversionPattern<IREE::HAL::CommandBufferFillBufferOp> { public: @@ -492,6 +567,11 @@ .insert<VMImportOpConversion<IREE::HAL::CommandBufferExecutionBarrierOp>>( context, importSymbols, typeConverter, "hal.command_buffer.execution_barrier"); + patterns.insert<CommandBufferFlushBufferOpConversion>( + context, importSymbols, typeConverter, "hal.command_buffer.flush_buffer"); + patterns.insert<CommandBufferInvalidateBufferOpConversion>( + context, importSymbols, typeConverter, + "hal.command_buffer.invalidate_buffer"); patterns.insert<CommandBufferFillBufferOpConversion>( context, importSymbols, typeConverter, "hal.command_buffer.fill_buffer"); patterns.insert<CommandBufferUpdateBufferOpConversion>(
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/test/command_buffer_ops.mlir b/compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/test/command_buffer_ops.mlir index 1bd04e7..25cff86 100644 --- a/compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/test/command_buffer_ops.mlir +++ b/compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/test/command_buffer_ops.mlir
@@ -49,6 +49,74 @@ // ----- +// CHECK-LABEL: @command_buffer_flush_buffer +// CHECK-SAME: (%[[CMD:.+]]: !vm.ref<!hal.command_buffer>, %[[BUFFER:.+]]: !vm.ref<!hal.buffer>) +util.func public @command_buffer_flush_buffer( + %cmd: !hal.command_buffer, + %buffer: !hal.buffer +) { + %c100 = arith.constant 100 : index + %c200 = arith.constant 200 : index + // CHECK-DAG: %[[UNUSED_SLOT:.+]] = vm.const.i32.zero + // CHECK: vm.call @hal.command_buffer.flush_buffer(%[[CMD]], %[[UNUSED_SLOT]], %[[BUFFER]], %c100, %c200) + hal.command_buffer.flush_buffer<%cmd : !hal.command_buffer> + target(%buffer : !hal.buffer)[%c100, %c200] + util.return +} + +// ----- + +// CHECK-LABEL: @command_buffer_flush_buffer_indirect +// CHECK-SAME: (%[[CMD:.+]]: !vm.ref<!hal.command_buffer>, %[[BUFFER_SLOT:.+]]: i32) +util.func public @command_buffer_flush_buffer_indirect( + %cmd: !hal.command_buffer, + %buffer_slot: index +) { + %c100 = arith.constant 100 : index + %c200 = arith.constant 200 : index + // CHECK-DAG: %[[NULL_BUFFER:.+]] = vm.const.ref.zero : !vm.ref<!hal.buffer> + // CHECK: vm.call @hal.command_buffer.flush_buffer(%[[CMD]], %[[BUFFER_SLOT]], %[[NULL_BUFFER]], %c100, %c200) + hal.command_buffer.flush_buffer<%cmd : !hal.command_buffer> + target(%buffer_slot : index)[%c100, %c200] + util.return +} + +// ----- + +// CHECK-LABEL: @command_buffer_invalidate_buffer +// CHECK-SAME: (%[[CMD:.+]]: !vm.ref<!hal.command_buffer>, %[[BUFFER:.+]]: !vm.ref<!hal.buffer>) +util.func public @command_buffer_invalidate_buffer( + %cmd: !hal.command_buffer, + %buffer: !hal.buffer +) { + %c100 = arith.constant 100 : index + %c200 = arith.constant 200 : index + // CHECK-DAG: %[[UNUSED_SLOT:.+]] = vm.const.i32.zero + // CHECK: vm.call @hal.command_buffer.invalidate_buffer(%[[CMD]], %[[UNUSED_SLOT]], %[[BUFFER]], %c100, %c200) + hal.command_buffer.invalidate_buffer<%cmd : !hal.command_buffer> + target(%buffer : !hal.buffer)[%c100, %c200] + util.return +} + +// ----- + +// CHECK-LABEL: @command_buffer_invalidate_buffer_indirect +// CHECK-SAME: (%[[CMD:.+]]: !vm.ref<!hal.command_buffer>, %[[BUFFER_SLOT:.+]]: i32) +util.func public @command_buffer_invalidate_buffer_indirect( + %cmd: !hal.command_buffer, + %buffer_slot: index +) { + %c100 = arith.constant 100 : index + %c200 = arith.constant 200 : index + // CHECK-DAG: %[[NULL_BUFFER:.+]] = vm.const.ref.zero : !vm.ref<!hal.buffer> + // CHECK: vm.call @hal.command_buffer.invalidate_buffer(%[[CMD]], %[[BUFFER_SLOT]], %[[NULL_BUFFER]], %c100, %c200) + hal.command_buffer.invalidate_buffer<%cmd : !hal.command_buffer> + target(%buffer_slot : index)[%c100, %c200] + util.return +} + +// ----- + // CHECK-LABEL: @command_buffer_fill_buffer_i8 util.func public @command_buffer_fill_buffer_i8( %arg0: !hal.command_buffer,
diff --git a/compiler/src/iree/compiler/Dialect/HAL/IR/HALOps.td b/compiler/src/iree/compiler/Dialect/HAL/IR/HALOps.td index 0a99931..0f50c56 100644 --- a/compiler/src/iree/compiler/Dialect/HAL/IR/HALOps.td +++ b/compiler/src/iree/compiler/Dialect/HAL/IR/HALOps.td
@@ -1645,6 +1645,51 @@ // TODO(benvanik): event ops. +def HAL_CommandBufferFlushBufferOp : HAL_Op<"command_buffer.flush_buffer"> { + let summary = [{Makes prior writes available to external consumers.}]; + let description = [{ + Makes prior writes to the target buffer range available to external + consumers. + }]; + + let arguments = (ins + HAL_CommandBuffer:$command_buffer, + AnyTypeOf<[Index, HAL_BufferType]>:$target_buffer, + HAL_DeviceSize:$target_offset, + HAL_DeviceSize:$length + ); + + let assemblyFormat = [{ + `<` $command_buffer `:` type($command_buffer) `>` + `target` `(` $target_buffer `:` type($target_buffer) `)` + `` `[` $target_offset `,` $length `]` + attr-dict-with-keyword + }]; +} + +def HAL_CommandBufferInvalidateBufferOp + : HAL_Op<"command_buffer.invalidate_buffer"> { + let summary = [{Makes external writes visible to subsequent consumers.}]; + let description = [{ + Makes external writes to the target buffer range visible to subsequent + consumers. + }]; + + let arguments = (ins + HAL_CommandBuffer:$command_buffer, + AnyTypeOf<[Index, HAL_BufferType]>:$target_buffer, + HAL_DeviceSize:$target_offset, + HAL_DeviceSize:$length + ); + + let assemblyFormat = [{ + `<` $command_buffer `:` type($command_buffer) `>` + `target` `(` $target_buffer `:` type($target_buffer) `)` + `` `[` $target_offset `,` $length `]` + attr-dict-with-keyword + }]; +} + def HAL_CommandBufferFillBufferOp : HAL_Op<"command_buffer.fill_buffer"> { let summary = [{Command buffer buffer fill recording operation.}]; let description = [{
diff --git a/compiler/src/iree/compiler/Dialect/HAL/IR/test/command_buffer_ops.mlir b/compiler/src/iree/compiler/Dialect/HAL/IR/test/command_buffer_ops.mlir index 357fee3..83e49bb 100644 --- a/compiler/src/iree/compiler/Dialect/HAL/IR/test/command_buffer_ops.mlir +++ b/compiler/src/iree/compiler/Dialect/HAL/IR/test/command_buffer_ops.mlir
@@ -53,6 +53,42 @@ // ----- +// CHECK-LABEL: @command_buffer_flush_buffer +// CHECK-SAME: (%[[CMD:.+]]: !hal.command_buffer, +// CHECK-SAME: %[[BUFFER:.+]]: !hal.buffer, +// CHECK-SAME: %[[OFFSET:.+]]: index, %[[LENGTH:.+]]: index) +util.func public @command_buffer_flush_buffer( + %cmd: !hal.command_buffer, + %buffer: !hal.buffer, + %offset: index, + %length: index) { + // CHECK: hal.command_buffer.flush_buffer<%[[CMD]] : !hal.command_buffer> + // CHECK-SAME: target(%[[BUFFER]] : !hal.buffer)[%[[OFFSET]], %[[LENGTH]]] + hal.command_buffer.flush_buffer<%cmd : !hal.command_buffer> + target(%buffer : !hal.buffer)[%offset, %length] + util.return +} + +// ----- + +// CHECK-LABEL: @command_buffer_invalidate_buffer +// CHECK-SAME: (%[[CMD:.+]]: !hal.command_buffer, +// CHECK-SAME: %[[BUFFER:.+]]: !hal.buffer, +// CHECK-SAME: %[[OFFSET:.+]]: index, %[[LENGTH:.+]]: index) +util.func public @command_buffer_invalidate_buffer( + %cmd: !hal.command_buffer, + %buffer: !hal.buffer, + %offset: index, + %length: index) { + // CHECK: hal.command_buffer.invalidate_buffer<%[[CMD]] : !hal.command_buffer> + // CHECK-SAME: target(%[[BUFFER]] : !hal.buffer)[%[[OFFSET]], %[[LENGTH]]] + hal.command_buffer.invalidate_buffer<%cmd : !hal.command_buffer> + target(%buffer : !hal.buffer)[%offset, %length] + util.return +} + +// ----- + // CHECK-LABEL: @command_buffer_fill_buffer // CHECK-SAME: (%[[CMD:.+]]: !hal.command_buffer, // CHECK-SAME: %[[BUFFER:.+]]: !hal.buffer,
diff --git a/compiler/src/iree/compiler/Dialect/HAL/hal.imports.mlir b/compiler/src/iree/compiler/Dialect/HAL/hal.imports.mlir index 6bad000..2279e49 100644 --- a/compiler/src/iree/compiler/Dialect/HAL/hal.imports.mlir +++ b/compiler/src/iree/compiler/Dialect/HAL/hal.imports.mlir
@@ -269,6 +269,24 @@ %buffer_slot : i32 ) +// Makes prior writes available to external consumers. +vm.import private @command_buffer.flush_buffer( + %command_buffer : !vm.ref<!hal.command_buffer>, + %target_buffer_slot : i32, + %target_buffer : !vm.ref<!hal.buffer>, + %target_offset : i64, + %length : i64 +) + +// Makes external writes visible to subsequent consumers. +vm.import private @command_buffer.invalidate_buffer( + %command_buffer : !vm.ref<!hal.command_buffer>, + %target_buffer_slot : i32, + %target_buffer : !vm.ref<!hal.buffer>, + %target_offset : i64, + %length : i64 +) + // Fills the target buffer with the given repeating value. // NOTE: order slightly differs from op in order to get better arg alignment. vm.import private @command_buffer.fill_buffer(
diff --git a/runtime/src/iree/hal/BUILD.bazel b/runtime/src/iree/hal/BUILD.bazel index 021b83a..feb5b18 100644 --- a/runtime/src/iree/hal/BUILD.bazel +++ b/runtime/src/iree/hal/BUILD.bazel
@@ -162,6 +162,17 @@ ) iree_runtime_cc_test( + name = "command_buffer_test", + srcs = ["command_buffer_test.cc"], + deps = [ + ":hal", + "//runtime/src/iree/base", + "//runtime/src/iree/testing:gtest", + "//runtime/src/iree/testing:gtest_main", + ], +) + +iree_runtime_cc_test( name = "pool_test", srcs = ["pool_test.cc"], deps = [
diff --git a/runtime/src/iree/hal/CMakeLists.txt b/runtime/src/iree/hal/CMakeLists.txt index a6b2e95..601ad64 100644 --- a/runtime/src/iree/hal/CMakeLists.txt +++ b/runtime/src/iree/hal/CMakeLists.txt
@@ -142,6 +142,18 @@ iree_cc_test( NAME + command_buffer_test + SRCS + "command_buffer_test.cc" + DEPS + ::hal + iree::base + iree::testing::gtest + iree::testing::gtest_main +) + +iree_cc_test( + NAME pool_test SRCS "pool_test.cc"
diff --git a/runtime/src/iree/hal/command_buffer.c b/runtime/src/iree/hal/command_buffer.c index 986d0d5..311709d 100644 --- a/runtime/src/iree/hal/command_buffer.c +++ b/runtime/src/iree/hal/command_buffer.c
@@ -498,6 +498,65 @@ return status; } +IREE_API_EXPORT iree_status_t +iree_hal_command_buffer_flush_buffer(iree_hal_command_buffer_t* command_buffer, + iree_hal_buffer_ref_t target_ref) { + IREE_ASSERT_ARGUMENT(command_buffer); + if (target_ref.length == 0) { + // No-op flush. All other validation is skipped. + return iree_ok_status(); + } + IREE_TRACE_ZONE_BEGIN(z0); + IF_VALIDATING(command_buffer, { + IREE_RETURN_AND_END_ZONE_IF_ERROR( + z0, iree_hal_command_buffer_flush_buffer_validation( + command_buffer, VALIDATION_STATE(command_buffer), target_ref)); + }); + iree_status_t status; + if (_VTABLE_DISPATCH(command_buffer, flush_buffer)) { + status = _VTABLE_DISPATCH(command_buffer, flush_buffer)(command_buffer, + target_ref); + } else { + // FIXME: Not all command-buffer backends implement flush_buffer yet. + // Stream-to-HAL lowering remains disabled until backend execution semantics + // are supported. + status = iree_make_status(IREE_STATUS_UNIMPLEMENTED, + "command buffer does not support flush_buffer"); + } + IREE_TRACE_ZONE_END(z0); + return status; +} + +IREE_API_EXPORT iree_status_t iree_hal_command_buffer_invalidate_buffer( + iree_hal_command_buffer_t* command_buffer, + iree_hal_buffer_ref_t target_ref) { + IREE_ASSERT_ARGUMENT(command_buffer); + if (target_ref.length == 0) { + // No-op invalidate. All other validation is skipped. + return iree_ok_status(); + } + IREE_TRACE_ZONE_BEGIN(z0); + IF_VALIDATING(command_buffer, { + IREE_RETURN_AND_END_ZONE_IF_ERROR( + z0, iree_hal_command_buffer_invalidate_buffer_validation( + command_buffer, VALIDATION_STATE(command_buffer), target_ref)); + }); + iree_status_t status; + if (_VTABLE_DISPATCH(command_buffer, invalidate_buffer)) { + status = _VTABLE_DISPATCH(command_buffer, invalidate_buffer)(command_buffer, + target_ref); + } else { + // FIXME: Not all command-buffer backends implement invalidate_buffer yet. + // Stream-to-HAL lowering remains disabled until backend execution semantics + // are supported. + status = + iree_make_status(IREE_STATUS_UNIMPLEMENTED, + "command buffer does not support invalidate_buffer"); + } + IREE_TRACE_ZONE_END(z0); + return status; +} + IREE_API_EXPORT iree_status_t iree_hal_command_buffer_fill_buffer( iree_hal_command_buffer_t* command_buffer, iree_hal_buffer_ref_t target_ref, const void* pattern, iree_host_size_t pattern_length,
diff --git a/runtime/src/iree/hal/command_buffer.h b/runtime/src/iree/hal/command_buffer.h index a73b7f2..7a78a2b 100644 --- a/runtime/src/iree/hal/command_buffer.h +++ b/runtime/src/iree/hal/command_buffer.h
@@ -937,6 +937,24 @@ iree_hal_command_buffer_t* command_buffer, iree_hal_buffer_ref_t buffer_ref, iree_hal_memory_advise_flags_t flags, uint64_t arg0, uint64_t arg1); +// Performs cache maintenance on the target buffer range to make prior writes +// available to external consumers. This operation does not establish execution +// dependencies. Callers must use execution barriers to order preceding writes +// before the flush when ordering is required. Recording order alone may not +// provide this ordering on implementations that execute commands concurrently. +IREE_API_EXPORT iree_status_t +iree_hal_command_buffer_flush_buffer(iree_hal_command_buffer_t* command_buffer, + iree_hal_buffer_ref_t target_ref); + +// Performs cache maintenance on the target buffer range to make external writes +// visible to subsequent consumers. This operation does not establish execution +// dependencies. Callers must use execution barriers to order subsequent reads +// after the invalidate when ordering is required. Recording order alone may not +// provide this ordering on implementations that execute commands concurrently. +IREE_API_EXPORT iree_status_t iree_hal_command_buffer_invalidate_buffer( + iree_hal_command_buffer_t* command_buffer, + iree_hal_buffer_ref_t target_ref); + // Fills the target buffer with the given repeating value. // Expects that |pattern_length| is one of 1, 2, or 4 and that the offset and // length are aligned to the natural alignment of the value. @@ -1144,6 +1162,14 @@ iree_hal_buffer_ref_t buffer_ref, iree_hal_memory_advise_flags_t flags, uint64_t arg0, uint64_t arg1); + iree_status_t(IREE_API_PTR* flush_buffer)( + iree_hal_command_buffer_t* command_buffer, + iree_hal_buffer_ref_t target_ref); + + iree_status_t(IREE_API_PTR* invalidate_buffer)( + iree_hal_command_buffer_t* command_buffer, + iree_hal_buffer_ref_t target_ref); + iree_status_t(IREE_API_PTR* fill_buffer)( iree_hal_command_buffer_t* command_buffer, iree_hal_buffer_ref_t target_ref, const void* pattern,
diff --git a/runtime/src/iree/hal/command_buffer_test.cc b/runtime/src/iree/hal/command_buffer_test.cc new file mode 100644 index 0000000..a2bc5d5 --- /dev/null +++ b/runtime/src/iree/hal/command_buffer_test.cc
@@ -0,0 +1,321 @@ +// Copyright 2026 The IREE Authors +// +// Licensed under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#include "iree/base/api.h" +#include "iree/hal/api.h" +#include "iree/testing/gtest.h" +#include "iree/testing/status_matchers.h" + +namespace iree { +namespace { + +struct FakeCommandBuffer { + iree_hal_command_buffer_t base; + int flush_count = 0; + int invalidate_count = 0; + iree_hal_buffer_ref_t flush_ref = {}; + iree_hal_buffer_ref_t invalidate_ref = {}; + iree_status_code_t flush_status = IREE_STATUS_OK; + iree_status_code_t invalidate_status = IREE_STATUS_OK; +}; + +struct FakeBuffer { + iree_hal_buffer_t base; + uint8_t storage[64] = {}; +}; + +static void IREE_API_PTR DestroyBuffer(iree_hal_buffer_t* /*buffer*/) {} + +static const iree_hal_buffer_vtable_t kFakeBufferVTable = { + /*.recycle=*/iree_hal_buffer_recycle, + /*.destroy=*/DestroyBuffer, +}; + +static void IREE_API_PTR Destroy(iree_hal_command_buffer_t* command_buffer) { + // The fixture owns the fake; only the validation storage is heap allocated. + iree_allocator_free(iree_allocator_system(), + command_buffer->validation_state); +} + +static iree_status_t IREE_API_PTR +BeginOrEnd(iree_hal_command_buffer_t* /*command_buffer*/) { + return iree_ok_status(); +} + +static iree_status_t IREE_API_PTR +Flush(iree_hal_command_buffer_t* command_buffer, + iree_hal_buffer_ref_t target_ref) { + auto* fake = reinterpret_cast<FakeCommandBuffer*>(command_buffer); + ++fake->flush_count; + fake->flush_ref = target_ref; + return iree_status_from_code(fake->flush_status); +} + +static iree_status_t IREE_API_PTR +Invalidate(iree_hal_command_buffer_t* command_buffer, + iree_hal_buffer_ref_t target_ref) { + auto* fake = reinterpret_cast<FakeCommandBuffer*>(command_buffer); + ++fake->invalidate_count; + fake->invalidate_ref = target_ref; + return iree_status_from_code(fake->invalidate_status); +} + +class CommandBufferTest : public ::testing::Test { + protected: + void SetUp() override { + vtable_.destroy = Destroy; + vtable_.begin = BeginOrEnd; + vtable_.end = BeginOrEnd; + vtable_.flush_buffer = Flush; + vtable_.invalidate_buffer = Invalidate; + } + + void TearDown() override { + if (initialized_) iree_hal_command_buffer_release(&fake_.base); + } + + iree_status_t Initialize(iree_hal_command_buffer_mode_t mode = + IREE_HAL_COMMAND_BUFFER_MODE_UNVALIDATED, + iree_hal_command_category_t categories = + IREE_HAL_COMMAND_CATEGORY_TRANSFER, + iree_host_size_t binding_capacity = 0) { + void* validation_state = nullptr; + iree_host_size_t validation_size = + iree_hal_command_buffer_validation_state_size(mode, binding_capacity); + if (validation_size) { + IREE_RETURN_IF_ERROR(iree_allocator_malloc( + iree_allocator_system(), validation_size, &validation_state)); + } + iree_hal_command_buffer_initialize( + /*device_allocator=*/nullptr, mode, categories, + IREE_HAL_QUEUE_AFFINITY_ANY, binding_capacity, validation_state, + &vtable_, &fake_.base); + initialized_ = true; + return iree_ok_status(); + } + + static void ExpectBufferRefEq(iree_hal_buffer_ref_t expected, + iree_hal_buffer_ref_t actual) { + EXPECT_EQ(static_cast<uint32_t>(expected.reserved), + static_cast<uint32_t>(actual.reserved)); + EXPECT_EQ(static_cast<uint32_t>(expected.buffer_slot), + static_cast<uint32_t>(actual.buffer_slot)); + EXPECT_EQ(expected.buffer, actual.buffer); + EXPECT_EQ(expected.offset, actual.offset); + EXPECT_EQ(expected.length, actual.length); + } + + FakeCommandBuffer fake_ = {}; + iree_hal_command_buffer_vtable_t vtable_ = {}; + bool initialized_ = false; +}; + +TEST_F(CommandBufferTest, FlushDispatchesExactBufferRef) { + IREE_ASSERT_OK(Initialize()); + // Unvalidated dispatch only: the buffer is an identity token, never accessed. + iree_hal_buffer_t buffer = {}; + auto ref = iree_hal_make_buffer_ref(&buffer, 37, 59); + IREE_ASSERT_OK(iree_hal_command_buffer_flush_buffer(&fake_.base, ref)); + EXPECT_EQ(fake_.flush_count, 1); + EXPECT_EQ(fake_.invalidate_count, 0); + ExpectBufferRefEq(ref, fake_.flush_ref); +} + +TEST_F(CommandBufferTest, InvalidateDispatchesExactBufferRef) { + IREE_ASSERT_OK(Initialize()); + // Unvalidated dispatch only: the buffer is an identity token, never accessed. + iree_hal_buffer_t buffer = {}; + auto ref = iree_hal_make_buffer_ref(&buffer, 41, 61); + IREE_ASSERT_OK(iree_hal_command_buffer_invalidate_buffer(&fake_.base, ref)); + EXPECT_EQ(fake_.invalidate_count, 1); + EXPECT_EQ(fake_.flush_count, 0); + ExpectBufferRefEq(ref, fake_.invalidate_ref); +} + +TEST_F(CommandBufferTest, FlushPropagatesCallbackError) { + IREE_ASSERT_OK(Initialize()); + fake_.flush_status = IREE_STATUS_ABORTED; + IREE_EXPECT_STATUS_IS( + IREE_STATUS_ABORTED, + iree_hal_command_buffer_flush_buffer( + &fake_.base, iree_hal_make_indirect_buffer_ref(3, 17, 29))); + EXPECT_EQ(fake_.flush_count, 1); +} + +TEST_F(CommandBufferTest, InvalidatePropagatesCallbackError) { + IREE_ASSERT_OK(Initialize()); + fake_.invalidate_status = IREE_STATUS_DATA_LOSS; + IREE_EXPECT_STATUS_IS( + IREE_STATUS_DATA_LOSS, + iree_hal_command_buffer_invalidate_buffer( + &fake_.base, iree_hal_make_indirect_buffer_ref(3, 17, 29))); + EXPECT_EQ(fake_.invalidate_count, 1); +} + +TEST_F(CommandBufferTest, ZeroLengthFlushIsNoOp) { + IREE_ASSERT_OK(Initialize()); + fake_.flush_status = IREE_STATUS_ABORTED; + IREE_EXPECT_OK(iree_hal_command_buffer_flush_buffer( + &fake_.base, iree_hal_make_indirect_buffer_ref(3, 17, 0))); + EXPECT_EQ(fake_.flush_count, 0); +} + +TEST_F(CommandBufferTest, ZeroLengthInvalidateIsNoOp) { + IREE_ASSERT_OK(Initialize()); + fake_.invalidate_status = IREE_STATUS_DATA_LOSS; + IREE_EXPECT_OK(iree_hal_command_buffer_invalidate_buffer( + &fake_.base, iree_hal_make_indirect_buffer_ref(3, 17, 0))); + EXPECT_EQ(fake_.invalidate_count, 0); +} + +TEST_F(CommandBufferTest, FlushWithoutBackendCallbackIsUnimplemented) { + vtable_.flush_buffer = nullptr; + IREE_ASSERT_OK(Initialize()); + IREE_EXPECT_STATUS_IS( + IREE_STATUS_UNIMPLEMENTED, + iree_hal_command_buffer_flush_buffer( + &fake_.base, iree_hal_make_indirect_buffer_ref(3, 17, 29))); + EXPECT_EQ(fake_.flush_count, 0); +} + +TEST_F(CommandBufferTest, InvalidateWithoutBackendCallbackIsUnimplemented) { + vtable_.invalidate_buffer = nullptr; + IREE_ASSERT_OK(Initialize()); + IREE_EXPECT_STATUS_IS( + IREE_STATUS_UNIMPLEMENTED, + iree_hal_command_buffer_invalidate_buffer( + &fake_.base, iree_hal_make_indirect_buffer_ref(3, 17, 29))); + EXPECT_EQ(fake_.invalidate_count, 0); +} + +TEST_F(CommandBufferTest, ValidationRejectsUnsupportedCommandCategory) { + if (iree_hal_command_buffer_validation_state_size( + /*mode=*/0, /*binding_capacity=*/0) == 0) { + GTEST_SKIP() << "Command buffer validation is compiled out"; + } + IREE_ASSERT_OK(Initialize(/*mode=*/0, IREE_HAL_COMMAND_CATEGORY_DISPATCH)); + IREE_ASSERT_OK(iree_hal_command_buffer_begin(&fake_.base)); + auto ref = iree_hal_make_indirect_buffer_ref(0, 17, 29); + IREE_EXPECT_STATUS_IS(IREE_STATUS_FAILED_PRECONDITION, + iree_hal_command_buffer_flush_buffer(&fake_.base, ref)); + IREE_EXPECT_STATUS_IS( + IREE_STATUS_FAILED_PRECONDITION, + iree_hal_command_buffer_invalidate_buffer(&fake_.base, ref)); + EXPECT_EQ(fake_.flush_count, 0); + EXPECT_EQ(fake_.invalidate_count, 0); + IREE_EXPECT_OK(iree_hal_command_buffer_end(&fake_.base)); +} + +TEST_F(CommandBufferTest, FlushValidationRejectsNullBinding) { + if (iree_hal_command_buffer_validation_state_size( + /*mode=*/0, /*binding_capacity=*/1) == 0) { + GTEST_SKIP() << "Command buffer validation is compiled out"; + } + IREE_ASSERT_OK(Initialize(/*mode=*/0, IREE_HAL_COMMAND_CATEGORY_TRANSFER, + /*binding_capacity=*/1)); + IREE_ASSERT_OK(iree_hal_command_buffer_begin(&fake_.base)); + IREE_ASSERT_OK(iree_hal_command_buffer_flush_buffer( + &fake_.base, iree_hal_make_indirect_buffer_ref(0, 8, 16))); + IREE_ASSERT_OK(iree_hal_command_buffer_end(&fake_.base)); + + iree_hal_buffer_binding_t binding = {}; + iree_hal_buffer_binding_table_t binding_table = {}; + binding_table.count = 1; + binding_table.bindings = &binding; + IREE_EXPECT_STATUS_IS( + IREE_STATUS_INVALID_ARGUMENT, + iree_hal_command_buffer_validate_submission(&fake_.base, binding_table)); +} + +TEST_F(CommandBufferTest, InvalidateValidationRejectsNullBinding) { + if (iree_hal_command_buffer_validation_state_size( + /*mode=*/0, /*binding_capacity=*/1) == 0) { + GTEST_SKIP() << "Command buffer validation is compiled out"; + } + IREE_ASSERT_OK(Initialize(/*mode=*/0, IREE_HAL_COMMAND_CATEGORY_TRANSFER, + /*binding_capacity=*/1)); + IREE_ASSERT_OK(iree_hal_command_buffer_begin(&fake_.base)); + IREE_ASSERT_OK(iree_hal_command_buffer_invalidate_buffer( + &fake_.base, iree_hal_make_indirect_buffer_ref(0, 8, 16))); + IREE_ASSERT_OK(iree_hal_command_buffer_end(&fake_.base)); + + iree_hal_buffer_binding_t binding = {}; + iree_hal_buffer_binding_table_t binding_table = {}; + binding_table.count = 1; + binding_table.bindings = &binding; + IREE_EXPECT_STATUS_IS( + IREE_STATUS_INVALID_ARGUMENT, + iree_hal_command_buffer_validate_submission(&fake_.base, binding_table)); +} + +TEST_F(CommandBufferTest, FlushValidationRejectsUndersizedBinding) { + if (iree_hal_command_buffer_validation_state_size( + /*mode=*/0, /*binding_capacity=*/1) == 0) { + GTEST_SKIP() << "Command buffer validation is compiled out"; + } + IREE_ASSERT_OK(Initialize(/*mode=*/0, IREE_HAL_COMMAND_CATEGORY_TRANSFER, + /*binding_capacity=*/1)); + IREE_ASSERT_OK(iree_hal_command_buffer_begin(&fake_.base)); + IREE_ASSERT_OK(iree_hal_command_buffer_flush_buffer( + &fake_.base, iree_hal_make_indirect_buffer_ref(0, 8, 16))); + IREE_ASSERT_OK(iree_hal_command_buffer_end(&fake_.base)); + + FakeBuffer buffer = {}; + iree_hal_buffer_initialize( + /*placement=*/{}, /*allocated_buffer=*/&buffer.base, + /*allocation_size=*/sizeof(buffer.storage), /*byte_offset=*/0, + /*byte_length=*/sizeof(buffer.storage), + IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE, IREE_HAL_MEMORY_ACCESS_ALL, + IREE_HAL_BUFFER_USAGE_DEFAULT, &kFakeBufferVTable, &buffer.base); + + iree_hal_buffer_binding_t binding = {}; + binding.buffer = &buffer.base; + binding.offset = 0; + binding.length = 16; + iree_hal_buffer_binding_table_t binding_table = {}; + binding_table.count = 1; + binding_table.bindings = &binding; + IREE_EXPECT_STATUS_IS( + IREE_STATUS_OUT_OF_RANGE, + iree_hal_command_buffer_validate_submission(&fake_.base, binding_table)); + iree_hal_buffer_release(&buffer.base); +} + +TEST_F(CommandBufferTest, InvalidateValidationRejectsUndersizedBinding) { + if (iree_hal_command_buffer_validation_state_size( + /*mode=*/0, /*binding_capacity=*/1) == 0) { + GTEST_SKIP() << "Command buffer validation is compiled out"; + } + IREE_ASSERT_OK(Initialize(/*mode=*/0, IREE_HAL_COMMAND_CATEGORY_TRANSFER, + /*binding_capacity=*/1)); + IREE_ASSERT_OK(iree_hal_command_buffer_begin(&fake_.base)); + IREE_ASSERT_OK(iree_hal_command_buffer_invalidate_buffer( + &fake_.base, iree_hal_make_indirect_buffer_ref(0, 8, 16))); + IREE_ASSERT_OK(iree_hal_command_buffer_end(&fake_.base)); + + FakeBuffer buffer = {}; + iree_hal_buffer_initialize( + /*placement=*/{}, /*allocated_buffer=*/&buffer.base, + /*allocation_size=*/sizeof(buffer.storage), /*byte_offset=*/0, + /*byte_length=*/sizeof(buffer.storage), + IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE, IREE_HAL_MEMORY_ACCESS_ALL, + IREE_HAL_BUFFER_USAGE_DEFAULT, &kFakeBufferVTable, &buffer.base); + + iree_hal_buffer_binding_t binding = {}; + binding.buffer = &buffer.base; + binding.offset = 0; + binding.length = 16; + iree_hal_buffer_binding_table_t binding_table = {}; + binding_table.count = 1; + binding_table.bindings = &binding; + IREE_EXPECT_STATUS_IS( + IREE_STATUS_OUT_OF_RANGE, + iree_hal_command_buffer_validate_submission(&fake_.base, binding_table)); + iree_hal_buffer_release(&buffer.base); +} + +} // namespace +} // namespace iree
diff --git a/runtime/src/iree/hal/command_buffer_validation.c b/runtime/src/iree/hal/command_buffer_validation.c index ed96fb3..138bae6 100644 --- a/runtime/src/iree/hal/command_buffer_validation.c +++ b/runtime/src/iree/hal/command_buffer_validation.c
@@ -101,9 +101,15 @@ const iree_hal_command_buffer_validation_state_t* validation_state, iree_hal_buffer_binding_t binding, iree_hal_buffer_binding_requirements_t requirements) { - // Check for binding presence. - if (requirements.usage == IREE_HAL_BUFFER_USAGE_NONE) { - // Binding slot is unused and its value in the table is ignored. + // Check for binding presence. A slot is unused only when it has no + // requirements at all; some operations require a buffer without imposing a + // specific usage or access mode. + if (requirements.required_compatibility == 0 && + requirements.usage == IREE_HAL_BUFFER_USAGE_NONE && + requirements.access == IREE_HAL_MEMORY_ACCESS_NONE && + requirements.type == IREE_HAL_MEMORY_TYPE_NONE && + requirements.max_byte_offset == 0 && + requirements.min_byte_alignment == 0) { return iree_ok_status(); } else if (!binding.buffer) { // Binding is used and required. @@ -112,18 +118,23 @@ "binding table slot requires a buffer but none was provided"); } - // Ensure the buffer is compatible with the device. + // Ensure the buffer is compatible with the device when the operation + // requires specific queue compatibility. // NOTE: this check is very slow! We may want to disable this outside of debug // mode or try to fast path it if the buffer is known-good. - IREE_RETURN_IF_ERROR(iree_hal_command_buffer_validate_buffer_compatibility( - command_buffer, validation_state, binding.buffer, - requirements.required_compatibility, requirements.usage)); + if (requirements.required_compatibility) { + IREE_RETURN_IF_ERROR(iree_hal_command_buffer_validate_buffer_compatibility( + command_buffer, validation_state, binding.buffer, + requirements.required_compatibility, requirements.usage)); + } // Verify buffer compatibility. IREE_RETURN_IF_ERROR(iree_hal_buffer_validate_usage( iree_hal_buffer_allowed_usage(binding.buffer), requirements.usage)); - IREE_RETURN_IF_ERROR(iree_hal_buffer_validate_access( - iree_hal_buffer_allowed_access(binding.buffer), requirements.access)); + if (requirements.access != IREE_HAL_MEMORY_ACCESS_NONE) { + IREE_RETURN_IF_ERROR(iree_hal_buffer_validate_access( + iree_hal_buffer_allowed_access(binding.buffer), requirements.access)); + } IREE_RETURN_IF_ERROR(iree_hal_buffer_validate_memory_type( iree_hal_buffer_memory_type(binding.buffer), requirements.type)); @@ -352,6 +363,40 @@ return iree_ok_status(); } +iree_status_t iree_hal_command_buffer_flush_buffer_validation( + iree_hal_command_buffer_t* command_buffer, + iree_hal_command_buffer_validation_state_t* validation_state, + iree_hal_buffer_ref_t target_ref) { + IREE_RETURN_IF_ERROR(iree_hal_command_buffer_validate_categories( + command_buffer, validation_state, IREE_HAL_COMMAND_CATEGORY_TRANSFER)); + + const iree_hal_buffer_binding_requirements_t target_reqs = { + .type = IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE, + .max_byte_offset = target_ref.offset + target_ref.length, + }; + IREE_RETURN_IF_ERROR(iree_hal_command_buffer_validate_buffer_requirements( + command_buffer, validation_state, target_ref, target_reqs)); + + return iree_ok_status(); +} + +iree_status_t iree_hal_command_buffer_invalidate_buffer_validation( + iree_hal_command_buffer_t* command_buffer, + iree_hal_command_buffer_validation_state_t* validation_state, + iree_hal_buffer_ref_t target_ref) { + IREE_RETURN_IF_ERROR(iree_hal_command_buffer_validate_categories( + command_buffer, validation_state, IREE_HAL_COMMAND_CATEGORY_TRANSFER)); + + const iree_hal_buffer_binding_requirements_t target_reqs = { + .type = IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE, + .max_byte_offset = target_ref.offset + target_ref.length, + }; + IREE_RETURN_IF_ERROR(iree_hal_command_buffer_validate_buffer_requirements( + command_buffer, validation_state, target_ref, target_reqs)); + + return iree_ok_status(); +} + iree_status_t iree_hal_command_buffer_fill_buffer_validation( iree_hal_command_buffer_t* command_buffer, iree_hal_command_buffer_validation_state_t* validation_state,
diff --git a/runtime/src/iree/hal/command_buffer_validation.h b/runtime/src/iree/hal/command_buffer_validation.h index a91c01c..95770b9 100644 --- a/runtime/src/iree/hal/command_buffer_validation.h +++ b/runtime/src/iree/hal/command_buffer_validation.h
@@ -39,8 +39,8 @@ // TODO(benvanik): current pipeline layout/descriptor set layout info. // TODO(benvanik): valid push constant bit ranges. // Requirements for each binding table entry. - // Unused slots in the binding table will have IREE_HAL_BUFFER_USAGE_NONE and - // are ignored if set when executed. + // Unused slots have no compatibility, usage, access, memory type, range, or + // alignment requirements and are ignored if set when executed. iree_hal_buffer_binding_requirements_t binding_requirements[0]; } iree_hal_command_buffer_validation_state_t; @@ -105,6 +105,16 @@ iree_hal_buffer_ref_t buffer_ref, iree_hal_memory_advise_flags_t flags, uint64_t arg0, uint64_t arg1); +iree_status_t iree_hal_command_buffer_flush_buffer_validation( + iree_hal_command_buffer_t* command_buffer, + iree_hal_command_buffer_validation_state_t* validation_state, + iree_hal_buffer_ref_t target_ref); + +iree_status_t iree_hal_command_buffer_invalidate_buffer_validation( + iree_hal_command_buffer_t* command_buffer, + iree_hal_command_buffer_validation_state_t* validation_state, + iree_hal_buffer_ref_t target_ref); + iree_status_t iree_hal_command_buffer_fill_buffer_validation( iree_hal_command_buffer_t* command_buffer, iree_hal_command_buffer_validation_state_t* validation_state,
diff --git a/runtime/src/iree/hal/local/BUILD.bazel b/runtime/src/iree/hal/local/BUILD.bazel index 3f5e1be..78b5ca7 100644 --- a/runtime/src/iree/hal/local/BUILD.bazel +++ b/runtime/src/iree/hal/local/BUILD.bazel
@@ -127,6 +127,18 @@ ], ) +iree_runtime_cc_test( + name = "inline_command_buffer_test", + srcs = ["inline_command_buffer_test.cc"], + deps = [ + ":local", + "//runtime/src/iree/base", + "//runtime/src/iree/hal", + "//runtime/src/iree/testing:gtest", + "//runtime/src/iree/testing:gtest_main", + ], +) + iree_runtime_cc_library( name = "inline_dispatch", srcs = ["inline_dispatch.c"],
diff --git a/runtime/src/iree/hal/local/CMakeLists.txt b/runtime/src/iree/hal/local/CMakeLists.txt index bf4ad5c..de5c8a8 100644 --- a/runtime/src/iree/hal/local/CMakeLists.txt +++ b/runtime/src/iree/hal/local/CMakeLists.txt
@@ -143,6 +143,19 @@ PUBLIC ) +iree_cc_test( + NAME + inline_command_buffer_test + SRCS + "inline_command_buffer_test.cc" + DEPS + ::local + iree::base + iree::hal + iree::testing::gtest + iree::testing::gtest_main +) + iree_cc_library( NAME inline_dispatch
diff --git a/runtime/src/iree/hal/local/inline_command_buffer.c b/runtime/src/iree/hal/local/inline_command_buffer.c index e67a1a2..1aaf1a1 100644 --- a/runtime/src/iree/hal/local/inline_command_buffer.c +++ b/runtime/src/iree/hal/local/inline_command_buffer.c
@@ -357,6 +357,44 @@ } //===----------------------------------------------------------------------===// +// iree_hal_command_buffer_flush_buffer +//===----------------------------------------------------------------------===// + +static iree_status_t iree_hal_inline_command_buffer_flush_buffer( + iree_hal_command_buffer_t* base_command_buffer, + iree_hal_buffer_ref_t target_ref) { + if (target_ref.length == 0) return iree_ok_status(); + + iree_hal_buffer_mapping_t mapping; + IREE_RETURN_IF_ERROR( + iree_hal_buffer_map_range(target_ref.buffer, IREE_HAL_MAPPING_MODE_SCOPED, + IREE_HAL_MEMORY_ACCESS_WRITE, target_ref.offset, + target_ref.length, &mapping)); + iree_status_t status = + iree_hal_buffer_mapping_flush_range(&mapping, 0, IREE_HAL_WHOLE_BUFFER); + return iree_status_join(status, iree_hal_buffer_unmap_range(&mapping)); +} + +//===----------------------------------------------------------------------===// +// iree_hal_command_buffer_invalidate_buffer +//===----------------------------------------------------------------------===// + +static iree_status_t iree_hal_inline_command_buffer_invalidate_buffer( + iree_hal_command_buffer_t* base_command_buffer, + iree_hal_buffer_ref_t target_ref) { + if (target_ref.length == 0) return iree_ok_status(); + + iree_hal_buffer_mapping_t mapping; + IREE_RETURN_IF_ERROR( + iree_hal_buffer_map_range(target_ref.buffer, IREE_HAL_MAPPING_MODE_SCOPED, + IREE_HAL_MEMORY_ACCESS_READ, target_ref.offset, + target_ref.length, &mapping)); + iree_status_t status = iree_hal_buffer_mapping_invalidate_range( + &mapping, 0, IREE_HAL_WHOLE_BUFFER); + return iree_status_join(status, iree_hal_buffer_unmap_range(&mapping)); +} + +//===----------------------------------------------------------------------===// // iree_hal_command_buffer_fill_buffer //===----------------------------------------------------------------------===// @@ -624,6 +662,8 @@ .reset_event = iree_hal_inline_command_buffer_reset_event, .wait_events = iree_hal_inline_command_buffer_wait_events, .advise_buffer = iree_hal_inline_command_buffer_advise_buffer, + .flush_buffer = iree_hal_inline_command_buffer_flush_buffer, + .invalidate_buffer = iree_hal_inline_command_buffer_invalidate_buffer, .fill_buffer = iree_hal_inline_command_buffer_fill_buffer, .update_buffer = iree_hal_inline_command_buffer_update_buffer, .copy_buffer = iree_hal_inline_command_buffer_copy_buffer,
diff --git a/runtime/src/iree/hal/local/inline_command_buffer_test.cc b/runtime/src/iree/hal/local/inline_command_buffer_test.cc new file mode 100644 index 0000000..f67577c --- /dev/null +++ b/runtime/src/iree/hal/local/inline_command_buffer_test.cc
@@ -0,0 +1,256 @@ +// Copyright 2026 The IREE Authors +// +// Licensed under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#include "iree/hal/local/inline_command_buffer.h" + +#include "iree/base/api.h" +#include "iree/hal/api.h" +#include "iree/testing/gtest.h" +#include "iree/testing/status_matchers.h" + +namespace iree::hal::local { +namespace { + +struct FakeBuffer { + iree_hal_buffer_t base; + uint8_t storage[256] = {}; + int map_count = 0; + int unmap_count = 0; + int flush_count = 0; + int invalidate_count = 0; + iree_hal_mapping_mode_t mapping_mode = 0; + iree_hal_memory_access_t memory_access = 0; + iree_device_size_t mapped_offset = 0; + iree_device_size_t mapped_length = 0; + iree_device_size_t flush_offset = 0; + iree_device_size_t flush_length = 0; + iree_device_size_t invalidate_offset = 0; + iree_device_size_t invalidate_length = 0; + iree_status_code_t map_status = IREE_STATUS_OK; + iree_status_code_t unmap_status = IREE_STATUS_OK; + iree_status_code_t flush_status = IREE_STATUS_OK; + iree_status_code_t invalidate_status = IREE_STATUS_OK; +}; + +static void IREE_API_PTR Destroy(iree_hal_buffer_t* /*buffer*/) { + // The fixture owns both the buffer and its storage; neither may be freed + // here. +} + +static iree_status_t IREE_API_PTR +MapRange(iree_hal_buffer_t* base_buffer, iree_hal_mapping_mode_t mapping_mode, + iree_hal_memory_access_t memory_access, iree_device_size_t byte_offset, + iree_device_size_t byte_length, iree_hal_buffer_mapping_t* mapping) { + auto* buffer = reinterpret_cast<FakeBuffer*>(base_buffer); + ++buffer->map_count; + buffer->mapping_mode = mapping_mode; + buffer->memory_access = memory_access; + buffer->mapped_offset = byte_offset; + buffer->mapped_length = byte_length; + if (buffer->map_status != IREE_STATUS_OK) { + return iree_status_from_code(buffer->map_status); + } + // Preserve the HAL-populated mapping metadata and provide real storage so + // iree_hal_buffer_unmap_range dispatches to our unmap callback. + mapping->contents = + iree_make_byte_span(buffer->storage + byte_offset, byte_length); + return iree_ok_status(); +} + +static iree_status_t IREE_API_PTR +UnmapRange(iree_hal_buffer_t* base_buffer, iree_device_size_t byte_offset, + iree_device_size_t byte_length, iree_hal_buffer_mapping_t* mapping) { + auto* buffer = reinterpret_cast<FakeBuffer*>(base_buffer); + ++buffer->unmap_count; + EXPECT_EQ(buffer->mapped_offset, byte_offset); + EXPECT_EQ(buffer->mapped_length, byte_length); + EXPECT_EQ(buffer->storage + byte_offset, mapping->contents.data); + return iree_status_from_code(buffer->unmap_status); +} + +static iree_status_t IREE_API_PTR +InvalidateRange(iree_hal_buffer_t* base_buffer, iree_device_size_t byte_offset, + iree_device_size_t byte_length) { + auto* buffer = reinterpret_cast<FakeBuffer*>(base_buffer); + ++buffer->invalidate_count; + buffer->invalidate_offset = byte_offset; + buffer->invalidate_length = byte_length; + EXPECT_EQ(buffer->map_count, 1); + EXPECT_EQ(buffer->unmap_count, 0); + return iree_status_from_code(buffer->invalidate_status); +} + +static iree_status_t IREE_API_PTR FlushRange(iree_hal_buffer_t* base_buffer, + iree_device_size_t byte_offset, + iree_device_size_t byte_length) { + auto* buffer = reinterpret_cast<FakeBuffer*>(base_buffer); + ++buffer->flush_count; + buffer->flush_offset = byte_offset; + buffer->flush_length = byte_length; + EXPECT_EQ(buffer->map_count, 1); + EXPECT_EQ(buffer->unmap_count, 0); + return iree_status_from_code(buffer->flush_status); +} + +static const iree_hal_buffer_vtable_t kFakeBufferVTable = { + /*.recycle=*/iree_hal_buffer_recycle, + /*.destroy=*/Destroy, + /*.map_range=*/MapRange, + /*.unmap_range=*/UnmapRange, + /*.invalidate_range=*/InvalidateRange, + /*.flush_range=*/FlushRange, +}; + +class InlineCommandBufferTest : public ::testing::Test { + protected: + void SetUp() override { + iree_hal_buffer_initialize( + /*placement=*/{}, /*allocated_buffer=*/&buffer_.base, + /*allocation_size=*/sizeof(buffer_.storage), /*byte_offset=*/0, + /*byte_length=*/sizeof(buffer_.storage), + IREE_HAL_MEMORY_TYPE_HOST_VISIBLE, + IREE_HAL_MEMORY_ACCESS_READ | IREE_HAL_MEMORY_ACCESS_WRITE, + IREE_HAL_BUFFER_USAGE_MAPPING_SCOPED | IREE_HAL_BUFFER_USAGE_TRANSFER, + &kFakeBufferVTable, &buffer_.base); + IREE_ASSERT_OK(iree_hal_inline_command_buffer_create( + /*device_allocator=*/nullptr, + IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT | + IREE_HAL_COMMAND_BUFFER_MODE_ALLOW_INLINE_EXECUTION | + IREE_HAL_COMMAND_BUFFER_MODE_UNVALIDATED, + IREE_HAL_COMMAND_CATEGORY_TRANSFER, IREE_HAL_QUEUE_AFFINITY_ANY, + /*binding_capacity=*/0, iree_allocator_system(), &command_buffer_)); + IREE_ASSERT_OK(iree_hal_command_buffer_begin(command_buffer_)); + } + + void TearDown() override { + if (command_buffer_) { + IREE_EXPECT_OK(iree_hal_command_buffer_end(command_buffer_)); + iree_hal_command_buffer_release(command_buffer_); + } + iree_hal_buffer_release(&buffer_.base); + } + + void ExpectCounts(int map, int flush, int invalidate, int unmap) { + EXPECT_EQ(buffer_.map_count, map); + EXPECT_EQ(buffer_.flush_count, flush); + EXPECT_EQ(buffer_.invalidate_count, invalidate); + EXPECT_EQ(buffer_.unmap_count, unmap); + } + + FakeBuffer buffer_ = {}; + iree_hal_command_buffer_t* command_buffer_ = nullptr; +}; + +TEST_F(InlineCommandBufferTest, FlushMapsCorrectRangeAndAccess) { + IREE_ASSERT_OK(iree_hal_command_buffer_flush_buffer( + command_buffer_, iree_hal_make_buffer_ref(&buffer_.base, 37, 59))); + // All callbacks must have completed before the command returns, without end + // or submission of the command buffer. + ExpectCounts(/*map=*/1, /*flush=*/1, /*invalidate=*/0, /*unmap=*/1); + EXPECT_EQ(buffer_.mapping_mode, IREE_HAL_MAPPING_MODE_SCOPED); + EXPECT_EQ(buffer_.memory_access, IREE_HAL_MEMORY_ACCESS_WRITE); + EXPECT_EQ(buffer_.mapped_offset, 37); + EXPECT_EQ(buffer_.mapped_length, 59); + EXPECT_EQ(buffer_.flush_offset, 37); + EXPECT_EQ(buffer_.flush_length, 59); +} + +TEST_F(InlineCommandBufferTest, InvalidateMapsCorrectRangeAndAccess) { + IREE_ASSERT_OK(iree_hal_command_buffer_invalidate_buffer( + command_buffer_, iree_hal_make_buffer_ref(&buffer_.base, 41, 61))); + // Check synchronous execution before ending or submitting the command buffer. + ExpectCounts(/*map=*/1, /*flush=*/0, /*invalidate=*/1, /*unmap=*/1); + EXPECT_EQ(buffer_.mapping_mode, IREE_HAL_MAPPING_MODE_SCOPED); + EXPECT_EQ(buffer_.memory_access, IREE_HAL_MEMORY_ACCESS_READ); + EXPECT_EQ(buffer_.mapped_offset, 41); + EXPECT_EQ(buffer_.mapped_length, 61); + EXPECT_EQ(buffer_.invalidate_offset, 41); + EXPECT_EQ(buffer_.invalidate_length, 61); +} + +TEST_F(InlineCommandBufferTest, FlushUsesMappedRange) { + IREE_ASSERT_OK(iree_hal_command_buffer_flush_buffer( + command_buffer_, + iree_hal_make_buffer_ref(&buffer_.base, 53, IREE_HAL_WHOLE_BUFFER))); + // Mapping-relative 0/WHOLE_BUFFER must resolve to the mapped tail, not the + // entire allocation or an untranslated WHOLE_BUFFER sentinel. + ExpectCounts(/*map=*/1, /*flush=*/1, /*invalidate=*/0, /*unmap=*/1); + EXPECT_EQ(buffer_.mapped_offset, 53); + EXPECT_EQ(buffer_.mapped_length, sizeof(buffer_.storage) - 53); + EXPECT_EQ(buffer_.flush_offset, buffer_.mapped_offset); + EXPECT_EQ(buffer_.flush_length, buffer_.mapped_length); +} + +TEST_F(InlineCommandBufferTest, InvalidateUsesMappedRange) { + IREE_ASSERT_OK(iree_hal_command_buffer_invalidate_buffer( + command_buffer_, + iree_hal_make_buffer_ref(&buffer_.base, 67, IREE_HAL_WHOLE_BUFFER))); + // The invalidate helper must also translate mapping-relative 0/WHOLE_BUFFER. + ExpectCounts(/*map=*/1, /*flush=*/0, /*invalidate=*/1, /*unmap=*/1); + EXPECT_EQ(buffer_.mapped_offset, 67); + EXPECT_EQ(buffer_.mapped_length, sizeof(buffer_.storage) - 67); + EXPECT_EQ(buffer_.invalidate_offset, buffer_.mapped_offset); + EXPECT_EQ(buffer_.invalidate_length, buffer_.mapped_length); +} + +TEST_F(InlineCommandBufferTest, FlushErrorPropagatesAndStillUnmaps) { + buffer_.flush_status = IREE_STATUS_ABORTED; + IREE_EXPECT_STATUS_IS( + IREE_STATUS_ABORTED, + iree_hal_command_buffer_flush_buffer( + command_buffer_, iree_hal_make_buffer_ref(&buffer_.base, 37, 59))); + ExpectCounts(/*map=*/1, /*flush=*/1, /*invalidate=*/0, /*unmap=*/1); +} + +TEST_F(InlineCommandBufferTest, InvalidateErrorPropagatesAndStillUnmaps) { + buffer_.invalidate_status = IREE_STATUS_DATA_LOSS; + IREE_EXPECT_STATUS_IS( + IREE_STATUS_DATA_LOSS, + iree_hal_command_buffer_invalidate_buffer( + command_buffer_, iree_hal_make_buffer_ref(&buffer_.base, 41, 61))); + ExpectCounts(/*map=*/1, /*flush=*/0, /*invalidate=*/1, /*unmap=*/1); +} + +TEST_F(InlineCommandBufferTest, + FlushMapErrorPropagatesWithoutFurtherCallbacks) { + buffer_.map_status = IREE_STATUS_UNAVAILABLE; + IREE_EXPECT_STATUS_IS( + IREE_STATUS_UNAVAILABLE, + iree_hal_command_buffer_flush_buffer( + command_buffer_, iree_hal_make_buffer_ref(&buffer_.base, 37, 59))); + ExpectCounts(/*map=*/1, /*flush=*/0, /*invalidate=*/0, /*unmap=*/0); +} + +TEST_F(InlineCommandBufferTest, + InvalidateMapErrorPropagatesWithoutFurtherCallbacks) { + buffer_.map_status = IREE_STATUS_UNAVAILABLE; + IREE_EXPECT_STATUS_IS( + IREE_STATUS_UNAVAILABLE, + iree_hal_command_buffer_invalidate_buffer( + command_buffer_, iree_hal_make_buffer_ref(&buffer_.base, 41, 61))); + ExpectCounts(/*map=*/1, /*flush=*/0, /*invalidate=*/0, /*unmap=*/0); +} + +TEST_F(InlineCommandBufferTest, FlushUnmapErrorPropagates) { + buffer_.unmap_status = IREE_STATUS_INTERNAL; + IREE_EXPECT_STATUS_IS( + IREE_STATUS_INTERNAL, + iree_hal_command_buffer_flush_buffer( + command_buffer_, iree_hal_make_buffer_ref(&buffer_.base, 37, 59))); + ExpectCounts(/*map=*/1, /*flush=*/1, /*invalidate=*/0, /*unmap=*/1); +} + +TEST_F(InlineCommandBufferTest, InvalidateUnmapErrorPropagates) { + buffer_.unmap_status = IREE_STATUS_INTERNAL; + IREE_EXPECT_STATUS_IS( + IREE_STATUS_INTERNAL, + iree_hal_command_buffer_invalidate_buffer( + command_buffer_, iree_hal_make_buffer_ref(&buffer_.base, 41, 61))); + ExpectCounts(/*map=*/1, /*flush=*/0, /*invalidate=*/1, /*unmap=*/1); +} + +} // namespace +} // namespace iree::hal::local
diff --git a/runtime/src/iree/hal/utils/BUILD.bazel b/runtime/src/iree/hal/utils/BUILD.bazel index ee0f715..8de98c4 100644 --- a/runtime/src/iree/hal/utils/BUILD.bazel +++ b/runtime/src/iree/hal/utils/BUILD.bazel
@@ -70,6 +70,19 @@ ], ) +iree_runtime_cc_test( + name = "deferred_command_buffer_test", + srcs = ["deferred_command_buffer_test.cc"], + deps = [ + ":deferred_command_buffer", + "//runtime/src/iree/base", + "//runtime/src/iree/base/internal:arena", + "//runtime/src/iree/hal", + "//runtime/src/iree/testing:gtest", + "//runtime/src/iree/testing:gtest_main", + ], +) + iree_runtime_cc_library( name = "executable_debug_info", srcs = ["executable_debug_info.c"],
diff --git a/runtime/src/iree/hal/utils/CMakeLists.txt b/runtime/src/iree/hal/utils/CMakeLists.txt index 42c1de8..54fe20b 100644 --- a/runtime/src/iree/hal/utils/CMakeLists.txt +++ b/runtime/src/iree/hal/utils/CMakeLists.txt
@@ -82,6 +82,20 @@ PUBLIC ) +iree_cc_test( + NAME + deferred_command_buffer_test + SRCS + "deferred_command_buffer_test.cc" + DEPS + ::deferred_command_buffer + iree::base + iree::base::internal::arena + iree::hal + iree::testing::gtest + iree::testing::gtest_main +) + iree_cc_library( NAME executable_debug_info
diff --git a/runtime/src/iree/hal/utils/deferred_command_buffer.c b/runtime/src/iree/hal/utils/deferred_command_buffer.c index 034a4c7..ed49c5e 100644 --- a/runtime/src/iree/hal/utils/deferred_command_buffer.c +++ b/runtime/src/iree/hal/utils/deferred_command_buffer.c
@@ -19,6 +19,8 @@ IREE_HAL_CMD_RESET_EVENT, IREE_HAL_CMD_WAIT_EVENTS, IREE_HAL_CMD_ADVISE_BUFFER, + IREE_HAL_CMD_FLUSH_BUFFER, + IREE_HAL_CMD_INVALIDATE_BUFFER, IREE_HAL_CMD_FILL_BUFFER, IREE_HAL_CMD_UPDATE_BUFFER, IREE_HAL_CMD_COPY_BUFFER, @@ -478,6 +480,80 @@ } //===----------------------------------------------------------------------===// +// IREE_HAL_CMD_FLUSH_BUFFER +//===----------------------------------------------------------------------===// + +typedef struct iree_hal_cmd_flush_buffer_t { + iree_hal_cmd_header_t header; + iree_hal_buffer_ref_t target_ref; +} iree_hal_cmd_flush_buffer_t; + +static iree_status_t iree_hal_deferred_command_buffer_flush_buffer( + iree_hal_command_buffer_t* base_command_buffer, + iree_hal_buffer_ref_t target_ref) { + iree_hal_deferred_command_buffer_t* command_buffer = + iree_hal_deferred_command_buffer_cast(base_command_buffer); + iree_hal_cmd_list_t* cmd_list = &command_buffer->cmd_list; + if (target_ref.buffer) { + IREE_RETURN_IF_ERROR(iree_hal_resource_set_insert( + command_buffer->resource_set, 1, &target_ref.buffer)); + } + iree_hal_cmd_flush_buffer_t* cmd = NULL; + IREE_RETURN_IF_ERROR(iree_hal_cmd_list_append_command( + cmd_list, IREE_HAL_CMD_FLUSH_BUFFER, sizeof(*cmd), (void**)&cmd)); + cmd->target_ref = target_ref; + return iree_ok_status(); +} + +static iree_status_t iree_hal_deferred_command_buffer_apply_flush_buffer( + iree_hal_command_buffer_t* target_command_buffer, + iree_hal_buffer_binding_table_t binding_table, + const iree_hal_cmd_flush_buffer_t* cmd) { + iree_hal_buffer_ref_t target_ref; + IREE_RETURN_IF_ERROR(iree_hal_buffer_binding_table_resolve_ref( + binding_table, cmd->target_ref, &target_ref)); + return iree_hal_command_buffer_flush_buffer(target_command_buffer, + target_ref); +} + +//===----------------------------------------------------------------------===// +// IREE_HAL_CMD_INVALIDATE_BUFFER +//===----------------------------------------------------------------------===// + +typedef struct iree_hal_cmd_invalidate_buffer_t { + iree_hal_cmd_header_t header; + iree_hal_buffer_ref_t target_ref; +} iree_hal_cmd_invalidate_buffer_t; + +static iree_status_t iree_hal_deferred_command_buffer_invalidate_buffer( + iree_hal_command_buffer_t* base_command_buffer, + iree_hal_buffer_ref_t target_ref) { + iree_hal_deferred_command_buffer_t* command_buffer = + iree_hal_deferred_command_buffer_cast(base_command_buffer); + iree_hal_cmd_list_t* cmd_list = &command_buffer->cmd_list; + if (target_ref.buffer) { + IREE_RETURN_IF_ERROR(iree_hal_resource_set_insert( + command_buffer->resource_set, 1, &target_ref.buffer)); + } + iree_hal_cmd_invalidate_buffer_t* cmd = NULL; + IREE_RETURN_IF_ERROR(iree_hal_cmd_list_append_command( + cmd_list, IREE_HAL_CMD_INVALIDATE_BUFFER, sizeof(*cmd), (void**)&cmd)); + cmd->target_ref = target_ref; + return iree_ok_status(); +} + +static iree_status_t iree_hal_deferred_command_buffer_apply_invalidate_buffer( + iree_hal_command_buffer_t* target_command_buffer, + iree_hal_buffer_binding_table_t binding_table, + const iree_hal_cmd_invalidate_buffer_t* cmd) { + iree_hal_buffer_ref_t target_ref; + IREE_RETURN_IF_ERROR(iree_hal_buffer_binding_table_resolve_ref( + binding_table, cmd->target_ref, &target_ref)); + return iree_hal_command_buffer_invalidate_buffer(target_command_buffer, + target_ref); +} + +//===----------------------------------------------------------------------===// // IREE_HAL_CMD_FILL_BUFFER //===----------------------------------------------------------------------===// @@ -787,6 +863,10 @@ iree_hal_deferred_command_buffer_apply_wait_events, [IREE_HAL_CMD_ADVISE_BUFFER] = (iree_hal_cmd_apply_fn_t) iree_hal_deferred_command_buffer_apply_advise_buffer, + [IREE_HAL_CMD_FLUSH_BUFFER] = (iree_hal_cmd_apply_fn_t) + iree_hal_deferred_command_buffer_apply_flush_buffer, + [IREE_HAL_CMD_INVALIDATE_BUFFER] = (iree_hal_cmd_apply_fn_t) + iree_hal_deferred_command_buffer_apply_invalidate_buffer, [IREE_HAL_CMD_FILL_BUFFER] = (iree_hal_cmd_apply_fn_t) iree_hal_deferred_command_buffer_apply_fill_buffer, [IREE_HAL_CMD_UPDATE_BUFFER] = (iree_hal_cmd_apply_fn_t) @@ -846,6 +926,8 @@ .reset_event = iree_hal_deferred_command_buffer_reset_event, .wait_events = iree_hal_deferred_command_buffer_wait_events, .advise_buffer = iree_hal_deferred_command_buffer_advise_buffer, + .flush_buffer = iree_hal_deferred_command_buffer_flush_buffer, + .invalidate_buffer = iree_hal_deferred_command_buffer_invalidate_buffer, .fill_buffer = iree_hal_deferred_command_buffer_fill_buffer, .update_buffer = iree_hal_deferred_command_buffer_update_buffer, .copy_buffer = iree_hal_deferred_command_buffer_copy_buffer,
diff --git a/runtime/src/iree/hal/utils/deferred_command_buffer_test.cc b/runtime/src/iree/hal/utils/deferred_command_buffer_test.cc new file mode 100644 index 0000000..a5a5952 --- /dev/null +++ b/runtime/src/iree/hal/utils/deferred_command_buffer_test.cc
@@ -0,0 +1,264 @@ +// Copyright 2026 The IREE Authors +// +// Licensed under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#include "iree/hal/utils/deferred_command_buffer.h" + +#include "iree/base/api.h" +#include "iree/base/internal/arena.h" +#include "iree/hal/api.h" +#include "iree/testing/gtest.h" +#include "iree/testing/status_matchers.h" + +namespace iree { +namespace hal { +namespace { + +struct FakeBuffer { + iree_hal_buffer_t base; + int* destroy_count; +}; + +static void IREE_API_PTR DestroyBuffer(iree_hal_buffer_t* base_buffer) { + auto* buffer = reinterpret_cast<FakeBuffer*>(base_buffer); + ++*buffer->destroy_count; + delete buffer; +} + +static const iree_hal_buffer_vtable_t kBufferVTable = { + /*.recycle=*/iree_hal_buffer_recycle, + /*.destroy=*/DestroyBuffer, +}; + +struct FakeTargetCommandBuffer { + iree_hal_command_buffer_t base; + int flush_count = 0; + int invalidate_count = 0; + iree_hal_buffer_ref_t flush_ref = {}; + iree_hal_buffer_ref_t invalidate_ref = {}; + iree_status_code_t flush_status = IREE_STATUS_OK; + iree_status_code_t invalidate_status = IREE_STATUS_OK; +}; + +static void IREE_API_PTR DestroyTarget(iree_hal_command_buffer_t* base) { + delete reinterpret_cast<FakeTargetCommandBuffer*>(base); +} + +static iree_status_t IREE_API_PTR +BeginOrEndTarget(iree_hal_command_buffer_t* /*base*/) { + return iree_ok_status(); +} + +static iree_status_t IREE_API_PTR Flush(iree_hal_command_buffer_t* base, + iree_hal_buffer_ref_t ref) { + auto* target = reinterpret_cast<FakeTargetCommandBuffer*>(base); + ++target->flush_count; + target->flush_ref = ref; + return iree_status_from_code(target->flush_status); +} + +static iree_status_t IREE_API_PTR Invalidate(iree_hal_command_buffer_t* base, + iree_hal_buffer_ref_t ref) { + auto* target = reinterpret_cast<FakeTargetCommandBuffer*>(base); + ++target->invalidate_count; + target->invalidate_ref = ref; + return iree_status_from_code(target->invalidate_status); +} + +class DeferredCommandBufferTest : public ::testing::Test { + protected: + void SetUp() override { + iree_arena_block_pool_initialize(16 * 1024, iree_allocator_system(), + &block_pool_); + target_vtable_.destroy = DestroyTarget; + target_vtable_.begin = BeginOrEndTarget; + target_vtable_.end = BeginOrEndTarget; + target_vtable_.flush_buffer = Flush; + target_vtable_.invalidate_buffer = Invalidate; + target_ = new FakeTargetCommandBuffer{}; + iree_hal_command_buffer_initialize( + /*device_allocator=*/nullptr, + IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT | + IREE_HAL_COMMAND_BUFFER_MODE_UNVALIDATED, + IREE_HAL_COMMAND_CATEGORY_TRANSFER, IREE_HAL_QUEUE_AFFINITY_ANY, + /*binding_capacity=*/0, /*validation_state=*/nullptr, &target_vtable_, + &target_->base); + + auto* buffer = new FakeBuffer{}; + buffer->destroy_count = &buffer_destroy_count_; + iree_hal_buffer_initialize( + /*placement=*/{}, /*allocated_buffer=*/&buffer->base, + /*allocation_size=*/4096, /*byte_offset=*/0, /*byte_length=*/4096, + IREE_HAL_MEMORY_TYPE_HOST_VISIBLE, IREE_HAL_MEMORY_ACCESS_ALL, + IREE_HAL_BUFFER_USAGE_TRANSFER, &kBufferVTable, &buffer->base); + buffer_ = &buffer->base; + + IREE_ASSERT_OK(iree_hal_deferred_command_buffer_create( + /*device_allocator=*/nullptr, + IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT | + IREE_HAL_COMMAND_BUFFER_MODE_UNVALIDATED, + IREE_HAL_COMMAND_CATEGORY_TRANSFER, IREE_HAL_QUEUE_AFFINITY_ANY, + /*binding_capacity=*/4, &block_pool_, iree_allocator_system(), + &command_buffer_)); + } + + void TearDown() override { + iree_hal_command_buffer_release(command_buffer_); + iree_hal_command_buffer_release(&target_->base); + iree_hal_buffer_release(buffer_); + EXPECT_EQ(buffer_destroy_count_, 1); + // The pool must outlive the deferred command buffer and its resource set. + iree_arena_block_pool_deinitialize(&block_pool_); + } + + using RecordFn = iree_status_t (*)(iree_hal_command_buffer_t*, + iree_hal_buffer_ref_t); + + iree_status_t Record(RecordFn record, iree_hal_buffer_ref_t ref) { + IREE_RETURN_IF_ERROR(iree_hal_command_buffer_begin(command_buffer_)); + IREE_RETURN_IF_ERROR(record(command_buffer_, ref)); + ExpectCounts(0, 0); + IREE_RETURN_IF_ERROR(iree_hal_command_buffer_end(command_buffer_)); + ExpectCounts(0, 0); + return iree_ok_status(); + } + + iree_status_t Apply(iree_hal_buffer_binding_table_t bindings = + iree_hal_buffer_binding_table_empty()) { + return iree_hal_deferred_command_buffer_apply(command_buffer_, + &target_->base, bindings); + } + + void ExpectCounts(int flush, int invalidate) { + EXPECT_EQ(target_->flush_count, flush); + EXPECT_EQ(target_->invalidate_count, invalidate); + } + + static void ExpectRef(iree_hal_buffer_ref_t actual, iree_hal_buffer_t* buffer, + iree_device_size_t offset, iree_device_size_t length) { + EXPECT_EQ(actual.buffer, buffer); + EXPECT_EQ(actual.offset, offset); + EXPECT_EQ(actual.length, length); + EXPECT_EQ(static_cast<uint32_t>(actual.buffer_slot), 0u); + EXPECT_EQ(static_cast<uint32_t>(actual.reserved), 0u); + } + + void CheckRetention(RecordFn record, bool flush) { + auto* recorded_buffer = buffer_; + IREE_ASSERT_OK( + Record(record, iree_hal_make_buffer_ref(recorded_buffer, 137, 251))); + EXPECT_EQ(buffer_destroy_count_, 0); + iree_hal_buffer_release(buffer_); + buffer_ = nullptr; + ASSERT_EQ(buffer_destroy_count_, 0); + + IREE_ASSERT_OK(Apply()); + ExpectCounts(flush ? 1 : 0, flush ? 0 : 1); + ASSERT_EQ(buffer_destroy_count_, 0); + ExpectRef(flush ? target_->flush_ref : target_->invalidate_ref, + recorded_buffer, 137, 251); + + // Even one-shot apply keeps direct resources until deferred destruction. + iree_hal_command_buffer_release(command_buffer_); + command_buffer_ = nullptr; + EXPECT_EQ(buffer_destroy_count_, 1); + } + + iree_arena_block_pool_t block_pool_ = {}; + iree_hal_command_buffer_vtable_t target_vtable_ = {}; + FakeTargetCommandBuffer* target_ = nullptr; + iree_hal_command_buffer_t* command_buffer_ = nullptr; + int buffer_destroy_count_ = 0; + iree_hal_buffer_t* buffer_ = nullptr; +}; + +TEST_F(DeferredCommandBufferTest, DirectFlushExecutesOnlyDuringApply) { + // Also verifies that recording and replay preserve the exact direct range. + IREE_ASSERT_OK(Record(iree_hal_command_buffer_flush_buffer, + iree_hal_make_buffer_ref(buffer_, 137, 251))); + IREE_ASSERT_OK(Apply()); + ExpectCounts(1, 0); + ExpectRef(target_->flush_ref, buffer_, 137, 251); +} + +TEST_F(DeferredCommandBufferTest, DirectInvalidateExecutesOnlyDuringApply) { + IREE_ASSERT_OK(Record(iree_hal_command_buffer_invalidate_buffer, + iree_hal_make_buffer_ref(buffer_, 193, 347))); + IREE_ASSERT_OK(Apply()); + ExpectCounts(0, 1); + ExpectRef(target_->invalidate_ref, buffer_, 193, 347); +} + +TEST_F(DeferredCommandBufferTest, IndirectFlushResolvesBindingAtApply) { + IREE_ASSERT_OK(Record(iree_hal_command_buffer_flush_buffer, + iree_hal_make_indirect_buffer_ref(3, 37, 59))); + // The binding is supplied only after recording has completed. + iree_hal_buffer_binding_t bindings[4] = {}; + bindings[3] = {buffer_, 101, 503}; + IREE_ASSERT_OK(Apply({4, bindings})); + ExpectCounts(1, 0); + ExpectRef(target_->flush_ref, buffer_, 138, 59); +} + +TEST_F(DeferredCommandBufferTest, IndirectInvalidateResolvesBindingAtApply) { + IREE_ASSERT_OK(Record(iree_hal_command_buffer_invalidate_buffer, + iree_hal_make_indirect_buffer_ref(3, 41, 61))); + iree_hal_buffer_binding_t bindings[4] = {}; + bindings[3] = {buffer_, 103, 509}; + IREE_ASSERT_OK(Apply({4, bindings})); + ExpectCounts(0, 1); + ExpectRef(target_->invalidate_ref, buffer_, 144, 61); +} + +TEST_F(DeferredCommandBufferTest, MissingIndirectFlushBindingFailsDuringApply) { + IREE_ASSERT_OK(Record(iree_hal_command_buffer_flush_buffer, + iree_hal_make_indirect_buffer_ref(3, 37, 59))); + // An empty table resolves to a null ref successfully. A nonempty table with + // an out-of-bounds slot makes resolve_ref return OUT_OF_RANGE. + iree_hal_buffer_binding_t binding = {buffer_, 101, 503}; + IREE_EXPECT_STATUS_IS(IREE_STATUS_OUT_OF_RANGE, Apply({1, &binding})); + ExpectCounts(0, 0); +} + +TEST_F(DeferredCommandBufferTest, + MissingIndirectInvalidateBindingFailsDuringApply) { + IREE_ASSERT_OK(Record(iree_hal_command_buffer_invalidate_buffer, + iree_hal_make_indirect_buffer_ref(3, 41, 61))); + iree_hal_buffer_binding_t binding = {buffer_, 103, 509}; + IREE_EXPECT_STATUS_IS(IREE_STATUS_OUT_OF_RANGE, Apply({1, &binding})); + ExpectCounts(0, 0); +} + +TEST_F(DeferredCommandBufferTest, FlushTargetErrorPropagatesFromApply) { + target_->flush_status = IREE_STATUS_ABORTED; + IREE_ASSERT_OK(Record(iree_hal_command_buffer_flush_buffer, + iree_hal_make_buffer_ref(buffer_, 137, 251))); + IREE_EXPECT_STATUS_IS(IREE_STATUS_ABORTED, Apply()); + ExpectCounts(1, 0); + ExpectRef(target_->flush_ref, buffer_, 137, 251); +} + +TEST_F(DeferredCommandBufferTest, InvalidateTargetErrorPropagatesFromApply) { + target_->invalidate_status = IREE_STATUS_DATA_LOSS; + IREE_ASSERT_OK(Record(iree_hal_command_buffer_invalidate_buffer, + iree_hal_make_buffer_ref(buffer_, 193, 347))); + IREE_EXPECT_STATUS_IS(IREE_STATUS_DATA_LOSS, Apply()); + ExpectCounts(0, 1); + ExpectRef(target_->invalidate_ref, buffer_, 193, 347); +} + +TEST_F(DeferredCommandBufferTest, + DirectFlushBufferRetainedUntilDeferredDestruction) { + CheckRetention(iree_hal_command_buffer_flush_buffer, /*flush=*/true); +} + +TEST_F(DeferredCommandBufferTest, + DirectInvalidateBufferRetainedUntilDeferredDestruction) { + CheckRetention(iree_hal_command_buffer_invalidate_buffer, /*flush=*/false); +} + +} // namespace +} // namespace hal +} // namespace iree
diff --git a/runtime/src/iree/modules/hal/exports.inl b/runtime/src/iree/modules/hal/exports.inl index 75c47bd..8515dd6 100644 --- a/runtime/src/iree/modules/hal/exports.inl +++ b/runtime/src/iree/modules/hal/exports.inl
@@ -61,6 +61,8 @@ EXPORT_FN("command_buffer.execution_barrier", iree_hal_module_command_buffer_execution_barrier, riiI, v) EXPORT_FN("command_buffer.fill_buffer", iree_hal_module_command_buffer_fill_buffer, rrIIiIiI, v) EXPORT_FN("command_buffer.finalize", iree_hal_module_command_buffer_finalize, r, v) +EXPORT_FN("command_buffer.flush_buffer", iree_hal_module_command_buffer_flush_buffer, rirII, v) +EXPORT_FN("command_buffer.invalidate_buffer", iree_hal_module_command_buffer_invalidate_buffer, rirII, v) EXPORT_FN("command_buffer.update_buffer", iree_hal_module_command_buffer_update_buffer, rrIrIIiI, v) EXPORT_FN("device.allocator", iree_hal_module_device_allocator, r, r)
diff --git a/runtime/src/iree/modules/hal/module.c b/runtime/src/iree/modules/hal/module.c index f7cf848..f5f63a9 100644 --- a/runtime/src/iree/modules/hal/module.c +++ b/runtime/src/iree/modules/hal/module.c
@@ -941,6 +941,38 @@ flags, arg0, arg1); } +IREE_VM_ABI_EXPORT(iree_hal_module_command_buffer_flush_buffer, // + iree_hal_module_state_t, // + rirII, v) { + iree_hal_command_buffer_t* command_buffer = NULL; + IREE_RETURN_IF_ERROR( + iree_hal_command_buffer_check_deref(args->r0, &command_buffer)); + uint32_t target_buffer_slot = (uint32_t)args->i1; + iree_device_size_t target_offset = iree_hal_cast_device_size(args->i3); + iree_device_size_t length = iree_hal_cast_device_size(args->i4); + iree_hal_buffer_ref_t target_ref = iree_hal_make_indirect_buffer_ref( + target_buffer_slot, target_offset, length); + IREE_RETURN_IF_ERROR( + iree_hal_buffer_check_deref_or_null(args->r2, &target_ref.buffer)); + return iree_hal_command_buffer_flush_buffer(command_buffer, target_ref); +} + +IREE_VM_ABI_EXPORT(iree_hal_module_command_buffer_invalidate_buffer, // + iree_hal_module_state_t, // + rirII, v) { + iree_hal_command_buffer_t* command_buffer = NULL; + IREE_RETURN_IF_ERROR( + iree_hal_command_buffer_check_deref(args->r0, &command_buffer)); + uint32_t target_buffer_slot = (uint32_t)args->i1; + iree_device_size_t target_offset = iree_hal_cast_device_size(args->i3); + iree_device_size_t length = iree_hal_cast_device_size(args->i4); + iree_hal_buffer_ref_t target_ref = iree_hal_make_indirect_buffer_ref( + target_buffer_slot, target_offset, length); + IREE_RETURN_IF_ERROR( + iree_hal_buffer_check_deref_or_null(args->r2, &target_ref.buffer)); + return iree_hal_command_buffer_invalidate_buffer(command_buffer, target_ref); +} + IREE_VM_ABI_EXPORT(iree_hal_module_command_buffer_fill_buffer, // iree_hal_module_state_t, // rrIIiIiI, v) {
diff --git a/runtime/src/iree/vm/shims.c b/runtime/src/iree/vm/shims.c index 99d764f..01923f2 100644 --- a/runtime/src/iree/vm/shims.c +++ b/runtime/src/iree/vm/shims.c
@@ -72,6 +72,7 @@ IREE_VM_ABI_DEFINE_SHIM(rriirII, v); IREE_VM_ABI_DEFINE_SHIM(rrIirII, v); IREE_VM_ABI_DEFINE_SHIM(rrIrIIiI, v); +IREE_VM_ABI_DEFINE_SHIM(rirII, v); IREE_VM_ABI_DEFINE_SHIM(riirIrIII, v); IREE_VM_ABI_DEFINE_SHIM(rrIii, v); IREE_VM_ABI_DEFINE_SHIM(rrrIii, v);
diff --git a/runtime/src/iree/vm/shims.h b/runtime/src/iree/vm/shims.h index abaae20..e2fd511 100644 --- a/runtime/src/iree/vm/shims.h +++ b/runtime/src/iree/vm/shims.h
@@ -475,6 +475,14 @@ int64_t i7; }); +IREE_VM_ABI_FIXED_STRUCT(rirII, { + iree_vm_ref_t r0; + int32_t i1; + iree_vm_ref_t r2; + int64_t i3; + int64_t i4; +}); + IREE_VM_ABI_FIXED_STRUCT(riirIrIII, { iree_vm_ref_t r0; int32_t i1; @@ -863,6 +871,7 @@ IREE_VM_ABI_DECLARE_SHIM(rriirII, v); IREE_VM_ABI_DECLARE_SHIM(rrIirII, v); IREE_VM_ABI_DECLARE_SHIM(rrIrIIiI, v); +IREE_VM_ABI_DECLARE_SHIM(rirII, v); IREE_VM_ABI_DECLARE_SHIM(riirIrIII, v); IREE_VM_ABI_DECLARE_SHIM(rrIii, v); IREE_VM_ABI_DECLARE_SHIM(rrrIii, v);