Making iree_vm_list_resize work and switching ref cconv behavior (#4858)
The new C HAL module was the first user of this function - now there
are tests. The new shims also assumed that arguments were not retained
(as doing so requires cleanup) but the dispatch code was double-retaining
them. I'm surprised things worked as well as they did like that o_o
Fixes #4852.
diff --git a/iree/base/CMakeLists.txt b/iree/base/CMakeLists.txt
index a01db40..048bf22 100644
--- a/iree/base/CMakeLists.txt
+++ b/iree/base/CMakeLists.txt
@@ -266,8 +266,7 @@
DEPS
::core_headers
DEFINES
- # TODO(#2114): Change the mode to 2.
- "IREE_TRACING_MODE=1"
+ "IREE_TRACING_MODE=2"
PUBLIC
)
else()
diff --git a/iree/hal/local/arena.c b/iree/hal/local/arena.c
index 7b9529a..dabc746 100644
--- a/iree/hal/local/arena.c
+++ b/iree/hal/local/arena.c
@@ -52,6 +52,8 @@
iree_status_t iree_arena_block_pool_acquire(iree_arena_block_pool_t* block_pool,
iree_arena_block_t** out_block) {
+ IREE_TRACE_ZONE_BEGIN(z0);
+
iree_arena_block_t* block =
iree_atomic_arena_block_slist_pop(&block_pool->available_slist);
@@ -63,23 +65,28 @@
// that's fine - it's just one block and the contention means there's likely
// to be a need for more anyway.
uint8_t* block_base = NULL;
- IREE_RETURN_IF_ERROR(iree_allocator_malloc(block_pool->block_allocator,
- block_pool->total_block_size,
- (void**)&block_base));
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0, iree_allocator_malloc(block_pool->block_allocator,
+ block_pool->total_block_size,
+ (void**)&block_base));
block = (iree_arena_block_t*)(block_base + (block_pool->total_block_size -
sizeof(iree_arena_block_t)));
}
block->next = NULL;
*out_block = block;
+
+ IREE_TRACE_ZONE_END(z0);
return iree_ok_status();
}
void iree_arena_block_pool_release(iree_arena_block_pool_t* block_pool,
iree_arena_block_t* block_head,
iree_arena_block_t* block_tail) {
+ IREE_TRACE_ZONE_BEGIN(z0);
iree_atomic_arena_block_slist_concat(&block_pool->available_slist, block_head,
block_tail);
+ IREE_TRACE_ZONE_END(z0);
}
//===----------------------------------------------------------------------===//
diff --git a/iree/modules/hal/hal_module.c b/iree/modules/hal/hal_module.c
index 06c0baf..02f31fe 100644
--- a/iree/modules/hal/hal_module.c
+++ b/iree/modules/hal/hal_module.c
@@ -226,8 +226,8 @@
// Block and wait for the semaphore to be signaled (or fail).
status = iree_hal_semaphore_wait_with_deadline(semaphore, 1ull,
IREE_TIME_INFINITE_FUTURE);
+ iree_hal_semaphore_release(semaphore);
if (!iree_status_is_ok(status)) {
- iree_hal_semaphore_release(semaphore);
return status;
}
diff --git a/iree/tools/iree-benchmark-module-main.cc b/iree/tools/iree-benchmark-module-main.cc
index 0bc1825..b567c1d 100644
--- a/iree/tools/iree-benchmark-module-main.cc
+++ b/iree/tools/iree-benchmark-module-main.cc
@@ -125,21 +125,17 @@
// benchmarking.
class IREEBenchmark {
public:
- IREEBenchmark()
- : instance_(nullptr),
- device_(nullptr),
- hal_module_(nullptr),
- context_(nullptr),
- input_module_(nullptr){};
+ IREEBenchmark() = default;
+
~IREEBenchmark() {
IREE_TRACE_SCOPE0("IREEBenchmark::dtor");
// Order matters.
inputs_.reset();
+ iree_vm_context_release(context_);
iree_vm_module_release(hal_module_);
iree_vm_module_release(input_module_);
iree_hal_device_release(device_);
- iree_vm_context_release(context_);
iree_vm_instance_release(instance_);
};
@@ -247,11 +243,11 @@
}
std::string module_data_;
- iree_vm_instance_t* instance_;
- iree_hal_device_t* device_;
- iree_vm_module_t* hal_module_;
- iree_vm_context_t* context_;
- iree_vm_module_t* input_module_;
+ iree_vm_instance_t* instance_ = nullptr;
+ iree_hal_device_t* device_ = nullptr;
+ iree_vm_module_t* hal_module_ = nullptr;
+ iree_vm_context_t* context_ = nullptr;
+ iree_vm_module_t* input_module_ = nullptr;
iree::vm::ref<iree_vm_list_t> inputs_;
};
} // namespace
diff --git a/iree/vm/bytecode_dispatch.c b/iree/vm/bytecode_dispatch.c
index 86c6509..96292a0 100644
--- a/iree/vm/bytecode_dispatch.c
+++ b/iree/vm/bytecode_dispatch.c
@@ -402,8 +402,7 @@
} break;
case IREE_VM_CCONV_TYPE_REF: {
uint16_t src_reg = src_reg_list->registers[reg_i++];
- iree_vm_ref_retain_or_move(
- src_reg & IREE_REF_REGISTER_MOVE_BIT,
+ iree_vm_ref_assign(
&caller_registers.ref[src_reg & caller_registers.ref_mask],
(iree_vm_ref_t*)p);
p += sizeof(iree_vm_ref_t);
@@ -447,8 +446,7 @@
} break;
case IREE_VM_CCONV_TYPE_REF: {
uint16_t src_reg = src_reg_list->registers[reg_i++];
- iree_vm_ref_retain_or_move(
- src_reg & IREE_REF_REGISTER_MOVE_BIT,
+ iree_vm_ref_assign(
&caller_registers.ref[src_reg & caller_registers.ref_mask],
(iree_vm_ref_t*)p);
p += sizeof(iree_vm_ref_t);
diff --git a/iree/vm/list.c b/iree/vm/list.c
index 64fbdc8..24e3158 100644
--- a/iree/vm/list.c
+++ b/iree/vm/list.c
@@ -70,21 +70,27 @@
iree_host_size_t offset,
iree_host_size_t length) {
switch (list->storage_mode) {
- case IREE_VM_LIST_STORAGE_MODE_VALUE:
- // Nothing special, freeing the storage is all we need.
+ case IREE_VM_LIST_STORAGE_MODE_VALUE: {
+ void* base_ptr =
+ (void*)((uintptr_t)list->storage + offset * list->element_size);
+ memset(base_ptr, 0, length * list->element_size);
break;
+ }
case IREE_VM_LIST_STORAGE_MODE_REF: {
iree_vm_ref_t* ref_storage = (iree_vm_ref_t*)list->storage;
- for (iree_host_size_t i = offset; i < length; ++i) {
+ for (iree_host_size_t i = offset; i < offset + length; ++i) {
iree_vm_ref_release(&ref_storage[i]);
}
break;
}
case IREE_VM_LIST_STORAGE_MODE_VARIANT: {
iree_vm_variant_t* variant_storage = (iree_vm_variant_t*)list->storage;
- for (iree_host_size_t i = offset; i < length; ++i) {
+ for (iree_host_size_t i = offset; i < offset + length; ++i) {
if (iree_vm_type_def_is_ref(&variant_storage[i].type)) {
iree_vm_ref_release(&variant_storage[i].ref);
+ memset(&variant_storage[i].type, 0, sizeof(variant_storage[i].type));
+ } else {
+ memset(&variant_storage[i], 0, sizeof(variant_storage[i]));
}
}
break;
@@ -243,7 +249,7 @@
return iree_ok_status();
} else if (new_size < list->count) {
// Truncating.
- iree_vm_list_reset_range(list, new_size + 1, list->count - new_size);
+ iree_vm_list_reset_range(list, new_size, list->count - new_size);
list->count = new_size;
} else if (new_size > list->capacity) {
// Extending beyond capacity.
diff --git a/iree/vm/list_test.cc b/iree/vm/list_test.cc
index 4423cbc..3bad3ac 100644
--- a/iree/vm/list_test.cc
+++ b/iree/vm/list_test.cc
@@ -229,11 +229,161 @@
iree_vm_list_release(list);
}
-// TODO(benvanik): test resize value.
+// Tests the behavior of resize for truncation and extension on primitives.
+TEST_F(VMListTest, ResizeI32) {
+ iree_vm_type_def_t element_type =
+ iree_vm_type_def_make_value_type(IREE_VM_VALUE_TYPE_I32);
+ iree_host_size_t initial_capacity = 4;
+ iree_vm_list_t* list = nullptr;
+ IREE_ASSERT_OK(iree_vm_list_create(&element_type, initial_capacity,
+ iree_allocator_system(), &list));
+ EXPECT_LE(initial_capacity, iree_vm_list_capacity(list));
+ EXPECT_EQ(0, iree_vm_list_size(list));
-// TODO(benvanik): test resize ref.
+ // Extend and zero-initialize.
+ IREE_ASSERT_OK(iree_vm_list_resize(list, 5));
+ for (iree_host_size_t i = 0; i < 5; ++i) {
+ iree_vm_value_t value;
+ IREE_ASSERT_OK(
+ iree_vm_list_get_value_as(list, i, IREE_VM_VALUE_TYPE_I32, &value));
+ EXPECT_EQ(0, value.i32);
+ }
-// TODO(benvanik): test resize variant.
+ // Overwrite with [0, 5).
+ for (iree_host_size_t i = 0; i < 5; ++i) {
+ iree_vm_value_t value = iree_vm_value_make_i32((int32_t)i);
+ IREE_ASSERT_OK(iree_vm_list_set_value(list, i, &value));
+ }
+
+ // Truncate to [0, 2) and then extend again.
+ // This ensures that we test the primitive clearing path during cleanup:
+ // [int, int, int, int, int]
+ // |___________| <- truncation region
+ IREE_ASSERT_OK(iree_vm_list_resize(list, 2));
+ IREE_ASSERT_OK(iree_vm_list_resize(list, 5));
+
+ // Ensure that elements 2+ are zeroed after having been reset while 0 and 1
+ // are still valid as before.
+ for (iree_host_size_t i = 0; i < 2; ++i) {
+ iree_vm_value_t value;
+ IREE_ASSERT_OK(
+ iree_vm_list_get_value_as(list, i, IREE_VM_VALUE_TYPE_I32, &value));
+ EXPECT_EQ(i, value.i32);
+ }
+ for (iree_host_size_t i = 2; i < 5; ++i) {
+ iree_vm_value_t value;
+ IREE_ASSERT_OK(
+ iree_vm_list_get_value_as(list, i, IREE_VM_VALUE_TYPE_I32, &value));
+ EXPECT_EQ(0, value.i32);
+ }
+
+ iree_vm_list_release(list);
+}
+
+// Tests the behavior of resize for truncation and extension on refs.
+TEST_F(VMListTest, ResizeRef) {
+ iree_vm_type_def_t element_type =
+ iree_vm_type_def_make_ref_type(test_a_type_id());
+ iree_host_size_t initial_capacity = 4;
+ iree_vm_list_t* list = nullptr;
+ IREE_ASSERT_OK(iree_vm_list_create(&element_type, initial_capacity,
+ iree_allocator_system(), &list));
+ EXPECT_LE(initial_capacity, iree_vm_list_capacity(list));
+ EXPECT_EQ(0, iree_vm_list_size(list));
+
+ // Extend and zero-initialize.
+ IREE_ASSERT_OK(iree_vm_list_resize(list, 5));
+ for (iree_host_size_t i = 0; i < 5; ++i) {
+ iree_vm_ref_t ref_a{0};
+ IREE_ASSERT_OK(iree_vm_list_get_ref_assign(list, i, &ref_a));
+ EXPECT_TRUE(iree_vm_ref_is_null(&ref_a));
+ }
+
+ // Overwrite with [0, 5).
+ for (iree_host_size_t i = 0; i < 5; ++i) {
+ iree_vm_ref_t ref_a = MakeRef<A>((float)i);
+ IREE_ASSERT_OK(iree_vm_list_set_ref_move(list, i, &ref_a));
+ }
+
+ // Truncate to [0, 2) and then extend again.
+ // This ensures that we test the ref path during cleanup:
+ // [ref, ref, ref, ref, ref]
+ // |___________| <- truncation region
+ IREE_ASSERT_OK(iree_vm_list_resize(list, 2));
+ IREE_ASSERT_OK(iree_vm_list_resize(list, 5));
+
+ // Ensure that elements 2+ are reset after having been reset while 0 and 1
+ // are still valid as before.
+ for (iree_host_size_t i = 0; i < 2; ++i) {
+ iree_vm_ref_t ref_a{0};
+ IREE_ASSERT_OK(iree_vm_list_get_ref_retain(list, i, &ref_a));
+ EXPECT_TRUE(test_a_isa(ref_a));
+ auto* a = test_a_deref(ref_a);
+ EXPECT_EQ(i, a->data());
+ iree_vm_ref_release(&ref_a);
+ }
+ for (iree_host_size_t i = 2; i < 5; ++i) {
+ iree_vm_ref_t ref_a{0};
+ IREE_ASSERT_OK(iree_vm_list_get_ref_assign(list, i, &ref_a));
+ EXPECT_TRUE(iree_vm_ref_is_null(&ref_a));
+ }
+
+ iree_vm_list_release(list);
+}
+
+// Tests the behavior of resize for truncation and extension on variants.
+TEST_F(VMListTest, ResizeVariant) {
+ iree_vm_type_def_t element_type = iree_vm_type_def_make_variant_type();
+ iree_host_size_t initial_capacity = 4;
+ iree_vm_list_t* list = nullptr;
+ IREE_ASSERT_OK(iree_vm_list_create(&element_type, initial_capacity,
+ iree_allocator_system(), &list));
+ EXPECT_LE(initial_capacity, iree_vm_list_capacity(list));
+ EXPECT_EQ(0, iree_vm_list_size(list));
+
+ // Extend and zero-initialize.
+ IREE_ASSERT_OK(iree_vm_list_resize(list, 5));
+ for (iree_host_size_t i = 0; i < 5; ++i) {
+ iree_vm_variant_t value = iree_vm_variant_empty();
+ IREE_ASSERT_OK(iree_vm_list_get_variant(list, i, &value));
+ EXPECT_TRUE(iree_vm_variant_is_empty(value));
+ }
+
+ // Overwrite with [0, 5) in mixed types.
+ for (iree_host_size_t i = 0; i < 4; ++i) {
+ iree_vm_ref_t ref_a = MakeRef<A>((float)i);
+ IREE_ASSERT_OK(iree_vm_list_set_ref_move(list, i, &ref_a));
+ }
+ for (iree_host_size_t i = 4; i < 5; ++i) {
+ iree_vm_value_t value = iree_vm_value_make_i32((int32_t)i);
+ IREE_ASSERT_OK(iree_vm_list_set_value(list, i, &value));
+ }
+
+ // Truncate to [0, 2) and then extend again.
+ // This ensures that we test the variant path during cleanup:
+ // [ref, ref, ref, ref, int]
+ // |___________| <- truncation region
+ IREE_ASSERT_OK(iree_vm_list_resize(list, 2));
+ IREE_ASSERT_OK(iree_vm_list_resize(list, 5));
+
+ // Ensure that elements 2+ are reset after having been reset while 0 and 1
+ // are still valid as before.
+ for (iree_host_size_t i = 0; i < 2; ++i) {
+ iree_vm_ref_t ref_a{0};
+ IREE_ASSERT_OK(iree_vm_list_get_ref_retain(list, i, &ref_a));
+ EXPECT_TRUE(test_a_isa(ref_a));
+ auto* a = test_a_deref(ref_a);
+ EXPECT_EQ(i, a->data());
+ iree_vm_ref_release(&ref_a);
+ }
+ for (iree_host_size_t i = 2; i < 5; ++i) {
+ iree_vm_variant_t value = iree_vm_variant_empty();
+ IREE_ASSERT_OK(iree_vm_list_get_variant(list, i, &value));
+ EXPECT_TRUE(iree_vm_variant_is_empty(value));
+ }
+
+ iree_vm_list_release(list);
+}
// TODO(benvanik): test value get/set.
diff --git a/iree/vm/module.h b/iree/vm/module.h
index d8635e0..8eabc0a 100644
--- a/iree/vm/module.h
+++ b/iree/vm/module.h
@@ -192,9 +192,8 @@
// Argument buffer in the format described above.
// This is only read on beginning the function and need not live beyond that.
//
- // Refs contained will be moved into the target function or released if not
- // needed. Callers must ensure they move or retain arguments when populating
- // the arguments buffer.
+ // Refs contained are retained by the caller and callees must retain them if
+ // they need them to live beyond the call.
iree_byte_span_t arguments;
// Storage for the result buffer; assumed undefined and then populated with
diff --git a/iree/vm/module_abi_packing.h b/iree/vm/module_abi_packing.h
index 0b00b30..d679964 100644
--- a/iree/vm/module_abi_packing.h
+++ b/iree/vm/module_abi_packing.h
@@ -362,7 +362,7 @@
struct ParamUnpack<opaque_ref> {
using storage_type = opaque_ref;
static void Load(Status& status, params_ptr_t& ptr, storage_type& out_param) {
- iree_vm_ref_move(reinterpret_cast<iree_vm_ref_t*>(ptr), &out_param);
+ iree_vm_ref_retain(reinterpret_cast<iree_vm_ref_t*>(ptr), &out_param);
ptr += sizeof(iree_vm_ref_t);
}
};
@@ -376,7 +376,7 @@
auto* reg_ptr = reinterpret_cast<iree_vm_ref_t*>(ptr);
ptr += sizeof(iree_vm_ref_t);
if (reg_ptr->type == ref_type_descriptor<T>::get()->type) {
- out_param = vm::assign_ref(reinterpret_cast<T*>(reg_ptr->ptr));
+ out_param = vm::retain_ref(reinterpret_cast<T*>(reg_ptr->ptr));
memset(reg_ptr, 0, sizeof(*reg_ptr));
} else if (IREE_UNLIKELY(reg_ptr->type != IREE_VM_REF_TYPE_NULL)) {
status =
@@ -401,7 +401,7 @@
auto* reg_ptr = reinterpret_cast<iree_vm_ref_t*>(ptr);
ptr += sizeof(iree_vm_ref_t);
if (reg_ptr->type == ref_type_descriptor<T>::get()->type) {
- out_param = vm::assign_ref(reinterpret_cast<T*>(reg_ptr->ptr));
+ out_param = vm::retain_ref(reinterpret_cast<T*>(reg_ptr->ptr));
memset(reg_ptr, 0, sizeof(*reg_ptr));
} else if (IREE_UNLIKELY(reg_ptr->type != IREE_VM_REF_TYPE_NULL)) {
status =
diff --git a/iree/vm/type_def.h b/iree/vm/type_def.h
index 5592494..c3294d6 100644
--- a/iree/vm/type_def.h
+++ b/iree/vm/type_def.h
@@ -87,6 +87,7 @@
{ {IREE_VM_VALUE_TYPE_NONE, IREE_VM_REF_TYPE_NULL}, {0}, }
#define iree_vm_variant_is_value(v) iree_vm_type_def_is_value(&(v).type)
#define iree_vm_variant_is_ref(v) iree_vm_type_def_is_ref(&(v).type)
+#define iree_vm_variant_is_empty(v) iree_vm_type_def_is_variant(&(v).type)
#ifdef __cplusplus
} // extern "C"
