[HAL] Fix inline execution models (initializer inlining and dispatch ABI) (#24827)
The inline execution models were broken end-to-end in both the compiler
and runtime.
On the compiler side, the `HAL_Loader` and `HAL_Inline` dialects did not
support inlining, preventing `CombineInitializers` pass from combining
initializers containing their ops. This broke `inline-dynamic`
compilation generally, and stateful programs using mutable globals under
both `inline-dynamic` and `inline-static`.
On the runtime side, the `hal_loader` dispatch shim incorrectly
accounted for padding in the fixed portion of the VM cconv arguments by
using the C struct size, causing `inline-dynamic` dispatches to fail
with an argument/result signature mismatch.
The runtime also leaked the inline HAL storage buffer wrapper on
release. VM buffer destruction also assumed that every buffer was
created by `iree_vm_buffer_create/clone` as a single aligned allocation
containing both the `iree_vm_buffer_t` handle and its data. Inline HAL
instead embeds the handle inside a wrapper struct that maps an
`iree_hal_buffer_t`, initializing it in place with
`iree_vm_buffer_initialize`. Such buffers must invoke the wrapper’s
custom allocator to release the underlying HAL buffer and free the
wrapper, rather than calling `iree_allocator_free_aligned` on the
embedded handle. Since nothing else in the handle distinguishes the two
cases at destroy time, create/clone now tag their buffers with
`IREE_VM_BUFFER_ACCESS_COALLOCATED` (set internally), and destroy takes
the free-aligned path only for tagged buffers; all others pass the
handle to their allocator.
This PR adds the required inliner interfaces, fixes the dispatch
argument handling, distinguishes co-allocated and in-place-initialized
VM buffers during destruction, and adds compiler/runtime coverage for
the affected execution models and buffer lifetimes.
Assisted-by: Claude Code
---------
Signed-off-by: Pooja Hemashekar <hemashekar@roofline.ai>
diff --git a/compiler/src/iree/compiler/Modules/HAL/Inline/IR/BUILD.bazel b/compiler/src/iree/compiler/Modules/HAL/Inline/IR/BUILD.bazel
index 52bdf36..86e7d74 100644
--- a/compiler/src/iree/compiler/Modules/HAL/Inline/IR/BUILD.bazel
+++ b/compiler/src/iree/compiler/Modules/HAL/Inline/IR/BUILD.bazel
@@ -75,6 +75,7 @@
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:FunctionInterfaces",
"@llvm-project//mlir:IR",
+ "@llvm-project//mlir:InliningUtils",
"@llvm-project//mlir:Parser",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
diff --git a/compiler/src/iree/compiler/Modules/HAL/Inline/IR/HALInlineDialect.cpp b/compiler/src/iree/compiler/Modules/HAL/Inline/IR/HALInlineDialect.cpp
index 461a557..98deb01 100644
--- a/compiler/src/iree/compiler/Modules/HAL/Inline/IR/HALInlineDialect.cpp
+++ b/compiler/src/iree/compiler/Modules/HAL/Inline/IR/HALInlineDialect.cpp
@@ -14,6 +14,7 @@
#include "mlir/IR/DialectImplementation.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/Parser/Parser.h"
+#include "mlir/Transforms/InliningUtils.h"
namespace mlir::iree_compiler::IREE::HAL::Inline {
@@ -41,11 +42,30 @@
}
};
+struct HALInlineInlinerInterface : DialectInlinerInterface {
+ using DialectInlinerInterface::DialectInlinerInterface;
+
+ bool isLegalToInline(Operation *call, Operation *callable,
+ bool wouldBeCloned) const final {
+ return true;
+ }
+
+ bool isLegalToInline(Region *dest, Region *src, bool wouldBeCloned,
+ IRMapping &valueMapping) const final {
+ return true;
+ }
+
+ bool isLegalToInline(Operation *op, Region *dest, bool wouldBeCloned,
+ IRMapping &valueMapping) const final {
+ return true;
+ }
+};
+
} // namespace
HALInlineDialect::HALInlineDialect(MLIRContext *context)
: Dialect(getDialectNamespace(), context, TypeID::get<HALInlineDialect>()) {
- addInterfaces<HALInlineToVMConversionInterface>();
+ addInterfaces<HALInlineToVMConversionInterface, HALInlineInlinerInterface>();
#define GET_OP_LIST
addOperations<
diff --git a/compiler/src/iree/compiler/Modules/HAL/Loader/IR/BUILD.bazel b/compiler/src/iree/compiler/Modules/HAL/Loader/IR/BUILD.bazel
index 0f575de..2135610 100644
--- a/compiler/src/iree/compiler/Modules/HAL/Loader/IR/BUILD.bazel
+++ b/compiler/src/iree/compiler/Modules/HAL/Loader/IR/BUILD.bazel
@@ -75,6 +75,7 @@
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:FunctionInterfaces",
"@llvm-project//mlir:IR",
+ "@llvm-project//mlir:InliningUtils",
"@llvm-project//mlir:Parser",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
diff --git a/compiler/src/iree/compiler/Modules/HAL/Loader/IR/HALLoaderDialect.cpp b/compiler/src/iree/compiler/Modules/HAL/Loader/IR/HALLoaderDialect.cpp
index e34c1e8..067930e 100644
--- a/compiler/src/iree/compiler/Modules/HAL/Loader/IR/HALLoaderDialect.cpp
+++ b/compiler/src/iree/compiler/Modules/HAL/Loader/IR/HALLoaderDialect.cpp
@@ -14,6 +14,7 @@
#include "mlir/IR/DialectImplementation.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/Parser/Parser.h"
+#include "mlir/Transforms/InliningUtils.h"
namespace mlir::iree_compiler::IREE::HAL::Loader {
@@ -41,11 +42,30 @@
}
};
+struct HALLoaderInlinerInterface : DialectInlinerInterface {
+ using DialectInlinerInterface::DialectInlinerInterface;
+
+ bool isLegalToInline(Operation *call, Operation *callable,
+ bool wouldBeCloned) const final {
+ return true;
+ }
+
+ bool isLegalToInline(Region *dest, Region *src, bool wouldBeCloned,
+ IRMapping &valueMapping) const final {
+ return true;
+ }
+
+ bool isLegalToInline(Operation *op, Region *dest, bool wouldBeCloned,
+ IRMapping &valueMapping) const final {
+ return true;
+ }
+};
+
} // namespace
HALLoaderDialect::HALLoaderDialect(MLIRContext *context)
: Dialect(getDialectNamespace(), context, TypeID::get<HALLoaderDialect>()) {
- addInterfaces<HALLoaderToVMConversionInterface>();
+ addInterfaces<HALLoaderToVMConversionInterface, HALLoaderInlinerInterface>();
#define GET_OP_LIST
addOperations<
diff --git a/runtime/src/iree/modules/hal/inline/module.c b/runtime/src/iree/modules/hal/inline/module.c
index 350cf32..eb3aaf0 100644
--- a/runtime/src/iree/modules/hal/inline/module.c
+++ b/runtime/src/iree/modules/hal/inline/module.c
@@ -87,6 +87,9 @@
IREE_RETURN_AND_END_ZONE_IF_ERROR(
z0, iree_allocator_malloc(host_allocator, sizeof(*storage),
(void**)&storage));
+ storage->host_allocator = host_allocator;
+ storage->hal_buffer = hal_buffer;
+ iree_hal_buffer_retain(hal_buffer);
// Map the HAL buffer into host-accessible memory. It almost always is but
// it's possible the buffer we were passed was allocated on a real device that
diff --git a/runtime/src/iree/modules/hal/loader/module.c b/runtime/src/iree/modules/hal/loader/module.c
index 063190d..7b53bdc 100644
--- a/runtime/src/iree/modules/hal/loader/module.c
+++ b/runtime/src/iree/modules/hal/loader/module.c
@@ -341,18 +341,24 @@
void* IREE_RESTRICT module_state) {
// TODO(benvanik): support multiple variadic segments in one call.
// For now we inline what it would do in a very painful way.
+ //
+ // NOTE: sizeof(iree_vm_abi_rIiii_t) may include trailing struct padding
+ // that is not present before the following variadic segment. Size the
+ // fixed prefix through its final field instead.
+ const iree_host_size_t params_size =
+ offsetof(iree_vm_abi_rIiii_t, i4) +
+ IREE_VM_ABI_FIELD_SIZE(iree_vm_abi_rIiii_t, i4);
bool args_ok = true;
if (args_storage.data_length <
- (sizeof(iree_vm_abi_rIiii_t) + sizeof(iree_vm_size_t) +
- sizeof(iree_vm_size_t))) {
+ (params_size + sizeof(iree_vm_size_t) + sizeof(iree_vm_size_t))) {
// Can't fit even with zero lengths.
args_ok = false;
}
- iree_hal_loader_dispatch_args_t args = {
- .params = *(const iree_vm_abi_rIiii_t*)args_storage.data,
- };
+ iree_hal_loader_dispatch_args_t args;
+ memset(&args, 0, sizeof(args));
if (args_ok) {
- const uint8_t* constants_ptr = args_storage.data + sizeof(args.params);
+ memcpy(&args.params, args_storage.data, params_size);
+ const uint8_t* constants_ptr = args_storage.data + params_size;
args.constant_count = *(const iree_vm_size_t*)constants_ptr;
args.constants = (const uint32_t*)(constants_ptr + sizeof(iree_vm_size_t));
const uint8_t* bindings_ptr =
diff --git a/runtime/src/iree/modules/hal/module.c b/runtime/src/iree/modules/hal/module.c
index 1b86137..6c51b13 100644
--- a/runtime/src/iree/modules/hal/module.c
+++ b/runtime/src/iree/modules/hal/module.c
@@ -1100,6 +1100,14 @@
args->constant_count * sizeof(uint32_t)),
bindings, (iree_hal_dispatch_flags_t)args->flags);
}
+// The shim below uses sizeof(iree_vm_abi_rrIiiiI_t) as the offset of the
+// following variadic segment; that is only correct while the struct has no
+// trailing padding as the VM packs cconv arguments unpadded.
+static_assert(sizeof(iree_vm_abi_rrIiiiI_t) ==
+ offsetof(iree_vm_abi_rrIiiiI_t, i6) +
+ IREE_VM_ABI_FIELD_SIZE(iree_vm_abi_rrIiiiI_t, i6),
+ "iree_vm_abi_rrIiiiI_t must have no trailing padding; size "
+ "through the final field instead of using sizeof");
static iree_status_t iree_hal_module_command_buffer_dispatch_shim(
iree_vm_stack_t* IREE_RESTRICT stack, iree_vm_native_function_flags_t flags,
iree_byte_span_t args_storage, iree_byte_span_t rets_storage,
@@ -1222,6 +1230,14 @@
args->constant_count * sizeof(uint32_t)),
bindings, flags);
}
+// The shim below uses sizeof(iree_vm_abi_rrIirII_t) as the offset of the
+// following variadic segment; that is only correct while the struct has no
+// trailing padding as the VM packs cconv arguments unpadded.
+static_assert(sizeof(iree_vm_abi_rrIirII_t) ==
+ offsetof(iree_vm_abi_rrIirII_t, i6) +
+ IREE_VM_ABI_FIELD_SIZE(iree_vm_abi_rrIirII_t, i6),
+ "iree_vm_abi_rrIirII_t must have no trailing padding; size "
+ "through the final field instead of using sizeof");
static iree_status_t iree_hal_module_command_buffer_dispatch_indirect_shim(
iree_vm_stack_t* IREE_RESTRICT stack, iree_vm_native_function_flags_t flags,
iree_byte_span_t args_storage, iree_byte_span_t rets_storage,
diff --git a/runtime/src/iree/vm/buffer.c b/runtime/src/iree/vm/buffer.c
index f903a16..d4e8b6e 100644
--- a/runtime/src/iree/vm/buffer.c
+++ b/runtime/src/iree/vm/buffer.c
@@ -125,7 +125,8 @@
memset(data_ptr, 0, prefix_size - sizeof(*buffer)); // padding
iree_byte_span_t target_span =
iree_make_byte_span(data_ptr + prefix_size, length);
- iree_vm_buffer_initialize(access, target_span, allocator, buffer);
+ iree_vm_buffer_initialize(access | IREE_VM_BUFFER_ACCESS_COALLOCATED,
+ target_span, allocator, buffer);
*out_buffer = buffer;
IREE_TRACE_ZONE_END(z0);
@@ -140,10 +141,16 @@
// Module-owned buffers are embedded in the bytecode module. Reaching zero
// external references releases the owner retain taken by the first retain.
iree_vm_module_release(buffer->storage_module);
- } else {
+ } else if (iree_all_bits_set(buffer->access,
+ IREE_VM_BUFFER_ACCESS_COALLOCATED)) {
// Buffers are stored as [prefix | data]; freeing the prefix is all we need
// to do to free it all.
iree_allocator_free_aligned(buffer->allocator, buffer);
+ } else {
+ // Buffer handle was initialized in-place over storage owned elsewhere;
+ // notify the allocator so any wrapping storage can be released. Note that
+ // it must not be treated as an aligned allocation.
+ iree_allocator_free(buffer->allocator, buffer);
}
IREE_TRACE_ZONE_END(z0);
@@ -263,7 +270,8 @@
memset(data_ptr, 0, prefix_size - sizeof(*buffer)); // padding
iree_byte_span_t target_span =
iree_make_byte_span(data_ptr + prefix_size, length);
- iree_vm_buffer_initialize(access, target_span, allocator, buffer);
+ iree_vm_buffer_initialize(access | IREE_VM_BUFFER_ACCESS_COALLOCATED,
+ target_span, allocator, buffer);
// Copy the data from the source buffer.
memcpy(target_span.data, source_span.data, target_span.data_length);
diff --git a/runtime/src/iree/vm/buffer.h b/runtime/src/iree/vm/buffer.h
index 18fc8d6..88084fa 100644
--- a/runtime/src/iree/vm/buffer.h
+++ b/runtime/src/iree/vm/buffer.h
@@ -34,6 +34,12 @@
IREE_VM_BUFFER_ACCESS_ORIGIN_GUEST = 1u << 2,
// Buffer references external host memory with an unknown lifetime.
IREE_VM_BUFFER_ACCESS_ORIGIN_HOST = 1u << 3,
+
+ // Buffer handle and data are co-allocated in a single aligned allocation
+ // and the handle pointer must be freed with iree_allocator_free_aligned.
+ // Set internally by iree_vm_buffer_create/iree_vm_buffer_clone; buffers
+ // initialized in-place over existing storage must not set this.
+ IREE_VM_BUFFER_ACCESS_COALLOCATED = 1u << 4,
};
typedef uint32_t iree_vm_buffer_access_t;
diff --git a/runtime/src/iree/vm/buffer_test.cc b/runtime/src/iree/vm/buffer_test.cc
index a3801b2..94b7145 100644
--- a/runtime/src/iree/vm/buffer_test.cc
+++ b/runtime/src/iree/vm/buffer_test.cc
@@ -52,6 +52,42 @@
ASSERT_TRUE(did_free);
}
+// Buffers initialized in-place may use their allocator as a destruction
+// callback for wrapping storage and must not take the aligned free path.
+TEST_F(VMBufferTest, InitializedBufferDestroyNotifiesAllocator) {
+ void* freed_ptr = nullptr;
+ iree_allocator_t test_allocator = {
+ /*.self=*/&freed_ptr,
+ /*.ctl=*/
+ +[](void* self, iree_allocator_command_t command, const void* params,
+ void** inout_ptr) {
+ if (command == IREE_ALLOCATOR_COMMAND_FREE) {
+ *(void**)self = *inout_ptr;
+ }
+ return iree_ok_status();
+ },
+ };
+
+ uint32_t data[] = {0, 1, 2, 3};
+ // Mimic the inline HAL wrapper. The zeroed prefix makes an incorrect
+ // iree_allocator_free_aligned call read a null base pointer and skip the
+ // free.
+ struct {
+ void* prefix[2];
+ iree_vm_buffer_t buffer;
+ } storage;
+ memset(&storage, 0, sizeof(storage));
+ iree_vm_buffer_initialize(
+ IREE_VM_BUFFER_ACCESS_MUTABLE | IREE_VM_BUFFER_ACCESS_ORIGIN_HOST,
+ iree_make_byte_span(data, sizeof(data)), test_allocator, &storage.buffer);
+
+ ASSERT_EQ(freed_ptr, nullptr);
+ iree_vm_buffer_release(&storage.buffer);
+ // The destroy path passes the non-NULL buffer handle so wrapper allocators
+ // are notified even when the buffer data is empty.
+ ASSERT_EQ(freed_ptr, &storage.buffer);
+}
+
typedef struct test_module_t {
// VM module interface used for retain/release testing.
iree_vm_module_t interface;
diff --git a/tests/compiler_driver/BUILD.bazel b/tests/compiler_driver/BUILD.bazel
index fafafa8..a463f5d 100644
--- a/tests/compiler_driver/BUILD.bazel
+++ b/tests/compiler_driver/BUILD.bazel
@@ -22,6 +22,7 @@
"executable_benchmarks.mlir",
"hal_executable.mlir",
"inline_dynamic_hal_executable.mlir",
+ "inline_hal_mutable_global.mlir",
"inline_static_hal_executable.mlir",
"precompile.mlir",
"preprocessing_flags.mlir",
diff --git a/tests/compiler_driver/CMakeLists.txt b/tests/compiler_driver/CMakeLists.txt
index 5ebc0fc..df0a623 100644
--- a/tests/compiler_driver/CMakeLists.txt
+++ b/tests/compiler_driver/CMakeLists.txt
@@ -17,6 +17,7 @@
"executable_benchmarks.mlir"
"hal_executable.mlir"
"inline_dynamic_hal_executable.mlir"
+ "inline_hal_mutable_global.mlir"
"inline_static_hal_executable.mlir"
"precompile.mlir"
"preprocessing_flags.mlir"
diff --git a/tests/compiler_driver/inline_dynamic_hal_executable.mlir b/tests/compiler_driver/inline_dynamic_hal_executable.mlir
index c2e7cd0..536ac43 100644
--- a/tests/compiler_driver/inline_dynamic_hal_executable.mlir
+++ b/tests/compiler_driver/inline_dynamic_hal_executable.mlir
@@ -18,3 +18,16 @@
// CHECK-NOT: hal.pipeline_layout
// CHECK-NOT: hal.semaphore
// CHECK-NOT: hal.executable private
+
+// Compiling through the VM lowering requires inlining the executable-loading
+// initializer into the combined initializer.
+// RUN: iree-compile \
+// RUN: --compile-to=vm \
+// RUN: --iree-execution-model=inline-dynamic \
+// RUN: --iree-hal-target-device=local \
+// RUN: --iree-hal-local-target-device-backends=vmvx %s | FileCheck %s --check-prefix=CHECK-VM
+
+// CHECK-VM-LABEL: vm.func private @__init()
+// CHECK-VM: vm.call @hal_loader.executable.query_support
+// CHECK-VM: vm.call @hal_loader.executable.load
+// CHECK-VM: vm.global.store.ref %{{.+}}, @simple_mul_dispatch_0
diff --git a/tests/compiler_driver/inline_hal_mutable_global.mlir b/tests/compiler_driver/inline_hal_mutable_global.mlir
new file mode 100644
index 0000000..40bd4f0
--- /dev/null
+++ b/tests/compiler_driver/inline_hal_mutable_global.mlir
@@ -0,0 +1,36 @@
+// Mutable globals require an initializer that allocates their backing buffer
+// with hal_inline ops; compiling through the VM lowering requires inlining
+// that initializer into the combined initializer.
+
+// RUN: iree-compile \
+// RUN: --compile-to=vm \
+// RUN: --iree-execution-model=inline-dynamic \
+// RUN: --iree-hal-target-device=local \
+// RUN: --iree-hal-local-target-device-backends=vmvx %s | \
+// RUN: FileCheck %s --check-prefix=CHECK-DYNAMIC
+
+// CHECK-DYNAMIC-LABEL: vm.func private @__init()
+// CHECK-DYNAMIC-DAG: vm.call @hal_inline.buffer.allocate
+// CHECK-DYNAMIC-DAG: vm.global.store.ref %{{.+}}, @state
+// CHECK-DYNAMIC-DAG: vm.call @hal_loader.executable.load
+// CHECK-DYNAMIC: vm.global.store.ref %{{.+}}, @update_dispatch_0
+
+// RUN: iree-compile \
+// RUN: --compile-to=vm \
+// RUN: --iree-execution-model=inline-static \
+// RUN: --iree-hal-target-device=local \
+// RUN: --iree-hal-local-target-device-backends=vmvx-inline %s | \
+// RUN: FileCheck %s --check-prefix=CHECK-STATIC
+
+// CHECK-STATIC-LABEL: vm.func private @__init()
+// CHECK-STATIC: vm.call @hal_inline.buffer.allocate
+// CHECK-STATIC: vm.global.store.ref %{{.+}}, @state
+
+util.global private mutable @state = dense<1.0> : tensor<4xf32>
+
+func.func @update(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+ %state = util.global.load @state : tensor<4xf32>
+ %sum = arith.addf %state, %arg0 : tensor<4xf32>
+ util.global.store %sum, @state : tensor<4xf32>
+ return %sum : tensor<4xf32>
+}
diff --git a/tools/test/BUILD.bazel b/tools/test/BUILD.bazel
index d71a773..ed42d7b 100644
--- a/tools/test/BUILD.bazel
+++ b/tools/test/BUILD.bazel
@@ -49,6 +49,7 @@
"iree-replay-tools.txt",
"iree-run-mlir.mlir",
"iree-run-module-expected.mlir",
+ "iree-run-module-inline-dynamic.mlir",
"iree-run-module-inputs.mlir",
"iree-run-module-multi.mlir",
"iree-run-module-outputs.mlir",
diff --git a/tools/test/CMakeLists.txt b/tools/test/CMakeLists.txt
index 04e11b4..effeaf4 100644
--- a/tools/test/CMakeLists.txt
+++ b/tools/test/CMakeLists.txt
@@ -38,6 +38,7 @@
"iree-run-mlir.mlir"
"iree-run-module-expected.mlir"
"iree-run-module-in-place.mlir"
+ "iree-run-module-inline-dynamic.mlir"
"iree-run-module-inputs.mlir"
"iree-run-module-multi.mlir"
"iree-run-module-outputs.mlir"
diff --git a/tools/test/iree-run-module-inline-dynamic.mlir b/tools/test/iree-run-module-inline-dynamic.mlir
new file mode 100644
index 0000000..08b1fa9
--- /dev/null
+++ b/tools/test/iree-run-module-inline-dynamic.mlir
@@ -0,0 +1,13 @@
+// Executes an inline HAL (no device) module via the hal_loader module.
+// RUN: (iree-compile --iree-execution-model=inline-dynamic --iree-hal-target-device=local --iree-hal-local-target-device-backends=vmvx %s | \
+// RUN: iree-run-module --module=- --function=abs --input="2xf32=-2 3") | FileCheck %s
+// RUN: (iree-compile --iree-execution-model=inline-dynamic --iree-hal-target-device=local --iree-hal-local-target-device-backends=llvm-cpu %s | \
+// RUN: iree-run-module --module=- --function=abs --input="2xf32=-2 3") | FileCheck %s
+
+// CHECK-LABEL: EXEC @abs
+func.func @abs(%input : tensor<2xf32>) -> (tensor<2xf32>) {
+ %result = math.absf %input : tensor<2xf32>
+ return %result : tensor<2xf32>
+}
+// CHECK: result[0]: hal.buffer_view
+// CHECK-NEXT: 2xf32=2 3
