Changing hal.device.query.i32 -> hal.device.query.i64.
Now that we use i64 values for things it makes more sense to do this and
then truncate to i32 when needed. This will allow us to return things
like maximum semaphore values, maximum allocation sizes, free memory, etc.
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/ConvertDeviceOps.cpp b/compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/ConvertDeviceOps.cpp
index 9b60477..5bbdbf2 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/ConvertDeviceOps.cpp
+++ b/compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/ConvertDeviceOps.cpp
@@ -12,8 +12,8 @@
namespace mlir {
namespace iree_compiler {
-// Rewrites a hal.device.query of an i1 value to a hal.device.query of i32 with
-// a truncation on the result.
+// Rewrites a hal.device.query of an i1/i16/i32 value to a hal.device.query of
+// i64 with a truncation on the result.
class DeviceQueryIntCastOpConversion
: public OpConversionPattern<IREE::HAL::DeviceQueryOp> {
public:
@@ -26,34 +26,34 @@
ConversionPatternRewriter &rewriter) const override {
// We only deal with in-dialect conversions to i32 in this pattern.
auto targetType = op.value().getType();
- if (targetType.isInteger(32)) return failure();
+ if (targetType.isInteger(64)) return failure();
if (!targetType.isIntOrIndex()) return failure();
- // Query as I32.
+ // Query as I64.
// Note that due to type conversion we need to handle the default logic
- // ourselves instead of allowing the i32 do the same. We could let it handle
+ // ourselves instead of allowing the i64 do the same. We could let it handle
// things but then we are generating more IR that may prevent other
// canonicalizations (a select of i1 to i1 is easier to handle).
auto queryOp = rewriter.create<IREE::HAL::DeviceQueryOp>(
- op.getLoc(), rewriter.getI1Type(), rewriter.getI32Type(),
+ op.getLoc(), rewriter.getI1Type(), rewriter.getI64Type(),
adaptor.device(), op.categoryAttr(), op.keyAttr(), Attribute{});
auto ok = queryOp.ok();
auto value = queryOp.value();
// Truncate or extend based on the target type.
if (targetType.isIndex()) {
- // i32 -> index cast.
+ // i64 -> index cast.
value = rewriter.createOrFold<arith::IndexCastOp>(op.getLoc(), targetType,
value);
} else if (targetType.isa<IntegerType>()) {
- // i32 -> {integer} cast.
+ // i64 -> {integer} cast.
if (targetType.getIntOrFloatBitWidth() <
value.getType().getIntOrFloatBitWidth()) {
- // i32 -> narrowing cast.
+ // i64 -> narrowing cast.
value = rewriter.createOrFold<arith::TruncIOp>(op.getLoc(), targetType,
value);
} else {
- // i32 -> widening cast.
+ // i64 -> widening cast.
value = rewriter.createOrFold<arith::ExtUIOp>(op.getLoc(), targetType,
value);
}
@@ -75,10 +75,10 @@
}
};
-class DeviceQueryI32OpConversion
+class DeviceQueryI64OpConversion
: public OpConversionPattern<IREE::HAL::DeviceQueryOp> {
public:
- DeviceQueryI32OpConversion(MLIRContext *context, SymbolTable &importSymbols,
+ DeviceQueryI64OpConversion(MLIRContext *context, SymbolTable &importSymbols,
TypeConverter &typeConverter, StringRef importName)
: OpConversionPattern(typeConverter, context) {
importOp = importSymbols.lookup<IREE::VM::ImportOp>(importName);
@@ -88,7 +88,7 @@
LogicalResult matchAndRewrite(
IREE::HAL::DeviceQueryOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
- if (!op.value().getType().isInteger(32)) return failure();
+ if (!op.value().getType().isInteger(64)) return failure();
auto results =
rewriteToCall(op, adaptor, importOp, *getTypeConverter(), rewriter);
if (!results.hasValue()) return failure();
@@ -97,9 +97,9 @@
if (op.default_value().hasValue()) {
value = rewriter.createOrFold<arith::SelectOp>(
op.getLoc(), ok, value,
- rewriter.createOrFold<IREE::VM::ConstI32Op>(op.getLoc(),
+ rewriter.createOrFold<IREE::VM::ConstI64Op>(op.getLoc(),
op.default_valueAttr()));
- ok = rewriter.createOrFold<IREE::VM::ConstI32Op>(op.getLoc(), 1);
+ ok = rewriter.createOrFold<IREE::VM::ConstI64Op>(op.getLoc(), 1);
}
rewriter.replaceOp(op, {ok, value});
return success();
@@ -117,8 +117,8 @@
context, importSymbols, typeConverter, "hal.device.allocator");
patterns.insert<DeviceQueryIntCastOpConversion>(context, typeConverter);
- patterns.insert<DeviceQueryI32OpConversion>(
- context, importSymbols, typeConverter, "hal.device.query.i32");
+ patterns.insert<DeviceQueryI64OpConversion>(
+ context, importSymbols, typeConverter, "hal.device.query.i64");
}
} // namespace iree_compiler
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/test/device_ops.mlir b/compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/test/device_ops.mlir
index 3f39bc3..56056db 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/test/device_ops.mlir
+++ b/compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/test/device_ops.mlir
@@ -10,14 +10,42 @@
// -----
+// CHECK-LABEL: @device_query_i64
+// CHECK-SAME: (%[[DEVICE:.+]]: !vm.ref<!hal.device>)
+func.func @device_query_i64(%device: !hal.device) -> (i1, i64) {
+ // CHECK-DAG: %[[NS:.+]] = vm.rodata.inline "_utf8_sys_
+ // CHECK-DAG: %[[KEY:.+]] = vm.rodata.inline "_utf8_foo_
+ // CHECK: %[[RET:.+]]:2 = vm.call @hal.device.query.i64(%[[DEVICE]], %[[NS]], %[[KEY]]) {nosideeffects} : (!vm.ref<!hal.device>, !vm.buffer, !vm.buffer) -> (i32, i64)
+ %ok, %value = hal.device.query<%device : !hal.device> key("sys" :: "foo") : i1, i64
+ // CHECK: return %[[RET]]#0, %[[RET]]#1
+ return %ok, %value : i1, i64
+}
+
+// -----
+
+// CHECK-LABEL: @device_query_i64_default
+// CHECK-SAME: (%[[DEVICE:.+]]: !vm.ref<!hal.device>)
+func.func @device_query_i64_default(%device: !hal.device) -> i64 {
+ // CHECK-DAG: %[[NS:.+]] = vm.rodata.inline "_utf8_sys_
+ // CHECK-DAG: %[[KEY:.+]] = vm.rodata.inline "_utf8_foo_
+ // CHECK: %[[RET:.+]]:2 = vm.call @hal.device.query.i64(%[[DEVICE]], %[[NS]], %[[KEY]]) {nosideeffects} : (!vm.ref<!hal.device>, !vm.buffer, !vm.buffer) -> (i32, i64)
+ %ok, %value = hal.device.query<%device : !hal.device> key("sys" :: "foo") : i1, i64 = 123 : i64
+ // CHECK: %[[OUT:.+]] = vm.select.i64 %[[RET]]#0, %[[RET]]#1, %c123 : i64
+ // CHECK: return %[[OUT]]
+ return %value : i64
+}
+
+// -----
+
// CHECK-LABEL: @device_query_i32
// CHECK-SAME: (%[[DEVICE:.+]]: !vm.ref<!hal.device>)
func.func @device_query_i32(%device: !hal.device) -> (i1, i32) {
// CHECK-DAG: %[[NS:.+]] = vm.rodata.inline "_utf8_sys_
// CHECK-DAG: %[[KEY:.+]] = vm.rodata.inline "_utf8_foo_
- // CHECK: %[[RET:.+]]:2 = vm.call @hal.device.query.i32(%[[DEVICE]], %[[NS]], %[[KEY]]) {nosideeffects} : (!vm.ref<!hal.device>, !vm.buffer, !vm.buffer) -> (i32, i32)
+ // CHECK: %[[RET:.+]]:2 = vm.call @hal.device.query.i64(%[[DEVICE]], %[[NS]], %[[KEY]]) {nosideeffects} : (!vm.ref<!hal.device>, !vm.buffer, !vm.buffer) -> (i32, i64)
+ // CHECK: %[[RET_I32:.+]] = vm.trunc.i64.i32 %[[RET]]#1 : i64 -> i32
%ok, %value = hal.device.query<%device : !hal.device> key("sys" :: "foo") : i1, i32
- // CHECK: return %[[RET]]#0, %[[RET]]#1
+ // CHECK: return %[[RET]]#0, %[[RET_I32]]
return %ok, %value : i1, i32
}
@@ -28,9 +56,10 @@
func.func @device_query_i32_default(%device: !hal.device) -> i32 {
// CHECK-DAG: %[[NS:.+]] = vm.rodata.inline "_utf8_sys_
// CHECK-DAG: %[[KEY:.+]] = vm.rodata.inline "_utf8_foo_
- // CHECK: %[[RET:.+]]:2 = vm.call @hal.device.query.i32(%[[DEVICE]], %[[NS]], %[[KEY]]) {nosideeffects} : (!vm.ref<!hal.device>, !vm.buffer, !vm.buffer) -> (i32, i32)
+ // CHECK: %[[RET:.+]]:2 = vm.call @hal.device.query.i64(%[[DEVICE]], %[[NS]], %[[KEY]]) {nosideeffects} : (!vm.ref<!hal.device>, !vm.buffer, !vm.buffer) -> (i32, i64)
+ // CHECK: %[[RET_I32:.+]] = vm.trunc.i64.i32 %[[RET]]#1 : i64 -> i32
%ok, %value = hal.device.query<%device : !hal.device> key("sys" :: "foo") : i1, i32 = 123 : i32
- // CHECK: %[[OUT:.+]] = vm.select.i32 %[[RET]]#0, %[[RET]]#1, %c123 : i32
+ // CHECK: %[[OUT:.+]] = vm.select.i32 %[[RET]]#0, %[[RET_I32]], %c123 : i32
// CHECK: return %[[OUT]]
return %value : i32
}
@@ -42,9 +71,10 @@
func.func @device_query_i1(%device: !hal.device) -> (i1, i1) {
// CHECK-DAG: %[[NS:.+]] = vm.rodata.inline "_utf8_sys_
// CHECK-DAG: %[[KEY:.+]] = vm.rodata.inline "_utf8_foo_
- // CHECK: %[[RET:.+]]:2 = vm.call @hal.device.query.i32(%[[DEVICE]], %[[NS]], %[[KEY]]) {nosideeffects} : (!vm.ref<!hal.device>, !vm.buffer, !vm.buffer) -> (i32, i32)
+ // CHECK: %[[RET:.+]]:2 = vm.call @hal.device.query.i64(%[[DEVICE]], %[[NS]], %[[KEY]]) {nosideeffects} : (!vm.ref<!hal.device>, !vm.buffer, !vm.buffer) -> (i32, i64)
+ // CHECK: %[[RET_I32:.+]] = vm.trunc.i64.i32 %[[RET]]#1 : i64 -> i32
%ok, %value = hal.device.query<%device : !hal.device> key("sys" :: "foo") : i1, i1
- // CHECK: %[[I1:.+]] = vm.and.i32 %[[RET]]#1, %c1 : i32
+ // CHECK: %[[I1:.+]] = vm.and.i32 %[[RET_I32]], %c1 : i32
// CHECK: return %[[RET]]#0, %[[I1]]
return %ok, %value : i1, i1
}
@@ -56,9 +86,10 @@
func.func @device_query_i1_default(%device: !hal.device) -> i1 {
// CHECK-DAG: %[[NS:.+]] = vm.rodata.inline "_utf8_sys_
// CHECK-DAG: %[[KEY:.+]] = vm.rodata.inline "_utf8_foo_
- // CHECK: %[[RET:.+]]:2 = vm.call @hal.device.query.i32(%[[DEVICE]], %[[NS]], %[[KEY]]) {nosideeffects} : (!vm.ref<!hal.device>, !vm.buffer, !vm.buffer) -> (i32, i32)
+ // CHECK: %[[RET:.+]]:2 = vm.call @hal.device.query.i64(%[[DEVICE]], %[[NS]], %[[KEY]]) {nosideeffects} : (!vm.ref<!hal.device>, !vm.buffer, !vm.buffer) -> (i32, i64)
+ // CHECK: %[[RET_I32:.+]] = vm.trunc.i64.i32 %[[RET]]#1 : i64 -> i32
%ok, %value = hal.device.query<%device : !hal.device> key("sys" :: "foo") : i1, i1 = 1 : i1
- // CHECK: %[[I1:.+]] = vm.and.i32 %[[RET]]#1, %c1 : i32
+ // CHECK: %[[I1:.+]] = vm.and.i32 %[[RET_I32]], %c1 : i32
// CHECK: %[[OUT:.+]] = vm.select.i32 %[[RET]]#0, %[[I1]], %c1
// CHECK: return %[[OUT]]
return %value : i1
diff --git a/compiler/src/iree/compiler/Dialect/HAL/hal.imports.mlir b/compiler/src/iree/compiler/Dialect/HAL/hal.imports.mlir
index e95e1c4..e7fa7d2 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/hal.imports.mlir
+++ b/compiler/src/iree/compiler/Dialect/HAL/hal.imports.mlir
@@ -300,11 +300,11 @@
attributes {nosideeffects}
// Returns a tuple of (ok, value) for the given configuration key.
-vm.import @device.query.i32(
+vm.import @device.query.i64(
%device : !vm.ref<!hal.device>,
%category : !vm.buffer,
%key : !vm.buffer
-) -> (i32, i32)
+) -> (i32, i64)
attributes {nosideeffects}
//===----------------------------------------------------------------------===//
diff --git a/experimental/rocm/rocm_device.c b/experimental/rocm/rocm_device.c
index 12e373c..7ddd495 100644
--- a/experimental/rocm/rocm_device.c
+++ b/experimental/rocm/rocm_device.c
@@ -154,9 +154,9 @@
return device->device_allocator;
}
-static iree_status_t iree_hal_rocm_device_query_i32(
+static iree_status_t iree_hal_rocm_device_query_i64(
iree_hal_device_t* base_device, iree_string_view_t category,
- iree_string_view_t key, int32_t* out_value) {
+ iree_string_view_t key, int64_t* out_value) {
// iree_hal_rocm_device_t* device = iree_hal_rocm_device_cast(base_device);
*out_value = 0;
@@ -310,7 +310,7 @@
.host_allocator = iree_hal_rocm_device_host_allocator,
.device_allocator = iree_hal_rocm_device_allocator,
.trim = iree_hal_rocm_device_trim,
- .query_i32 = iree_hal_rocm_device_query_i32,
+ .query_i64 = iree_hal_rocm_device_query_i64,
.create_command_buffer = iree_hal_rocm_device_create_command_buffer,
.create_descriptor_set = iree_hal_rocm_device_create_descriptor_set,
.create_descriptor_set_layout =
diff --git a/runtime/src/iree/hal/device.c b/runtime/src/iree/hal/device.c
index 5e46953..7cad656 100644
--- a/runtime/src/iree/hal/device.c
+++ b/runtime/src/iree/hal/device.c
@@ -45,9 +45,9 @@
return status;
}
-IREE_API_EXPORT iree_status_t iree_hal_device_query_i32(
+IREE_API_EXPORT iree_status_t iree_hal_device_query_i64(
iree_hal_device_t* device, iree_string_view_t category,
- iree_string_view_t key, int32_t* out_value) {
+ iree_string_view_t key, int64_t* out_value) {
IREE_ASSERT_ARGUMENT(device);
IREE_ASSERT_ARGUMENT(out_value);
@@ -58,7 +58,7 @@
return iree_ok_status();
}
- return _VTABLE_DISPATCH(device, query_i32)(device, category, key, out_value);
+ return _VTABLE_DISPATCH(device, query_i64)(device, category, key, out_value);
}
IREE_API_EXPORT iree_hal_semaphore_compatibility_t
diff --git a/runtime/src/iree/hal/device.h b/runtime/src/iree/hal/device.h
index 3228e5d..c5b53f8 100644
--- a/runtime/src/iree/hal/device.h
+++ b/runtime/src/iree/hal/device.h
@@ -229,11 +229,11 @@
IREE_API_EXPORT
iree_status_t iree_hal_device_trim(iree_hal_device_t* device);
-// Queries a configuration value as an int32_t.
+// Queries a configuration value as an int64_t.
// The |category| and |key| will be provided to the device driver to interpret
// in a device-specific way and if recognized the value will be converted to an
-// int32_t and returned in |out_value|. Fails if the value represented by the
-// key is not convertable (overflows a 32-bit integer, not a number, etc).
+// int64_t and returned in |out_value|. Fails if the value represented by the
+// key is not convertable.
//
// This is roughly equivalent to the `sysconf` linux syscall
// (https://man7.org/linux/man-pages/man3/sysconf.3.html) in that the exact
@@ -248,9 +248,9 @@
//
// Returned values must remain the same for the lifetime of the device as
// callers may cache them to avoid redundant calls.
-IREE_API_EXPORT iree_status_t iree_hal_device_query_i32(
+IREE_API_EXPORT iree_status_t iree_hal_device_query_i64(
iree_hal_device_t* device, iree_string_view_t category,
- iree_string_view_t key, int32_t* out_value);
+ iree_string_view_t key, int64_t* out_value);
// Queries in what ways the given |semaphore| may be used with |device|.
IREE_API_EXPORT iree_hal_semaphore_compatibility_t
@@ -413,10 +413,10 @@
iree_status_t(IREE_API_PTR* trim)(iree_hal_device_t* device);
- iree_status_t(IREE_API_PTR* query_i32)(iree_hal_device_t* device,
+ iree_status_t(IREE_API_PTR* query_i64)(iree_hal_device_t* device,
iree_string_view_t category,
iree_string_view_t key,
- int32_t* out_value);
+ int64_t* out_value);
iree_status_t(IREE_API_PTR* create_command_buffer)(
iree_hal_device_t* device, iree_hal_command_buffer_mode_t mode,
diff --git a/runtime/src/iree/hal/drivers/cuda/cuda_device.c b/runtime/src/iree/hal/drivers/cuda/cuda_device.c
index 91027ab..6fe8f13 100644
--- a/runtime/src/iree/hal/drivers/cuda/cuda_device.c
+++ b/runtime/src/iree/hal/drivers/cuda/cuda_device.c
@@ -208,9 +208,9 @@
return iree_hal_allocator_trim(device->device_allocator);
}
-static iree_status_t iree_hal_cuda_device_query_i32(
+static iree_status_t iree_hal_cuda_device_query_i64(
iree_hal_device_t* base_device, iree_string_view_t category,
- iree_string_view_t key, int32_t* out_value) {
+ iree_string_view_t key, int64_t* out_value) {
// iree_hal_cuda_device_t* device = iree_hal_cuda_device_cast(base_device);
*out_value = 0;
@@ -397,7 +397,7 @@
.host_allocator = iree_hal_cuda_device_host_allocator,
.device_allocator = iree_hal_cuda_device_allocator,
.trim = iree_hal_cuda_device_trim,
- .query_i32 = iree_hal_cuda_device_query_i32,
+ .query_i64 = iree_hal_cuda_device_query_i64,
.create_command_buffer = iree_hal_cuda_device_create_command_buffer,
.create_descriptor_set = iree_hal_cuda_device_create_descriptor_set,
.create_descriptor_set_layout =
diff --git a/runtime/src/iree/hal/drivers/local_sync/sync_device.c b/runtime/src/iree/hal/drivers/local_sync/sync_device.c
index a9ee632..a59f5c7 100644
--- a/runtime/src/iree/hal/drivers/local_sync/sync_device.c
+++ b/runtime/src/iree/hal/drivers/local_sync/sync_device.c
@@ -139,9 +139,9 @@
return iree_hal_allocator_trim(device->device_allocator);
}
-static iree_status_t iree_hal_sync_device_query_i32(
+static iree_status_t iree_hal_sync_device_query_i64(
iree_hal_device_t* base_device, iree_string_view_t category,
- iree_string_view_t key, int32_t* out_value) {
+ iree_string_view_t key, int64_t* out_value) {
iree_hal_sync_device_t* device = iree_hal_sync_device_cast(base_device);
*out_value = 0;
@@ -314,7 +314,7 @@
.host_allocator = iree_hal_sync_device_host_allocator,
.device_allocator = iree_hal_sync_device_allocator,
.trim = iree_hal_sync_device_trim,
- .query_i32 = iree_hal_sync_device_query_i32,
+ .query_i64 = iree_hal_sync_device_query_i64,
.create_command_buffer = iree_hal_sync_device_create_command_buffer,
.create_descriptor_set = iree_hal_sync_device_create_descriptor_set,
.create_descriptor_set_layout =
diff --git a/runtime/src/iree/hal/drivers/local_task/task_device.c b/runtime/src/iree/hal/drivers/local_task/task_device.c
index 717fcf9..261790b 100644
--- a/runtime/src/iree/hal/drivers/local_task/task_device.c
+++ b/runtime/src/iree/hal/drivers/local_task/task_device.c
@@ -187,9 +187,9 @@
return iree_hal_allocator_trim(device->device_allocator);
}
-static iree_status_t iree_hal_task_device_query_i32(
+static iree_status_t iree_hal_task_device_query_i64(
iree_hal_device_t* base_device, iree_string_view_t category,
- iree_string_view_t key, int32_t* out_value) {
+ iree_string_view_t key, int64_t* out_value) {
iree_hal_task_device_t* device = iree_hal_task_device_cast(base_device);
*out_value = 0;
@@ -204,13 +204,13 @@
} else if (iree_string_view_equal(category,
iree_make_cstring_view("hal.device"))) {
if (iree_string_view_equal(key, iree_make_cstring_view("concurrency"))) {
- *out_value = (int32_t)device->queue_count;
+ *out_value = (int64_t)device->queue_count;
return iree_ok_status();
}
} else if (iree_string_view_equal(category,
iree_make_cstring_view("hal.dispatch"))) {
if (iree_string_view_equal(key, iree_make_cstring_view("concurrency"))) {
- *out_value = (int32_t)iree_task_executor_worker_count(device->executor);
+ *out_value = (int64_t)iree_task_executor_worker_count(device->executor);
return iree_ok_status();
}
}
@@ -375,7 +375,7 @@
.host_allocator = iree_hal_task_device_host_allocator,
.device_allocator = iree_hal_task_device_allocator,
.trim = iree_hal_task_device_trim,
- .query_i32 = iree_hal_task_device_query_i32,
+ .query_i64 = iree_hal_task_device_query_i64,
.create_command_buffer = iree_hal_task_device_create_command_buffer,
.create_descriptor_set = iree_hal_task_device_create_descriptor_set,
.create_descriptor_set_layout =
diff --git a/runtime/src/iree/hal/drivers/vulkan/vulkan_device.cc b/runtime/src/iree/hal/drivers/vulkan/vulkan_device.cc
index eae283d..af0eb41 100644
--- a/runtime/src/iree/hal/drivers/vulkan/vulkan_device.cc
+++ b/runtime/src/iree/hal/drivers/vulkan/vulkan_device.cc
@@ -919,9 +919,9 @@
return iree_hal_allocator_trim(device->device_allocator);
}
-static iree_status_t iree_hal_vulkan_device_query_i32(
+static iree_status_t iree_hal_vulkan_device_query_i64(
iree_hal_device_t* base_device, iree_string_view_t category,
- iree_string_view_t key, int32_t* out_value) {
+ iree_string_view_t key, int64_t* out_value) {
// iree_hal_vulkan_device_t* device =
// iree_hal_vulkan_device_cast(base_device);
*out_value = 0;
@@ -1130,7 +1130,7 @@
/*.host_allocator=*/iree_hal_vulkan_device_host_allocator,
/*.device_allocator=*/iree_hal_vulkan_device_allocator,
/*.trim=*/iree_hal_vulkan_device_trim,
- /*.query_i32=*/iree_hal_vulkan_device_query_i32,
+ /*.query_i64=*/iree_hal_vulkan_device_query_i64,
/*.create_command_buffer=*/iree_hal_vulkan_device_create_command_buffer,
/*.create_descriptor_set=*/iree_hal_vulkan_device_create_descriptor_set,
/*.create_descriptor_set_layout=*/
diff --git a/runtime/src/iree/modules/hal/exports.inl b/runtime/src/iree/modules/hal/exports.inl
index 96208f6..101dc82 100644
--- a/runtime/src/iree/modules/hal/exports.inl
+++ b/runtime/src/iree/modules/hal/exports.inl
@@ -61,7 +61,7 @@
EXPORT_FN("descriptor_set_layout.create", iree_hal_module_descriptor_set_layout_create, riCiiD, r)
EXPORT_FN("device.allocator", iree_hal_module_device_allocator, r, r)
-EXPORT_FN("device.query.i32", iree_hal_module_device_query_i32, rrr, ii)
+EXPORT_FN("device.query.i64", iree_hal_module_device_query_i64, rrr, iI)
EXPORT_FN("ex.shared_device", iree_hal_module_ex_shared_device, v, r)
EXPORT_FN("ex.submit_and_wait", iree_hal_module_ex_submit_and_wait, rr, v)
diff --git a/runtime/src/iree/modules/hal/module.c b/runtime/src/iree/modules/hal/module.c
index f6c7387..268f371 100644
--- a/runtime/src/iree/modules/hal/module.c
+++ b/runtime/src/iree/modules/hal/module.c
@@ -1184,9 +1184,9 @@
return iree_ok_status();
}
-IREE_VM_ABI_EXPORT(iree_hal_module_device_query_i32, //
+IREE_VM_ABI_EXPORT(iree_hal_module_device_query_i64, //
iree_hal_module_state_t, //
- rrr, ii) {
+ rrr, iI) {
iree_hal_device_t* device = NULL;
IREE_RETURN_IF_ERROR(iree_hal_device_check_deref(args->r0, &device));
iree_vm_buffer_t* category = NULL;
@@ -1196,11 +1196,11 @@
IREE_RETURN_IF_ERROR(iree_vm_buffer_check_deref(args->r2, &key));
iree_string_view_t key_str = iree_vm_buffer_as_string(key);
- int32_t value = 0;
+ int64_t value = 0;
iree_status_t query_status =
- iree_hal_device_query_i32(device, category_str, key_str, &value);
+ iree_hal_device_query_i64(device, category_str, key_str, &value);
rets->i0 = iree_status_consume_code(query_status) == IREE_STATUS_OK ? 1 : 0;
- rets->i1 = (int32_t)value;
+ rets->i1 = value;
return iree_ok_status();
}
diff --git a/runtime/src/iree/vm/shims.c b/runtime/src/iree/vm/shims.c
index 86ac8e8..da737ba 100644
--- a/runtime/src/iree/vm/shims.c
+++ b/runtime/src/iree/vm/shims.c
@@ -45,7 +45,7 @@
IREE_VM_ABI_DEFINE_SHIM(rr, r);
IREE_VM_ABI_DEFINE_SHIM(rr, v);
IREE_VM_ABI_DEFINE_SHIM(rr, ii);
-IREE_VM_ABI_DEFINE_SHIM(rrr, ii);
+IREE_VM_ABI_DEFINE_SHIM(rrr, iI);
IREE_VM_ABI_DEFINE_SHIM(rrCirIID, r);
IREE_VM_ABI_DEFINE_SHIM(rriCiD, v);
IREE_VM_ABI_DEFINE_SHIM(rriiCID, v);
diff --git a/runtime/src/iree/vm/shims.h b/runtime/src/iree/vm/shims.h
index 1b9011e..00d880c 100644
--- a/runtime/src/iree/vm/shims.h
+++ b/runtime/src/iree/vm/shims.h
@@ -496,7 +496,7 @@
IREE_VM_ABI_DECLARE_SHIM(rr, r);
IREE_VM_ABI_DECLARE_SHIM(rr, v);
IREE_VM_ABI_DECLARE_SHIM(rr, ii);
-IREE_VM_ABI_DECLARE_SHIM(rrr, ii);
+IREE_VM_ABI_DECLARE_SHIM(rrr, iI);
IREE_VM_ABI_DECLARE_SHIM(rrCirIID, r);
IREE_VM_ABI_DECLARE_SHIM(rriCiD, v);
IREE_VM_ABI_DECLARE_SHIM(rriiCID, v);
