Adding queue affinity arg to `hal.command_buffer.create`.
The C API has taken this for awhile as a hint to implementations as to
which queues a command buffer may be executed on. It's legal for this to
always be "any" but we may want to scope things more tightly in the
future to e.g. replicate command buffers on multiple logical devices
represented as queues or manage NUMA hinting.
diff --git a/compiler/src/iree/compiler/Codegen/SPIRV/test/link_executables.mlir b/compiler/src/iree/compiler/Codegen/SPIRV/test/link_executables.mlir
index cda0bbb..7d2977e 100644
--- a/compiler/src/iree/compiler/Codegen/SPIRV/test/link_executables.mlir
+++ b/compiler/src/iree/compiler/Codegen/SPIRV/test/link_executables.mlir
@@ -81,7 +81,8 @@
} {
%c0 = arith.constant 0 : index
%device = hal.devices.get %c0 : !hal.device
- %cmd = hal.command_buffer.create device(%device : !hal.device) mode("OneShot") categories("Transfer|Dispatch") : !hal.command_buffer attributes {
+ %affinity = arith.constant -1 : i64
+ %cmd = hal.command_buffer.create device(%device : !hal.device) mode("OneShot") categories("Transfer|Dispatch") affinity(%affinity) : !hal.command_buffer attributes {
testing.op.a = @dispatch_0,
testing.op.b = @dispatch_0::@spirv,
testing.op.c = @dispatch_0::@spirv::@dispatch_0
@@ -101,7 +102,8 @@
util.initializer {
%c0 = arith.constant 0 : index
%device = hal.devices.get %c0 : !hal.device
- %cmd = hal.command_buffer.create device(%device : !hal.device) mode("OneShot") categories("Transfer|Dispatch") : !hal.command_buffer
+ %affinity = arith.constant -1 : i64
+ %cmd = hal.command_buffer.create device(%device : !hal.device) mode("OneShot") categories("Transfer|Dispatch") affinity(%affinity) : !hal.command_buffer
%c1 = arith.constant 1 : index
%dispatch_0_exe = hal.executable.lookup device(%device : !hal.device) executable(@dispatch_0) : !hal.executable
%dispatch_1_exe = hal.executable.lookup device(%device : !hal.device) executable(@dispatch_1) : !hal.executable
@@ -291,7 +293,8 @@
func.func @two_target_environments() -> () {
%c0 = arith.constant 0 : index
%device = hal.devices.get %c0 : !hal.device
- %cmd = hal.command_buffer.create device(%device : !hal.device) mode("OneShot") categories("Transfer|Dispatch") : !hal.command_buffer
+ %affinity = arith.constant -1 : i64
+ %cmd = hal.command_buffer.create device(%device : !hal.device) mode("OneShot") categories("Transfer|Dispatch") affinity(%affinity) : !hal.command_buffer
%c1 = arith.constant 1 : index
%dispatch_0_exe = hal.executable.lookup device(%device : !hal.device) executable(@dispatch_0) : !hal.executable
%dispatch_1_exe = hal.executable.lookup device(%device : !hal.device) executable(@dispatch_1) : !hal.executable
diff --git a/compiler/src/iree/compiler/Codegen/VMVX/test/link_executables.mlir b/compiler/src/iree/compiler/Codegen/VMVX/test/link_executables.mlir
index d6d6d15..af83f1b 100644
--- a/compiler/src/iree/compiler/Codegen/VMVX/test/link_executables.mlir
+++ b/compiler/src/iree/compiler/Codegen/VMVX/test/link_executables.mlir
@@ -74,7 +74,8 @@
} {
%c0 = arith.constant 0 : index
%device = hal.devices.get %c0 : !hal.device
- %cmd = hal.command_buffer.create device(%device : !hal.device) mode("OneShot") categories("Transfer|Dispatch") : !hal.command_buffer attributes {
+ %affinity = arith.constant -1 : i64
+ %cmd = hal.command_buffer.create device(%device : !hal.device) mode("OneShot") categories("Transfer|Dispatch") affinity(%affinity) : !hal.command_buffer attributes {
testing.op.a = @dispatch_0,
testing.op.b = @dispatch_0::@vmvx,
testing.op.c = @dispatch_0::@vmvx::@dispatch_0
@@ -94,7 +95,8 @@
util.initializer {
%c0 = arith.constant 0 : index
%device = hal.devices.get %c0 : !hal.device
- %cmd = hal.command_buffer.create device(%device : !hal.device) mode("OneShot") categories("Transfer|Dispatch") : !hal.command_buffer
+ %affinity = arith.constant -1 : i64
+ %cmd = hal.command_buffer.create device(%device : !hal.device) mode("OneShot") categories("Transfer|Dispatch") affinity(%affinity) : !hal.command_buffer
%c1 = arith.constant 1 : index
%dispatch_0_exe = hal.executable.lookup device(%device : !hal.device) executable(@dispatch_0) : !hal.executable
%dispatch_1_exe = hal.executable.lookup device(%device : !hal.device) executable(@dispatch_1) : !hal.executable
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 6071614..8299939 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/ConvertCommandBufferOps.cpp
+++ b/compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/ConvertCommandBufferOps.cpp
@@ -51,6 +51,7 @@
if (!categoriesValue.has_value())
return failure();
callOperands.append(categoriesValue.value());
+ callOperands.push_back(adaptor.getQueueAffinity());
if (adaptor.getBindingCapacity()) {
callOperands.push_back(castToImportType(adaptor.getBindingCapacity(),
rewriter.getI32Type(), rewriter));
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 9a3e26d..7402f62 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
@@ -1,18 +1,20 @@
// RUN: iree-opt --split-input-file --iree-vm-conversion --canonicalize --iree-vm-target-index-bits=32 %s | FileCheck %s
// CHECK-LABEL: @command_buffer_create
-util.func public @command_buffer_create(%arg0: !hal.device) {
- // CHECK: %ref = vm.call @hal.command_buffer.create(%arg0, %c1, %c3, %zero) : (!vm.ref<!hal.device>, i32, i32, i32) -> !vm.ref<!hal.command_buffer>
- %cmd = hal.command_buffer.create device(%arg0 : !hal.device) mode("OneShot") categories("Transfer|Dispatch") : !hal.command_buffer
+// CHECK-SAME: (%[[DEVICE:.+]]: !vm.ref<!hal.device>, %[[AFFINITY:.+]]: i64)
+util.func public @command_buffer_create(%device: !hal.device, %affinity: i64) {
+ // CHECK: = vm.call @hal.command_buffer.create(%[[DEVICE]], %c1, %c3, %[[AFFINITY]], %zero) : (!vm.ref<!hal.device>, i32, i32, i64, i32) -> !vm.ref<!hal.command_buffer>
+ %cmd = hal.command_buffer.create device(%device : !hal.device) mode("OneShot") categories("Transfer|Dispatch") affinity(%affinity) : !hal.command_buffer
util.return
}
// -----
// CHECK-LABEL: @command_buffer_create_bindings
-util.func public @command_buffer_create_bindings(%arg0: !hal.device, %arg1: index) {
- // CHECK: %ref = vm.call @hal.command_buffer.create(%arg0, %c1, %c3, %arg1) : (!vm.ref<!hal.device>, i32, i32, i32) -> !vm.ref<!hal.command_buffer>
- %cmd = hal.command_buffer.create device(%arg0 : !hal.device) mode("OneShot") categories("Transfer|Dispatch") bindings(%arg1) : !hal.command_buffer
+// CHECK-SAME: (%[[DEVICE:.+]]: !vm.ref<!hal.device>, %[[AFFINITY:.+]]: i64, %[[CAPACITY:.+]]: i32)
+util.func public @command_buffer_create_bindings(%device: !hal.device, %affinity: i64, %capacity: index) {
+ // CHECK: = vm.call @hal.command_buffer.create(%[[DEVICE]], %c1, %c3, %[[AFFINITY]], %[[CAPACITY]]) : (!vm.ref<!hal.device>, i32, i32, i64, i32) -> !vm.ref<!hal.command_buffer>
+ %cmd = hal.command_buffer.create device(%device : !hal.device) mode("OneShot") categories("Transfer|Dispatch") affinity(%affinity) bindings(%capacity) : !hal.command_buffer
util.return
}
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Patterns.cpp b/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Patterns.cpp
index 771e7d4..274b348 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Patterns.cpp
+++ b/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Patterns.cpp
@@ -939,7 +939,8 @@
rewriter
.create<IREE::HAL::CommandBufferCreateOp>(
loc, rewriter.getType<IREE::HAL::CommandBufferType>(), device,
- modes, commandCategories, /*binding_capacity=*/Value{})
+ modes, commandCategories, queueAffinity,
+ /*binding_capacity=*/Value{})
.getResult();
mapping->mapCommandBuffer(executeOp, commandBuffer);
diff --git a/compiler/src/iree/compiler/Dialect/HAL/IR/HALOps.td b/compiler/src/iree/compiler/Dialect/HAL/IR/HALOps.td
index 74d0b56..68a4c5f 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/IR/HALOps.td
+++ b/compiler/src/iree/compiler/Dialect/HAL/IR/HALOps.td
@@ -1163,6 +1163,7 @@
HAL_Device:$device,
HAL_CommandBufferModeBitfieldAttr:$modes,
HAL_CommandCategoryBitfieldAttr:$command_categories,
+ HAL_DeviceQueueAffinity:$queue_affinity,
Optional<Index>:$binding_capacity
);
let results = (outs
@@ -1173,6 +1174,7 @@
`device` `(` $device `:` type($device) `)`
`mode` `(` $modes `)`
`categories` `(` $command_categories `)`
+ `affinity` `(` $queue_affinity `)`
(`bindings` `(` $binding_capacity^ `)`)?
`:` type($result)
attr-dict-with-keyword
diff --git a/compiler/src/iree/compiler/Dialect/HAL/IR/test/command_buffer_folding.mlir b/compiler/src/iree/compiler/Dialect/HAL/IR/test/command_buffer_folding.mlir
index a61ea4f..3adbce8 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/IR/test/command_buffer_folding.mlir
+++ b/compiler/src/iree/compiler/Dialect/HAL/IR/test/command_buffer_folding.mlir
@@ -1,11 +1,12 @@
// RUN: iree-opt --split-input-file --canonicalize %s | iree-opt --split-input-file | FileCheck %s
// CHECK-LABEL: @skip_command_buffer_device
-// CHECK-SAME: (%[[DEVICE:.+]]: !hal.device)
-util.func public @skip_command_buffer_device(%device: !hal.device) -> !hal.executable {
+// CHECK-SAME: (%[[DEVICE:.+]]: !hal.device, %[[AFFINITY:.+]]: i64)
+util.func public @skip_command_buffer_device(%device: !hal.device, %affinity: i64) -> !hal.executable {
%cmd = hal.command_buffer.create device(%device : !hal.device)
mode(OneShot)
- categories("Transfer|Dispatch") : !hal.command_buffer
+ categories("Transfer|Dispatch")
+ affinity(%affinity) : !hal.command_buffer
// CHECK-NOT: hal.command_buffer.device
// CHECK: = hal.executable.lookup device(%[[DEVICE]] : !hal.device)
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 dc16d45..c3348d9 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
@@ -1,15 +1,17 @@
// RUN: iree-opt --split-input-file %s | FileCheck %s
// CHECK-LABEL: @command_buffer_create
-// CHECK-SAME: (%[[DEVICE:.+]]: !hal.device)
-util.func public @command_buffer_create(%device: !hal.device) {
+// CHECK-SAME: (%[[DEVICE:.+]]: !hal.device, %[[AFFINITY:.+]]: i64)
+util.func public @command_buffer_create(%device: !hal.device, %affinity: i64) {
// CHECK: %cmd = hal.command_buffer.create
// CHECK-SAME: device(%[[DEVICE]] : !hal.device)
// CHECK-SAME: mode(OneShot)
- // CHECK-SAME: categories("Transfer|Dispatch") : !hal.command_buffer
+ // CHECK-SAME: categories("Transfer|Dispatch")
+ // CHECK-SAME: affinity(%[[AFFINITY]]) : !hal.command_buffer
%cmd = hal.command_buffer.create device(%device : !hal.device)
mode(OneShot)
- categories("Transfer|Dispatch") : !hal.command_buffer
+ categories("Transfer|Dispatch")
+ affinity(%affinity) : !hal.command_buffer
util.return
}
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/DumpExecutableBenchmarks.cpp b/compiler/src/iree/compiler/Dialect/HAL/Transforms/DumpExecutableBenchmarks.cpp
index 588f560..4ec14bf 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/DumpExecutableBenchmarks.cpp
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/DumpExecutableBenchmarks.cpp
@@ -256,6 +256,7 @@
// TODO(multi-device): support multiple devices in benchmark generation.
// For now we should just use the affinityAttr to resolve the device.
Value device = IREE::HAL::DeviceType::resolveAny(loc, funcBuilder);
+ Value queueAffinity = funcBuilder.create<arith::ConstantIntOp>(loc, -1, 64);
// Create and begin command buffer.
// TODO(benvanik): reuse the command buffer (initialize once and store).
@@ -267,6 +268,7 @@
.create<IREE::HAL::CommandBufferCreateOp>(
loc, funcBuilder.getType<IREE::HAL::CommandBufferType>(), device,
commandBufferModes, IREE::HAL::CommandCategoryBitfield::Dispatch,
+ queueAffinity,
/*binding_capacity=*/Value{})
.getResult();
@@ -379,7 +381,6 @@
IREE::HAL::FenceFlagBitfield::None);
// Queue execution.
- auto queueAffinity = funcBuilder.create<arith::ConstantIntOp>(loc, -1, 64);
funcBuilder.create<IREE::HAL::DeviceQueueExecuteOp>(
loc, device, queueAffinity, waitFence, signalFence,
ValueRange{commandBuffer});
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/convert_to_hal.mlir b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/convert_to_hal.mlir
index 3bf0a7b..d504c5e 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/convert_to_hal.mlir
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/convert_to_hal.mlir
@@ -109,7 +109,7 @@
// CHECK: %[[CMD:.+]] = hal.command_buffer.create
// CHECK-SAME: device(%[[DEVICE]] : !hal.device)
// CHECK-SAME: mode("OneShot|AllowInlineExecution")
- // CHECK-SAME: categories("Transfer|Dispatch") : !hal.command_buffer
+ // CHECK-SAME: categories("Transfer|Dispatch")
%timepoint = stream.cmd.execute
with(%arg0_resource as %arg0_capture: !stream.resource<external>{%c16},
%arg1_resource as %arg1_capture: !stream.resource<external>{%c16},
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/fixup_legacy_sync.mlir b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/fixup_legacy_sync.mlir
index f47bcd2..29de091 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/fixup_legacy_sync.mlir
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/fixup_legacy_sync.mlir
@@ -5,9 +5,9 @@
module attributes {hal.device.targets = [#hal.device.target<"vulkan", {legacy_sync}>]} {
// CHECK-LABEL: @command_buffer_reusable
-util.func public @command_buffer_reusable(%arg0: !hal.device) {
- // CHECK: hal.command_buffer.create device(%arg0 : !hal.device) mode("None")
- %cmd = hal.command_buffer.create device(%arg0 : !hal.device) mode("None") categories("Transfer|Dispatch") : !hal.command_buffer
+util.func public @command_buffer_reusable(%device: !hal.device, %affinity: i64) {
+ // CHECK: hal.command_buffer.create device(%{{.+}} : !hal.device) mode("None")
+ %cmd = hal.command_buffer.create device(%device : !hal.device) mode("None") categories("Transfer|Dispatch") affinity(%affinity) : !hal.command_buffer
util.return
}
} // module
@@ -18,9 +18,9 @@
module attributes {hal.device.targets = [#hal.device.target<"vulkan", {legacy_sync}>]} {
// CHECK-LABEL: @command_buffer_oneshot
-util.func public @command_buffer_oneshot(%arg0: !hal.device) {
- // CHECK: hal.command_buffer.create device(%arg0 : !hal.device) mode("OneShot|AllowInlineExecution")
- %cmd = hal.command_buffer.create device(%arg0 : !hal.device) mode(OneShot) categories("Transfer|Dispatch") : !hal.command_buffer
+util.func public @command_buffer_oneshot(%device: !hal.device, %affinity: i64) {
+ // CHECK: hal.command_buffer.create device(%{{.+}} : !hal.device) mode("OneShot|AllowInlineExecution")
+ %cmd = hal.command_buffer.create device(%device : !hal.device) mode(OneShot) categories("Transfer|Dispatch") affinity(%affinity) : !hal.command_buffer
util.return
}
} // module
@@ -34,9 +34,9 @@
#hal.device.target<"vulkan", {}>
]} {
// CHECK-LABEL: @legacy_mode_not_required
-util.func public @legacy_mode_not_required(%arg0: !hal.device) {
- // CHECK: hal.command_buffer.create device(%arg0 : !hal.device) mode(OneShot)
- %cmd = hal.command_buffer.create device(%arg0 : !hal.device) mode(OneShot) categories("Transfer|Dispatch") : !hal.command_buffer
+util.func public @legacy_mode_not_required(%device: !hal.device, %affinity: i64) {
+ // CHECK: hal.command_buffer.create device(%{{.+}} : !hal.device) mode(OneShot)
+ %cmd = hal.command_buffer.create device(%device : !hal.device) mode(OneShot) categories("Transfer|Dispatch") affinity(%affinity) : !hal.command_buffer
util.return
}
} // module
@@ -51,7 +51,7 @@
]} {
// CHECK-LABEL: @mixed_legacy_mode_required
util.func public @mixed_legacy_mode_required(%device: !hal.device, %wait: !hal.fence, %cmd: !hal.command_buffer, %signal: !hal.fence) {
- %affinity = arith.constant 0 : i64
+ %affinity = arith.constant 1 : i64
// CHECK: hal.fence.await
// CHECK: hal.device.queue.execute
// CHECK: hal.fence.await
@@ -71,7 +71,7 @@
// CHECK-LABEL: @blocking_execute
// CHECK-SAME: (%[[DEVICE:.+]]: !hal.device, %[[WAIT:.+]]: !hal.fence, %[[CMD:.+]]: !hal.command_buffer, %[[SIGNAL:.+]]: !hal.fence)
util.func public @blocking_execute(%device: !hal.device, %wait: !hal.fence, %cmd: !hal.command_buffer, %signal: !hal.fence) {
- %affinity = arith.constant 0 : i64
+ %affinity = arith.constant 1 : i64
// CHECK-DAG: %[[NULL:.+]] = util.null : !hal.fence
// CHECK-DAG: hal.fence.await until([%[[WAIT]]])
// CHECK-NEXT: hal.device.queue.execute<%[[DEVICE]] : !hal.device>
diff --git a/compiler/src/iree/compiler/Dialect/HAL/hal.imports.mlir b/compiler/src/iree/compiler/Dialect/HAL/hal.imports.mlir
index 6731942..de1ed18 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/hal.imports.mlir
+++ b/compiler/src/iree/compiler/Dialect/HAL/hal.imports.mlir
@@ -199,6 +199,7 @@
%device : !vm.ref<!hal.device>,
%modes : i32,
%command_categories : i32,
+ %queue_affinity : i64,
%binding_capacity : i32
) -> !vm.ref<!hal.command_buffer>
diff --git a/runtime/src/iree/modules/hal/exports.inl b/runtime/src/iree/modules/hal/exports.inl
index b808785..bd5adee 100644
--- a/runtime/src/iree/modules/hal/exports.inl
+++ b/runtime/src/iree/modules/hal/exports.inl
@@ -49,7 +49,7 @@
EXPORT_FN("command_buffer.begin_debug_group", iree_hal_module_command_buffer_begin_debug_group, rr, v)
EXPORT_FN("command_buffer.collective", iree_hal_module_command_buffer_collective, rriirIIrIII, v)
EXPORT_FN("command_buffer.copy_buffer", iree_hal_module_command_buffer_copy_buffer, rrIrII, v)
-EXPORT_FN("command_buffer.create", iree_hal_module_command_buffer_create, riii, r)
+EXPORT_FN("command_buffer.create", iree_hal_module_command_buffer_create, riiIi, r)
EXPORT_FN("command_buffer.dispatch", iree_hal_module_command_buffer_dispatch, rriiii, v)
EXPORT_FN("command_buffer.dispatch.indirect", iree_hal_module_command_buffer_dispatch_indirect, rrirI, v)
EXPORT_FN("command_buffer.end_debug_group", iree_hal_module_command_buffer_end_debug_group, r, v)
diff --git a/runtime/src/iree/modules/hal/module.c b/runtime/src/iree/modules/hal/module.c
index 4b84671..52b3fe3 100644
--- a/runtime/src/iree/modules/hal/module.c
+++ b/runtime/src/iree/modules/hal/module.c
@@ -670,19 +670,21 @@
IREE_VM_ABI_EXPORT(iree_hal_module_command_buffer_create, //
iree_hal_module_state_t, //
- riii, r) {
+ riiIi, r) {
iree_hal_device_t* device = NULL;
IREE_RETURN_IF_ERROR(iree_hal_device_check_deref(args->r0, &device));
iree_hal_command_buffer_mode_t modes =
(iree_hal_command_buffer_mode_t)args->i1;
iree_hal_command_category_t command_categories =
(iree_hal_command_category_t)args->i2;
- iree_host_size_t binding_capacity = (iree_host_size_t)args->i3;
+ iree_hal_queue_affinity_t queue_affinity =
+ (iree_hal_queue_affinity_t)args->i3;
+ iree_host_size_t binding_capacity = (iree_host_size_t)args->i4;
iree_hal_command_buffer_t* command_buffer = NULL;
IREE_RETURN_IF_ERROR(iree_hal_command_buffer_create(
- device, modes, command_categories, IREE_HAL_QUEUE_AFFINITY_ANY,
- binding_capacity, &command_buffer));
+ device, modes, command_categories, queue_affinity, binding_capacity,
+ &command_buffer));
iree_status_t status = iree_hal_command_buffer_begin(command_buffer);
if (iree_status_is_ok(status)) {
diff --git a/runtime/src/iree/vm/shims.c b/runtime/src/iree/vm/shims.c
index 9c8b68e..1d5744d 100644
--- a/runtime/src/iree/vm/shims.c
+++ b/runtime/src/iree/vm/shims.c
@@ -42,6 +42,7 @@
IREE_VM_ABI_DEFINE_SHIM(riii, r);
IREE_VM_ABI_DEFINE_SHIM(riiI, r);
IREE_VM_ABI_DEFINE_SHIM(riii, v);
+IREE_VM_ABI_DEFINE_SHIM(riiIi, r);
IREE_VM_ABI_DEFINE_SHIM(rIiiI, r);
IREE_VM_ABI_DEFINE_SHIM(riIiirII, r);
IREE_VM_ABI_DEFINE_SHIM(rriirIIrIII, v);
diff --git a/runtime/src/iree/vm/shims.h b/runtime/src/iree/vm/shims.h
index ab2f012..195daf8 100644
--- a/runtime/src/iree/vm/shims.h
+++ b/runtime/src/iree/vm/shims.h
@@ -313,6 +313,14 @@
int32_t i4;
});
+IREE_VM_ABI_FIXED_STRUCT(riiIi, {
+ iree_vm_ref_t r0;
+ int32_t i1;
+ int32_t i2;
+ int64_t i3;
+ int32_t i4;
+});
+
IREE_VM_ABI_FIXED_STRUCT(riiI, {
iree_vm_ref_t r0;
int32_t i1;
@@ -659,6 +667,7 @@
IREE_VM_ABI_DECLARE_SHIM(riii, r);
IREE_VM_ABI_DECLARE_SHIM(riiI, r);
IREE_VM_ABI_DECLARE_SHIM(riii, v);
+IREE_VM_ABI_DECLARE_SHIM(riiIi, r);
IREE_VM_ABI_DECLARE_SHIM(rIiiI, r);
IREE_VM_ABI_DECLARE_SHIM(riIiirII, r);
IREE_VM_ABI_DECLARE_SHIM(rriirIIrIII, v);
