Merge pull request #8006 from not-jenni:main-to-google PiperOrigin-RevId: 419843023
diff --git a/CMakeLists.txt b/CMakeLists.txt index 13ecbb0..3e67822 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt
@@ -25,6 +25,8 @@ enable_language(ASM) endif() +include(CMakeDependentOption) + #------------------------------------------------------------------------------- # Project component configuration #------------------------------------------------------------------------------- @@ -40,6 +42,7 @@ option(IREE_BUILD_SAMPLES "Builds IREE sample projects." ON) option(IREE_BUILD_TRACY "Builds tracy server tools." OFF) +option(IREE_BUILD_LEGACY_JAX "Builds the legacy JAX Python API" ON) option(IREE_BUILD_TENSORFLOW_ALL "Builds all TensorFlow compiler frontends." OFF) option(IREE_BUILD_TENSORFLOW_COMPILER "Builds TensorFlow compiler frontend." "${IREE_BUILD_TENSORFLOW_ALL}") option(IREE_BUILD_TFLITE_COMPILER "Builds the TFLite compiler frontend." "${IREE_BUILD_TENSORFLOW_ALL}") @@ -79,7 +82,7 @@ # Derived flags based on primary options #------------------------------------------------------------------------------- -option(IREE_ENABLE_EMITC "Enables MLIR EmitC dependencies." ${IREE_BUILD_COMPILER}) +cmake_dependent_option(IREE_ENABLE_EMITC "Enables MLIR EmitC dependencies." ON ${IREE_BUILD_COMPILER} OFF) #------------------------------------------------------------------------------- # Target and backend configuration @@ -110,15 +113,15 @@ option(IREE_HAL_DRIVER_VMVX_SYNC "Enables the 'vmvx-sync' runtime HAL driver" ${IREE_HAL_DRIVER_DEFAULTS}) option(IREE_HAL_DRIVER_VULKAN "Enables the 'vulkan' runtime HAL driver" ${IREE_HAL_DRIVER_VULKAN_DEFAULT}) -option(IREE_TARGET_BACKEND_CUDA "Enables the 'cuda' compiler target backend" ${IREE_BUILD_COMPILER}) -option(IREE_TARGET_BACKEND_DYLIB_LLVM_AOT "Enables the 'dylib-llvm-aot' compiler target backend" ${IREE_BUILD_COMPILER}) -option(IREE_TARGET_BACKEND_METAL_SPIRV "Enables the 'metal-spirv' compiler target backend" ${IREE_BUILD_COMPILER}) -option(IREE_TARGET_BACKEND_ROCM "Enables the 'rocm' compiler target backend" ${IREE_BUILD_COMPILER}) -option(IREE_TARGET_BACKEND_VMVX "Enables the 'vmvx' compiler target backend" ${IREE_BUILD_COMPILER}) -option(IREE_TARGET_BACKEND_VULKAN_SPIRV "Enables the 'vulkan-spirv' compiler target backend" ${IREE_BUILD_COMPILER}) -option(IREE_TARGET_BACKEND_WASM_LLVM_AOT "Enables the 'wasm-llvm-aot' compiler target backend" ${IREE_BUILD_COMPILER}) +cmake_dependent_option(IREE_TARGET_BACKEND_CUDA "Enables the 'cuda' compiler target backend" ON ${IREE_BUILD_COMPILER} OFF) +cmake_dependent_option(IREE_TARGET_BACKEND_DYLIB_LLVM_AOT "Enables the 'dylib-llvm-aot' compiler target backend" ON ${IREE_BUILD_COMPILER} OFF) +cmake_dependent_option(IREE_TARGET_BACKEND_METAL_SPIRV "Enables the 'metal-spirv' compiler target backend" ON ${IREE_BUILD_COMPILER} OFF) +cmake_dependent_option(IREE_TARGET_BACKEND_ROCM "Enables the 'rocm' compiler target backend" ON ${IREE_BUILD_COMPILER} OFF) +cmake_dependent_option(IREE_TARGET_BACKEND_VMVX "Enables the 'vmvx' compiler target backend" ON ${IREE_BUILD_COMPILER} OFF) +cmake_dependent_option(IREE_TARGET_BACKEND_VULKAN_SPIRV "Enables the 'vulkan-spirv' compiler target backend" ON ${IREE_BUILD_COMPILER} OFF) +cmake_dependent_option(IREE_TARGET_BACKEND_WASM_LLVM_AOT "Enables the 'wasm-llvm-aot' compiler target backend" ON ${IREE_BUILD_COMPILER} OFF) # Disable WebGPU by default - it has complex deps and is under development. -option(IREE_TARGET_BACKEND_WEBGPU "Enables the 'webgpu' compiler target backend" OFF) +cmake_dependent_option(IREE_TARGET_BACKEND_WEBGPU "Enables the 'webgpu' compiler target backend" OFF ${IREE_BUILD_COMPILER} OFF) message(VERBOSE "IREE build runtime HAL driver 'cuda': ${IREE_HAL_DRIVER_CUDA}") message(VERBOSE "IREE build runtime HAL driver 'dylib': ${IREE_HAL_DRIVER_DYLIB}")
diff --git a/bindings/python/CMakeLists.txt b/bindings/python/CMakeLists.txt index f019d28..77bf1ce 100644 --- a/bindings/python/CMakeLists.txt +++ b/bindings/python/CMakeLists.txt
@@ -11,4 +11,8 @@ # Namespace packages. add_subdirectory(iree/runtime) -add_subdirectory(iree/jax) + +if(IREE_BUILD_LEGACY_JAX) + message(STATUS "Building legacy JAX API") + add_subdirectory(iree/jax) +endif()
diff --git a/bindings/python/iree/runtime/function.py b/bindings/python/iree/runtime/function.py index 826fac9..a7e3874 100644 --- a/bindings/python/iree/runtime/function.py +++ b/bindings/python/iree/runtime/function.py
@@ -98,54 +98,55 @@ call_trace = None # type: Optional[tracing.CallTrace] if self._tracer: call_trace = self._tracer.start_call(self._vm_function) + try: + # Initialize the capacity to our total number of args, since we should + # be below that when doing a flat invocation. May want to be more + # conservative here when considering nesting. + inv = Invocation(self._device) + ret_descs = self._ret_descs - # Initialize the capacity to our total number of args, since we should - # be below that when doing a flat invocation. May want to be more - # conservative here when considering nesting. - inv = Invocation(self._device) - ret_descs = self._ret_descs + # Merge keyword args in by name->position mapping. + if kwargs: + args = list(args) + len_delta = self._max_named_arg_index - len(args) + 1 + if len_delta > 0: + # Fill in MissingArgument placeholders before arranging kwarg input. + # Any remaining placeholders will fail arity checks later on. + args.extend([MissingArgument] * len_delta) - # Merge keyword args in by name->position mapping. - if kwargs: - args = list(args) - len_delta = self._max_named_arg_index - len(args) + 1 - if len_delta > 0: - # Fill in MissingArgument placeholders before arranging kwarg input. - # Any remaining placeholders will fail arity checks later on. - args.extend([MissingArgument] * len_delta) + for kwarg_key, kwarg_value in kwargs.items(): + try: + kwarg_index = self._named_arg_indices[kwarg_key] + except KeyError: + raise ArgumentError(f"specified kwarg '{kwarg_key}' is unknown") + args[kwarg_index] = kwarg_value - for kwarg_key, kwarg_value in kwargs.items(): - try: - kwarg_index = self._named_arg_indices[kwarg_key] - except KeyError: - raise ArgumentError(f"specified kwarg '{kwarg_key}' is unknown") - args[kwarg_index] = kwarg_value + arg_list = VmVariantList(len(args)) + ret_list = VmVariantList(len(ret_descs) if ret_descs is not None else 1) + _merge_python_sequence_to_vm(inv, arg_list, args, self._arg_descs) + if call_trace: + call_trace.add_vm_list(arg_list, "args") + self._vm_context.invoke(self._vm_function, arg_list, ret_list) + if call_trace: + call_trace.add_vm_list(ret_list, "results") - arg_list = VmVariantList(len(args)) - ret_list = VmVariantList(len(ret_descs) if ret_descs is not None else 1) - _merge_python_sequence_to_vm(inv, arg_list, args, self._arg_descs) - if call_trace: - call_trace.add_vm_list(arg_list, "args") - self._vm_context.invoke(self._vm_function, arg_list, ret_list) - if call_trace: - call_trace.add_vm_list(ret_list, "results") - - # Un-inline the results to align with reflection, as needed. - reflection_aligned_ret_list = ret_list - if self._has_inlined_results: - reflection_aligned_ret_list = VmVariantList(1) - reflection_aligned_ret_list.push_list(ret_list) - returns = _extract_vm_sequence_to_python(inv, reflection_aligned_ret_list, - ret_descs) - if call_trace: - call_trace.end_call() - return_arity = len(returns) - if return_arity == 1: - return returns[0] - elif return_arity == 0: - return None - else: - return tuple(returns) + # Un-inline the results to align with reflection, as needed. + reflection_aligned_ret_list = ret_list + if self._has_inlined_results: + reflection_aligned_ret_list = VmVariantList(1) + reflection_aligned_ret_list.push_list(ret_list) + returns = _extract_vm_sequence_to_python(inv, reflection_aligned_ret_list, + ret_descs) + return_arity = len(returns) + if return_arity == 1: + return returns[0] + elif return_arity == 0: + return None + else: + return tuple(returns) + finally: + if call_trace: + call_trace.end_call() def _parse_abi_dict(self, vm_function: VmFunction): reflection = vm_function.reflection
diff --git a/bindings/python/iree/runtime/hal.h b/bindings/python/iree/runtime/hal.h index 011e9d6..72debbb 100644 --- a/bindings/python/iree/runtime/hal.h +++ b/bindings/python/iree/runtime/hal.h
@@ -134,11 +134,12 @@ static HalMappedMemory Create(HalBufferView& bv) { iree_hal_buffer_t* buffer = iree_hal_buffer_view_buffer(bv.raw_ptr()); iree_device_size_t byte_length = iree_hal_buffer_byte_length(buffer); - iree_hal_buffer_mapping_t mapped_memory; - CheckApiStatus(iree_hal_buffer_map_range( - buffer, IREE_HAL_MEMORY_ACCESS_READ, - 0 /* element_offset */, byte_length, &mapped_memory), - "Could not map memory"); + iree_hal_buffer_mapping_t mapped_memory = {{0}}; + CheckApiStatus( + iree_hal_buffer_map_range(buffer, IREE_HAL_MAPPING_MODE_SCOPED, + IREE_HAL_MEMORY_ACCESS_READ, 0, byte_length, + &mapped_memory), + "Could not map memory"); return HalMappedMemory(mapped_memory, bv.raw_ptr()); } @@ -168,8 +169,8 @@ } private: - iree_hal_buffer_mapping_t mapped_memory_; - iree_hal_buffer_view_t* bv_; + iree_hal_buffer_mapping_t mapped_memory_ = {{0}}; + iree_hal_buffer_view_t* bv_ = nullptr; }; void SetupHalBindings(pybind11::module m);
diff --git a/bindings/python/iree/runtime/vm.cc b/bindings/python/iree/runtime/vm.cc index d3d9dfa..15a6365 100644 --- a/bindings/python/iree/runtime/vm.cc +++ b/bindings/python/iree/runtime/vm.cc
@@ -218,16 +218,15 @@ // TODO(laurenzo): Expand to other layouts as needed. // TODO(laurenzo): Wrap and retain original buffer (depends_on_pyobject=true). iree_hal_buffer_t* raw_buffer; - CheckApiStatus(iree_hal_allocator_allocate_buffer( - device.allocator(), - static_cast<iree_hal_memory_type_t>( - IREE_HAL_MEMORY_TYPE_HOST_LOCAL | - IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE), - IREE_HAL_BUFFER_USAGE_ALL, py_view.len, &raw_buffer), - "Failed to allocate device visible buffer"); CheckApiStatus( - iree_hal_buffer_write_data(raw_buffer, 0, py_view.buf, py_view.len), - "Error writing to input buffer"); + iree_hal_allocator_allocate_buffer( + device.allocator(), + static_cast<iree_hal_memory_type_t>( + IREE_HAL_MEMORY_TYPE_HOST_LOCAL | + IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE), + IREE_HAL_BUFFER_USAGE_ALL, py_view.len, + iree_make_const_byte_span(py_view.buf, py_view.len), &raw_buffer), + "Failed to allocate device visible buffer"); // Only capture the reference to the exporting object (incrementing it) // once guaranteed successful. @@ -390,10 +389,11 @@ // Map memory. iree_device_size_t byte_length = iree_hal_buffer_byte_length(raw_buffer); - iree_hal_buffer_mapping_t mapped_memory; + iree_hal_buffer_mapping_t mapped_memory = {{0}}; CheckApiStatus(iree_hal_buffer_map_range( - raw_buffer, IREE_HAL_MEMORY_ACCESS_READ, - 0 /* element_offset */, byte_length, &mapped_memory), + raw_buffer, IREE_HAL_MAPPING_MODE_SCOPED, + IREE_HAL_MEMORY_ACCESS_READ, 0 /* element_offset */, + byte_length, &mapped_memory), "Could not map memory"); record["contents"] = py::bytes(reinterpret_cast<const char*>(mapped_memory.contents.data), @@ -489,10 +489,11 @@ // Map memory. iree_device_size_t byte_length = iree_hal_buffer_byte_length(buffer.raw_ptr()); - iree_hal_buffer_mapping_t mapped_memory; + iree_hal_buffer_mapping_t mapped_memory = {{0}}; CheckApiStatus(iree_hal_buffer_map_range( - buffer.raw_ptr(), IREE_HAL_MEMORY_ACCESS_READ, - 0 /* element_offset */, byte_length, &mapped_memory), + buffer.raw_ptr(), IREE_HAL_MAPPING_MODE_SCOPED, + IREE_HAL_MEMORY_ACCESS_READ, 0 /* element_offset */, + byte_length, &mapped_memory), "Could not map memory"); // Turn the mapping into a python object that retains until the array is
diff --git a/bindings/tflite/tensor.c b/bindings/tflite/tensor.c index 1fdb16d..0112fd6 100644 --- a/bindings/tflite/tensor.c +++ b/bindings/tflite/tensor.c
@@ -140,7 +140,8 @@ iree_hal_allocator_allocate_buffer( buffer_allocator, IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL | IREE_HAL_MEMORY_TYPE_HOST_VISIBLE, - IREE_HAL_BUFFER_USAGE_ALL, allocation_size, &tensor->buffer)); + IREE_HAL_BUFFER_USAGE_ALL, allocation_size, + iree_const_byte_span_empty(), &tensor->buffer)); // Map the buffer memory immediately. The tflite API doesn't let us know if // this is a buffer the user will actually touch or some state buffer that is @@ -149,7 +150,8 @@ // puts potential errors in the same easy to find place. IREE_RETURN_AND_END_ZONE_IF_ERROR( z0, - iree_hal_buffer_map_range(tensor->buffer, IREE_HAL_MEMORY_ACCESS_ALL, 0, + iree_hal_buffer_map_range(tensor->buffer, IREE_HAL_MAPPING_MODE_SCOPED, + IREE_HAL_MEMORY_ACCESS_ALL, 0, IREE_WHOLE_BUFFER, &tensor->buffer_mapping)); IREE_TRACE_ZONE_END(z0); @@ -173,10 +175,10 @@ iree_device_size_t byte_offset = 0; iree_device_size_t byte_length = IREE_WHOLE_BUFFER; IREE_RETURN_AND_END_ZONE_IF_ERROR( - z0, - iree_hal_buffer_map_range( - buffer, IREE_HAL_MEMORY_ACCESS_READ | IREE_HAL_MEMORY_ACCESS_WRITE, - byte_offset, byte_length, &tensor->buffer_mapping)); + z0, iree_hal_buffer_map_range( + buffer, IREE_HAL_MAPPING_MODE_SCOPED, + IREE_HAL_MEMORY_ACCESS_READ | IREE_HAL_MEMORY_ACCESS_WRITE, + byte_offset, byte_length, &tensor->buffer_mapping)); // Retain the buffer view until discarded/reset. tensor->buffer = buffer;
diff --git a/build_tools/buildkite/cmake/android/arm64-v8a/benchmark2.yml b/build_tools/buildkite/cmake/android/arm64-v8a/benchmark2.yml index 1cc771a..ed945de 100644 --- a/build_tools/buildkite/cmake/android/arm64-v8a/benchmark2.yml +++ b/build_tools/buildkite/cmake/android/arm64-v8a/benchmark2.yml
@@ -63,26 +63,6 @@ - "trace-captures-galaxy-pixel6-pro-${BUILDKITE_BUILD_NUMBER}.tgz" timeout_in_minutes: "60" - - label: "Benchmark on Galaxy S20 (exynos-990, mali-g77)" - commands: - - "git clean -fdx" - - "buildkite-agent artifact download --step Build benchmark-suites-${BUILDKITE_BUILD_NUMBER}.tgz ./" - - "buildkite-agent artifact download --step Build iree-android-tools-${BUILDKITE_BUILD_NUMBER}.tgz ./" - - "wget https://storage.googleapis.com/iree-shared-files/tracy-capture-058e8901.tgz" - - "tar -xzvf benchmark-suites-${BUILDKITE_BUILD_NUMBER}.tgz" - - "tar -xzvf iree-android-tools-${BUILDKITE_BUILD_NUMBER}.tgz" - - "tar -xzvf tracy-capture-058e8901.tgz" - - "python3 build_tools/benchmarks/run_benchmarks_on_android.py --normal_benchmark_tool=build-android/iree/tools/iree-benchmark-module --traced_benchmark_tool=build-android-trace/iree/tools/iree-benchmark-module --trace_capture_tool=tracy-capture -o benchmark-results-galaxy-s20-${BUILDKITE_BUILD_NUMBER}.json --capture_tarball=trace-captures-galaxy-s20-${BUILDKITE_BUILD_NUMBER}.tgz --verbose build-host/" - if: "build.pull_request.id == null || (build.pull_request.labels includes 'buildkite:benchmark')" - agents: - - "android-soc=exynos-990" - - "android-version=11" - - "queue=benchmark-android" - artifact_paths: - - "benchmark-results-galaxy-s20-${BUILDKITE_BUILD_NUMBER}.json" - - "trace-captures-galaxy-s20-${BUILDKITE_BUILD_NUMBER}.tgz" - timeout_in_minutes: "60" - - wait - label: "Comment benchmark results on pull request"
diff --git a/experimental/rocm/CMakeLists.txt b/experimental/rocm/CMakeLists.txt index 69943a7..0b4438d 100644 --- a/experimental/rocm/CMakeLists.txt +++ b/experimental/rocm/CMakeLists.txt
@@ -53,6 +53,7 @@ iree::base::internal::synchronization iree::base::tracing iree::hal + iree::hal::utils::buffer_transfer iree::schemas::rocm_executable_def_c_fbs PUBLIC )
diff --git a/experimental/rocm/rocm_allocator.c b/experimental/rocm/rocm_allocator.c index 093cd21..85a6a3d 100644 --- a/experimental/rocm/rocm_allocator.c +++ b/experimental/rocm/rocm_allocator.c
@@ -16,6 +16,7 @@ typedef struct iree_hal_rocm_allocator_t { iree_hal_resource_t resource; + iree_hal_device_t* base_device; iree_hal_rocm_context_wrapper_t* context; IREE_STATISTICS(iree_hal_allocator_statistics_t statistics;) @@ -30,8 +31,9 @@ } iree_status_t iree_hal_rocm_allocator_create( - iree_hal_rocm_context_wrapper_t* context, + iree_hal_device_t* base_device, iree_hal_rocm_context_wrapper_t* context, iree_hal_allocator_t** out_allocator) { + IREE_ASSERT_ARGUMENT(base_device); IREE_ASSERT_ARGUMENT(context); IREE_TRACE_ZONE_BEGIN(z0); iree_hal_rocm_allocator_t* allocator = NULL; @@ -67,6 +69,11 @@ return allocator->context->host_allocator; } +static iree_status_t iree_hal_rocm_allocator_trim( + iree_hal_allocator_t* base_allocator) { + return iree_ok_status(); +} + static void iree_hal_rocm_allocator_query_statistics( iree_hal_allocator_t* base_allocator, iree_hal_allocator_statistics_t* out_statistics) { @@ -123,7 +130,7 @@ static iree_status_t iree_hal_rocm_allocator_allocate_buffer( iree_hal_allocator_t* base_allocator, iree_hal_memory_type_t memory_type, iree_hal_buffer_usage_t allowed_usage, iree_host_size_t allocation_size, - iree_hal_buffer_t** out_buffer) { + iree_const_byte_span_t initial_data, iree_hal_buffer_t** out_buffer) { iree_hal_rocm_allocator_t* allocator = iree_hal_rocm_allocator_cast(base_allocator); // Guard against the corner case where the requested buffer size is 0. The @@ -160,6 +167,7 @@ } } + iree_hal_buffer_t* buffer = NULL; if (iree_status_is_ok(status)) { status = iree_hal_rocm_buffer_wrap( (iree_hal_allocator_t*)allocator, memory_type, @@ -167,12 +175,30 @@ /*byte_offset=*/0, /*byte_length=*/allocation_size, device_ptr, host_ptr, out_buffer); } + + // Copy the initial contents into the buffer. This may require staging. + if (iree_status_is_ok(status) && + !iree_const_byte_span_is_empty(initial_data)) { + status = iree_hal_device_transfer_range( + allocator->base_device, + iree_hal_make_host_transfer_buffer_span((void*)initial_data.data, + initial_data.data_length), + 0, iree_hal_make_device_transfer_buffer(buffer), 0, + initial_data.data_length, IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, + iree_infinite_timeout()); + } + if (iree_status_is_ok(status)) { IREE_STATISTICS(iree_hal_allocator_statistics_record_alloc( &allocator->statistics, memory_type, allocation_size)); + *out_buffer = buffer; } else { - iree_hal_rocm_buffer_free(allocator->context, memory_type, device_ptr, - host_ptr); + if (!buffer) { + iree_hal_rocm_buffer_free(allocator->context, memory_type, device_ptr, + host_ptr); + } else { + iree_hal_buffer_release(buffer); + } } return status; } @@ -206,6 +232,7 @@ static const iree_hal_allocator_vtable_t iree_hal_rocm_allocator_vtable = { .destroy = iree_hal_rocm_allocator_destroy, .host_allocator = iree_hal_rocm_allocator_host_allocator, + .trim = iree_hal_rocm_allocator_trim, .query_statistics = iree_hal_rocm_allocator_query_statistics, .query_buffer_compatibility = iree_hal_rocm_allocator_query_buffer_compatibility,
diff --git a/experimental/rocm/rocm_allocator.h b/experimental/rocm/rocm_allocator.h index a2a89ea..9156b70 100644 --- a/experimental/rocm/rocm_allocator.h +++ b/experimental/rocm/rocm_allocator.h
@@ -18,7 +18,7 @@ // Create a ROCM allocator. iree_status_t iree_hal_rocm_allocator_create( - iree_hal_rocm_context_wrapper_t* context, + iree_hal_device_t* base_device, iree_hal_rocm_context_wrapper_t* context, iree_hal_allocator_t** out_allocator); #ifdef __cplusplus
diff --git a/experimental/rocm/rocm_buffer.c b/experimental/rocm/rocm_buffer.c index b39a40e..047f251 100644 --- a/experimental/rocm/rocm_buffer.c +++ b/experimental/rocm/rocm_buffer.c
@@ -68,14 +68,16 @@ iree_hal_buffer_t* base_buffer, iree_hal_mapping_mode_t mapping_mode, iree_hal_memory_access_t memory_access, iree_device_size_t local_byte_offset, iree_device_size_t local_byte_length, - void** out_data_ptr) { + iree_hal_buffer_mapping_t* mapping) { iree_hal_rocm_buffer_t* buffer = iree_hal_rocm_buffer_cast(base_buffer); - if (!iree_all_bits_set(buffer->base.memory_type, - IREE_HAL_MEMORY_TYPE_HOST_VISIBLE)) { - return iree_make_status(IREE_STATUS_INTERNAL, - "trying to map memory not host visible"); - } + // TODO(benvanik): add upload/download for unmapped buffers. + IREE_RETURN_IF_ERROR(iree_hal_buffer_validate_memory_type( + iree_hal_buffer_memory_type(base_buffer), + IREE_HAL_MEMORY_TYPE_HOST_VISIBLE)); + IREE_RETURN_IF_ERROR( + iree_hal_buffer_validate_usage(iree_hal_buffer_allowed_usage(base_buffer), + IREE_HAL_BUFFER_USAGE_MAPPING)); uint8_t* data_ptr = (uint8_t*)(buffer->host_ptr) + local_byte_offset; // If we mapped for discard scribble over the bytes. This is not a mandated @@ -87,14 +89,16 @@ memset(data_ptr, 0xCD, local_byte_length); } #endif // !NDEBUG - *out_data_ptr = data_ptr; + + mapping->contents = iree_make_byte_span(data_ptr, local_byte_length); return iree_ok_status(); } -static void iree_hal_rocm_buffer_unmap_range( +static iree_status_t iree_hal_rocm_buffer_unmap_range( iree_hal_buffer_t* base_buffer, iree_device_size_t local_byte_offset, - iree_device_size_t local_byte_length, void* data_ptr) { - // nothing to do. + iree_device_size_t local_byte_length, iree_hal_buffer_mapping_t* mapping) { + // Nothing to do (today). + return iree_ok_status(); } static iree_status_t iree_hal_rocm_buffer_invalidate_range(
diff --git a/experimental/rocm/rocm_device.c b/experimental/rocm/rocm_device.c index bfa86aa..6666bd4 100644 --- a/experimental/rocm/rocm_device.c +++ b/experimental/rocm/rocm_device.c
@@ -21,6 +21,7 @@ #include "experimental/rocm/rocm_event.h" #include "experimental/rocm/status_util.h" #include "iree/base/tracing.h" +#include "iree/hal/utils/buffer_transfer.h" //===----------------------------------------------------------------------===// // iree_hal_rocm_device_t @@ -91,7 +92,8 @@ device->context_wrapper.host_allocator = host_allocator; device->context_wrapper.syms = syms; iree_status_t status = iree_hal_rocm_allocator_create( - &device->context_wrapper, &device->device_allocator); + (iree_hal_device_t*)device, &device->context_wrapper, + &device->device_allocator); if (iree_status_is_ok(status)) { *out_device = (iree_hal_device_t*)device; } else { @@ -168,6 +170,12 @@ (int)category.size, category.data, (int)key.size, key.data); } +static iree_status_t iree_hal_rocm_device_trim(iree_hal_device_t* base_device) { + iree_hal_rocm_device_t* device = iree_hal_rocm_device_cast(base_device); + iree_arena_block_pool_trim(&device->block_pool); + return iree_hal_allocator_trim(device->device_allocator); +} + static iree_status_t iree_hal_rocm_device_create_command_buffer( iree_hal_device_t* base_device, iree_hal_command_buffer_mode_t mode, iree_hal_command_category_t command_categories, @@ -289,6 +297,7 @@ .id = iree_hal_rocm_device_id, .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, .create_command_buffer = iree_hal_rocm_device_create_command_buffer, .create_descriptor_set = iree_hal_rocm_device_create_descriptor_set, @@ -298,6 +307,7 @@ .create_executable_cache = iree_hal_rocm_device_create_executable_cache, .create_executable_layout = iree_hal_rocm_device_create_executable_layout, .create_semaphore = iree_hal_rocm_device_create_semaphore, + .transfer_range = iree_hal_device_submit_transfer_range_and_wait, .queue_submit = iree_hal_rocm_device_queue_submit, .submit_and_wait = iree_hal_rocm_device_submit_and_wait, .wait_semaphores = iree_hal_rocm_device_wait_semaphores,
diff --git a/iree/base/allocator.h b/iree/base/allocator.h index ac06123..ea2acf6 100644 --- a/iree/base/allocator.h +++ b/iree/base/allocator.h
@@ -69,6 +69,15 @@ return v; } +static inline iree_byte_span_t iree_byte_span_empty() { + iree_byte_span_t v = {NULL, 0}; + return v; +} + +static bool iree_byte_span_is_empty(iree_byte_span_t span) { + return span.data == NULL || span.data_length == 0; +} + // A span of constant bytes (ala std::span of const uint8_t). typedef struct iree_const_byte_span_t { const uint8_t* data; @@ -81,6 +90,15 @@ return v; } +static inline iree_const_byte_span_t iree_const_byte_span_empty() { + iree_const_byte_span_t v = {NULL, 0}; + return v; +} + +static bool iree_const_byte_span_is_empty(iree_const_byte_span_t span) { + return span.data == NULL || span.data_length == 0; +} + //===----------------------------------------------------------------------===// // Totally shady stack allocation //===----------------------------------------------------------------------===//
diff --git a/iree/base/internal/arena.c b/iree/base/internal/arena.c index 791b24a..81853d4 100644 --- a/iree/base/internal/arena.c +++ b/iree/base/internal/arena.c
@@ -18,22 +18,32 @@ void iree_arena_block_pool_initialize(iree_host_size_t total_block_size, iree_allocator_t block_allocator, iree_arena_block_pool_t* out_block_pool) { + IREE_TRACE_ZONE_BEGIN(z0); + memset(out_block_pool, 0, sizeof(*out_block_pool)); out_block_pool->total_block_size = total_block_size; out_block_pool->usable_block_size = total_block_size - sizeof(iree_arena_block_t); out_block_pool->block_allocator = block_allocator; iree_atomic_arena_block_slist_initialize(&out_block_pool->available_slist); + + IREE_TRACE_ZONE_END(z0); } void iree_arena_block_pool_deinitialize(iree_arena_block_pool_t* block_pool) { + IREE_TRACE_ZONE_BEGIN(z0); + // Since all blocks must have been released we can just reuse trim (today) as // it doesn't retain any blocks. iree_arena_block_pool_trim(block_pool); iree_atomic_arena_block_slist_deinitialize(&block_pool->available_slist); + + IREE_TRACE_ZONE_END(z0); } void iree_arena_block_pool_trim(iree_arena_block_pool_t* block_pool) { + IREE_TRACE_ZONE_BEGIN(z0); + iree_arena_block_t* head = NULL; iree_atomic_arena_block_slist_flush( &block_pool->available_slist, @@ -43,6 +53,8 @@ head = head->next; iree_allocator_free(block_pool->block_allocator, ptr); } + + IREE_TRACE_ZONE_END(z0); } iree_status_t iree_arena_block_pool_acquire(iree_arena_block_pool_t* block_pool, @@ -64,8 +76,7 @@ z0, iree_allocator_malloc_uninitialized(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 = (iree_arena_block_t*)(block_base + block_pool->usable_block_size); } block->next = NULL; @@ -99,6 +110,8 @@ } void iree_arena_reset(iree_arena_allocator_t* arena) { + IREE_TRACE_ZONE_BEGIN(z0); + if (arena->allocation_head != NULL) { iree_arena_oversized_allocation_t* head = arena->allocation_head; do { @@ -114,6 +127,8 @@ arena->block_head = NULL; arena->block_tail = NULL; } + + IREE_TRACE_ZONE_END(z0); } iree_status_t iree_arena_allocate(iree_arena_allocator_t* arena, @@ -127,16 +142,20 @@ // Oversized allocation that can't be handled by the block pool. We'll // allocate directly from the system allocator and track it ourselves for // freeing during reset. + IREE_TRACE_ZONE_BEGIN(z0); iree_host_size_t allocation_size = sizeof(iree_arena_oversized_allocation_t) + byte_length; iree_arena_oversized_allocation_t* allocation = NULL; - IREE_RETURN_IF_ERROR(iree_allocator_malloc_uninitialized( - block_pool->block_allocator, allocation_size, (void**)&allocation)); + IREE_RETURN_AND_END_ZONE_IF_ERROR( + z0, + iree_allocator_malloc_uninitialized( + block_pool->block_allocator, allocation_size, (void**)&allocation)); allocation->next = arena->allocation_head; arena->allocation_head = allocation; arena->total_allocation_size += allocation_size; arena->used_allocation_size += byte_length; *out_ptr = (uint8_t*)allocation + sizeof(iree_arena_oversized_allocation_t); + IREE_TRACE_ZONE_END(z0); return iree_ok_status(); } @@ -148,14 +167,16 @@ // Check to see if the current block (if any) has space - if not, get another. if (arena->block_head == NULL || arena->block_bytes_remaining < aligned_length) { + IREE_TRACE_ZONE_BEGIN(z0); iree_arena_block_t* block = NULL; - IREE_RETURN_IF_ERROR( - iree_arena_block_pool_acquire(arena->block_pool, &block)); + IREE_RETURN_AND_END_ZONE_IF_ERROR( + z0, iree_arena_block_pool_acquire(arena->block_pool, &block)); block->next = arena->block_head; arena->block_head = block; if (!arena->block_tail) arena->block_tail = block; arena->total_allocation_size += block_pool->total_block_size; arena->block_bytes_remaining = block_pool->usable_block_size; + IREE_TRACE_ZONE_END(z0); } // Slice out the allocation from the current block.
diff --git a/iree/base/status.c b/iree/base/status.c index 4942eb6..8fe228e 100644 --- a/iree/base/status.c +++ b/iree/base/status.c
@@ -525,11 +525,23 @@ return iree_ok_status(); } +IREE_API_EXPORT iree_status_t iree_status_join(iree_status_t base_status, + iree_status_t new_status) { + // TODO(benvanik): annotate |base_status| with |new_status| so we see it? + // This is intended for failure handling and usually the first failure is the + // root cause and most important to see. + if (!iree_status_is_ok(base_status)) { + iree_status_ignore(new_status); + return base_status; + } + return new_status; +} + IREE_API_EXPORT IREE_ATTRIBUTE_NORETURN void iree_status_abort( iree_status_t status) { + iree_status_fprint(stderr, status); IREE_ASSERT(!iree_status_is_ok(status), "only valid to call with failing status codes"); - iree_status_fprint(stderr, status); iree_status_free(status); abort(); }
diff --git a/iree/base/status.h b/iree/base/status.h index 81baad3..a26352c 100644 --- a/iree/base/status.h +++ b/iree/base/status.h
@@ -379,6 +379,16 @@ // Returns an OK status that can be used when chaining. IREE_API_EXPORT iree_status_t iree_status_ignore(iree_status_t status); +// Returns a new status that is |base_status| if not OK and otherwise returns +// |new_status|. This allows for chaining failure handling code that may also +// return statuses. +// +// Example: +// iree_status_t status = do_something(); +// return iree_status_join(status, do_cleanup()); +IREE_API_EXPORT iree_status_t iree_status_join(iree_status_t base_status, + iree_status_t new_status); + // Aborts the program with a failing |status|. // This will trigger a SIGABRT. It's best not to use this at all outside of // demos or tools.
diff --git a/iree/compiler/Bindings/Native/Transforms/WrapEntryPoints.cpp b/iree/compiler/Bindings/Native/Transforms/WrapEntryPoints.cpp index e2dcefd..a3873b8 100644 --- a/iree/compiler/Bindings/Native/Transforms/WrapEntryPoints.cpp +++ b/iree/compiler/Bindings/Native/Transforms/WrapEntryPoints.cpp
@@ -128,6 +128,25 @@ auto *entryBlock = wrapperFuncOp.addEntryBlock(); auto entryBuilder = OpBuilder::atBlockBegin(entryBlock); + // Build a map of result value to the argument that has its backing storage. + SmallVector<Value> resultStorages; + resultStorages.resize(resultTypes.size()); + for (unsigned i = 0; i < inputTypes.size(); ++i) { + auto outputAttr = + entryFuncOp.getArgAttrOfType<IntegerAttr>(i, "iree.abi.output"); + if (!outputAttr) continue; + // Today all outputs need to be a !hal.buffer - we could change this + // in the future to be something more generalized. + auto storageArg = entryBlock->getArgument(i); + if (!storageArg.getType().isa<IREE::HAL::BufferType>()) { + entryFuncOp.emitError() + << "storage argument " << i << " has an invalid type " + << storageArg.getType() << "; must be a !hal.buffer"; + return {}; + } + resultStorages[outputAttr.getInt()] = storageArg; + } + // Marshal arguments. SmallVector<Value> arguments; for (auto arg : llvm::enumerate(entryBlock->getArguments())) { @@ -152,8 +171,12 @@ auto oldType = entryFuncType.getResult(result.index()); auto newType = wrapperFuncType.getResult(result.index()); if (oldType.isa<TensorType>()) { - results.push_back(entryBuilder.createOrFold<IREE::HAL::TensorExportOp>( - entryFuncOp.getLoc(), newType, result.value())); + auto dynamicDims = IREE::Util::buildDynamicDimsForValue( + entryFuncOp.getLoc(), result.value(), entryBuilder); + results.push_back(entryBuilder.create<IREE::HAL::TensorExportOp>( + entryFuncOp.getLoc(), newType, result.value(), + TypeAttr::get(result.value().getType()), dynamicDims, + resultStorages[result.index()])); } else { results.push_back(result.value()); }
diff --git a/iree/compiler/Bindings/Native/Transforms/test/wrap_entry_points.mlir b/iree/compiler/Bindings/Native/Transforms/test/wrap_entry_points.mlir index 1762358..a56f0c3 100644 --- a/iree/compiler/Bindings/Native/Transforms/test/wrap_entry_points.mlir +++ b/iree/compiler/Bindings/Native/Transforms/test/wrap_entry_points.mlir
@@ -11,11 +11,11 @@ // CHECK-NEXT: %[[ARG0_TENSOR:.+]] = hal.tensor.import %[[ARG0]] : !hal.buffer_view -> tensor<?x8x8x3xf32>{%[[ARG0_DIM0]]} // CHECK-NEXT: %[[ARG1_DIM0:.+]] = hal.buffer_view.dim<%[[ARG1]] : !hal.buffer_view>[0] : index // CHECK-NEXT: %[[ARG1_TENSOR:.+]] = hal.tensor.import %[[ARG1]] : !hal.buffer_view -> tensor<?x8x8x3xf32>{%[[ARG1_DIM0]]} -// CHECK-NEXT: %[[RET_TENSOR:.+]]:2 = call @_dynamicEntry(%[[ARG0_TENSOR]], %[[ARG1_TENSOR]]) -// CHECK: %[[RET0_DIM0:.+]] = tensor.dim %[[RET_TENSOR]]#0, %c0{{.*}} : tensor<?x8x8x3xf32> -// CHECK-NEXT: %[[RET0_VIEW:.+]] = hal.tensor.export %[[RET_TENSOR]]#0 : tensor<?x8x8x3xf32>{%[[RET0_DIM0]]} -> !hal.buffer_view -// CHECK: %[[RET1_DIM0:.+]] = tensor.dim %[[RET_TENSOR]]#1, %c0{{.*}} : tensor<?x8x8x3xf32> -// CHECK-NEXT: %[[RET1_VIEW:.+]] = hal.tensor.export %[[RET_TENSOR]]#1 : tensor<?x8x8x3xf32>{%[[RET1_DIM0]]} -> !hal.buffer_view +// CHECK-NEXT: %[[RET_TENSORS:.+]]:2 = call @_dynamicEntry(%[[ARG0_TENSOR]], %[[ARG1_TENSOR]]) +// CHECK: %[[RET0_DIM0:.+]] = tensor.dim %[[RET_TENSORS]]#0, %c0{{.*}} : tensor<?x8x8x3xf32> +// CHECK-NEXT: %[[RET0_VIEW:.+]] = hal.tensor.export %[[RET_TENSORS]]#0 : tensor<?x8x8x3xf32>{%[[RET0_DIM0]]} -> !hal.buffer_view +// CHECK: %[[RET1_DIM0:.+]] = tensor.dim %[[RET_TENSORS]]#1, %c0{{.*}} : tensor<?x8x8x3xf32> +// CHECK-NEXT: %[[RET1_VIEW:.+]] = hal.tensor.export %[[RET_TENSORS]]#1 : tensor<?x8x8x3xf32>{%[[RET1_DIM0]]} -> !hal.buffer_view // CHECK-NEXT: return %[[RET0_VIEW]], %[[RET1_VIEW]] : !hal.buffer_view, !hal.buffer_view // CHECK-NEXT: } @@ -29,6 +29,34 @@ // ----- +// CHECK-LABEL: func @outputStorage( +// CHECK-SAME: %[[ARG0:.+]]: !hal.buffer_view, +// CHECK-SAME: %[[RET1_STORAGE:.+]]: !hal.buffer +// CHECK-SAME: -> ( +// CHECK-SAME: !hal.buffer_view, !hal.buffer_view +// CHECK-SAME: ) attributes { +// CHECK-SAME: iree.abi.stub +// CHECK-SAME: } { +// CHECK-NEXT: %[[ARG0_DIM0:.+]] = hal.buffer_view.dim<%[[ARG0]] : !hal.buffer_view>[0] : index +// CHECK-NEXT: %[[ARG0_TENSOR:.+]] = hal.tensor.import %[[ARG0]] : !hal.buffer_view -> tensor<?x8x8x3xf32>{%[[ARG0_DIM0]]} +// CHECK-NEXT: %[[RET_TENSORS:.+]]:2 = call @_outputStorage(%[[ARG0_TENSOR]], %[[RET1_STORAGE]]) +// CHECK: %[[RET0_DIM0:.+]] = tensor.dim %[[RET_TENSORS]]#0, %c0{{.*}} : tensor<?x8x8x3xf32> +// CHECK-NEXT: %[[RET0_VIEW:.+]] = hal.tensor.export %[[RET_TENSORS]]#0 : tensor<?x8x8x3xf32>{%[[RET0_DIM0]]} -> !hal.buffer_view +// CHECK: %[[RET1_DIM0:.+]] = tensor.dim %[[RET_TENSORS]]#1, %c0{{.*}} : tensor<?x8x8x3xf32> +// CHECK-NEXT: %[[RET1_VIEW:.+]] = hal.tensor.export %[[RET_TENSORS]]#1 into %[[RET1_STORAGE]] : tensor<?x8x8x3xf32>{%[[RET1_DIM0]]} -> !hal.buffer_view +// CHECK-NEXT: return %[[RET0_VIEW]], %[[RET1_VIEW]] : !hal.buffer_view, !hal.buffer_view +// CHECK-NEXT: } + +// CHECK-LABEL: func private @_outputStorage( +func @outputStorage(%arg0: tensor<?x8x8x3xf32>, %ret1: !hal.buffer {iree.abi.output = 1 : index}) -> + (tensor<?x8x8x3xf32>, tensor<?x8x8x3xf32>) { + %0 = "mhlo.add"(%arg0, %arg0) : (tensor<?x8x8x3xf32>, tensor<?x8x8x3xf32>) -> tensor<?x8x8x3xf32> + %1 = "mhlo.add"(%0, %arg0) : (tensor<?x8x8x3xf32>, tensor<?x8x8x3xf32>) -> tensor<?x8x8x3xf32> + return %0, %1 : tensor<?x8x8x3xf32>, tensor<?x8x8x3xf32> +} + +// ----- + // CHECK-LABEL: func @wrappedAlready // CHECK-SAME: (%arg0: !hal.buffer_view) -> !hal.buffer_view // CHECK-SAME: attributes {iree.abi.stub}
diff --git a/iree/compiler/Bindings/TFLite/Transforms/WrapEntryPoints.cpp b/iree/compiler/Bindings/TFLite/Transforms/WrapEntryPoints.cpp index 3ba6149..a74fa5c 100644 --- a/iree/compiler/Bindings/TFLite/Transforms/WrapEntryPoints.cpp +++ b/iree/compiler/Bindings/TFLite/Transforms/WrapEntryPoints.cpp
@@ -543,7 +543,7 @@ } callResults.push_back(entryBuilder.create<IREE::HAL::TensorExportOp>( result.getLoc(), bufferType, result, outputDynamicDims.tensorType, - dynamicDims)); + dynamicDims, /*target_storage=*/nullptr)); for (auto it : llvm::zip(dynamicDims, outputDynamicDims.globalOps)) { auto dynamicDim = std::get<0>(it); auto globalOp = std::get<1>(it);
diff --git a/iree/compiler/ConstEval/JitGlobals.cpp b/iree/compiler/ConstEval/JitGlobals.cpp index 08c8998..21947c91 100644 --- a/iree/compiler/ConstEval/JitGlobals.cpp +++ b/iree/compiler/ConstEval/JitGlobals.cpp
@@ -13,7 +13,7 @@ #include "iree/compiler/Dialect/Stream/Transforms/Passes.h" #include "iree/compiler/Dialect/Util/IR/UtilOps.h" #include "iree/compiler/Dialect/Util/Transforms/Passes.h" -#include "iree/compiler/Dialect/VM/Target/Bytecode/TranslationFlags.h" +#include "iree/compiler/Dialect/VM/Target/Bytecode/BytecodeModuleTarget.h" #include "iree/compiler/Dialect/VM/Transforms/Passes.h" #include "iree/compiler/Utils/PassUtils.h" #include "llvm/ADT/DenseSet.h"
diff --git a/iree/compiler/ConstEval/Runtime.cpp b/iree/compiler/ConstEval/Runtime.cpp index 4818679..cf6d155 100644 --- a/iree/compiler/ConstEval/Runtime.cpp +++ b/iree/compiler/ConstEval/Runtime.cpp
@@ -230,15 +230,17 @@ iree_hal_buffer_t* buffer = iree_hal_buffer_view_buffer(bufferView); // Map the memory and construct. - Attribute convertedAttr; + // TODO(benvanik): fallback to alloc + iree_hal_buffer_read_data if + // mapping is not available. Today with the CPU backends it's always + // possible but would not work with accelerators. iree_hal_buffer_mapping_t mapping; - IREE_CHECK_OK( - iree_hal_buffer_map_range(buffer, IREE_HAL_MEMORY_ACCESS_READ, - /*byte_offset=*/0, length, &mapping)); + IREE_CHECK_OK(iree_hal_buffer_map_range( + buffer, IREE_HAL_MAPPING_MODE_SCOPED, IREE_HAL_MEMORY_ACCESS_READ, + /*byte_offset=*/0, length, &mapping)); MutableArrayRef<char> rawBufferArray( reinterpret_cast<char*>(mapping.contents.data), mapping.contents.data_length); - convertedAttr = + auto convertedAttr = createAttributeFromRawData(loc, tensorType, rawBufferArray); iree_hal_buffer_unmap_range(&mapping); return convertedAttr; @@ -299,13 +301,17 @@ } Runtime::Runtime() { - registry = iree_hal_driver_registry_default(); + IREE_CHECK_OK( + iree_hal_driver_registry_allocate(iree_allocator_system(), ®istry)); IREE_CHECK_OK(iree_hal_vmvx_driver_module_register(registry)); IREE_CHECK_OK(iree_hal_module_register_types()); IREE_CHECK_OK(iree_vm_instance_create(iree_allocator_system(), &instance)); } -Runtime::~Runtime() { iree_vm_instance_release(instance); } +Runtime::~Runtime() { + iree_vm_instance_release(instance); + iree_hal_driver_registry_free(registry); +} Runtime& Runtime::getInstance() { static Runtime instance;
diff --git a/iree/compiler/Dialect/Flow/Transforms/BUILD b/iree/compiler/Dialect/Flow/Transforms/BUILD index e52a14d..b9f93b6 100644 --- a/iree/compiler/Dialect/Flow/Transforms/BUILD +++ b/iree/compiler/Dialect/Flow/Transforms/BUILD
@@ -31,6 +31,7 @@ cc_library( name = "Transforms", srcs = [ + "CleanupNumericNarrowing.cpp", "ConvertConv2D1x1ToMatmulPass.cpp", "ConvertConv2DToImg2ColPass.cpp", "ConvertLinalgMatmulToMmt4D.cpp", @@ -40,8 +41,10 @@ "DispatchLinalgOnTensors.cpp", "ExportBenchmarkFuncs.cpp", "FusionOfTensorOps.cpp", + "InferNumericNarrowing.cpp", "InjectDispatchTracing.cpp", "InterchangeGenericOps.cpp", + "OptimizeNumerics.cpp", "OutlineDispatchRegions.cpp", "PadLinalgOps.cpp", "PadTensorToSubTensorInsert.cpp", @@ -64,6 +67,9 @@ "//iree/compiler/Dialect/Flow/Conversion/TensorToFlow", "//iree/compiler/Dialect/Flow/IR", "//iree/compiler/Dialect/HAL/IR", + "//iree/compiler/Dialect/Util/Analysis", + "//iree/compiler/Dialect/Util/Analysis/Attributes", + "//iree/compiler/Dialect/Util/Analysis/DFX", "//iree/compiler/Dialect/Util/IR", "//iree/compiler/Dialect/Util/Transforms", "//iree/compiler/Utils", @@ -71,8 +77,10 @@ "//llvm-external-projects/iree-dialects:IREELinalgExtTransforms", "@llvm-project//llvm:Support", "@llvm-project//mlir:Affine", + "@llvm-project//mlir:ArithmeticDialect", "@llvm-project//mlir:DialectUtils", "@llvm-project//mlir:IR", + "@llvm-project//mlir:LinalgInterfaces", "@llvm-project//mlir:LinalgOps", "@llvm-project//mlir:LinalgTransforms", "@llvm-project//mlir:MemRefDialect",
diff --git a/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt b/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt index e381f09..60e371b 100644 --- a/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt +++ b/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
@@ -28,6 +28,7 @@ "Passes.h.inc" "TypeConverter.h" SRCS + "CleanupNumericNarrowing.cpp" "ConvertConv2D1x1ToMatmulPass.cpp" "ConvertConv2DToImg2ColPass.cpp" "ConvertLinalgMatmulToMmt4D.cpp" @@ -37,8 +38,10 @@ "DispatchLinalgOnTensors.cpp" "ExportBenchmarkFuncs.cpp" "FusionOfTensorOps.cpp" + "InferNumericNarrowing.cpp" "InjectDispatchTracing.cpp" "InterchangeGenericOps.cpp" + "OptimizeNumerics.cpp" "OutlineDispatchRegions.cpp" "PadLinalgOps.cpp" "PadTensorToSubTensorInsert.cpp" @@ -55,6 +58,7 @@ IREELinalgExtPasses LLVMSupport MLIRAffine + MLIRArithmetic MLIRIR MLIRLinalg MLIRLinalgTransforms @@ -71,6 +75,9 @@ iree::compiler::Dialect::Flow::Conversion::TensorToFlow iree::compiler::Dialect::Flow::IR iree::compiler::Dialect::HAL::IR + iree::compiler::Dialect::Util::Analysis + iree::compiler::Dialect::Util::Analysis::Attributes + iree::compiler::Dialect::Util::Analysis::DFX iree::compiler::Dialect::Util::IR iree::compiler::Dialect::Util::Transforms iree::compiler::Utils
diff --git a/iree/compiler/Dialect/Flow/Transforms/CleanupNumericNarrowing.cpp b/iree/compiler/Dialect/Flow/Transforms/CleanupNumericNarrowing.cpp new file mode 100644 index 0000000..4b200e8 --- /dev/null +++ b/iree/compiler/Dialect/Flow/Transforms/CleanupNumericNarrowing.cpp
@@ -0,0 +1,37 @@ +// Copyright 2021 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/compiler/Dialect/Flow/Transforms/PassDetail.h" +#include "iree/compiler/Dialect/Flow/Transforms/Passes.h" +#include "iree/compiler/Dialect/Util/IR/UtilOps.h" + +namespace mlir { +namespace iree_compiler { +namespace IREE { +namespace Flow { + +namespace { + +class CleanupNumericNarrowingPass + : public CleanupNumericNarrowingBase<CleanupNumericNarrowingPass> { + void runOnOperation() override { + getOperation()->walk([](IREE::Util::NumericOptionalNarrowOp op) { + op.getResult().replaceAllUsesWith(op.getOperand()); + op->erase(); + }); + } +}; + +} // namespace + +std::unique_ptr<Pass> createCleanupNumericNarrowingPass() { + return std::make_unique<CleanupNumericNarrowingPass>(); +} + +} // namespace Flow +} // namespace IREE +} // namespace iree_compiler +} // namespace mlir
diff --git a/iree/compiler/Dialect/Flow/Transforms/InferNumericNarrowing.cpp b/iree/compiler/Dialect/Flow/Transforms/InferNumericNarrowing.cpp new file mode 100644 index 0000000..1c3663d --- /dev/null +++ b/iree/compiler/Dialect/Flow/Transforms/InferNumericNarrowing.cpp
@@ -0,0 +1,142 @@ +// Copyright 2021 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/compiler/Dialect/Flow/Transforms/PassDetail.h" +#include "iree/compiler/Dialect/Flow/Transforms/Passes.h" +#include "iree/compiler/Dialect/Util/Analysis/Attributes/Range.h" +#include "iree/compiler/Dialect/Util/Analysis/DFX/Solver.h" +#include "iree/compiler/Dialect/Util/Analysis/DFX/State.h" +#include "iree/compiler/Dialect/Util/Analysis/Explorer.h" +#include "iree/compiler/Dialect/Util/IR/UtilDialect.h" +#include "iree/compiler/Dialect/Util/IR/UtilOps.h" +#include "llvm/ADT/SmallPtrSet.h" +#include "llvm/Support/Debug.h" +#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h" + +using llvm::SmallPtrSet; + +namespace mlir { +namespace iree_compiler { +namespace IREE { +namespace Flow { + +namespace { + +IntegerType deriveIntegerTypeFromRange(MLIRContext *context, int64_t minValue, + int64_t maxValue) { + // Clamp min/max to span 0. + const int64_t zero = 0; + minValue = std::min(zero, minValue); + maxValue = std::max(zero, maxValue); + bool isSigned; + if (minValue < 0) { + // For signed, make symmetric from -N:N-1 + isSigned = true; + maxValue = std::max(std::abs(minValue) - 1, maxValue); + minValue = std::min(-maxValue - 1, minValue); + } else { + isSigned = false; + } + int64_t n = maxValue - minValue + 1; + int64_t numBits = std::ceil(std::log2(n)); + + return IntegerType::get(context, numBits, + isSigned + ? IntegerType::SignednessSemantics::Signed + : IntegerType::SignednessSemantics::Unsigned); +} + +class InferNumericNarrowingPass + : public InferNumericNarrowingBase<InferNumericNarrowingPass> { + void getDependentDialects(DialectRegistry ®istry) const override { + registry.insert<IREE::Util::UtilDialect>(); + } + + void runOnOperation() override { + auto probePoints = collectProbePoints(); + + Explorer explorer(getOperation(), TraversalAction::SHALLOW); + llvm::BumpPtrAllocator allocator; + DFX::Solver solver(explorer, allocator); + + // Prime with probe points. + for (Value probePoint : probePoints) { + solver.getOrCreateElementFor<IREE::Util::FloatRangeValueElement>( + Position::forValue(probePoint)); + } + + // Solve. + if (failed(solver.run())) { + return signalPassFailure(); + } + + // Annotate. + for (Value probePoint : probePoints) { + auto *elt = solver.lookupElementFor<IREE::Util::FloatRangeValueElement>( + Position::forValue(probePoint)); + if (!elt) { + // Not valid analysis. + continue; + } + + applyAnnotation(probePoint, elt->getKnown()); + } + } + + SmallPtrSet<Value, 8> collectProbePoints() { + SmallPtrSet<Value, 8> probePoints; + getOperation()->walk([&](Operation *op) { + if (auto linalgOp = llvm::dyn_cast<linalg::LinalgOp>(op)) { + for (Value input : linalgOp.inputs()) { + probePoints.insert(input); + } + for (Value output : linalgOp.outputs()) { + probePoints.insert(output); + } + } + }); + return probePoints; + } + + void applyAnnotation(Value probePoint, IREE::Util::FloatRangeStats stats) { + if (stats.isTruncated() && stats.isFinite()) { + // Integer annotation. + applyIntegerAnnotation(probePoint, stats); + } + } + + void applyIntegerAnnotation(Value probePoint, + IREE::Util::FloatRangeStats stats) { + auto context = probePoint.getContext(); + auto minValue = static_cast<int64_t>(stats.minValue); + auto maxValue = static_cast<int64_t>(stats.maxValue); + IntegerType type = + deriveIntegerTypeFromRange(probePoint.getContext(), minValue, maxValue); + + // Insert the annotation. + OpBuilder builder(context); + builder.setInsertionPointAfterValue(probePoint); + Optional<std::pair<int64_t, int64_t>> range; + // i0 values cannot parse any values so omit. + if (type.getWidth() != 0) { + range = std::make_pair(minValue, maxValue); + } + auto annotationOp = builder.create<IREE::Util::NumericOptionalNarrowOp>( + probePoint.getLoc(), probePoint, type, range); + probePoint.replaceAllUsesExcept(annotationOp, annotationOp); + } +}; + +} // namespace + +std::unique_ptr<Pass> createInferNumericNarrowingPass() { + return std::make_unique<InferNumericNarrowingPass>(); +} + +} // namespace Flow +} // namespace IREE +} // namespace iree_compiler +} // namespace mlir
diff --git a/iree/compiler/Dialect/Flow/Transforms/OptimizeNumerics.cpp b/iree/compiler/Dialect/Flow/Transforms/OptimizeNumerics.cpp new file mode 100644 index 0000000..4c9ac77 --- /dev/null +++ b/iree/compiler/Dialect/Flow/Transforms/OptimizeNumerics.cpp
@@ -0,0 +1,283 @@ +// Copyright 2021 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/compiler/Dialect/Flow/Transforms/PassDetail.h" +#include "iree/compiler/Dialect/Flow/Transforms/Passes.h" +#include "iree/compiler/Dialect/Util/IR/UtilOps.h" +#include "llvm/Support/Debug.h" +#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" + +namespace mlir { +namespace iree_compiler { +namespace IREE { +namespace Flow { + +namespace { + +int getNextPotBitWidth(int bitWidth, int minBitWidth = 8) { + for (int i = minBitWidth;; i *= 2) { + if (i >= bitWidth) return i; + } +} + +Type withNewElementType(Type origType, Type elementType) { + if (auto st = origType.dyn_cast<ShapedType>()) { + return st.clone(elementType); + } else { + return elementType; + } +} + +Type makeLowPType(Type origType, int bitWidth) { + auto *context = origType.getContext(); + auto elementType = IntegerType::get(context, bitWidth); + return withNewElementType(origType, elementType); +} + +Value castNumeric(Value origValue, Type toType, bool isSigned, + OpBuilder &builder) { + Location loc = origValue.getLoc(); + Type origElementType = getElementTypeOrSelf(origValue.getType()); + Type toElementType = getElementTypeOrSelf(toType); + + if (origElementType.isa<FloatType>() && toElementType.isa<IntegerType>()) { + if (isSigned) { + return builder.create<arith::FPToSIOp>(loc, toType, origValue); + } else { + return builder.create<arith::FPToUIOp>(loc, toType, origValue); + } + } else if (origElementType.isa<IntegerType>() && + toElementType.isa<FloatType>()) { + if (isSigned) { + return builder.create<arith::SIToFPOp>(loc, toType, origValue); + } else { + return builder.create<arith::UIToFPOp>(loc, toType, origValue); + } + } else { + // If we need int<->int and float<->float, implement those cases. Since + // this is just needed for things in this file, it is ok to leave it + // under implemented. + llvm_unreachable("unsupported numeric cast"); + } +} + +struct NarrowParams { + static Optional<NarrowParams> forValue(Value value) { + if (auto narrowOp = + llvm::dyn_cast_or_null<IREE::Util::NumericOptionalNarrowOp>( + value.getDefiningOp())) { + NarrowParams params; + params.producer = narrowOp.operand(); + params.fromType = value.getType(); + params.toElementType = narrowOp.semantic_type(); + params.range = narrowOp.getIntegerRange(); + + return params; + } + return {}; + } + + bool isFromFloat() { return getElementTypeOrSelf(fromType).isa<FloatType>(); } + + bool isToInteger() { return toElementType.isa<IntegerType>(); } + + bool isToSigned() { return toElementType.cast<IntegerType>().isSigned(); } + + int getToBitWidth() { return toElementType.cast<IntegerType>().getWidth(); } + + Value producer; + Type fromType; + Type toElementType; + Optional<std::pair<int64_t, int64_t>> range; +}; + +// Eliminates a cast produced by an init_tensor by just initializing to that +// type directly. +struct LinalgInitTensorCast + : OpInterfaceRewritePattern<IREE::Util::NumericCastOpInterface> { + using OpInterfaceRewritePattern::OpInterfaceRewritePattern; + + LogicalResult matchAndRewrite(IREE::Util::NumericCastOpInterface castOp, + PatternRewriter &rewriter) const override { + auto initTensorOp = castOp.getInput().getDefiningOp<linalg::InitTensorOp>(); + if (!initTensorOp) return failure(); + Type resultType = castOp.getCasted().getType(); + + rewriter.replaceOpWithNewOp<linalg::InitTensorOp>( + castOp, resultType, initTensorOp.sizes(), initTensorOp.static_sizes()); + return success(); + } +}; + +// For a cast produced by a fill, rewrites the cast to be on the fill operands. +struct LinalgFillCast + : public OpInterfaceRewritePattern<IREE::Util::NumericCastOpInterface> { + using OpInterfaceRewritePattern::OpInterfaceRewritePattern; + + LogicalResult matchAndRewrite(IREE::Util::NumericCastOpInterface castOp, + PatternRewriter &rewriter) const override { + auto loc = castOp.getLoc(); + auto fillOp = castOp.getInput().getDefiningOp<linalg::FillOp>(); + if (!fillOp) return failure(); + Type toElementType = getElementTypeOrSelf(castOp.getCastedType()); + + Value fillInput = fillOp.value(); + Value fillInit = fillOp.output(); + fillInput = castOp + .cloneWithInput( + rewriter, + withNewElementType(fillInput.getType(), toElementType), + fillInput) + .getCasted(); + fillInit = + castOp + .cloneWithInput( + rewriter, withNewElementType(fillInit.getType(), toElementType), + fillInit) + .getCasted(); + Value fillResult = + rewriter.create<linalg::FillOp>(loc, fillInput, fillInit).result(); + rewriter.replaceOp(castOp, fillResult); + return success(); + } +}; + +// For narrowable inputs, selects +struct LinalgFpMatmulToLowP : public OpRewritePattern<linalg::MatmulOp> { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(linalg::MatmulOp matmulOp, + PatternRewriter &rewriter) const override { + Location loc = matmulOp.getLoc(); + Type origResultType = matmulOp.getResult(0).getType(); + auto lhsParams = NarrowParams::forValue(matmulOp.inputs()[0]); + auto rhsParams = NarrowParams::forValue(matmulOp.inputs()[1]); + auto accumParams = NarrowParams::forValue(matmulOp.outputs()[0]); + if (!lhsParams || !rhsParams || !accumParams) { + return rewriter.notifyMatchFailure(matmulOp, "no narrowing annotations"); + } + + // TODO(#7987): This could be more flexible, allowing mix and match + // integer/float types. + if (!lhsParams->isFromFloat() || !rhsParams->isFromFloat()) { + return rewriter.notifyMatchFailure(matmulOp, "not from floating point"); + } + + // TODO(#7987): Could support partial conversion to integer. + if (!lhsParams->isToInteger() || !rhsParams->isToInteger() || + !accumParams->isToInteger()) { + return rewriter.notifyMatchFailure(matmulOp, "not to an integer type"); + } + + int lhsBitWidth = lhsParams->getToBitWidth(); + int rhsBitWidth = rhsParams->getToBitWidth(); + + // Handle signed/unsigned mismatch. + // TODO(#7987): Implement a proper unsigned->signed widening. + bool isSigned; + if (lhsParams->isToSigned() != rhsParams->isToSigned()) { + // Mixed signed/unsigned. Promote to signed. + isSigned = true; + if (!lhsParams->isToSigned()) { + lhsBitWidth += 1; + } + if (!rhsParams->isToSigned()) { + rhsBitWidth += 1; + } + } else { + // Uniform signed/unsigned. + isSigned = lhsParams->isToSigned(); + } + + // Round up to a suitable POT width. + lhsBitWidth = getNextPotBitWidth(lhsBitWidth); + rhsBitWidth = getNextPotBitWidth(rhsBitWidth); + + // Promote accumulator to match signedness. + int accumBitWidth = accumParams->getToBitWidth(); + if (isSigned && !accumParams->isToSigned()) { + // TODO(#7987): A proper unsigned widening based on range. + accumBitWidth += 1; + } + + // Determine an appropriate accumulator size. + // TODO(#7987): Apply the clamp of: + // lhsBitWidth + rhsBitWidth + log2_ceil(contraction_dim + 1) to determine + // the accumulator size. Note: Can drop the +1 if one of lhs/rhs is signed + // and symmetric (i.e. does not use the asymmetric lower bound). + if (lhsBitWidth > 8 || rhsBitWidth > 8) { + return rewriter.notifyMatchFailure(matmulOp, "outside of low-p range"); + } + accumBitWidth = getNextPotBitWidth(accumBitWidth, 32); + if (accumBitWidth > 32) { + return rewriter.notifyMatchFailure(matmulOp, "accumulator > 32 bits"); + } + + Type lhsLowPType = makeLowPType(lhsParams->fromType, lhsBitWidth); + Type rhsLowPType = makeLowPType(rhsParams->fromType, rhsBitWidth); + Type accumLowPType = makeLowPType(accumParams->fromType, accumBitWidth); + + // Replace the matmul op. + Value newLhs = + castNumeric(lhsParams->producer, lhsLowPType, isSigned, rewriter); + Value newRhs = + castNumeric(rhsParams->producer, rhsLowPType, isSigned, rewriter); + Value newAccum = + castNumeric(accumParams->producer, accumLowPType, isSigned, rewriter); + Value newResult; + + if (isSigned) { + newResult = rewriter + .create<linalg::MatmulOp>(loc, ValueRange{newLhs, newRhs}, + ValueRange{newAccum}) + .getResult(0); + } else { + newResult = rewriter + .create<linalg::MatmulUnsignedOp>( + loc, ValueRange{newLhs, newRhs}, ValueRange{newAccum}) + .getResult(0); + } + + // Cast back. + newResult = castNumeric(newResult, origResultType, isSigned, rewriter); + rewriter.replaceOp(matmulOp, ValueRange{newResult}); + + return success(); + } +}; + +class OptimizeNumericsPass : public OptimizeNumericsBase<OptimizeNumericsPass> { + void runOnOperation() override { + MLIRContext *context = &getContext(); + RewritePatternSet patterns(context); + + // Precision reduction. + patterns.insert<LinalgFpMatmulToLowP>(context); + + // Cast propagation. + patterns.insert<LinalgInitTensorCast>(context); + patterns.insert<LinalgFillCast>(context); + + if (failed(applyPatternsAndFoldGreedily(getOperation(), + std::move(patterns)))) { + return signalPassFailure(); + } + } +}; + +} // namespace + +std::unique_ptr<Pass> createOptimizeNumericsPass() { + return std::make_unique<OptimizeNumericsPass>(); +} + +} // namespace Flow +} // namespace IREE +} // namespace iree_compiler +} // namespace mlir
diff --git a/iree/compiler/Dialect/Flow/Transforms/Passes.h b/iree/compiler/Dialect/Flow/Transforms/Passes.h index 340ae1d..fb9854a 100644 --- a/iree/compiler/Dialect/Flow/Transforms/Passes.h +++ b/iree/compiler/Dialect/Flow/Transforms/Passes.h
@@ -57,6 +57,10 @@ // Input canonicalization and legalization //===----------------------------------------------------------------------===// +// Cleans up any numeric narrowing ops inserted by +// iree-flow-infer-numeric-narrowing. +std::unique_ptr<Pass> createCleanupNumericNarrowingPass(); + /// Creates a pass to convert linalg convolution ops with 1x1 kernels into /// linalg.matmul std::unique_ptr<Pass> createConvertConv2D1x1ToMatmulPass(); @@ -76,6 +80,10 @@ /// Creates a pass to fuse Linalg operations on tensors. std::unique_ptr<Pass> createFusionOfTensorOpsPass(); +/// Infers and inserts util.numeric.optional_narrow ops at points that may be +/// beneficial. +std::unique_ptr<Pass> createInferNumericNarrowingPass(); + /// Create a pass to interchange generic ops to force the reduction loop to be /// the most inner loops. std::unique_ptr<Pass> createInterchangeGenericOpsPass(); @@ -87,6 +95,10 @@ // equivalent flow ops. std::unique_ptr<Pass> createConvertToFlowAfterDispatchFormation(); +// Optimizes numerics given annotations added via +// iree-flow-infer-numeric-narrowing. +std::unique_ptr<Pass> createOptimizeNumericsPass(); + // Promote I1 tensor constants to I8 tensors to match later operations. std::unique_ptr<OperationPass<mlir::FuncOp>> createPromoteI1ToI8Pass();
diff --git a/iree/compiler/Dialect/Flow/Transforms/Passes.td b/iree/compiler/Dialect/Flow/Transforms/Passes.td index 2d9644d..a1e0775 100644 --- a/iree/compiler/Dialect/Flow/Transforms/Passes.td +++ b/iree/compiler/Dialect/Flow/Transforms/Passes.td
@@ -9,6 +9,12 @@ include "mlir/Pass/PassBase.td" +def CleanupNumericNarrowing : + Pass<"iree-flow-cleanup-numeric-narrowing", ""> { + let summary = "Cleans up any numeric narrowing ops inserted by iree-flow-infer-numeric-narrowing"; + let constructor = "mlir::iree_compiler::IREE::Flow::createCleanupNumericNarrowingPass()"; +} + def ConvertConv2D1x1ConvToMatmul : Pass<"iree-flow-convert-conv2d-1x1-to-matmul", ""> { let summary = "Convert linalg convolution ops with 1x1 kernels into linalg matrix multiplication ops."; @@ -57,6 +63,12 @@ let constructor = "mlir::iree_compiler::IREE::Flow::createFusionOfTensorOpsPass()"; } +def InferNumericNarrowing : + Pass<"iree-flow-infer-numeric-narrowing", ""> { + let summary = "Infers and inserts util.numeric.optional_narrow ops at points that may be beneficial"; + let constructor = "mlir::iree_compiler::IREE::Flow::createInferNumericNarrowingPass()"; +} + def InjectDispatchTracing : Pass<"iree-flow-inject-dispatch-tracing", ""> { let summary = "Injects dispatch region tracing."; @@ -69,6 +81,12 @@ let constructor = "mlir::iree_compiler::IREE::Flow::createInterchangeGenericOpsPass()"; } +def OptimizeNumerics : + Pass<"iree-flow-optimize-numerics", ""> { + let summary = "Optimizes numerics given annotations added via iree-flow-infer-numeric-narrowing"; + let constructor = "mlir::iree_compiler::IREE::Flow::createOptimizeNumericsPass()"; +} + def OutlineDispatchRegions : Pass<"iree-flow-outline-dispatch-regions", "mlir::ModuleOp"> { let summary = "Outlines dispatch regions into executables";
diff --git a/iree/compiler/Dialect/Flow/Transforms/test/BUILD b/iree/compiler/Dialect/Flow/Transforms/test/BUILD index b66b655..6423883 100644 --- a/iree/compiler/Dialect/Flow/Transforms/test/BUILD +++ b/iree/compiler/Dialect/Flow/Transforms/test/BUILD
@@ -17,6 +17,7 @@ name = "lit", srcs = enforce_glob( [ + "cleanup_numeric_narrowing.mlir", "conv1x1_to_matmul.mlir", "conv2d_to_img2col.mlir", "convert_linalg_tensor_ops_after.mlir", @@ -26,9 +27,11 @@ "dispatch_linalg_on_tensors_elementwise.mlir", "dispatch_linalg_on_tensors_fusion.mlir", "export_benchmark_funcs.mlir", + "infer_numeric_narrowing.mlir", "inject_dispatch_tracing.mlir", "interchange_generic_ops.mlir", "matmul_to_mmt4d.mlir", + "optimize_numerics.mlir", "outline_dispatch_regions.mlir", "pad_linalg_ops.mlir", "pad_tensor_to_tensor.mlir",
diff --git a/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt b/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt index 205634c..9183bd0 100644 --- a/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt +++ b/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt
@@ -14,6 +14,7 @@ NAME lit SRCS + "cleanup_numeric_narrowing.mlir" "conv1x1_to_matmul.mlir" "conv2d_to_img2col.mlir" "convert_linalg_tensor_ops_after.mlir" @@ -23,9 +24,11 @@ "dispatch_linalg_on_tensors_elementwise.mlir" "dispatch_linalg_on_tensors_fusion.mlir" "export_benchmark_funcs.mlir" + "infer_numeric_narrowing.mlir" "inject_dispatch_tracing.mlir" "interchange_generic_ops.mlir" "matmul_to_mmt4d.mlir" + "optimize_numerics.mlir" "outline_dispatch_regions.mlir" "pad_linalg_ops.mlir" "pad_tensor_to_tensor.mlir"
diff --git a/iree/compiler/Dialect/Flow/Transforms/test/cleanup_numeric_narrowing.mlir b/iree/compiler/Dialect/Flow/Transforms/test/cleanup_numeric_narrowing.mlir new file mode 100644 index 0000000..2d2ea4b --- /dev/null +++ b/iree/compiler/Dialect/Flow/Transforms/test/cleanup_numeric_narrowing.mlir
@@ -0,0 +1,8 @@ +// RUN: iree-opt -iree-flow-cleanup-numeric-narrowing %s | IreeFileCheck %s + +// CHECK-LABEL: @remove_inferences +func @remove_inferences(%arg0 : tensor<5x3xf32>) -> tensor<5x3xf32> { + %0 = util.numeric.optional_narrow %arg0 : tensor<5x3xf32> as ui3 {max_value = 5 : ui3, min_value = 5 : ui3} + // CHECK: return %arg0 + return %0 : tensor<5x3xf32> +}
diff --git a/iree/compiler/Dialect/Flow/Transforms/test/infer_numeric_narrowing.mlir b/iree/compiler/Dialect/Flow/Transforms/test/infer_numeric_narrowing.mlir new file mode 100644 index 0000000..47bf850 --- /dev/null +++ b/iree/compiler/Dialect/Flow/Transforms/test/infer_numeric_narrowing.mlir
@@ -0,0 +1,75 @@ +// RUN: iree-opt -iree-flow-infer-numeric-narrowing %s | IreeFileCheck %s +// This does not test all of the analysis logic, just that the annotations +// are inserted at proper points in the right way. Probe points checked: +// - Every operand of a LinalgOp + +// CHECK-LABEL: @probe_linalg_op +// Checks as a by-product: +// - Infering ui0 for [0, 0] range +// - Infering unsigned for >= 0 range +func @probe_linalg_op(%arg0 : tensor<5x3xf32>) -> tensor<5x1xf32> { + // CHECK-DAG: %[[RHS:.*]] = arith.constant dense + // CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32 + // CHECK-DAG: util.numeric.optional_narrow %[[ZERO]] : f32 as ui0 + // CHECK-DAG: util.numeric.optional_narrow %[[RHS]] : tensor<3x1xf32> as ui7 {max_value = 127 : ui7, min_value = 0 : ui7} + // CHECK-DAG: %[[FILL:.*]] = linalg.fill + // CHECK-DAG: util.numeric.optional_narrow %[[FILL]] : tensor<5x1xf32> as ui0 + %rhs = arith.constant dense< + [[3.900000e+01], [0.000000e+00], [1.270000e+02]]> : tensor<3x1xf32> + %init_value = arith.constant 0.000000e+00 : f32 + %0 = linalg.init_tensor [5, 1] : tensor<5x1xf32> + %1 = linalg.fill(%init_value, %0) : f32, tensor<5x1xf32> -> tensor<5x1xf32> + %2 = linalg.matmul ins(%arg0, %rhs : tensor<5x3xf32>, tensor<3x1xf32>) outs(%1 : tensor<5x1xf32>) -> tensor<5x1xf32> + return %2 : tensor<5x1xf32> +} + +// CHECK-LABEL: @infer_symmetric_signed +// CHECK: util.numeric.optional_narrow %{{.*}} : tensor<3x1xf32> as si8 {max_value = 127 : si8, min_value = -39 : si8} +func @infer_symmetric_signed(%arg0 : tensor<5x3xf32>) -> tensor<5x1xf32> { + %rhs = arith.constant dense< + [[-3.900000e+01], [0.000000e+00], [1.270000e+02]]> : tensor<3x1xf32> + %init_value = arith.constant 0.000000e+00 : f32 + %0 = linalg.init_tensor [5, 1] : tensor<5x1xf32> + %1 = linalg.fill(%init_value, %0) : f32, tensor<5x1xf32> -> tensor<5x1xf32> + %2 = linalg.matmul ins(%arg0, %rhs : tensor<5x3xf32>, tensor<3x1xf32>) outs(%1 : tensor<5x1xf32>) -> tensor<5x1xf32> + return %2 : tensor<5x1xf32> +} + +// CHECK-LABEL: @infer_i1_signed +// Signed i1 is a silly boundary condition worth checking. +// CHECK: util.numeric.optional_narrow %{{.*}} : tensor<3x1xf32> as si1 {max_value = 0 : si1, min_value = -1 : si1} +func @infer_i1_signed(%arg0 : tensor<5x3xf32>) -> tensor<5x1xf32> { + %rhs = arith.constant dense< + [[0.000000e+00], [0.000000e+00], [-1.000000e+00]]> : tensor<3x1xf32> + %init_value = arith.constant 0.000000e+00 : f32 + %0 = linalg.init_tensor [5, 1] : tensor<5x1xf32> + %1 = linalg.fill(%init_value, %0) : f32, tensor<5x1xf32> -> tensor<5x1xf32> + %2 = linalg.matmul ins(%arg0, %rhs : tensor<5x3xf32>, tensor<3x1xf32>) outs(%1 : tensor<5x1xf32>) -> tensor<5x1xf32> + return %2 : tensor<5x1xf32> +} + +// CHECK-LABEL: @infer_positive_non_straddling_zero +// A range that does not straddle zero is a special case in the code. +// CHECK: util.numeric.optional_narrow %{{.*}} : tensor<3x1xf32> as ui2 {max_value = 2 : ui2, min_value = 1 : ui2} +func @infer_positive_non_straddling_zero(%arg0 : tensor<5x3xf32>) -> tensor<5x1xf32> { + %rhs = arith.constant dense< + [[1.000000e+00], [1.000000e+00], [2.000000e+00]]> : tensor<3x1xf32> + %init_value = arith.constant 0.000000e+00 : f32 + %0 = linalg.init_tensor [5, 1] : tensor<5x1xf32> + %1 = linalg.fill(%init_value, %0) : f32, tensor<5x1xf32> -> tensor<5x1xf32> + %2 = linalg.matmul ins(%arg0, %rhs : tensor<5x3xf32>, tensor<3x1xf32>) outs(%1 : tensor<5x1xf32>) -> tensor<5x1xf32> + return %2 : tensor<5x1xf32> +} + +// CHECK-LABEL: @infer_negative_non_straddling_zero +// A range that does not straddle zero is a special case in the code. +// CHECK: util.numeric.optional_narrow %{{.*}} : tensor<3x1xf32> as si2 {max_value = -1 : si2, min_value = -2 : si2} +func @infer_negative_non_straddling_zero(%arg0 : tensor<5x3xf32>) -> tensor<5x1xf32> { + %rhs = arith.constant dense< + [[-1.000000e+00], [-1.000000e+00], [-2.000000e+00]]> : tensor<3x1xf32> + %init_value = arith.constant 0.000000e+00 : f32 + %0 = linalg.init_tensor [5, 1] : tensor<5x1xf32> + %1 = linalg.fill(%init_value, %0) : f32, tensor<5x1xf32> -> tensor<5x1xf32> + %2 = linalg.matmul ins(%arg0, %rhs : tensor<5x3xf32>, tensor<3x1xf32>) outs(%1 : tensor<5x1xf32>) -> tensor<5x1xf32> + return %2 : tensor<5x1xf32> +}
diff --git a/iree/compiler/Dialect/Flow/Transforms/test/optimize_numerics.mlir b/iree/compiler/Dialect/Flow/Transforms/test/optimize_numerics.mlir new file mode 100644 index 0000000..cc8e368 --- /dev/null +++ b/iree/compiler/Dialect/Flow/Transforms/test/optimize_numerics.mlir
@@ -0,0 +1,77 @@ +// RUN: iree-opt -iree-flow-optimize-numerics %s | IreeFileCheck %s + +// CHECK-LABEL: @matmul_i8_i8_i32_unsigned +func @matmul_i8_i8_i32_unsigned(%arg0 : tensor<5x3xf32>, %arg1 : tensor<3x1xf32>, %arg2 : tensor<5x1xf32>) -> tensor<5x1xf32> { + // CHECK: %[[LHS:.*]] = arith.fptoui %arg0 : tensor<5x3xf32> to tensor<5x3xi8> + // CHECK: %[[RHS:.*]] = arith.fptoui %arg1 : tensor<3x1xf32> to tensor<3x1xi8> + // CHECK: %[[INIT:.*]] = arith.fptoui %arg2 : tensor<5x1xf32> to tensor<5x1xi32> + %lhs = util.numeric.optional_narrow %arg0 : tensor<5x3xf32> as ui7 {max_value = 127 : ui7, min_value = 0 : ui7} + %rhs = util.numeric.optional_narrow %arg1 : tensor<3x1xf32> as ui7 {max_value = 127 : ui7, min_value = 0 : ui7} + %init = util.numeric.optional_narrow %arg2 : tensor<5x1xf32> as ui0 + // CHECK: %[[RESULT:.*]] = linalg.matmul_unsigned ins(%[[LHS]], %[[RHS]] : tensor<5x3xi8>, tensor<3x1xi8>) outs(%[[INIT]] : tensor<5x1xi32>) + %2 = linalg.matmul ins(%lhs, %rhs : tensor<5x3xf32>, tensor<3x1xf32>) outs(%init : tensor<5x1xf32>) -> tensor<5x1xf32> + // CHECK: arith.uitofp %[[RESULT]] : tensor<5x1xi32> to tensor<5x1xf32> + return %2 : tensor<5x1xf32> +} + +// CHECK-LABEL: @matmul_i8_i8_i32_signed +func @matmul_i8_i8_i32_signed(%arg0 : tensor<5x3xf32>, %arg1 : tensor<3x1xf32>, %arg2 : tensor<5x1xf32>) -> tensor<5x1xf32> { + // CHECK: %[[LHS:.*]] = arith.fptosi %arg0 : tensor<5x3xf32> to tensor<5x3xi8> + // CHECK: %[[RHS:.*]] = arith.fptosi %arg1 : tensor<3x1xf32> to tensor<3x1xi8> + // CHECK: %[[INIT:.*]] = arith.fptosi %arg2 : tensor<5x1xf32> to tensor<5x1xi32> + %lhs = util.numeric.optional_narrow %arg0 : tensor<5x3xf32> as ui7 {max_value = 127 : ui7, min_value = 0 : ui7} + %rhs = util.numeric.optional_narrow %arg1 : tensor<3x1xf32> as si8 {max_value = 127 : si8, min_value = -127 : si8} + %init = util.numeric.optional_narrow %arg2 : tensor<5x1xf32> as ui0 + // CHECK: %[[RESULT:.*]] = linalg.matmul ins(%[[LHS]], %[[RHS]] : tensor<5x3xi8>, tensor<3x1xi8>) outs(%[[INIT]] : tensor<5x1xi32>) + %2 = linalg.matmul ins(%lhs, %rhs : tensor<5x3xf32>, tensor<3x1xf32>) outs(%init : tensor<5x1xf32>) -> tensor<5x1xf32> + // CHECK: arith.sitofp %[[RESULT]] : tensor<5x1xi32> to tensor<5x1xf32> + return %2 : tensor<5x1xf32> +} + +// CHECK-LABEL: @matmul_i4_i4_i32_signed +// For now we clamp this to i8 +func @matmul_i4_i4_i32_signed(%arg0 : tensor<5x3xf32>, %arg1 : tensor<3x1xf32>, %arg2 : tensor<5x1xf32>) -> tensor<5x1xf32> { + // CHECK: %[[LHS:.*]] = arith.fptosi %arg0 : tensor<5x3xf32> to tensor<5x3xi8> + // CHECK: %[[RHS:.*]] = arith.fptosi %arg1 : tensor<3x1xf32> to tensor<3x1xi8> + // CHECK: %[[INIT:.*]] = arith.fptosi %arg2 : tensor<5x1xf32> to tensor<5x1xi32> + %lhs = util.numeric.optional_narrow %arg0 : tensor<5x3xf32> as si4 {max_value = 7 : si4, min_value = -7 : si4} + %rhs = util.numeric.optional_narrow %arg1 : tensor<3x1xf32> as si4 {max_value = 3 : si4, min_value = -7 : si4} + %init = util.numeric.optional_narrow %arg2 : tensor<5x1xf32> as ui0 + // CHECK: %[[RESULT:.*]] = linalg.matmul ins(%[[LHS]], %[[RHS]] : tensor<5x3xi8>, tensor<3x1xi8>) outs(%[[INIT]] : tensor<5x1xi32>) + %2 = linalg.matmul ins(%lhs, %rhs : tensor<5x3xf32>, tensor<3x1xf32>) outs(%init : tensor<5x1xf32>) -> tensor<5x1xf32> + // CHECK: arith.sitofp %[[RESULT]] : tensor<5x1xi32> to tensor<5x1xf32> + return %2 : tensor<5x1xf32> +} + +// CHECK-LABEL: @matmul_reject_gt_8bit +// We may relax this restriction at some point but for right now we have it +// because less analysis is needed to prove safety. +// CHECK-NOT: fptosi +func @matmul_reject_gt_8bit(%arg0 : tensor<5x3xf32>, %arg1 : tensor<3x1xf32>, %arg2 : tensor<5x1xf32>) -> tensor<5x1xf32> { + %lhs = util.numeric.optional_narrow %arg0 : tensor<5x3xf32> as ui9 {max_value = 312 : ui9, min_value = 0 : ui9} + %rhs = util.numeric.optional_narrow %arg1 : tensor<3x1xf32> as si8 {max_value = 127 : si8, min_value = -127 : si8} + %init = util.numeric.optional_narrow %arg2 : tensor<5x1xf32> as ui0 + // CHECK: linalg.matmul {{.*}} -> tensor<5x1xf32> + %2 = linalg.matmul ins(%lhs, %rhs : tensor<5x3xf32>, tensor<3x1xf32>) outs(%init : tensor<5x1xf32>) -> tensor<5x1xf32> + return %2 : tensor<5x1xf32> +} + +// CHECK-LABEL: @cast_fill +func @cast_fill(%arg0 : f32, %arg1 : tensor<3xf32>) -> tensor<3xi8> { + // CHECK: %[[SCALAR:.*]] = arith.fptosi %arg0 : f32 to i8 + // CHECK: %[[INIT:.*]] = arith.fptosi %arg1 : tensor<3xf32> to tensor<3xi8> + // CHECK: %[[RESULT:.*]] = linalg.fill(%[[SCALAR]], %[[INIT]]) : i8, tensor<3xi8> -> tensor<3xi8> + // CHECK: return %[[RESULT]] + %0 = linalg.fill(%arg0, %arg1) : f32, tensor<3xf32> -> tensor<3xf32> + %1 = arith.fptosi %0 : tensor<3xf32> to tensor<3xi8> + return %1 : tensor<3xi8> +} + +// CHECK-LABEL: @cast_init +func @cast_init() -> tensor<5x9xi8> { + // CHECK: %[[RESULT:.*]] = linalg.init_tensor [5, 9] : tensor<5x9xi8> + // CHECK: return %[[RESULT]] + %0 = linalg.init_tensor [5, 9] : tensor<5x9xf32> + %1 = arith.fptosi %0 : tensor<5x9xf32> to tensor<5x9xi8> + return %1 : tensor<5x9xi8> +}
diff --git a/iree/compiler/Dialect/HAL/Conversion/HALToVM/ConvertHALToVM.cpp b/iree/compiler/Dialect/HAL/Conversion/HALToVM/ConvertHALToVM.cpp index d37e1cf..3f406a7 100644 --- a/iree/compiler/Dialect/HAL/Conversion/HALToVM/ConvertHALToVM.cpp +++ b/iree/compiler/Dialect/HAL/Conversion/HALToVM/ConvertHALToVM.cpp
@@ -138,7 +138,7 @@ } static PassRegistration<ConvertHALToVMPass> pass([] { - auto options = IREE::VM::getTargetOptionsFromFlags(); + auto options = IREE::VM::TargetOptions::FromFlags::get(); return std::make_unique<ConvertHALToVMPass>(options); });
diff --git a/iree/compiler/Dialect/HAL/Conversion/Passes.h b/iree/compiler/Dialect/HAL/Conversion/Passes.h index 8052a7f..92734e2 100644 --- a/iree/compiler/Dialect/HAL/Conversion/Passes.h +++ b/iree/compiler/Dialect/HAL/Conversion/Passes.h
@@ -17,7 +17,7 @@ IREE::VM::TargetOptions targetOptions); inline void registerHALConversionPasses() { - createConvertHALToVMPass(IREE::VM::getTargetOptionsFromFlags()); + createConvertHALToVMPass(IREE::VM::TargetOptions::FromFlags::get()); } } // namespace iree_compiler
diff --git a/iree/compiler/Dialect/HAL/IR/HALOps.cpp b/iree/compiler/Dialect/HAL/IR/HALOps.cpp index 76c96d2..7e6024c 100644 --- a/iree/compiler/Dialect/HAL/IR/HALOps.cpp +++ b/iree/compiler/Dialect/HAL/IR/HALOps.cpp
@@ -269,7 +269,7 @@ auto dynamicDims = IREE::Util::buildDynamicDimsForValue(result.location, source, builder); build(builder, result, resultType, source, TypeAttr::get(source.getType()), - dynamicDims); + dynamicDims, /*target_storage=*/nullptr); } Value TensorExportOp::getTiedResult(unsigned resultIndex) {
diff --git a/iree/compiler/Dialect/HAL/IR/HALOps.td b/iree/compiler/Dialect/HAL/IR/HALOps.td index 9d7d063..06df00b 100644 --- a/iree/compiler/Dialect/HAL/IR/HALOps.td +++ b/iree/compiler/Dialect/HAL/IR/HALOps.td
@@ -115,6 +115,7 @@ } def HAL_TensorExportOp : HAL_PureOp<"tensor.export", [ + AttrSizedOperandSegments, DeclareOpInterfaceMethods<Util_TiedOpInterface, [ "getTiedResult", "getTiedResultOperandIndex", @@ -135,19 +136,27 @@ dynamically shaped values must have the same number of dynamic dimensions. This allows for casting between rank-0 and rank-N types, different element types, etc. + + An optional `target_storage` buffer can be provided to hold the exported + result. The export will fail at runtime if the storage is null or if it has + insufficient capacity to store the output. The storage must be + device-visible and defined for transfer-target and dispatch usage. }]; let arguments = (ins AnyTensor:$source, TypeAttr:$source_encoding, - HAL_ShapeDynamicDims:$source_dims + HAL_ShapeDynamicDims:$source_dims, + Optional<HAL_Buffer>:$target_storage ); let results = (outs AnyTypeOf<[HAL_Buffer, HAL_BufferView]>:$target ); let assemblyFormat = [{ - $source `:` + $source + (`into` $target_storage^)? + `:` custom<TypeAlias>($source_encoding, type($source)) (`{` $source_dims^ `}`)? `->` type($target)
diff --git a/iree/compiler/Dialect/HAL/IR/test/tensor_ops.mlir b/iree/compiler/Dialect/HAL/IR/test/tensor_ops.mlir index 6498bc6..2baeaea 100644 --- a/iree/compiler/Dialect/HAL/IR/test/tensor_ops.mlir +++ b/iree/compiler/Dialect/HAL/IR/test/tensor_ops.mlir
@@ -24,3 +24,12 @@ %0 = hal.tensor.export %arg0 : tensor<?x3xf32> as tensor<?x3xi32>{%arg1} -> !hal.buffer_view return %0 : !hal.buffer_view } + +// ----- + +// CHECK-LABEL: @tensorExportInPlace +func @tensorExportInPlace(%arg0: tensor<?x3xi32>, %arg1 : index, %arg2: !hal.buffer) -> !hal.buffer_view { + // CHECK: hal.tensor.export %arg0 into %arg2 : tensor<?x3xf32> as tensor<?x3xi32>{%arg1} -> !hal.buffer_view + %0 = hal.tensor.export %arg0 into %arg2 : tensor<?x3xf32> as tensor<?x3xi32>{%arg1} -> !hal.buffer_view + return %0 : !hal.buffer_view +}
diff --git a/iree/compiler/Dialect/HAL/Target/BUILD b/iree/compiler/Dialect/HAL/Target/BUILD index f804469..3cae4a0 100644 --- a/iree/compiler/Dialect/HAL/Target/BUILD +++ b/iree/compiler/Dialect/HAL/Target/BUILD
@@ -25,6 +25,7 @@ "//iree/compiler/Dialect/HAL/IR", "//iree/compiler/Dialect/HAL/Utils", "//iree/compiler/Dialect/Util/IR", + "//iree/compiler/Utils", "@llvm-project//llvm:Support", "@llvm-project//mlir:IR", "@llvm-project//mlir:Pass",
diff --git a/iree/compiler/Dialect/HAL/Target/CMakeLists.txt b/iree/compiler/Dialect/HAL/Target/CMakeLists.txt index 2227647..658b678 100644 --- a/iree/compiler/Dialect/HAL/Target/CMakeLists.txt +++ b/iree/compiler/Dialect/HAL/Target/CMakeLists.txt
@@ -29,6 +29,7 @@ iree::compiler::Dialect::HAL::IR iree::compiler::Dialect::HAL::Utils iree::compiler::Dialect::Util::IR + iree::compiler::Utils PUBLIC )
diff --git a/iree/compiler/Dialect/HAL/Target/TargetBackend.cpp b/iree/compiler/Dialect/HAL/Target/TargetBackend.cpp index 7774efc..cc7f5c8 100644 --- a/iree/compiler/Dialect/HAL/Target/TargetBackend.cpp +++ b/iree/compiler/Dialect/HAL/Target/TargetBackend.cpp
@@ -17,7 +17,7 @@ namespace IREE { namespace HAL { -TargetOptions getTargetOptionsFromFlags() { +void TargetOptions::bindOptions(OptionsBinder &binder) { static llvm::cl::OptionCategory halTargetOptionsCategory( "IREE HAL executable target options"); @@ -25,15 +25,10 @@ // TranslateExecutablesPass. Pass registery is also staticly // initialized, so targetBackendsFlags needs to be here to be initialized // first. - static llvm::cl::list<std::string> *targetBackendsFlag = - new llvm::cl::list<std::string>{ - "iree-hal-target-backends", - llvm::cl::desc("Target backends for executable compilation"), - llvm::cl::ZeroOrMore, llvm::cl::cat(halTargetOptionsCategory)}; - - TargetOptions targetOptions; - targetOptions.targets = *targetBackendsFlag; - return targetOptions; + binder.list<std::string>( + "iree-hal-target-backends", targets, + llvm::cl::desc("Target backends for executable compilation"), + llvm::cl::ZeroOrMore, llvm::cl::cat(halTargetOptionsCategory)); } // Renames |op| within |moduleOp| with a new name that is unique within both
diff --git a/iree/compiler/Dialect/HAL/Target/TargetBackend.h b/iree/compiler/Dialect/HAL/Target/TargetBackend.h index 6caf416..5e17d0e 100644 --- a/iree/compiler/Dialect/HAL/Target/TargetBackend.h +++ b/iree/compiler/Dialect/HAL/Target/TargetBackend.h
@@ -14,6 +14,7 @@ #include "iree/compiler/Dialect/Flow/IR/FlowOps.h" #include "iree/compiler/Dialect/HAL/IR/HALOps.h" #include "iree/compiler/Dialect/HAL/Utils/DeviceSwitchBuilder.h" +#include "iree/compiler/Utils/OptionUtils.h" #include "llvm/ADT/StringMap.h" #include "llvm/ADT/StringRef.h" #include "mlir/IR/Dialect.h" @@ -35,12 +36,10 @@ // the best we can do is a coarse flag as to whether source maps should be // embedded, however we could be much better here on the TargetBackend // interface. + void bindOptions(OptionsBinder &binder); + using FromFlags = OptionsFromFlags<TargetOptions>; }; -// Returns a TargetOptions struct initialized with the -// --iree-hal-target-* flags. -TargetOptions getTargetOptionsFromFlags(); - // HAL executable target backend interface. // Multiple backends can be registered and targeted during a single compilation. // The flow->hal conversion process will use registered TargetBackend interfaces
diff --git a/iree/compiler/Dialect/HAL/Target/VMVX/VMVXTarget.cpp b/iree/compiler/Dialect/HAL/Target/VMVX/VMVXTarget.cpp index 8b3fa52..5d99ca2 100644 --- a/iree/compiler/Dialect/HAL/Target/VMVX/VMVXTarget.cpp +++ b/iree/compiler/Dialect/HAL/Target/VMVX/VMVXTarget.cpp
@@ -59,7 +59,7 @@ OpPassManager &nestedModulePM = passManager.nest<ModuleOp>(); // TODO(benvanik): derive these from a vm target triple. - auto vmOptions = IREE::VM::getTargetOptionsFromFlags(); + auto vmOptions = IREE::VM::TargetOptions::FromFlags::get(); vmOptions.f32Extension = true; vmOptions.optimizeForStackSize = false; IREE::VM::buildVMTransformPassPipeline(nestedModulePM, vmOptions);
diff --git a/iree/compiler/Dialect/HAL/Target/WebGPU/WebGPUTarget.cpp b/iree/compiler/Dialect/HAL/Target/WebGPU/WebGPUTarget.cpp index 0cf66f0..8fd908d 100644 --- a/iree/compiler/Dialect/HAL/Target/WebGPU/WebGPUTarget.cpp +++ b/iree/compiler/Dialect/HAL/Target/WebGPU/WebGPUTarget.cpp
@@ -124,8 +124,10 @@ SymbolTable::setSymbolName(entryPointFunc, symbolName); // We only have one shader module right now, so all point to index 0. - // TODO(#7824): Support multiple shader modules per executable - entryPointOrdinals[entryPointOp.ordinal().getZExtValue()] = 0; + // TODO(#7824): Support multiple shader modules per executable. + uint64_t ordinal = + entryPointOp.ordinal().getValueOr(APInt(64, 0)).getZExtValue(); + entryPointOrdinals[ordinal] = 0; } // Serialize the spirv::ModuleOp into binary format.
diff --git a/iree/compiler/Dialect/HAL/Transforms/MaterializeResourceCaches.cpp b/iree/compiler/Dialect/HAL/Transforms/MaterializeResourceCaches.cpp index 822e7fa..7ed1636 100644 --- a/iree/compiler/Dialect/HAL/Transforms/MaterializeResourceCaches.cpp +++ b/iree/compiler/Dialect/HAL/Transforms/MaterializeResourceCaches.cpp
@@ -294,7 +294,7 @@ } static PassRegistration<MaterializeResourceCachesPass> pass([] { - auto options = getTargetOptionsFromFlags(); + auto options = TargetOptions::FromFlags::get(); return std::make_unique<MaterializeResourceCachesPass>(options); });
diff --git a/iree/compiler/Dialect/HAL/Transforms/Passes.cpp b/iree/compiler/Dialect/HAL/Transforms/Passes.cpp index 0d77197..52e4515 100644 --- a/iree/compiler/Dialect/HAL/Transforms/Passes.cpp +++ b/iree/compiler/Dialect/HAL/Transforms/Passes.cpp
@@ -210,8 +210,8 @@ "iree-hal-transformation-pipeline", "Runs the full IREE HAL dialect transformation pipeline", [](OpPassManager &passManager, const TransformOptions &transformOptions) { - buildHALTransformPassPipeline(passManager, getTargetOptionsFromFlags(), - transformOptions); + buildHALTransformPassPipeline( + passManager, TargetOptions::FromFlags::get(), transformOptions); }); }
diff --git a/iree/compiler/Dialect/HAL/Transforms/Passes.h b/iree/compiler/Dialect/HAL/Transforms/Passes.h index 1b9eafd..8ab5e34 100644 --- a/iree/compiler/Dialect/HAL/Transforms/Passes.h +++ b/iree/compiler/Dialect/HAL/Transforms/Passes.h
@@ -134,7 +134,7 @@ inline void registerHALPasses() { registerHALTransformPassPipeline(); - auto targetOptions = getTargetOptionsFromFlags(); + auto targetOptions = TargetOptions::FromFlags::get(); createAssignTargetDevicesPass({}); createBenchmarkBatchDispatchesPass(/*repeatCount=*/1); createConvertToHALPass();
diff --git a/iree/compiler/Dialect/Stream/Conversion/HALToStream/ConvertHALToStream.cpp b/iree/compiler/Dialect/Stream/Conversion/HALToStream/ConvertHALToStream.cpp index fe93673..d8128c8 100644 --- a/iree/compiler/Dialect/Stream/Conversion/HALToStream/ConvertHALToStream.cpp +++ b/iree/compiler/Dialect/Stream/Conversion/HALToStream/ConvertHALToStream.cpp
@@ -141,18 +141,53 @@ auto externalType = rewriter.getType<IREE::Stream::ResourceType>( IREE::Stream::Lifetime::External); auto exportSource = adaptor.source(); - if (source.resource.getType() != externalType) { - exportSource = rewriter.create<IREE::Stream::AsyncTransferOp>( - op.getLoc(), externalType, source.resource, source.resourceSize, - source.resourceSize, - /*source_affinity=*/nullptr, - /*result_affinity=*/nullptr); + auto exportSize = source.resourceSize; + if (adaptor.target_storage()) { + // Query the target storage buffer length; we will only populate up to + // what is required for the output. + auto storageSize = + rewriter + .create<IREE::HAL::BufferLengthOp>( + op.getLoc(), rewriter.getIndexType(), op.target_storage()) + .result(); + + // Import the target storage as a resource that we can use as an update + // target. We overwrite the contents and just cast the storage to the + // target type so we know we can update it. + auto importOp = rewriter.create<IREE::Stream::TensorImportOp>( + op.getLoc(), externalType, adaptor.target_storage(), + TypeAttr::get(sourceType), adaptor.source_dims(), storageSize, + /*affinity=*/nullptr); + + // Copy the source value into the imported target storage. + auto zeroOffset = rewriter.create<arith::ConstantIndexOp>(op.getLoc(), 0); + auto updateOp = rewriter.create<IREE::Stream::AsyncUpdateOp>( + op.getLoc(), externalType, importOp.result(), importOp.result_size(), + zeroOffset, source.resourceSize, source.resource, source.resourceSize, + /*affinity=*/nullptr); + + // Export the updated resource. + // NOTE: the buffer size wrapped in the buffer view is the full size of + // the input buffer. This is so that we don't insert a data dependency on + // sparse operations or data-dependent dynamic shape dimensions. + exportSource = updateOp.result(); + exportSize = updateOp.target_size(); + } else { + // Exporting a produced value - transfer our source value to an externally + // usable resource and directly export it. This will cause an allocation. + if (source.resource.getType() != externalType) { + exportSource = rewriter.create<IREE::Stream::AsyncTransferOp>( + op.getLoc(), externalType, source.resource, source.resourceSize, + source.resourceSize, + /*source_affinity=*/nullptr, + /*result_affinity=*/nullptr); + } } // Export (stream resource to buffer view). rewriter.replaceOpWithNewOp<IREE::Stream::TensorExportOp>( op, targetType, exportSource, TypeAttr::get(sourceType), - adaptor.source_dims(), source.resourceSize, + adaptor.source_dims(), exportSize, /*affinity=*/nullptr); return success(); }
diff --git a/iree/compiler/Dialect/Stream/Conversion/HALToStream/test/abi_ops.mlir b/iree/compiler/Dialect/Stream/Conversion/HALToStream/test/abi_ops.mlir index d7d8849..5f8e181 100644 --- a/iree/compiler/Dialect/Stream/Conversion/HALToStream/test/abi_ops.mlir +++ b/iree/compiler/Dialect/Stream/Conversion/HALToStream/test/abi_ops.mlir
@@ -47,3 +47,21 @@ // CHECK: return %[[RESULT]] return %0 : !hal.buffer_view } + +// ----- + +// CHECK-LABEL: @exportBufferViewInPlace +// CHECK-SAME: (%[[TENSOR:.+]]: !stream.resource<*>, %[[SIZE:.+]]: index, %[[DIM0:.+]]: index, %[[DIM1:.+]]: index, %[[STORAGE:.+]]: !hal.buffer) +func @exportBufferViewInPlace(%tensor: tensor<?x?x4xf32>, %dim0: index, %dim1: index, %storage: !hal.buffer) -> !hal.buffer_view { + // CHECK: %[[STORAGE_LENGTH:.+]] = hal.buffer.length<%[[STORAGE]] + // CHECK-NEXT: %[[STORAGE_IMPORT:.+]] = stream.tensor.import %[[STORAGE]] + // CHECK-SAME: : !hal.buffer -> tensor<?x?x4xf32>{%[[DIM0]], %[[DIM1]]} in !stream.resource<external>{%[[STORAGE_LENGTH]]} + // CHECK-NEXT: %[[STORAGE_UPDATE:.+]] = stream.async.update %[[TENSOR]], %[[STORAGE_IMPORT]][%c0 to %[[SIZE]]] + // CHECK-SAME: : !stream.resource<*>{%[[SIZE]]} -> %[[STORAGE_IMPORT]] as !stream.resource<external>{%[[STORAGE_LENGTH]]} + // CHECK-NEXT: %[[STORAGE_RESULT:.+]] = stream.tensor.export %[[STORAGE_UPDATE]] : + // CHECK-SAME: tensor<?x?x4xf32>{%[[DIM0]], %[[DIM1]]} in !stream.resource<external>{%[[STORAGE_LENGTH]]} + // CHECK-SAME: -> !hal.buffer_view + %0 = hal.tensor.export %tensor into %storage : tensor<?x?x4xf32>{%dim0, %dim1} -> !hal.buffer_view + // CHECK: return %[[STORAGE_RESULT]] + return %0 : !hal.buffer_view +}
diff --git a/iree/compiler/Dialect/Util/Analysis/Attributes/Range.h b/iree/compiler/Dialect/Util/Analysis/Attributes/Range.h index aca1cb8..d8c640e 100644 --- a/iree/compiler/Dialect/Util/Analysis/Attributes/Range.h +++ b/iree/compiler/Dialect/Util/Analysis/Attributes/Range.h
@@ -56,6 +56,12 @@ return static_cast<TruncationFlag>(std::max(lhs, rhs)); } + // Whether the range is known to only contain values that have been + // truncated to exclude fractional bits. + bool isTruncated() { return truncationFlag == TRUNC; } + + bool isFinite() { return std::isfinite(minValue) && std::isfinite(maxValue); } + // Reset to initial state. void reset() { *this = FloatRangeStats(); }
diff --git a/iree/compiler/Dialect/Util/IR/UtilDialect.cpp b/iree/compiler/Dialect/Util/IR/UtilDialect.cpp index 1b3e2a4..58b4e7a 100644 --- a/iree/compiler/Dialect/Util/IR/UtilDialect.cpp +++ b/iree/compiler/Dialect/Util/IR/UtilDialect.cpp
@@ -143,6 +143,45 @@ results.insert<FoldDimOp<tensor::DimOp>>(getContext()); } +//===----------------------------------------------------------------------===// +// Interface external models +//===----------------------------------------------------------------------===// + +namespace { + +// Since all details of the interface are provided via default implementations, +// we can just have one templated external model to apply per op, vs one +// explicit model per op. +struct GenericNumericCastExternalModel { + template <typename OpTy> + struct ExternalModel + : public NumericCastOpInterface::ExternalModel<ExternalModel<OpTy>, + OpTy> {}; + + template <typename OpTy> + static void add(DialectRegistry ®istry) { + registry.addOpInterface<OpTy, ExternalModel<OpTy>>(); + } + + template <typename OpTy1, typename OpTy2, typename... More> + static void add(DialectRegistry ®istry) { + add<OpTy1>(registry); + add<OpTy2, More...>(registry); + } +}; + +} // namespace + +void registerUtilExternalModels(DialectRegistry ®istry) { + // Must ensure that any dependent dialects are registered. + registry.insert<arith::ArithmeticDialect>(); + + GenericNumericCastExternalModel::add< + arith::BitcastOp, arith::ExtFOp, arith::ExtUIOp, arith::ExtSIOp, + arith::FPToSIOp, arith::FPToUIOp, arith::IndexCastOp, arith::TruncFOp, + arith::TruncIOp, arith::SIToFPOp, arith::UIToFPOp>(registry); +} + } // namespace Util } // namespace IREE } // namespace iree_compiler
diff --git a/iree/compiler/Dialect/Util/IR/UtilDialect.h b/iree/compiler/Dialect/Util/IR/UtilDialect.h index 8a07a6e..8e08e28 100644 --- a/iree/compiler/Dialect/Util/IR/UtilDialect.h +++ b/iree/compiler/Dialect/Util/IR/UtilDialect.h
@@ -36,6 +36,8 @@ void registerTypes(); }; +void registerUtilExternalModels(DialectRegistry& registry); + } // namespace Util } // namespace IREE } // namespace iree_compiler
diff --git a/iree/compiler/Dialect/Util/IR/UtilInterfaces.td b/iree/compiler/Dialect/Util/IR/UtilInterfaces.td index c85926f..b415059 100644 --- a/iree/compiler/Dialect/Util/IR/UtilInterfaces.td +++ b/iree/compiler/Dialect/Util/IR/UtilInterfaces.td
@@ -118,6 +118,101 @@ } //===----------------------------------------------------------------------===// +// IREE::Util::NumericCastOpInterface +//===----------------------------------------------------------------------===// + +def Util_NumericCastOpInterface : OpInterface<"NumericCastOpInterface"> { + let cppNamespace = "::mlir::iree_compiler::IREE::Util"; + + let description = [{ + Applied to numeric casting ops which can convert between different numeric + types or shaped-types thereof. Example ops include `fptosi`, `trunci`, etc. + Treating these generically allows us to perform various cast movement + optimizations. + + Conforming operations must: + * Have no attributes. + * Have one operand and one result. + * Be able to operate on supported scalar types: IntegerType, FloatType, + IndexType. + * Be able to operate on tensors/vectors of supported scalar types. + * Have a builder that takes (Type, Value). + }]; + + let methods = [ + InterfaceMethod< + /*desc=*/[{ + Gets the input value. + }], + /*retTy=*/"Value", + /*methodName=*/"getInput", + /*args=*/(ins), + /*methodBody=*/[{}], + /*defaultImplementation=*/[{ + return $_op->getOperand(0); + }] + >, + + InterfaceMethod< + /*desc=*/[{ + Gets the input type. + }], + /*retTy=*/"Type", + /*methodName=*/"getInputType", + /*args=*/(ins), + /*methodBody=*/[{}], + /*defaultImplementation=*/[{ + return $_op->getOperand(0).getType(); + }] + >, + + InterfaceMethod< + /*desc=*/[{ + Gets the result casted value. + }], + /*retTy=*/"Value", + /*methodName=*/"getCasted", + /*args=*/(ins), + /*methodBody=*/[{}], + /*defaultImplementation=*/[{ + return $_op->getResult(0); + }] + >, + InterfaceMethod< + /*desc=*/[{ + Gets the result casted type. + }], + /*retTy=*/"Type", + /*methodName=*/"getCastedType", + /*args=*/(ins), + /*methodBody=*/[{}], + /*defaultImplementation=*/[{ + return $_op->getResult(0).getType(); + }] + >, + InterfaceMethod< + /*desc=*/[{ + Clones the operation with a new result type and input. + Note that it is generally not legal to generically change the scalar + type of the input or the result, but it is legal to transform between + a scalar type and a tensor/vector with an element type of the original + scalar type (and vica-versa). + }], + /*retTy=*/"NumericCastOpInterface", + /*methodName=*/"cloneWithInput", + /*args=*/(ins "OpBuilder &":$builder, "Type":$resultType, "Value":$input), + /*methodBody=*/[{}], + /*defaultImplementation=*/[{ + return llvm::cast<NumericCastOpInterface>( + builder.create<ConcreteOp>($_op->getLoc(), resultType, input) + .getOperation()); + }] + >, + ]; +} + + +//===----------------------------------------------------------------------===// // IREE::Util::TiedOpInterface //===----------------------------------------------------------------------===//
diff --git a/iree/compiler/Dialect/Util/IR/UtilOps.cpp b/iree/compiler/Dialect/Util/IR/UtilOps.cpp index f555f63..9e0b3ff 100644 --- a/iree/compiler/Dialect/Util/IR/UtilOps.cpp +++ b/iree/compiler/Dialect/Util/IR/UtilOps.cpp
@@ -663,6 +663,24 @@ } //===----------------------------------------------------------------------===// +// Numeric ops +//===----------------------------------------------------------------------===// + +Optional<std::pair<int64_t, int64_t>> +NumericOptionalNarrowOp::getIntegerRange() { + if (!min_value() || !max_value()) return {}; + bool signExtend = isSigned(); + // Note: Cannot sign extend 0 bit values. + int64_t minValue = signExtend && min_value()->getBitWidth() > 0 + ? min_value()->getSExtValue() + : min_value()->getZExtValue(); + int64_t maxValue = signExtend && max_value()->getBitWidth() > 0 + ? max_value()->getSExtValue() + : max_value()->getZExtValue(); + return std::make_pair(minValue, maxValue); +} + +//===----------------------------------------------------------------------===// // Structural ops //===----------------------------------------------------------------------===//
diff --git a/iree/compiler/Dialect/Util/IR/UtilOps.td b/iree/compiler/Dialect/Util/IR/UtilOps.td index 7b5131a..f3561de 100644 --- a/iree/compiler/Dialect/Util/IR/UtilOps.td +++ b/iree/compiler/Dialect/Util/IR/UtilOps.td
@@ -96,6 +96,81 @@ } //===----------------------------------------------------------------------===// +// Data type conversions +//===----------------------------------------------------------------------===// + +def Util_NumericOptionalNarrowOp : Util_PureOp<"numeric.optional_narrow", [ + SameOperandsAndResultType +]> { + let summary = "memorializes an optional numeric narrowing that is valid"; + let description = [{ + Serves as a placeholder for points in the computation where an optional + numeric narrowing can be performed without loss of information. Such ops + can guide optimization passes wishing to perform precision reduction. + + In addition to the operand and result type, this op takes an additional + `semantic_type` attribute representing the semantic target type which can + be: + * FloatType + * Signed IntegerType + * Unsigned IntegerType + + Note that this `semantic_type` must be a sign-carrying integer if using an + integer type and cannot be IndexType (i.e. it can be used to indicate a + possible narrowing of an IndexType to a specific integer). + + If the operand is a TensorType, then the result must be a TensorType. The + `semantic_type` constrains the element type. + + Optionally, the minimum and maximum integer values (for integer semantic + types) are tracked if known. + }]; + + let arguments = (ins + AnyTypeOf<[Util_Element, Util_Tensor]>:$operand, + TypeAttr:$semantic_type, + OptionalAttr<APIntAttr>:$min_value, + OptionalAttr<APIntAttr>:$max_value + ); + let results = (outs + AnyTypeOf<[Util_Element, Util_Tensor]>:$result + ); + + let assemblyFormat = [{ + $operand `:` type($operand) `as` $semantic_type attr-dict + }]; + + let builders = [ + OpBuilder<(ins + "Value":$operand, + "Type":$type, + "Optional<std::pair<int64_t, int64_t>>":$integerRange + ), + [{ + IntegerAttr minValueAttr; + IntegerAttr maxValueAttr; + if (integerRange) { + minValueAttr = $_builder.getIntegerAttr(type, integerRange->first); + maxValueAttr = $_builder.getIntegerAttr(type, integerRange->second); + } + build($_builder, $_state, operand.getType(), operand, TypeAttr::get(type), + minValueAttr, maxValueAttr); + }]>, + ]; + + let extraClassDeclaration = [{ + bool isSigned() { + if (auto integerType = getType().dyn_cast<IntegerType>()) { + return !integerType.isUnsigned(); + } + return true; + } + + Optional<std::pair<int64_t, int64_t>> getIntegerRange(); + }]; +} + +//===----------------------------------------------------------------------===// // Range arithmetic //===----------------------------------------------------------------------===//
diff --git a/iree/compiler/Dialect/Util/IR/test/BUILD b/iree/compiler/Dialect/Util/IR/test/BUILD index 457b1e6..84f17fb 100644 --- a/iree/compiler/Dialect/Util/IR/test/BUILD +++ b/iree/compiler/Dialect/Util/IR/test/BUILD
@@ -26,6 +26,7 @@ "hint_folding.mlir", "hint_ops.mlir", "list_ops.mlir", + "numeric_ops.mlir", "range_folding.mlir", "range_ops.mlir", "structural_folding.mlir",
diff --git a/iree/compiler/Dialect/Util/IR/test/CMakeLists.txt b/iree/compiler/Dialect/Util/IR/test/CMakeLists.txt index 3fa6522..c2f61f8 100644 --- a/iree/compiler/Dialect/Util/IR/test/CMakeLists.txt +++ b/iree/compiler/Dialect/Util/IR/test/CMakeLists.txt
@@ -23,6 +23,7 @@ "hint_folding.mlir" "hint_ops.mlir" "list_ops.mlir" + "numeric_ops.mlir" "range_folding.mlir" "range_ops.mlir" "structural_folding.mlir"
diff --git a/iree/compiler/Dialect/Util/IR/test/numeric_ops.mlir b/iree/compiler/Dialect/Util/IR/test/numeric_ops.mlir new file mode 100644 index 0000000..0481412 --- /dev/null +++ b/iree/compiler/Dialect/Util/IR/test/numeric_ops.mlir
@@ -0,0 +1,19 @@ +// RUN: iree-opt -split-input-file %s | IreeFileCheck %s + +func @optional_convert_scalar(%arg0 : i32) -> i32 { + // CHECK: util.numeric.optional_narrow %arg0 : i32 as si8 + %0 = util.numeric.optional_narrow %arg0 : i32 as si8 + return %0 : i32 +} + +func @optional_convert_tensor(%arg0 : tensor<f32>) -> tensor<f32> { + // CHECK: util.numeric.optional_narrow %arg0 : tensor<f32> as si8 + %0 = util.numeric.optional_narrow %arg0 : tensor<f32> as si8 + return %0 : tensor<f32> +} + +func @optional_convert_zero(%arg0 : i32) -> i32 { + // CHECK: util.numeric.optional_narrow %arg0 : i32 as ui0 + %0 = util.numeric.optional_narrow %arg0 : i32 as ui0 + return %0 : i32 +}
diff --git a/iree/compiler/Dialect/VM/Conversion/BUILD b/iree/compiler/Dialect/VM/Conversion/BUILD index 4b90b33..6d49265 100644 --- a/iree/compiler/Dialect/VM/Conversion/BUILD +++ b/iree/compiler/Dialect/VM/Conversion/BUILD
@@ -28,6 +28,7 @@ deps = [ "//iree/compiler/Dialect/Util/IR", "//iree/compiler/Dialect/VM/IR", + "//iree/compiler/Utils", "@llvm-project//llvm:Support", "@llvm-project//mlir:IR", "@llvm-project//mlir:Parser",
diff --git a/iree/compiler/Dialect/VM/Conversion/CMakeLists.txt b/iree/compiler/Dialect/VM/Conversion/CMakeLists.txt index e0f872a..3cb1337 100644 --- a/iree/compiler/Dialect/VM/Conversion/CMakeLists.txt +++ b/iree/compiler/Dialect/VM/Conversion/CMakeLists.txt
@@ -32,6 +32,7 @@ MLIRTransforms iree::compiler::Dialect::Util::IR iree::compiler::Dialect::VM::IR + iree::compiler::Utils PUBLIC )
diff --git a/iree/compiler/Dialect/VM/Conversion/StandardToVM/ConvertStandardToVMTest.cpp b/iree/compiler/Dialect/VM/Conversion/StandardToVM/ConvertStandardToVMTest.cpp index c3ee86c..dc755da 100644 --- a/iree/compiler/Dialect/VM/Conversion/StandardToVM/ConvertStandardToVMTest.cpp +++ b/iree/compiler/Dialect/VM/Conversion/StandardToVM/ConvertStandardToVMTest.cpp
@@ -40,7 +40,7 @@ mlir::arith::ArithmeticDialect>(); IREE::VM::TypeConverter typeConverter( - IREE::VM::getTargetOptionsFromFlags()); + IREE::VM::TargetOptions::FromFlags::get()); OwningRewritePatternList patterns(&getContext()); populateStandardToVMPatterns(&getContext(), typeConverter, patterns);
diff --git a/iree/compiler/Dialect/VM/Conversion/TargetOptions.cpp b/iree/compiler/Dialect/VM/Conversion/TargetOptions.cpp index 2f74c57..a27e3cd 100644 --- a/iree/compiler/Dialect/VM/Conversion/TargetOptions.cpp +++ b/iree/compiler/Dialect/VM/Conversion/TargetOptions.cpp
@@ -13,61 +13,35 @@ namespace IREE { namespace VM { -TargetOptions getTargetOptionsFromFlags() { +void TargetOptions::bindOptions(OptionsBinder &binder) { static llvm::cl::OptionCategory vmTargetOptionsCategory( "IREE VM target options"); - static auto *indexBitsFlag = new llvm::cl::opt<int>{ - "iree-vm-target-index-bits", - llvm::cl::init(32), - llvm::cl::desc("Bit width of index types."), - llvm::cl::cat(vmTargetOptionsCategory), - }; - static auto *i64ExtensionFlag = new llvm::cl::opt<bool>{ - "iree-vm-target-extension-i64", - llvm::cl::init(false), - llvm::cl::desc("Support i64 target opcode extensions."), - llvm::cl::cat(vmTargetOptionsCategory), - }; - static auto *f32ExtensionFlag = new llvm::cl::opt<bool>{ - "iree-vm-target-extension-f32", - llvm::cl::init(true), - llvm::cl::desc("Support f32 target opcode extensions."), - llvm::cl::cat(vmTargetOptionsCategory), - }; - static auto *f64ExtensionFlag = new llvm::cl::opt<bool>{ - "iree-vm-target-extension-f64", - llvm::cl::init(false), - llvm::cl::desc("Support f64 target opcode extensions."), - llvm::cl::cat(vmTargetOptionsCategory), - }; - static auto *truncateUnsupportedIntegersFlag = new llvm::cl::opt<bool>{ - "iree-vm-target-truncate-unsupported-integers", - llvm::cl::init(true), - llvm::cl::desc("Truncate i64 to i32 when unsupported."), - llvm::cl::cat(vmTargetOptionsCategory), - }; - static auto *truncateUnsupportedFloatsFlag = new llvm::cl::opt<bool>{ - "iree-vm-target-truncate-unsupported-floats", - llvm::cl::init(true), - llvm::cl::desc("Truncate f64 to f32 when unsupported."), - llvm::cl::cat(vmTargetOptionsCategory), - }; - - TargetOptions targetOptions; - targetOptions.indexBits = *indexBitsFlag; - if (*i64ExtensionFlag) { - targetOptions.i64Extension = true; - } - if (*f32ExtensionFlag) { - targetOptions.f32Extension = true; - } - if (*f64ExtensionFlag) { - targetOptions.f64Extension = true; - } - targetOptions.truncateUnsupportedIntegers = *truncateUnsupportedIntegersFlag; - targetOptions.truncateUnsupportedFloats = *truncateUnsupportedFloatsFlag; - return targetOptions; + binder.opt<int>("iree-vm-target-index-bits", indexBits, + llvm::cl::desc("Bit width of index types."), + llvm::cl::cat(vmTargetOptionsCategory)); + binder.opt<bool>("iree-vm-target-extension-i64", i64Extension, + llvm::cl::desc("Support i64 target opcode extensions."), + llvm::cl::cat(vmTargetOptionsCategory)); + binder.opt<bool>("iree-vm-target-extension-f32", f32Extension, + llvm::cl::desc("Support f32 target opcode extensions."), + llvm::cl::cat(vmTargetOptionsCategory)); + binder.opt<bool>("iree-vm-target-extension-f64", f64Extension, + llvm::cl::desc("Support f64 target opcode extensions."), + llvm::cl::cat(vmTargetOptionsCategory)); + binder.opt<bool>("iree-vm-target-truncate-unsupported-integers", + truncateUnsupportedIntegers, + llvm::cl::desc("Truncate i64 to i32 when unsupported."), + llvm::cl::cat(vmTargetOptionsCategory)); + binder.opt<bool>("iree-vm-target-truncate-unsupported-floats", + truncateUnsupportedFloats, + llvm::cl::desc("Truncate f64 to f32 when unsupported."), + llvm::cl::cat(vmTargetOptionsCategory)); + binder.opt<bool>( + "iree-vm-target-optimize-for-stack-size", optimizeForStackSize, + llvm::cl::desc( + "Prefer optimizations that reduce VM stack usage over performance."), + llvm::cl::cat(vmTargetOptionsCategory)); } } // namespace VM
diff --git a/iree/compiler/Dialect/VM/Conversion/TargetOptions.h b/iree/compiler/Dialect/VM/Conversion/TargetOptions.h index 5e02b3f..82d3f97 100644 --- a/iree/compiler/Dialect/VM/Conversion/TargetOptions.h +++ b/iree/compiler/Dialect/VM/Conversion/TargetOptions.h
@@ -7,6 +7,7 @@ #ifndef IREE_COMPILER_DIALECT_VM_CONVERSION_TARGETOPTIONS_H_ #define IREE_COMPILER_DIALECT_VM_CONVERSION_TARGETOPTIONS_H_ +#include "iree/compiler/Utils/OptionUtils.h" #include "mlir/Transforms/DialectConversion.h" namespace mlir { @@ -32,7 +33,7 @@ // Whether the i64 extension is enabled in the target VM. bool i64Extension = false; // Whether the f32 extension is enabled in the target VM. - bool f32Extension = false; + bool f32Extension = true; // Whether the f64 extension is enabled in the target VM. bool f64Extension = false; @@ -45,11 +46,10 @@ // Prefer optimizations that reduce VM stack usage over performance. bool optimizeForStackSize = true; -}; -// Returns a TargetOptions struct initialized with the -// --iree-vm-target-* flags. -TargetOptions getTargetOptionsFromFlags(); + void bindOptions(OptionsBinder &binder); + using FromFlags = OptionsFromFlags<TargetOptions>; +}; } // namespace VM } // namespace IREE
diff --git a/iree/compiler/Dialect/VM/Target/Bytecode/BUILD b/iree/compiler/Dialect/VM/Target/Bytecode/BUILD index a62850a..3914d3c 100644 --- a/iree/compiler/Dialect/VM/Target/Bytecode/BUILD +++ b/iree/compiler/Dialect/VM/Target/Bytecode/BUILD
@@ -12,12 +12,10 @@ "BytecodeModuleTarget.cpp", "DebugDatabaseBuilder.cpp", "DebugDatabaseBuilder.h", - "TranslationFlags.cpp", "TranslationRegistration.cpp", ], hdrs = [ "BytecodeModuleTarget.h", - "TranslationFlags.h", ], deps = [ "//iree/compiler/Dialect/Util/IR",
diff --git a/iree/compiler/Dialect/VM/Target/Bytecode/BytecodeModuleTarget.cpp b/iree/compiler/Dialect/VM/Target/Bytecode/BytecodeModuleTarget.cpp index a4ee999..0f2cc96 100644 --- a/iree/compiler/Dialect/VM/Target/Bytecode/BytecodeModuleTarget.cpp +++ b/iree/compiler/Dialect/VM/Target/Bytecode/BytecodeModuleTarget.cpp
@@ -543,6 +543,7 @@ static LogicalResult buildFlatBufferModule(BytecodeTargetOptions targetOptions, IREE::VM::ModuleOp moduleOp, SmallVector<ZIPFileRef> &zipFileRefs, + bool emitPolyglotZip, FlatbufferBuilder &fbb) { // Start the buffer so that we can begin recording data prior to the root // table (which we do at the very end). This does not change the layout of the @@ -603,8 +604,7 @@ for (auto rodataOp : llvm::reverse(rodataOps)) { // Only include rodata entries in the ZIP if they are file-like. This // prevents all of our string tables from getting included. - bool includeInZIP = - targetOptions.emitPolyglotZip && rodataOp.mime_type().hasValue(); + bool includeInZIP = emitPolyglotZip && rodataOp.mime_type().hasValue(); // Embed the rodata contents. size_t alignment = @@ -774,6 +774,9 @@ LogicalResult translateModuleToBytecode(IREE::VM::ModuleOp moduleOp, BytecodeTargetOptions targetOptions, llvm::raw_ostream &output) { + bool emitPolyglotZip = + targetOptions.emitPolyglotZip && + targetOptions.outputFormat == BytecodeOutputFormat::kFlatBufferBinary; moduleOp.getContext()->getOrLoadDialect<IREE::Util::UtilDialect>(); uint64_t startOffset = output.tell(); @@ -821,8 +824,8 @@ // can be large bulk data. FlatbufferBuilder fbb; SmallVector<ZIPFileRef> zipFileRefs; - if (failed( - buildFlatBufferModule(targetOptions, moduleOp, zipFileRefs, fbb))) { + if (failed(buildFlatBufferModule(targetOptions, moduleOp, zipFileRefs, + emitPolyglotZip, fbb))) { return moduleOp.emitError() << "failed to build FlatBuffer BytecodeModuleDef"; } @@ -852,7 +855,7 @@ } output.flush(); - if (targetOptions.emitPolyglotZip) { + if (emitPolyglotZip) { // Append the ZIP central directory to the end of the output. // We have to do this here as we need to have flushed the flatbuffer // contents to the output so that we have their final absolute addresses. @@ -875,6 +878,48 @@ return translateModuleToBytecode(*moduleOps.begin(), targetOptions, output); } +void BytecodeTargetOptions::bindOptions(OptionsBinder &binder) { + static llvm::cl::OptionCategory vmBytecodeOptionsCategory( + "IREE VM bytecode options"); + + binder.opt<BytecodeOutputFormat>( + "iree-vm-bytecode-module-output-format", outputFormat, + llvm::cl::cat(vmBytecodeOptionsCategory), + llvm::cl::desc("Output format the bytecode module is written in"), + llvm::cl::values( + clEnumValN(BytecodeOutputFormat::kFlatBufferBinary, + "flatbuffer-binary", "Binary FlatBuffer file"), + clEnumValN(BytecodeOutputFormat::kFlatBufferText, "flatbuffer-text", + "Text FlatBuffer file, debug-only"), + clEnumValN(BytecodeOutputFormat::kMlirText, "mlir-text", + "MLIR module file in the VM dialect"), + clEnumValN(BytecodeOutputFormat::kAnnotatedMlirText, + "annotated-mlir-text", + "MLIR module file in the VM dialect with annotations"))); + binder.opt<bool>( + "iree-vm-bytecode-module-optimize", optimize, + llvm::cl::cat(vmBytecodeOptionsCategory), + llvm::cl::desc("Optimizes the VM module with CSE/inlining/etc prior to " + "serialization")); + binder.opt<std::string>( + "iree-vm-bytecode-source-listing", sourceListing, + llvm::cl::cat(vmBytecodeOptionsCategory), + llvm::cl::desc( + "Dump a VM MLIR file and annotate source locations with it")); + binder.opt<bool>("iree-vm-bytecode-module-strip-source-map", stripSourceMap, + llvm::cl::cat(vmBytecodeOptionsCategory), + llvm::cl::desc("Strips the source map from the module")); + binder.opt<bool>("iree-vm-bytecode-module-strip-debug-ops", stripDebugOps, + llvm::cl::cat(vmBytecodeOptionsCategory), + llvm::cl::desc("Strips debug-only ops from the module")); + binder.opt<bool>( + "iree-vm-emit-polyglot-zip", emitPolyglotZip, + llvm::cl::cat(vmBytecodeOptionsCategory), + llvm::cl::desc( + "Enables output files to be viewed as zip files for debugging " + "(only applies to binary targets)")); +} + } // namespace VM } // namespace IREE } // namespace iree_compiler
diff --git a/iree/compiler/Dialect/VM/Target/Bytecode/BytecodeModuleTarget.h b/iree/compiler/Dialect/VM/Target/Bytecode/BytecodeModuleTarget.h index 57b31dd..8ef8a92 100644 --- a/iree/compiler/Dialect/VM/Target/Bytecode/BytecodeModuleTarget.h +++ b/iree/compiler/Dialect/VM/Target/Bytecode/BytecodeModuleTarget.h
@@ -8,6 +8,7 @@ #define IREE_COMPILER_DIALECT_VM_TARGET_BYTECODE_BYTECODEMODULETARGET_H_ #include "iree/compiler/Dialect/VM/IR/VMOps.h" +#include "iree/compiler/Utils/OptionUtils.h" #include "llvm/Support/raw_ostream.h" #include "mlir/IR/BuiltinOps.h" #include "mlir/Support/LogicalResult.h" @@ -52,6 +53,9 @@ // Enables the output .vmfb to be inspected as a ZIP file. // This is only useful for debugging and should be disabled otherwise. bool emitPolyglotZip = false; + + void bindOptions(OptionsBinder &binder); + using FromFlags = OptionsFromFlags<BytecodeTargetOptions>; }; // Translates a vm.module to a bytecode module flatbuffer.
diff --git a/iree/compiler/Dialect/VM/Target/Bytecode/CMakeLists.txt b/iree/compiler/Dialect/VM/Target/Bytecode/CMakeLists.txt index 1315acf..fcd6fb7 100644 --- a/iree/compiler/Dialect/VM/Target/Bytecode/CMakeLists.txt +++ b/iree/compiler/Dialect/VM/Target/Bytecode/CMakeLists.txt
@@ -15,14 +15,12 @@ Bytecode HDRS "BytecodeModuleTarget.h" - "TranslationFlags.h" SRCS "BytecodeEncoder.cpp" "BytecodeEncoder.h" "BytecodeModuleTarget.cpp" "DebugDatabaseBuilder.cpp" "DebugDatabaseBuilder.h" - "TranslationFlags.cpp" "TranslationRegistration.cpp" DEPS LLVMSupport
diff --git a/iree/compiler/Dialect/VM/Target/Bytecode/TranslationFlags.cpp b/iree/compiler/Dialect/VM/Target/Bytecode/TranslationFlags.cpp deleted file mode 100644 index 815bfc9..0000000 --- a/iree/compiler/Dialect/VM/Target/Bytecode/TranslationFlags.cpp +++ /dev/null
@@ -1,86 +0,0 @@ -// Copyright 2019 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/compiler/Dialect/VM/Target/Bytecode/TranslationFlags.h" - -#include "llvm/Support/CommandLine.h" - -namespace mlir { -namespace iree_compiler { -namespace IREE { -namespace VM { - -// TODO(b/145140345): fix LLVM category registration with ASAN. -// static llvm::cl::OptionCategory vmBytecodeOptionsCategory( -// "IREE VM bytecode options"); - -static llvm::cl::opt<BytecodeOutputFormat> outputFormatFlag{ - "iree-vm-bytecode-module-output-format", - llvm::cl::desc("Output format the bytecode module is written in"), - llvm::cl::init(BytecodeOutputFormat::kFlatBufferBinary), - llvm::cl::values( - clEnumValN(BytecodeOutputFormat::kFlatBufferBinary, "flatbuffer-binary", - "Binary FlatBuffer file"), - clEnumValN(BytecodeOutputFormat::kFlatBufferText, "flatbuffer-text", - "Text FlatBuffer file, debug-only"), - clEnumValN(BytecodeOutputFormat::kMlirText, "mlir-text", - "MLIR module file in the VM dialect"), - clEnumValN(BytecodeOutputFormat::kAnnotatedMlirText, - "annotated-mlir-text", - "MLIR module file in the VM dialect with annotations")), -}; - -static llvm::cl::opt<bool> optimizeFlag{ - "iree-vm-bytecode-module-optimize", - llvm::cl::desc( - "Optimizes the VM module with CSE/inlining/etc prior to serialization"), - llvm::cl::init(true), -}; - -static llvm::cl::opt<std::string> sourceListingFlag{ - "iree-vm-bytecode-source-listing", - llvm::cl::desc("Dump a VM MLIR file and annotate source locations with it"), - llvm::cl::init(""), -}; - -static llvm::cl::opt<bool> stripSourceMapFlag{ - "iree-vm-bytecode-module-strip-source-map", - llvm::cl::desc("Strips the source map from the module"), - llvm::cl::init(false), -}; - -static llvm::cl::opt<bool> stripDebugOpsFlag{ - "iree-vm-bytecode-module-strip-debug-ops", - llvm::cl::desc("Strips debug-only ops from the module"), - llvm::cl::init(false), -}; - -static llvm::cl::opt<bool> emitPolyglotZipFlag{ - "iree-vm-emit-polyglot-zip", - llvm::cl::desc( - "Enables output files to be viewed as zip files for debugging"), - llvm::cl::init(true), -}; - -BytecodeTargetOptions getBytecodeTargetOptionsFromFlags() { - BytecodeTargetOptions targetOptions; - targetOptions.outputFormat = outputFormatFlag; - targetOptions.optimize = optimizeFlag; - targetOptions.sourceListing = sourceListingFlag; - targetOptions.stripSourceMap = stripSourceMapFlag; - targetOptions.stripDebugOps = stripDebugOpsFlag; - targetOptions.emitPolyglotZip = emitPolyglotZipFlag; - if (outputFormatFlag != BytecodeOutputFormat::kFlatBufferBinary) { - // Only allow binary output formats to also be .zip files. - targetOptions.emitPolyglotZip = false; - } - return targetOptions; -} - -} // namespace VM -} // namespace IREE -} // namespace iree_compiler -} // namespace mlir
diff --git a/iree/compiler/Dialect/VM/Target/Bytecode/TranslationFlags.h b/iree/compiler/Dialect/VM/Target/Bytecode/TranslationFlags.h deleted file mode 100644 index 1835dab..0000000 --- a/iree/compiler/Dialect/VM/Target/Bytecode/TranslationFlags.h +++ /dev/null
@@ -1,26 +0,0 @@ -// Copyright 2019 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 - -#ifndef IREE_COMPILER_DIALECT_VM_TARGET_BYTECODE_TRANSLATIONFLAGS_H_ -#define IREE_COMPILER_DIALECT_VM_TARGET_BYTECODE_TRANSLATIONFLAGS_H_ - -#include "iree/compiler/Dialect/VM/Target/Bytecode/BytecodeModuleTarget.h" - -namespace mlir { -namespace iree_compiler { -namespace IREE { -namespace VM { - -// Returns a BytecodeTargetOptions struct initialized with the -// --iree-vm-bytecode-* flags. -BytecodeTargetOptions getBytecodeTargetOptionsFromFlags(); - -} // namespace VM -} // namespace IREE -} // namespace iree_compiler -} // namespace mlir - -#endif // IREE_COMPILER_DIALECT_VM_TARGET_BYTECODE_TRANSLATIONFLAGS_H_
diff --git a/iree/compiler/Dialect/VM/Target/Bytecode/TranslationRegistration.cpp b/iree/compiler/Dialect/VM/Target/Bytecode/TranslationRegistration.cpp index 77568f1..88367c4 100644 --- a/iree/compiler/Dialect/VM/Target/Bytecode/TranslationRegistration.cpp +++ b/iree/compiler/Dialect/VM/Target/Bytecode/TranslationRegistration.cpp
@@ -5,7 +5,6 @@ // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception #include "iree/compiler/Dialect/VM/Target/Bytecode/BytecodeModuleTarget.h" -#include "iree/compiler/Dialect/VM/Target/Bytecode/TranslationFlags.h" #include "mlir/IR/BuiltinOps.h" #include "mlir/IR/Visitors.h" #include "mlir/Translation.h" @@ -20,7 +19,7 @@ "iree-vm-ir-to-bytecode-module", [](mlir::ModuleOp moduleOp, llvm::raw_ostream &output) { return translateModuleToBytecode( - moduleOp, getBytecodeTargetOptionsFromFlags(), output); + moduleOp, BytecodeTargetOptions::FromFlags::get(), output); }); }
diff --git a/iree/compiler/Dialect/VM/Transforms/Conversion.cpp b/iree/compiler/Dialect/VM/Transforms/Conversion.cpp index c190deb..cc55193 100644 --- a/iree/compiler/Dialect/VM/Transforms/Conversion.cpp +++ b/iree/compiler/Dialect/VM/Transforms/Conversion.cpp
@@ -158,7 +158,7 @@ static PassRegistration<ConversionPass> pass( [] { - auto options = getTargetOptionsFromFlags(); + auto options = TargetOptions::FromFlags::get(); return std::make_unique<ConversionPass>(options); });
diff --git a/iree/compiler/Dialect/VM/Transforms/Passes.cpp b/iree/compiler/Dialect/VM/Transforms/Passes.cpp index 8dd1981..db2720a 100644 --- a/iree/compiler/Dialect/VM/Transforms/Passes.cpp +++ b/iree/compiler/Dialect/VM/Transforms/Passes.cpp
@@ -53,7 +53,8 @@ "iree-vm-transformation-pipeline", "Runs the full IREE VM dialect transformation pipeline", [](OpPassManager &passManager) { - buildVMTransformPassPipeline(passManager, getTargetOptionsFromFlags()); + buildVMTransformPassPipeline(passManager, + TargetOptions::FromFlags::get()); }); }
diff --git a/iree/compiler/Dialect/VM/Transforms/Passes.h b/iree/compiler/Dialect/VM/Transforms/Passes.h index 83bd4c9..eff235a 100644 --- a/iree/compiler/Dialect/VM/Transforms/Passes.h +++ b/iree/compiler/Dialect/VM/Transforms/Passes.h
@@ -93,7 +93,7 @@ //===----------------------------------------------------------------------===// inline void registerVMPasses() { - auto targetOptions = getTargetOptionsFromFlags(); + auto targetOptions = TargetOptions::FromFlags::get(); registerVMTransformPassPipeline(); createConversionPass(targetOptions); createHoistInlinedRodataPass(); @@ -104,7 +104,7 @@ } inline void registerVMTestPasses() { - getTargetOptionsFromFlags(); + TargetOptions::FromFlags::get(); createConvertStandardToVMTestPass(); }
diff --git a/iree/compiler/InputConversion/Common/IREEImportPublic.cpp b/iree/compiler/InputConversion/Common/IREEImportPublic.cpp index 895bf00..21096bb 100644 --- a/iree/compiler/InputConversion/Common/IREEImportPublic.cpp +++ b/iree/compiler/InputConversion/Common/IREEImportPublic.cpp
@@ -110,7 +110,8 @@ if (!resultType) return failure(); rewriter.replaceOpWithNewOp<IREE::HAL::TensorExportOp>( srcOp, resultType, adaptor.source(), - TypeAttr::get(adaptor.source().getType()), adaptor.source_dims()); + TypeAttr::get(adaptor.source().getType()), adaptor.source_dims(), + /*target_storage=*/nullptr); return success(); } };
diff --git a/iree/compiler/Translation/HALExecutable.cpp b/iree/compiler/Translation/HALExecutable.cpp index 1b9fc71..957515c 100644 --- a/iree/compiler/Translation/HALExecutable.cpp +++ b/iree/compiler/Translation/HALExecutable.cpp
@@ -52,7 +52,7 @@ mlir::registerPassManagerCLOptions(); // Convert into the final target-specific module definition and serialize it. - auto executableOptions = IREE::HAL::getTargetOptionsFromFlags(); + auto executableOptions = IREE::HAL::TargetOptions::FromFlags::get(); auto result = translateFromMLIRToHALExecutable(moduleOp, executableOptions); if (failed(result)) { return result;
diff --git a/iree/compiler/Translation/IREEVM.cpp b/iree/compiler/Translation/IREEVM.cpp index 7e22013..a9eeb7f 100644 --- a/iree/compiler/Translation/IREEVM.cpp +++ b/iree/compiler/Translation/IREEVM.cpp
@@ -13,7 +13,6 @@ #include "iree/compiler/Dialect/HAL/Transforms/Passes.h" #include "iree/compiler/Dialect/Stream/Transforms/Passes.h" #include "iree/compiler/Dialect/Util/Transforms/Passes.h" -#include "iree/compiler/Dialect/VM/Target/Bytecode/TranslationFlags.h" #include "iree/compiler/Dialect/VM/Transforms/Passes.h" #include "iree/compiler/InputConversion/Common/Passes.h" #include "iree/compiler/InputConversion/MHLO/Passes.h" @@ -31,73 +30,55 @@ namespace mlir { namespace iree_compiler { -static BindingOptions getBindingOptionsFromFlags() { +void BindingOptions::bindOptions(OptionsBinder &binder) { static llvm::cl::OptionCategory bindingOptionsCategory( "IREE translation binding support options"); - - static llvm::cl::opt<bool> *bindingsNativeFlag = new llvm::cl::opt<bool>{ - "iree-native-bindings-support", + binder.opt<bool>( + "iree-native-bindings-support", native, llvm::cl::desc( "Include runtime support for native IREE ABI-compatible bindings"), - llvm::cl::init(true), llvm::cl::cat(bindingOptionsCategory)}; - - static llvm::cl::opt<bool> *bindingsTFLiteFlag = new llvm::cl::opt<bool>{ - "iree-tflite-bindings-support", + llvm::cl::cat(bindingOptionsCategory)); + binder.opt<bool>( + "iree-tflite-bindings-support", tflite, llvm::cl::desc( "Include runtime support for the IREE TFLite compatibility bindings"), - llvm::cl::init(false), llvm::cl::cat(bindingOptionsCategory)}; - - BindingOptions bindingOptions; - bindingOptions.native = *bindingsNativeFlag; - bindingOptions.tflite = *bindingsTFLiteFlag; - return bindingOptions; + llvm::cl::cat(bindingOptionsCategory)); } -static InputDialectOptions getInputDialectOptionsFromFlags() { +void InputDialectOptions::bindOptions(OptionsBinder &binder) { static llvm::cl::OptionCategory inputDialectOptions( "IREE options for controlling the input transformations to apply"); - static llvm::cl::opt<InputDialectOptions::Type> *typeFlag = - new llvm::cl::opt<InputDialectOptions::Type>{ - "iree-input-type", llvm::cl::desc("IREE input type"), - llvm::cl::values( - clEnumValN(InputDialectOptions::Type::none, "none", - "No input dialect transformation"), - clEnumValN(InputDialectOptions::Type::tosa, "tosa", - "Legalize from TOSA ops"), - clEnumValN(InputDialectOptions::Type::mhlo, "mhlo", - "Legalize from MHLO ops"), - clEnumValN( - InputDialectOptions::Type::xla, "xla", - "Legalize from MHLO ops (with XLA cleanup preprocessing)")), - llvm::cl::init(InputDialectOptions::Type::none), - llvm::cl::cat(inputDialectOptions)}; - - InputDialectOptions options; - options.type = *typeFlag; - return options; + binder.opt<InputDialectOptions::Type>( + "iree-input-type", type, llvm::cl::desc("IREE input type"), + llvm::cl::values( + clEnumValN(InputDialectOptions::Type::none, "none", + "No input dialect transformation"), + clEnumValN(InputDialectOptions::Type::tosa, "tosa", + "Legalize from TOSA ops"), + clEnumValN(InputDialectOptions::Type::mhlo, "mhlo", + "Legalize from MHLO ops"), + clEnumValN( + InputDialectOptions::Type::xla, "xla", + "Legalize from MHLO ops (with XLA cleanup preprocessing)")), + llvm::cl::cat(inputDialectOptions)); } -static HighLevelOptimizationOptions getHighLevelOptimizationOptionsFromFlags() { +void HighLevelOptimizationOptions::bindOptions(OptionsBinder &binder) { static llvm::cl::OptionCategory category( "IREE options for controlling high level optimizations"); - static llvm::cl::opt<bool> *constEval = new llvm::cl::opt<bool>{ - "iree-const-eval", + binder.opt<bool>( + "iree-const-eval", constEval, llvm::cl::desc("Enables eager evaluation of constants using the full " "compiler and runtime"), - llvm::cl::init(false), llvm::cl::cat(category)}; - static llvm::cl::opt<bool> *constExprHoisting = new llvm::cl::opt<bool>{ - "iree-const-expr-hoisting", + llvm::cl::cat(category)); + binder.opt<bool>( + "iree-const-expr-hoisting", constExprHoisting, llvm::cl::desc( "Hoists the results of latent constant expressions into immutable " "global initializers for evaluation at program load"), - llvm::cl::init(false), llvm::cl::cat(category)}; - - HighLevelOptimizationOptions options; - options.constEval = *constEval; - options.constExprHoisting = *constExprHoisting; - return options; + llvm::cl::cat(category)); } void buildIREEVMTransformPassPipeline( @@ -152,10 +133,10 @@ void buildDefaultIREEVMTransformPassPipeline(OpPassManager &passManager) { buildIREEVMTransformPassPipeline( - getBindingOptionsFromFlags(), getInputDialectOptionsFromFlags(), - getHighLevelOptimizationOptionsFromFlags(), - IREE::HAL::getTargetOptionsFromFlags(), - IREE::VM::getTargetOptionsFromFlags(), passManager); + BindingOptions::FromFlags::get(), InputDialectOptions::FromFlags::get(), + HighLevelOptimizationOptions::FromFlags::get(), + IREE::HAL::TargetOptions::FromFlags::get(), + IREE::VM::TargetOptions::FromFlags::get(), passManager); } void registerIREEVMTransformPassPipeline() { @@ -199,13 +180,14 @@ static LogicalResult translateFromMLIRToVMBytecodeModuleWithFlags( ModuleOp moduleOp, llvm::raw_ostream &output) { mlir::registerPassManagerCLOptions(); - auto bindingOptions = getBindingOptionsFromFlags(); - auto inputOptions = getInputDialectOptionsFromFlags(); + auto bindingOptions = BindingOptions::FromFlags::get(); + auto inputOptions = InputDialectOptions::FromFlags::get(); auto highLevelOptimizationOptions = - getHighLevelOptimizationOptionsFromFlags(); - auto halTargetOptions = IREE::HAL::getTargetOptionsFromFlags(); - auto vmTargetOptions = IREE::VM::getTargetOptionsFromFlags(); - auto bytecodeTargetOptions = IREE::VM::getBytecodeTargetOptionsFromFlags(); + HighLevelOptimizationOptions::FromFlags::get(); + auto halTargetOptions = IREE::HAL::TargetOptions::FromFlags::get(); + auto vmTargetOptions = IREE::VM::TargetOptions::FromFlags::get(); + auto bytecodeTargetOptions = + IREE::VM::BytecodeTargetOptions::FromFlags::get(); auto result = translateFromMLIRToVM(moduleOp, bindingOptions, inputOptions, highLevelOptimizationOptions, halTargetOptions, vmTargetOptions); @@ -223,12 +205,12 @@ static LogicalResult translateFromMLIRToVMCModuleWithFlags( ModuleOp moduleOp, llvm::raw_ostream &output) { mlir::registerPassManagerCLOptions(); - auto bindingOptions = getBindingOptionsFromFlags(); - auto inputOptions = getInputDialectOptionsFromFlags(); + auto bindingOptions = BindingOptions::FromFlags::get(); + auto inputOptions = InputDialectOptions::FromFlags::get(); auto highLevelOptimizationOptions = - getHighLevelOptimizationOptionsFromFlags(); - auto halTargetOptions = IREE::HAL::getTargetOptionsFromFlags(); - auto vmTargetOptions = IREE::VM::getTargetOptionsFromFlags(); + HighLevelOptimizationOptions::FromFlags::get(); + auto halTargetOptions = IREE::HAL::TargetOptions::FromFlags::get(); + auto vmTargetOptions = IREE::VM::TargetOptions::FromFlags::get(); auto cTargetOptions = IREE::VM::getCTargetOptionsFromFlags(); auto result = translateFromMLIRToVM(moduleOp, bindingOptions, inputOptions, highLevelOptimizationOptions, @@ -243,9 +225,12 @@ #endif // IREE_HAVE_EMITC_DIALECT void registerIREEVMTranslationFlags() { - getBindingOptionsFromFlags(); - getInputDialectOptionsFromFlags(); - getHighLevelOptimizationOptionsFromFlags(); + BindingOptions::FromFlags::get(); + InputDialectOptions::FromFlags::get(); + HighLevelOptimizationOptions::FromFlags::get(); + IREE::HAL::TargetOptions::FromFlags::get(); + IREE::VM::TargetOptions::FromFlags::get(); + IREE::VM::BytecodeTargetOptions::FromFlags::get(); } void registerIREEVMTranslation() {
diff --git a/iree/compiler/Translation/IREEVM.h b/iree/compiler/Translation/IREEVM.h index 287dc6c..dbf601f 100644 --- a/iree/compiler/Translation/IREEVM.h +++ b/iree/compiler/Translation/IREEVM.h
@@ -10,6 +10,7 @@ #include "iree/compiler/Dialect/HAL/Target/TargetRegistry.h" #include "iree/compiler/Dialect/VM/Conversion/TargetOptions.h" #include "iree/compiler/Dialect/VM/Target/Bytecode/BytecodeModuleTarget.h" +#include "iree/compiler/Utils/OptionUtils.h" #include "llvm/Support/raw_ostream.h" #include "mlir/IR/BuiltinOps.h" #include "mlir/Pass/PassManager.h" @@ -29,6 +30,9 @@ // Whether to include runtime support functions required for the IREE TFLite // API compatibility bindings. bool tflite = false; + + void bindOptions(OptionsBinder &binder); + using FromFlags = OptionsFromFlags<BindingOptions>; }; // The transformation to apply to the input prior to main compiler execution. @@ -54,6 +58,9 @@ xla, }; Type type = Type::none; + + void bindOptions(OptionsBinder &binder); + using FromFlags = OptionsFromFlags<InputDialectOptions>; }; // Options controlling high level optimizations. @@ -64,6 +71,9 @@ // Enables recursive evaluation of immutable globals using the compiler // and runtime. bool constEval = false; + + void bindOptions(OptionsBinder &binder); + using FromFlags = OptionsFromFlags<HighLevelOptimizationOptions>; }; // Builds the translation pipeline with defaults.
diff --git a/iree/compiler/Utils/BUILD b/iree/compiler/Utils/BUILD index 0406bcb..188bd41 100644 --- a/iree/compiler/Utils/BUILD +++ b/iree/compiler/Utils/BUILD
@@ -18,6 +18,7 @@ "ConversionUtils.cpp", "FlatbufferUtils.cpp", "GraphUtils.cpp", + "OptionUtils.cpp", "PassUtils.cpp", "TracingUtils.cpp", ], @@ -26,6 +27,7 @@ "FlatbufferUtils.h", "GraphUtils.h", "IndexSet.h", + "OptionUtils.h", "PassUtils.h", "PatternUtils.h", "TracingUtils.h",
diff --git a/iree/compiler/Utils/CMakeLists.txt b/iree/compiler/Utils/CMakeLists.txt index a103507..56d7e2e 100644 --- a/iree/compiler/Utils/CMakeLists.txt +++ b/iree/compiler/Utils/CMakeLists.txt
@@ -18,6 +18,7 @@ "FlatbufferUtils.h" "GraphUtils.h" "IndexSet.h" + "OptionUtils.h" "PassUtils.h" "PatternUtils.h" "TracingUtils.h" @@ -25,6 +26,7 @@ "ConversionUtils.cpp" "FlatbufferUtils.cpp" "GraphUtils.cpp" + "OptionUtils.cpp" "PassUtils.cpp" "TracingUtils.cpp" DEPS
diff --git a/iree/compiler/Utils/OptionUtils.cpp b/iree/compiler/Utils/OptionUtils.cpp new file mode 100644 index 0000000..2a2a7f2 --- /dev/null +++ b/iree/compiler/Utils/OptionUtils.cpp
@@ -0,0 +1,94 @@ +// Copyright 2022 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/compiler/Utils/OptionUtils.h" + +#include "llvm/Support/ManagedStatic.h" + +namespace mlir { +namespace iree_compiler { + +void OptionsBinder::addGlobalOption(std::unique_ptr<llvm::cl::Option> option) { + static llvm::ManagedStatic<std::vector<std::unique_ptr<llvm::cl::Option>>> + globalOptions; + globalOptions->push_back(std::move(option)); +} + +LogicalResult OptionsBinder::parseArguments(int argc, const char *const *argv, + ErrorCallback onError) { + assert(scope && "can only parse arguments for local scoped binder"); + for (int i = 0; i < argc; ++i) { + llvm::StringRef arg(argv[i]); + llvm::StringRef nameVal; + if (arg.startswith("--")) { + nameVal = arg.drop_front(2); + } else if (arg.startswith("-")) { + nameVal = arg.drop_front(1); + } else { + // Pure positional options not supported. + if (onError) { + onError("pure positional arguments not supported (prefix with '--')"); + } + return failure(); + } + + // Split name and value. + llvm::StringRef name; + llvm::StringRef value; + size_t eqPos = nameVal.find("="); + if (eqPos == llvm::StringRef::npos) { + name = nameVal; + } else { + name = nameVal.take_front(eqPos); + value = nameVal.drop_front(eqPos + 1); + } + + // Find the option. + auto foundIt = scope->OptionsMap.find(name); + if (foundIt == scope->OptionsMap.end()) { + if (onError) { + std::string message("option not found: "); + message.append(name.begin(), name.end()); + onError(message); + } + return failure(); + } + llvm::cl::Option *option = foundIt->second; + + if (llvm::cl::ProvidePositionalOption(option, value, argc)) { + // Error. + if (onError) { + std::string message("option parse error for: "); + message.append(name.begin(), name.end()); + message.append("="); + message.append(value.begin(), value.end()); + onError(message); + } + return failure(); + } + } + + return success(); +} + +llvm::SmallVector<std::string> OptionsBinder::printArguments( + bool nonDefaultOnly) { + llvm::SmallVector<std::string> values; + for (auto &info : localOptions) { + if (!info.print) continue; + if (nonDefaultOnly && !info.isChanged()) continue; + + std::string s; + llvm::raw_string_ostream os(s); + info.print(os); + os.flush(); + values.push_back(std::move(s)); + } + return values; +} + +} // namespace iree_compiler +} // namespace mlir
diff --git a/iree/compiler/Utils/OptionUtils.h b/iree/compiler/Utils/OptionUtils.h new file mode 100644 index 0000000..2e7fae5 --- /dev/null +++ b/iree/compiler/Utils/OptionUtils.h
@@ -0,0 +1,230 @@ +// Copyright 2022 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 + +#ifndef IREE_COMPILER_UTILS_FLAG_UTILS_H +#define IREE_COMPILER_UTILS_FLAG_UTILS_H + +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/Support/CommandLine.h" +#include "llvm/Support/raw_ostream.h" +#include "mlir/Support/LogicalResult.h" + +namespace mlir { +namespace iree_compiler { + +// Base class that can bind named options to fields of structs. +// +// Typically use by adding the following to your struct: +// void bindOptions(OptionsBinder &binder); +// using FromFlags = OptionsFromFlags<MyStruct>; +// +// Then you can get the struct as initialized from global CL options as: +// MyStruct::FromFlags::get() +// +// Such use is referred to as a "global binder". You can also create a +// local binder, which does not interact with global flags by calling the +// local() static factory. When in this mode, the lifetime of any bound +// structures must exceed uses of the binder to parse or print (ie. via +// parseArguments() or printArguments()). +// +// The underlying LLVM command line support is quite flexible and all esoteric +// features are not supported here. Consider that supported structs can define +// options of built-in scalar types (string, ints, bool, etc) and enums. Lists +// of built-in scalar types are also supported. +class OptionsBinder { + public: + static OptionsBinder global() { return OptionsBinder(); } + + static OptionsBinder local() { + return OptionsBinder(std::make_unique<llvm::cl::SubCommand>()); + } + + template <typename T, typename V, typename... Mods> + void opt(llvm::StringRef name, V &value, Mods... Ms) { + if (!scope) { + // Bind global options. + auto opt = std::make_unique<llvm::cl::opt<T, /*ExternalStorage=*/true>>( + name, llvm::cl::location(value), llvm::cl::init(value), + std::forward<Mods>(Ms)...); + addGlobalOption(std::move(opt)); + } else { + // Bind local options. + auto option = + std::make_unique<llvm::cl::opt<T, /*ExternalStorage=*/true>>( + name, llvm::cl::sub(*scope), llvm::cl::location(value), + llvm::cl::init(value), std::forward<Mods>(Ms)...); + auto printCallback = + makePrintCallback(option->ArgStr, option->getParser(), &value); + auto changedCallback = makeChangedCallback(&value); + localOptions.push_back( + LocalOptionInfo{std::move(option), printCallback, changedCallback}); + } + } + + template <typename T, typename V, typename... Mods> + void list(llvm::StringRef name, V &value, Mods... Ms) { + if (!scope) { + // Bind global options. + auto list = + std::make_unique<llvm::cl::list<T>>(name, std::forward<Mods>(Ms)...); + // Since list does not support external storage, hook the callback + // and use it to update. + list->setCallback( + [&value](const T &newElement) { value.push_back(newElement); }); + addGlobalOption(std::move(list)); + } else { + // Bind local options. + auto list = std::make_unique<llvm::cl::list<T>>( + name, llvm::cl::sub(*scope), std::forward<Mods>(Ms)...); + auto printCallback = + makeListPrintCallback(list->ArgStr, list->getParser(), &value); + auto changedCallback = makeListChangedCallback(&value); + // Since list does not support external storage, hook the callback + // and use it to update. + list->setCallback( + [&value](const T &newElement) { value.push_back(newElement); }); + + localOptions.push_back( + LocalOptionInfo{std::move(list), printCallback, changedCallback}); + } + } + + // For a local binder, parses a sequence of flags of the usual form on + // command lines. + using ErrorCallback = std::function<void(llvm::StringRef message)>; + LogicalResult parseArguments(int argc, const char *const *argv, + ErrorCallback onError = nullptr); + + // Prints any flag values that differ from their default. + // Flags print in the order declared, which preserves some notion of grouping + // and is stable. + llvm::SmallVector<std::string> printArguments(bool nonDefaultOnly = false); + + private: + struct LocalOptionInfo { + using ChangedCallback = std::function<bool()>; + using PrintCallback = std::function<void(llvm::raw_ostream &)>; + std::unique_ptr<llvm::cl::Option> option; + PrintCallback print; + ChangedCallback isChanged; + }; + + OptionsBinder() = default; + OptionsBinder(std::unique_ptr<llvm::cl::SubCommand> scope) + : scope(std::move(scope)) {} + void addGlobalOption(std::unique_ptr<llvm::cl::Option> option); + + // LLVM makes a half-hearted (i.e. "best effort" == "no effort") attempt to + // handle non-enumerated generic value based options, but the generic + // comparisons are not reliably implemented. Simplify our lives by only + // supporting int convertible values (i.e. enums), which we can restrict + // ourselves to. + // Scalar enum print specialization. + template <typename V, typename ParserTy> + static auto makePrintCallback(llvm::StringRef optionName, ParserTy &parser, + V *value) + -> decltype(static_cast<llvm::cl::generic_parser_base &>(parser), + static_cast<int>(*value), LocalOptionInfo::PrintCallback()) { + return [optionName, &parser, value](llvm::raw_ostream &os) { + StringRef valueName("<unknown>"); + for (unsigned i = 0; i < parser.getNumOptions(); ++i) { + V cmpValue = static_cast<const llvm::cl::OptionValue<V> &>( + parser.getOptionValue(i)) + .getValue(); + if (cmpValue == *value) { + valueName = parser.getOption(i); + break; + } + } + os << "--" << optionName << "=" << valueName; + }; + } + + // Basic scalar print specialization. + template <typename V, typename ParserTy> + static auto makePrintCallback(llvm::StringRef optionName, ParserTy &parser, + V *value) + -> decltype(static_cast<llvm::cl::basic_parser<V> &>(parser), + LocalOptionInfo::PrintCallback()) { + return [optionName, value](llvm::raw_ostream &os) { + os << "--" << optionName << "=" << *value; + }; + } + + // Bool scalar print specialization. + template <typename ParserTy> + static auto makePrintCallback(llvm::StringRef optionName, ParserTy &parser, + bool *value) + -> decltype(static_cast<llvm::cl::basic_parser<bool> &>(parser), + LocalOptionInfo::PrintCallback()) { + return [optionName, value](llvm::raw_ostream &os) { + os << "--" << optionName << "="; + if (*value) { + os << "true"; + } else { + os << "false"; + } + }; + } + + // Scalar changed specialization. + template <typename V> + static LocalOptionInfo::ChangedCallback makeChangedCallback(V *currentValue) { + // Capture the current value as the initial value. + V initialValue = *currentValue; + return [currentValue, initialValue]() -> bool { + return *currentValue != initialValue; + }; + } + + // List changed specialization. + template <typename V> + static LocalOptionInfo::ChangedCallback makeListChangedCallback( + V *currentValue) { + return [currentValue]() -> bool { return !currentValue->empty(); }; + } + + // List basic print specialization. This is the only one we provide so far. + // Add others if the compiler tells you to. + template <typename ListTy, typename ParserTy, + typename V = typename ListTy::value_type> + static auto makeListPrintCallback(llvm::StringRef optionName, + ParserTy &parser, ListTy *values) + -> decltype(static_cast<llvm::cl::basic_parser<V> &>(parser), + LocalOptionInfo::PrintCallback()) { + return [optionName, values](llvm::raw_ostream &os) { + os << "--" << optionName << "="; + for (auto it : llvm::enumerate(*values)) { + if (it.index() > 0) os << ","; + os << it.value(); + } + }; + } + + std::unique_ptr<llvm::cl::SubCommand> scope; + llvm::SmallVector<LocalOptionInfo> localOptions; +}; + +template <typename DerivedTy> +class OptionsFromFlags { + public: + static DerivedTy &get() { + struct InitializedTy : DerivedTy { + InitializedTy() { + OptionsBinder binder = OptionsBinder::global(); + DerivedTy::bindOptions(binder); + } + }; + static InitializedTy singleton; + return singleton; + } +}; + +} // namespace iree_compiler +} // namespace mlir + +#endif // IREE_COMPILER_UTILS_FLAG_UTILS_H
diff --git a/iree/hal/BUILD b/iree/hal/BUILD index c87a07d..ece606e 100644 --- a/iree/hal/BUILD +++ b/iree/hal/BUILD
@@ -30,6 +30,8 @@ "buffer_heap_impl.h", "buffer_view.c", "buffer_view.h", + "buffer_view_util.c", + "buffer_view_util.h", "command_buffer.c", "command_buffer.h", "command_buffer_validation.c",
diff --git a/iree/hal/CMakeLists.txt b/iree/hal/CMakeLists.txt index 53528cb..eb0afe7 100644 --- a/iree/hal/CMakeLists.txt +++ b/iree/hal/CMakeLists.txt
@@ -25,6 +25,8 @@ "buffer_heap_impl.h" "buffer_view.c" "buffer_view.h" + "buffer_view_util.c" + "buffer_view_util.h" "command_buffer.c" "command_buffer.h" "command_buffer_validation.c"
diff --git a/iree/hal/allocator.c b/iree/hal/allocator.c index 1f92b36..c6f4794 100644 --- a/iree/hal/allocator.c +++ b/iree/hal/allocator.c
@@ -53,6 +53,15 @@ return _VTABLE_DISPATCH(allocator, host_allocator)(allocator); } +IREE_API_EXPORT +iree_status_t iree_hal_allocator_trim(iree_hal_allocator_t* allocator) { + IREE_ASSERT_ARGUMENT(allocator); + IREE_TRACE_ZONE_BEGIN(z0); + iree_status_t status = _VTABLE_DISPATCH(allocator, trim)(allocator); + IREE_TRACE_ZONE_END(z0); + return status; +} + IREE_API_EXPORT void iree_hal_allocator_query_statistics( iree_hal_allocator_t* allocator, iree_hal_allocator_statistics_t* out_statistics) { @@ -77,13 +86,14 @@ IREE_API_EXPORT iree_status_t iree_hal_allocator_allocate_buffer( iree_hal_allocator_t* allocator, iree_hal_memory_type_t memory_type, iree_hal_buffer_usage_t allowed_usage, iree_host_size_t allocation_size, - iree_hal_buffer_t** out_buffer) { + iree_const_byte_span_t initial_data, iree_hal_buffer_t** out_buffer) { IREE_ASSERT_ARGUMENT(allocator); IREE_ASSERT_ARGUMENT(out_buffer); *out_buffer = NULL; IREE_TRACE_ZONE_BEGIN(z0); iree_status_t status = _VTABLE_DISPATCH(allocator, allocate_buffer)( - allocator, memory_type, allowed_usage, allocation_size, out_buffer); + allocator, memory_type, allowed_usage, allocation_size, initial_data, + out_buffer); IREE_TRACE_ZONE_END(z0); return status; }
diff --git a/iree/hal/allocator.h b/iree/hal/allocator.h index 9d68e31..a7d9d1d 100644 --- a/iree/hal/allocator.h +++ b/iree/hal/allocator.h
@@ -95,6 +95,10 @@ IREE_API_EXPORT iree_allocator_t iree_hal_allocator_host_allocator(const iree_hal_allocator_t* allocator); +// Trims cached/unused pooled buffers, if any. +IREE_API_EXPORT +iree_status_t iree_hal_allocator_trim(iree_hal_allocator_t* allocator); + // Queries the aggregate statistics from the allocator since creation. // Thread-safe; statistics are captured at the time the call is made. // @@ -121,6 +125,9 @@ iree_hal_buffer_usage_t intended_usage, iree_device_size_t allocation_size); // Allocates a buffer from the allocator. +// If |initial_data| is provided then the bytes will be copied into the device +// buffer. To avoid the copy when constant data is used prefer +// iree_hal_allocator_wrap_buffer when available. // Fails if the memory type requested for the given usage cannot be serviced. // Callers can use iree_hal_allocator_can_allocate to decide their memory use // strategy. @@ -138,7 +145,7 @@ IREE_API_EXPORT iree_status_t iree_hal_allocator_allocate_buffer( iree_hal_allocator_t* allocator, iree_hal_memory_type_t memory_type, iree_hal_buffer_usage_t allowed_usage, iree_host_size_t allocation_size, - iree_hal_buffer_t** out_buffer); + iree_const_byte_span_t initial_data, iree_hal_buffer_t** out_buffer); // Wraps an existing host allocation in a buffer. // @@ -187,14 +194,13 @@ //===----------------------------------------------------------------------===// typedef struct iree_hal_allocator_vtable_t { - // << HAL C porting in progress >> - IREE_API_UNSTABLE - void(IREE_API_PTR* destroy)(iree_hal_allocator_t* allocator); iree_allocator_t(IREE_API_PTR* host_allocator)( const iree_hal_allocator_t* allocator); + iree_status_t(IREE_API_PTR* trim)(iree_hal_allocator_t* allocator); + void(IREE_API_PTR* query_statistics)( iree_hal_allocator_t* allocator, iree_hal_allocator_statistics_t* out_statistics); @@ -208,7 +214,7 @@ iree_status_t(IREE_API_PTR* allocate_buffer)( iree_hal_allocator_t* allocator, iree_hal_memory_type_t memory_type, iree_hal_buffer_usage_t allowed_usage, iree_host_size_t allocation_size, - iree_hal_buffer_t** out_buffer); + iree_const_byte_span_t initial_data, iree_hal_buffer_t** out_buffer); iree_status_t(IREE_API_PTR* wrap_buffer)( iree_hal_allocator_t* allocator, iree_hal_memory_type_t memory_type, @@ -219,6 +225,7 @@ void(IREE_API_PTR* deallocate_buffer)(iree_hal_allocator_t* allocator, iree_hal_buffer_t* buffer); } iree_hal_allocator_vtable_t; +IREE_HAL_ASSERT_VTABLE_LAYOUT(iree_hal_allocator_vtable_t); IREE_API_EXPORT void iree_hal_allocator_destroy( iree_hal_allocator_t* allocator);
diff --git a/iree/hal/allocator_heap.c b/iree/hal/allocator_heap.c index 7954299..7513803 100644 --- a/iree/hal/allocator_heap.c +++ b/iree/hal/allocator_heap.c
@@ -81,6 +81,11 @@ return allocator->host_allocator; } +static iree_status_t iree_hal_heap_allocator_trim( + iree_hal_allocator_t* base_allocator) { + return iree_ok_status(); +} + static void iree_hal_heap_allocator_query_statistics( iree_hal_allocator_t* base_allocator, iree_hal_allocator_statistics_t* out_statistics) { @@ -148,7 +153,7 @@ static iree_status_t iree_hal_heap_allocator_allocate_buffer( iree_hal_allocator_t* base_allocator, iree_hal_memory_type_t memory_type, iree_hal_buffer_usage_t allowed_usage, iree_host_size_t allocation_size, - iree_hal_buffer_t** out_buffer) { + iree_const_byte_span_t initial_data, iree_hal_buffer_t** out_buffer) { iree_hal_heap_allocator_t* allocator = iree_hal_heap_allocator_cast(base_allocator); @@ -160,10 +165,24 @@ // Allocate the buffer (both the wrapper and the contents). iree_hal_heap_allocator_statistics_t* statistics = NULL; IREE_STATISTICS(statistics = &allocator->statistics); - return iree_hal_heap_buffer_create(base_allocator, statistics, memory_type, - allowed_access, allowed_usage, - allocation_size, allocator->data_allocator, - allocator->host_allocator, out_buffer); + iree_hal_buffer_t* buffer = NULL; + IREE_RETURN_IF_ERROR(iree_hal_heap_buffer_create( + base_allocator, statistics, memory_type, allowed_access, allowed_usage, + allocation_size, allocator->data_allocator, allocator->host_allocator, + &buffer)); + + iree_status_t status = iree_ok_status(); + if (!iree_const_byte_span_is_empty(initial_data)) { + status = iree_hal_buffer_write_data(buffer, 0, initial_data.data, + initial_data.data_length); + } + + if (iree_status_is_ok(status)) { + *out_buffer = buffer; + } else { + iree_hal_buffer_release(buffer); + } + return status; } static iree_status_t iree_hal_heap_allocator_wrap_buffer( @@ -190,6 +209,7 @@ static const iree_hal_allocator_vtable_t iree_hal_heap_allocator_vtable = { .destroy = iree_hal_heap_allocator_destroy, .host_allocator = iree_hal_heap_allocator_host_allocator, + .trim = iree_hal_heap_allocator_trim, .query_statistics = iree_hal_heap_allocator_query_statistics, .query_buffer_compatibility = iree_hal_heap_allocator_query_buffer_compatibility,
diff --git a/iree/hal/api.h b/iree/hal/api.h index 33e2bdd..0ca7171 100644 --- a/iree/hal/api.h +++ b/iree/hal/api.h
@@ -12,6 +12,7 @@ #include "iree/hal/allocator.h" // IWYU pragma: export #include "iree/hal/buffer.h" // IWYU pragma: export #include "iree/hal/buffer_view.h" // IWYU pragma: export +#include "iree/hal/buffer_view_util.h" // IWYU pragma: export #include "iree/hal/command_buffer.h" // IWYU pragma: export #include "iree/hal/descriptor_set.h" // IWYU pragma: export #include "iree/hal/descriptor_set_layout.h" // IWYU pragma: export
diff --git a/iree/hal/buffer.c b/iree/hal/buffer.c index b83d809..81876ed 100644 --- a/iree/hal/buffer.c +++ b/iree/hal/buffer.c
@@ -115,18 +115,18 @@ iree_hal_buffer_t* buffer, iree_hal_mapping_mode_t mapping_mode, iree_hal_memory_access_t memory_access, iree_device_size_t local_byte_offset, iree_device_size_t local_byte_length, - void** out_data_ptr) { + iree_hal_buffer_mapping_t* mapping) { return _VTABLE_DISPATCH(buffer->allocated_buffer, map_range)( buffer->allocated_buffer, mapping_mode, memory_access, local_byte_offset, - local_byte_length, out_data_ptr); + local_byte_length, mapping); } -static void iree_hal_subspan_buffer_unmap_range( +static iree_status_t iree_hal_subspan_buffer_unmap_range( iree_hal_buffer_t* buffer, iree_device_size_t local_byte_offset, - iree_device_size_t local_byte_length, void* data_ptr) { - if (!buffer->allocated_buffer) return; - _VTABLE_DISPATCH(buffer->allocated_buffer, unmap_range) - (buffer->allocated_buffer, local_byte_offset, local_byte_length, data_ptr); + iree_device_size_t local_byte_length, iree_hal_buffer_mapping_t* mapping) { + if (!buffer->allocated_buffer) return iree_ok_status(); + return _VTABLE_DISPATCH(buffer->allocated_buffer, unmap_range)( + buffer->allocated_buffer, local_byte_offset, local_byte_length, mapping); } static iree_status_t iree_hal_subspan_buffer_invalidate_range( @@ -481,6 +481,10 @@ return buffer->allowed_usage; } +//===----------------------------------------------------------------------===// +// Transfer +//===----------------------------------------------------------------------===// + IREE_API_EXPORT iree_status_t iree_hal_buffer_zero(iree_hal_buffer_t* buffer, iree_device_size_t byte_offset, iree_device_size_t byte_length) { @@ -507,18 +511,18 @@ } IREE_TRACE_ZONE_BEGIN(z0); - iree_hal_buffer_mapping_t target_mapping; + iree_hal_buffer_mapping_t target_mapping = {{0}}; IREE_RETURN_AND_END_ZONE_IF_ERROR( - z0, - iree_hal_buffer_map_range(buffer, IREE_HAL_MEMORY_ACCESS_DISCARD_WRITE, - byte_offset, byte_length, &target_mapping)); + z0, iree_hal_buffer_map_range(buffer, IREE_HAL_MAPPING_MODE_SCOPED, + IREE_HAL_MEMORY_ACCESS_DISCARD_WRITE, + byte_offset, byte_length, &target_mapping)); if (byte_length == IREE_WHOLE_BUFFER) { byte_length = target_mapping.contents.data_length; } if (IREE_UNLIKELY((byte_offset % pattern_length) != 0) || IREE_UNLIKELY((byte_length % pattern_length) != 0)) { - iree_hal_buffer_unmap_range(&target_mapping); + iree_status_ignore(iree_hal_buffer_unmap_range(&target_mapping)); IREE_TRACE_ZONE_END(z0); return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, "attempting to fill a range with %zu byte values " @@ -572,7 +576,8 @@ status = iree_hal_buffer_flush_range(&target_mapping, 0, IREE_WHOLE_BUFFER); } - iree_hal_buffer_unmap_range(&target_mapping); + status = + iree_status_join(status, iree_hal_buffer_unmap_range(&target_mapping)); IREE_TRACE_ZONE_END(z0); return status; } @@ -587,11 +592,12 @@ IREE_ASSERT_ARGUMENT(target_buffer); IREE_TRACE_ZONE_BEGIN(z0); - iree_hal_buffer_mapping_t source_mapping; + IREE_TRACE_ZONE_APPEND_VALUE(z0, data_length); + iree_hal_buffer_mapping_t source_mapping = {{0}}; IREE_RETURN_AND_END_ZONE_IF_ERROR( - z0, - iree_hal_buffer_map_range(source_buffer, IREE_HAL_MEMORY_ACCESS_READ, - source_offset, data_length, &source_mapping)); + z0, iree_hal_buffer_map_range(source_buffer, IREE_HAL_MAPPING_MODE_SCOPED, + IREE_HAL_MEMORY_ACCESS_READ, source_offset, + data_length, &source_mapping)); memcpy(target_buffer, source_mapping.contents.data, data_length); @@ -610,11 +616,13 @@ IREE_ASSERT_ARGUMENT(source_buffer); IREE_TRACE_ZONE_BEGIN(z0); + IREE_TRACE_ZONE_APPEND_VALUE(z0, data_length); iree_hal_buffer_mapping_t target_mapping; IREE_RETURN_AND_END_ZONE_IF_ERROR( - z0, iree_hal_buffer_map_range( - target_buffer, IREE_HAL_MEMORY_ACCESS_DISCARD_WRITE, - target_offset, data_length, &target_mapping)); + z0, + iree_hal_buffer_map_range(target_buffer, IREE_HAL_MAPPING_MODE_SCOPED, + IREE_HAL_MEMORY_ACCESS_DISCARD_WRITE, + target_offset, data_length, &target_mapping)); memcpy(target_mapping.contents.data, source_buffer, data_length); @@ -650,19 +658,21 @@ } IREE_TRACE_ZONE_BEGIN(z0); + IREE_TRACE_ZONE_APPEND_VALUE(z0, data_length); // Map source, which may have IREE_WHOLE_BUFFER length. iree_hal_buffer_mapping_t source_mapping; IREE_RETURN_AND_END_ZONE_IF_ERROR( - z0, - iree_hal_buffer_map_range(source_buffer, IREE_HAL_MEMORY_ACCESS_READ, - source_offset, data_length, &source_mapping)); + z0, iree_hal_buffer_map_range(source_buffer, IREE_HAL_MAPPING_MODE_SCOPED, + IREE_HAL_MEMORY_ACCESS_READ, source_offset, + data_length, &source_mapping)); // Map target, which may also have IREE_WHOLE_BUFFER length. iree_hal_buffer_mapping_t target_mapping; - iree_status_t status = iree_hal_buffer_map_range( - target_buffer, IREE_HAL_MEMORY_ACCESS_DISCARD_WRITE, target_offset, - data_length, &target_mapping); + iree_status_t status = + iree_hal_buffer_map_range(target_buffer, IREE_HAL_MAPPING_MODE_SCOPED, + IREE_HAL_MEMORY_ACCESS_DISCARD_WRITE, + target_offset, data_length, &target_mapping); if (!iree_status_is_ok(status)) { iree_hal_buffer_unmap_range(&source_mapping); IREE_TRACE_ZONE_END(z0); @@ -682,8 +692,10 @@ adjusted_data_length = target_mapping.contents.data_length; } - // Elide zero length copies. - if (adjusted_data_length == 0) { + // Elide zero length copies. It's been expensive to get to this point just to + // bail but we need to have mapped to resolve IREE_WHOLE_BUFFERs that may + // result in zero lengths. + if (IREE_UNLIKELY(adjusted_data_length == 0)) { IREE_TRACE_ZONE_END(z0); return iree_ok_status(); } @@ -704,126 +716,123 @@ } //===----------------------------------------------------------------------===// -// Mapping / iree_hal_buffer_mapping_impl_t +// Mapping //===----------------------------------------------------------------------===// -typedef struct iree_hal_buffer_mapping_impl_t { - // Must be first (as in iree_hal_buffer_mapping_t). - // Stores both the offset data pointer and the byte_length of the mapping. - iree_byte_span_t contents; - // Retained buffer providing the backing storage for the mapping. - iree_hal_buffer_t* backing_buffer; - // Byte offset within the buffer where the mapped data begins. - iree_device_size_t byte_offset; - // Used for validation only. - iree_hal_memory_access_t allowed_access; - uint32_t reserved0; // unused - uint64_t reserved1; // unused -} iree_hal_buffer_mapping_impl_t; - -// We overlay the impl onto the external iree_hal_buffer_mapping_t struct; -// ensure we match the fields that are exposed. -static_assert(sizeof(iree_hal_buffer_mapping_impl_t) <= - sizeof(iree_hal_buffer_mapping_t), - "buffer mapping impl must fit inside the external struct"); -static_assert(offsetof(iree_hal_buffer_mapping_impl_t, contents) == - offsetof(iree_hal_buffer_mapping_t, contents), - "contents byte span must match the external struct offset"); - IREE_API_EXPORT iree_status_t iree_hal_buffer_map_range( - iree_hal_buffer_t* buffer, iree_hal_memory_access_t memory_access, - iree_device_size_t byte_offset, iree_device_size_t byte_length, + iree_hal_buffer_t* 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* out_buffer_mapping) { IREE_ASSERT_ARGUMENT(buffer); IREE_ASSERT_ARGUMENT(out_buffer_mapping); - memset(out_buffer_mapping, 0, sizeof(*out_buffer_mapping)); - IREE_RETURN_IF_ERROR(iree_hal_buffer_validate_memory_type( - iree_hal_buffer_memory_type(buffer), IREE_HAL_MEMORY_TYPE_HOST_VISIBLE)); - IREE_RETURN_IF_ERROR(iree_hal_buffer_validate_access( - iree_hal_buffer_allowed_access(buffer), memory_access)); - IREE_RETURN_IF_ERROR(iree_hal_buffer_validate_usage( - iree_hal_buffer_allowed_usage(buffer), IREE_HAL_BUFFER_USAGE_MAPPING)); - - iree_hal_buffer_mapping_impl_t* buffer_mapping = - (iree_hal_buffer_mapping_impl_t*)out_buffer_mapping; - buffer_mapping->backing_buffer = buffer; - buffer_mapping->allowed_access = memory_access; - iree_device_size_t data_length; - IREE_RETURN_IF_ERROR(iree_hal_buffer_calculate_range( - iree_hal_buffer_byte_offset(buffer), iree_hal_buffer_byte_length(buffer), - byte_offset, byte_length, &buffer_mapping->byte_offset, &data_length)); - buffer_mapping->contents.data_length = data_length; - - // TODO(benvanik): add mode arg to the HAL API. - iree_hal_mapping_mode_t mapping_mode = IREE_HAL_MAPPING_MODE_SCOPED; - IREE_TRACE_ZONE_BEGIN(z0); + memset(out_buffer_mapping, 0, sizeof(*out_buffer_mapping)); + + IREE_RETURN_AND_END_ZONE_IF_ERROR( + z0, iree_hal_buffer_validate_access( + iree_hal_buffer_allowed_access(buffer), memory_access)); + + // Persistent mapping requires the buffer was allocated to support it. + const bool is_persistent = + iree_all_bits_set(mapping_mode, IREE_HAL_MAPPING_MODE_PERSISTENT); + if (is_persistent) { + IREE_RETURN_AND_END_ZONE_IF_ERROR(z0, + iree_hal_buffer_validate_memory_type( + iree_hal_buffer_memory_type(buffer), + IREE_HAL_MEMORY_TYPE_HOST_VISIBLE)); + IREE_RETURN_AND_END_ZONE_IF_ERROR( + z0, + iree_hal_buffer_validate_usage(iree_hal_buffer_allowed_usage(buffer), + IREE_HAL_BUFFER_USAGE_MAPPING)); + } + + iree_device_size_t local_byte_offset = 0; + iree_device_size_t local_byte_length = 0; + IREE_RETURN_AND_END_ZONE_IF_ERROR( + z0, iree_hal_buffer_calculate_range( + iree_hal_buffer_byte_offset(buffer), + iree_hal_buffer_byte_length(buffer), byte_offset, byte_length, + &local_byte_offset, &local_byte_length)); + + out_buffer_mapping->buffer = buffer; + out_buffer_mapping->impl.allowed_access = memory_access; + out_buffer_mapping->impl.is_persistent = is_persistent ? 1 : 0; + out_buffer_mapping->impl.byte_offset = local_byte_offset; + iree_status_t status = _VTABLE_DISPATCH(buffer, map_range)( - buffer, mapping_mode, buffer_mapping->allowed_access, - buffer_mapping->byte_offset, buffer_mapping->contents.data_length, - (void**)&buffer_mapping->contents.data); + buffer, mapping_mode, memory_access, out_buffer_mapping->impl.byte_offset, + local_byte_length, out_buffer_mapping); + + if (iree_status_is_ok(status)) { + // Scoped mappings retain the buffer until unmapped. + if (!is_persistent) iree_hal_buffer_retain(buffer); + } else { + memset(out_buffer_mapping, 0, sizeof(*out_buffer_mapping)); + } + IREE_TRACE_ZONE_END(z0); return status; } -IREE_API_EXPORT void iree_hal_buffer_unmap_range( - iree_hal_buffer_mapping_t* base_buffer_mapping) { - IREE_ASSERT_ARGUMENT(base_buffer_mapping); - iree_hal_buffer_mapping_impl_t* buffer_mapping = - (iree_hal_buffer_mapping_impl_t*)base_buffer_mapping; - iree_hal_buffer_t* buffer = buffer_mapping->backing_buffer; - if (!buffer) return; +IREE_API_EXPORT iree_status_t +iree_hal_buffer_unmap_range(iree_hal_buffer_mapping_t* buffer_mapping) { + IREE_ASSERT_ARGUMENT(buffer_mapping); + iree_hal_buffer_t* buffer = buffer_mapping->buffer; + if (!buffer) return iree_ok_status(); IREE_TRACE_ZONE_BEGIN(z0); - _VTABLE_DISPATCH(buffer, unmap_range) - (buffer, buffer_mapping->byte_offset, buffer_mapping->contents.data_length, - buffer_mapping->contents.data); + + iree_status_t status = _VTABLE_DISPATCH(buffer, unmap_range)( + buffer, buffer_mapping->impl.byte_offset, + buffer_mapping->contents.data_length, buffer_mapping); + + if (!buffer_mapping->impl.is_persistent) { + iree_hal_buffer_release(buffer); + } + memset(buffer_mapping, 0, sizeof(*buffer_mapping)); + IREE_TRACE_ZONE_END(z0); + return status; } IREE_API_EXPORT iree_status_t iree_hal_buffer_invalidate_range( - iree_hal_buffer_mapping_t* base_buffer_mapping, - iree_device_size_t byte_offset, iree_device_size_t byte_length) { - IREE_ASSERT_ARGUMENT(base_buffer_mapping); - iree_hal_buffer_mapping_impl_t* buffer_mapping = - (iree_hal_buffer_mapping_impl_t*)base_buffer_mapping; - iree_hal_buffer_t* buffer = buffer_mapping->backing_buffer; + iree_hal_buffer_mapping_t* buffer_mapping, iree_device_size_t byte_offset, + iree_device_size_t byte_length) { + IREE_ASSERT_ARGUMENT(buffer_mapping); + iree_hal_buffer_t* buffer = buffer_mapping->buffer; IREE_RETURN_IF_ERROR(iree_hal_buffer_validate_access( - buffer_mapping->allowed_access, IREE_HAL_MEMORY_ACCESS_READ)); + buffer_mapping->impl.allowed_access, IREE_HAL_MEMORY_ACCESS_READ)); IREE_RETURN_IF_ERROR(iree_hal_buffer_calculate_range( - buffer_mapping->byte_offset, buffer_mapping->contents.data_length, + buffer_mapping->impl.byte_offset, buffer_mapping->contents.data_length, byte_offset, byte_length, &byte_offset, &byte_length)); return _VTABLE_DISPATCH(buffer, invalidate_range)(buffer, byte_offset, byte_length); } IREE_API_EXPORT iree_status_t iree_hal_buffer_flush_range( - iree_hal_buffer_mapping_t* base_buffer_mapping, - iree_device_size_t byte_offset, iree_device_size_t byte_length) { - IREE_ASSERT_ARGUMENT(base_buffer_mapping); - iree_hal_buffer_mapping_impl_t* buffer_mapping = - (iree_hal_buffer_mapping_impl_t*)base_buffer_mapping; - iree_hal_buffer_t* buffer = buffer_mapping->backing_buffer; + iree_hal_buffer_mapping_t* buffer_mapping, iree_device_size_t byte_offset, + iree_device_size_t byte_length) { + IREE_ASSERT_ARGUMENT(buffer_mapping); + iree_hal_buffer_t* buffer = buffer_mapping->buffer; IREE_RETURN_IF_ERROR(iree_hal_buffer_validate_access( - buffer_mapping->allowed_access, IREE_HAL_MEMORY_ACCESS_WRITE)); + buffer_mapping->impl.allowed_access, IREE_HAL_MEMORY_ACCESS_WRITE)); IREE_RETURN_IF_ERROR(iree_hal_buffer_calculate_range( - buffer_mapping->byte_offset, buffer_mapping->contents.data_length, + buffer_mapping->impl.byte_offset, buffer_mapping->contents.data_length, byte_offset, byte_length, &byte_offset, &byte_length)); return _VTABLE_DISPATCH(buffer, flush_range)(buffer, byte_offset, byte_length); } IREE_API_EXPORT iree_status_t iree_hal_buffer_mapping_subspan( - iree_hal_buffer_mapping_t* base_buffer_mapping, + iree_hal_buffer_mapping_t* buffer_mapping, iree_hal_memory_access_t memory_access, iree_device_size_t byte_offset, iree_device_size_t byte_length, iree_byte_span_t* out_span) { - IREE_ASSERT_ARGUMENT(base_buffer_mapping); - iree_hal_buffer_mapping_impl_t* buffer_mapping = - (iree_hal_buffer_mapping_impl_t*)base_buffer_mapping; + IREE_ASSERT_ARGUMENT(buffer_mapping); IREE_ASSERT_ARGUMENT(out_span); memset(out_span, 0, sizeof(*out_span)); IREE_RETURN_IF_ERROR(iree_hal_buffer_validate_access( - buffer_mapping->allowed_access, memory_access)); - iree_device_size_t data_length; + buffer_mapping->impl.allowed_access, memory_access)); + iree_device_size_t data_length = 0; IREE_RETURN_IF_ERROR(iree_hal_buffer_calculate_range( 0, buffer_mapping->contents.data_length, byte_offset, byte_length, &byte_offset, &data_length));
diff --git a/iree/hal/buffer.h b/iree/hal/buffer.h index 326efa9..82dfece 100644 --- a/iree/hal/buffer.h +++ b/iree/hal/buffer.h
@@ -140,7 +140,8 @@ // accesses will happen via command buffers. IREE_HAL_BUFFER_USAGE_TRANSFER = 1u << 1, - // The buffer can be mapped by the host application for reading and writing. + // The buffer can be mapped by the host application for reading and writing + // without a copy. // // As mapping may require placement in special address ranges or system // calls to enable visibility the driver can use the presence (or lack of) @@ -169,12 +170,48 @@ IREE_HAL_BUFFER_OVERLAP_COMPLETE, } iree_hal_buffer_overlap_t; +// A bitfield specifying buffer transfer behavior. +enum iree_hal_transfer_buffer_flag_bits_t { + // TODO(benvanik): flags controlling blocking, flushing, invalidation, and + // persistence. We may also want to set a bit that causes failure on emulated + // transfers that would otherwise be really expensive. + IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT = 0, +}; +typedef uint32_t iree_hal_transfer_buffer_flags_t; + +// Determines buffer mapping behavior. enum iree_hal_mapping_mode_bits_t { + // Buffers are mapped as part of a scoped map-access-unmap sequence. + // If there are any in-flight operations using the buffer contents are + // undefined though they may deceivingly still seem correct under certain + // implementations. IREE_HAL_MAPPING_MODE_SCOPED = 1u << 0, + + // Buffers are mapped persistently and concurrently accessible by both the + // host and device. Mapping happens once and so long as there are any live + // mappings the buffer will remain accessible. Not all implementations or + // buffer memory types support this, and even ones that do may not support + // coherent cross-device sharing. IREE_HAL_MAPPING_MODE_PERSISTENT = 1u << 1, }; typedef uint32_t iree_hal_mapping_mode_t; +// Implementation-specific mapping data. +typedef struct iree_hal_buffer_mapping_impl_t { + // Byte offset within the buffer where the mapped data begins. + iree_device_size_t byte_offset; + // Used for validation only. + iree_hal_memory_access_t allowed_access; + // Tracking flags. + uint32_t is_persistent : 1; + uint32_t reserved_flags : 31; + // Backing implementation data. + // For backends that require additional tracking (shadow data structures/etc) + // this can be used to store references to them for the duration of the + // mapping. + uint64_t reserved[1]; +} iree_hal_buffer_mapping_impl_t; + // Reference to a buffer's mapped memory. typedef struct iree_hal_buffer_mapping_t { // Contents of the buffer. Behavior is undefined if an access is performed @@ -185,8 +222,16 @@ // accessed. iree_byte_span_t contents; + // Buffer providing the backing storage for the mapping. + // When mapped with IREE_HAL_MAPPING_MODE_SCOPED the buffer will be retained + // until it is unmapped. When mapped with IREE_HAL_MAPPING_MODE_PERSISTENT the + // caller is responsible for retaining the buffer. + struct iree_hal_buffer_t* buffer; + // Used internally - do not modify. - uint64_t reserved[4]; + // Implementations are allowed to use the reserved fields for their own + // storage but should otherwise ignore the remaining parts. + iree_hal_buffer_mapping_impl_t impl; } iree_hal_buffer_mapping_t; // Formats a memory type bitfield as a string. @@ -389,7 +434,8 @@ // Copies data from the provided |source_buffer| into the |target_buffer|. // // Requires that both buffers have the IREE_HAL_BUFFER_USAGE_MAPPING bit set. -// The byte range in |target_buffer| will be flushed if needed. +// The byte range in |target_buffer| will be flushed if needed. Both buffers +// need not come from the same device. // // It is strongly recommended that buffer operations are performed on transfer // queues; using this synchronous function may incur additional cache flushes @@ -411,16 +457,21 @@ // invalidate the byte range they want to access to update the visibility of the // mapped memory. IREE_API_EXPORT iree_status_t iree_hal_buffer_map_range( - iree_hal_buffer_t* buffer, iree_hal_memory_access_t memory_access, - iree_device_size_t byte_offset, iree_device_size_t byte_length, + iree_hal_buffer_t* 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* out_buffer_mapping); // Unmaps the buffer as was previously mapped to |buffer_mapping|. // // If the buffer is not IREE_HAL_MEMORY_TYPE_HOST_COHERENT then the caller must // flush the byte range they want to make available to other threads/devices. -IREE_API_EXPORT void iree_hal_buffer_unmap_range( - iree_hal_buffer_mapping_t* buffer_mapping); +// +// May fail, though unlikely to do so for read-only mapping and the result can +// be safely ignored using iree_status_ignore. If writing then users must check +// the status to ensure their writes succeeded. +IREE_API_EXPORT iree_status_t +iree_hal_buffer_unmap_range(iree_hal_buffer_mapping_t* buffer_mapping); // Invalidates ranges of non-coherent memory from the host caches. // This guarantees that device writes to the memory ranges provided are @@ -486,9 +537,6 @@ //===----------------------------------------------------------------------===// typedef struct iree_hal_buffer_vtable_t { - // << HAL C porting in progress >> - IREE_API_UNSTABLE - void(IREE_API_PTR* destroy)(iree_hal_buffer_t* buffer); iree_status_t(IREE_API_PTR* map_range)(iree_hal_buffer_t* buffer, @@ -496,12 +544,12 @@ iree_hal_memory_access_t memory_access, iree_device_size_t local_byte_offset, iree_device_size_t local_byte_length, - void** out_data_ptr); + iree_hal_buffer_mapping_t* mapping); - void(IREE_API_PTR* unmap_range)(iree_hal_buffer_t* buffer, - iree_device_size_t local_byte_offset, - iree_device_size_t local_byte_length, - void* data_ptr); + iree_status_t(IREE_API_PTR* unmap_range)(iree_hal_buffer_t* buffer, + iree_device_size_t local_byte_offset, + iree_device_size_t local_byte_length, + iree_hal_buffer_mapping_t* mapping); iree_status_t(IREE_API_PTR* invalidate_range)( iree_hal_buffer_t* buffer, iree_device_size_t local_byte_offset, @@ -511,6 +559,7 @@ iree_hal_buffer_t* buffer, iree_device_size_t local_byte_offset, iree_device_size_t local_byte_length); } iree_hal_buffer_vtable_t; +IREE_HAL_ASSERT_VTABLE_LAYOUT(iree_hal_buffer_vtable_t); struct iree_hal_buffer_t { iree_hal_resource_t resource;
diff --git a/iree/hal/buffer_heap.c b/iree/hal/buffer_heap.c index 6649a8c..7486af2 100644 --- a/iree/hal/buffer_heap.c +++ b/iree/hal/buffer_heap.c
@@ -174,9 +174,10 @@ iree_hal_buffer_t* base_buffer, iree_hal_mapping_mode_t mapping_mode, iree_hal_memory_access_t memory_access, iree_device_size_t local_byte_offset, iree_device_size_t local_byte_length, - void** out_data_ptr) { + iree_hal_buffer_mapping_t* mapping) { iree_hal_heap_buffer_t* buffer = (iree_hal_heap_buffer_t*)base_buffer; - *out_data_ptr = buffer->data.data + local_byte_offset; + mapping->contents = iree_make_byte_span(buffer->data.data + local_byte_offset, + local_byte_length); // If we mapped for discard scribble over the bytes. This is not a mandated // behavior but it will make debugging issues easier. Alternatively for @@ -184,17 +185,18 @@ // would only work if the entire buffer was discarded. #ifndef NDEBUG if (iree_any_bit_set(memory_access, IREE_HAL_MEMORY_ACCESS_DISCARD)) { - memset(*out_data_ptr, 0xCD, local_byte_length); + memset(mapping->contents.data, 0xCD, local_byte_length); } #endif // !NDEBUG return iree_ok_status(); } -static void iree_hal_heap_buffer_unmap_range( +static iree_status_t iree_hal_heap_buffer_unmap_range( iree_hal_buffer_t* base_buffer, iree_device_size_t local_byte_offset, - iree_device_size_t local_byte_length, void* data_ptr) { + iree_device_size_t local_byte_length, iree_hal_buffer_mapping_t* mapping) { // No-op here as we always have the pointer. + return iree_ok_status(); } static iree_status_t iree_hal_heap_buffer_invalidate_range(
diff --git a/iree/hal/buffer_view.c b/iree/hal/buffer_view.c index 3c32045..c338235 100644 --- a/iree/hal/buffer_view.c +++ b/iree/hal/buffer_view.c
@@ -6,14 +6,11 @@ #include "iree/hal/buffer_view.h" -#include <inttypes.h> -#include <stdbool.h> - #include "iree/base/api.h" #include "iree/base/tracing.h" #include "iree/hal/allocator.h" +#include "iree/hal/buffer_view_util.h" #include "iree/hal/resource.h" -#include "iree/hal/string_util.h" struct iree_hal_buffer_view_t { iree_atomic_ref_count_t ref_count; @@ -94,131 +91,6 @@ IREE_TRACE_ZONE_END(z0); } -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_allocate_buffer( - iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, - iree_host_size_t shape_rank, iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, - iree_hal_buffer_usage_t allowed_usage, - iree_hal_buffer_view_t** out_buffer_view) { - IREE_ASSERT_ARGUMENT(allocator); - IREE_ASSERT_ARGUMENT(out_buffer_view); - IREE_TRACE_ZONE_BEGIN(z0); - - iree_device_size_t allocation_size = 0; - iree_status_t status = iree_hal_buffer_compute_view_size( - shape, shape_rank, element_type, encoding_type, &allocation_size); - - iree_hal_buffer_t* buffer = NULL; - if (iree_status_is_ok(status)) { - status = iree_hal_allocator_allocate_buffer( - allocator, memory_type, allowed_usage, allocation_size, &buffer); - } - - if (iree_status_is_ok(status)) { - status = iree_hal_buffer_view_create( - buffer, shape, shape_rank, element_type, encoding_type, - iree_hal_allocator_host_allocator(allocator), out_buffer_view); - } - - iree_hal_buffer_release(buffer); - IREE_TRACE_ZONE_END(z0); - return status; -} - -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_clone_heap_buffer( - iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, - iree_host_size_t shape_rank, iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, - iree_hal_buffer_usage_t allowed_usage, iree_const_byte_span_t data, - iree_hal_buffer_view_t** out_buffer_view) { - IREE_ASSERT_ARGUMENT(allocator); - IREE_ASSERT_ARGUMENT(out_buffer_view); - IREE_TRACE_ZONE_BEGIN(z0); - - // Allocate the buffer. - iree_hal_buffer_view_t* buffer_view = NULL; - IREE_RETURN_AND_END_ZONE_IF_ERROR( - z0, iree_hal_buffer_view_allocate_buffer( - allocator, shape, shape_rank, element_type, encoding_type, - memory_type, allowed_usage, &buffer_view)); - - // Copy all of the data into it in the worst way possible. - // If you find yourself coming here from profiling: - // Don't clone data. Allocate and then populate it in-place. - // -or- - // Schedule sequences of DMA transfers using iree_hal_command_buffer_t. - iree_status_t status = iree_hal_buffer_write_data( - iree_hal_buffer_view_buffer(buffer_view), 0, data.data, data.data_length); - - if (iree_status_is_ok(status)) { - *out_buffer_view = buffer_view; - } else { - iree_hal_buffer_view_release(buffer_view); - } - IREE_TRACE_ZONE_END(z0); - return status; -} - -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_wrap_heap_buffer( - iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, - iree_host_size_t shape_rank, iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, - iree_hal_memory_access_t allowed_access, - iree_hal_buffer_usage_t allowed_usage, iree_byte_span_t data, - iree_allocator_t data_allocator, iree_hal_buffer_view_t** out_buffer_view) { - IREE_ASSERT_ARGUMENT(allocator); - IREE_ASSERT_ARGUMENT(out_buffer_view); - IREE_TRACE_ZONE_BEGIN(z0); - - // NOTE: this will fail if the data cannot be imported into the allocator. - iree_hal_buffer_t* buffer = NULL; - iree_status_t status = iree_hal_allocator_wrap_buffer( - allocator, memory_type, allowed_access, allowed_usage, data, - data_allocator, &buffer); - - if (iree_status_is_ok(status)) { - status = iree_hal_buffer_view_create( - buffer, shape, shape_rank, element_type, encoding_type, - iree_hal_allocator_host_allocator(allocator), out_buffer_view); - } - - iree_hal_buffer_release(buffer); - IREE_TRACE_ZONE_END(z0); - return status; -} - -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_wrap_or_clone_heap_buffer( - iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, - iree_host_size_t shape_rank, iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, - iree_hal_memory_access_t allowed_access, - iree_hal_buffer_usage_t allowed_usage, iree_byte_span_t data, - iree_allocator_t data_allocator, iree_hal_buffer_view_t** out_buffer_view) { - IREE_ASSERT_ARGUMENT(allocator); - - // Not all HAL implementations support wrapping buffers, and of those that do - // some may only support it in special situations such as when the buffer is - // not DEVICE_VISIBLE. The user application can query whether the wrapping is - // possible and decide to use alternative means of upload if it is not; we - // make no policy (other than validity) over what's best here. - iree_hal_buffer_compatibility_t compatibility = - iree_hal_allocator_query_buffer_compatibility( - allocator, memory_type, allowed_usage, IREE_HAL_BUFFER_USAGE_MAPPING, - (iree_device_size_t)data.data_length); - bool wrap_allowed = iree_all_bits_set( - compatibility, IREE_HAL_BUFFER_COMPATIBILITY_IMPORTABLE); - if (wrap_allowed) { - return iree_hal_buffer_view_wrap_heap_buffer( - allocator, shape, shape_rank, element_type, encoding_type, memory_type, - allowed_access, allowed_usage, data, data_allocator, out_buffer_view); - } else { - return iree_hal_buffer_view_clone_heap_buffer( - allocator, shape, shape_rank, element_type, encoding_type, memory_type, - allowed_usage, iree_make_const_byte_span(data.data, data.data_length), - out_buffer_view); - } -} - IREE_API_EXPORT iree_hal_buffer_t* iree_hal_buffer_view_buffer( const iree_hal_buffer_view_t* buffer_view) { IREE_ASSERT_ARGUMENT(buffer_view); @@ -287,6 +159,10 @@ IREE_ASSERT_ARGUMENT(shape); if (shape_rank != buffer_view->shape_rank) { + // Rank changes require reallocation of the structure as we inline the + // shape dimensions. We could lighten this restriction to allow for rank + // reduction but knowing that rank changes aren't allowed is easier than + // remembering all the conditions in which they may be. return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, "buffer view reshapes must have the same rank; " "target=%zu, existing=%zu", @@ -357,406 +233,3 @@ buffer_view->encoding_type, start_indices, indices_count, lengths, lengths_count, out_start_offset, out_length); } - -IREE_API_EXPORT iree_status_t iree_hal_buffer_compute_view_size( - const iree_hal_dim_t* shape, iree_host_size_t shape_rank, - iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, - iree_device_size_t* out_allocation_size) { - IREE_ASSERT_ARGUMENT(!shape_rank || shape); - IREE_ASSERT_ARGUMENT(out_allocation_size); - *out_allocation_size = 0; - - iree_device_size_t byte_length = 0; - - switch (encoding_type) { - case IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR: { - if (IREE_UNLIKELY(iree_hal_element_bit_count(element_type) == 0) || - IREE_UNLIKELY(!iree_hal_element_is_byte_aligned(element_type))) { - return iree_make_status( - IREE_STATUS_INVALID_ARGUMENT, - "opaque and sub-byte aligned element types cannot be indexed"); - } - byte_length = iree_hal_element_dense_byte_count(element_type); - for (iree_host_size_t i = 0; i < shape_rank; ++i) { - byte_length *= shape[i]; - } - break; - } - default: - return iree_make_status(IREE_STATUS_UNIMPLEMENTED, - "unimplemented encoding type size calculation"); - } - - *out_allocation_size = byte_length; - return iree_ok_status(); -} - -IREE_API_EXPORT iree_status_t iree_hal_buffer_compute_view_offset( - const iree_hal_dim_t* shape, iree_host_size_t shape_rank, - iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, const iree_hal_dim_t* indices, - iree_host_size_t indices_count, iree_device_size_t* out_offset) { - IREE_ASSERT_ARGUMENT(shape); - IREE_ASSERT_ARGUMENT(indices); - IREE_ASSERT_ARGUMENT(out_offset); - *out_offset = 0; - if (IREE_UNLIKELY(encoding_type != IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR)) { - return iree_make_status( - IREE_STATUS_INVALID_ARGUMENT, - "only dense encodings support view range computation"); - } else if (IREE_UNLIKELY(iree_hal_element_bit_count(element_type) == 0) || - IREE_UNLIKELY(!iree_hal_element_is_byte_aligned(element_type))) { - return iree_make_status( - IREE_STATUS_INVALID_ARGUMENT, - "opaque and sub-byte aligned element types cannot be indexed"); - } else if (IREE_UNLIKELY(shape_rank != indices_count)) { - return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, - "shape rank/indices mismatch: %zu != %zu", - shape_rank, indices_count); - } - - iree_device_size_t offset = 0; - for (iree_host_size_t i = 0; i < indices_count; ++i) { - if (IREE_UNLIKELY(indices[i] >= shape[i])) { - return iree_make_status(IREE_STATUS_OUT_OF_RANGE, - "index[%zu] out of bounds: %d >= %d", i, - indices[i], shape[i]); - } - iree_device_size_t axis_offset = indices[i]; - for (iree_host_size_t j = i + 1; j < shape_rank; ++j) { - axis_offset *= shape[j]; - } - offset += axis_offset; - } - offset *= iree_hal_element_dense_byte_count(element_type); - - *out_offset = offset; - return iree_ok_status(); -} - -IREE_API_EXPORT iree_status_t iree_hal_buffer_compute_view_range( - const iree_hal_dim_t* shape, iree_host_size_t shape_rank, - iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, const iree_hal_dim_t* start_indices, - iree_host_size_t indices_count, const iree_hal_dim_t* lengths, - iree_host_size_t lengths_count, iree_device_size_t* out_start_offset, - iree_device_size_t* out_length) { - IREE_ASSERT_ARGUMENT(shape); - IREE_ASSERT_ARGUMENT(start_indices); - IREE_ASSERT_ARGUMENT(lengths); - IREE_ASSERT_ARGUMENT(out_start_offset); - IREE_ASSERT_ARGUMENT(out_length); - *out_start_offset = 0; - *out_length = 0; - if (IREE_UNLIKELY(encoding_type != IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR)) { - return iree_make_status( - IREE_STATUS_INVALID_ARGUMENT, - "only dense encodings support view range computation"); - } else if (IREE_UNLIKELY(iree_hal_element_bit_count(element_type) == 0) || - IREE_UNLIKELY(!iree_hal_element_is_byte_aligned(element_type))) { - return iree_make_status( - IREE_STATUS_INVALID_ARGUMENT, - "opaque and sub-byte aligned element types cannot be indexed"); - } else if (IREE_UNLIKELY(indices_count != lengths_count)) { - return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, - "indices/lengths mismatch: %zu != %zu", - indices_count, lengths_count); - } else if (IREE_UNLIKELY(shape_rank != indices_count)) { - return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, - "shape rank/indices mismatch: %zu != %zu", - shape_rank, indices_count); - } - - iree_hal_dim_t* end_indices = - iree_alloca(shape_rank * sizeof(iree_hal_dim_t)); - iree_device_size_t element_size = - iree_hal_element_dense_byte_count(element_type); - iree_device_size_t subspan_length = element_size; - for (iree_host_size_t i = 0; i < lengths_count; ++i) { - subspan_length *= lengths[i]; - end_indices[i] = start_indices[i] + lengths[i] - 1; - } - - iree_device_size_t start_byte_offset = 0; - IREE_RETURN_IF_ERROR(iree_hal_buffer_compute_view_offset( - shape, shape_rank, element_type, encoding_type, start_indices, - indices_count, &start_byte_offset)); - iree_device_size_t end_byte_offset = 0; - IREE_RETURN_IF_ERROR(iree_hal_buffer_compute_view_offset( - shape, shape_rank, element_type, encoding_type, end_indices, shape_rank, - &end_byte_offset)); - - // Non-contiguous regions not yet implemented. Will be easier to detect when - // we have strides. - iree_device_size_t offset_length = - end_byte_offset - start_byte_offset + element_size; - if (subspan_length != offset_length) { - return iree_make_status( - IREE_STATUS_UNIMPLEMENTED, - "non-contiguous range region computation not implemented"); - } - - *out_start_offset = start_byte_offset; - *out_length = subspan_length; - return iree_ok_status(); -} - -static iree_status_t iree_hal_buffer_view_parse_impl( - iree_string_view_t value, iree_hal_allocator_t* buffer_allocator, - iree_hal_buffer_view_t** out_buffer_view) { - // Strip whitespace that may come along (linefeeds/etc). - value = iree_string_view_trim(value); - value = iree_string_view_strip_prefix(value, IREE_SV("\"")); - value = iree_string_view_strip_suffix(value, IREE_SV("\"")); - if (iree_string_view_is_empty(value)) { - // Empty lines are invalid; need at least the shape/type information. - *out_buffer_view = NULL; - return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, "empty string input"); - } - - // The part of the string corresponding to the shape, e.g. 1x2x3. - iree_string_view_t shape_str = iree_string_view_empty(); - // The part of the string corresponding to the type, e.g. f32 - iree_string_view_t type_str = iree_string_view_empty(); - // The part of the string corresponding to the buffer data, e.g. 1 2 3 4 5 6 - iree_string_view_t data_str = iree_string_view_empty(); - - iree_string_view_t shape_and_type_str = value; - iree_string_view_split(value, '=', &shape_and_type_str, &data_str); - iree_host_size_t last_x_index = iree_string_view_find_last_of( - shape_and_type_str, IREE_SV("x"), IREE_STRING_VIEW_NPOS); - if (last_x_index == IREE_STRING_VIEW_NPOS) { - // Scalar. - type_str = shape_and_type_str; - } else { - // Has a shape. - shape_str = iree_string_view_substr(shape_and_type_str, 0, last_x_index); - type_str = iree_string_view_substr(shape_and_type_str, last_x_index + 1, - IREE_STRING_VIEW_NPOS); - } - - // AxBxC... - iree_host_size_t shape_rank = 0; - iree_status_t shape_result = - iree_hal_parse_shape(shape_str, 0, NULL, &shape_rank); - if (!iree_status_is_ok(shape_result) && - !iree_status_is_out_of_range(shape_result)) { - return shape_result; - } else if (shape_rank > 128) { - return iree_make_status( - IREE_STATUS_RESOURCE_EXHAUSTED, - "a shape rank of %zu is just a little bit excessive, eh?", shape_rank); - } - shape_result = iree_status_ignore(shape_result); - iree_hal_dim_t* shape = - (iree_hal_dim_t*)iree_alloca(shape_rank * sizeof(iree_hal_dim_t)); - IREE_RETURN_IF_ERROR( - iree_hal_parse_shape(shape_str, shape_rank, shape, &shape_rank)); - - // f32, i32, etc - iree_hal_element_type_t element_type = IREE_HAL_ELEMENT_TYPE_NONE; - IREE_RETURN_IF_ERROR(iree_hal_parse_element_type(type_str, &element_type)); - - // TODO(benvanik): allow specifying the encoding. - iree_hal_encoding_type_t encoding_type = - IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR; - - // Allocate the buffer we will parse into from the provided allocator. - iree_device_size_t buffer_length = 0; - IREE_RETURN_IF_ERROR(iree_hal_buffer_compute_view_size( - shape, shape_rank, element_type, encoding_type, &buffer_length)); - iree_hal_buffer_t* buffer = NULL; - IREE_RETURN_IF_ERROR(iree_hal_allocator_allocate_buffer( - buffer_allocator, - IREE_HAL_MEMORY_TYPE_HOST_LOCAL | IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE, - IREE_HAL_BUFFER_USAGE_TRANSFER | IREE_HAL_BUFFER_USAGE_MAPPING | - IREE_HAL_BUFFER_USAGE_DISPATCH, - buffer_length, &buffer)); - - // Parse the elements directly into the buffer. - iree_hal_buffer_mapping_t buffer_mapping; - iree_status_t status = - iree_hal_buffer_map_range(buffer, IREE_HAL_MEMORY_ACCESS_DISCARD_WRITE, 0, - buffer_length, &buffer_mapping); - if (!iree_status_is_ok(status)) { - iree_hal_buffer_release(buffer); - return status; - } - status = iree_hal_parse_buffer_elements(data_str, element_type, - buffer_mapping.contents); - iree_hal_buffer_unmap_range(&buffer_mapping); - if (!iree_status_is_ok(status)) { - iree_hal_buffer_release(buffer); - return status; - } - - // Wrap and pass ownership of the buffer to the buffer view. - status = iree_hal_buffer_view_create( - buffer, shape, shape_rank, element_type, encoding_type, - iree_hal_allocator_host_allocator(buffer_allocator), out_buffer_view); - iree_hal_buffer_release(buffer); - return status; -} - -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_parse( - iree_string_view_t value, iree_hal_allocator_t* buffer_allocator, - iree_hal_buffer_view_t** out_buffer_view) { - IREE_ASSERT_ARGUMENT(buffer_allocator); - IREE_ASSERT_ARGUMENT(out_buffer_view); - *out_buffer_view = NULL; - IREE_TRACE_ZONE_BEGIN(z0); - iree_status_t status = - iree_hal_buffer_view_parse_impl(value, buffer_allocator, out_buffer_view); - IREE_TRACE_ZONE_END(z0); - return status; -} - -#define APPEND_CHAR(c) \ - { \ - if (buffer) { \ - if (buffer_length < buffer_capacity - 1) { \ - buffer[buffer_length] = c; \ - buffer[buffer_length + 1] = '\0'; \ - } else { \ - buffer = NULL; \ - } \ - } \ - ++buffer_length; \ - } - -static iree_status_t iree_hal_buffer_view_format_impl( - const iree_hal_buffer_view_t* buffer_view, - iree_host_size_t max_element_count, iree_host_size_t buffer_capacity, - char* buffer, iree_host_size_t* out_buffer_length) { - if (out_buffer_length) { - *out_buffer_length = 0; - } - if (buffer && buffer_capacity) { - buffer[0] = 0; - } - - iree_host_size_t buffer_length = 0; - if (iree_hal_buffer_view_shape_rank(buffer_view) > 0) { - // Shape: 1x2x3 - iree_host_size_t shape_length = 0; - iree_status_t status = iree_hal_format_shape( - iree_hal_buffer_view_shape_dims(buffer_view), - iree_hal_buffer_view_shape_rank(buffer_view), - buffer ? buffer_capacity - buffer_length : 0, - buffer ? buffer + buffer_length : NULL, &shape_length); - buffer_length += shape_length; - if (iree_status_is_out_of_range(status)) { - status = iree_status_ignore(status); - buffer = NULL; - } else if (!iree_status_is_ok(status)) { - return status; - } - - // Separator: <shape>x<format> - APPEND_CHAR('x'); - } - - // Element type: f32 - iree_host_size_t element_type_length = 0; - iree_status_t status = iree_hal_format_element_type( - iree_hal_buffer_view_element_type(buffer_view), - buffer ? buffer_capacity - buffer_length : 0, - buffer ? buffer + buffer_length : NULL, &element_type_length); - buffer_length += element_type_length; - if (iree_status_is_out_of_range(status)) { - status = iree_status_ignore(status); - buffer = NULL; - } else if (!iree_status_is_ok(status)) { - return status; - } - - // TODO(benvanik): allow printing the encoding. - - // Separator: <meta>=<value> - APPEND_CHAR('='); - - // Buffer contents: 0 1 2 3 ... - iree_hal_buffer_mapping_t buffer_mapping; - IREE_RETURN_IF_ERROR(iree_hal_buffer_map_range( - iree_hal_buffer_view_buffer(buffer_view), IREE_HAL_MEMORY_ACCESS_READ, 0, - IREE_WHOLE_BUFFER, &buffer_mapping)); - iree_host_size_t elements_length = 0; - status = iree_hal_format_buffer_elements( - iree_make_const_byte_span(buffer_mapping.contents.data, - buffer_mapping.contents.data_length), - iree_hal_buffer_view_shape_dims(buffer_view), - iree_hal_buffer_view_shape_rank(buffer_view), - iree_hal_buffer_view_element_type(buffer_view), max_element_count, - buffer ? buffer_capacity - buffer_length : 0, - buffer ? buffer + buffer_length : NULL, &elements_length); - buffer_length += elements_length; - iree_hal_buffer_unmap_range(&buffer_mapping); - if (iree_status_is_out_of_range(status)) { - status = iree_status_ignore(status); - buffer = NULL; - } else if (!iree_status_is_ok(status)) { - return status; - } - - if (out_buffer_length) { - *out_buffer_length = buffer_length; - } - return buffer ? iree_ok_status() - : iree_status_from_code(IREE_STATUS_OUT_OF_RANGE); -} - -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_format( - const iree_hal_buffer_view_t* buffer_view, - iree_host_size_t max_element_count, iree_host_size_t buffer_capacity, - char* buffer, iree_host_size_t* out_buffer_length) { - IREE_ASSERT_ARGUMENT(buffer_view); - IREE_TRACE_ZONE_BEGIN(z0); - iree_status_t status = iree_hal_buffer_view_format_impl( - buffer_view, max_element_count, buffer_capacity, buffer, - out_buffer_length); - IREE_TRACE_ZONE_END(z0); - return status; -} - -// TODO(benvanik): streaming all the way down (needs string_util updates). -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_fprint( - FILE* file, const iree_hal_buffer_view_t* buffer_view, - iree_host_size_t max_element_count) { - IREE_ASSERT_ARGUMENT(file); - IREE_ASSERT_ARGUMENT(buffer_view); - IREE_TRACE_ZONE_BEGIN(z0); - - // Query the string length (in characters). - iree_host_size_t buffer_length = 0; - iree_status_t status = iree_hal_buffer_view_format( - buffer_view, max_element_count, 0, NULL, &buffer_length); - if (!iree_status_is_out_of_range(status)) { - IREE_TRACE_ZONE_END(z0); - return status; - } - - // Allocate scratch space to format in to. - // We should be streaming. - iree_allocator_t host_allocator = buffer_view->host_allocator; - iree_host_size_t buffer_capacity = buffer_length + 1; // NUL - char* buffer = NULL; - status = - iree_allocator_malloc(host_allocator, buffer_capacity, (void**)&buffer); - - // Format the buffer into the string storage. - if (iree_status_is_ok(status)) { - status = - iree_hal_buffer_view_format(buffer_view, max_element_count, - buffer_capacity, buffer, &buffer_length); - } - - // Dump to the file. - if (iree_status_is_ok(status)) { - fprintf(file, "%.*s", (int)buffer_length, buffer); - } - - iree_allocator_free(host_allocator, buffer); - IREE_TRACE_ZONE_END(z0); - return status; -}
diff --git a/iree/hal/buffer_view.h b/iree/hal/buffer_view.h index 9a0c623..a6ee1d6 100644 --- a/iree/hal/buffer_view.h +++ b/iree/hal/buffer_view.h
@@ -7,9 +7,7 @@ #ifndef IREE_HAL_BUFFER_VIEW_H_ #define IREE_HAL_BUFFER_VIEW_H_ -#include <stdbool.h> #include <stdint.h> -#include <stdio.h> #include "iree/base/api.h" #include "iree/hal/buffer.h" @@ -156,35 +154,6 @@ typedef int32_t iree_hal_dim_t; //===----------------------------------------------------------------------===// -// Buffer view math -//===----------------------------------------------------------------------===// - -// Calculates the allocation size of a buffer view. -IREE_API_EXPORT iree_status_t iree_hal_buffer_compute_view_size( - const iree_hal_dim_t* shape, iree_host_size_t shape_rank, - iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, - iree_device_size_t* out_allocation_size); - -// Calculates a byte offset into a buffer at the given indices. -// Only works with densely-packed representations. -IREE_API_EXPORT iree_status_t iree_hal_buffer_compute_view_offset( - const iree_hal_dim_t* shape, iree_host_size_t shape_rank, - iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, const iree_hal_dim_t* indices, - size_t indices_count, iree_device_size_t* out_offset); - -// Calculates a byte range into a buffer of the given contiguous range. -// Only works with densely-packed representations. -IREE_API_EXPORT iree_status_t iree_hal_buffer_compute_view_range( - const iree_hal_dim_t* shape, iree_host_size_t shape_rank, - iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, const iree_hal_dim_t* start_indices, - iree_host_size_t indices_count, const iree_hal_dim_t* lengths, - iree_host_size_t lengths_count, iree_device_size_t* out_start_offset, - iree_device_size_t* out_length); - -//===----------------------------------------------------------------------===// // iree_hal_buffer_view_t //===----------------------------------------------------------------------===// @@ -204,69 +173,6 @@ iree_hal_encoding_type_t encoding_type, iree_allocator_t host_allocator, iree_hal_buffer_view_t** out_buffer_view); -// Allocates a buffer from |allocator| and wraps it in a buffer view. -// This is equivalent to: -// 1. iree_hal_buffer_compute_view_size -// 2. iree_hal_allocator_allocate_buffer -// 3. iree_hal_buffer_view_create -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_allocate_buffer( - iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, - iree_host_size_t shape_rank, iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, - iree_hal_buffer_usage_t allowed_usage, - iree_hal_buffer_view_t** out_buffer_view); - -// Clones a host buffer using |allocator| and wraps it in a buffer view. -// This is equivalent to: -// 1. iree_hal_allocator_allocate_buffer -// 2. iree_hal_buffer_write_data -// 3. iree_hal_buffer_view_create -// -// Always prefer allocating a device buffer and populating it in place. -// If cloning multiple buffers it is better to use iree_hal_command_buffer_ts to -// batch up the memory transfer operations. -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_clone_heap_buffer( - iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, - iree_host_size_t shape_rank, iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, - iree_hal_buffer_usage_t allowed_usage, iree_const_byte_span_t data, - iree_hal_buffer_view_t** out_buffer_view); - -// Imports a host buffer using |allocator| and wraps it in a buffer view. -// This is equivalent to: -// 1. iree_hal_allocator_wrap_buffer -// 2. iree_hal_buffer_view_create -// -// NOTE: not all buffers can be imported and not all allocators support -// importing. See iree_hal_allocator_wrap_buffer for more information. -// Fails if the buffer cannot be imported. -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_wrap_heap_buffer( - iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, - iree_host_size_t shape_rank, iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, - iree_hal_memory_access_t allowed_access, - iree_hal_buffer_usage_t allowed_usage, iree_byte_span_t data, - iree_allocator_t data_allocator, iree_hal_buffer_view_t** out_buffer_view); - -// Tries to import a host buffer using |allocator| and wrap it in a buffer view. -// If the buffer cannot be imported then a new buffer will be allocated and the -// source data will be copied into it. -// This is equivalent to: -// if iree_hal_allocator_query_buffer_compatibility ok: -// 1. iree_hal_allocator_wrap_buffer -// 2. iree_hal_buffer_view_create -// else: -// 1. iree_hal_allocator_allocate_buffer -// 2. iree_hal_buffer_write_data -// 3. iree_hal_buffer_view_create -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_wrap_or_clone_heap_buffer( - iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, - iree_host_size_t shape_rank, iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, - iree_hal_memory_access_t allowed_access, - iree_hal_buffer_usage_t allowed_usage, iree_byte_span_t data, - iree_allocator_t data_allocator, iree_hal_buffer_view_t** out_buffer_view); - // Retains the given |buffer_view| for the caller. IREE_API_EXPORT void iree_hal_buffer_view_retain( iree_hal_buffer_view_t* buffer_view); @@ -351,42 +257,6 @@ const iree_hal_dim_t* lengths, iree_host_size_t lengths_count, iree_device_size_t* out_start_offset, iree_device_size_t* out_length); -// Parses a serialized set of buffer elements in the canonical tensor format -// (the same as produced by iree_hal_buffer_view_format). The underlying buffer -// will be allocated with |buffer_allocator| as a host-local/device-visible -// buffer. -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_parse( - iree_string_view_t value, iree_hal_allocator_t* buffer_allocator, - iree_hal_buffer_view_t** out_buffer_view); - -// TODO(#5413): enum for printing mode (include shape, precision). - -// Converts buffer view elements into a fully-specified string-form format like -// `2x4xi16=[[1 2][3 4]]`. -// -// |max_element_count| can be used to limit the total number of elements printed -// when the count may be large. Elided elements will be replaced with `...`. -// -// |buffer_capacity| defines the size of |buffer| in bytes and -// |out_buffer_length| will return the string length in characters. Returns -// IREE_STATUS_OUT_OF_RANGE if the buffer capacity is insufficient to hold the -// formatted elements and |out_buffer_length| will contain the required size. -// -// Follows the standard API string formatting rules. See iree/base/api.h. -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_format( - const iree_hal_buffer_view_t* buffer_view, - iree_host_size_t max_element_count, iree_host_size_t buffer_capacity, - char* buffer, iree_host_size_t* out_buffer_length); - -// Prints buffer view elements into a fully-specified string-form format like -// `2x4xi16=[[1 2][3 4]]`. -// -// |max_element_count| can be used to limit the total number of elements printed -// when the count may be large. Elided elements will be replaced with `...`. -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_fprint( - FILE* file, const iree_hal_buffer_view_t* buffer_view, - iree_host_size_t max_element_count); - //===----------------------------------------------------------------------===// // iree_hal_buffer_view_t implementation details //===----------------------------------------------------------------------===//
diff --git a/iree/hal/buffer_view_util.c b/iree/hal/buffer_view_util.c new file mode 100644 index 0000000..ae2cdac --- /dev/null +++ b/iree/hal/buffer_view_util.c
@@ -0,0 +1,626 @@ +// Copyright 2020 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/buffer_view_util.h" + +#include <inttypes.h> +#include <stdbool.h> + +#include "iree/base/api.h" +#include "iree/base/tracing.h" +#include "iree/hal/allocator.h" +#include "iree/hal/resource.h" +#include "iree/hal/string_util.h" + +//===----------------------------------------------------------------------===// +// Buffer view math +//===----------------------------------------------------------------------===// + +IREE_API_EXPORT iree_status_t iree_hal_buffer_compute_view_size( + const iree_hal_dim_t* shape, iree_host_size_t shape_rank, + iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, + iree_device_size_t* out_allocation_size) { + IREE_ASSERT_ARGUMENT(!shape_rank || shape); + IREE_ASSERT_ARGUMENT(out_allocation_size); + *out_allocation_size = 0; + + iree_device_size_t byte_length = 0; + + switch (encoding_type) { + case IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR: { + if (IREE_UNLIKELY(iree_hal_element_bit_count(element_type) == 0) || + IREE_UNLIKELY(!iree_hal_element_is_byte_aligned(element_type))) { + return iree_make_status( + IREE_STATUS_INVALID_ARGUMENT, + "opaque and sub-byte aligned element types cannot be indexed"); + } + byte_length = iree_hal_element_dense_byte_count(element_type); + for (iree_host_size_t i = 0; i < shape_rank; ++i) { + byte_length *= shape[i]; + } + break; + } + default: + return iree_make_status(IREE_STATUS_UNIMPLEMENTED, + "unimplemented encoding type size calculation"); + } + + *out_allocation_size = byte_length; + return iree_ok_status(); +} + +IREE_API_EXPORT iree_status_t iree_hal_buffer_compute_view_offset( + const iree_hal_dim_t* shape, iree_host_size_t shape_rank, + iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, const iree_hal_dim_t* indices, + iree_host_size_t indices_count, iree_device_size_t* out_offset) { + IREE_ASSERT_ARGUMENT(shape); + IREE_ASSERT_ARGUMENT(indices); + IREE_ASSERT_ARGUMENT(out_offset); + *out_offset = 0; + if (IREE_UNLIKELY(encoding_type != IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR)) { + return iree_make_status( + IREE_STATUS_INVALID_ARGUMENT, + "only dense encodings support view range computation"); + } else if (IREE_UNLIKELY(iree_hal_element_bit_count(element_type) == 0) || + IREE_UNLIKELY(!iree_hal_element_is_byte_aligned(element_type))) { + return iree_make_status( + IREE_STATUS_INVALID_ARGUMENT, + "opaque and sub-byte aligned element types cannot be indexed"); + } else if (IREE_UNLIKELY(shape_rank != indices_count)) { + return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, + "shape rank/indices mismatch: %zu != %zu", + shape_rank, indices_count); + } + + iree_device_size_t offset = 0; + for (iree_host_size_t i = 0; i < indices_count; ++i) { + if (IREE_UNLIKELY(indices[i] >= shape[i])) { + return iree_make_status(IREE_STATUS_OUT_OF_RANGE, + "index[%zu] out of bounds: %d >= %d", i, + indices[i], shape[i]); + } + iree_device_size_t axis_offset = indices[i]; + for (iree_host_size_t j = i + 1; j < shape_rank; ++j) { + axis_offset *= shape[j]; + } + offset += axis_offset; + } + offset *= iree_hal_element_dense_byte_count(element_type); + + *out_offset = offset; + return iree_ok_status(); +} + +IREE_API_EXPORT iree_status_t iree_hal_buffer_compute_view_range( + const iree_hal_dim_t* shape, iree_host_size_t shape_rank, + iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, const iree_hal_dim_t* start_indices, + iree_host_size_t indices_count, const iree_hal_dim_t* lengths, + iree_host_size_t lengths_count, iree_device_size_t* out_start_offset, + iree_device_size_t* out_length) { + IREE_ASSERT_ARGUMENT(shape); + IREE_ASSERT_ARGUMENT(start_indices); + IREE_ASSERT_ARGUMENT(lengths); + IREE_ASSERT_ARGUMENT(out_start_offset); + IREE_ASSERT_ARGUMENT(out_length); + *out_start_offset = 0; + *out_length = 0; + if (IREE_UNLIKELY(encoding_type != IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR)) { + return iree_make_status( + IREE_STATUS_INVALID_ARGUMENT, + "only dense encodings support view range computation"); + } else if (IREE_UNLIKELY(iree_hal_element_bit_count(element_type) == 0) || + IREE_UNLIKELY(!iree_hal_element_is_byte_aligned(element_type))) { + return iree_make_status( + IREE_STATUS_INVALID_ARGUMENT, + "opaque and sub-byte aligned element types cannot be indexed"); + } else if (IREE_UNLIKELY(indices_count != lengths_count)) { + return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, + "indices/lengths mismatch: %zu != %zu", + indices_count, lengths_count); + } else if (IREE_UNLIKELY(shape_rank != indices_count)) { + return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, + "shape rank/indices mismatch: %zu != %zu", + shape_rank, indices_count); + } + + iree_hal_dim_t* end_indices = + iree_alloca(shape_rank * sizeof(iree_hal_dim_t)); + iree_device_size_t element_size = + iree_hal_element_dense_byte_count(element_type); + iree_device_size_t subspan_length = element_size; + for (iree_host_size_t i = 0; i < lengths_count; ++i) { + subspan_length *= lengths[i]; + end_indices[i] = start_indices[i] + lengths[i] - 1; + } + + iree_device_size_t start_byte_offset = 0; + IREE_RETURN_IF_ERROR(iree_hal_buffer_compute_view_offset( + shape, shape_rank, element_type, encoding_type, start_indices, + indices_count, &start_byte_offset)); + iree_device_size_t end_byte_offset = 0; + IREE_RETURN_IF_ERROR(iree_hal_buffer_compute_view_offset( + shape, shape_rank, element_type, encoding_type, end_indices, shape_rank, + &end_byte_offset)); + + // Non-contiguous regions not yet implemented. Will be easier to detect when + // we have strides. + iree_device_size_t offset_length = + end_byte_offset - start_byte_offset + element_size; + if (subspan_length != offset_length) { + return iree_make_status( + IREE_STATUS_UNIMPLEMENTED, + "non-contiguous range region computation not implemented"); + } + + *out_start_offset = start_byte_offset; + *out_length = subspan_length; + return iree_ok_status(); +} + +//===----------------------------------------------------------------------===// +// Buffer view allocation and generation +//===----------------------------------------------------------------------===// + +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_allocate_buffer( + iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, + iree_host_size_t shape_rank, iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, + iree_hal_buffer_usage_t allowed_usage, iree_const_byte_span_t initial_data, + iree_hal_buffer_view_t** out_buffer_view) { + IREE_ASSERT_ARGUMENT(allocator); + IREE_ASSERT_ARGUMENT(out_buffer_view); + IREE_TRACE_ZONE_BEGIN(z0); + + iree_device_size_t allocation_size = 0; + iree_status_t status = iree_hal_buffer_compute_view_size( + shape, shape_rank, element_type, encoding_type, &allocation_size); + + iree_hal_buffer_t* buffer = NULL; + if (iree_status_is_ok(status)) { + status = iree_hal_allocator_allocate_buffer(allocator, memory_type, + allowed_usage, allocation_size, + initial_data, &buffer); + } + + if (iree_status_is_ok(status)) { + status = iree_hal_buffer_view_create( + buffer, shape, shape_rank, element_type, encoding_type, + iree_hal_allocator_host_allocator(allocator), out_buffer_view); + } + + iree_hal_buffer_release(buffer); + IREE_TRACE_ZONE_END(z0); + return status; +} + +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_wrap_heap_buffer( + iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, + iree_host_size_t shape_rank, iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, + iree_hal_memory_access_t allowed_access, + iree_hal_buffer_usage_t allowed_usage, iree_byte_span_t data, + iree_allocator_t data_allocator, iree_hal_buffer_view_t** out_buffer_view) { + IREE_ASSERT_ARGUMENT(allocator); + IREE_ASSERT_ARGUMENT(out_buffer_view); + IREE_TRACE_ZONE_BEGIN(z0); + + // NOTE: this will fail if the data cannot be imported into the allocator. + iree_hal_buffer_t* buffer = NULL; + iree_status_t status = iree_hal_allocator_wrap_buffer( + allocator, memory_type, allowed_access, allowed_usage, data, + data_allocator, &buffer); + + if (iree_status_is_ok(status)) { + status = iree_hal_buffer_view_create( + buffer, shape, shape_rank, element_type, encoding_type, + iree_hal_allocator_host_allocator(allocator), out_buffer_view); + } + + iree_hal_buffer_release(buffer); + IREE_TRACE_ZONE_END(z0); + return status; +} + +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_wrap_or_clone_heap_buffer( + iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, + iree_host_size_t shape_rank, iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, + iree_hal_memory_access_t allowed_access, + iree_hal_buffer_usage_t allowed_usage, iree_byte_span_t data, + iree_allocator_t data_allocator, iree_hal_buffer_view_t** out_buffer_view) { + IREE_ASSERT_ARGUMENT(allocator); + + // Not all HAL implementations support wrapping buffers, and of those that do + // some may only support it in special situations such as when the buffer is + // not DEVICE_VISIBLE. The user application can query whether the wrapping is + // possible and decide to use alternative means of upload if it is not; we + // make no policy (other than validity) over what's best here. + iree_hal_buffer_compatibility_t compatibility = + iree_hal_allocator_query_buffer_compatibility( + allocator, memory_type, allowed_usage, IREE_HAL_BUFFER_USAGE_MAPPING, + (iree_device_size_t)data.data_length); + bool wrap_allowed = iree_all_bits_set( + compatibility, IREE_HAL_BUFFER_COMPATIBILITY_IMPORTABLE); + if (wrap_allowed) { + return iree_hal_buffer_view_wrap_heap_buffer( + allocator, shape, shape_rank, element_type, encoding_type, memory_type, + allowed_access, allowed_usage, data, data_allocator, out_buffer_view); + } else { + return iree_hal_buffer_view_allocate_buffer( + allocator, shape, shape_rank, element_type, encoding_type, memory_type, + allowed_usage, iree_make_const_byte_span(data.data, data.data_length), + out_buffer_view); + } +} + +static iree_status_t iree_hal_buffer_view_generate_buffer_in_situ( + iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, + iree_host_size_t shape_rank, iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, + iree_hal_buffer_usage_t allowed_usage, + iree_hal_buffer_view_generator_callback_t callback, void* user_data, + iree_hal_buffer_view_t** out_buffer_view) { + // Allocate the buffer view and entire buffer contents with the target memory + // type and the mapping bits. + iree_hal_buffer_view_t* buffer_view = NULL; + IREE_RETURN_IF_ERROR(iree_hal_buffer_view_allocate_buffer( + allocator, shape, shape_rank, element_type, encoding_type, memory_type, + allowed_usage | IREE_HAL_BUFFER_USAGE_MAPPING, + iree_const_byte_span_empty(), &buffer_view)); + + // Map the buffer into host-visible memory. + iree_hal_buffer_mapping_t buffer_mapping = {{0}}; + iree_status_t status = iree_hal_buffer_map_range( + iree_hal_buffer_view_buffer(buffer_view), IREE_HAL_MAPPING_MODE_SCOPED, + IREE_HAL_MEMORY_ACCESS_DISCARD_WRITE, 0, IREE_WHOLE_BUFFER, + &buffer_mapping); + + // Generate using the callback directly into the buffer. + if (iree_status_is_ok(status)) { + status = callback(&buffer_mapping, user_data); + } + + status = + iree_status_join(status, iree_hal_buffer_unmap_range(&buffer_mapping)); + if (iree_status_is_ok(status)) { + *out_buffer_view = buffer_view; + } else { + iree_hal_buffer_view_release(buffer_view); + } + return status; +} + +static iree_status_t iree_hal_buffer_view_generate_buffer_on_host( + iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, + iree_host_size_t shape_rank, iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, + iree_hal_buffer_usage_t allowed_usage, iree_device_size_t allocation_size, + iree_hal_buffer_view_generator_callback_t callback, void* user_data, + iree_hal_buffer_view_t** out_buffer_view) { + // Allocate the host memory and generate the contents. + iree_allocator_t host_allocator = + iree_hal_allocator_host_allocator(allocator); + void* host_ptr = NULL; + IREE_RETURN_IF_ERROR( + iree_allocator_malloc(host_allocator, allocation_size, &host_ptr)); + iree_hal_buffer_mapping_t mapping = { + .contents = iree_make_byte_span(host_ptr, allocation_size), + }; + iree_status_t status = callback(&mapping, user_data); + if (!iree_status_is_ok(status)) { + iree_allocator_free(host_allocator, host_ptr); + return status; + } + + // Try to wrap the host allocation to avoid the extra allocation and copy - + // this call will either hang on to the memory or do the copy and immediately + // free it. + return iree_hal_buffer_view_wrap_or_clone_heap_buffer( + allocator, shape, shape_rank, element_type, encoding_type, memory_type, + IREE_HAL_MEMORY_ACCESS_ALL, allowed_usage, + iree_make_byte_span(host_ptr, allocation_size), host_allocator, + out_buffer_view); +} + +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_generate_buffer( + iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, + iree_host_size_t shape_rank, iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, + iree_hal_buffer_usage_t allowed_usage, + iree_hal_buffer_view_generator_callback_t callback, void* user_data, + iree_hal_buffer_view_t** out_buffer_view) { + IREE_ASSERT_ARGUMENT(allocator); + IREE_ASSERT_ARGUMENT(callback); + IREE_ASSERT_ARGUMENT(out_buffer_view); + IREE_TRACE_ZONE_BEGIN(z0); + + // Compute how large of an allocation we need to hold the whole view. + iree_device_size_t allocation_size = 0; + IREE_RETURN_AND_END_ZONE_IF_ERROR( + z0, iree_hal_buffer_compute_view_size(shape, shape_rank, element_type, + encoding_type, &allocation_size)); + + // If we can create the requested memory type with mapping then we'll do that + // and avoid needing to allocate the staging memory. If we can't get that + // memory type (or the allocator doesn't want us using it) then we'll fall + // back to allocation -> generation -> copy. + iree_hal_buffer_compatibility_t compatibility = + iree_hal_allocator_query_buffer_compatibility( + allocator, memory_type, allowed_usage, IREE_HAL_BUFFER_USAGE_MAPPING, + allocation_size); + bool is_mappable = iree_all_bits_set( + compatibility, IREE_HAL_BUFFER_COMPATIBILITY_ALLOCATABLE); + + iree_status_t status = iree_ok_status(); + if (is_mappable) { + // Compatible with allocate -> map -> generate. + status = iree_hal_buffer_view_generate_buffer_in_situ( + allocator, shape, shape_rank, element_type, encoding_type, memory_type, + allowed_usage, callback, user_data, out_buffer_view); + } else { + // Allocate host-local memory first and generate into that. + status = iree_hal_buffer_view_generate_buffer_on_host( + allocator, shape, shape_rank, element_type, encoding_type, memory_type, + allowed_usage, allocation_size, callback, user_data, out_buffer_view); + } + + IREE_TRACE_ZONE_END(z0); + return status; +} + +//===----------------------------------------------------------------------===// +// Buffer view parsing and printing +//===----------------------------------------------------------------------===// + +typedef struct iree_hal_buffer_view_parse_params_t { + iree_string_view_t data_str; + iree_hal_element_type_t element_type; +} iree_hal_buffer_view_parse_params_t; +static iree_status_t iree_hal_buffer_view_parse_into( + iree_hal_buffer_mapping_t* mapping, void* user_data) { + iree_hal_buffer_view_parse_params_t* params = + (iree_hal_buffer_view_parse_params_t*)user_data; + return iree_hal_parse_buffer_elements(params->data_str, params->element_type, + mapping->contents); +} + +static iree_status_t iree_hal_buffer_view_parse_impl( + iree_string_view_t value, iree_hal_allocator_t* buffer_allocator, + iree_hal_buffer_view_t** out_buffer_view) { + // Strip whitespace that may come along (linefeeds/etc). + value = iree_string_view_trim(value); + value = iree_string_view_strip_prefix(value, IREE_SV("\"")); + value = iree_string_view_strip_suffix(value, IREE_SV("\"")); + if (iree_string_view_is_empty(value)) { + // Empty lines are invalid; need at least the shape/type information. + *out_buffer_view = NULL; + return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, "empty string input"); + } + + // The part of the string corresponding to the shape, e.g. 1x2x3. + iree_string_view_t shape_str = iree_string_view_empty(); + // The part of the string corresponding to the type, e.g. f32 + iree_string_view_t type_str = iree_string_view_empty(); + // The part of the string corresponding to the buffer data, e.g. 1 2 3 4 5 6 + iree_string_view_t data_str = iree_string_view_empty(); + + iree_string_view_t shape_and_type_str = value; + iree_string_view_split(value, '=', &shape_and_type_str, &data_str); + iree_host_size_t last_x_index = iree_string_view_find_last_of( + shape_and_type_str, IREE_SV("x"), IREE_STRING_VIEW_NPOS); + if (last_x_index == IREE_STRING_VIEW_NPOS) { + // Scalar. + type_str = shape_and_type_str; + } else { + // Has a shape. + shape_str = iree_string_view_substr(shape_and_type_str, 0, last_x_index); + type_str = iree_string_view_substr(shape_and_type_str, last_x_index + 1, + IREE_STRING_VIEW_NPOS); + } + + // AxBxC... + iree_host_size_t shape_rank = 0; + iree_status_t shape_result = + iree_hal_parse_shape(shape_str, 0, NULL, &shape_rank); + if (!iree_status_is_ok(shape_result) && + !iree_status_is_out_of_range(shape_result)) { + return shape_result; + } else if (shape_rank > 128) { + return iree_make_status( + IREE_STATUS_RESOURCE_EXHAUSTED, + "a shape rank of %zu is just a little bit excessive, eh?", shape_rank); + } + shape_result = iree_status_ignore(shape_result); + iree_hal_dim_t* shape = + (iree_hal_dim_t*)iree_alloca(shape_rank * sizeof(iree_hal_dim_t)); + IREE_RETURN_IF_ERROR( + iree_hal_parse_shape(shape_str, shape_rank, shape, &shape_rank)); + + // f32, i32, etc + iree_hal_element_type_t element_type = IREE_HAL_ELEMENT_TYPE_NONE; + IREE_RETURN_IF_ERROR(iree_hal_parse_element_type(type_str, &element_type)); + + // TODO(benvanik): allow specifying the encoding. + iree_hal_encoding_type_t encoding_type = + IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR; + + // Allocate the buffer from the provided allocator and parse directly into it. + iree_hal_buffer_view_parse_params_t params = { + .data_str = data_str, + .element_type = element_type, + }; + return iree_hal_buffer_view_generate_buffer( + buffer_allocator, shape, shape_rank, element_type, encoding_type, + IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL | IREE_HAL_MEMORY_TYPE_HOST_VISIBLE, + IREE_HAL_BUFFER_USAGE_DISPATCH | IREE_HAL_BUFFER_USAGE_TRANSFER | + IREE_HAL_BUFFER_USAGE_MAPPING, + iree_hal_buffer_view_parse_into, ¶ms, out_buffer_view); +} + +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_parse( + iree_string_view_t value, iree_hal_allocator_t* buffer_allocator, + iree_hal_buffer_view_t** out_buffer_view) { + IREE_ASSERT_ARGUMENT(buffer_allocator); + IREE_ASSERT_ARGUMENT(out_buffer_view); + *out_buffer_view = NULL; + IREE_TRACE_ZONE_BEGIN(z0); + iree_status_t status = + iree_hal_buffer_view_parse_impl(value, buffer_allocator, out_buffer_view); + IREE_TRACE_ZONE_END(z0); + return status; +} + +#define APPEND_CHAR(c) \ + { \ + if (buffer) { \ + if (buffer_length < buffer_capacity - 1) { \ + buffer[buffer_length] = c; \ + buffer[buffer_length + 1] = '\0'; \ + } else { \ + buffer = NULL; \ + } \ + } \ + ++buffer_length; \ + } + +static iree_status_t iree_hal_buffer_view_format_impl( + const iree_hal_buffer_view_t* buffer_view, + iree_host_size_t max_element_count, iree_host_size_t buffer_capacity, + char* buffer, iree_host_size_t* out_buffer_length) { + if (out_buffer_length) { + *out_buffer_length = 0; + } + if (buffer && buffer_capacity) { + buffer[0] = 0; + } + + iree_host_size_t buffer_length = 0; + if (iree_hal_buffer_view_shape_rank(buffer_view) > 0) { + // Shape: 1x2x3 + iree_host_size_t shape_length = 0; + iree_status_t status = iree_hal_format_shape( + iree_hal_buffer_view_shape_dims(buffer_view), + iree_hal_buffer_view_shape_rank(buffer_view), + buffer ? buffer_capacity - buffer_length : 0, + buffer ? buffer + buffer_length : NULL, &shape_length); + buffer_length += shape_length; + if (iree_status_is_out_of_range(status)) { + status = iree_status_ignore(status); + buffer = NULL; + } else if (!iree_status_is_ok(status)) { + return status; + } + + // Separator: <shape>x<format> + APPEND_CHAR('x'); + } + + // Element type: f32 + iree_host_size_t element_type_length = 0; + iree_status_t status = iree_hal_format_element_type( + iree_hal_buffer_view_element_type(buffer_view), + buffer ? buffer_capacity - buffer_length : 0, + buffer ? buffer + buffer_length : NULL, &element_type_length); + buffer_length += element_type_length; + if (iree_status_is_out_of_range(status)) { + status = iree_status_ignore(status); + buffer = NULL; + } else if (!iree_status_is_ok(status)) { + return status; + } + + // TODO(benvanik): allow printing the encoding. + + // Separator: <meta>=<value> + APPEND_CHAR('='); + + // Buffer contents: 0 1 2 3 ... + iree_hal_buffer_mapping_t buffer_mapping = {{0}}; + IREE_RETURN_IF_ERROR(iree_hal_buffer_map_range( + iree_hal_buffer_view_buffer(buffer_view), IREE_HAL_MAPPING_MODE_SCOPED, + IREE_HAL_MEMORY_ACCESS_READ, 0, IREE_WHOLE_BUFFER, &buffer_mapping)); + iree_host_size_t elements_length = 0; + status = iree_hal_format_buffer_elements( + iree_make_const_byte_span(buffer_mapping.contents.data, + buffer_mapping.contents.data_length), + iree_hal_buffer_view_shape_dims(buffer_view), + iree_hal_buffer_view_shape_rank(buffer_view), + iree_hal_buffer_view_element_type(buffer_view), max_element_count, + buffer ? buffer_capacity - buffer_length : 0, + buffer ? buffer + buffer_length : NULL, &elements_length); + buffer_length += elements_length; + status = + iree_status_join(status, iree_hal_buffer_unmap_range(&buffer_mapping)); + if (iree_status_is_out_of_range(status)) { + status = iree_status_ignore(status); + buffer = NULL; + } else if (!iree_status_is_ok(status)) { + return status; + } + + if (out_buffer_length) { + *out_buffer_length = buffer_length; + } + return buffer ? iree_ok_status() + : iree_status_from_code(IREE_STATUS_OUT_OF_RANGE); +} + +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_format( + const iree_hal_buffer_view_t* buffer_view, + iree_host_size_t max_element_count, iree_host_size_t buffer_capacity, + char* buffer, iree_host_size_t* out_buffer_length) { + IREE_ASSERT_ARGUMENT(buffer_view); + IREE_TRACE_ZONE_BEGIN(z0); + iree_status_t status = iree_hal_buffer_view_format_impl( + buffer_view, max_element_count, buffer_capacity, buffer, + out_buffer_length); + IREE_TRACE_ZONE_END(z0); + return status; +} + +// TODO(benvanik): streaming all the way down (needs string_util updates). +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_fprint( + FILE* file, const iree_hal_buffer_view_t* buffer_view, + iree_host_size_t max_element_count, iree_allocator_t host_allocator) { + IREE_ASSERT_ARGUMENT(file); + IREE_ASSERT_ARGUMENT(buffer_view); + IREE_TRACE_ZONE_BEGIN(z0); + + // Query the string length (in characters). + iree_host_size_t buffer_length = 0; + iree_status_t status = iree_hal_buffer_view_format( + buffer_view, max_element_count, 0, NULL, &buffer_length); + if (!iree_status_is_out_of_range(status)) { + IREE_TRACE_ZONE_END(z0); + return status; + } + + // Allocate scratch space to format in to. + // We should be streaming. + iree_host_size_t buffer_capacity = buffer_length + 1; // NUL + char* buffer = NULL; + status = + iree_allocator_malloc(host_allocator, buffer_capacity, (void**)&buffer); + + // Format the buffer into the string storage. + if (iree_status_is_ok(status)) { + status = + iree_hal_buffer_view_format(buffer_view, max_element_count, + buffer_capacity, buffer, &buffer_length); + } + + // Dump to the file. + if (iree_status_is_ok(status)) { + fprintf(file, "%.*s", (int)buffer_length, buffer); + } + + iree_allocator_free(host_allocator, buffer); + IREE_TRACE_ZONE_END(z0); + return status; +}
diff --git a/iree/hal/buffer_view_util.h b/iree/hal/buffer_view_util.h new file mode 100644 index 0000000..08bf361 --- /dev/null +++ b/iree/hal/buffer_view_util.h
@@ -0,0 +1,184 @@ +// Copyright 2021 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 + +#ifndef IREE_HAL_BUFFER_VIEW_UTIL_H_ +#define IREE_HAL_BUFFER_VIEW_UTIL_H_ + +#include <stdbool.h> +#include <stdint.h> +#include <stdio.h> + +#include "iree/base/api.h" +#include "iree/hal/buffer_view.h" + +#ifdef __cplusplus +extern "C" { +#endif // __cplusplus + +//===----------------------------------------------------------------------===// +// Buffer view math +//===----------------------------------------------------------------------===// + +// Calculates the allocation size of a buffer view. +IREE_API_EXPORT iree_status_t iree_hal_buffer_compute_view_size( + const iree_hal_dim_t* shape, iree_host_size_t shape_rank, + iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, + iree_device_size_t* out_allocation_size); + +// Calculates a byte offset into a buffer at the given indices. +// Only works with densely-packed representations. +IREE_API_EXPORT iree_status_t iree_hal_buffer_compute_view_offset( + const iree_hal_dim_t* shape, iree_host_size_t shape_rank, + iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, const iree_hal_dim_t* indices, + size_t indices_count, iree_device_size_t* out_offset); + +// Calculates a byte range into a buffer of the given contiguous range. +// Only works with densely-packed representations. +IREE_API_EXPORT iree_status_t iree_hal_buffer_compute_view_range( + const iree_hal_dim_t* shape, iree_host_size_t shape_rank, + iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, const iree_hal_dim_t* start_indices, + iree_host_size_t indices_count, const iree_hal_dim_t* lengths, + iree_host_size_t lengths_count, iree_device_size_t* out_start_offset, + iree_device_size_t* out_length); + +//===----------------------------------------------------------------------===// +// Buffer view allocation and generation +//===----------------------------------------------------------------------===// + +// Allocates a buffer from |allocator| and wraps it in a buffer view. +// +// This is equivalent to: +// 1. iree_hal_buffer_compute_view_size +// 2. iree_hal_allocator_allocate_buffer +// 3. iree_hal_buffer_view_create +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_allocate_buffer( + iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, + iree_host_size_t shape_rank, iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, + iree_hal_buffer_usage_t allowed_usage, iree_const_byte_span_t initial_data, + iree_hal_buffer_view_t** out_buffer_view); + +// Imports a host buffer using |allocator| and wraps it in a buffer view. +// +// This is equivalent to: +// 1. iree_hal_allocator_wrap_buffer +// 2. iree_hal_buffer_view_create +// +// NOTE: not all buffers can be imported and not all allocators support +// importing. See iree_hal_allocator_wrap_buffer for more information. +// Fails if the buffer cannot be imported. +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_wrap_heap_buffer( + iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, + iree_host_size_t shape_rank, iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, + iree_hal_memory_access_t allowed_access, + iree_hal_buffer_usage_t allowed_usage, iree_byte_span_t data, + iree_allocator_t data_allocator, iree_hal_buffer_view_t** out_buffer_view); + +// Tries to import a host buffer using |allocator| and wrap it in a buffer view. +// If the buffer cannot be imported then a new buffer will be allocated and the +// source data will be copied into it. +// +// This is equivalent to: +// if iree_hal_allocator_query_buffer_compatibility ok: +// 1. iree_hal_allocator_wrap_buffer +// 2. iree_hal_buffer_view_create +// else: +// 1. iree_hal_allocator_allocate_buffer +// 2. iree_hal_buffer_write_data +// 3. iree_hal_buffer_view_create +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_wrap_or_clone_heap_buffer( + iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, + iree_host_size_t shape_rank, iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, + iree_hal_memory_access_t allowed_access, + iree_hal_buffer_usage_t allowed_usage, iree_byte_span_t data, + iree_allocator_t data_allocator, iree_hal_buffer_view_t** out_buffer_view); + +typedef iree_status_t(IREE_API_PTR* iree_hal_buffer_view_generator_callback_t)( + iree_hal_buffer_mapping_t* mapping, void* user_data); + +// Generates a buffer view with its initial contents produced by a callback. +// When host and device memory are shared this allows direct generation into the +// target device buffer. If not shared this can avoid expensive transfer mapping +// operations at the cost of a transient host memory allocation. The mapped host +// pointer passed to the callback is only valid within the callback. +// +// Buffers allocated like this do not need the IREE_HAL_BUFFER_USAGE_MAPPING bit +// set; it will be added automatically if the allocator needs it and otherwise +// the memory can remain unmappable (and thus fully device isolated). +// +// As this _may_ require allocation of the entire buffer content in host memory +// it is always preferable to stage and issue copy commands via the device +// queue. Even better is to do all generation on-device via dispatches without +// the need to ever transfer. Usage of this method should be limited to times +// where device-side generation isn't possible or memory consumption is not a +// concern. +// +// This is equivalent to: +// 1. iree_hal_buffer_compute_view_size +// 2. iree_hal_allocator_allocate_buffer +// 3. iree_hal_buffer_map_range + callback + iree_hal_buffer_unmap_range +// 4. iree_hal_buffer_view_create +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_generate_buffer( + iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, + iree_host_size_t shape_rank, iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, + iree_hal_buffer_usage_t allowed_usage, + iree_hal_buffer_view_generator_callback_t callback, void* user_data, + iree_hal_buffer_view_t** out_buffer_view); + +//===----------------------------------------------------------------------===// +// Buffer view parsing and printing +//===----------------------------------------------------------------------===// + +// Parses a serialized set of buffer elements in the canonical tensor format +// (the same as produced by iree_hal_buffer_view_format). The underlying buffer +// will be allocated with |buffer_allocator| as a host-local/device-visible +// buffer. +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_parse( + iree_string_view_t value, iree_hal_allocator_t* buffer_allocator, + iree_hal_buffer_view_t** out_buffer_view); + +// TODO(#5413): enum for printing mode (include shape, precision). + +// Converts buffer view elements into a fully-specified string-form format like +// `2x4xi16=[[1 2][3 4]]`. +// +// |max_element_count| can be used to limit the total number of elements printed +// when the count may be large. Elided elements will be replaced with `...`. +// +// |buffer_capacity| defines the size of |buffer| in bytes and +// |out_buffer_length| will return the string length in characters. Returns +// IREE_STATUS_OUT_OF_RANGE if the buffer capacity is insufficient to hold the +// formatted elements and |out_buffer_length| will contain the required size. +// +// Follows the standard API string formatting rules. See iree/base/api.h. +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_format( + const iree_hal_buffer_view_t* buffer_view, + iree_host_size_t max_element_count, iree_host_size_t buffer_capacity, + char* buffer, iree_host_size_t* out_buffer_length); + +// Prints buffer view elements into a fully-specified string-form format like +// `2x4xi16=[[1 2][3 4]]`. +// +// |max_element_count| can be used to limit the total number of elements printed +// when the count may be large. Elided elements will be replaced with `...`. +// +// |host_allocator| will be used for any transient allocations required while +// printing. +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_fprint( + FILE* file, const iree_hal_buffer_view_t* buffer_view, + iree_host_size_t max_element_count, iree_allocator_t host_allocator); + +#ifdef __cplusplus +} // extern "C" +#endif // __cplusplus + +#endif // IREE_HAL_BUFFER_VIEW_UTIL_H_
diff --git a/iree/hal/command_buffer.c b/iree/hal/command_buffer.c index b34df20..e4c7fdf 100644 --- a/iree/hal/command_buffer.c +++ b/iree/hal/command_buffer.c
@@ -454,3 +454,70 @@ IREE_TRACE_ZONE_END(z0); return status; } + +//===----------------------------------------------------------------------===// +// Utilities for command buffer creation +//===----------------------------------------------------------------------===// + +IREE_API_EXPORT iree_status_t iree_hal_create_transfer_command_buffer( + iree_hal_device_t* device, iree_hal_command_buffer_mode_t mode, + iree_hal_queue_affinity_t queue_affinity, iree_host_size_t transfer_count, + const iree_hal_transfer_command_t* transfer_commands, + iree_hal_command_buffer_t** out_command_buffer) { + IREE_TRACE_ZONE_BEGIN(z0); + + iree_hal_command_buffer_t* command_buffer = NULL; + IREE_RETURN_AND_END_ZONE_IF_ERROR( + z0, iree_hal_command_buffer_create(device, mode, + IREE_HAL_COMMAND_CATEGORY_TRANSFER, + queue_affinity, &command_buffer)); + + iree_status_t status = iree_hal_command_buffer_begin(command_buffer); + if (iree_status_is_ok(status)) { + for (iree_host_size_t i = 0; i < transfer_count; ++i) { + const iree_hal_transfer_command_t* transfer_command = + &transfer_commands[i]; + switch (transfer_command->type) { + case IREE_HAL_TRANSFER_COMMAND_TYPE_FILL: + status = iree_hal_command_buffer_fill_buffer( + command_buffer, transfer_command->fill.target_buffer, + transfer_command->fill.target_offset, + transfer_command->fill.length, transfer_command->fill.pattern, + transfer_command->fill.pattern_length); + break; + case IREE_HAL_TRANSFER_COMMAND_TYPE_COPY: + status = iree_hal_command_buffer_copy_buffer( + command_buffer, transfer_command->copy.source_buffer, + transfer_command->copy.source_offset, + transfer_command->copy.target_buffer, + transfer_command->copy.target_offset, + transfer_command->copy.length); + break; + case IREE_HAL_TRANSFER_COMMAND_TYPE_UPDATE: + status = iree_hal_command_buffer_update_buffer( + command_buffer, transfer_command->update.source_buffer, + transfer_command->update.source_offset, + transfer_command->update.target_buffer, + transfer_command->update.target_offset, + transfer_command->update.length); + break; + default: + status = iree_make_status(IREE_STATUS_INVALID_ARGUMENT, + "unknown transfer_commands[%zu] type %d", i, + (int)transfer_command->type); + break; + } + if (!iree_status_is_ok(status)) break; + } + } + status = + iree_status_join(status, iree_hal_command_buffer_end(command_buffer)); + + if (iree_status_is_ok(status)) { + *out_command_buffer = command_buffer; + } else { + iree_hal_command_buffer_release(command_buffer); + } + IREE_TRACE_ZONE_END(z0); + return status; +}
diff --git a/iree/hal/command_buffer.h b/iree/hal/command_buffer.h index e9d2c69..8f6d8fe 100644 --- a/iree/hal/command_buffer.h +++ b/iree/hal/command_buffer.h
@@ -229,6 +229,12 @@ // TODO(benvanik): valid push constant bit ranges. } iree_hal_command_buffer_validation_state_t; +// Maximum size of any update in iree_hal_command_buffer_update_buffer. +// 64KB is the limit on Vulkan and we uniformly use that today across all +// targets as to not need too much command buffer memory. +#define IREE_HAL_COMMAND_BUFFER_MAX_UPDATE_SIZE \ + ((iree_device_size_t)(64 * 1024)) + //===----------------------------------------------------------------------===// // iree_hal_command_buffer_t //===----------------------------------------------------------------------===// @@ -237,11 +243,13 @@ // Commands are recorded by the implementation for later submission to command // queues. // -// Buffers and synchronization objects referenced must remain valid and not be +// Buffers, events, and programs referenced must remain valid and not be // modified or read while there are commands in-flight. The usual flow is to -// populate input buffers, Dispatch using those buffers, wait on a Semaphore -// until the buffers are guaranteed to no longer be in use, and then reuse or -// release the buffers. +// populate input buffers, dispatch using those buffers, wait on a semaphore +// until the buffers are guaranteed to no longer be in use, and then reuse the +// buffers. Lifetimes are managed by the command buffer and all used resources +// will be retained for as long as the command buffer is live or until it is +// reset. // // Errors that can be recognized when operations are enqueued will be returned // immediately, such as invalid argument errors. Errors that can only be @@ -493,6 +501,68 @@ iree_hal_buffer_t* workgroups_buffer, iree_device_size_t workgroups_offset); //===----------------------------------------------------------------------===// +// Utilities for command buffer creation +//===----------------------------------------------------------------------===// + +// Defines a transfer command operation. +typedef enum iree_hal_transfer_command_type_t { + // iree_hal_command_buffer_fill_buffer + IREE_HAL_TRANSFER_COMMAND_TYPE_FILL = 0u, + // iree_hal_command_buffer_copy_buffer + IREE_HAL_TRANSFER_COMMAND_TYPE_COPY = 1u, + // iree_hal_command_buffer_update_buffer + IREE_HAL_TRANSFER_COMMAND_TYPE_UPDATE = 2u, +} iree_hal_transfer_command_type_t; + +// Represents a single transfer command within a batch of commands. +typedef struct iree_hal_transfer_command_t { + // The type of the command selecting which of the payload data is used. + iree_hal_transfer_command_type_t type; + union { + // IREE_HAL_TRANSFER_COMMAND_TYPE_FILL + struct { + iree_hal_buffer_t* target_buffer; + iree_device_size_t target_offset; + iree_device_size_t length; + const void* pattern; + iree_host_size_t pattern_length; + } fill; + // IREE_HAL_TRANSFER_COMMAND_TYPE_COPY + struct { + iree_hal_buffer_t* source_buffer; + iree_device_size_t source_offset; + iree_hal_buffer_t* target_buffer; + iree_device_size_t target_offset; + iree_device_size_t length; + } copy; + // IREE_HAL_TRANSFER_COMMAND_TYPE_UPDATE + struct { + const void* source_buffer; + iree_host_size_t source_offset; + iree_hal_buffer_t* target_buffer; + iree_device_size_t target_offset; + iree_device_size_t length; + } update; + }; +} iree_hal_transfer_command_t; + +// Builds a command buffer containing a recording of all |transfer_commands|. +// All buffers must be compatible with |device| and ranges must not overlap +// (same as with memcpy). All commands are executed concurrently with no +// barriers. The provided commands and any referenced data needs only remain +// live during recording, while all referenced buffers must be kept live by +// the caller until the command buffer has completed execution. +// +// This is just a utility to make it easier to quickly construct batches of +// transfer operations. If more control is required then record the command +// buffer as normal. +IREE_API_EXPORT iree_status_t iree_hal_create_transfer_command_buffer( + iree_hal_device_t* device, iree_hal_command_buffer_mode_t mode, + iree_hal_queue_affinity_t queue_affinity, iree_host_size_t transfer_count, + const iree_hal_transfer_command_t* transfer_commands, + iree_hal_command_buffer_t** out_command_buffer); + +//===----------------------------------------------------------------------===// // iree_hal_command_buffer_t validation wrapper //===----------------------------------------------------------------------===// @@ -517,9 +587,6 @@ //===----------------------------------------------------------------------===// typedef struct iree_hal_command_buffer_vtable_t { - // << HAL C porting in progress >> - IREE_API_UNSTABLE - void(IREE_API_PTR* destroy)(iree_hal_command_buffer_t* command_buffer); void*(IREE_API_PTR* dyn_cast)(iree_hal_command_buffer_t* command_buffer, @@ -612,6 +679,7 @@ iree_hal_buffer_t* workgroups_buffer, iree_device_size_t workgroups_offset); } iree_hal_command_buffer_vtable_t; +IREE_HAL_ASSERT_VTABLE_LAYOUT(iree_hal_command_buffer_vtable_t); struct iree_hal_command_buffer_t { iree_hal_resource_t resource;
diff --git a/iree/hal/command_buffer_validation.c b/iree/hal/command_buffer_validation.c index 56e6cc0..871045d 100644 --- a/iree/hal/command_buffer_validation.c +++ b/iree/hal/command_buffer_validation.c
@@ -73,7 +73,8 @@ return iree_make_status( IREE_STATUS_PERMISSION_DENIED, "requested buffer usage is not supported for the buffer on this queue; " - "buffer allows %.*s, operation requires %.*s", + "buffer allows %.*s, operation requires %.*s (allocator compatibility " + "mismatch)", (int)allowed_usage_str.size, allowed_usage_str.data, (int)intended_usage_str.size, intended_usage_str.data); }
diff --git a/iree/hal/cts/allocator_test.h b/iree/hal/cts/allocator_test.h index 506d6ee..d4d4c50 100644 --- a/iree/hal/cts/allocator_test.h +++ b/iree/hal/cts/allocator_test.h
@@ -85,7 +85,8 @@ iree_hal_buffer_t* buffer; IREE_ASSERT_OK(iree_hal_allocator_allocate_buffer( - device_allocator_, memory_type, buffer_usage, kAllocationSize, &buffer)); + device_allocator_, memory_type, buffer_usage, kAllocationSize, + iree_const_byte_span_empty(), &buffer)); // At a mimimum, the requested memory type should be respected. // Additional bits may be optionally set depending on the allocator. @@ -108,7 +109,7 @@ iree_hal_buffer_t* buffer; IREE_ASSERT_OK(iree_hal_allocator_allocate_buffer( device_allocator_, memory_type, buffer_usage, /*allocation_size=*/0, - &buffer)); + iree_const_byte_span_empty(), &buffer)); iree_hal_buffer_release(buffer); }
diff --git a/iree/hal/cts/buffer_mapping_test.h b/iree/hal/cts/buffer_mapping_test.h index 106b24d..f8229c1 100644 --- a/iree/hal/cts/buffer_mapping_test.h +++ b/iree/hal/cts/buffer_mapping_test.h
@@ -46,7 +46,8 @@ iree_hal_buffer_t* buffer; IREE_ASSERT_OK(iree_hal_allocator_allocate_buffer( - device_allocator_, memory_type, buffer_usage, kAllocationSize, &buffer)); + device_allocator_, memory_type, buffer_usage, kAllocationSize, + iree_const_byte_span_empty(), &buffer)); EXPECT_TRUE( iree_all_bits_set(iree_hal_buffer_memory_type(buffer), memory_type)); @@ -63,7 +64,8 @@ iree_hal_buffer_t* buffer; IREE_ASSERT_OK(iree_hal_allocator_allocate_buffer( - device_allocator_, memory_type, buffer_usage, kAllocationSize, &buffer)); + device_allocator_, memory_type, buffer_usage, kAllocationSize, + iree_const_byte_span_empty(), &buffer)); IREE_ASSERT_OK(iree_hal_buffer_zero(buffer, /*byte_offset=*/0, /*byte_length=*/kAllocationSize)); @@ -85,7 +87,8 @@ iree_hal_buffer_t* buffer; IREE_ASSERT_OK(iree_hal_allocator_allocate_buffer( - device_allocator_, memory_type, buffer_usage, kAllocationSize, &buffer)); + device_allocator_, memory_type, buffer_usage, kAllocationSize, + iree_const_byte_span_empty(), &buffer)); IREE_ASSERT_OK(iree_hal_buffer_zero(buffer, /*byte_offset=*/0, /*byte_length=*/kAllocationSize)); @@ -113,7 +116,8 @@ iree_hal_buffer_t* buffer; IREE_ASSERT_OK(iree_hal_allocator_allocate_buffer( - device_allocator_, memory_type, buffer_usage, kAllocationSize, &buffer)); + device_allocator_, memory_type, buffer_usage, kAllocationSize, + iree_const_byte_span_empty(), &buffer)); uint8_t fill_value = 0x07; IREE_ASSERT_OK(iree_hal_buffer_fill(buffer, /*byte_offset=*/0, @@ -138,7 +142,8 @@ iree_hal_buffer_t* buffer; IREE_ASSERT_OK(iree_hal_allocator_allocate_buffer( - device_allocator_, memory_type, buffer_usage, kAllocationSize, &buffer)); + device_allocator_, memory_type, buffer_usage, kAllocationSize, + iree_const_byte_span_empty(), &buffer)); uint8_t fill_value = 0x07; std::vector<uint8_t> reference_buffer(kAllocationSize); @@ -162,10 +167,10 @@ iree_hal_buffer_t* buffer_b; IREE_ASSERT_OK(iree_hal_allocator_allocate_buffer( device_allocator_, memory_type, buffer_usage, kAllocationSize, - &buffer_a)); + iree_const_byte_span_empty(), &buffer_a)); IREE_ASSERT_OK(iree_hal_allocator_allocate_buffer( device_allocator_, memory_type, buffer_usage, kAllocationSize, - &buffer_b)); + iree_const_byte_span_empty(), &buffer_b)); uint8_t fill_value = 0x07; IREE_ASSERT_OK(iree_hal_buffer_fill(buffer_a, /*byte_offset=*/0,
diff --git a/iree/hal/cts/command_buffer_test.h b/iree/hal/cts/command_buffer_test.h index 60df42d..4bd0d7b 100644 --- a/iree/hal/cts/command_buffer_test.h +++ b/iree/hal/cts/command_buffer_test.h
@@ -44,7 +44,8 @@ IREE_CHECK_OK(iree_hal_allocator_allocate_buffer( iree_hal_device_allocator(device_), IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL | IREE_HAL_MEMORY_TYPE_HOST_VISIBLE, - IREE_HAL_BUFFER_USAGE_ALL, buffer_size, &device_buffer)); + IREE_HAL_BUFFER_USAGE_ALL, buffer_size, iree_const_byte_span_empty(), + &device_buffer)); IREE_CHECK_OK(iree_hal_command_buffer_begin(command_buffer)); // Start with a zero fill on the entire buffer... @@ -141,26 +142,28 @@ IREE_HAL_COMMAND_CATEGORY_TRANSFER, IREE_HAL_QUEUE_AFFINITY_ANY, &command_buffer)); + uint8_t i8_val = 0x54; + std::vector<uint8_t> reference_buffer(kBufferSize); + std::memset(reference_buffer.data(), i8_val, kBufferSize); + // Create and fill a host buffer. iree_hal_buffer_t* host_buffer; IREE_ASSERT_OK(iree_hal_allocator_allocate_buffer( device_allocator_, IREE_HAL_MEMORY_TYPE_HOST_VISIBLE | IREE_HAL_MEMORY_TYPE_HOST_CACHED | IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE, - IREE_HAL_BUFFER_USAGE_ALL, kBufferSize, &host_buffer)); - uint8_t i8_val = 0x54; - IREE_ASSERT_OK(iree_hal_buffer_fill(host_buffer, /*byte_offset=*/0, - /*byte_length=*/kBufferSize, &i8_val, - /*pattern_length=*/sizeof(i8_val))); - std::vector<uint8_t> reference_buffer(kBufferSize); - std::memset(reference_buffer.data(), i8_val, kBufferSize); + IREE_HAL_BUFFER_USAGE_ALL, kBufferSize, + iree_make_const_byte_span(reference_buffer.data(), + reference_buffer.size()), + &host_buffer)); // Create a device buffer. iree_hal_buffer_t* device_buffer; IREE_ASSERT_OK(iree_hal_allocator_allocate_buffer( device_allocator_, IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL | IREE_HAL_MEMORY_TYPE_HOST_VISIBLE, - IREE_HAL_BUFFER_USAGE_ALL, kBufferSize, &device_buffer)); + IREE_HAL_BUFFER_USAGE_ALL, kBufferSize, iree_const_byte_span_empty(), + &device_buffer)); // Copy the host buffer to the device buffer. IREE_ASSERT_OK(iree_hal_command_buffer_begin(command_buffer)); @@ -187,33 +190,34 @@ } TEST_P(command_buffer_test, CopySubBuffer) { - iree_hal_command_buffer_t* command_buffer; + iree_hal_command_buffer_t* command_buffer = NULL; IREE_ASSERT_OK(iree_hal_command_buffer_create( device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT, IREE_HAL_COMMAND_CATEGORY_TRANSFER, IREE_HAL_QUEUE_AFFINITY_ANY, &command_buffer)); - iree_hal_buffer_t* device_buffer; + iree_hal_buffer_t* device_buffer = NULL; IREE_ASSERT_OK(iree_hal_allocator_allocate_buffer( device_allocator_, IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL | IREE_HAL_MEMORY_TYPE_HOST_VISIBLE, - IREE_HAL_BUFFER_USAGE_ALL, kBufferSize, &device_buffer)); + IREE_HAL_BUFFER_USAGE_ALL, kBufferSize, iree_const_byte_span_empty(), + &device_buffer)); + + uint8_t i8_val = 0x88; + std::vector<uint8_t> reference_buffer(kBufferSize); + std::memset(reference_buffer.data() + 8, i8_val, kBufferSize / 2 - 4); // Create another host buffer with a smaller size. - iree_hal_buffer_t* host_buffer; + std::vector<uint8_t> host_buffer_data(kBufferSize, i8_val); + iree_hal_buffer_t* host_buffer = NULL; IREE_ASSERT_OK(iree_hal_allocator_allocate_buffer( device_allocator_, IREE_HAL_MEMORY_TYPE_HOST_VISIBLE | IREE_HAL_MEMORY_TYPE_HOST_CACHED | IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE, - IREE_HAL_BUFFER_USAGE_ALL, kBufferSize / 2, &host_buffer)); - - // Fill the host buffer. - uint8_t i8_val = 0x88; - IREE_ASSERT_OK(iree_hal_buffer_fill(host_buffer, /*byte_offset=*/0, - /*byte_length=*/kBufferSize / 2, &i8_val, - /*pattern_length=*/sizeof(i8_val))); - std::vector<uint8_t> reference_buffer(kBufferSize); - std::memset(reference_buffer.data() + 8, i8_val, kBufferSize / 2 - 4); + IREE_HAL_BUFFER_USAGE_ALL, host_buffer_data.size() / 2, + iree_make_const_byte_span(host_buffer_data.data(), + host_buffer_data.size() / 2), + &host_buffer)); // Copy the host buffer to the device buffer; zero fill the untouched bytes. uint8_t zero_val = 0x0;
diff --git a/iree/hal/cts/cts_test_base.h b/iree/hal/cts/cts_test_base.h index 1269646..611926d 100644 --- a/iree/hal/cts/cts_test_base.h +++ b/iree/hal/cts/cts_test_base.h
@@ -41,7 +41,7 @@ iree_status_t status = TryGetDriver(driver_name, &driver); if (iree_status_is_unavailable(status)) { iree_status_free(status); - IREE_LOG(WARNING) << "Skipping test as << '" << driver_name + IREE_LOG(WARNING) << "Skipping test as '" << driver_name << "' driver is unavailable"; GTEST_SKIP(); return;
diff --git a/iree/hal/cuda/BUILD b/iree/hal/cuda/BUILD index c3d6f02..3ab575b 100644 --- a/iree/hal/cuda/BUILD +++ b/iree/hal/cuda/BUILD
@@ -65,7 +65,9 @@ "//iree/base/internal:synchronization", "//iree/base/internal/flatcc:parsing", "//iree/hal", + "//iree/hal/utils:buffer_transfer", "//iree/hal/utils:deferred_command_buffer", + "//iree/hal/utils:resource_set", "//iree/schemas:cuda_executable_def_c_fbs", ], )
diff --git a/iree/hal/cuda/CMakeLists.txt b/iree/hal/cuda/CMakeLists.txt index d3bd43b..6fe645f 100644 --- a/iree/hal/cuda/CMakeLists.txt +++ b/iree/hal/cuda/CMakeLists.txt
@@ -57,7 +57,9 @@ iree::base::internal::synchronization iree::base::tracing iree::hal + iree::hal::utils::buffer_transfer iree::hal::utils::deferred_command_buffer + iree::hal::utils::resource_set iree::schemas::cuda_executable_def_c_fbs PUBLIC )
diff --git a/iree/hal/cuda/cuda_allocator.c b/iree/hal/cuda/cuda_allocator.c index 7470282..6ec25e7 100644 --- a/iree/hal/cuda/cuda_allocator.c +++ b/iree/hal/cuda/cuda_allocator.c
@@ -16,10 +16,11 @@ typedef struct iree_hal_cuda_allocator_t { iree_hal_resource_t resource; + iree_hal_device_t* base_device; iree_hal_cuda_context_wrapper_t* context; CUdevice device; - bool supports_concurrent_managed_access; CUstream stream; + bool supports_concurrent_managed_access; IREE_STATISTICS(iree_hal_allocator_statistics_t statistics;) } iree_hal_cuda_allocator_t; @@ -33,8 +34,9 @@ } iree_status_t iree_hal_cuda_allocator_create( - iree_hal_cuda_context_wrapper_t* context, CUdevice device, CUstream stream, - iree_hal_allocator_t** out_allocator) { + iree_hal_device_t* base_device, iree_hal_cuda_context_wrapper_t* context, + CUdevice device, CUstream stream, iree_hal_allocator_t** out_allocator) { + IREE_ASSERT_ARGUMENT(base_device); IREE_ASSERT_ARGUMENT(context); IREE_TRACE_ZONE_BEGIN(z0); @@ -65,6 +67,7 @@ if (iree_status_is_ok(status)) { iree_hal_resource_initialize(&iree_hal_cuda_allocator_vtable, &allocator->resource); + allocator->base_device = base_device; allocator->context = context; allocator->device = device; allocator->stream = stream; @@ -96,6 +99,11 @@ return allocator->context->host_allocator; } +static iree_status_t iree_hal_cuda_allocator_trim( + iree_hal_allocator_t* base_allocator) { + return iree_ok_status(); +} + static void iree_hal_cuda_allocator_query_statistics( iree_hal_allocator_t* base_allocator, iree_hal_allocator_statistics_t* out_statistics) { @@ -151,7 +159,7 @@ static iree_status_t iree_hal_cuda_allocator_allocate_buffer( iree_hal_allocator_t* base_allocator, iree_hal_memory_type_t memory_type, iree_hal_buffer_usage_t allowed_usage, iree_host_size_t allocation_size, - iree_hal_buffer_t** out_buffer) { + iree_const_byte_span_t initial_data, iree_hal_buffer_t** out_buffer) { iree_hal_cuda_allocator_t* allocator = iree_hal_cuda_allocator_cast(base_allocator); // Guard against the corner case where the requested buffer size is 0. The @@ -210,19 +218,39 @@ cuMemHostGetDevicePointer(&device_ptr, host_ptr, /*flags=*/0)); } } + + iree_hal_buffer_t* buffer = NULL; if (iree_status_is_ok(status)) { status = iree_hal_cuda_buffer_wrap( (iree_hal_allocator_t*)allocator, memory_type, IREE_HAL_MEMORY_ACCESS_ALL, allowed_usage, allocation_size, /*byte_offset=*/0, - /*byte_length=*/allocation_size, device_ptr, host_ptr, out_buffer); + /*byte_length=*/allocation_size, device_ptr, host_ptr, &buffer); } + + // Copy the initial contents into the buffer. This may require staging. + if (iree_status_is_ok(status) && + !iree_const_byte_span_is_empty(initial_data)) { + status = iree_hal_device_transfer_range( + allocator->base_device, + iree_hal_make_host_transfer_buffer_span((void*)initial_data.data, + initial_data.data_length), + 0, iree_hal_make_device_transfer_buffer(buffer), 0, + initial_data.data_length, IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, + iree_infinite_timeout()); + } + if (iree_status_is_ok(status)) { IREE_STATISTICS(iree_hal_allocator_statistics_record_alloc( &allocator->statistics, memory_type, allocation_size)); + *out_buffer = buffer; } else { - iree_hal_cuda_buffer_free(allocator->context, memory_type, device_ptr, - host_ptr); + if (!buffer) { + iree_hal_cuda_buffer_free(allocator->context, memory_type, device_ptr, + host_ptr); + } else { + iree_hal_buffer_release(buffer); + } } return status; } @@ -255,6 +283,7 @@ static const iree_hal_allocator_vtable_t iree_hal_cuda_allocator_vtable = { .destroy = iree_hal_cuda_allocator_destroy, .host_allocator = iree_hal_cuda_allocator_host_allocator, + .trim = iree_hal_cuda_allocator_trim, .query_statistics = iree_hal_cuda_allocator_query_statistics, .query_buffer_compatibility = iree_hal_cuda_allocator_query_buffer_compatibility,
diff --git a/iree/hal/cuda/cuda_allocator.h b/iree/hal/cuda/cuda_allocator.h index 0fe2739..4f22579 100644 --- a/iree/hal/cuda/cuda_allocator.h +++ b/iree/hal/cuda/cuda_allocator.h
@@ -18,8 +18,8 @@ // Create a cuda allocator. iree_status_t iree_hal_cuda_allocator_create( - iree_hal_cuda_context_wrapper_t* context, CUdevice device, CUstream stream, - iree_hal_allocator_t** out_allocator); + iree_hal_device_t* base_device, iree_hal_cuda_context_wrapper_t* context, + CUdevice device, CUstream stream, iree_hal_allocator_t** out_allocator); #ifdef __cplusplus } // extern "C"
diff --git a/iree/hal/cuda/cuda_buffer.c b/iree/hal/cuda/cuda_buffer.c index db2a8aa..b9a56ef 100644 --- a/iree/hal/cuda/cuda_buffer.c +++ b/iree/hal/cuda/cuda_buffer.c
@@ -68,14 +68,16 @@ iree_hal_buffer_t* base_buffer, iree_hal_mapping_mode_t mapping_mode, iree_hal_memory_access_t memory_access, iree_device_size_t local_byte_offset, iree_device_size_t local_byte_length, - void** out_data_ptr) { + iree_hal_buffer_mapping_t* mapping) { iree_hal_cuda_buffer_t* buffer = iree_hal_cuda_buffer_cast(base_buffer); - if (!iree_all_bits_set(buffer->base.memory_type, - IREE_HAL_MEMORY_TYPE_HOST_VISIBLE)) { - return iree_make_status(IREE_STATUS_INTERNAL, - "trying to map memory not host visible"); - } + // TODO(benvanik): add upload/download for unmapped buffers. + IREE_RETURN_IF_ERROR(iree_hal_buffer_validate_memory_type( + iree_hal_buffer_memory_type(base_buffer), + IREE_HAL_MEMORY_TYPE_HOST_VISIBLE)); + IREE_RETURN_IF_ERROR( + iree_hal_buffer_validate_usage(iree_hal_buffer_allowed_usage(base_buffer), + IREE_HAL_BUFFER_USAGE_MAPPING)); uint8_t* data_ptr = (uint8_t*)(buffer->host_ptr) + local_byte_offset; // If we mapped for discard scribble over the bytes. This is not a mandated @@ -87,14 +89,16 @@ memset(data_ptr, 0xCD, local_byte_length); } #endif // !NDEBUG - *out_data_ptr = data_ptr; + + mapping->contents = iree_make_byte_span(data_ptr, local_byte_length); return iree_ok_status(); } -static void iree_hal_cuda_buffer_unmap_range( +static iree_status_t iree_hal_cuda_buffer_unmap_range( iree_hal_buffer_t* base_buffer, iree_device_size_t local_byte_offset, - iree_device_size_t local_byte_length, void* data_ptr) { - // nothing to do. + iree_device_size_t local_byte_length, iree_hal_buffer_mapping_t* mapping) { + // Nothing to do (today). + return iree_ok_status(); } static iree_status_t iree_hal_cuda_buffer_invalidate_range(
diff --git a/iree/hal/cuda/cuda_device.c b/iree/hal/cuda/cuda_device.c index 32fbd82..677fb5b 100644 --- a/iree/hal/cuda/cuda_device.c +++ b/iree/hal/cuda/cuda_device.c
@@ -23,6 +23,7 @@ #include "iree/hal/cuda/nop_executable_cache.h" #include "iree/hal/cuda/status_util.h" #include "iree/hal/cuda/stream_command_buffer.h" +#include "iree/hal/utils/buffer_transfer.h" #include "iree/hal/utils/deferred_command_buffer.h" //===----------------------------------------------------------------------===// @@ -108,7 +109,8 @@ device->context_wrapper.syms = syms; iree_status_t status = iree_hal_cuda_allocator_create( - &device->context_wrapper, cu_device, stream, &device->device_allocator); + (iree_hal_device_t*)device, &device->context_wrapper, cu_device, stream, + &device->device_allocator); device->command_buffer_mode = params->command_buffer_mode; if (iree_status_is_ok(status) && @@ -198,6 +200,12 @@ return device->device_allocator; } +static iree_status_t iree_hal_cuda_device_trim(iree_hal_device_t* base_device) { + iree_hal_cuda_device_t* device = iree_hal_cuda_device_cast(base_device); + iree_arena_block_pool_trim(&device->block_pool); + return iree_hal_allocator_trim(device->device_allocator); +} + static iree_status_t iree_hal_cuda_device_query_i32( iree_hal_device_t* base_device, iree_string_view_t category, iree_string_view_t key, int32_t* out_value) { @@ -229,7 +237,7 @@ case IREE_HAL_CUDA_COMMAND_BUFFER_MODE_GRAPH: return iree_hal_cuda_graph_command_buffer_create( base_device, &device->context_wrapper, mode, command_categories, - queue_affinity, out_command_buffer); + queue_affinity, &device->block_pool, out_command_buffer); case IREE_HAL_CUDA_COMMAND_BUFFER_MODE_STREAM: return iree_hal_deferred_command_buffer_create( base_device, mode, command_categories, &device->block_pool, @@ -316,8 +324,8 @@ } } } - // TODO(thomasraoux): Conservatively syncronize after every submit until we - // support semaphores. + // TODO(thomasraoux): implement semaphores - for now this conservatively + // synchronizes after every submit. CUDA_RETURN_IF_ERROR(device->context_wrapper.syms, cuStreamSynchronize(device->stream), "cuStreamSynchronize"); @@ -363,6 +371,7 @@ .id = iree_hal_cuda_device_id, .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, .create_command_buffer = iree_hal_cuda_device_create_command_buffer, .create_descriptor_set = iree_hal_cuda_device_create_descriptor_set, @@ -372,6 +381,7 @@ .create_executable_cache = iree_hal_cuda_device_create_executable_cache, .create_executable_layout = iree_hal_cuda_device_create_executable_layout, .create_semaphore = iree_hal_cuda_device_create_semaphore, + .transfer_range = iree_hal_device_submit_transfer_range_and_wait, .queue_submit = iree_hal_cuda_device_queue_submit, .submit_and_wait = iree_hal_cuda_device_submit_and_wait, .wait_semaphores = iree_hal_cuda_device_wait_semaphores,
diff --git a/iree/hal/cuda/dynamic_symbols_test.cc b/iree/hal/cuda/dynamic_symbols_test.cc index 47862d3..ab5136c 100644 --- a/iree/hal/cuda/dynamic_symbols_test.cc +++ b/iree/hal/cuda/dynamic_symbols_test.cc
@@ -27,6 +27,8 @@ iree_status_t status = iree_hal_cuda_dynamic_symbols_initialize( iree_allocator_system(), &symbols); if (!iree_status_is_ok(status)) { + iree_status_fprint(stderr, status); + iree_status_ignore(status); std::cerr << "Symbols cannot be loaded, skipping test."; GTEST_SKIP(); }
diff --git a/iree/hal/cuda/graph_command_buffer.c b/iree/hal/cuda/graph_command_buffer.c index d55ac91..093bbcb 100644 --- a/iree/hal/cuda/graph_command_buffer.c +++ b/iree/hal/cuda/graph_command_buffer.c
@@ -17,6 +17,7 @@ #include "iree/hal/cuda/executable_layout.h" #include "iree/hal/cuda/native_executable.h" #include "iree/hal/cuda/status_util.h" +#include "iree/hal/utils/resource_set.h" #define IREE_HAL_CUDA_MAX_BINDING_COUNT 64 // Kernel arguments contains binding and push constants. @@ -28,6 +29,11 @@ typedef struct iree_hal_cuda_graph_command_buffer_t { iree_hal_command_buffer_t base; iree_hal_cuda_context_wrapper_t* context; + iree_arena_block_pool_t* block_pool; + + // Maintains a reference to all resources used within the command buffer. + // Reset on each begin. + iree_hal_resource_set_t* resource_set; CUgraph graph; CUgraphExec exec; @@ -54,14 +60,13 @@ iree_hal_command_buffer_mode_t mode, iree_hal_command_category_t command_categories, iree_hal_queue_affinity_t queue_affinity, + iree_arena_block_pool_t* block_pool, iree_hal_command_buffer_t** out_command_buffer) { IREE_ASSERT_ARGUMENT(context); + IREE_ASSERT_ARGUMENT(block_pool); IREE_ASSERT_ARGUMENT(out_command_buffer); IREE_TRACE_ZONE_BEGIN(z0); - CUgraph graph = NULL; - CUDA_RETURN_IF_ERROR(context->syms, cuGraphCreate(&graph, /*flags=*/0), - "cuGraphCreate"); iree_hal_cuda_graph_command_buffer_t* command_buffer = NULL; size_t total_size = sizeof(*command_buffer) + IREE_HAL_CUDA_MAX_KERNEL_ARG * sizeof(void*) + @@ -73,7 +78,8 @@ device, mode, command_categories, queue_affinity, &iree_hal_cuda_graph_command_buffer_vtable, &command_buffer->base); command_buffer->context = context; - command_buffer->graph = graph; + command_buffer->block_pool = block_pool; + command_buffer->graph = NULL; command_buffer->exec = NULL; command_buffer->last_node = NULL; @@ -84,29 +90,46 @@ command_buffer->current_descriptor[i] = &device_ptrs[i]; } - *out_command_buffer = &command_buffer->base; - } else { - context->syms->cuGraphDestroy(graph); + status = iree_hal_resource_set_allocate(block_pool, + &command_buffer->resource_set); } + if (iree_status_is_ok(status)) { + *out_command_buffer = &command_buffer->base; + } else { + iree_hal_command_buffer_release(&command_buffer->base); + } IREE_TRACE_ZONE_END(z0); return status; } +static void iree_hal_cuda_graph_command_buffer_reset( + iree_hal_cuda_graph_command_buffer_t* command_buffer) { + if (command_buffer->graph != NULL) { + CUDA_IGNORE_ERROR(command_buffer->context->syms, + cuGraphDestroy(command_buffer->graph)); + command_buffer->graph = NULL; + } + + if (command_buffer->exec != NULL) { + CUDA_IGNORE_ERROR(command_buffer->context->syms, + cuGraphExecDestroy(command_buffer->exec)); + command_buffer->exec = NULL; + } + + command_buffer->last_node = NULL; + + iree_hal_resource_set_reset(command_buffer->resource_set); +} + static void iree_hal_cuda_graph_command_buffer_destroy( iree_hal_command_buffer_t* base_command_buffer) { iree_hal_cuda_graph_command_buffer_t* command_buffer = iree_hal_cuda_graph_command_buffer_cast(base_command_buffer); IREE_TRACE_ZONE_BEGIN(z0); - if (command_buffer->graph != NULL) { - CUDA_IGNORE_ERROR(command_buffer->context->syms, - cuGraphDestroy(command_buffer->graph)); - } - if (command_buffer->exec != NULL) { - CUDA_IGNORE_ERROR(command_buffer->context->syms, - cuGraphExecDestroy(command_buffer->exec)); - } + iree_hal_cuda_graph_command_buffer_reset(command_buffer); + iree_hal_resource_set_free(command_buffer->resource_set); iree_allocator_free(command_buffer->context->host_allocator, command_buffer); IREE_TRACE_ZONE_END(z0); @@ -136,7 +159,17 @@ static iree_status_t iree_hal_cuda_graph_command_buffer_begin( iree_hal_command_buffer_t* base_command_buffer) { - // Nothing to do. + iree_hal_cuda_graph_command_buffer_t* command_buffer = + iree_hal_cuda_graph_command_buffer_cast(base_command_buffer); + + // Reset any prior recorded commands. + iree_hal_cuda_graph_command_buffer_reset(command_buffer); + + // Create a new empty graph to record into. + CUDA_RETURN_IF_ERROR(command_buffer->context->syms, + cuGraphCreate(&command_buffer->graph, /*flags=*/0), + "cuGraphCreate"); + return iree_ok_status(); } @@ -145,23 +178,24 @@ iree_hal_cuda_graph_command_buffer_t* command_buffer = iree_hal_cuda_graph_command_buffer_cast(base_command_buffer); - size_t num_nodes; - CUDA_RETURN_IF_ERROR(command_buffer->context->syms, - cuGraphGetNodes(command_buffer->graph, NULL, &num_nodes), - "cuGraphGetNodes"); + // Reset state used during recording. + command_buffer->last_node = NULL; - CUgraphNode error_node; + // Compile the graph. + CUgraphNode error_node = NULL; iree_status_t status = CU_RESULT_TO_STATUS(command_buffer->context->syms, cuGraphInstantiate(&command_buffer->exec, command_buffer->graph, &error_node, /*logBuffer=*/NULL, - /* bufferSize=*/0)); + /*bufferSize=*/0)); if (iree_status_is_ok(status)) { + // No longer need the source graph used for construction. CUDA_IGNORE_ERROR(command_buffer->context->syms, cuGraphDestroy(command_buffer->graph)); + command_buffer->graph = NULL; } - command_buffer->graph = NULL; + return iree_ok_status(); } @@ -257,6 +291,9 @@ iree_hal_cuda_graph_command_buffer_t* command_buffer = iree_hal_cuda_graph_command_buffer_cast(base_command_buffer); + IREE_RETURN_IF_ERROR(iree_hal_resource_set_insert( + command_buffer->resource_set, 1, &target_buffer)); + CUdeviceptr target_device_buffer = iree_hal_cuda_buffer_device_pointer( iree_hal_buffer_allocated_buffer(target_buffer)); target_offset += iree_hal_buffer_byte_offset(target_buffer); @@ -297,6 +334,10 @@ iree_hal_cuda_graph_command_buffer_t* command_buffer = iree_hal_cuda_graph_command_buffer_cast(base_command_buffer); + const iree_hal_buffer_t* buffers[2] = {source_buffer, target_buffer}; + IREE_RETURN_IF_ERROR( + iree_hal_resource_set_insert(command_buffer->resource_set, 2, buffers)); + CUdeviceptr target_device_buffer = iree_hal_cuda_buffer_device_pointer( iree_hal_buffer_allocated_buffer(target_buffer)); target_offset += iree_hal_buffer_byte_offset(target_buffer); @@ -376,16 +417,19 @@ } qsort(binding_used, binding_count, sizeof(iree_hal_cuda_binding_mapping_t), compare_binding_index); - assert(binding_count < IREE_HAL_CUDA_MAX_BINDING_COUNT && - "binding count larger than the max expected."); + IREE_ASSERT_LT(binding_count, IREE_HAL_CUDA_MAX_BINDING_COUNT, + "binding count larger than the max expected"); for (iree_host_size_t i = 0; i < binding_count; i++) { - iree_hal_descriptor_set_binding_t binding = bindings[binding_used[i].index]; + const iree_hal_descriptor_set_binding_t* binding = + &bindings[binding_used[i].index]; CUdeviceptr device_ptr = iree_hal_cuda_buffer_device_pointer( - iree_hal_buffer_allocated_buffer(binding.buffer)) + - iree_hal_buffer_byte_offset(binding.buffer) + binding.offset; + iree_hal_buffer_allocated_buffer(binding->buffer)) + + iree_hal_buffer_byte_offset(binding->buffer) + binding->offset; *((CUdeviceptr*)command_buffer->current_descriptor[i + base_binding]) = device_ptr; + IREE_RETURN_IF_ERROR(iree_hal_resource_set_insert( + command_buffer->resource_set, 1, &binding->buffer)); } return iree_ok_status(); } @@ -406,6 +450,8 @@ uint32_t workgroup_x, uint32_t workgroup_y, uint32_t workgroup_z) { iree_hal_cuda_graph_command_buffer_t* command_buffer = iree_hal_cuda_graph_command_buffer_cast(base_command_buffer); + IREE_RETURN_IF_ERROR(iree_hal_resource_set_insert( + command_buffer->resource_set, 1, &executable)); iree_hal_executable_layout_t* layout = iree_hal_cuda_executable_get_layout(executable, entry_point); iree_host_size_t num_constants =
diff --git a/iree/hal/cuda/graph_command_buffer.h b/iree/hal/cuda/graph_command_buffer.h index 7ad5251..8ef4fda 100644 --- a/iree/hal/cuda/graph_command_buffer.h +++ b/iree/hal/cuda/graph_command_buffer.h
@@ -17,12 +17,18 @@ extern "C" { #endif // __cplusplus -// Creates a cuda graph. +typedef struct iree_arena_block_pool_t iree_arena_block_pool_t; + +// Creates a command buffer that records into a CUDA graph. +// +// NOTE: the |block_pool| must remain live for the lifetime of the command +// buffers that use it. iree_status_t iree_hal_cuda_graph_command_buffer_create( iree_hal_device_t* device, iree_hal_cuda_context_wrapper_t* context, iree_hal_command_buffer_mode_t mode, iree_hal_command_category_t command_categories, iree_hal_queue_affinity_t queue_affinity, + iree_arena_block_pool_t* block_pool, iree_hal_command_buffer_t** out_command_buffer); // Returns true if |command_buffer| is a CUDA graph-based command buffer.
diff --git a/iree/hal/descriptor_set.h b/iree/hal/descriptor_set.h index 6753492..11c7957 100644 --- a/iree/hal/descriptor_set.h +++ b/iree/hal/descriptor_set.h
@@ -88,11 +88,9 @@ //===----------------------------------------------------------------------===// typedef struct iree_hal_descriptor_set_vtable_t { - // << HAL C porting in progress >> - IREE_API_UNSTABLE - void(IREE_API_PTR* destroy)(iree_hal_descriptor_set_t* descriptor_set); } iree_hal_descriptor_set_vtable_t; +IREE_HAL_ASSERT_VTABLE_LAYOUT(iree_hal_descriptor_set_vtable_t); IREE_API_EXPORT void iree_hal_descriptor_set_destroy( iree_hal_descriptor_set_t* descriptor_set);
diff --git a/iree/hal/descriptor_set_layout.h b/iree/hal/descriptor_set_layout.h index 68c1248..36e3940 100644 --- a/iree/hal/descriptor_set_layout.h +++ b/iree/hal/descriptor_set_layout.h
@@ -89,12 +89,10 @@ //===----------------------------------------------------------------------===// typedef struct iree_hal_descriptor_set_layout_vtable_t { - // << HAL C porting in progress >> - IREE_API_UNSTABLE - void(IREE_API_PTR* destroy)( iree_hal_descriptor_set_layout_t* descriptor_set_layout); } iree_hal_descriptor_set_layout_vtable_t; +IREE_HAL_ASSERT_VTABLE_LAYOUT(iree_hal_descriptor_set_layout_vtable_t); IREE_API_EXPORT void iree_hal_descriptor_set_layout_destroy( iree_hal_descriptor_set_layout_t* descriptor_set_layout);
diff --git a/iree/hal/device.c b/iree/hal/device.c index bf9d26d..648b3b9 100644 --- a/iree/hal/device.c +++ b/iree/hal/device.c
@@ -7,6 +7,9 @@ #include "iree/hal/device.h" #include "iree/base/tracing.h" +#include "iree/hal/allocator.h" +#include "iree/hal/buffer.h" +#include "iree/hal/command_buffer.h" #include "iree/hal/detail.h" #include "iree/hal/resource.h" @@ -33,6 +36,15 @@ return _VTABLE_DISPATCH(device, device_allocator)(device); } +IREE_API_EXPORT +iree_status_t iree_hal_device_trim(iree_hal_device_t* device) { + IREE_ASSERT_ARGUMENT(device); + IREE_TRACE_ZONE_BEGIN(z0); + iree_status_t status = _VTABLE_DISPATCH(device, trim)(device); + IREE_TRACE_ZONE_END(z0); + return status; +} + IREE_API_EXPORT iree_status_t iree_hal_device_query_i32( iree_hal_device_t* device, iree_string_view_t category, iree_string_view_t key, int32_t* out_value) { @@ -49,6 +61,200 @@ return _VTABLE_DISPATCH(device, query_i32)(device, category, key, out_value); } +// Performs a synchronous host->device or device->host transfer. +static iree_status_t iree_hal_device_transfer_buffer( + iree_hal_device_t* device, iree_hal_buffer_t* source_buffer, + iree_hal_memory_type_t memory_type, iree_hal_buffer_usage_t allowed_usage, + iree_hal_buffer_t** out_target_buffer) { + IREE_ASSERT_ARGUMENT(device); + IREE_ASSERT_ARGUMENT(source_buffer); + IREE_ASSERT_ARGUMENT(out_target_buffer); + *out_target_buffer = NULL; + IREE_TRACE_ZONE_BEGIN(z0); + + if (iree_all_bits_set(iree_hal_buffer_memory_type(source_buffer), + memory_type) && + iree_all_bits_set(iree_hal_buffer_allowed_usage(source_buffer), + allowed_usage)) { + // Source is already usable for the intended purposes - avoid the copy. + *out_target_buffer = source_buffer; + iree_hal_buffer_retain(source_buffer); + IREE_TRACE_ZONE_END(z0); + return iree_ok_status(); + } + + iree_device_size_t allocation_size = + iree_hal_buffer_allocation_size(source_buffer); + IREE_TRACE_ZONE_APPEND_VALUE(z0, allocation_size); + + // Allocate a new buffer of the desired memory type. It needs to be at least + // large enough to hold the requested data. + iree_hal_buffer_t* target_buffer = NULL; + IREE_RETURN_AND_END_ZONE_IF_ERROR( + z0, iree_hal_allocator_allocate_buffer( + iree_hal_device_allocator(device), memory_type, + allowed_usage | IREE_HAL_BUFFER_USAGE_TRANSFER, allocation_size, + iree_const_byte_span_empty(), &target_buffer)); + + // Perform the transfer and wait for it to complete. + const iree_hal_transfer_command_t transfer_command = { + .type = IREE_HAL_TRANSFER_COMMAND_TYPE_COPY, + .copy = + { + .source_buffer = source_buffer, + .source_offset = 0, + .target_buffer = target_buffer, + .target_offset = 0, + .length = IREE_WHOLE_BUFFER, + }, + }; + iree_status_t status = iree_hal_device_transfer_and_wait( + device, /*wait_semaphore=*/NULL, + /*wait_value=*/0ull, 1, &transfer_command, iree_infinite_timeout()); + + if (iree_status_is_ok(status)) { + *out_target_buffer = target_buffer; + } else { + iree_hal_buffer_release(target_buffer); + } + IREE_TRACE_ZONE_END(z0); + return status; +} + +IREE_API_EXPORT iree_status_t iree_hal_device_transfer_to_device( + iree_hal_device_t* device, iree_hal_buffer_t* source_buffer, + iree_hal_buffer_usage_t allowed_usage, + iree_hal_buffer_t** out_target_buffer) { + return iree_hal_device_transfer_buffer(device, source_buffer, + IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL, + allowed_usage, out_target_buffer); +} + +IREE_API_EXPORT iree_status_t iree_hal_device_transfer_to_host( + iree_hal_device_t* device, iree_hal_buffer_t* source_buffer, + iree_hal_buffer_t** out_target_buffer) { + iree_hal_buffer_usage_t allowed_usage = + IREE_HAL_BUFFER_USAGE_TRANSFER | IREE_HAL_BUFFER_USAGE_MAPPING; + return iree_hal_device_transfer_buffer(device, source_buffer, + IREE_HAL_MEMORY_TYPE_HOST_LOCAL, + allowed_usage, out_target_buffer); +} + +IREE_API_EXPORT iree_status_t iree_hal_device_transfer_range( + iree_hal_device_t* device, iree_hal_transfer_buffer_t source, + iree_device_size_t source_offset, iree_hal_transfer_buffer_t target, + iree_device_size_t target_offset, iree_device_size_t data_length, + iree_hal_transfer_buffer_flags_t flags, iree_timeout_t timeout) { + if (data_length == 0) { + return iree_ok_status(); // No-op. + } + + // host->host is not allowed. We may want to support this one day to allow for + // parallelized copies and such, however the validation code differs quite a + // bit and it'd be better to have this as part of a task system API. + bool is_source_host = source.device_buffer == NULL; + bool is_target_host = target.device_buffer == NULL; + if (is_source_host && is_target_host) { + return iree_make_status( + IREE_STATUS_INVALID_ARGUMENT, + "cannot perform host->host transfers via this API, use memcpy/memmove"); + } + + // Check for overlap - like memcpy we require that the two ranges don't have + // any overlap as we may use memcpy. This only matters if the buffers are + // both device buffers - host and device should never alias: behavior is + // undefined if a user tries to pass a mapped device pointer as if it was a + // host pointer. + if (!is_source_host && !is_target_host && + iree_hal_buffer_test_overlap(source.device_buffer, source_offset, + data_length, target.device_buffer, + target_offset, data_length) != + IREE_HAL_BUFFER_OVERLAP_DISJOINT) { + return iree_make_status( + IREE_STATUS_INVALID_ARGUMENT, + "source and target ranges must not overlap within the same buffer"); + } + + IREE_TRACE_ZONE_BEGIN(z0); + IREE_TRACE_ZONE_APPEND_TEXT( + z0, is_source_host ? "h2d" : (is_target_host ? "d2h" : "d2d")); + IREE_TRACE_ZONE_APPEND_VALUE(z0, data_length); + + // Defer to the backing implementation. + iree_status_t status = _VTABLE_DISPATCH(device, transfer_range)( + device, source, source_offset, target, target_offset, data_length, flags, + timeout); + + IREE_TRACE_ZONE_END(z0); + return status; +} + +IREE_API_EXPORT iree_status_t iree_hal_device_transfer_and_wait( + iree_hal_device_t* device, iree_hal_semaphore_t* wait_semaphore, + uint64_t wait_value, iree_host_size_t transfer_count, + const iree_hal_transfer_command_t* transfer_commands, + iree_timeout_t timeout) { + IREE_ASSERT_ARGUMENT(device); + IREE_ASSERT_ARGUMENT(!transfer_count || transfer_commands); + IREE_TRACE_ZONE_BEGIN(z0); + + // We only want to allow inline execution if we have not been instructed to + // wait on a semaphore and it hasn't yet been signaled. + iree_hal_command_buffer_mode_t mode = IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT; + if (wait_semaphore) { + uint64_t current_value = 0ull; + IREE_RETURN_AND_END_ZONE_IF_ERROR( + z0, iree_hal_semaphore_query(wait_semaphore, ¤t_value)); + if (current_value >= wait_value) { + mode |= IREE_HAL_COMMAND_BUFFER_MODE_ALLOW_INLINE_EXECUTION; + } + } else { + mode |= IREE_HAL_COMMAND_BUFFER_MODE_ALLOW_INLINE_EXECUTION; + } + + // Create a command buffer performing all of the transfer operations. + iree_hal_command_buffer_t* command_buffer = NULL; + IREE_RETURN_AND_END_ZONE_IF_ERROR( + z0, iree_hal_create_transfer_command_buffer( + device, mode, IREE_HAL_QUEUE_AFFINITY_ANY, transfer_count, + transfer_commands, &command_buffer)); + + // Perform a full submit-and-wait. On devices with multiple queues this can + // run out-of-order/overlapped with other work and return earlier than device + // idle. + iree_hal_semaphore_t* fence_semaphore = NULL; + iree_status_t status = + iree_hal_semaphore_create(device, 0ull, &fence_semaphore); + uint64_t signal_value = 1ull; + if (iree_status_is_ok(status)) { + iree_hal_submission_batch_t batch = { + .wait_semaphores = + { + .count = wait_semaphore != NULL ? 1 : 0, + .semaphores = &wait_semaphore, + .payload_values = &wait_value, + }, + .command_buffer_count = 1, + .command_buffers = &command_buffer, + .signal_semaphores = + { + .count = 1, + .semaphores = &fence_semaphore, + .payload_values = &signal_value, + }, + }; + status = iree_hal_device_submit_and_wait( + device, IREE_HAL_COMMAND_CATEGORY_TRANSFER, IREE_HAL_QUEUE_AFFINITY_ANY, + 1, &batch, fence_semaphore, signal_value, timeout); + } + + iree_hal_command_buffer_release(command_buffer); + iree_hal_semaphore_release(fence_semaphore); + + IREE_TRACE_ZONE_END(z0); + return status; +} + // Validates that the submission is well-formed. static iree_status_t iree_hal_device_validate_submission( iree_host_size_t batch_count, const iree_hal_submission_batch_t* batches) {
diff --git a/iree/hal/device.h b/iree/hal/device.h index 1d77d27..8dbf9b9 100644 --- a/iree/hal/device.h +++ b/iree/hal/device.h
@@ -74,6 +74,41 @@ iree_string_view_t name; } iree_hal_device_info_t; +// A transfer source or destination. +typedef struct iree_hal_transfer_buffer_t { + // A host-allocated void* buffer. + iree_byte_span_t host_buffer; + // A device-allocated buffer (may be of any memory type). + iree_hal_buffer_t* device_buffer; +} iree_hal_transfer_buffer_t; + +static inline iree_hal_transfer_buffer_t iree_hal_make_host_transfer_buffer( + iree_byte_span_t host_buffer) { + iree_hal_transfer_buffer_t transfer_buffer = { + host_buffer, + NULL, + }; + return transfer_buffer; +} + +static inline iree_hal_transfer_buffer_t +iree_hal_make_host_transfer_buffer_span(void* ptr, iree_host_size_t length) { + iree_hal_transfer_buffer_t transfer_buffer = { + iree_make_byte_span(ptr, length), + NULL, + }; + return transfer_buffer; +} + +static inline iree_hal_transfer_buffer_t iree_hal_make_device_transfer_buffer( + iree_hal_buffer_t* device_buffer) { + iree_hal_transfer_buffer_t transfer_buffer = { + iree_byte_span_empty(), + device_buffer, + }; + return transfer_buffer; +} + // A list of semaphores and their corresponding payloads. // When signaling each semaphore will be set to the new payload value provided. // When waiting each semaphore must reach or exceed the payload value. @@ -144,6 +179,12 @@ IREE_API_EXPORT iree_hal_allocator_t* iree_hal_device_allocator( iree_hal_device_t* device); +// Trims pools and caches used by the HAL to the minimum required for live +// allocations. This can be used on low-memory conditions or when +// suspending/parking instances. +IREE_API_EXPORT +iree_status_t iree_hal_device_trim(iree_hal_device_t* device); + // Queries a configuration value as an int32_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 @@ -167,6 +208,80 @@ iree_hal_device_t* device, iree_string_view_t category, iree_string_view_t key, int32_t* out_value); +// Synchronously transfers the given |source_buffer| to a device-local +// buffer returned in |out_target_buffer|. Callers must release the target +// buffer when no longer used. If the source buffer is already device-local it +// will be returned without an allocation or copy occurring. +// +// This utility may incur signficant overhead and is present for simple tooling +// and prototypes; when transfering multiple buffers users should always prefer +// asynchronous command buffers submitted to device queues. Note too that the +// entire buffer is transferred: if reading back smaller portions it is better +// to perform these as ranged transfers to avoid the amount of data that needs +// to be moved. +IREE_API_EXPORT iree_status_t iree_hal_device_transfer_to_device( + iree_hal_device_t* device, iree_hal_buffer_t* source_buffer, + iree_hal_buffer_usage_t allowed_usage, + iree_hal_buffer_t** out_target_buffer); + +// Synchronously transfers the given |source_buffer| to a host-local +// buffer returned in |out_target_buffer|. Callers must release the target +// buffer when no longer used. If the source buffer is already host-local it +// will be returned without an allocation or copy occurring. +// +// This utility may incur signficant overhead and is present for simple tooling +// and prototypes; when transfering multiple buffers users should always prefer +// asynchronous command buffers submitted to device queues. +IREE_API_EXPORT iree_status_t iree_hal_device_transfer_to_host( + iree_hal_device_t* device, iree_hal_buffer_t* source_buffer, + iree_hal_buffer_t** out_target_buffer); + +// Synchronously copies data from |source| into |target|. +// +// Supports host->device, device->host, and device->device transfer, +// including across devices. This method will never fail based on device +// capabilities but may incur some extreme transient allocations and copies in +// order to perform the transfer. +// +// The ordering of the transfer is undefined with respect to queue execution on +// the source or target device; some may require full device flushes in order to +// perform this operation while others may immediately perform it while there is +// still work outstanding. +// +// It is strongly recommended that buffer operations are performed on transfer +// queues; using this synchronous function may incur additional cache flushes +// and synchronous blocking behavior and is not supported on all buffer types. +// See iree_hal_command_buffer_copy_buffer. +IREE_API_EXPORT iree_status_t iree_hal_device_transfer_range( + iree_hal_device_t* device, iree_hal_transfer_buffer_t source, + iree_device_size_t source_offset, iree_hal_transfer_buffer_t target, + iree_device_size_t target_offset, iree_device_size_t data_length, + iree_hal_transfer_buffer_flags_t flags, iree_timeout_t timeout); + +// Synchronously executes one or more transfer operations against a queue. +// All buffers must be compatible with |device| and ranges must not overlap +// (same as with memcpy). +// +// This is a blocking operation and may incur significant overheads as +// internally it issues a command buffer with the transfer operations and waits +// for it to complete. Users should do that themselves so that the work can be +// issued concurrently and batched effectively. This is only useful as a +// fallback for implementations that require it or tools where things like I/O +// are transferred without worrying about performance. When submitting other +// work it's preferable to use iree_hal_create_transfer_command_buffer and a +// normal queue submission that allows for more fine-grained sequencing and +// amortizes the submission cost by batching other work. +// +// The transfer will begin after the optional |wait_semaphore| reaches +// |wait_value|. Behavior is undefined if no semaphore is provided and there are +// in-flight operations concurrently using the buffer ranges. +// Returns only after all transfers have completed and been flushed. +IREE_API_EXPORT iree_status_t iree_hal_device_transfer_and_wait( + iree_hal_device_t* device, iree_hal_semaphore_t* wait_semaphore, + uint64_t wait_value, iree_host_size_t transfer_count, + const iree_hal_transfer_command_t* transfer_commands, + iree_timeout_t timeout); + // Submits one or more batches of work to a device queue. // // The queue is selected based on the flags set in |command_categories| and the @@ -244,9 +359,6 @@ //===----------------------------------------------------------------------===// typedef struct iree_hal_device_vtable_t { - // << HAL C porting in progress >> - IREE_API_UNSTABLE - void(IREE_API_PTR* destroy)(iree_hal_device_t* device); iree_string_view_t(IREE_API_PTR* id)(iree_hal_device_t* device); @@ -255,6 +367,8 @@ iree_hal_allocator_t*(IREE_API_PTR* device_allocator)( iree_hal_device_t* device); + 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_string_view_t category, iree_string_view_t key, @@ -296,6 +410,12 @@ iree_hal_device_t* device, uint64_t initial_value, iree_hal_semaphore_t** out_semaphore); + iree_status_t(IREE_API_PTR* transfer_range)( + iree_hal_device_t* device, iree_hal_transfer_buffer_t source, + iree_device_size_t source_offset, iree_hal_transfer_buffer_t target, + iree_device_size_t target_offset, iree_device_size_t data_length, + iree_hal_transfer_buffer_flags_t flags, iree_timeout_t timeout); + iree_status_t(IREE_API_PTR* queue_submit)( iree_hal_device_t* device, iree_hal_command_category_t command_categories, iree_hal_queue_affinity_t queue_affinity, iree_host_size_t batch_count, @@ -315,6 +435,7 @@ iree_status_t(IREE_API_PTR* wait_idle)(iree_hal_device_t* device, iree_timeout_t timeout); } iree_hal_device_vtable_t; +IREE_HAL_ASSERT_VTABLE_LAYOUT(iree_hal_device_vtable_t); IREE_API_EXPORT void iree_hal_device_destroy(iree_hal_device_t* device);
diff --git a/iree/hal/driver.h b/iree/hal/driver.h index 8cd6076..65cbd66 100644 --- a/iree/hal/driver.h +++ b/iree/hal/driver.h
@@ -94,9 +94,6 @@ //===----------------------------------------------------------------------===// typedef struct iree_hal_driver_vtable_t { - // << HAL C porting in progress >> - IREE_API_UNSTABLE - void(IREE_API_PTR* destroy)(iree_hal_driver_t* driver); iree_status_t(IREE_API_PTR* query_available_devices)( @@ -109,6 +106,7 @@ iree_allocator_t allocator, iree_hal_device_t** out_device); } iree_hal_driver_vtable_t; +IREE_HAL_ASSERT_VTABLE_LAYOUT(iree_hal_driver_vtable_t); IREE_API_EXPORT void iree_hal_driver_destroy(iree_hal_driver_t* driver);
diff --git a/iree/hal/event.h b/iree/hal/event.h index 911e50f..a6ea312 100644 --- a/iree/hal/event.h +++ b/iree/hal/event.h
@@ -51,11 +51,9 @@ //===----------------------------------------------------------------------===// typedef struct iree_hal_event_vtable_t { - // << HAL C porting in progress >> - IREE_API_UNSTABLE - void(IREE_API_PTR* destroy)(iree_hal_event_t* event); } iree_hal_event_vtable_t; +IREE_HAL_ASSERT_VTABLE_LAYOUT(iree_hal_event_vtable_t); IREE_API_EXPORT void iree_hal_event_destroy(iree_hal_event_t* event);
diff --git a/iree/hal/executable.h b/iree/hal/executable.h index 7cabcf4..561ed3a 100644 --- a/iree/hal/executable.h +++ b/iree/hal/executable.h
@@ -51,11 +51,9 @@ //===----------------------------------------------------------------------===// typedef struct iree_hal_executable_vtable_t { - // << HAL C porting in progress >> - IREE_API_UNSTABLE - void(IREE_API_PTR* destroy)(iree_hal_executable_t* executable); } iree_hal_executable_vtable_t; +IREE_HAL_ASSERT_VTABLE_LAYOUT(iree_hal_executable_vtable_t); IREE_API_EXPORT void iree_hal_executable_destroy( iree_hal_executable_t* executable);
diff --git a/iree/hal/executable_cache.h b/iree/hal/executable_cache.h index d75cfcb..50945f8 100644 --- a/iree/hal/executable_cache.h +++ b/iree/hal/executable_cache.h
@@ -171,9 +171,6 @@ //===----------------------------------------------------------------------===// typedef struct iree_hal_executable_cache_vtable_t { - // << HAL C porting in progress >> - IREE_API_UNSTABLE - void(IREE_API_PTR* destroy)(iree_hal_executable_cache_t* executable_cache); bool(IREE_API_PTR* can_prepare_format)( @@ -186,6 +183,7 @@ const iree_hal_executable_spec_t* executable_spec, iree_hal_executable_t** out_executable); } iree_hal_executable_cache_vtable_t; +IREE_HAL_ASSERT_VTABLE_LAYOUT(iree_hal_executable_cache_vtable_t); IREE_API_EXPORT void iree_hal_executable_cache_destroy( iree_hal_executable_cache_t* executable_cache);
diff --git a/iree/hal/executable_layout.h b/iree/hal/executable_layout.h index 1fbcd48..7fa1a21 100644 --- a/iree/hal/executable_layout.h +++ b/iree/hal/executable_layout.h
@@ -64,11 +64,9 @@ //===----------------------------------------------------------------------===// typedef struct iree_hal_executable_layout_vtable_t { - // << HAL C porting in progress >> - IREE_API_UNSTABLE - void(IREE_API_PTR* destroy)(iree_hal_executable_layout_t* executable_layout); } iree_hal_executable_layout_vtable_t; +IREE_HAL_ASSERT_VTABLE_LAYOUT(iree_hal_executable_layout_vtable_t); IREE_API_EXPORT void iree_hal_executable_layout_destroy( iree_hal_executable_layout_t* executable_layout);
diff --git a/iree/hal/local/BUILD b/iree/hal/local/BUILD index 3b6cd65..7b06009 100644 --- a/iree/hal/local/BUILD +++ b/iree/hal/local/BUILD
@@ -103,6 +103,7 @@ "//iree/base/internal:arena", "//iree/base/internal:synchronization", "//iree/hal", + "//iree/hal/utils:buffer_transfer", ], ) @@ -168,6 +169,8 @@ "//iree/base/internal:synchronization", "//iree/base/internal:wait_handle", "//iree/hal", + "//iree/hal/utils:buffer_transfer", + "//iree/hal/utils:resource_set", "//iree/task", ], )
diff --git a/iree/hal/local/CMakeLists.txt b/iree/hal/local/CMakeLists.txt index 42ad0c1..42fb8d2 100644 --- a/iree/hal/local/CMakeLists.txt +++ b/iree/hal/local/CMakeLists.txt
@@ -99,6 +99,7 @@ iree::base::internal::synchronization iree::base::tracing iree::hal + iree::hal::utils::buffer_transfer PUBLIC ) @@ -155,6 +156,8 @@ iree::base::internal::wait_handle iree::base::tracing iree::hal + iree::hal::utils::buffer_transfer + iree::hal::utils::resource_set iree::task PUBLIC )
diff --git a/iree/hal/local/executable_library_benchmark.c b/iree/hal/local/executable_library_benchmark.c index 2605928..1002f4a 100644 --- a/iree/hal/local/executable_library_benchmark.c +++ b/iree/hal/local/executable_library_benchmark.c
@@ -274,9 +274,10 @@ iree_hal_buffer_t* buffer = iree_hal_buffer_view_buffer(buffer_views[i]); iree_device_size_t buffer_length = iree_hal_buffer_view_byte_length(buffer_views[i]); - iree_hal_buffer_mapping_t buffer_mapping; + iree_hal_buffer_mapping_t buffer_mapping = {{0}}; IREE_RETURN_IF_ERROR(iree_hal_buffer_map_range( - buffer, IREE_HAL_MEMORY_ACCESS_READ | IREE_HAL_MEMORY_ACCESS_WRITE, 0, + buffer, IREE_HAL_MAPPING_MODE_PERSISTENT, + IREE_HAL_MEMORY_ACCESS_READ | IREE_HAL_MEMORY_ACCESS_WRITE, 0, buffer_length, &buffer_mapping)); binding_ptrs[i] = buffer_mapping.contents.data; binding_lengths[i] = (size_t)buffer_mapping.contents.data_length;
diff --git a/iree/hal/local/inline_command_buffer.c b/iree/hal/local/inline_command_buffer.c index 5819de7..07d95c5 100644 --- a/iree/hal/local/inline_command_buffer.c +++ b/iree/hal/local/inline_command_buffer.c
@@ -345,10 +345,11 @@ iree_host_size_t binding_ordinal = binding_base + bindings[i].binding; // TODO(benvanik): track mapping so we can properly map/unmap/flush/etc. - iree_hal_buffer_mapping_t buffer_mapping; + iree_hal_buffer_mapping_t buffer_mapping = {{0}}; IREE_RETURN_IF_ERROR(iree_hal_buffer_map_range( - bindings[i].buffer, IREE_HAL_MEMORY_ACCESS_ANY, bindings[i].offset, - bindings[i].length, &buffer_mapping)); + bindings[i].buffer, IREE_HAL_MAPPING_MODE_PERSISTENT, + IREE_HAL_MEMORY_ACCESS_ANY, bindings[i].offset, bindings[i].length, + &buffer_mapping)); command_buffer->state.full_bindings[binding_ordinal] = buffer_mapping.contents.data; command_buffer->state.full_binding_lengths[binding_ordinal] = @@ -475,10 +476,11 @@ iree_hal_buffer_t* workgroups_buffer, iree_device_size_t workgroups_offset) { // TODO(benvanik): track mapping so we can properly map/unmap/flush/etc. - iree_hal_buffer_mapping_t buffer_mapping; + iree_hal_buffer_mapping_t buffer_mapping = {{0}}; IREE_RETURN_IF_ERROR(iree_hal_buffer_map_range( - workgroups_buffer, IREE_HAL_MEMORY_ACCESS_READ, workgroups_offset, - 3 * sizeof(uint32_t), &buffer_mapping)); + workgroups_buffer, IREE_HAL_MAPPING_MODE_PERSISTENT, + IREE_HAL_MEMORY_ACCESS_READ, workgroups_offset, 3 * sizeof(uint32_t), + &buffer_mapping)); iree_hal_vec3_t workgroup_count = *(const iree_hal_vec3_t*)buffer_mapping.contents.data; return iree_hal_inline_command_buffer_dispatch(
diff --git a/iree/hal/local/sync_device.c b/iree/hal/local/sync_device.c index 59a6eff..58b5029 100644 --- a/iree/hal/local/sync_device.c +++ b/iree/hal/local/sync_device.c
@@ -18,6 +18,7 @@ #include "iree/hal/local/local_executable_layout.h" #include "iree/hal/local/sync_event.h" #include "iree/hal/local/sync_semaphore.h" +#include "iree/hal/utils/buffer_transfer.h" typedef struct iree_hal_sync_device_t { iree_hal_resource_t resource; @@ -133,6 +134,11 @@ return device->device_allocator; } +static iree_status_t iree_hal_sync_device_trim(iree_hal_device_t* base_device) { + iree_hal_sync_device_t* device = iree_hal_sync_device_cast(base_device); + return iree_hal_allocator_trim(device->device_allocator); +} + static iree_status_t iree_hal_sync_device_query_i32( iree_hal_device_t* base_device, iree_string_view_t category, iree_string_view_t key, int32_t* out_value) { @@ -288,6 +294,7 @@ .id = iree_hal_sync_device_id, .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, .create_command_buffer = iree_hal_sync_device_create_command_buffer, .create_descriptor_set = iree_hal_sync_device_create_descriptor_set, @@ -297,6 +304,7 @@ .create_executable_cache = iree_hal_sync_device_create_executable_cache, .create_executable_layout = iree_hal_sync_device_create_executable_layout, .create_semaphore = iree_hal_sync_device_create_semaphore, + .transfer_range = iree_hal_device_transfer_mappable_range, .queue_submit = iree_hal_sync_device_queue_submit, .submit_and_wait = iree_hal_sync_device_submit_and_wait, .wait_semaphores = iree_hal_sync_device_wait_semaphores,
diff --git a/iree/hal/local/task_command_buffer.c b/iree/hal/local/task_command_buffer.c index 051eade..ed5cb96 100644 --- a/iree/hal/local/task_command_buffer.c +++ b/iree/hal/local/task_command_buffer.c
@@ -17,6 +17,7 @@ #include "iree/hal/local/local_descriptor_set_layout.h" #include "iree/hal/local/local_executable.h" #include "iree/hal/local/local_executable_layout.h" +#include "iree/hal/utils/resource_set.h" #include "iree/task/affinity_set.h" #include "iree/task/list.h" #include "iree/task/submission.h" @@ -43,6 +44,10 @@ // Arena used for all allocations; references the shared device block pool. iree_arena_allocator_t arena; + // Maintains a reference to all resources used within the command buffer. + // Reset on each begin. + iree_hal_resource_set_t* resource_set; + // One or more tasks at the root of the command buffer task DAG. // These tasks are all able to execute concurrently and will be the initial // ready task set in the submission. @@ -139,7 +144,13 @@ iree_task_list_initialize(&command_buffer->root_tasks); iree_task_list_initialize(&command_buffer->leaf_tasks); memset(&command_buffer->state, 0, sizeof(command_buffer->state)); + status = iree_hal_resource_set_allocate(block_pool, + &command_buffer->resource_set); + } + if (iree_status_is_ok(status)) { *out_command_buffer = &command_buffer->base; + } else { + iree_hal_command_buffer_release(&command_buffer->base); } IREE_TRACE_ZONE_END(z0); @@ -151,6 +162,7 @@ memset(&command_buffer->state, 0, sizeof(command_buffer->state)); iree_task_list_discard(&command_buffer->leaf_tasks); iree_task_list_discard(&command_buffer->root_tasks); + iree_hal_resource_set_reset(command_buffer->resource_set); iree_arena_reset(&command_buffer->arena); } @@ -163,6 +175,7 @@ iree_hal_task_command_buffer_reset(command_buffer); iree_arena_deinitialize(&command_buffer->arena); + iree_hal_resource_set_free(command_buffer->resource_set); iree_allocator_free(host_allocator, command_buffer); IREE_TRACE_ZONE_END(z0); @@ -493,14 +506,14 @@ const iree_hal_cmd_fill_buffer_t* cmd = (const iree_hal_cmd_fill_buffer_t*)user_context; IREE_TRACE_ZONE_BEGIN(z0); - uint32_t length_per_slice = tile_context->workgroup_size[0]; - IREE_TRACE_ZONE_APPEND_VALUE(z0, length_per_slice); + uint32_t length_per_slice = tile_context->workgroup_size[0]; iree_device_size_t slice_offset = tile_context->workgroup_xyz[0] * length_per_slice; iree_device_size_t remaining_length = cmd->length - slice_offset; iree_device_size_t slice_length = iree_min(length_per_slice, remaining_length); + IREE_TRACE_ZONE_APPEND_VALUE(z0, (uint64_t)slice_length); iree_status_t status = iree_hal_buffer_fill( cmd->target_buffer, cmd->target_offset + slice_offset, slice_length, @@ -518,6 +531,9 @@ iree_hal_task_command_buffer_t* command_buffer = iree_hal_task_command_buffer_cast(base_command_buffer); + IREE_RETURN_IF_ERROR(iree_hal_resource_set_insert( + command_buffer->resource_set, 1, &target_buffer)); + iree_hal_cmd_fill_buffer_t* cmd = NULL; IREE_RETURN_IF_ERROR( iree_arena_allocate(&command_buffer->arena, sizeof(*cmd), (void**)&cmd)); @@ -577,6 +593,9 @@ iree_hal_task_command_buffer_t* command_buffer = iree_hal_task_command_buffer_cast(base_command_buffer); + IREE_RETURN_IF_ERROR(iree_hal_resource_set_insert( + command_buffer->resource_set, 1, &target_buffer)); + iree_host_size_t total_cmd_size = sizeof(iree_hal_cmd_update_buffer_t) + length; @@ -626,14 +645,14 @@ const iree_hal_cmd_copy_buffer_t* cmd = (const iree_hal_cmd_copy_buffer_t*)user_context; IREE_TRACE_ZONE_BEGIN(z0); - uint32_t length_per_slice = tile_context->workgroup_size[0]; - IREE_TRACE_ZONE_APPEND_VALUE(z0, length_per_slice); + uint32_t length_per_slice = tile_context->workgroup_size[0]; iree_device_size_t slice_offset = tile_context->workgroup_xyz[0] * length_per_slice; iree_device_size_t remaining_length = cmd->length - slice_offset; iree_device_size_t slice_length = iree_min(length_per_slice, remaining_length); + IREE_TRACE_ZONE_APPEND_VALUE(z0, (uint64_t)slice_length); iree_status_t status = iree_hal_buffer_copy_data( cmd->source_buffer, cmd->source_offset + slice_offset, cmd->target_buffer, @@ -651,6 +670,10 @@ iree_hal_task_command_buffer_t* command_buffer = iree_hal_task_command_buffer_cast(base_command_buffer); + const iree_hal_buffer_t* buffers[2] = {source_buffer, target_buffer}; + IREE_RETURN_IF_ERROR( + iree_hal_resource_set_insert(command_buffer->resource_set, 2, buffers)); + iree_hal_cmd_copy_buffer_t* cmd = NULL; IREE_RETURN_IF_ERROR( iree_arena_allocate(&command_buffer->arena, sizeof(*cmd), (void**)&cmd)); @@ -732,11 +755,16 @@ } iree_host_size_t binding_ordinal = binding_base + bindings[i].binding; + // TODO(benvanik): batch insert by getting the resources in their own list. + IREE_RETURN_IF_ERROR(iree_hal_resource_set_insert( + command_buffer->resource_set, 1, &bindings[i].buffer)); + // TODO(benvanik): track mapping so we can properly map/unmap/flush/etc. - iree_hal_buffer_mapping_t buffer_mapping; + iree_hal_buffer_mapping_t buffer_mapping = {{0}}; IREE_RETURN_IF_ERROR(iree_hal_buffer_map_range( - bindings[i].buffer, IREE_HAL_MEMORY_ACCESS_ANY, bindings[i].offset, - bindings[i].length, &buffer_mapping)); + bindings[i].buffer, IREE_HAL_MAPPING_MODE_PERSISTENT, + IREE_HAL_MEMORY_ACCESS_ANY, bindings[i].offset, bindings[i].length, + &buffer_mapping)); command_buffer->state.bindings[binding_ordinal] = buffer_mapping.contents.data; command_buffer->state.binding_lengths[binding_ordinal] = @@ -921,6 +949,10 @@ iree_hal_command_buffer_t* base_command_buffer, iree_hal_executable_t* executable, int32_t entry_point, uint32_t workgroup_x, uint32_t workgroup_y, uint32_t workgroup_z) { + iree_hal_task_command_buffer_t* command_buffer = + iree_hal_task_command_buffer_cast(base_command_buffer); + IREE_RETURN_IF_ERROR(iree_hal_resource_set_insert( + command_buffer->resource_set, 1, &executable)); iree_hal_cmd_dispatch_t* cmd = NULL; return iree_hal_task_command_buffer_build_dispatch( base_command_buffer, executable, entry_point, workgroup_x, workgroup_y, @@ -932,11 +964,19 @@ iree_hal_executable_t* executable, int32_t entry_point, iree_hal_buffer_t* workgroups_buffer, iree_device_size_t workgroups_offset) { + iree_hal_task_command_buffer_t* command_buffer = + iree_hal_task_command_buffer_cast(base_command_buffer); + + const void* resources[2] = {executable, workgroups_buffer}; + IREE_RETURN_IF_ERROR( + iree_hal_resource_set_insert(command_buffer->resource_set, 2, resources)); + // TODO(benvanik): track mapping so we can properly map/unmap/flush/etc. - iree_hal_buffer_mapping_t buffer_mapping; + iree_hal_buffer_mapping_t buffer_mapping = {{0}}; IREE_RETURN_IF_ERROR(iree_hal_buffer_map_range( - workgroups_buffer, IREE_HAL_MEMORY_ACCESS_READ, workgroups_offset, - 3 * sizeof(uint32_t), &buffer_mapping)); + workgroups_buffer, IREE_HAL_MAPPING_MODE_PERSISTENT, + IREE_HAL_MEMORY_ACCESS_READ, workgroups_offset, 3 * sizeof(uint32_t), + &buffer_mapping)); iree_hal_cmd_dispatch_t* cmd = NULL; IREE_RETURN_IF_ERROR(iree_hal_task_command_buffer_build_dispatch(
diff --git a/iree/hal/local/task_device.c b/iree/hal/local/task_device.c index be91249..4b0e1ed 100644 --- a/iree/hal/local/task_device.c +++ b/iree/hal/local/task_device.c
@@ -21,6 +21,7 @@ #include "iree/hal/local/task_event.h" #include "iree/hal/local/task_queue.h" #include "iree/hal/local/task_semaphore.h" +#include "iree/hal/utils/buffer_transfer.h" #define IREE_HAL_LOCAL_TASK_EVENT_POOL_CAPACITY 32 @@ -192,6 +193,14 @@ return device->device_allocator; } +static iree_status_t iree_hal_task_device_trim(iree_hal_device_t* base_device) { + iree_hal_task_device_t* device = iree_hal_task_device_cast(base_device); + iree_arena_block_pool_trim(&device->small_block_pool); + iree_arena_block_pool_trim(&device->large_block_pool); + iree_task_executor_trim(device->executor); + return iree_hal_allocator_trim(device->device_allocator); +} + static iree_status_t iree_hal_task_device_query_i32( iree_hal_device_t* base_device, iree_string_view_t category, iree_string_view_t key, int32_t* out_value) { @@ -350,6 +359,7 @@ .id = iree_hal_task_device_id, .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, .create_command_buffer = iree_hal_task_device_create_command_buffer, .create_descriptor_set = iree_hal_task_device_create_descriptor_set, @@ -359,6 +369,7 @@ .create_executable_cache = iree_hal_task_device_create_executable_cache, .create_executable_layout = iree_hal_task_device_create_executable_layout, .create_semaphore = iree_hal_task_device_create_semaphore, + .transfer_range = iree_hal_device_transfer_mappable_range, .queue_submit = iree_hal_task_device_queue_submit, .submit_and_wait = iree_hal_task_device_submit_and_wait, .wait_semaphores = iree_hal_task_device_wait_semaphores,
diff --git a/iree/hal/resource.h b/iree/hal/resource.h index 3f4c794..0f7abbe 100644 --- a/iree/hal/resource.h +++ b/iree/hal/resource.h
@@ -7,6 +7,7 @@ #ifndef IREE_HAL_RESOURCE_H_ #define IREE_HAL_RESOURCE_H_ +#include <assert.h> #include <stdbool.h> #include <stdint.h> @@ -33,6 +34,7 @@ iree_atomic_ref_count_t ref_count; // Opaque vtable for the resource object. + // Must start with iree_hal_resource_vtable_t at offset 0. // // NOTE: this field may be hidden in the future. Only use this for // IREE_HAL_VTABLE_DISPATCH and not equality/direct dereferencing. @@ -41,12 +43,54 @@ // TODO(benvanik): debug string/logging utilities. } iree_hal_resource_t; +// Base vtable for all resources. +// This provides the base functions required to generically manipulate resources +// of various types. +// +// This must be aliased at offset 0 of all typed vtables: +// typedef struct iree_hal_foo_vtable_t { +// void(IREE_API_PTR* destroy)(...); +// void(IREE_API_PTR* foo_method)(...); +// } iree_hal_foo_vtable_t; +typedef struct iree_hal_resource_vtable_t { + // Destroys the resource upon the final reference being released. + // The resource pointer must be assumed invalid upon return from the function + // (even if in some implementations its returned to a pool and still live). + void(IREE_API_PTR* destroy)(iree_hal_resource_t* resource); +} iree_hal_resource_vtable_t; + +// Verifies that the vtable has the right resource sub-vtable. +#define IREE_HAL_ASSERT_VTABLE_LAYOUT(vtable_type) \ + static_assert(offsetof(vtable_type, destroy) == 0, \ + "iree_hal_resource_vtable_t must be at offset 0"); + +// Initializes the base resource type. static inline void iree_hal_resource_initialize( const void* vtable, iree_hal_resource_t* out_resource) { iree_atomic_ref_count_init(&out_resource->ref_count); out_resource->vtable = vtable; } +// Retains a resource for the caller. +static inline void iree_hal_resource_retain(const void* any_resource) { + iree_hal_resource_t* resource = (iree_hal_resource_t*)any_resource; + if (IREE_LIKELY(resource)) { + iree_atomic_ref_count_inc(&resource->ref_count); + } +} + +// Releases a resource and destroys it if there are no more references. +// This routes through the vtable and can disable optimizations; always prefer +// to use the type-specific release functions (such as iree_hal_buffer_release) +// to allow for more optimizations and better compile-time type safety. +static inline void iree_hal_resource_release(const void* any_resource) { + iree_hal_resource_t* resource = (iree_hal_resource_t*)any_resource; + if (IREE_LIKELY(resource) && + iree_atomic_ref_count_dec(&resource->ref_count) == 1) { + ((iree_hal_resource_vtable_t*)resource->vtable)->destroy(resource); + } +} + // Returns true if the |resource| has the given |vtable| type. // This is *not* a way to ensure that an instance is of a specific type but // instead that it has a compatible vtable. This is because LTO may very rarely
diff --git a/iree/hal/semaphore.h b/iree/hal/semaphore.h index 5eded01..afc8959 100644 --- a/iree/hal/semaphore.h +++ b/iree/hal/semaphore.h
@@ -114,9 +114,6 @@ //===----------------------------------------------------------------------===// typedef struct iree_hal_semaphore_vtable_t { - // << HAL C porting in progress >> - IREE_API_UNSTABLE - void(IREE_API_PTR* destroy)(iree_hal_semaphore_t* semaphore); iree_status_t(IREE_API_PTR* query)(iree_hal_semaphore_t* semaphore, @@ -129,6 +126,7 @@ iree_status_t(IREE_API_PTR* wait)(iree_hal_semaphore_t* semaphore, uint64_t value, iree_timeout_t timeout); } iree_hal_semaphore_vtable_t; +IREE_HAL_ASSERT_VTABLE_LAYOUT(iree_hal_semaphore_vtable_t); IREE_API_EXPORT void iree_hal_semaphore_destroy( iree_hal_semaphore_t* semaphore);
diff --git a/iree/hal/utils/BUILD b/iree/hal/utils/BUILD index 234d6d3..42e00dc 100644 --- a/iree/hal/utils/BUILD +++ b/iree/hal/utils/BUILD
@@ -11,14 +11,64 @@ ) cc_library( + name = "buffer_transfer", + srcs = ["buffer_transfer.c"], + hdrs = ["buffer_transfer.h"], + visibility = ["//visibility:public"], + deps = [ + "//iree/base", + "//iree/base:tracing", + "//iree/hal", + ], +) + +cc_library( name = "deferred_command_buffer", srcs = ["deferred_command_buffer.c"], hdrs = ["deferred_command_buffer.h"], visibility = ["//visibility:public"], deps = [ + ":resource_set", + "//iree/base", + "//iree/base:tracing", + "//iree/base/internal:arena", + "//iree/hal", + ], +) + +cc_library( + name = "resource_set", + srcs = ["resource_set.c"], + hdrs = ["resource_set.h"], + visibility = ["//visibility:public"], + deps = [ "//iree/base", "//iree/base:tracing", "//iree/base/internal:arena", "//iree/hal", ], ) + +cc_test( + name = "resource_set_benchmark", + srcs = ["resource_set_benchmark.c"], + deps = [ + ":resource_set", + "//iree/base", + "//iree/base/internal:prng", + "//iree/hal", + "//iree/testing:benchmark", + ], +) + +cc_test( + name = "resource_set_test", + srcs = ["resource_set_test.cc"], + deps = [ + ":resource_set", + "//iree/base", + "//iree/hal", + "//iree/testing:gtest", + "//iree/testing:gtest_main", + ], +)
diff --git a/iree/hal/utils/CMakeLists.txt b/iree/hal/utils/CMakeLists.txt index 6709717..cd18b36 100644 --- a/iree/hal/utils/CMakeLists.txt +++ b/iree/hal/utils/CMakeLists.txt
@@ -12,12 +12,27 @@ iree_cc_library( NAME + buffer_transfer + HDRS + "buffer_transfer.h" + SRCS + "buffer_transfer.c" + DEPS + iree::base + iree::base::tracing + iree::hal + PUBLIC +) + +iree_cc_library( + NAME deferred_command_buffer HDRS "deferred_command_buffer.h" SRCS "deferred_command_buffer.c" DEPS + ::resource_set iree::base iree::base::internal::arena iree::base::tracing @@ -25,4 +40,45 @@ PUBLIC ) +iree_cc_library( + NAME + resource_set + HDRS + "resource_set.h" + SRCS + "resource_set.c" + DEPS + iree::base + iree::base::internal::arena + iree::base::tracing + iree::hal + PUBLIC +) + +iree_cc_test( + NAME + resource_set_benchmark + SRCS + "resource_set_benchmark.c" + DEPS + ::resource_set + iree::base + iree::base::internal::prng + iree::hal + iree::testing::benchmark +) + +iree_cc_test( + NAME + resource_set_test + SRCS + "resource_set_test.cc" + DEPS + ::resource_set + iree::base + iree::hal + iree::testing::gtest + iree::testing::gtest_main +) + ### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/iree/hal/utils/buffer_transfer.c b/iree/hal/utils/buffer_transfer.c new file mode 100644 index 0000000..258e36d --- /dev/null +++ b/iree/hal/utils/buffer_transfer.c
@@ -0,0 +1,351 @@ +// Copyright 2021 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/buffer_transfer.h" + +#include "iree/base/tracing.h" + +//===----------------------------------------------------------------------===// +// iree_hal_device_transfer_range implementations +//===----------------------------------------------------------------------===// + +IREE_API_EXPORT iree_status_t iree_hal_device_submit_transfer_range_and_wait( + iree_hal_device_t* device, iree_hal_transfer_buffer_t source, + iree_device_size_t source_offset, iree_hal_transfer_buffer_t target, + iree_device_size_t target_offset, iree_device_size_t data_length, + iree_hal_transfer_buffer_flags_t flags, iree_timeout_t timeout) { + // If the source and target are both mappable into host memory (or are host + // memory) then we can use the fast zero-alloc path. This may actually be + // slower than doing a device queue transfer depending on the size of the data + // and where the memory lives. For example, if we have two device buffers in + // device-local host-visible memory we'd be performing the transfer by pulling + // all the memory to the CPU and pushing it back again. + // TODO(benvanik): check for device-local -> device-local and avoid mapping. + bool is_source_mappable = + !source.device_buffer || + iree_all_bits_set(iree_hal_buffer_allowed_usage(source.device_buffer), + IREE_HAL_BUFFER_USAGE_MAPPING); + bool is_target_mappable = + !target.device_buffer || + iree_all_bits_set(iree_hal_buffer_allowed_usage(target.device_buffer), + IREE_HAL_BUFFER_USAGE_MAPPING); + if (is_source_mappable && is_target_mappable) { + return iree_hal_device_transfer_mappable_range( + device, source, source_offset, target, target_offset, data_length, + flags, timeout); + } + + // If the source is a host buffer under 64KB then we can do a more efficient + // (though still relatively costly) update instead of needing a staging + // buffer. + if (!source.device_buffer && target.device_buffer && + data_length <= IREE_HAL_COMMAND_BUFFER_MAX_UPDATE_SIZE) { + const iree_hal_transfer_command_t transfer_command = { + .type = IREE_HAL_TRANSFER_COMMAND_TYPE_UPDATE, + .update = + { + .source_buffer = source.host_buffer.data, + .source_offset = source_offset, + .target_buffer = target.device_buffer, + .target_offset = target_offset, + .length = data_length, + }, + }; + return iree_hal_device_transfer_and_wait(device, /*wait_semaphore=*/NULL, + /*wait_value=*/0ull, 1, + &transfer_command, timeout); + } + + iree_status_t status = iree_ok_status(); + + // Allocate the staging buffer for upload to the device. + iree_hal_buffer_t* source_buffer = source.device_buffer; + if (!source_buffer) { + // Allocate staging memory with a copy of the host data. We only initialize + // the portion being transferred. + // TODO(benvanik): use wrap_buffer if supported to avoid the + // allocation/copy. + status = iree_hal_allocator_allocate_buffer( + iree_hal_device_allocator(device), + IREE_HAL_MEMORY_TYPE_HOST_LOCAL | IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE, + IREE_HAL_BUFFER_USAGE_TRANSFER | IREE_HAL_BUFFER_USAGE_MAPPING, + data_length, + iree_make_const_byte_span(source.host_buffer.data + source_offset, + data_length), + &source_buffer); + source_offset = 0; + } + + // Allocate the staging buffer for download from the device. + iree_hal_buffer_t* target_buffer = target.device_buffer; + if (!target_buffer) { + // Allocate uninitialized staging memory for the transfer target. + // We only allocate enough for the portion we are transfering. + // TODO(benvanik): use wrap_buffer if supported to avoid the + // allocation/copy. + status = iree_hal_allocator_allocate_buffer( + iree_hal_device_allocator(device), + IREE_HAL_MEMORY_TYPE_HOST_LOCAL | IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE, + IREE_HAL_BUFFER_USAGE_TRANSFER | IREE_HAL_BUFFER_USAGE_MAPPING, + data_length, iree_const_byte_span_empty(), &target_buffer); + target_offset = 0; + } + + // Issue synchronous device copy. + if (iree_status_is_ok(status)) { + const iree_hal_transfer_command_t transfer_command = { + .type = IREE_HAL_TRANSFER_COMMAND_TYPE_COPY, + .copy = + { + .source_buffer = source_buffer, + .source_offset = source_offset, + .target_buffer = target_buffer, + .target_offset = target_offset, + .length = data_length, + }, + }; + status = iree_hal_device_transfer_and_wait(device, /*wait_semaphore=*/NULL, + /*wait_value=*/0ull, 1, + &transfer_command, timeout); + } + + // Read back the staging buffer into memory, if needed. + if (iree_status_is_ok(status) && !target.device_buffer) { + status = iree_hal_buffer_read_data(target_buffer, 0, + target.host_buffer.data, data_length); + } + + // Discard staging buffers, if they were required. + if (!source.device_buffer) iree_hal_buffer_release(source_buffer); + if (!target.device_buffer) iree_hal_buffer_release(target_buffer); + + return status; +} + +IREE_API_EXPORT iree_status_t iree_hal_device_transfer_mappable_range( + iree_hal_device_t* device, iree_hal_transfer_buffer_t source, + iree_device_size_t source_offset, iree_hal_transfer_buffer_t target, + iree_device_size_t target_offset, iree_device_size_t data_length, + iree_hal_transfer_buffer_flags_t flags, iree_timeout_t timeout) { + iree_status_t status = iree_ok_status(); + + iree_hal_buffer_mapping_t source_mapping = {{0}}; + if (iree_status_is_ok(status)) { + if (source.device_buffer) { + status = iree_hal_buffer_map_range( + source.device_buffer, IREE_HAL_MAPPING_MODE_SCOPED, + IREE_HAL_MEMORY_ACCESS_READ, source_offset, data_length, + &source_mapping); + } else { + source_mapping = (iree_hal_buffer_mapping_t){ + .contents = source.host_buffer, + }; + } + } + + iree_hal_buffer_mapping_t target_mapping = {{0}}; + if (iree_status_is_ok(status)) { + if (target.device_buffer) { + status = iree_hal_buffer_map_range( + target.device_buffer, IREE_HAL_MAPPING_MODE_SCOPED, + IREE_HAL_MEMORY_ACCESS_DISCARD_WRITE, target_offset, data_length, + &target_mapping); + } else { + target_mapping = (iree_hal_buffer_mapping_t){ + .contents = target.host_buffer, + }; + } + } + + iree_device_size_t adjusted_data_length = 0; + if (iree_status_is_ok(status)) { + // Adjust the data length based on the min we have. + if (data_length == IREE_WHOLE_BUFFER) { + // Whole buffer copy requested - that could mean either, so take the min. + adjusted_data_length = iree_min(source_mapping.contents.data_length, + target_mapping.contents.data_length); + } else { + // Specific length requested - validate that we have matching lengths. + IREE_ASSERT_EQ(source_mapping.contents.data_length, + target_mapping.contents.data_length); + adjusted_data_length = target_mapping.contents.data_length; + } + + // Perform the copy, assuming there's anything to do. + if (adjusted_data_length != 0) { + memcpy(target_mapping.contents.data, source_mapping.contents.data, + adjusted_data_length); + } + } + + if (source.device_buffer) { + status = + iree_status_join(status, iree_hal_buffer_unmap_range(&source_mapping)); + } + if (target.device_buffer) { + if (adjusted_data_length > 0 && + !iree_all_bits_set(iree_hal_buffer_memory_type(target.device_buffer), + IREE_HAL_MEMORY_TYPE_HOST_COHERENT)) { + status = iree_status_join( + status, iree_hal_buffer_flush_range(&target_mapping, 0, + adjusted_data_length)); + } + status = + iree_status_join(status, iree_hal_buffer_unmap_range(&target_mapping)); + } + return status; +} + +//===----------------------------------------------------------------------===// +// iree_hal_buffer_map_range implementations +//===----------------------------------------------------------------------===// + +typedef struct iree_hal_emulated_buffer_mapping_t { + iree_hal_buffer_t* host_local_buffer; + iree_hal_buffer_mapping_t host_local_mapping; +} iree_hal_emulated_buffer_mapping_t; + +IREE_API_EXPORT iree_status_t iree_hal_buffer_emulated_map_range( + iree_hal_device_t* device, iree_hal_buffer_t* buffer, + iree_hal_mapping_mode_t mapping_mode, + iree_hal_memory_access_t memory_access, + iree_device_size_t local_byte_offset, iree_device_size_t local_byte_length, + iree_hal_buffer_mapping_t* mapping) { + IREE_ASSERT_ARGUMENT(device); + IREE_ASSERT_ARGUMENT(buffer); + IREE_ASSERT_ARGUMENT(mapping); + + iree_hal_allocator_t* device_allocator = iree_hal_device_allocator(device); + iree_allocator_t host_allocator = iree_hal_device_host_allocator(device); + + // We can't perform persistent mapping with this as we need to manage the + // staging buffer lifetime. + if (IREE_UNLIKELY(mapping_mode == IREE_HAL_MAPPING_MODE_PERSISTENT)) { + return iree_make_status( + IREE_STATUS_INVALID_ARGUMENT, + "emulated buffer mapping only possible with scoped mappings"); + } + + // No implementation should be using this emulated method with memory that is + // allocated as mappable. + if (IREE_UNLIKELY(iree_all_bits_set(iree_hal_buffer_memory_type(buffer), + IREE_HAL_BUFFER_USAGE_MAPPING))) { + return iree_make_status( + IREE_STATUS_FAILED_PRECONDITION, + "emulated buffer mapping should not be used with mappable buffers"); + } + + IREE_TRACE_ZONE_BEGIN(z0); + IREE_TRACE_ZONE_APPEND_VALUE(z0, (uint64_t)local_byte_length); + + // NOTE: this is assuming that the host is going to be doing a lot of work + // on the mapped memory and wants read/write caching and such. If the user + // wants write combining on device memory and other things they should ensure + // this emulated mapping path is not hit. + + // Create a transient struct we use to track the emulated operation. + // We could pack this into the mapping but this composes better - it's small + // and pooled by the host allocator anyway. + iree_hal_emulated_buffer_mapping_t* emulation_state = NULL; + IREE_RETURN_AND_END_ZONE_IF_ERROR( + z0, iree_allocator_malloc(host_allocator, sizeof(*emulation_state), + (void**)&emulation_state)); + + // Allocate the buffer we'll be using to stage our copy of the device memory. + // All devices should be able to satisfy this host-local + mapping request. + iree_status_t status = iree_hal_allocator_allocate_buffer( + device_allocator, IREE_HAL_MEMORY_TYPE_HOST_LOCAL, + IREE_HAL_BUFFER_USAGE_TRANSFER | IREE_HAL_BUFFER_USAGE_MAPPING, + local_byte_length, iree_const_byte_span_empty(), + &emulation_state->host_local_buffer); + + // We need to capture a copy of the device buffer to work with; unless the + // user was nice and said they don't care about the contents with the DISCARD + // bit. Ideally we'd also enable invalidate_range to specify subranges we want + // to map. + if (iree_status_is_ok(status) && + !iree_all_bits_set(memory_access, IREE_HAL_MEMORY_ACCESS_DISCARD)) { + // Download (device->host) the data. + status = iree_hal_device_transfer_range( + device, iree_hal_make_device_transfer_buffer(mapping->buffer), + local_byte_offset, + iree_hal_make_device_transfer_buffer( + emulation_state->host_local_buffer), + 0, local_byte_length, IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, + iree_infinite_timeout()); + } + + if (iree_status_is_ok(status)) { + // Map the scratch buffer: map-ception. + // Code-wise it looks like this may loop back onto this emulated path + // but no implementation should be using this emulation if they have host + // local IREE_HAL_BUFFER_USAGE_MAPPING memory - and we check that above. + status = iree_hal_buffer_map_range(emulation_state->host_local_buffer, + IREE_HAL_MAPPING_MODE_SCOPED, + memory_access, 0, local_byte_length, + &emulation_state->host_local_mapping); + } + + // Retain the scratch buffer for the duration of the mapping. + if (iree_status_is_ok(status)) { + // Note that we are giving back the host-local mapped contents to the user - + // they don't need to know it's from our staging buffer. + mapping->contents = emulation_state->host_local_mapping.contents; + mapping->impl.reserved[0] = (uint64_t)((uintptr_t)emulation_state); + } else { + status = iree_status_join( + status, + iree_hal_buffer_unmap_range(&emulation_state->host_local_mapping)); + iree_hal_buffer_release(emulation_state->host_local_buffer); + iree_allocator_free(host_allocator, emulation_state); + } + IREE_TRACE_ZONE_END(z0); + return status; +} + +IREE_API_EXPORT iree_status_t iree_hal_buffer_emulated_unmap_range( + iree_hal_device_t* device, iree_hal_buffer_t* buffer, + iree_device_size_t local_byte_offset, iree_device_size_t local_byte_length, + iree_hal_buffer_mapping_t* mapping) { + IREE_ASSERT_ARGUMENT(device); + IREE_ASSERT_ARGUMENT(buffer); + IREE_ASSERT_ARGUMENT(mapping); + IREE_TRACE_ZONE_BEGIN(z0); + IREE_TRACE_ZONE_APPEND_VALUE(z0, (uint64_t)local_byte_length); + iree_hal_emulated_buffer_mapping_t* emulation_state = + (iree_hal_emulated_buffer_mapping_t*)((uintptr_t) + mapping->impl.reserved[0]); + IREE_ASSERT_NE(emulation_state, NULL); + + // Unmap the scratch buffer first to make it available for copying (if + // needed). + iree_status_t status = + iree_hal_buffer_unmap_range(&emulation_state->host_local_mapping); + + // If we were writing then we'll need to flush the range. + // Ideally we'd keep track of this on the mapping itself based on the user's + // calls to flush_range to limit how much we need to transfer. + if (iree_status_is_ok(status) && + iree_all_bits_set(mapping->impl.allowed_access, + IREE_HAL_MEMORY_ACCESS_WRITE)) { + // Upload (host->device) the data. + status = iree_hal_device_transfer_range( + device, + iree_hal_make_device_transfer_buffer( + emulation_state->host_local_buffer), + 0, iree_hal_make_device_transfer_buffer(mapping->buffer), + local_byte_offset, local_byte_length, + IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, iree_infinite_timeout()); + } + + // Deallocate the scratch buffer and our emulation state. + iree_hal_buffer_release(emulation_state->host_local_buffer); + iree_allocator_t host_allocator = iree_hal_device_host_allocator(device); + iree_allocator_free(host_allocator, emulation_state); + + IREE_TRACE_ZONE_END(z0); + return status; +}
diff --git a/iree/hal/utils/buffer_transfer.h b/iree/hal/utils/buffer_transfer.h new file mode 100644 index 0000000..2daac0a --- /dev/null +++ b/iree/hal/utils/buffer_transfer.h
@@ -0,0 +1,69 @@ +// Copyright 2021 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 + +#ifndef IREE_HAL_UTILS_BUFFER_TRANSFER_H_ +#define IREE_HAL_UTILS_BUFFER_TRANSFER_H_ + +#include "iree/base/api.h" +#include "iree/hal/api.h" + +#ifdef __cplusplus +extern "C" { +#endif // __cplusplus + +//===----------------------------------------------------------------------===// +// iree_hal_device_transfer_range implementations +//===----------------------------------------------------------------------===// + +// Performs a full transfer operation on a device transfer queue. +// This creates a transfer command buffer, submits it against the device, and +// waits for it to complete synchronously. Implementations that can do this +// cheaper are encouraged to do so. +// +// Precondition: source and target do not overlap. +IREE_API_EXPORT iree_status_t iree_hal_device_submit_transfer_range_and_wait( + iree_hal_device_t* device, iree_hal_transfer_buffer_t source, + iree_device_size_t source_offset, iree_hal_transfer_buffer_t target, + iree_device_size_t target_offset, iree_device_size_t data_length, + iree_hal_transfer_buffer_flags_t flags, iree_timeout_t timeout); + +// Generic implementation of iree_hal_device_transfer_range for when the buffers +// are mappable. In certain implementations even if buffers are mappable it's +// often cheaper to still use the full queue transfers: instead of wasting CPU +// cycles copying the memory (and possible PCIe round-trips) letting the device +// do it is effectively free. +// +// Precondition: source and target do not overlap. +IREE_API_EXPORT iree_status_t iree_hal_device_transfer_mappable_range( + iree_hal_device_t* device, iree_hal_transfer_buffer_t source, + iree_device_size_t source_offset, iree_hal_transfer_buffer_t target, + iree_device_size_t target_offset, iree_device_size_t data_length, + iree_hal_transfer_buffer_flags_t flags, iree_timeout_t timeout); + +//===----------------------------------------------------------------------===// +// iree_hal_buffer_map_range implementations +//===----------------------------------------------------------------------===// + +// Generic implementation of iree_hal_buffer_map_range and unmap_range for when +// the buffer is not mappable and a full device transfer is required. This will +// allocate additional host-local buffers and submit copy commands. +// Implementations able to do this more efficiently should do so. +IREE_API_EXPORT iree_status_t iree_hal_buffer_emulated_map_range( + iree_hal_device_t* device, iree_hal_buffer_t* buffer, + iree_hal_mapping_mode_t mapping_mode, + iree_hal_memory_access_t memory_access, + iree_device_size_t local_byte_offset, iree_device_size_t local_byte_length, + iree_hal_buffer_mapping_t* mapping); +IREE_API_EXPORT iree_status_t iree_hal_buffer_emulated_unmap_range( + iree_hal_device_t* device, iree_hal_buffer_t* buffer, + iree_device_size_t local_byte_offset, iree_device_size_t local_byte_length, + iree_hal_buffer_mapping_t* mapping); + +#ifdef __cplusplus +} // extern "C" +#endif // __cplusplus + +#endif // IREE_HAL_UTILS_BUFFER_TRANSFER_H_
diff --git a/iree/hal/utils/deferred_command_buffer.c b/iree/hal/utils/deferred_command_buffer.c index 8f2a05a..347a222 100644 --- a/iree/hal/utils/deferred_command_buffer.c +++ b/iree/hal/utils/deferred_command_buffer.c
@@ -8,6 +8,7 @@ #include "iree/base/internal/arena.h" #include "iree/base/tracing.h" +#include "iree/hal/utils/resource_set.h" //===----------------------------------------------------------------------===// // Command recording structures @@ -134,6 +135,12 @@ typedef struct iree_hal_deferred_command_buffer_t { iree_hal_command_buffer_t base; iree_allocator_t host_allocator; + + // Maintains a reference to all resources used within the command buffer. + // Reset on each begin. + iree_hal_resource_set_t* resource_set; + + // All commands in encoding order. iree_hal_cmd_list_t cmd_list; } iree_hal_deferred_command_buffer_t; @@ -165,9 +172,16 @@ &iree_hal_deferred_command_buffer_vtable, &command_buffer->base); command_buffer->host_allocator = host_allocator; iree_hal_cmd_list_initialize(block_pool, &command_buffer->cmd_list); + + status = iree_hal_resource_set_allocate(block_pool, + &command_buffer->resource_set); } - *out_command_buffer = &command_buffer->base; + if (iree_status_is_ok(status)) { + *out_command_buffer = &command_buffer->base; + } else { + iree_hal_command_buffer_destroy(&command_buffer->base); + } IREE_TRACE_ZONE_END(z0); return status; } @@ -180,6 +194,7 @@ IREE_TRACE_ZONE_BEGIN(z0); iree_hal_cmd_list_deinitialize(&command_buffer->cmd_list); + iree_hal_resource_set_free(command_buffer->resource_set); iree_allocator_free(host_allocator, command_buffer); IREE_TRACE_ZONE_END(z0); @@ -199,6 +214,7 @@ iree_hal_deferred_command_buffer_t* command_buffer = iree_hal_deferred_command_buffer_cast(base_command_buffer); iree_hal_cmd_list_reset(&command_buffer->cmd_list); + iree_hal_resource_set_reset(command_buffer->resource_set); return iree_ok_status(); } @@ -280,8 +296,11 @@ static iree_status_t iree_hal_deferred_command_buffer_signal_event( iree_hal_command_buffer_t* base_command_buffer, iree_hal_event_t* event, iree_hal_execution_stage_t source_stage_mask) { - iree_hal_cmd_list_t* cmd_list = - &iree_hal_deferred_command_buffer_cast(base_command_buffer)->cmd_list; + 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; + IREE_RETURN_IF_ERROR( + iree_hal_resource_set_insert(command_buffer->resource_set, 1, &event)); iree_hal_cmd_signal_event_t* cmd = NULL; IREE_RETURN_IF_ERROR(iree_hal_cmd_list_append_command( cmd_list, IREE_HAL_CMD_SIGNAL_EVENT, sizeof(*cmd), (void**)&cmd)); @@ -310,8 +329,11 @@ static iree_status_t iree_hal_deferred_command_buffer_reset_event( iree_hal_command_buffer_t* base_command_buffer, iree_hal_event_t* event, iree_hal_execution_stage_t source_stage_mask) { - iree_hal_cmd_list_t* cmd_list = - &iree_hal_deferred_command_buffer_cast(base_command_buffer)->cmd_list; + 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; + IREE_RETURN_IF_ERROR( + iree_hal_resource_set_insert(command_buffer->resource_set, 1, &event)); iree_hal_cmd_reset_event_t* cmd = NULL; IREE_RETURN_IF_ERROR(iree_hal_cmd_list_append_command( cmd_list, IREE_HAL_CMD_RESET_EVENT, sizeof(*cmd), (void**)&cmd)); @@ -352,8 +374,11 @@ const iree_hal_memory_barrier_t* memory_barriers, iree_host_size_t buffer_barrier_count, const iree_hal_buffer_barrier_t* buffer_barriers) { - iree_hal_cmd_list_t* cmd_list = - &iree_hal_deferred_command_buffer_cast(base_command_buffer)->cmd_list; + 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; + IREE_RETURN_IF_ERROR(iree_hal_resource_set_insert( + command_buffer->resource_set, event_count, events)); iree_hal_cmd_wait_events_t* cmd = NULL; IREE_RETURN_IF_ERROR(iree_hal_cmd_list_append_command( cmd_list, IREE_HAL_CMD_WAIT_EVENTS, @@ -402,8 +427,11 @@ static iree_status_t iree_hal_deferred_command_buffer_discard_buffer( iree_hal_command_buffer_t* base_command_buffer, iree_hal_buffer_t* buffer) { - iree_hal_cmd_list_t* cmd_list = - &iree_hal_deferred_command_buffer_cast(base_command_buffer)->cmd_list; + 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; + IREE_RETURN_IF_ERROR( + iree_hal_resource_set_insert(command_buffer->resource_set, 1, &buffer)); iree_hal_cmd_discard_buffer_t* cmd = NULL; IREE_RETURN_IF_ERROR(iree_hal_cmd_list_append_command( cmd_list, IREE_HAL_CMD_DISCARD_BUFFER, sizeof(*cmd), (void**)&cmd)); @@ -436,13 +464,16 @@ iree_hal_buffer_t* target_buffer, iree_device_size_t target_offset, iree_device_size_t length, const void* pattern, iree_host_size_t pattern_length) { - iree_hal_cmd_list_t* cmd_list = - &iree_hal_deferred_command_buffer_cast(base_command_buffer)->cmd_list; + 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; iree_hal_cmd_fill_buffer_t* cmd = NULL; if (pattern_length > sizeof(cmd->pattern)) { return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, "fill patterns must be < 8 bytes"); } + IREE_RETURN_IF_ERROR(iree_hal_resource_set_insert( + command_buffer->resource_set, 1, &target_buffer)); IREE_RETURN_IF_ERROR(iree_hal_cmd_list_append_command( cmd_list, IREE_HAL_CMD_FILL_BUFFER, sizeof(*cmd), (void**)&cmd)); cmd->target_buffer = target_buffer; @@ -477,8 +508,11 @@ iree_hal_command_buffer_t* base_command_buffer, const void* source_buffer, iree_host_size_t source_offset, iree_hal_buffer_t* target_buffer, iree_device_size_t target_offset, iree_device_size_t length) { - iree_hal_cmd_list_t* cmd_list = - &iree_hal_deferred_command_buffer_cast(base_command_buffer)->cmd_list; + 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; + IREE_RETURN_IF_ERROR(iree_hal_resource_set_insert( + command_buffer->resource_set, 1, &target_buffer)); iree_hal_cmd_update_buffer_t* cmd = NULL; IREE_RETURN_IF_ERROR(iree_hal_cmd_list_append_command( cmd_list, IREE_HAL_CMD_UPDATE_BUFFER, @@ -517,8 +551,12 @@ iree_hal_buffer_t* source_buffer, iree_device_size_t source_offset, iree_hal_buffer_t* target_buffer, iree_device_size_t target_offset, iree_device_size_t length) { - iree_hal_cmd_list_t* cmd_list = - &iree_hal_deferred_command_buffer_cast(base_command_buffer)->cmd_list; + 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; + const void* buffers[2] = {source_buffer, target_buffer}; + IREE_RETURN_IF_ERROR( + iree_hal_resource_set_insert(command_buffer->resource_set, 2, buffers)); iree_hal_cmd_copy_buffer_t* cmd = NULL; IREE_RETURN_IF_ERROR(iree_hal_cmd_list_append_command( cmd_list, IREE_HAL_CMD_COPY_BUFFER, sizeof(*cmd), (void**)&cmd)); @@ -554,8 +592,11 @@ iree_hal_command_buffer_t* base_command_buffer, iree_hal_executable_layout_t* executable_layout, iree_host_size_t offset, const void* values, iree_host_size_t values_length) { - iree_hal_cmd_list_t* cmd_list = - &iree_hal_deferred_command_buffer_cast(base_command_buffer)->cmd_list; + 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; + IREE_RETURN_IF_ERROR(iree_hal_resource_set_insert( + command_buffer->resource_set, 1, &executable_layout)); iree_hal_cmd_push_constants_t* cmd = NULL; IREE_RETURN_IF_ERROR(iree_hal_cmd_list_append_command( cmd_list, IREE_HAL_CMD_PUSH_CONSTANTS, @@ -592,8 +633,15 @@ iree_hal_executable_layout_t* executable_layout, uint32_t set, iree_host_size_t binding_count, const iree_hal_descriptor_set_binding_t* bindings) { - iree_hal_cmd_list_t* cmd_list = - &iree_hal_deferred_command_buffer_cast(base_command_buffer)->cmd_list; + 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; + IREE_RETURN_IF_ERROR(iree_hal_resource_set_insert( + command_buffer->resource_set, 1, &executable_layout)); + for (iree_host_size_t i = 0; i < binding_count; ++i) { + IREE_RETURN_IF_ERROR(iree_hal_resource_set_insert( + command_buffer->resource_set, 1, &bindings[i].buffer)); + } iree_hal_cmd_push_descriptor_set_t* cmd = NULL; IREE_RETURN_IF_ERROR(iree_hal_cmd_list_append_command( cmd_list, IREE_HAL_CMD_PUSH_DESCRIPTOR_SET, @@ -632,8 +680,12 @@ iree_hal_descriptor_set_t* descriptor_set, iree_host_size_t dynamic_offset_count, const iree_device_size_t* dynamic_offsets) { - iree_hal_cmd_list_t* cmd_list = - &iree_hal_deferred_command_buffer_cast(base_command_buffer)->cmd_list; + 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; + const void* resources[2] = {executable_layout, descriptor_set}; + IREE_RETURN_IF_ERROR( + iree_hal_resource_set_insert(command_buffer->resource_set, 2, resources)); iree_hal_cmd_bind_descriptor_set_t* cmd = NULL; IREE_RETURN_IF_ERROR(iree_hal_cmd_list_append_command( cmd_list, IREE_HAL_CMD_BIND_DESCRIPTOR_SET, @@ -673,8 +725,11 @@ iree_hal_command_buffer_t* base_command_buffer, iree_hal_executable_t* executable, int32_t entry_point, uint32_t workgroup_x, uint32_t workgroup_y, uint32_t workgroup_z) { - iree_hal_cmd_list_t* cmd_list = - &iree_hal_deferred_command_buffer_cast(base_command_buffer)->cmd_list; + 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; + IREE_RETURN_IF_ERROR(iree_hal_resource_set_insert( + command_buffer->resource_set, 1, &executable)); iree_hal_cmd_dispatch_t* cmd = NULL; IREE_RETURN_IF_ERROR(iree_hal_cmd_list_append_command( cmd_list, IREE_HAL_CMD_DISPATCH, sizeof(*cmd), (void**)&cmd)); @@ -711,8 +766,12 @@ iree_hal_executable_t* executable, int32_t entry_point, iree_hal_buffer_t* workgroups_buffer, iree_device_size_t workgroups_offset) { - iree_hal_cmd_list_t* cmd_list = - &iree_hal_deferred_command_buffer_cast(base_command_buffer)->cmd_list; + 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; + const void* resources[2] = {executable, workgroups_buffer}; + IREE_RETURN_IF_ERROR( + iree_hal_resource_set_insert(command_buffer->resource_set, 2, resources)); iree_hal_cmd_dispatch_indirect_t* cmd = NULL; IREE_RETURN_IF_ERROR(iree_hal_cmd_list_append_command( cmd_list, IREE_HAL_CMD_DISPATCH_INDIRECT, sizeof(*cmd), (void**)&cmd));
diff --git a/iree/hal/utils/resource_set.c b/iree/hal/utils/resource_set.c new file mode 100644 index 0000000..14e5871 --- /dev/null +++ b/iree/hal/utils/resource_set.c
@@ -0,0 +1,276 @@ +// Copyright 2022 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/resource_set.h" + +#include "iree/base/tracing.h" + +// Inlines the first chunk into the block using all of the remaining space. +// This is a special case chunk that is released back to the pool with the +// resource set and lets us avoid an additional allocation. +static void iree_hal_resource_set_setup_inline_chunk( + iree_hal_resource_set_t* set) { + uint8_t* block_ptr = (uint8_t*)set + sizeof(*set); + iree_hal_resource_set_chunk_t* inlined_chunk = + (iree_hal_resource_set_chunk_t*)block_ptr; + inlined_chunk->flags = IREE_HAL_RESOURCE_SET_CHUNK_FLAG_INLINE; + inlined_chunk->capacity = (set->block_pool->total_block_size - sizeof(*set) - + sizeof(*inlined_chunk)) / + sizeof(iree_hal_resource_t*); + inlined_chunk->capacity = iree_min(inlined_chunk->capacity, + IREE_HAL_RESOURCE_SET_CHUNK_MAX_CAPACITY); + inlined_chunk->count = 0; + set->chunk_head = inlined_chunk; +} + +IREE_API_EXPORT iree_status_t iree_hal_resource_set_allocate( + iree_arena_block_pool_t* block_pool, iree_hal_resource_set_t** out_set) { + IREE_TRACE_ZONE_BEGIN(z0); + + // We could allow larger sizes (would require widening the capacity/count + // fields in the chunk) but in real usage having even 64k is a bit too much. + IREE_ASSERT_LE(block_pool->total_block_size, 64 * 1024, + "keep block sizes small for resource sets"); + + // Acquire block and place the set struct at the head. + iree_arena_block_t* block = NULL; + IREE_RETURN_AND_END_ZONE_IF_ERROR( + z0, iree_arena_block_pool_acquire(block_pool, &block)); + uint8_t* block_ptr = (uint8_t*)block - block_pool->usable_block_size; + iree_hal_resource_set_t* set = (iree_hal_resource_set_t*)block_ptr; + memset(set, 0, sizeof(*set)); + set->block_pool = block_pool; + iree_hal_resource_set_setup_inline_chunk(set); + + *out_set = set; + IREE_TRACE_ZONE_END(z0); + return iree_ok_status(); +} + +static void iree_hal_resource_set_release_blocks(iree_hal_resource_set_t* set, + bool preserve_set) { + // Release all resources in all chunks and stitch together the blocks in a + // linked list. We do this first so that we can release all of the chunks back + // to the block pool in one operation. Ideally we'd maintain the linked list + // in our chunks but there's some weirdness with prefix/suffix header/footers + // that isn't worth the complexity. + iree_arena_block_t* block_head = NULL; + iree_arena_block_t* block_tail = NULL; + iree_hal_resource_set_chunk_t* chunk = set->chunk_head; + while (chunk) { + // Release all resources in the chunk. + for (iree_host_size_t i = 0; i < chunk->count; ++i) { + iree_hal_resource_release(chunk->resources[i]); + } + // Consume the chunk and add it to the block pool release linked list. + iree_hal_resource_set_chunk_t* next_chunk = chunk->next_chunk; + iree_arena_block_t* block = NULL; + if (iree_hal_resource_set_chunk_is_stored_inline(chunk)) { + // This is the inlined first chunk that also stores the set header. + // If we are not freeing the set then we don't release the block back to + // the pool. + if (preserve_set) { + // Don't release the block. + break; + } else { + block = (iree_arena_block_t*)((uint8_t*)set + + set->block_pool->usable_block_size); + next_chunk = NULL; + } + } else { + // A chunk acquired after the set was acquired. + block = (iree_arena_block_t*)((uint8_t*)chunk + + set->block_pool->usable_block_size); + } + block->next = block_head; + block_head = block; + if (!block_tail) block_tail = block; + chunk = next_chunk; + } + + // Release all blocks back to the block pool in one operation. + // NOTE: this invalidates the |set| memory. + iree_arena_block_pool_t* block_pool = set->block_pool; + iree_arena_block_pool_release(block_pool, block_head, block_tail); +} + +IREE_API_EXPORT void iree_hal_resource_set_free(iree_hal_resource_set_t* set) { + IREE_TRACE_ZONE_BEGIN(z0); + + // Release all resources and the arena block used by the set. + // The set pointer is invalid after this call returns. + iree_hal_resource_set_release_blocks(set, /*preserve_set=*/false); + + IREE_TRACE_ZONE_END(z0); +} + +IREE_API_EXPORT void iree_hal_resource_set_reset(iree_hal_resource_set_t* set) { + IREE_TRACE_ZONE_BEGIN(z0); + + // Release all resources and the blocks besides the base set. + iree_hal_resource_set_release_blocks(set, /*preserve_set=*/true); + + // Reset the set state. + memset(set->mru, 0, sizeof(set->mru)); + iree_hal_resource_set_setup_inline_chunk(set); + + IREE_TRACE_ZONE_END(z0); +} + +// Retains |resource| and adds it to the main |set| list. +static iree_status_t iree_hal_resource_set_insert_retain( + iree_hal_resource_set_t* set, iree_hal_resource_t* resource) { + iree_hal_resource_set_chunk_t* chunk = set->chunk_head; + if (IREE_UNLIKELY(chunk->count + 1 > chunk->capacity)) { + // Ran out of room in the current chunk - acquire a new one and link it into + // the list of chunks. + iree_arena_block_t* block = NULL; + IREE_RETURN_IF_ERROR( + iree_arena_block_pool_acquire(set->block_pool, &block)); + chunk = + (iree_hal_resource_set_chunk_t*)((uint8_t*)block - + set->block_pool->usable_block_size); + chunk->next_chunk = set->chunk_head; + set->chunk_head = chunk; + chunk->capacity = (set->block_pool->total_block_size - sizeof(*chunk)) / + sizeof(iree_hal_resource_t*); + chunk->capacity = + iree_min(chunk->capacity, IREE_HAL_RESOURCE_SET_CHUNK_MAX_CAPACITY); + chunk->count = 0; + } + + // Retain and insert into the chunk. + chunk->resources[chunk->count++] = resource; + iree_hal_resource_retain(resource); + return iree_ok_status(); +} + +// Scans the lookaside for the resource pointer and updates the order if found. +// If the resource was not found then it will be inserted into the main list as +// well as the MRU. +// +// This performs a full scan over the MRU and if the resource is found will +// move the resource to the front of the list before returning. Otherwise the +// resource will be retained in the main source-of-truth list. +// +// Example (hit): +// +----+----+----+----+ +// | AA | BB | CC | DD | resource: CC +// +----+----+----+----+ +// scan mru to find CC: +// found at mru[2] +// shift prefix down 1: +// +----+----+----+----+ +// | AA | AA | BB | DD | +// +----+----+----+----+ +// insert resource at front: +// +----+----+----+----+ +// | CC | AA | BB | DD | +// +----+----+----+----+ +// +// Example (miss): +// +----+----+----+----+ +// | AA | BB | CC | DD | resource: EE +// +----+----+----+----+ +// scan mru to find EE: not found +// shift set down 1: +// +----+----+----+----+ +// | AA | AA | BB | CC | +// +----+----+----+----+ +// insert resource at front: +// +----+----+----+----+ +// | EE | AA | BB | CC | +// +----+----+----+----+ +// insert resource into main list +// +// The intent here is that we can model this behavior with SIMD ops to perform +// both the scan and update using comparison, extraction, and permutation. The +// best and worst case flows will load the entire MRU into registers from a +// single cache line, do all the scanning and shifting in registers, and then +// store back to the single cache line. +// +// Today, though, we leave this as an exercise to whoever comes across this :) +// Notes: +// As the MRU is a fixed size we can unroll it entirely and avoid any looping. +// On a 32-bit system with uint32x4_t we only need 4 registers. +// On a 64-bit system with uint64x2_t we also only need 4 registers - though +// the MRU has half as many entries and we may want to go >1 cache line. +// +// If we wanted to process more than one resource at a time we can specialize +// the code paths to handle 1/2/4/etc resources and process in batches with +// an optional remainder. This would increase the ratio of work performed on +// the loaded MRU registers before we do the shift/store. +// +// The tree sequence we likely want is something like: +// https://developer.arm.com/architectures/instruction-sets/intrinsics/vdupq_n_u32 +// https://developer.arm.com/architectures/instruction-sets/intrinsics/vceqq_u32 +// https://developer.arm.com/architectures/instruction-sets/intrinsics/vorrq_u32 +// https://developer.arm.com/architectures/instruction-sets/intrinsics/vmaxvq_u32 +// or +// https://developer.arm.com/architectures/instruction-sets/intrinsics/vdupq_n_u64 +// https://developer.arm.com/architectures/instruction-sets/intrinsics/vceqq_u64 +// https://developer.arm.com/architectures/instruction-sets/intrinsics/vorrq_u64 +// https://developer.arm.com/architectures/instruction-sets/intrinsics/vreinterpretq_u64_u32 +// https://developer.arm.com/architectures/instruction-sets/intrinsics/vmaxvq_u32 +// This would yield whether the pointer was found, but instead of maxing at +// the end we can use the produced mask to extract out a single register with +// which positions are hits and use that to then permute the registers into +// the proper order. At the end we could use a table instruction to remap and +// extract out a byte/bitmap of the indices that we need to insert into the +// main set. +// +// The shifting can be performed with +// https://developer.arm.com/architectures/instruction-sets/intrinsics/vextq_u32 +// https://developer.arm.com/architectures/instruction-sets/intrinsics/vextq_u64 +// This takes n low elements of LHS and rest from RHS and we can cascade them +// to shift down the whole MRU. +// +// We can use SIMDE as a rosetta stone for getting neon/avx/wasm/etc: +// https://github.com/simd-everywhere/simde/blob/master/simde/arm/neon/ceq.h#L591 +static iree_status_t iree_hal_resource_set_insert_1( + iree_hal_resource_set_t* set, iree_hal_resource_t* resource) { + // Scan and hope for a hit. + for (iree_host_size_t i = 0; i < IREE_ARRAYSIZE(set->mru); ++i) { + if (set->mru[i] != resource) continue; + // Hit - keep the list sorted by most->least recently used. + // We shift the MRU down to make room at index 0 and store the + // resource there. + if (i > 0) { + memmove(&set->mru[1], &set->mru[0], sizeof(set->mru[0]) * i); + set->mru[0] = resource; + } + return iree_ok_status(); + } + + // Miss - insert into the main list (slow path). + // Note that we do this before updating the MRU in case allocation fails - we + // don't want to keep the pointer around unless we've really retained it. + IREE_RETURN_IF_ERROR(iree_hal_resource_set_insert_retain(set, resource)); + + // Shift the MRU down and insert the new item at the head. + memmove(&set->mru[1], &set->mru[0], + sizeof(set->mru[0]) * (IREE_ARRAYSIZE(set->mru) - 1)); + set->mru[0] = resource; + + return iree_ok_status(); +} + +IREE_API_EXPORT iree_status_t +iree_hal_resource_set_insert(iree_hal_resource_set_t* set, + iree_host_size_t count, const void* resources) { + // For now we process one at a time. We should have a stride that lets us + // amortize the cost of doing the MRU update and insertion allocation by + // say slicing off 4/8/16/32 resources at a time etc. Today each miss that + // requires a full insertion goes down the whole path of checking chunk + // capacity and such. + iree_hal_resource_t* const* typed_resources = + (iree_hal_resource_t* const*)resources; + for (iree_host_size_t i = 0; i < count; ++i) { + IREE_RETURN_IF_ERROR( + iree_hal_resource_set_insert_1(set, typed_resources[i])); + } + return iree_ok_status(); +}
diff --git a/iree/hal/utils/resource_set.h b/iree/hal/utils/resource_set.h new file mode 100644 index 0000000..6f63ced --- /dev/null +++ b/iree/hal/utils/resource_set.h
@@ -0,0 +1,139 @@ +// Copyright 2022 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 + +#ifndef IREE_HAL_UTILS_RESOURCE_SET_H_ +#define IREE_HAL_UTILS_RESOURCE_SET_H_ + +#include "iree/base/api.h" +#include "iree/base/internal/arena.h" +#include "iree/hal/resource.h" + +#ifdef __cplusplus +extern "C" { +#endif // __cplusplus + +// Bit 0 of the next_chunk pointer indicates whether we are inlined into the +// resource set block - the chunks are always aligned and the bit is unused. +#define IREE_HAL_RESOURCE_SET_CHUNK_FLAG_INLINE 0x1 + +// Capacity is limited by how many bits we reserve for the count. +#define IREE_HAL_RESOURCE_SET_CHUNK_MAX_CAPACITY 0xFFFFu + +// A chunk of resources within a resource set. +// Chunks contain a fixed number of resources based on the block size of the +// pool the set was allocated from. +typedef struct iree_hal_resource_set_chunk_t { + // Next chunk in the chunk linked list. + // Bit 0 indicates whether this was an allocated block; 0 means that the + // chunk is stored within the parent resource set and should not be returned + // to the block pool. This works only because we know the blocks are allocated + // at an alignment >= 16 and we have a few bits to work with. + union { + struct iree_hal_resource_set_chunk_t* next_chunk; + uintptr_t flags; + }; + + // Retained resources - may be less than the capacity derived from the block + // pool block size. We keep the counts small here to reduce chunk overhead. We + // could recompute the capacity each time but at the point that we use even 1 + // byte we've already consumed 4 (or 8) thanks to padding and should make use + // of the rest. + uint16_t capacity; + uint16_t count; + iree_hal_resource_t* resources[]; +} iree_hal_resource_set_chunk_t; + +// Returns true if the chunk is stored inline in the parent resource set. +#define iree_hal_resource_set_chunk_is_stored_inline(chunk) \ + (((chunk)->flags & IREE_HAL_RESOURCE_SET_CHUNK_FLAG_INLINE) == \ + IREE_HAL_RESOURCE_SET_CHUNK_FLAG_INLINE) + +// Number of elements in the most-recently-used resource list of a set. +// The larger the number the greater the chance of having a hit but the more +// expensive every miss will be. +// +// To try to keep the MRU in cache we size this based on how many pointers will +// fit in a single cache line. This also makes it easier to author SIMD lookups +// as we'll (in-theory) be able to load the entries into SIMD registers. +// +// Values for the platforms we specify for: +// 32-bit: 64 / 4 = 16x4b ptrs (4 x uint32x4_t) +// 64-bit: 64 / 8 = 8x8b ptrs (4 x uint64x2_t) +// We could scale this up if we wanted but being able to unroll is nice. +#define IREE_HAL_RESOURCE_SET_MRU_SIZE \ + (iree_hardware_constructive_interference_size / sizeof(uintptr_t)) + +// "Efficient" append-only set for retaining a set of resources. +// This is a non-deterministic data structure that tries to reduce the amount of +// overhead involved in tracking a reasonably-sized set of resources (~dozens to +// hundreds). Set insertion may have false negatives and retain resources more +// than strictly required by trading off the expense of precisely detecting +// redundant insertions with the expense of an additional atomic operation. +// +// This tries to elide insertions by maintaining a most-recently-used list. +// This optimizes for temporal locality of resources used (the same executables, +// same buffers, etc) and is implemented to have a fixed cost regardless of +// whether the values are found and should hopefully trigger enough to avoid the +// subsequent full insertion that can introduce allocations and ref counting. +// The idea is that if we can keep the MRU in cache and spend a dozen cycles to +// manage it we only need to avoid a single cache miss that would occur doing +// the full insertion. We care here because this is on the critical path of +// command encoding and the parasitic cost of maintaining the set scales with +// the number of commands issued. This never needs to be free, only as fast as +// whatever user code may need to do to maintain proper lifetime - or as small +// in terms of code-size. +// +// **WARNING**: thread-unsafe insertion: it's assumed that sets are constructed +// by a single thread, sealed, and then released at once at a future time point. +// Multiple threads needing to insert into a set should have their own sets and +// then join them afterward. +typedef struct iree_hal_resource_set_t { + // A small MRUish list of resources for quickly deduplicating insertions. + // We use this to perform an O(k) comparison traded off with the cost of a + // miss that results in an atomic inc/dec. We shouldn't make this + // more expensive than the additional cost of the retain/release. + // + // This lives at the head of the struct as it's used in 100% of insertions and + // if we can get lucky with it staying in cache we reduce a lot of memory + // traffic. Once we spill the MRU and go to main memory to add the resource + // we're going to have a cache miss and this way we avoid two (one for the + // set and one for the chunk). + // + // TODO(benvanik): ensure alignment on the set - should be at + // iree_hardware_constructive_interference_size. + iree_hal_resource_t* mru[IREE_HAL_RESOURCE_SET_MRU_SIZE]; + + // Block pool used for allocating additional set storage slabs. + iree_arena_block_pool_t* block_pool; + + // Linked list of storage chunks. + iree_hal_resource_set_chunk_t* chunk_head; +} iree_hal_resource_set_t; + +// Allocates a new resource from the given |block_pool|. +// Resources can be inserted and are retained until the set is freed. +IREE_API_EXPORT iree_status_t iree_hal_resource_set_allocate( + iree_arena_block_pool_t* block_pool, iree_hal_resource_set_t** out_set); + +// Frees a resource set and releases all inserted resources. +// The |set| itself will be returned back to the block pool it was allocated +// from. +IREE_API_EXPORT void iree_hal_resource_set_free(iree_hal_resource_set_t* set); + +// Resets the set to its initial empty state by releasing all owned resources. +IREE_API_EXPORT void iree_hal_resource_set_reset(iree_hal_resource_set_t* set); + +// Inserts zero or more resources into the set. +// Each resource will be retained for at least the lifetime of the set. +IREE_API_EXPORT iree_status_t +iree_hal_resource_set_insert(iree_hal_resource_set_t* set, + iree_host_size_t count, const void* resources); + +#ifdef __cplusplus +} // extern "C" +#endif // __cplusplus + +#endif // IREE_HAL_UTILS_RESOURCE_SET_H_
diff --git a/iree/hal/utils/resource_set_benchmark.c b/iree/hal/utils/resource_set_benchmark.c new file mode 100644 index 0000000..5b22f97 --- /dev/null +++ b/iree/hal/utils/resource_set_benchmark.c
@@ -0,0 +1,287 @@ +// Copyright 2022 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 <stdint.h> +#include <stdio.h> +#include <stdlib.h> +#include <string.h> + +#include "iree/base/api.h" +#include "iree/base/internal/prng.h" +#include "iree/hal/api.h" +#include "iree/hal/utils/resource_set.h" +#include "iree/testing/benchmark.h" + +typedef struct iree_hal_test_resource_t { + iree_hal_resource_t resource; + iree_allocator_t host_allocator; +} iree_hal_test_resource_t; + +typedef struct iree_hal_test_resource_vtable_t { + void(IREE_API_PTR* destroy)(iree_hal_test_resource_t* resource); +} iree_hal_test_resource_vtable_t; +IREE_HAL_ASSERT_VTABLE_LAYOUT(iree_hal_test_resource_vtable_t); + +static const iree_hal_test_resource_vtable_t iree_hal_test_resource_vtable; + +static iree_status_t iree_hal_test_resource_create( + iree_allocator_t host_allocator, iree_hal_resource_t** out_resource) { + iree_hal_test_resource_t* test_resource = NULL; + IREE_RETURN_IF_ERROR(iree_allocator_malloc( + host_allocator, sizeof(*test_resource), (void**)&test_resource)); + iree_hal_resource_initialize(&iree_hal_test_resource_vtable, + &test_resource->resource); + test_resource->host_allocator = host_allocator; + *out_resource = (iree_hal_resource_t*)test_resource; + return iree_ok_status(); +} + +static void iree_hal_test_resource_destroy(iree_hal_test_resource_t* resource) { + iree_allocator_t host_allocator = resource->host_allocator; + iree_allocator_free(host_allocator, resource); +} + +static const iree_hal_test_resource_vtable_t iree_hal_test_resource_vtable = { + /*.destroy=*/iree_hal_test_resource_destroy, +}; + +// Tests init/deinit performance when 0+ resources are in the set. +// This is our worst-case with unique resources that never match the MRU. +// +// user_data is a count of elements to insert into each set. +static iree_status_t iree_hal_resource_set_benchmark_lifecycle_n( + const iree_benchmark_def_t* benchmark_def, + iree_benchmark_state_t* benchmark_state) { + iree_allocator_t host_allocator = benchmark_state->host_allocator; + + // Initialize the block pool we'll be serving from. + // Sized like we usually do it in the runtime for ~512-1024 elements. + iree_arena_block_pool_t block_pool; + iree_arena_block_pool_initialize(4096, host_allocator, &block_pool); + + // Allocate the resources we'll be using - we keep them live so that we are + // measuring just the retain/release and set times instead of the timing of + // resource creation/deletion. + uint32_t count = (uint32_t)(uintptr_t)benchmark_def->user_data; + iree_hal_resource_t** resources = NULL; + if (count > 0) { + IREE_CHECK_OK(iree_allocator_malloc(host_allocator, + sizeof(iree_hal_resource_t*) * count, + (void**)&resources)); + } + for (uint32_t i = 0; i < count; ++i) { + IREE_CHECK_OK(iree_hal_test_resource_create(host_allocator, &resources[i])); + } + + // Create/insert/delete lifecycle. + while (iree_benchmark_keep_running(benchmark_state, /*batch_count=*/1)) { + iree_hal_resource_set_t* set = NULL; + IREE_CHECK_OK(iree_hal_resource_set_allocate(&block_pool, &set)); + IREE_CHECK_OK(iree_hal_resource_set_insert(set, count, resources)); + iree_hal_resource_set_free(set); + } + + // Cleanup. + for (uint32_t i = 0; i < count; ++i) { + iree_hal_resource_release(resources[i]); + } + iree_allocator_free(host_allocator, resources); + iree_arena_block_pool_deinitialize(&block_pool); + + return iree_ok_status(); +} + +// Tests insertion performance when either the MRU is used (n < MRU size) or +// the worst-case performance when all resources are unique and guaranteed to +// miss the MRU. Expect to see a cliff where we spill the MRU. +// +// user_data is a count of unique elements to insert. +static iree_status_t iree_hal_resource_set_benchmark_insert_n( + const iree_benchmark_def_t* benchmark_def, + iree_benchmark_state_t* benchmark_state) { + iree_allocator_t host_allocator = benchmark_state->host_allocator; + + // Initialize the block pool we'll be serving from. + // Sized like we usually do it in the runtime for ~512-1024 elements. + iree_arena_block_pool_t block_pool; + iree_arena_block_pool_initialize(4096, host_allocator, &block_pool); + + // Create the empty set using the block pool for additional memory. + iree_hal_resource_set_t* set = NULL; + IREE_CHECK_OK(iree_hal_resource_set_allocate(&block_pool, &set)); + + // Allocate the resources we'll be using - we keep them live so that we are + // measuring just the retain/release and set times instead of the timing of + // resource creation/deletion. + uint32_t count = (uint32_t)(uintptr_t)benchmark_def->user_data; + iree_hal_resource_t** resources = NULL; + IREE_CHECK_OK(iree_allocator_malloc(host_allocator, + sizeof(iree_hal_resource_t*) * count, + (void**)&resources)); + for (uint32_t i = 0; i < count; ++i) { + IREE_CHECK_OK(iree_hal_test_resource_create(host_allocator, &resources[i])); + } + + // Insert the resources. After the first iteration these should all be hits. + while (iree_benchmark_keep_running(benchmark_state, /*batch_count=*/1)) { + IREE_CHECK_OK(iree_hal_resource_set_insert(set, count, resources)); + } + + // Cleanup. + for (uint32_t i = 0; i < count; ++i) { + iree_hal_resource_release(resources[i]); + } + iree_hal_resource_set_free(set); + iree_allocator_free(host_allocator, resources); + iree_arena_block_pool_deinitialize(&block_pool); + + return iree_ok_status(); +} + +// Tests insertion into the set in a randomized order. +// This lets us get a somewhat reasonable approximation of average performance. +// In reality what the compiler spits out is non-random and often just +// alternating A/B/C/B/A/C/A/B/C etc kind of sequences. +// +// This is the most important benchmark: if this is fast then we are :thumbsup:. +// +// user_data is a count of unique element pool to insert N times. The higher +// the pool size the more likely we are to miss the MRU. +static iree_status_t iree_hal_resource_set_benchmark_randomized_n( + const iree_benchmark_def_t* benchmark_def, + iree_benchmark_state_t* benchmark_state) { + iree_allocator_t host_allocator = benchmark_state->host_allocator; + + // Initialize the block pool we'll be serving from. + // Sized like we usually do it in the runtime for ~512-1024 elements. + iree_arena_block_pool_t block_pool; + iree_arena_block_pool_initialize(4096, host_allocator, &block_pool); + + // Allocate the resources we'll be using - we keep them live so that we are + // measuring just the retain/release and set times instead of the timing of + // resource creation/deletion. + uint32_t count = (uint32_t)(uintptr_t)benchmark_def->user_data; + iree_hal_resource_t** resources = NULL; + IREE_CHECK_OK(iree_allocator_malloc(host_allocator, + sizeof(iree_hal_resource_t*) * count, + (void**)&resources)); + for (uint32_t i = 0; i < count; ++i) { + IREE_CHECK_OK(iree_hal_test_resource_create(host_allocator, &resources[i])); + } + + // The same set is maintained; we'll eventually have all resources in the set + // and be testing the MRU hit %. + iree_hal_resource_set_t* set = NULL; + IREE_CHECK_OK(iree_hal_resource_set_allocate(&block_pool, &set)); + + // The PRNG we use to select the elements. + iree_prng_xoroshiro128_state_t prng = {0}; + iree_prng_xoroshiro128_initialize(123ull, &prng); + + // Insert N random resources into the set. To hide some of the overhead we do + // multiple insertions in each loop. + while (iree_benchmark_keep_running(benchmark_state, /*batch_count=*/256)) { + for (uint32_t i = 0; i < 256; ++i) { + uint32_t resource_idx = + iree_prng_xoroshiro128plus_next_uint32(&prng) % count; + iree_hal_resource_t* resource = resources[resource_idx]; + IREE_CHECK_OK(iree_hal_resource_set_insert(set, 1, &resource)); + } + } + + // Cleanup. + iree_hal_resource_set_free(set); + for (uint32_t i = 0; i < count; ++i) { + iree_hal_resource_release(resources[i]); + } + iree_allocator_free(host_allocator, resources); + iree_arena_block_pool_deinitialize(&block_pool); + + return iree_ok_status(); +} + +int main(int argc, char** argv) { + iree_benchmark_initialize(&argc, argv); + + // iree_hal_resource_set_benchmark_lifecycle_n + { + iree_benchmark_def_t benchmark_def = { + .flags = IREE_BENCHMARK_FLAG_MEASURE_PROCESS_CPU_TIME | + IREE_BENCHMARK_FLAG_USE_REAL_TIME, + .time_unit = IREE_BENCHMARK_UNIT_NANOSECOND, + .minimum_duration_ns = 0, + .iteration_count = 0, + .run = iree_hal_resource_set_benchmark_lifecycle_n, + }; + benchmark_def.user_data = (void*)0u; + iree_benchmark_register(iree_make_cstring_view("lifecycle_0"), + &benchmark_def); + benchmark_def.user_data = (void*)1u; + iree_benchmark_register(iree_make_cstring_view("lifecycle_1"), + &benchmark_def); + benchmark_def.user_data = (void*)256u; + iree_benchmark_register(iree_make_cstring_view("lifecycle_256"), + &benchmark_def); + benchmark_def.user_data = (void*)1024u; + iree_benchmark_register(iree_make_cstring_view("lifecycle_1024"), + &benchmark_def); + } + + // iree_hal_resource_set_benchmark_insert_n + { + iree_benchmark_def_t benchmark_def = { + .flags = IREE_BENCHMARK_FLAG_MEASURE_PROCESS_CPU_TIME | + IREE_BENCHMARK_FLAG_USE_REAL_TIME, + .time_unit = IREE_BENCHMARK_UNIT_NANOSECOND, + .minimum_duration_ns = 0, + .iteration_count = 0, + .run = iree_hal_resource_set_benchmark_insert_n, + }; + benchmark_def.user_data = (void*)1u; + iree_benchmark_register(iree_make_cstring_view("insert_1"), &benchmark_def); + benchmark_def.user_data = (void*)5u; + iree_benchmark_register(iree_make_cstring_view("insert_5"), &benchmark_def); + benchmark_def.user_data = (void*)32u; + iree_benchmark_register(iree_make_cstring_view("insert_32"), + &benchmark_def); + benchmark_def.user_data = (void*)64u; + iree_benchmark_register(iree_make_cstring_view("insert_64"), + &benchmark_def); + } + + // iree_hal_resource_set_benchmark_randomized_n + { + iree_benchmark_def_t benchmark_def = { + .flags = IREE_BENCHMARK_FLAG_MEASURE_PROCESS_CPU_TIME | + IREE_BENCHMARK_FLAG_USE_REAL_TIME, + .time_unit = IREE_BENCHMARK_UNIT_NANOSECOND, + .minimum_duration_ns = 0, + .iteration_count = 0, + .run = iree_hal_resource_set_benchmark_randomized_n, + }; + benchmark_def.user_data = (void*)1u; + iree_benchmark_register(iree_make_cstring_view("randomized_1"), + &benchmark_def); + benchmark_def.user_data = (void*)4u; + iree_benchmark_register(iree_make_cstring_view("randomized_4"), + &benchmark_def); + benchmark_def.user_data = (void*)8u; + iree_benchmark_register(iree_make_cstring_view("randomized_8"), + &benchmark_def); + benchmark_def.user_data = (void*)32u; + iree_benchmark_register(iree_make_cstring_view("randomized_32"), + &benchmark_def); + benchmark_def.user_data = (void*)256u; + iree_benchmark_register(iree_make_cstring_view("randomized_256"), + &benchmark_def); + benchmark_def.user_data = (void*)4096u; + iree_benchmark_register(iree_make_cstring_view("randomized_4096"), + &benchmark_def); + } + + iree_benchmark_run_specified(); + return 0; +}
diff --git a/iree/hal/utils/resource_set_test.cc b/iree/hal/utils/resource_set_test.cc new file mode 100644 index 0000000..021bb1b --- /dev/null +++ b/iree/hal/utils/resource_set_test.cc
@@ -0,0 +1,257 @@ +// Copyright 2022 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/resource_set.h" + +#include <cstddef> +#include <cstdint> +#include <memory> +#include <string> + +#include "iree/base/api.h" +#include "iree/hal/api.h" +#include "iree/testing/gtest.h" +#include "iree/testing/status_matchers.h" + +namespace iree { +namespace hal { +namespace { + +using ::iree::testing::status::IsOkAndHolds; +using ::iree::testing::status::StatusIs; +using ::testing::Eq; + +typedef struct iree_hal_test_resource_t { + iree_hal_resource_t resource; + iree_allocator_t host_allocator; + uint32_t index; + uint32_t* live_bitmap; +} iree_hal_test_resource_t; + +typedef struct iree_hal_test_resource_vtable_t { + void(IREE_API_PTR* destroy)(iree_hal_test_resource_t* resource); +} iree_hal_test_resource_vtable_t; +IREE_HAL_ASSERT_VTABLE_LAYOUT(iree_hal_test_resource_vtable_t); + +extern const iree_hal_test_resource_vtable_t iree_hal_test_resource_vtable; + +static iree_status_t iree_hal_test_resource_create( + uint32_t index, uint32_t* live_bitmap, iree_allocator_t host_allocator, + iree_hal_resource_t** out_resource) { + iree_hal_test_resource_t* test_resource = NULL; + IREE_RETURN_IF_ERROR(iree_allocator_malloc( + host_allocator, sizeof(*test_resource), (void**)&test_resource)); + iree_hal_resource_initialize(&iree_hal_test_resource_vtable, + &test_resource->resource); + test_resource->host_allocator = host_allocator; + test_resource->index = index; + test_resource->live_bitmap = live_bitmap; + *live_bitmap |= 1 << index; + *out_resource = (iree_hal_resource_t*)test_resource; + return iree_ok_status(); +} + +static void iree_hal_test_resource_destroy(iree_hal_test_resource_t* resource) { + iree_allocator_t host_allocator = resource->host_allocator; + *resource->live_bitmap &= ~(1 << resource->index); + iree_allocator_free(host_allocator, resource); +} + +const iree_hal_test_resource_vtable_t iree_hal_test_resource_vtable = { + /*.destroy=*/iree_hal_test_resource_destroy, +}; + +struct ResourceSetTest : public ::testing::Test { + // We could check the allocator to ensure all memory is freed if we wanted to + // reduce the reliance on asan. + iree_allocator_t host_allocator = iree_allocator_system(); + iree_arena_block_pool_t block_pool; + + void SetUp() override { + memset(&block_pool, 0, sizeof(block_pool)); + iree_arena_block_pool_initialize(128, host_allocator, &block_pool); + } + + void TearDown() override { + // This may assert (or at least trigger asan) if there are blocks + // outstanding. + iree_arena_block_pool_deinitialize(&block_pool); + } +}; + +using resource_set_ptr = std::unique_ptr<iree_hal_resource_set_t, + decltype(&iree_hal_resource_set_free)>; +static resource_set_ptr make_resource_set(iree_arena_block_pool_t* block_pool) { + iree_hal_resource_set_t* set = NULL; + IREE_CHECK_OK(iree_hal_resource_set_allocate(block_pool, &set)); + return resource_set_ptr(set, iree_hal_resource_set_free); +} + +// Tests a set that has no resources added to it. +TEST_F(ResourceSetTest, Empty) { + iree_hal_resource_set_t* set = NULL; + IREE_ASSERT_OK(iree_hal_resource_set_allocate(&block_pool, &set)); + iree_hal_resource_set_free(set); +} + +// Tests insertion of a single resource. +TEST_F(ResourceSetTest, Insert1) { + auto resource_set = make_resource_set(&block_pool); + + // Create test resource; it'll set its bit in the live_bitmap. + iree_hal_resource_t* resource = NULL; + uint32_t live_bitmap = 0u; + IREE_ASSERT_OK(iree_hal_test_resource_create(0, &live_bitmap, host_allocator, + &resource)); + EXPECT_EQ(live_bitmap, 1u); + + // Insert the resource and drop the reference; it should still be live as the + // set retains it. + IREE_ASSERT_OK( + iree_hal_resource_set_insert(resource_set.get(), 1, &resource)); + iree_hal_resource_release(resource); + EXPECT_EQ(live_bitmap, 1u); + + // Drop the set and expect the resource to be destroyed as it loses its last + // reference. + resource_set.reset(); + EXPECT_EQ(live_bitmap, 0u); +} + +// Tests inserting multiple resources at a time. +TEST_F(ResourceSetTest, Insert5) { + auto resource_set = make_resource_set(&block_pool); + + // Allocate 5 resources - this lets us test for special paths that may handle + // 4 at a time (to fit in SIMD registers) as well as the leftovers. + iree_hal_resource_t* resources[5] = {NULL}; + uint32_t live_bitmap = 0u; + for (iree_host_size_t i = 0; i < IREE_ARRAYSIZE(resources); ++i) { + IREE_ASSERT_OK(iree_hal_test_resource_create( + i, &live_bitmap, host_allocator, &resources[i])); + } + EXPECT_EQ(live_bitmap, 0x1Fu); + + // Transfer ownership of the resources to the set. + IREE_ASSERT_OK(iree_hal_resource_set_insert( + resource_set.get(), IREE_ARRAYSIZE(resources), resources)); + for (iree_host_size_t i = 0; i < IREE_ARRAYSIZE(resources); ++i) { + iree_hal_resource_release(resources[i]); + } + EXPECT_EQ(live_bitmap, 0x1Fu); + + // Ensure the set releases the resources. + resource_set.reset(); + EXPECT_EQ(live_bitmap, 0u); +} + +// Tests inserting enough resources to force set growth. This is ensured by +// choosing a sufficiently small block size such that even 32 elements triggers +// a growth. Of course, real usage should have at least ~4KB for the block size. +TEST_F(ResourceSetTest, InsertionGrowth) { + auto resource_set = make_resource_set(&block_pool); + + // Allocate 32 resources (one for each bit in our live map). + iree_hal_resource_t* resources[32] = {NULL}; + uint32_t live_bitmap = 0u; + for (iree_host_size_t i = 0; i < IREE_ARRAYSIZE(resources); ++i) { + IREE_ASSERT_OK(iree_hal_test_resource_create( + i, &live_bitmap, host_allocator, &resources[i])); + } + EXPECT_EQ(live_bitmap, 0xFFFFFFFFu); + + // Transfer ownership of the resources to the set. + IREE_ASSERT_OK(iree_hal_resource_set_insert( + resource_set.get(), IREE_ARRAYSIZE(resources), resources)); + for (iree_host_size_t i = 0; i < IREE_ARRAYSIZE(resources); ++i) { + iree_hal_resource_release(resources[i]); + } + EXPECT_EQ(live_bitmap, 0xFFFFFFFFu); + + // Ensure the set releases the resources. + resource_set.reset(); + EXPECT_EQ(live_bitmap, 0u); +} + +// Tests insertion of resources multiple times to verify the MRU works. +TEST_F(ResourceSetTest, RedundantInsertion) { + auto resource_set = make_resource_set(&block_pool); + + // Allocate 32 resources (one for each bit in our live map). + // We want to be able to miss in the MRU. + iree_hal_resource_t* resources[32] = {NULL}; + static_assert(IREE_ARRAYSIZE(resources) > IREE_HAL_RESOURCE_SET_MRU_SIZE, + "need to pick a value that lets us exceed the MRU capacity"); + uint32_t live_bitmap = 0u; + for (iree_host_size_t i = 0; i < IREE_ARRAYSIZE(resources); ++i) { + IREE_ASSERT_OK(iree_hal_test_resource_create( + i, &live_bitmap, host_allocator, &resources[i])); + } + EXPECT_EQ(live_bitmap, 0xFFFFFFFFu); + + // NOTE: the only requirement of the MRU is that it's _mostly_ MRU - we may + // for performance reasons make it a little fuzzy to avoid additional + // shuffling. Today it's always a proper MRU and we check the pointers here. + + // NOTE: the MRU size can vary across architectures; we know it should always + // be at least ~6 though so that's what we work with here. + static_assert(IREE_HAL_RESOURCE_SET_MRU_SIZE > 6, + "need at least enough elements to test with"); + + // Insert in sequence, MRU should contain: + // 31 30 29 28 27 ... + IREE_ASSERT_OK(iree_hal_resource_set_insert( + resource_set.get(), IREE_ARRAYSIZE(resources), resources)); + EXPECT_EQ(resource_set->mru[0], resources[31]); + EXPECT_EQ(resource_set->mru[1], resources[30]); + EXPECT_EQ(resource_set->mru[2], resources[29]); + EXPECT_EQ(resource_set->mru[3], resources[28]); + EXPECT_EQ(resource_set->mru[4], resources[27]); + + // Insert 31 again, MRU should remain the same as it's at the head. + IREE_ASSERT_OK( + iree_hal_resource_set_insert(resource_set.get(), 1, &resources[31])); + EXPECT_EQ(resource_set->mru[0], resources[31]); + EXPECT_EQ(resource_set->mru[1], resources[30]); + EXPECT_EQ(resource_set->mru[2], resources[29]); + EXPECT_EQ(resource_set->mru[3], resources[28]); + EXPECT_EQ(resource_set->mru[4], resources[27]); + + // Insert 28 again, MRU should be updated to move it to the front: + // 28 31 30 29 27 ... + IREE_ASSERT_OK( + iree_hal_resource_set_insert(resource_set.get(), 1, &resources[28])); + EXPECT_EQ(resource_set->mru[0], resources[28]); + EXPECT_EQ(resource_set->mru[1], resources[31]); + EXPECT_EQ(resource_set->mru[2], resources[30]); + EXPECT_EQ(resource_set->mru[3], resources[29]); + EXPECT_EQ(resource_set->mru[4], resources[27]); + + // Insert 0 again, which should be a miss as it fell off the end of the MRU: + // 0 28 31 30 29 27 ... + IREE_ASSERT_OK( + iree_hal_resource_set_insert(resource_set.get(), 1, &resources[0])); + EXPECT_EQ(resource_set->mru[0], resources[0]); + EXPECT_EQ(resource_set->mru[1], resources[28]); + EXPECT_EQ(resource_set->mru[2], resources[31]); + EXPECT_EQ(resource_set->mru[3], resources[30]); + EXPECT_EQ(resource_set->mru[4], resources[29]); + EXPECT_EQ(resource_set->mru[5], resources[27]); + + // Release all of the resources - they should still be owned by the set. + for (iree_host_size_t i = 0; i < IREE_ARRAYSIZE(resources); ++i) { + iree_hal_resource_release(resources[i]); + } + EXPECT_EQ(live_bitmap, 0xFFFFFFFFu); + + // Ensure the set releases the resources. + resource_set.reset(); + EXPECT_EQ(live_bitmap, 0u); +} + +} // namespace +} // namespace hal +} // namespace iree
diff --git a/iree/hal/vulkan/BUILD b/iree/hal/vulkan/BUILD index d531e68..40b56c1 100644 --- a/iree/hal/vulkan/BUILD +++ b/iree/hal/vulkan/BUILD
@@ -91,9 +91,12 @@ "//iree/base:logging", "//iree/base:tracing", "//iree/base/internal", + "//iree/base/internal:arena", "//iree/base/internal:synchronization", "//iree/base/internal/flatcc:parsing", "//iree/hal", + "//iree/hal/utils:buffer_transfer", + "//iree/hal/utils:resource_set", "//iree/hal/vulkan/builtin", "//iree/hal/vulkan/util:arena", "//iree/hal/vulkan/util:intrusive_list",
diff --git a/iree/hal/vulkan/CMakeLists.txt b/iree/hal/vulkan/CMakeLists.txt index caee945..37fe297 100644 --- a/iree/hal/vulkan/CMakeLists.txt +++ b/iree/hal/vulkan/CMakeLists.txt
@@ -79,11 +79,14 @@ iree::base::cc iree::base::core_headers iree::base::internal + iree::base::internal::arena iree::base::internal::flatcc::parsing iree::base::internal::synchronization iree::base::logging iree::base::tracing iree::hal + iree::hal::utils::buffer_transfer + iree::hal::utils::resource_set iree::hal::vulkan::builtin iree::hal::vulkan::util::arena iree::hal::vulkan::util::intrusive_list
diff --git a/iree/hal/vulkan/direct_command_buffer.cc b/iree/hal/vulkan/direct_command_buffer.cc index 9e6d799..015250c 100644 --- a/iree/hal/vulkan/direct_command_buffer.cc +++ b/iree/hal/vulkan/direct_command_buffer.cc
@@ -15,6 +15,7 @@ #include "iree/base/logging.h" #include "iree/base/status_cc.h" #include "iree/base/tracing.h" +#include "iree/hal/utils/resource_set.h" #include "iree/hal/vulkan/descriptor_set_arena.h" #include "iree/hal/vulkan/dynamic_symbols.h" #include "iree/hal/vulkan/native_descriptor_set.h" @@ -34,12 +35,17 @@ iree_hal_command_buffer_t base; VkDeviceHandle* logical_device; iree_hal_vulkan_tracing_context_t* tracing_context; + iree_arena_block_pool_t* block_pool; VkCommandPoolHandle* command_pool; VkCommandBuffer handle; DynamicSymbols* syms; + // Maintains a reference to all resources used within the command buffer. + // Reset on each begin. + iree_hal_resource_set_t* resource_set; + // TODO(benvanik): may grow large - should try to reclaim or reuse. DescriptorSetArena descriptor_set_arena; @@ -81,10 +87,12 @@ iree_hal_vulkan_tracing_context_t* tracing_context, iree::hal::vulkan::DescriptorPoolCache* descriptor_pool_cache, iree::hal::vulkan::BuiltinExecutables* builtin_executables, + iree_arena_block_pool_t* block_pool, iree_hal_command_buffer_t** out_command_buffer) { IREE_ASSERT_ARGUMENT(logical_device); IREE_ASSERT_ARGUMENT(command_pool); IREE_ASSERT_ARGUMENT(descriptor_pool_cache); + IREE_ASSERT_ARGUMENT(block_pool); IREE_ASSERT_ARGUMENT(out_command_buffer); IREE_TRACE_ZONE_BEGIN(z0); @@ -109,6 +117,7 @@ &iree_hal_vulkan_direct_command_buffer_vtable, &command_buffer->base); command_buffer->logical_device = logical_device; command_buffer->tracing_context = tracing_context; + command_buffer->block_pool = block_pool; command_buffer->command_pool = command_pool; command_buffer->handle = handle; command_buffer->syms = logical_device->syms().get(); @@ -118,7 +127,11 @@ new (&command_buffer->descriptor_set_group) DescriptorSetGroup(); command_buffer->builtin_executables = builtin_executables; + status = iree_hal_resource_set_allocate(block_pool, + &command_buffer->resource_set); + } + if (iree_status_is_ok(status)) { *out_command_buffer = &command_buffer->base; } else { command_pool->Free(handle); @@ -133,6 +146,7 @@ // NOTE: we require that command buffers not be recorded while they are // in-flight so this is safe. IREE_IGNORE_ERROR(command_buffer->descriptor_set_group.Reset()); + iree_hal_resource_set_reset(command_buffer->resource_set); } bool iree_hal_vulkan_direct_command_buffer_isa( @@ -164,6 +178,7 @@ command_buffer->descriptor_set_group.~DescriptorSetGroup(); command_buffer->descriptor_set_arena.~DescriptorSetArena(); + iree_hal_resource_set_reset(command_buffer->resource_set); iree_allocator_free(host_allocator, command_buffer); IREE_TRACE_ZONE_END(z0); @@ -399,6 +414,9 @@ iree_hal_vulkan_direct_command_buffer_t* command_buffer = iree_hal_vulkan_direct_command_buffer_cast(base_command_buffer); + IREE_RETURN_IF_ERROR( + iree_hal_resource_set_insert(command_buffer->resource_set, 1, &event)); + command_buffer->syms->vkCmdSetEvent( command_buffer->handle, iree_hal_vulkan_native_event_handle(event), iree_hal_vulkan_convert_pipeline_stage_flags(source_stage_mask)); @@ -412,6 +430,9 @@ iree_hal_vulkan_direct_command_buffer_t* command_buffer = iree_hal_vulkan_direct_command_buffer_cast(base_command_buffer); + IREE_RETURN_IF_ERROR( + iree_hal_resource_set_insert(command_buffer->resource_set, 1, &event)); + command_buffer->syms->vkCmdResetEvent( command_buffer->handle, iree_hal_vulkan_native_event_handle(event), iree_hal_vulkan_convert_pipeline_stage_flags(source_stage_mask)); @@ -433,6 +454,9 @@ iree_allocator_t host_allocator = command_buffer->logical_device->host_allocator(); + IREE_RETURN_IF_ERROR(iree_hal_resource_set_insert( + command_buffer->resource_set, event_count, events)); + iree_inline_array(VkEvent, event_handles, event_count, host_allocator); for (int i = 0; i < event_count; ++i) { *iree_inline_array_at(event_handles, i) = @@ -526,6 +550,9 @@ VkBuffer target_device_buffer = iree_hal_vulkan_vma_buffer_handle( iree_hal_buffer_allocated_buffer(target_buffer)); + IREE_RETURN_IF_ERROR(iree_hal_resource_set_insert( + command_buffer->resource_set, 1, &target_buffer)); + // vkCmdFillBuffer requires a 4 byte alignment for the offset, pattern, and // length. We use a polyfill here that fills the unaligned start and end of // fill operations, if needed. @@ -580,6 +607,9 @@ VkBuffer target_device_buffer = iree_hal_vulkan_vma_buffer_handle( iree_hal_buffer_allocated_buffer(target_buffer)); + IREE_RETURN_IF_ERROR(iree_hal_resource_set_insert( + command_buffer->resource_set, 1, &target_buffer)); + // Vulkan only allows updates of <= 65536 because you really, really, really // shouldn't do large updates like this (as it wastes command buffer space and // may be slower than just using write-through mapped memory). The @@ -613,6 +643,10 @@ VkBuffer target_device_buffer = iree_hal_vulkan_vma_buffer_handle( iree_hal_buffer_allocated_buffer(target_buffer)); + const iree_hal_buffer_t* buffers[2] = {source_buffer, target_buffer}; + IREE_RETURN_IF_ERROR( + iree_hal_resource_set_insert(command_buffer->resource_set, 2, buffers)); + VkBufferCopy region; region.srcOffset = iree_hal_buffer_byte_offset(source_buffer) + source_offset; region.dstOffset = iree_hal_buffer_byte_offset(target_buffer) + target_offset; @@ -655,6 +689,12 @@ iree_hal_vulkan_direct_command_buffer_t* command_buffer = iree_hal_vulkan_direct_command_buffer_cast(base_command_buffer); + // TODO(benvanik): batch insert by getting the resources in their own list. + for (iree_host_size_t i = 0; i < binding_count; ++i) { + IREE_RETURN_IF_ERROR(iree_hal_resource_set_insert( + command_buffer->resource_set, 1, &bindings[i].buffer)); + } + // Either allocate, update, and bind a descriptor set or use push descriptor // sets to use the command buffer pool when supported. return command_buffer->descriptor_set_arena.BindDescriptorSet( @@ -672,6 +712,9 @@ iree_allocator_t host_allocator = command_buffer->logical_device->host_allocator(); + IREE_RETURN_IF_ERROR(iree_hal_resource_set_insert( + command_buffer->resource_set, 1, &descriptor_set)); + // Vulkan takes uint32_t as the size here, unlike everywhere else. iree_inline_array(uint32_t, dynamic_offsets_i32, dynamic_offset_count, host_allocator); @@ -713,6 +756,9 @@ source_location.func_name.data, source_location.func_name.size); }); + IREE_RETURN_IF_ERROR(iree_hal_resource_set_insert( + command_buffer->resource_set, 1, &executable)); + // Get the compiled and linked pipeline for the specified entry point and // bind it to the command buffer. VkPipeline pipeline_handle = VK_NULL_HANDLE; @@ -739,6 +785,10 @@ iree_hal_vulkan_direct_command_buffer_t* command_buffer = iree_hal_vulkan_direct_command_buffer_cast(base_command_buffer); + const void* resources[2] = {executable, workgroups_buffer}; + IREE_RETURN_IF_ERROR(iree_hal_resource_set_insert( + command_buffer->resource_set, IREE_ARRAYSIZE(resources), resources)); + iree_hal_vulkan_source_location_t source_location; iree_hal_vulkan_native_executable_entry_point_source_location( executable, entry_point, &source_location);
diff --git a/iree/hal/vulkan/direct_command_buffer.h b/iree/hal/vulkan/direct_command_buffer.h index 071d369..57c15ad 100644 --- a/iree/hal/vulkan/direct_command_buffer.h +++ b/iree/hal/vulkan/direct_command_buffer.h
@@ -18,7 +18,12 @@ extern "C" { #endif // __cplusplus +typedef struct iree_arena_block_pool_t iree_arena_block_pool_t; + // Creates a command buffer that directly records into a VkCommandBuffer. +// +// NOTE: the |block_pool| must remain live for the lifetime of the command +// buffers that use it. iree_status_t iree_hal_vulkan_direct_command_buffer_allocate( iree_hal_device_t* device, iree::hal::vulkan::VkDeviceHandle* logical_device, @@ -29,6 +34,7 @@ iree_hal_vulkan_tracing_context_t* tracing_context, iree::hal::vulkan::DescriptorPoolCache* descriptor_pool_cache, iree::hal::vulkan::BuiltinExecutables* builtin_executables, + iree_arena_block_pool_t* block_pool, iree_hal_command_buffer_t** out_command_buffer); // Returns the native Vulkan VkCommandBuffer handle.
diff --git a/iree/hal/vulkan/tracing.cc b/iree/hal/vulkan/tracing.cc index 36dd3d4..9e0e64f 100644 --- a/iree/hal/vulkan/tracing.cc +++ b/iree/hal/vulkan/tracing.cc
@@ -17,7 +17,7 @@ // Total number of queries the per-queue query pool will contain. This // translates to the maximum number of outstanding queries before collection is // required. -#define IREE_HAL_VULKAN_TRACING_DEFAULT_QUERY_CAPACITY (64 * 1024) +#define IREE_HAL_VULKAN_TRACING_DEFAULT_QUERY_CAPACITY (16 * 1024) // Total number of queries that can be read back from the API in a single // collection. @@ -163,6 +163,7 @@ static void iree_hal_vulkan_tracing_query_calibration_timestamps( iree_hal_vulkan_tracing_context_t* context, uint64_t* out_cpu_time, uint64_t* out_gpu_time) { + IREE_TRACE_ZONE_BEGIN(z0); *out_cpu_time = 0; *out_gpu_time = 0; @@ -197,6 +198,8 @@ default: break; } + + IREE_TRACE_ZONE_END(z0); } // Populates |out_cpu_time| and |out_gpu_time| with calibrated timestamps. @@ -210,6 +213,12 @@ *out_cpu_time = 0; *out_gpu_time = 0; + IREE_TRACE_ZONE_BEGIN(z0); + IREE_TRACE_ZONE_APPEND_TEXT(z0, + context->time_domain == VK_TIME_DOMAIN_DEVICE_EXT + ? "VK_TIME_DOMAIN_DEVICE_EXT" + : "VK_TIME_DOMAIN_CLOCK_MONOTONIC_EXT"); + // Attempt to get a timestamp from both the device and the host at roughly the // same time. There's a gap between when we get control returned to use after // submitting and waiting for idle and that will be the slop we have in the @@ -234,6 +243,7 @@ // Reset the query used. iree_hal_vulkan_tracing_reset_query_pool(context, 0, 1); + IREE_TRACE_ZONE_END(z0); return; } @@ -258,12 +268,14 @@ timestamp_infos[1].pNext = NULL; timestamp_infos[1].timeDomain = context->time_domain; uint64_t max_deviations[IREE_HAL_VULKAN_TRACING_MAX_DEVIATION_PROBE_COUNT]; + IREE_TRACE_ZONE_BEGIN_NAMED(z1, "vkGetCalibratedTimestampsEXT"); for (iree_host_size_t i = 0; i < IREE_ARRAYSIZE(max_deviations); ++i) { uint64_t timestamps[2] = {0, 0}; syms->vkGetCalibratedTimestampsEXT( *context->logical_device, IREE_ARRAYSIZE(timestamps), timestamp_infos, timestamps, &max_deviations[i]); } + IREE_TRACE_ZONE_END(z1); uint64_t min_deviation = max_deviations[0]; for (iree_host_size_t i = 1; i < IREE_ARRAYSIZE(max_deviations); ++i) { min_deviation = iree_min(min_deviation, max_deviations[i]); @@ -273,6 +285,8 @@ iree_hal_vulkan_tracing_query_calibration_timestamps( context, &context->previous_cpu_time, out_gpu_time); *out_cpu_time = tracy::Profiler::GetTime(); + + IREE_TRACE_ZONE_END(z0); } // Performs a periodic calibration (if supported) and sends the data to tracy. @@ -282,6 +296,7 @@ void iree_hal_vulkan_tracing_perform_calibration( iree_hal_vulkan_tracing_context_t* context) { if (context->time_domain == VK_TIME_DOMAIN_DEVICE_EXT) return; + IREE_TRACE_ZONE_BEGIN(z0); uint64_t cpu_time = 0; uint64_t gpu_time = 0; @@ -300,28 +315,36 @@ tracy::MemWrite(&item->gpuCalibration.context, context->id); tracy::Profiler::QueueSerialFinish(); } + + IREE_TRACE_ZONE_END(z0); } // Prepares the VkQueryPool backing storage for our query ringbuffer. static void iree_hal_vulkan_tracing_prepare_query_pool( iree_hal_vulkan_tracing_context_t* context) { + IREE_TRACE_ZONE_BEGIN(z0); + // Create a query pool with the largest query capacity it can provide. VkQueryPoolCreateInfo pool_info; memset(&pool_info, 0, sizeof(pool_info)); pool_info.sType = VK_STRUCTURE_TYPE_QUERY_POOL_CREATE_INFO; pool_info.queryCount = IREE_HAL_VULKAN_TRACING_DEFAULT_QUERY_CAPACITY; pool_info.queryType = VK_QUERY_TYPE_TIMESTAMP; + IREE_TRACE_ZONE_APPEND_VALUE(z0, pool_info.queryCount); while (context->logical_device->syms()->vkCreateQueryPool( *context->logical_device, &pool_info, context->logical_device->allocator(), &context->query_pool) != VK_SUCCESS) { pool_info.queryCount /= 2; + IREE_TRACE_ZONE_APPEND_VALUE(z0, pool_info.queryCount); } context->query_capacity = pool_info.queryCount; // Perform initial reset of the query pool. All queries must be reset upon // creation before first use. iree_hal_vulkan_tracing_reset_query_pool(context, 0, context->query_capacity); + + IREE_TRACE_ZONE_END(z0); } // Prepares the Tracy-related GPU context that events are fed into. Each context @@ -329,6 +352,8 @@ static void iree_hal_vulkan_tracing_prepare_gpu_context( iree_hal_vulkan_tracing_context_t* context, VkPhysicalDevice physical_device, iree_string_view_t queue_name) { + IREE_TRACE_ZONE_BEGIN(z0); + // Allocate the process-unique GPU context ID. There's a max of 255 available; // if we are recreating devices a lot we may exceed that. Don't do that, or // wrap around and get weird (but probably still usable) numbers. @@ -386,6 +411,8 @@ tracy::MemWrite(&item->gpuContextNameFat.size, queue_name.size); tracy::Profiler::QueueSerialFinish(); } + + IREE_TRACE_ZONE_END(z0); } // Returns the best possible platform-supported time domain, falling back to
diff --git a/iree/hal/vulkan/vma_allocator.cc b/iree/hal/vulkan/vma_allocator.cc index 50c1dab..18e49fd 100644 --- a/iree/hal/vulkan/vma_allocator.cc +++ b/iree/hal/vulkan/vma_allocator.cc
@@ -20,6 +20,7 @@ typedef struct iree_hal_vulkan_vma_allocator_t { iree_hal_resource_t resource; + iree_hal_device_t* device; // unretained to avoid cycles iree_allocator_t host_allocator; VmaAllocator vma; @@ -82,11 +83,12 @@ iree_status_t iree_hal_vulkan_vma_allocator_create( VkInstance instance, VkPhysicalDevice physical_device, - VkDeviceHandle* logical_device, VmaRecordSettings record_settings, - iree_hal_allocator_t** out_allocator) { + VkDeviceHandle* logical_device, iree_hal_device_t* device, + VmaRecordSettings record_settings, iree_hal_allocator_t** out_allocator) { IREE_ASSERT_ARGUMENT(instance); IREE_ASSERT_ARGUMENT(physical_device); IREE_ASSERT_ARGUMENT(logical_device); + IREE_ASSERT_ARGUMENT(device); IREE_ASSERT_ARGUMENT(out_allocator); IREE_TRACE_ZONE_BEGIN(z0); @@ -98,6 +100,7 @@ iree_hal_resource_initialize(&iree_hal_vulkan_vma_allocator_vtable, &allocator->resource); allocator->host_allocator = host_allocator; + allocator->device = device; const auto& syms = logical_device->syms(); VmaVulkanFunctions vulkan_fns; @@ -188,6 +191,11 @@ return allocator->host_allocator; } +static iree_status_t iree_hal_vulkan_vma_allocator_trim( + iree_hal_allocator_t* base_allocator) { + return iree_ok_status(); +} + static void iree_hal_vulkan_vma_allocator_query_statistics( iree_hal_allocator_t* base_allocator, iree_hal_allocator_statistics_t* out_statistics) { @@ -215,11 +223,13 @@ iree_hal_buffer_compatibility_t compatibility = IREE_HAL_BUFFER_COMPATIBILITY_ALLOCATABLE; + // All buffers can be used as transfer source/dest. + if (iree_all_bits_set(intended_usage, IREE_HAL_BUFFER_USAGE_TRANSFER)) { + compatibility |= IREE_HAL_BUFFER_COMPATIBILITY_QUEUE_TRANSFER; + } + // Buffers can only be used on the queue if they are device visible. if (iree_all_bits_set(memory_type, IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE)) { - if (iree_all_bits_set(intended_usage, IREE_HAL_BUFFER_USAGE_TRANSFER)) { - compatibility |= IREE_HAL_BUFFER_COMPATIBILITY_QUEUE_TRANSFER; - } if (iree_all_bits_set(intended_usage, IREE_HAL_BUFFER_USAGE_DISPATCH)) { compatibility |= IREE_HAL_BUFFER_COMPATIBILITY_QUEUE_DISPATCH; } @@ -239,7 +249,8 @@ iree_hal_vulkan_vma_allocator_t* allocator, iree_hal_memory_type_t memory_type, iree_hal_buffer_usage_t allowed_usage, iree_hal_memory_access_t allowed_access, iree_host_size_t allocation_size, - VmaAllocationCreateFlags flags, iree_hal_buffer_t** out_buffer) { + iree_const_byte_span_t initial_data, VmaAllocationCreateFlags flags, + iree_hal_buffer_t** out_buffer) { // Guard against the corner case where the requested buffer size is 0. The // application is unlikely to do anything when requesting a 0-byte buffer; but // it can happen in real world use cases. So we should at least not crash. @@ -311,6 +322,9 @@ allocation_create_info.requiredFlags |= VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT; } + // TODO(benvanik): if on a unified memory system and initial data is present + // we could set the mapping bit and ensure a much more efficient upload. + VkBuffer handle = VK_NULL_HANDLE; VmaAllocation allocation = VK_NULL_HANDLE; VmaAllocationInfo allocation_info; @@ -319,18 +333,42 @@ &allocation, &allocation_info), "vmaCreateBuffer"); - return iree_hal_vulkan_vma_buffer_wrap( + iree_hal_buffer_t* buffer = NULL; + iree_status_t status = iree_hal_vulkan_vma_buffer_wrap( (iree_hal_allocator_t*)allocator, memory_type, allowed_access, allowed_usage, allocation_size, /*byte_offset=*/0, /*byte_length=*/allocation_size, allocator->vma, handle, allocation, - allocation_info, out_buffer); + allocation_info, &buffer); + if (!iree_status_is_ok(status)) { + vmaDestroyBuffer(allocator->vma, handle, allocation); + return status; + } + + // Copy the initial contents into the buffer. This may require staging. + if (iree_status_is_ok(status) && + !iree_const_byte_span_is_empty(initial_data)) { + status = iree_hal_device_transfer_range( + allocator->device, + iree_hal_make_host_transfer_buffer_span((void*)initial_data.data, + initial_data.data_length), + 0, iree_hal_make_device_transfer_buffer(buffer), 0, + initial_data.data_length, IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, + iree_infinite_timeout()); + } + + if (iree_status_is_ok(status)) { + *out_buffer = buffer; + } else { + iree_hal_buffer_release(buffer); + } + return status; } static iree_status_t iree_hal_vulkan_vma_allocator_allocate_buffer( iree_hal_allocator_t* base_allocator, iree_hal_memory_type_t memory_type, iree_hal_buffer_usage_t allowed_usage, iree_host_size_t allocation_size, - iree_hal_buffer_t** out_buffer) { + iree_const_byte_span_t initial_data, iree_hal_buffer_t** out_buffer) { iree_hal_vulkan_vma_allocator_t* allocator = iree_hal_vulkan_vma_allocator_cast(base_allocator); @@ -341,6 +379,7 @@ return iree_hal_vulkan_vma_allocator_allocate_internal( allocator, memory_type, allowed_usage, allowed_access, allocation_size, + initial_data, /*flags=*/0, out_buffer); } @@ -364,6 +403,7 @@ const iree_hal_allocator_vtable_t iree_hal_vulkan_vma_allocator_vtable = { /*.destroy=*/iree_hal_vulkan_vma_allocator_destroy, /*.host_allocator=*/iree_hal_vulkan_vma_allocator_host_allocator, + /*.trim=*/iree_hal_vulkan_vma_allocator_trim, /*.query_statistics=*/iree_hal_vulkan_vma_allocator_query_statistics, /*.query_buffer_compatibility=*/ iree_hal_vulkan_vma_allocator_query_buffer_compatibility,
diff --git a/iree/hal/vulkan/vma_allocator.h b/iree/hal/vulkan/vma_allocator.h index 71c068b..bc6a17e 100644 --- a/iree/hal/vulkan/vma_allocator.h +++ b/iree/hal/vulkan/vma_allocator.h
@@ -36,7 +36,8 @@ iree_status_t iree_hal_vulkan_vma_allocator_create( VkInstance instance, VkPhysicalDevice physical_device, iree::hal::vulkan::VkDeviceHandle* logical_device, - VmaRecordSettings record_settings, iree_hal_allocator_t** out_allocator); + iree_hal_device_t* device, VmaRecordSettings record_settings, + iree_hal_allocator_t** out_allocator); #ifdef __cplusplus } // extern "C"
diff --git a/iree/hal/vulkan/vma_buffer.cc b/iree/hal/vulkan/vma_buffer.cc index b69cd67..7bf517b 100644 --- a/iree/hal/vulkan/vma_buffer.cc +++ b/iree/hal/vulkan/vma_buffer.cc
@@ -103,15 +103,24 @@ iree_hal_buffer_t* base_buffer, iree_hal_mapping_mode_t mapping_mode, iree_hal_memory_access_t memory_access, iree_device_size_t local_byte_offset, iree_device_size_t local_byte_length, - void** out_data_ptr) { + iree_hal_buffer_mapping_t* mapping) { iree_hal_vulkan_vma_buffer_t* buffer = iree_hal_vulkan_vma_buffer_cast(base_buffer); + // TODO(benvanik): add upload/download for unmapped buffers. + IREE_RETURN_IF_ERROR(iree_hal_buffer_validate_memory_type( + iree_hal_buffer_memory_type(base_buffer), + IREE_HAL_MEMORY_TYPE_HOST_VISIBLE)); + IREE_RETURN_IF_ERROR( + iree_hal_buffer_validate_usage(iree_hal_buffer_allowed_usage(base_buffer), + IREE_HAL_BUFFER_USAGE_MAPPING)); + uint8_t* data_ptr = nullptr; VK_RETURN_IF_ERROR( vmaMapMemory(buffer->vma, buffer->allocation, (void**)&data_ptr), "vmaMapMemory"); - *out_data_ptr = data_ptr + local_byte_offset; + mapping->contents = + iree_make_byte_span(data_ptr + local_byte_offset, local_byte_length); // If we mapped for discard scribble over the bytes. This is not a mandated // behavior but it will make debugging issues easier. Alternatively for @@ -119,19 +128,20 @@ // would only work if the entire buffer was discarded. #ifndef NDEBUG if (iree_any_bit_set(memory_access, IREE_HAL_MEMORY_ACCESS_DISCARD)) { - memset(*out_data_ptr, 0xCD, local_byte_length); + memset(mapping->contents.data, 0xCD, local_byte_length); } #endif // !NDEBUG return iree_ok_status(); } -static void iree_hal_vulkan_vma_buffer_unmap_range( +static iree_status_t iree_hal_vulkan_vma_buffer_unmap_range( iree_hal_buffer_t* base_buffer, iree_device_size_t local_byte_offset, - iree_device_size_t local_byte_length, void* data_ptr) { + iree_device_size_t local_byte_length, iree_hal_buffer_mapping_t* mapping) { iree_hal_vulkan_vma_buffer_t* buffer = iree_hal_vulkan_vma_buffer_cast(base_buffer); vmaUnmapMemory(buffer->vma, buffer->allocation); + return iree_ok_status(); } static iree_status_t iree_hal_vulkan_vma_buffer_invalidate_range(
diff --git a/iree/hal/vulkan/vulkan_device.cc b/iree/hal/vulkan/vulkan_device.cc index 97cde8a..966e4e3 100644 --- a/iree/hal/vulkan/vulkan_device.cc +++ b/iree/hal/vulkan/vulkan_device.cc
@@ -11,8 +11,10 @@ #include <cstring> #include <vector> +#include "iree/base/internal/arena.h" #include "iree/base/internal/math.h" #include "iree/base/tracing.h" +#include "iree/hal/utils/buffer_transfer.h" #include "iree/hal/vulkan/api.h" #include "iree/hal/vulkan/builtin_executables.h" #include "iree/hal/vulkan/command_queue.h" @@ -361,6 +363,10 @@ VkCommandPoolHandle* dispatch_command_pool; VkCommandPoolHandle* transfer_command_pool; + // Block pool used for command buffers with a larger block size (as command + // buffers can contain inlined data uploads). + iree_arena_block_pool_t block_pool; + // Used only for emulated timeline semaphores. TimePointSemaphorePool* semaphore_pool; TimePointFencePool* fence_pool; @@ -562,6 +568,9 @@ device->logical_device = logical_device; device->logical_device->AddReference(); + iree_arena_block_pool_initialize(32 * 1024, host_allocator, + &device->block_pool); + // Point the queue storage into the new device allocation. The queues // themselves are populated device->queues = (CommandQueue**)buffer_ptr; @@ -582,8 +591,8 @@ VmaRecordSettings vma_record_settings; memset(&vma_record_settings, 0, sizeof(vma_record_settings)); iree_status_t status = iree_hal_vulkan_vma_allocator_create( - instance, physical_device, logical_device, vma_record_settings, - &device->device_allocator); + instance, physical_device, logical_device, (iree_hal_device_t*)device, + vma_record_settings, &device->device_allocator); // Create command pools for each queue family. If we don't have a transfer // queue then we'll ignore that one and just use the dispatch pool. @@ -666,6 +675,9 @@ // There should be no more buffers live that use the allocator. iree_hal_allocator_release(device->device_allocator); + // All arena blocks should have been returned. + iree_arena_block_pool_deinitialize(&device->block_pool); + // Finally, destroy the device. device->logical_device->ReleaseReference(); iree_hal_driver_release(device->driver); @@ -915,6 +927,13 @@ return device->device_allocator; } +static iree_status_t iree_hal_vulkan_device_trim( + iree_hal_device_t* base_device) { + iree_hal_vulkan_device_t* device = iree_hal_vulkan_device_cast(base_device); + iree_arena_block_pool_trim(&device->block_pool); + return iree_hal_allocator_trim(device->device_allocator); +} + static iree_status_t iree_hal_vulkan_device_query_i32( iree_hal_device_t* base_device, iree_string_view_t category, iree_string_view_t key, int32_t* out_value) { @@ -999,7 +1018,7 @@ base_device, device->logical_device, command_pool, mode, command_categories, queue_affinity, queue->tracing_context(), device->descriptor_pool_cache, device->builtin_executables, - out_command_buffer); + &device->block_pool, out_command_buffer); } static iree_status_t iree_hal_vulkan_device_create_descriptor_set( @@ -1120,6 +1139,7 @@ /*.id=*/iree_hal_vulkan_device_id, /*.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, /*.create_command_buffer=*/iree_hal_vulkan_device_create_command_buffer, /*.create_descriptor_set=*/iree_hal_vulkan_device_create_descriptor_set, @@ -1131,6 +1151,7 @@ /*.create_executable_layout=*/ iree_hal_vulkan_device_create_executable_layout, /*.create_semaphore=*/iree_hal_vulkan_device_create_semaphore, + /*.transfer_range=*/iree_hal_device_submit_transfer_range_and_wait, /*.queue_submit=*/iree_hal_vulkan_device_queue_submit, /*.submit_and_wait=*/ iree_hal_vulkan_device_submit_and_wait,
diff --git a/iree/modules/check/check_test.cc b/iree/modules/check/check_test.cc index 5bcd6a9..142dd6f 100644 --- a/iree/modules/check/check_test.cc +++ b/iree/modules/check/check_test.cc
@@ -80,18 +80,13 @@ num_elements *= dim; } ASSERT_EQ(contents.size(), num_elements); - vm::ref<iree_hal_buffer_t> buffer; - IREE_ASSERT_OK(iree_hal_allocator_allocate_buffer( - allocator_, - static_cast<iree_hal_memory_type_t>( - IREE_HAL_MEMORY_TYPE_HOST_LOCAL | - IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE), - IREE_HAL_BUFFER_USAGE_ALL, contents.size() * sizeof(int32_t), &buffer)); - IREE_ASSERT_OK(iree_hal_buffer_write_data( - buffer.get(), 0, contents.data(), contents.size() * sizeof(int32_t))); - IREE_ASSERT_OK(iree_hal_buffer_view_create( - buffer.get(), shape.data(), shape.size(), IREE_HAL_ELEMENT_TYPE_INT_32, - IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, iree_allocator_system(), + IREE_ASSERT_OK(iree_hal_buffer_view_allocate_buffer( + allocator_, shape.data(), shape.size(), IREE_HAL_ELEMENT_TYPE_INT_32, + IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, + IREE_HAL_MEMORY_TYPE_HOST_LOCAL | IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE, + IREE_HAL_BUFFER_USAGE_ALL, + iree_make_const_byte_span(contents.data(), + contents.size() * sizeof(int32_t)), &*out_buffer_view)); } @@ -103,20 +98,14 @@ num_elements *= dim; } ASSERT_EQ(contents.size(), num_elements); - vm::ref<iree_hal_buffer_t> buffer; - IREE_ASSERT_OK(iree_hal_allocator_allocate_buffer( - allocator_, - static_cast<iree_hal_memory_type_t>( - IREE_HAL_MEMORY_TYPE_HOST_LOCAL | - IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE), - IREE_HAL_BUFFER_USAGE_ALL, contents.size() * sizeof(uint16_t), - &buffer)); - IREE_ASSERT_OK(iree_hal_buffer_write_data( - buffer.get(), 0, contents.data(), contents.size() * sizeof(uint16_t))); - IREE_ASSERT_OK(iree_hal_buffer_view_create( - buffer.get(), shape.data(), shape.size(), - IREE_HAL_ELEMENT_TYPE_FLOAT_16, IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, - iree_allocator_system(), &*out_buffer_view)); + IREE_ASSERT_OK(iree_hal_buffer_view_allocate_buffer( + allocator_, shape.data(), shape.size(), IREE_HAL_ELEMENT_TYPE_FLOAT_16, + IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, + IREE_HAL_MEMORY_TYPE_HOST_LOCAL | IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE, + IREE_HAL_BUFFER_USAGE_ALL, + iree_make_const_byte_span(contents.data(), + contents.size() * sizeof(uint16_t)), + &*out_buffer_view)); } void CreateFloat32BufferView(iree::span<const float> contents, @@ -127,19 +116,14 @@ num_elements *= dim; } ASSERT_EQ(contents.size(), num_elements); - vm::ref<iree_hal_buffer_t> buffer; - IREE_ASSERT_OK(iree_hal_allocator_allocate_buffer( - allocator_, - static_cast<iree_hal_memory_type_t>( - IREE_HAL_MEMORY_TYPE_HOST_LOCAL | - IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE), - IREE_HAL_BUFFER_USAGE_ALL, contents.size() * sizeof(float), &buffer)); - IREE_ASSERT_OK(iree_hal_buffer_write_data(buffer.get(), 0, contents.data(), - contents.size() * sizeof(float))); - IREE_ASSERT_OK(iree_hal_buffer_view_create( - buffer.get(), shape.data(), shape.size(), - IREE_HAL_ELEMENT_TYPE_FLOAT_32, IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, - iree_allocator_system(), &*out_buffer_view)); + IREE_ASSERT_OK(iree_hal_buffer_view_allocate_buffer( + allocator_, shape.data(), shape.size(), IREE_HAL_ELEMENT_TYPE_FLOAT_32, + IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, + IREE_HAL_MEMORY_TYPE_HOST_LOCAL | IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE, + IREE_HAL_BUFFER_USAGE_ALL, + iree_make_const_byte_span(contents.data(), + contents.size() * sizeof(float)), + &*out_buffer_view)); } void CreateFloat64BufferView(iree::span<const double> contents, @@ -150,19 +134,14 @@ num_elements *= dim; } ASSERT_EQ(contents.size(), num_elements); - vm::ref<iree_hal_buffer_t> buffer; - IREE_ASSERT_OK(iree_hal_allocator_allocate_buffer( - allocator_, - static_cast<iree_hal_memory_type_t>( - IREE_HAL_MEMORY_TYPE_HOST_LOCAL | - IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE), - IREE_HAL_BUFFER_USAGE_ALL, contents.size() * sizeof(double), &buffer)); - IREE_ASSERT_OK(iree_hal_buffer_write_data( - buffer.get(), 0, contents.data(), contents.size() * sizeof(double))); - IREE_ASSERT_OK(iree_hal_buffer_view_create( - buffer.get(), shape.data(), shape.size(), - IREE_HAL_ELEMENT_TYPE_FLOAT_64, IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, - iree_allocator_system(), &*out_buffer_view)); + IREE_ASSERT_OK(iree_hal_buffer_view_allocate_buffer( + allocator_, shape.data(), shape.size(), IREE_HAL_ELEMENT_TYPE_FLOAT_64, + IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, + IREE_HAL_MEMORY_TYPE_HOST_LOCAL | IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE, + IREE_HAL_BUFFER_USAGE_ALL, + iree_make_const_byte_span(contents.data(), + contents.size() * sizeof(double)), + &*out_buffer_view)); } iree_status_t Invoke(const char* function_name) {
diff --git a/iree/modules/check/module.cc b/iree/modules/check/module.cc index 061fa71..9996e94 100644 --- a/iree/modules/check/module.cc +++ b/iree/modules/check/module.cc
@@ -184,13 +184,13 @@ iree_hal_buffer_view_element_type(view); iree_hal_buffer_t* buf = iree_hal_buffer_view_buffer(view); iree_device_size_t size = iree_hal_buffer_view_byte_length(view); - iree_hal_buffer_mapping_t mapped_memory; - IREE_RETURN_IF_ERROR( - iree_hal_buffer_map_range(buf, IREE_HAL_MEMORY_ACCESS_READ, - /*byte_offset=*/0, size, &mapped_memory)); + iree_hal_buffer_mapping_t mapped_memory = {{0}}; + IREE_RETURN_IF_ERROR(iree_hal_buffer_map_range( + buf, IREE_HAL_MAPPING_MODE_SCOPED, IREE_HAL_MEMORY_ACCESS_READ, + /*byte_offset=*/0, size, &mapped_memory)); IREE_RETURN_IF_ERROR( ::iree::ExpectAllTrue(mapped_memory.contents, element_type)); - iree_hal_buffer_unmap_range(&mapped_memory); + iree_status_ignore(iree_hal_buffer_unmap_range(&mapped_memory)); return OkStatus(); } @@ -220,23 +220,26 @@ iree_hal_element_type_t rhs_element_type = iree_hal_buffer_view_element_type(rhs); + // HACK: this is all broken and will leak. Let's kill this entire module + // please. + iree_hal_buffer_t* lhs_buf = iree_hal_buffer_view_buffer(lhs); - iree_hal_buffer_mapping_t lhs_mapped_memory; + iree_hal_buffer_mapping_t lhs_mapped_memory = {{0}}; IREE_RETURN_IF_ERROR(iree_hal_buffer_map_range( - lhs_buf, IREE_HAL_MEMORY_ACCESS_READ, + lhs_buf, IREE_HAL_MAPPING_MODE_SCOPED, IREE_HAL_MEMORY_ACCESS_READ, /*byte_offset=*/0, lhs_size, &lhs_mapped_memory)); iree_hal_buffer_t* rhs_buf = iree_hal_buffer_view_buffer(rhs); - iree_hal_buffer_mapping_t rhs_mapped_memory; + iree_hal_buffer_mapping_t rhs_mapped_memory = {{0}}; IREE_RETURN_IF_ERROR(iree_hal_buffer_map_range( - rhs_buf, IREE_HAL_MEMORY_ACCESS_READ, + rhs_buf, IREE_HAL_MAPPING_MODE_SCOPED, IREE_HAL_MEMORY_ACCESS_READ, /*byte_offset=*/0, rhs_size, &rhs_mapped_memory)); bool element_types_eq = lhs_element_type == rhs_element_type; bool shape_eq = lhs_shape == rhs_shape; bool contents_eq = EqByteSpan(lhs_mapped_memory.contents, rhs_mapped_memory.contents); - iree_hal_buffer_unmap_range(&lhs_mapped_memory); - iree_hal_buffer_unmap_range(&rhs_mapped_memory); + iree_status_ignore(iree_hal_buffer_unmap_range(&lhs_mapped_memory)); + iree_status_ignore(iree_hal_buffer_unmap_range(&rhs_mapped_memory)); if (!element_types_eq || !shape_eq || !contents_eq) { std::ostringstream os; @@ -297,14 +300,14 @@ iree_hal_buffer_view_element_type(rhs); iree_hal_buffer_t* lhs_buf = iree_hal_buffer_view_buffer(lhs); - iree_hal_buffer_mapping_t lhs_mapped_memory; + iree_hal_buffer_mapping_t lhs_mapped_memory = {{0}}; IREE_RETURN_IF_ERROR(iree_hal_buffer_map_range( - lhs_buf, IREE_HAL_MEMORY_ACCESS_READ, + lhs_buf, IREE_HAL_MAPPING_MODE_SCOPED, IREE_HAL_MEMORY_ACCESS_READ, /*byte_offset=*/0, lhs_size, &lhs_mapped_memory)); iree_hal_buffer_t* rhs_buf = iree_hal_buffer_view_buffer(rhs); - iree_hal_buffer_mapping_t rhs_mapped_memory; + iree_hal_buffer_mapping_t rhs_mapped_memory = {{0}}; IREE_RETURN_IF_ERROR(iree_hal_buffer_map_range( - rhs_buf, IREE_HAL_MEMORY_ACCESS_READ, + rhs_buf, IREE_HAL_MAPPING_MODE_SCOPED, IREE_HAL_MEMORY_ACCESS_READ, /*byte_offset=*/0, rhs_size, &rhs_mapped_memory)); bool element_types_eq = lhs_element_type == rhs_element_type; @@ -317,8 +320,8 @@ AlmostEqByteSpan(lhs_mapped_memory.contents, rhs_mapped_memory.contents, lhs_element_type)); } - iree_hal_buffer_unmap_range(&lhs_mapped_memory); - iree_hal_buffer_unmap_range(&rhs_mapped_memory); + iree_status_ignore(iree_hal_buffer_unmap_range(&lhs_mapped_memory)); + iree_status_ignore(iree_hal_buffer_unmap_range(&rhs_mapped_memory)); if (!element_types_eq || !shape_eq || !contents_could_be_almost_eq) { std::ostringstream os;
diff --git a/iree/modules/hal/module.c b/iree/modules/hal/module.c index ccd5cbc..b0b0a7f 100644 --- a/iree/modules/hal/module.c +++ b/iree/modules/hal/module.c
@@ -130,9 +130,6 @@ iree_hal_semaphore_t* submit_semaphore; uint64_t submit_value; - - void* deferred_lru[6]; - iree_vm_list_t* deferred_releases; } iree_hal_module_state_t; static void IREE_API_PTR iree_hal_module_destroy(void* base_module) { @@ -143,39 +140,56 @@ static iree_status_t IREE_API_PTR iree_hal_module_alloc_state(void* self, iree_allocator_t host_allocator, iree_vm_module_state_t** out_module_state) { + IREE_TRACE_ZONE_BEGIN(z0); + iree_hal_module_t* module = IREE_HAL_MODULE_CAST(self); iree_hal_module_state_t* state = NULL; - IREE_RETURN_IF_ERROR( + IREE_RETURN_AND_END_ZONE_IF_ERROR( + z0, iree_allocator_malloc(host_allocator, sizeof(*state), (void**)&state)); memset(state, 0, sizeof(*state)); state->host_allocator = host_allocator; state->shared_device = module->shared_device; iree_hal_device_retain(state->shared_device); - IREE_RETURN_IF_ERROR(iree_vm_list_create( - /*element_type=*/NULL, /*initial_capacity=*/512, state->host_allocator, - &state->deferred_releases)); - - IREE_RETURN_IF_ERROR(iree_hal_executable_cache_create( - state->shared_device, iree_string_view_empty(), - &state->executable_cache)); + IREE_RETURN_AND_END_ZONE_IF_ERROR( + z0, iree_hal_executable_cache_create(state->shared_device, + iree_string_view_empty(), + &state->executable_cache)); state->submit_value = 0ull; - IREE_RETURN_IF_ERROR(iree_hal_semaphore_create( - state->shared_device, state->submit_value, &state->submit_semaphore)); + IREE_RETURN_AND_END_ZONE_IF_ERROR( + z0, iree_hal_semaphore_create(state->shared_device, state->submit_value, + &state->submit_semaphore)); *out_module_state = (iree_vm_module_state_t*)state; + IREE_TRACE_ZONE_END(z0); return iree_ok_status(); } static void IREE_API_PTR iree_hal_module_free_state(void* self, iree_vm_module_state_t* module_state) { + IREE_TRACE_ZONE_BEGIN(z0); + iree_hal_module_state_t* state = (iree_hal_module_state_t*)module_state; iree_hal_semaphore_release(state->submit_semaphore); - iree_vm_list_release(state->deferred_releases); iree_hal_executable_cache_release(state->executable_cache); iree_hal_device_release(state->shared_device); iree_allocator_free(state->host_allocator, state); + + IREE_TRACE_ZONE_END(z0); +} + +static iree_status_t IREE_API_PTR iree_hal_module_notify( + void* self, iree_vm_module_state_t* module_state, iree_vm_signal_t signal) { + iree_hal_module_state_t* state = (iree_hal_module_state_t*)module_state; + switch (signal) { + case IREE_VM_SIGNAL_SUSPEND: + case IREE_VM_SIGNAL_LOW_MEMORY: + return iree_hal_device_trim(state->shared_device); + default: + return iree_ok_status(); + } } //===----------------------------------------------------------------------===// @@ -191,30 +205,6 @@ return iree_ok_status(); } -void iree_hal_module_ex_defer_release(iree_hal_module_state_t* state, - const iree_vm_ref_t value) { - // A bulk of the calls to this are for the same (or very recently same) - // objects, such as constant pool or transient buffer storage that may be - // bound 4-10 times per dispatch. This tiny LRU lets us avoid adding such - // repeated patterns in the common case. - for (iree_host_size_t i = 0; i < IREE_ARRAYSIZE(state->deferred_lru); ++i) { - if (state->deferred_lru[i] == value.ptr) { - // Hit - keep the list sorted by most->least recently used. - state->deferred_lru[i] = state->deferred_lru[0]; - state->deferred_lru[0] = value.ptr; - return; - } - } - // Miss - shift the list down and insert the new item at the head. - memmove(&state->deferred_lru[1], &state->deferred_lru[0], - sizeof(state->deferred_lru[0]) * - (IREE_ARRAYSIZE(state->deferred_lru) - 1)); - state->deferred_lru[0] = value.ptr; - - IREE_IGNORE_ERROR( - iree_vm_list_push_ref_retain(state->deferred_releases, &value)); -} - IREE_VM_ABI_EXPORT(iree_hal_module_ex_submit_and_wait, // iree_hal_module_state_t, // rr, v) { @@ -246,12 +236,6 @@ return status; } - // Drop all pending deferred releases (references to everything in flight). - // This will be replaced with resource sets in the future that are attached to - // each command buffer. - IREE_RETURN_IF_ERROR(iree_vm_list_resize(state->deferred_releases, 0)); - memset(state->deferred_lru, 0, sizeof(state->deferred_lru)); - return iree_ok_status(); } @@ -270,7 +254,8 @@ iree_hal_buffer_t* buffer = NULL; IREE_RETURN_IF_ERROR(iree_hal_allocator_allocate_buffer( - allocator, memory_types, buffer_usage, allocation_size, &buffer)); + allocator, memory_types, buffer_usage, allocation_size, + iree_const_byte_span_empty(), &buffer)); rets->r0 = iree_hal_buffer_move_ref(buffer); return iree_ok_status(); } @@ -389,18 +374,14 @@ iree_hal_buffer_t* buffer = NULL; IREE_RETURN_IF_ERROR( - iree_hal_allocator_allocate_buffer(allocator, memory_types, buffer_usage, - length, &buffer), + iree_hal_allocator_allocate_buffer( + allocator, memory_types, buffer_usage, length, + iree_make_const_byte_span(source->data.data + offset, length), + &buffer), "failed to allocate buffer of length %d", length); - iree_status_t status = - iree_hal_buffer_write_data(buffer, 0, source->data.data + offset, length); - if (iree_status_is_ok(status)) { - rets->r0 = iree_hal_buffer_move_ref(buffer); - } else { - iree_hal_buffer_release(buffer); - } - return status; + rets->r0 = iree_hal_buffer_move_ref(buffer); + return iree_ok_status(); } //===----------------------------------------------------------------------===// @@ -958,9 +939,6 @@ iree_vm_size_t length = (iree_vm_size_t)args->i3; uint32_t pattern = (uint32_t)args->i4; uint32_t pattern_length = (uint32_t)args->i5; - - iree_hal_module_ex_defer_release(state, args->r1); - return iree_hal_command_buffer_fill_buffer(command_buffer, target_buffer, target_offset, length, &pattern, pattern_length); @@ -979,10 +957,6 @@ IREE_RETURN_IF_ERROR(iree_hal_buffer_check_deref(args->r3, &target_buffer)); iree_vm_size_t target_offset = (iree_vm_size_t)args->i4; iree_vm_size_t length = (iree_vm_size_t)args->i5; - - iree_hal_module_ex_defer_release(state, args->r1); - iree_hal_module_ex_defer_release(state, args->r3); - return iree_hal_command_buffer_copy_buffer(command_buffer, source_buffer, source_offset, target_buffer, target_offset, length); @@ -1033,7 +1007,6 @@ bindings[i].binding = (uint32_t)args->a3[i].i0; bindings[i].offset = (iree_device_size_t)args->a3[i].i2; bindings[i].length = (iree_device_size_t)args->a3[i].i3; - iree_hal_module_ex_defer_release(state, args->a3[i].r1); } return iree_hal_command_buffer_push_descriptor_set( @@ -1057,9 +1030,6 @@ iree_device_size_t* dynamic_offsets = NULL; IREE_VM_ABI_VLA_STACK_CAST(args, a4_count, a4, iree_device_size_t, 64, &dynamic_offset_count, &dynamic_offsets); - - iree_hal_module_ex_defer_release(state, args->r3); - return iree_hal_command_buffer_bind_descriptor_set( command_buffer, executable_layout, set, descriptor_set, dynamic_offset_count, dynamic_offsets); @@ -1077,9 +1047,6 @@ uint32_t workgroup_x = (uint32_t)args->i3; uint32_t workgroup_y = (uint32_t)args->i4; uint32_t workgroup_z = (uint32_t)args->i5; - - iree_hal_module_ex_defer_release(state, args->r1); - return iree_hal_command_buffer_dispatch(command_buffer, executable, entry_point, workgroup_x, workgroup_y, workgroup_z); @@ -1098,10 +1065,6 @@ IREE_RETURN_IF_ERROR( iree_hal_buffer_check_deref(args->r3, &workgroups_buffer)); iree_vm_size_t workgroups_offset = (iree_vm_size_t)args->i4; - - iree_hal_module_ex_defer_release(state, args->r1); - iree_hal_module_ex_defer_release(state, args->r3); - return iree_hal_command_buffer_dispatch_indirect( command_buffer, executable, entry_point, workgroups_buffer, workgroups_offset); @@ -1419,6 +1382,7 @@ .destroy = iree_hal_module_destroy, .alloc_state = iree_hal_module_alloc_state, .free_state = iree_hal_module_free_state, + .notify = iree_hal_module_notify, }; // Allocate shared module state.
diff --git a/iree/runtime/demo/hello_world_explained.c b/iree/runtime/demo/hello_world_explained.c index cbdfb09..589a1d1 100644 --- a/iree/runtime/demo/hello_world_explained.c +++ b/iree/runtime/demo/hello_world_explained.c
@@ -183,6 +183,8 @@ // in other sessions depending on whether they share a compatible device. iree_hal_allocator_t* device_allocator = iree_runtime_session_device_allocator(session); + iree_allocator_t host_allocator = + iree_runtime_session_host_allocator(session); iree_status_t status = iree_ok_status(); { // %arg0: tensor<4xf32> @@ -212,7 +214,7 @@ } if (iree_status_is_ok(status)) { IREE_IGNORE_ERROR(iree_hal_buffer_view_fprint( - stdout, arg0, /*max_element_count=*/4096)); + stdout, arg0, /*max_element_count=*/4096, host_allocator)); // Add to the call inputs list (which retains the buffer view). status = iree_runtime_call_inputs_push_back_buffer_view(&call, arg0); } @@ -237,7 +239,7 @@ } if (iree_status_is_ok(status)) { IREE_IGNORE_ERROR(iree_hal_buffer_view_fprint( - stdout, arg1, /*max_element_count=*/4096)); + stdout, arg1, /*max_element_count=*/4096, host_allocator)); status = iree_runtime_call_inputs_push_back_buffer_view(&call, arg1); } iree_hal_buffer_view_release(arg1); @@ -259,8 +261,8 @@ if (iree_status_is_ok(status)) { // This prints the buffer view out but an application could read its // contents, pass it to another call, etc. - status = - iree_hal_buffer_view_fprint(stdout, ret0, /*max_element_count=*/4096); + status = iree_hal_buffer_view_fprint( + stdout, ret0, /*max_element_count=*/4096, host_allocator); } iree_hal_buffer_view_release(ret0);
diff --git a/iree/runtime/demo/hello_world_terse.c b/iree/runtime/demo/hello_world_terse.c index 3a10d32..b333e60 100644 --- a/iree/runtime/demo/hello_world_terse.c +++ b/iree/runtime/demo/hello_world_terse.c
@@ -87,8 +87,9 @@ IREE_HAL_MEMORY_ACCESS_READ, IREE_HAL_BUFFER_USAGE_ALL, iree_make_byte_span((void*)arg0_data, sizeof(arg0_data)), iree_allocator_null(), &arg0)); - IREE_CHECK_OK( - iree_hal_buffer_view_fprint(stdout, arg0, /*max_element_count=*/4096)); + IREE_CHECK_OK(iree_hal_buffer_view_fprint( + stdout, arg0, /*max_element_count=*/4096, + iree_runtime_session_host_allocator(session))); IREE_CHECK_OK(iree_runtime_call_inputs_push_back_buffer_view(&call, arg0)); iree_hal_buffer_view_release(arg0); @@ -106,8 +107,9 @@ IREE_HAL_MEMORY_ACCESS_READ, IREE_HAL_BUFFER_USAGE_ALL, iree_make_byte_span((void*)arg1_data, sizeof(arg1_data)), iree_allocator_null(), &arg1)); - IREE_CHECK_OK( - iree_hal_buffer_view_fprint(stdout, arg1, /*max_element_count=*/4096)); + IREE_CHECK_OK(iree_hal_buffer_view_fprint( + stdout, arg1, /*max_element_count=*/4096, + iree_runtime_session_host_allocator(session))); IREE_CHECK_OK(iree_runtime_call_inputs_push_back_buffer_view(&call, arg1)); iree_hal_buffer_view_release(arg1); @@ -118,8 +120,9 @@ // -> tensor<4xf32> iree_hal_buffer_view_t* ret0 = NULL; IREE_CHECK_OK(iree_runtime_call_outputs_pop_front_buffer_view(&call, &ret0)); - IREE_CHECK_OK( - iree_hal_buffer_view_fprint(stdout, ret0, /*max_element_count=*/4096)); + IREE_CHECK_OK(iree_hal_buffer_view_fprint( + stdout, ret0, /*max_element_count=*/4096, + iree_runtime_session_host_allocator(session))); iree_hal_buffer_view_release(ret0); iree_runtime_call_deinitialize(&call);
diff --git a/iree/runtime/session.c b/iree/runtime/session.c index 8704654..8128879 100644 --- a/iree/runtime/session.c +++ b/iree/runtime/session.c
@@ -171,6 +171,16 @@ return iree_hal_device_allocator(device); } +IREE_API_EXPORT iree_status_t +iree_runtime_session_trim(iree_runtime_session_t* session) { + IREE_ASSERT_ARGUMENT(session); + IREE_TRACE_ZONE_BEGIN(z0); + iree_status_t status = iree_vm_context_notify( + iree_runtime_session_context(session), IREE_VM_SIGNAL_LOW_MEMORY); + IREE_TRACE_ZONE_END(z0); + return status; +} + IREE_API_EXPORT iree_status_t iree_runtime_session_append_module( iree_runtime_session_t* session, iree_vm_module_t* module) { IREE_ASSERT_ARGUMENT(session);
diff --git a/iree/runtime/session.h b/iree/runtime/session.h index 960d7fb..ec0500c 100644 --- a/iree/runtime/session.h +++ b/iree/runtime/session.h
@@ -128,6 +128,13 @@ IREE_API_EXPORT iree_hal_allocator_t* iree_runtime_session_device_allocator( const iree_runtime_session_t* session); +// Trims transient/cached resources used by the session. +// Upon resuming these resources may be expensive to rematerialize/reload and +// as such this should only be called when it is known the resources will not +// be needed soon. +IREE_API_EXPORT iree_status_t +iree_runtime_session_trim(iree_runtime_session_t* session); + // Appends the given |module| to the context. // The module will be retained by the context. //
diff --git a/iree/samples/custom_modules/dialect/custom-translate-main.cc b/iree/samples/custom_modules/dialect/custom-translate-main.cc index c32f145..6e2372a 100644 --- a/iree/samples/custom_modules/dialect/custom-translate-main.cc +++ b/iree/samples/custom_modules/dialect/custom-translate-main.cc
@@ -61,7 +61,7 @@ mlir::registerMlirTranslations(); mlir::iree_compiler::registerIreeTranslations(); // Make sure command line options are registered. - (void)mlir::iree_compiler::IREE::HAL::getTargetOptionsFromFlags(); + (void)mlir::iree_compiler::IREE::HAL::TargetOptions::FromFlags::get(); // Register MLIRContext command-line options like // -mlir-print-op-on-diagnostic.
diff --git a/iree/samples/dynamic_shapes/main.c b/iree/samples/dynamic_shapes/main.c index f82737c..3808a33 100644 --- a/iree/samples/dynamic_shapes/main.c +++ b/iree/samples/dynamic_shapes/main.c
@@ -21,7 +21,7 @@ // * debugging some apparent memory corruption with the stack-local value iree_status_t status = iree_ok_status(); if (iree_status_is_ok(status)) { - status = iree_hal_buffer_view_clone_heap_buffer( + status = iree_hal_buffer_view_allocate_buffer( iree_runtime_session_device_allocator(session), arg0_shape, IREE_ARRAYSIZE(arg0_shape), IREE_HAL_ELEMENT_TYPE_SINT_32, IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, @@ -43,16 +43,10 @@ status = iree_runtime_call_outputs_pop_front_buffer_view(&call, &buffer_view); } - iree_hal_buffer_mapping_t buffer_mapping; if (iree_status_is_ok(status)) { - status = iree_hal_buffer_map_range(iree_hal_buffer_view_buffer(buffer_view), - IREE_HAL_MEMORY_ACCESS_READ, 0, - IREE_WHOLE_BUFFER, &buffer_mapping); + status = iree_hal_buffer_read_data(iree_hal_buffer_view_buffer(buffer_view), + 0, out_result, sizeof(*out_result)); } - if (iree_status_is_ok(status)) { - *out_result = *buffer_mapping.contents.data; - } - iree_hal_buffer_unmap_range(&buffer_mapping); iree_hal_buffer_view_release(buffer_view); iree_runtime_call_deinitialize(&call); @@ -73,7 +67,7 @@ // * debugging some apparent memory corruption with the stack-local value iree_status_t status = iree_ok_status(); if (iree_status_is_ok(status)) { - status = iree_hal_buffer_view_clone_heap_buffer( + status = iree_hal_buffer_view_allocate_buffer( iree_runtime_session_device_allocator(session), arg0_shape, IREE_ARRAYSIZE(arg0_shape), IREE_HAL_ELEMENT_TYPE_SINT_32, IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, @@ -113,7 +107,7 @@ // * debugging some apparent memory corruption with the stack-local value iree_status_t status = iree_ok_status(); if (iree_status_is_ok(status)) { - status = iree_hal_buffer_view_clone_heap_buffer( + status = iree_hal_buffer_view_allocate_buffer( iree_runtime_session_device_allocator(session), arg0_shape, IREE_ARRAYSIZE(arg0_shape), IREE_HAL_ELEMENT_TYPE_SINT_32, IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, @@ -204,8 +198,10 @@ status = reduce_sum_2d(session, input, 6, &result_buffer_view); if (iree_status_is_ok(status)) { fprintf(stdout, "reduce_sum_2d([[1, 2, 3], [10, 20, 30]]): "); - status = iree_hal_buffer_view_fprint(stdout, result_buffer_view, - /*max_element_count=*/4096); + status = iree_hal_buffer_view_fprint( + stdout, result_buffer_view, + /*max_element_count=*/4096, + iree_runtime_session_host_allocator(session)); fprintf(stdout, "\n"); } iree_hal_buffer_view_release(result_buffer_view); @@ -219,8 +215,10 @@ if (iree_status_is_ok(status)) { fprintf(stdout, "reduce_sum_2d([[1, 2, 3], [10, 20, 30], [100, 200, 300]]): "); - status = iree_hal_buffer_view_fprint(stdout, result_buffer_view, - /*max_element_count=*/4096); + status = iree_hal_buffer_view_fprint( + stdout, result_buffer_view, + /*max_element_count=*/4096, + iree_runtime_session_host_allocator(session)); fprintf(stdout, "\n"); } iree_hal_buffer_view_release(result_buffer_view); @@ -233,8 +231,10 @@ status = add_one(session, input, 3, &result_buffer_view); if (iree_status_is_ok(status)) { fprintf(stdout, "add_one([1, 10, 100]): "); - status = iree_hal_buffer_view_fprint(stdout, result_buffer_view, - /*max_element_count=*/64); + status = iree_hal_buffer_view_fprint( + stdout, result_buffer_view, + /*max_element_count=*/64, + iree_runtime_session_host_allocator(session)); fprintf(stdout, "\n"); } iree_hal_buffer_view_release(result_buffer_view);
diff --git a/iree/samples/ops/dynamic-mhlo-dot.mlir b/iree/samples/ops/dynamic-mhlo-dot.mlir deleted file mode 100644 index 46dce0e..0000000 --- a/iree/samples/ops/dynamic-mhlo-dot.mlir +++ /dev/null
@@ -1,4 +0,0 @@ -func @dot(%lhs: tensor<?x?xf32>, %rhs: tensor<?x?xf32>) -> tensor<?x?xf32> { - %0 = "mhlo.dot"(%lhs, %rhs) : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32> - return %0 : tensor<?x?xf32> -}
diff --git a/iree/samples/ops/mhlo-dot.mlir b/iree/samples/ops/mhlo-dot.mlir deleted file mode 100644 index afd5dfb..0000000 --- a/iree/samples/ops/mhlo-dot.mlir +++ /dev/null
@@ -1,4 +0,0 @@ -func @dot(%lhs: tensor<32x1024xf32>, %rhs: tensor<1024x64xf32>) -> tensor<32x64xf32> { - %0 = "mhlo.dot"(%lhs, %rhs) : (tensor<32x1024xf32>, tensor<1024x64xf32>) -> tensor<32x64xf32> - return %0 : tensor<32x64xf32> -}
diff --git a/iree/samples/simple_embedding/simple_embedding.c b/iree/samples/simple_embedding/simple_embedding.c index db9f2b6..94f2dee 100644 --- a/iree/samples/simple_embedding/simple_embedding.c +++ b/iree/samples/simple_embedding/simple_embedding.c
@@ -7,6 +7,10 @@ // A example of setting up the HAL module to run simple pointwise array // multiplication with the device implemented by different backends via // create_sample_driver(). +// +// NOTE: this file does not properly handle error cases and will leak on +// failure. Applications that are just going to exit()/abort() on failure can +// probably get away with the same thing but really should prefer not to. #include <stdio.h> @@ -66,42 +70,29 @@ IREE_RETURN_IF_ERROR(iree_vm_context_resolve_function( context, iree_make_cstring_view(kMainFunctionName), &main_function)); - // Allocate buffers that can be mapped on the CPU and that can also be used - // on the device. Not all devices support this, but the ones we have now do. - const int kElementCount = 4; - iree_hal_buffer_t* arg0_buffer = NULL; - iree_hal_buffer_t* arg1_buffer = NULL; - iree_hal_memory_type_t input_memory_type = - IREE_HAL_MEMORY_TYPE_HOST_LOCAL | IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE; - IREE_RETURN_IF_ERROR(iree_hal_allocator_allocate_buffer( - iree_hal_device_allocator(device), input_memory_type, - IREE_HAL_BUFFER_USAGE_ALL, sizeof(float) * kElementCount, &arg0_buffer)); - IREE_RETURN_IF_ERROR(iree_hal_allocator_allocate_buffer( - iree_hal_device_allocator(device), input_memory_type, - IREE_HAL_BUFFER_USAGE_ALL, sizeof(float) * kElementCount, &arg1_buffer)); + // Initial buffer contents for 4 * 2 = 8. + const float kFloat4[] = {4.0f, 4.0f, 4.0f, 4.0f}; + const float kFloat2[] = {2.0f, 2.0f, 2.0f, 2.0f}; - // Populate initial values for 4 * 2 = 8. - const float kFloat4 = 4.0f; - const float kFloat2 = 2.0f; - IREE_RETURN_IF_ERROR(iree_hal_buffer_fill(arg0_buffer, 0, IREE_WHOLE_BUFFER, - &kFloat4, sizeof(float))); - IREE_RETURN_IF_ERROR(iree_hal_buffer_fill(arg1_buffer, 0, IREE_WHOLE_BUFFER, - &kFloat2, sizeof(float))); - - // Wrap buffers in shaped buffer views. - iree_hal_dim_t shape[1] = {kElementCount}; + // Allocate buffers in device-local memory so that if the device has an + // independent address space they live on the fast side of the fence. + iree_hal_dim_t shape[1] = {IREE_ARRAYSIZE(kFloat4)}; iree_hal_buffer_view_t* arg0_buffer_view = NULL; iree_hal_buffer_view_t* arg1_buffer_view = NULL; - IREE_RETURN_IF_ERROR(iree_hal_buffer_view_create( - arg0_buffer, shape, IREE_ARRAYSIZE(shape), IREE_HAL_ELEMENT_TYPE_FLOAT_32, - IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, iree_allocator_system(), - &arg0_buffer_view)); - IREE_RETURN_IF_ERROR(iree_hal_buffer_view_create( - arg1_buffer, shape, IREE_ARRAYSIZE(shape), IREE_HAL_ELEMENT_TYPE_FLOAT_32, - IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, iree_allocator_system(), - &arg1_buffer_view)); - iree_hal_buffer_release(arg0_buffer); - iree_hal_buffer_release(arg1_buffer); + IREE_RETURN_IF_ERROR(iree_hal_buffer_view_allocate_buffer( + iree_hal_device_allocator(device), shape, IREE_ARRAYSIZE(shape), + IREE_HAL_ELEMENT_TYPE_FLOAT_32, IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, + IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL | IREE_HAL_MEMORY_TYPE_HOST_VISIBLE, + IREE_HAL_BUFFER_USAGE_DISPATCH | IREE_HAL_BUFFER_USAGE_TRANSFER | + IREE_HAL_BUFFER_USAGE_MAPPING, + iree_make_const_byte_span(kFloat4, sizeof(kFloat4)), &arg0_buffer_view)); + IREE_RETURN_IF_ERROR(iree_hal_buffer_view_allocate_buffer( + iree_hal_device_allocator(device), shape, IREE_ARRAYSIZE(shape), + IREE_HAL_ELEMENT_TYPE_FLOAT_32, IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, + IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL | IREE_HAL_MEMORY_TYPE_HOST_VISIBLE, + IREE_HAL_BUFFER_USAGE_DISPATCH | IREE_HAL_BUFFER_USAGE_TRANSFER | + IREE_HAL_BUFFER_USAGE_MAPPING, + iree_make_const_byte_span(kFloat2, sizeof(kFloat2)), &arg1_buffer_view)); // Setup call inputs with our buffers. iree_vm_list_t* inputs = NULL; @@ -142,16 +133,15 @@ } // Read back the results and ensure we got the right values. - iree_hal_buffer_mapping_t mapped_memory; - IREE_RETURN_IF_ERROR(iree_hal_buffer_map_range( - iree_hal_buffer_view_buffer(ret_buffer_view), IREE_HAL_MEMORY_ACCESS_READ, - 0, IREE_WHOLE_BUFFER, &mapped_memory)); - for (int i = 0; i < mapped_memory.contents.data_length / sizeof(float); ++i) { - if (((const float*)mapped_memory.contents.data)[i] != 8.0f) { + float results[] = {0.0f, 0.0f, 0.0f, 0.0f}; + IREE_RETURN_IF_ERROR( + iree_hal_buffer_read_data(iree_hal_buffer_view_buffer(ret_buffer_view), 0, + results, sizeof(results))); + for (iree_host_size_t i = 0; i < IREE_ARRAYSIZE(results); ++i) { + if (results[i] != 8.0f) { return iree_make_status(IREE_STATUS_UNKNOWN, "result mismatches"); } } - iree_hal_buffer_unmap_range(&mapped_memory); iree_vm_list_release(inputs); iree_vm_list_release(outputs);
diff --git a/iree/samples/static_library/static_library_demo.c b/iree/samples/static_library/static_library_demo.c index 637d01b..79c802e 100644 --- a/iree/samples/static_library/static_library_demo.c +++ b/iree/samples/static_library/static_library_demo.c
@@ -127,7 +127,7 @@ iree_hal_memory_type_t input_memory_type = IREE_HAL_MEMORY_TYPE_HOST_LOCAL | IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE; if (iree_status_is_ok(status)) { - status = iree_hal_buffer_view_clone_heap_buffer( + status = iree_hal_buffer_view_allocate_buffer( iree_hal_device_allocator(device), shape, IREE_ARRAYSIZE(shape), IREE_HAL_ELEMENT_TYPE_FLOAT_32, IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, input_memory_type, IREE_HAL_BUFFER_USAGE_ALL, @@ -136,7 +136,7 @@ &arg0_buffer_view); } if (iree_status_is_ok(status)) { - status = iree_hal_buffer_view_clone_heap_buffer( + status = iree_hal_buffer_view_allocate_buffer( iree_hal_device_allocator(device), shape, IREE_ARRAYSIZE(shape), IREE_HAL_ELEMENT_TYPE_FLOAT_32, IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, input_memory_type, IREE_HAL_BUFFER_USAGE_ALL, @@ -171,31 +171,22 @@ } // Read back the results and ensure we got the right values. - iree_hal_buffer_mapping_t mapped_memory; - memset(&mapped_memory, 0, sizeof(mapped_memory)); + float results[] = {0.0f, 0.0f, 0.0f, 0.0f}; if (iree_status_is_ok(status)) { - status = iree_hal_buffer_map_range( - iree_hal_buffer_view_buffer(ret_buffer_view), - IREE_HAL_MEMORY_ACCESS_READ, 0, IREE_WHOLE_BUFFER, &mapped_memory); + status = + iree_hal_buffer_read_data(iree_hal_buffer_view_buffer(ret_buffer_view), + 0, results, sizeof(results)); } if (iree_status_is_ok(status)) { - if (mapped_memory.contents.data_length / sizeof(float) != kElementCount) { - status = iree_make_status(IREE_STATUS_UNKNOWN, - "result does not match element count "); - } - } - if (iree_status_is_ok(status)) { - const float* data = (const float*)mapped_memory.contents.data; - for (iree_host_size_t i = 0; - i < mapped_memory.contents.data_length / sizeof(float); ++i) { - if (data[i] != 8.0f) { + for (iree_host_size_t i = 0; i < IREE_ARRAYSIZE(results); ++i) { + if (results[i] != 8.0f) { status = iree_make_status(IREE_STATUS_UNKNOWN, "result mismatches"); + break; } } } // Cleanup call and buffers. - iree_hal_buffer_unmap_range(&mapped_memory); iree_hal_buffer_view_release(ret_buffer_view); iree_runtime_call_deinitialize(&call);
diff --git a/iree/samples/variables_and_state/main.c b/iree/samples/variables_and_state/main.c index 302ff4e..aa08ad0 100644 --- a/iree/samples/variables_and_state/main.c +++ b/iree/samples/variables_and_state/main.c
@@ -23,16 +23,10 @@ status = iree_runtime_call_outputs_pop_front_buffer_view(&call, &buffer_view); } - iree_hal_buffer_mapping_t buffer_mapping; if (iree_status_is_ok(status)) { - status = iree_hal_buffer_map_range(iree_hal_buffer_view_buffer(buffer_view), - IREE_HAL_MEMORY_ACCESS_READ, 0, - IREE_WHOLE_BUFFER, &buffer_mapping); + status = iree_hal_buffer_read_data(iree_hal_buffer_view_buffer(buffer_view), + 0, out_value, sizeof(*out_value)); } - if (iree_status_is_ok(status)) { - *out_value = *buffer_mapping.contents.data; - } - iree_hal_buffer_unmap_range(&buffer_mapping); iree_hal_buffer_view_release(buffer_view); iree_runtime_call_deinitialize(&call); @@ -52,7 +46,7 @@ // * debugging some apparent memory corruption with the stack-local value iree_status_t status = iree_ok_status(); if (iree_status_is_ok(status)) { - status = iree_hal_buffer_view_clone_heap_buffer( + status = iree_hal_buffer_view_allocate_buffer( iree_runtime_session_device_allocator(session), /*shape=*/NULL, /*shape_rank=*/0, IREE_HAL_ELEMENT_TYPE_SINT_32, IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, @@ -84,7 +78,7 @@ // * debugging some apparent memory corruption with the stack-local value iree_status_t status = iree_ok_status(); if (iree_status_is_ok(status)) { - status = iree_hal_buffer_view_clone_heap_buffer( + status = iree_hal_buffer_view_allocate_buffer( iree_runtime_session_device_allocator(session), /*shape=*/NULL, /*shape_rank=*/0, IREE_HAL_ELEMENT_TYPE_SINT_32, IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR,
diff --git a/iree/samples/vision/iree-run-mnist-module.c b/iree/samples/vision/iree-run-mnist-module.c index bfd9d7c..d525bdd 100644 --- a/iree/samples/vision/iree-run-mnist-module.c +++ b/iree/samples/vision/iree-run-mnist-module.c
@@ -73,23 +73,22 @@ // Read back the results. The output of the mnist model is a 1x10 prediction // confidence values for each digit in [0, 9]. - iree_hal_buffer_mapping_t mapped_memory; - IREE_RETURN_IF_ERROR(iree_hal_buffer_map_range( - iree_hal_buffer_view_buffer(ret_buffer_view), IREE_HAL_MEMORY_ACCESS_READ, - 0, IREE_WHOLE_BUFFER, &mapped_memory)); + float predictions[1 * 10] = {0.0f}; + IREE_RETURN_IF_ERROR( + iree_hal_buffer_read_data(iree_hal_buffer_view_buffer(ret_buffer_view), 0, + predictions, sizeof(predictions))); + iree_hal_buffer_view_release(ret_buffer_view); + + // Get the highest index from the output. float result_val = FLT_MIN; int result_idx = 0; - const float* data_ptr = (const float*)mapped_memory.contents.data; - for (int i = 0; i < mapped_memory.contents.data_length / sizeof(float); ++i) { - if (data_ptr[i] > result_val) { - result_val = data_ptr[i]; + for (iree_host_size_t i = 0; i < IREE_ARRAYSIZE(predictions); ++i) { + if (predictions[i] > result_val) { + result_val = predictions[i]; result_idx = i; } } - iree_hal_buffer_unmap_range(&mapped_memory); - // Get the highest index from the output. fprintf(stdout, "Detected number: %d\n", result_idx); - iree_hal_buffer_view_release(ret_buffer_view); iree_runtime_call_deinitialize(&call); iree_runtime_session_release(session);
diff --git a/iree/samples/vulkan/vulkan_inference_gui.cc b/iree/samples/vulkan/vulkan_inference_gui.cc index 4939b93..9d895a1 100644 --- a/iree/samples/vulkan/vulkan_inference_gui.cc +++ b/iree/samples/vulkan/vulkan_inference_gui.cc
@@ -71,8 +71,10 @@ } // Setup window + // clang-format off SDL_WindowFlags window_flags = (SDL_WindowFlags)( SDL_WINDOW_VULKAN | SDL_WINDOW_RESIZABLE | SDL_WINDOW_ALLOW_HIGHDPI); + // clang-format on SDL_Window* window = SDL_CreateWindow( "IREE Samples - Vulkan Inference GUI", SDL_WINDOWPOS_CENTERED, SDL_WINDOWPOS_CENTERED, 1280, 720, window_flags); @@ -363,8 +365,6 @@ constexpr int32_t kElementCount = 4; iree_hal_allocator_t* allocator = iree_hal_device_allocator(iree_vk_device); - iree_hal_buffer_t* input0_buffer = nullptr; - iree_hal_buffer_t* input1_buffer = nullptr; iree_hal_memory_type_t input_memory_type = static_cast<iree_hal_memory_type_t>( IREE_HAL_MEMORY_TYPE_HOST_LOCAL | @@ -372,33 +372,25 @@ iree_hal_buffer_usage_t input_buffer_usage = static_cast<iree_hal_buffer_usage_t>( IREE_HAL_BUFFER_USAGE_ALL | IREE_HAL_BUFFER_USAGE_CONSTANT); - IREE_CHECK_OK(iree_hal_allocator_allocate_buffer( - allocator, input_memory_type, input_buffer_usage, - sizeof(float) * kElementCount, &input0_buffer)); - IREE_CHECK_OK(iree_hal_allocator_allocate_buffer( - allocator, input_memory_type, input_buffer_usage, - sizeof(float) * kElementCount, &input1_buffer)); - IREE_CHECK_OK(iree_hal_buffer_write_data(input0_buffer, 0, &input_x, - sizeof(input_x))); - IREE_CHECK_OK(iree_hal_buffer_write_data(input1_buffer, 0, &input_y, - sizeof(input_y))); // Wrap input buffers in buffer views. iree_hal_buffer_view_t* input0_buffer_view = nullptr; iree_hal_buffer_view_t* input1_buffer_view = nullptr; - IREE_CHECK_OK(iree_hal_buffer_view_create( - input0_buffer, + IREE_CHECK_OK(iree_hal_buffer_view_allocate_buffer( + allocator, /*shape=*/&kElementCount, /*shape_rank=*/1, IREE_HAL_ELEMENT_TYPE_FLOAT_32, - IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, iree_allocator_system(), - &input0_buffer_view)); - IREE_CHECK_OK(iree_hal_buffer_view_create( - input1_buffer, + IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, input_memory_type, + input_buffer_usage, + iree_make_const_byte_span(&input_x, sizeof(input_x)), + iree_allocator_system(), &input0_buffer_view)); + IREE_CHECK_OK(iree_hal_buffer_view_allocate_buffer( + allocator, /*shape=*/&kElementCount, /*shape_rank=*/1, IREE_HAL_ELEMENT_TYPE_FLOAT_32, - IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, iree_allocator_system(), - &input1_buffer_view)); - iree_hal_buffer_release(input0_buffer); - iree_hal_buffer_release(input1_buffer); + IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, input_memory_type, + input_buffer_usage, + iree_make_const_byte_span(&input_y, sizeof(input_y)), + iree_allocator_system(), &input1_buffer_view)); // Marshal inputs through a VM variant list. // [wait_semaphore|wait_value|arg0|arg1|signal_semaphore|signal_value] vm::ref<iree_vm_list_t> inputs;
diff --git a/iree/task/api.c b/iree/task/api.c index e1f3e86..03cb5df 100644 --- a/iree/task/api.c +++ b/iree/task/api.c
@@ -34,8 +34,9 @@ "threads that would otherwise need to perform the syscalls during\n" "coordination."); +// TODO(benvanik): enable this when we use it - though hopefully we don't! IREE_FLAG( - int32_t, task_worker_local_memory, 64 * 1024, + int32_t, task_worker_local_memory, 0, // 64 * 1024, "Specifies the bytes of per-worker local memory allocated for use by\n" "dispatched tiles. Tiles may use less than this but will fail to dispatch\n" "if they require more. Conceptually it is like a stack reservation and\n"
diff --git a/iree/task/executor.c b/iree/task/executor.c index 54de4ed..70c494f 100644 --- a/iree/task/executor.c +++ b/iree/task/executor.c
@@ -211,6 +211,15 @@ } } +void iree_task_executor_trim(iree_task_executor_t* executor) { + // TODO(benvanik): figure out a good way to do this; the pools require that + // no tasks are in-flight to trim but our caller can't reliably make that + // guarantee. We'd need some global executor lock that we did here and + // on submit - or rework pools to not have this limitation. + // iree_task_pool_trim(&executor->fence_task_pool); + // iree_task_pool_trim(&executor->dispatch_task_pool); +} + iree_status_t iree_task_executor_acquire_fence(iree_task_executor_t* executor, iree_task_scope_t* scope, iree_task_fence_t** out_fence) {
diff --git a/iree/task/executor.h b/iree/task/executor.h index c191602..1c06d9a 100644 --- a/iree/task/executor.h +++ b/iree/task/executor.h
@@ -321,6 +321,9 @@ // Releases the given |executor| from the caller. void iree_task_executor_release(iree_task_executor_t* executor); +// Trims pools and caches used by the executor and its workers. +void iree_task_executor_trim(iree_task_executor_t* executor); + // Acquires a fence for the given |scope| from the executor fence pool. iree_status_t iree_task_executor_acquire_fence(iree_task_executor_t* executor, iree_task_scope_t* scope,
diff --git a/iree/testing/benchmark_full.cc b/iree/testing/benchmark_full.cc index 16a3371..c01abf0 100644 --- a/iree/testing/benchmark_full.cc +++ b/iree/testing/benchmark_full.cc
@@ -114,10 +114,11 @@ const iree_benchmark_def_t* benchmark_def) { std::string name_str(name.data, name.size); std::string prefixed_str = "BM_" + name_str; + iree_benchmark_def_t cloned_def = *benchmark_def; auto* instance = benchmark::RegisterBenchmark( prefixed_str.c_str(), - [name_str, benchmark_def](benchmark::State& state) -> void { - iree_benchmark_run(name_str.c_str(), benchmark_def, state); + [name_str, cloned_def](benchmark::State& state) -> void { + iree_benchmark_run(name_str.c_str(), &cloned_def, state); }); if (iree_all_bits_set(benchmark_def->flags,
diff --git a/iree/tools/init_iree_dialects.h b/iree/tools/init_iree_dialects.h index e4c8277..0aa7cb1 100644 --- a/iree/tools/init_iree_dialects.h +++ b/iree/tools/init_iree_dialects.h
@@ -47,6 +47,7 @@ // clang-format on IREE::LinalgExt::registerTiledOpInterfaceExternalModels(registry); + IREE::Util::registerUtilExternalModels(registry); registerCodegenInterfaces(registry); }
diff --git a/iree/tools/iree-e2e-matmul-test.c b/iree/tools/iree-e2e-matmul-test.c index 4d622a0..8646827 100644 --- a/iree/tools/iree-e2e-matmul-test.c +++ b/iree/tools/iree-e2e-matmul-test.c
@@ -66,9 +66,10 @@ return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, "buffer_view is not dense row major"); } - iree_hal_buffer_mapping_t mapping; + iree_hal_buffer_mapping_t mapping = {{0}}; IREE_RETURN_IF_ERROR(iree_hal_buffer_map_range( - iree_hal_buffer_view_buffer(buffer_view), IREE_HAL_MEMORY_ACCESS_READ, 0, + iree_hal_buffer_view_buffer(buffer_view), + IREE_HAL_MAPPING_MODE_PERSISTENT, IREE_HAL_MEMORY_ACCESS_READ, 0, IREE_WHOLE_BUFFER, &mapping)); *data = mapping.contents.data; return iree_ok_status(); @@ -478,21 +479,23 @@ iree_hal_buffer_view_element_type(src), iree_hal_buffer_view_encoding_type(src), IREE_HAL_MEMORY_TYPE_HOST_LOCAL | IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE, - IREE_HAL_BUFFER_USAGE_ALL, dst); + IREE_HAL_BUFFER_USAGE_ALL, iree_const_byte_span_empty(), dst); } // Performs a deep copy of |src| into |dst|. Takes care of allocating |dst|. static iree_status_t copy_buffer(iree_hal_allocator_t* hal_allocator, iree_hal_buffer_view_t* src, iree_hal_buffer_view_t** dst) { - iree_hal_buffer_mapping_t src_mapping; + // TODO(benvanik): change this to use iree_hal_buffer_copy_data. Or something. + // I can't understand what all this code is doing. + iree_hal_buffer_mapping_t src_mapping = {{0}}; IREE_RETURN_IF_ERROR(iree_hal_buffer_map_range( - iree_hal_buffer_view_buffer(src), IREE_HAL_MEMORY_ACCESS_READ, 0, - IREE_WHOLE_BUFFER, &src_mapping)); + iree_hal_buffer_view_buffer(src), IREE_HAL_MAPPING_MODE_PERSISTENT, + IREE_HAL_MEMORY_ACCESS_READ, 0, IREE_WHOLE_BUFFER, &src_mapping)); iree_const_byte_span_t src_span; src_span.data = src_mapping.contents.data; src_span.data_length = src_mapping.contents.data_length; - return iree_hal_buffer_view_clone_heap_buffer( + return iree_hal_buffer_view_allocate_buffer( hal_allocator, iree_hal_buffer_view_shape_dims(src), iree_hal_buffer_view_shape_rank(src), iree_hal_buffer_view_element_type(src), @@ -551,7 +554,8 @@ // linalg.matmul. We need to preserve the original test inputs to run the // reference matmul on and to use in test failure logs. iree_vm_list_t* copy_of_input_list = NULL; - copy_list_of_buffer_views(device_allocator, input_list, ©_of_input_list); + IREE_CHECK_OK(copy_list_of_buffer_views(device_allocator, input_list, + ©_of_input_list)); // Invoke the function to produce the actual result. iree_vm_list_t* output_list = NULL;
diff --git a/iree/tools/iree-run-mlir-main.cc b/iree/tools/iree-run-mlir-main.cc index 11e4841..44d44d4 100644 --- a/iree/tools/iree-run-mlir-main.cc +++ b/iree/tools/iree-run-mlir-main.cc
@@ -45,7 +45,6 @@ #include "iree/base/tracing.h" #include "iree/compiler/Dialect/HAL/Target/TargetBackend.h" #include "iree/compiler/Dialect/VM/Target/Bytecode/BytecodeModuleTarget.h" -#include "iree/compiler/Dialect/VM/Target/Bytecode/TranslationFlags.h" #include "iree/compiler/Dialect/VM/Target/init_targets.h" #include "iree/compiler/Translation/IREEVM.h" #include "iree/hal/api.h" @@ -163,7 +162,7 @@ IREE_TRACE_SCOPE(); out_target_backends->clear(); auto target_backends = - mlir::iree_compiler::IREE::HAL::getTargetOptionsFromFlags().targets; + mlir::iree_compiler::IREE::HAL::TargetOptions::FromFlags::get().targets; if (target_backends.empty()) { iree_allocator_t host_allocator = iree_allocator_system(); iree_hal_driver_info_t* driver_infos = NULL; @@ -220,7 +219,7 @@ } auto bytecode_options = - mlir::iree_compiler::IREE::VM::getBytecodeTargetOptionsFromFlags(); + mlir::iree_compiler::IREE::VM::BytecodeTargetOptions::FromFlags::get(); std::string binary_contents; llvm::raw_string_ostream binary_output(binary_contents); if (failed(mlir::iree_compiler::IREE::VM::translateModuleToBytecode( @@ -502,7 +501,7 @@ mlir::iree_compiler::registerIREEVMTranslationFlags(); mlir::registerLLVMDialectTranslation(registry); // Make sure command line options are registered. - (void)mlir::iree_compiler::IREE::HAL::getTargetOptionsFromFlags(); + (void)mlir::iree_compiler::IREE::HAL::TargetOptions::FromFlags::get(); // Register MLIRContext command-line options like // -mlir-print-op-on-diagnostic.
diff --git a/iree/tools/iree_translate_lib.cc b/iree/tools/iree_translate_lib.cc index c4ce48f..f12725f 100644 --- a/iree/tools/iree_translate_lib.cc +++ b/iree/tools/iree_translate_lib.cc
@@ -53,7 +53,7 @@ mlir::registerMlirTranslations(); mlir::iree_compiler::registerIreeTranslations(); // Make sure command line options are registered. - (void)mlir::iree_compiler::IREE::HAL::getTargetOptionsFromFlags(); + (void)mlir::iree_compiler::IREE::HAL::TargetOptions::FromFlags::get(); // Register MLIRContext command-line options like // -mlir-print-op-on-diagnostic.
diff --git a/iree/tools/utils/image_util.c b/iree/tools/utils/image_util.c index 61e1486..94aaab0 100644 --- a/iree/tools/utils/image_util.c +++ b/iree/tools/utils/image_util.c
@@ -148,11 +148,27 @@ return result; } +typedef struct iree_tools_utils_buffer_view_load_params_t { + const uint8_t* pixel_data; + iree_host_size_t pixel_data_length; + const float* input_range; + iree_host_size_t input_range_length; +} iree_tools_utils_buffer_view_load_params_t; +static iree_status_t iree_tools_utils_buffer_view_load_image_rescaled( + iree_hal_buffer_mapping_t* mapping, void* user_data) { + iree_tools_utils_buffer_view_load_params_t* params = + (iree_tools_utils_buffer_view_load_params_t*)user_data; + return iree_tools_utils_pixel_rescaled_to_buffer( + params->pixel_data, params->pixel_data_length, params->input_range, + params->input_range_length, (float*)mapping->contents.data); +} + iree_status_t iree_tools_utils_buffer_view_from_image_rescaled( const iree_string_view_t filename, const iree_hal_dim_t* shape, iree_host_size_t shape_rank, iree_hal_element_type_t element_type, iree_hal_allocator_t* allocator, const float* input_range, - iree_host_size_t range_length, iree_hal_buffer_view_t** out_buffer_view) { + iree_host_size_t input_range_length, + iree_hal_buffer_view_t** out_buffer_view) { IREE_TRACE_ZONE_BEGIN(z0); *out_buffer_view = NULL; if (element_type != IREE_HAL_ELEMENT_TYPE_FLOAT_32) { @@ -161,42 +177,36 @@ "element type should be f32"); } - iree_status_t result; + // Classic row-major image layout. + iree_hal_encoding_type_t encoding_type = + IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR; + + // Load pixel data from the file into a new host memory allocation (the only + // interface stb_image provides). A real application would want to use the + // generation callback to directly decode the image into the target mapped + // device buffer. uint8_t* pixel_data = NULL; - iree_hal_buffer_t* buffer = NULL; - iree_host_size_t buffer_length; - iree_host_size_t element_byte = - iree_hal_element_dense_byte_count(element_type); - iree_hal_buffer_mapping_t mapped_memory; - result = iree_tools_utils_load_pixel_data( - filename, shape, shape_rank, element_type, &pixel_data, &buffer_length); - if (iree_status_is_ok(result)) { - result = iree_hal_allocator_allocate_buffer( - allocator, - IREE_HAL_MEMORY_TYPE_HOST_LOCAL | IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE, - IREE_HAL_BUFFER_USAGE_ALL, element_byte * buffer_length, &buffer); - } - if (iree_status_is_ok(result)) { - result = iree_hal_buffer_map_range( - buffer, IREE_HAL_MEMORY_ACCESS_DISCARD_WRITE, 0, - element_byte * buffer_length, &mapped_memory); - } - if (iree_status_is_ok(result)) { - // Need to normalize to the expected input range. - result = iree_tools_utils_pixel_rescaled_to_buffer( - pixel_data, buffer_length, input_range, range_length, - (float*)mapped_memory.contents.data); - iree_hal_buffer_unmap_range(&mapped_memory); - } - if (iree_status_is_ok(result)) { - iree_hal_encoding_type_t encoding_type = - IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR; - result = iree_hal_buffer_view_create( - buffer, shape, shape_rank, element_type, encoding_type, - iree_hal_allocator_host_allocator(allocator), out_buffer_view); - } - iree_hal_buffer_release(buffer); + iree_host_size_t buffer_length = 0; + IREE_RETURN_AND_END_ZONE_IF_ERROR( + z0, iree_tools_utils_load_pixel_data(filename, shape, shape_rank, + element_type, &pixel_data, + &buffer_length)); + + iree_tools_utils_buffer_view_load_params_t params = { + .pixel_data = pixel_data, + .pixel_data_length = buffer_length, + .input_range = input_range, + .input_range_length = input_range_length, + }; + iree_status_t status = iree_hal_buffer_view_generate_buffer( + allocator, shape, shape_rank, element_type, encoding_type, + IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL | IREE_HAL_MEMORY_TYPE_HOST_VISIBLE, + IREE_HAL_BUFFER_USAGE_DISPATCH | IREE_HAL_BUFFER_USAGE_TRANSFER | + IREE_HAL_BUFFER_USAGE_MAPPING, + iree_tools_utils_buffer_view_load_image_rescaled, ¶ms, + out_buffer_view); + stbi_image_free(pixel_data); IREE_TRACE_ZONE_END(z0); - return result; + return status; }
diff --git a/iree/tools/utils/image_util.h b/iree/tools/utils/image_util.h index 6c420d2..fa2e709 100644 --- a/iree/tools/utils/image_util.h +++ b/iree/tools/utils/image_util.h
@@ -54,9 +54,10 @@ const iree_string_view_t filename, const iree_hal_dim_t* shape, iree_host_size_t shape_rank, iree_hal_element_type_t element_type, iree_hal_allocator_t* allocator, const float* input_range, - iree_host_size_t range_length, iree_hal_buffer_view_t** out_buffer_view); + iree_host_size_t input_range_length, + iree_hal_buffer_view_t** out_buffer_view); -// Normalize uint8_t |pixel data| of the size |buffer_length| to float buffer +// Normalize uint8_t |pixel_data| of the size |buffer_length| to float buffer // |out_buffer| with the range |input_range|. // // float32_x = (uint8_x - 127.5) / 127.5 * input_scale + input_offset, where @@ -65,8 +66,9 @@ // // |out_buffer| needs to be allocated before the call. iree_status_t iree_tools_utils_pixel_rescaled_to_buffer( - const uint8_t* pixel_data, iree_host_size_t buffer_length, - const float* input_range, iree_host_size_t range_length, float* out_buffer); + const uint8_t* pixel_data, iree_host_size_t pixel_count, + const float* input_range, iree_host_size_t input_range_length, + float* out_buffer); #if __cplusplus }
diff --git a/iree/tools/utils/trace_replay.c b/iree/tools/utils/trace_replay.c index c84c9d4..827ee84 100644 --- a/iree/tools/utils/trace_replay.c +++ b/iree/tools/utils/trace_replay.c
@@ -465,10 +465,10 @@ // contents: !!binary | // AACAPwAAAEAAAEBAAACAQA== // ``` -static iree_status_t iree_trace_replay_parse_hal_buffer( +static iree_status_t iree_trace_replay_parse_hal_buffer_contents( iree_trace_replay_t* replay, yaml_document_t* document, yaml_node_t* contents_node, iree_hal_element_type_t element_type, - iree_hal_buffer_t* buffer) { + iree_hal_buffer_mapping_t* mapping) { if (!contents_node) { // Empty contents = zero fill. return iree_ok_status(); @@ -480,23 +480,16 @@ iree_string_view_t value = iree_string_view_trim(iree_yaml_node_as_string(contents_node)); - iree_hal_buffer_mapping_t mapping; - IREE_RETURN_IF_ERROR( - iree_hal_buffer_map_range(buffer, IREE_HAL_MEMORY_ACCESS_DISCARD_WRITE, 0, - IREE_WHOLE_BUFFER, &mapping)); - iree_status_t status = iree_ok_status(); if (strcmp(contents_node->tag, "tag:yaml.org,2002:binary") == 0) { - status = iree_yaml_base64_decode(value, mapping.contents); + return iree_yaml_base64_decode(value, mapping->contents); } else if (strcmp(contents_node->tag, "tag:yaml.org,2002:str") == 0) { - status = - iree_hal_parse_buffer_elements(value, element_type, mapping.contents); + return iree_hal_parse_buffer_elements(value, element_type, + mapping->contents); } else { - status = iree_make_status( - IREE_STATUS_UNIMPLEMENTED, "(%zu): unimplemented buffer encoding '%s'", - contents_node->start_mark.line, contents_node->tag); + return iree_make_status(IREE_STATUS_UNIMPLEMENTED, + "(%zu): unimplemented buffer encoding '%s'", + contents_node->start_mark.line, contents_node->tag); } - iree_hal_buffer_unmap_range(&mapping); - return status; } // Writes an element of the given |element_type| with the given integral |value| @@ -596,59 +589,127 @@ } // Generates the destination |buffer| using the generator specified by -// |contents_generator_node|. +// |generator_node|. static iree_status_t iree_trace_replay_generate_hal_buffer( iree_trace_replay_t* replay, yaml_document_t* document, - yaml_node_t* contents_generator_node, iree_hal_element_type_t element_type, - iree_hal_buffer_t* buffer, iree_hal_dim_t* shape, - iree_host_size_t shape_size) { - if (!contents_generator_node) { - return iree_ok_status(); - } else if (contents_generator_node->type != YAML_SCALAR_NODE) { + yaml_node_t* generator_node, iree_hal_element_type_t element_type, + const iree_hal_dim_t* shape, iree_host_size_t shape_rank, + iree_hal_buffer_mapping_t* mapping) { + if (generator_node->type != YAML_SCALAR_NODE) { return iree_make_status( IREE_STATUS_INVALID_ARGUMENT, "(%zu): expected scalar node for buffer contents_generator", - contents_generator_node->start_mark.line); + generator_node->start_mark.line); } - - iree_hal_buffer_mapping_t mapping; - IREE_RETURN_IF_ERROR( - iree_hal_buffer_map_range(buffer, IREE_HAL_MEMORY_ACCESS_DISCARD_WRITE, 0, - IREE_WHOLE_BUFFER, &mapping)); - iree_status_t status = iree_ok_status(); - if (strcmp(contents_generator_node->tag, "!tag:iree:identity_matrix") == 0) { - if (shape_size == 2) { - iree_hal_dim_t inner_size = shape[shape_size - 1]; - iree_trace_replay_generate_identity_matrix(element_type, mapping.contents, - inner_size); + if (strcmp(generator_node->tag, "!tag:iree:identity_matrix") == 0) { + if (shape_rank == 2) { + iree_hal_dim_t inner_size = shape[shape_rank - 1]; + iree_trace_replay_generate_identity_matrix(element_type, + mapping->contents, inner_size); } else { - status = iree_make_status( + return iree_make_status( IREE_STATUS_INVALID_ARGUMENT, "the identity_matrix generator is only for 2D shapes (matrices)"); } - } else if (strcmp(contents_generator_node->tag, + } else if (strcmp(generator_node->tag, "!tag:iree:fully_specified_pseudorandom") == 0) { // To enable pseudorandom tests that are both reproducible and invariant // under reordering and filtering testcases, the seed is explicitly // passed as argument in the contents_generator tag. - iree_string_view_t seed_str = iree_string_view_trim( - iree_yaml_node_as_string(contents_generator_node)); - uint32_t seed; + iree_string_view_t seed_str = + iree_string_view_trim(iree_yaml_node_as_string(generator_node)); + uint32_t seed = 0; if (iree_string_view_atoi_uint32(seed_str, &seed)) { iree_trace_replay_generate_fully_specified_pseudorandom_buffer( - element_type, mapping.contents, seed); + element_type, mapping->contents, seed); } else { - status = iree_make_status(IREE_STATUS_INVALID_ARGUMENT, - "could not parse the seed argument ('%s') of " - "the fully_specified_pseudorandom tag", - seed_str.data); + return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, + "could not parse the seed argument ('%s') of " + "the fully_specified_pseudorandom tag", + seed_str.data); } } else { - status = iree_make_status( + return iree_make_status( IREE_STATUS_UNIMPLEMENTED, "(%zu): unimplemented buffer generator '%s'", - contents_generator_node->start_mark.line, contents_generator_node->tag); + generator_node->start_mark.line, generator_node->tag); } - iree_hal_buffer_unmap_range(&mapping); + return iree_ok_status(); +} + +typedef struct iree_trace_replay_generation_params_t { + iree_trace_replay_t* replay; + yaml_document_t* document; + yaml_node_t* contents_node; + yaml_node_t* generator_node; + iree_hal_element_type_t element_type; + const iree_hal_dim_t* shape; + iree_host_size_t shape_rank; +} iree_trace_replay_generation_params_t; + +static iree_status_t iree_trace_replay_generate_hal_buffer_callback( + iree_hal_buffer_mapping_t* mapping, void* user_data) { + iree_trace_replay_generation_params_t* params = + (iree_trace_replay_generation_params_t*)user_data; + if (params->contents_node) { + return iree_trace_replay_parse_hal_buffer_contents( + params->replay, params->document, params->contents_node, + params->element_type, mapping); + } else { + return iree_trace_replay_generate_hal_buffer( + params->replay, params->document, params->generator_node, + params->element_type, params->shape, params->shape_rank, mapping); + } +} + +// Parses a !hal.buffer and appends it to |target_list|. +// +// ```yaml +// shape: +// - 4 +// element_type: 553648160 +// ``` +static iree_status_t iree_trace_replay_parse_hal_buffer( + iree_trace_replay_t* replay, yaml_document_t* document, + yaml_node_t* value_node, iree_vm_list_t* target_list) { + yaml_node_t* shape_node = NULL; + IREE_RETURN_IF_ERROR(iree_yaml_mapping_try_find( + document, value_node, iree_make_cstring_view("shape"), &shape_node)); + iree_hal_dim_t shape[16]; + iree_host_size_t shape_rank = 0; + IREE_RETURN_IF_ERROR(iree_trace_replay_parse_hal_shape( + replay, document, shape_node, IREE_ARRAYSIZE(shape), shape, &shape_rank)); + + yaml_node_t* element_type_node = NULL; + IREE_RETURN_IF_ERROR(iree_yaml_mapping_find( + document, value_node, iree_make_cstring_view("element_type"), + &element_type_node)); + iree_hal_element_type_t element_type = IREE_HAL_ELEMENT_TYPE_NONE; + IREE_RETURN_IF_ERROR(iree_trace_replay_parse_hal_element_type( + replay, document, element_type_node, &element_type)); + + yaml_node_t* encoding_type_node = NULL; + IREE_RETURN_IF_ERROR(iree_yaml_mapping_try_find( + document, value_node, iree_make_cstring_view("encoding_type"), + &encoding_type_node)); + iree_hal_encoding_type_t encoding_type = + IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR; + IREE_RETURN_IF_ERROR(iree_trace_replay_parse_hal_encoding_type( + replay, document, encoding_type_node, &encoding_type)); + + iree_device_size_t allocation_size = 0; + IREE_RETURN_IF_ERROR(iree_hal_buffer_compute_view_size( + shape, shape_rank, element_type, encoding_type, &allocation_size)); + + iree_hal_buffer_t* buffer = NULL; + IREE_RETURN_IF_ERROR(iree_hal_allocator_allocate_buffer( + iree_hal_device_allocator(replay->device), + IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL | IREE_HAL_MEMORY_TYPE_HOST_VISIBLE, + IREE_HAL_BUFFER_USAGE_ALL, allocation_size, iree_const_byte_span_empty(), + &buffer)); + + iree_vm_ref_t buffer_ref = iree_hal_buffer_move_ref(buffer); + iree_status_t status = iree_vm_list_push_ref_move(target_list, &buffer_ref); + iree_vm_ref_release(&buffer_ref); return status; } @@ -694,59 +755,76 @@ document, value_node, iree_make_cstring_view("contents"), &contents_node)); - yaml_node_t* contents_generator_node = NULL; + yaml_node_t* generator_node = NULL; IREE_RETURN_IF_ERROR(iree_yaml_mapping_try_find( document, value_node, iree_make_cstring_view("contents_generator"), - &contents_generator_node)); + &generator_node)); - if (contents_node && contents_generator_node) { + iree_hal_buffer_view_t* buffer_view = NULL; + if (contents_node && generator_node) { return iree_make_status( IREE_STATUS_INVALID_ARGUMENT, "(%zu): cannot have both contents and contents_generator", - contents_generator_node->start_mark.line); + generator_node->start_mark.line); + } else if (contents_node || generator_node) { + iree_trace_replay_generation_params_t params = { + .replay = replay, + .document = document, + .contents_node = contents_node, + .generator_node = generator_node, + .element_type = element_type, + .shape = shape, + .shape_rank = shape_rank, + }; + IREE_RETURN_IF_ERROR(iree_hal_buffer_view_generate_buffer( + iree_hal_device_allocator(replay->device), shape, shape_rank, + element_type, encoding_type, + IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL | IREE_HAL_MEMORY_TYPE_HOST_VISIBLE, + IREE_HAL_BUFFER_USAGE_DISPATCH | IREE_HAL_BUFFER_USAGE_TRANSFER | + IREE_HAL_BUFFER_USAGE_MAPPING, + iree_trace_replay_generate_hal_buffer_callback, ¶ms, &buffer_view)); + } else { + IREE_RETURN_IF_ERROR(iree_hal_buffer_view_allocate_buffer( + iree_hal_device_allocator(replay->device), shape, shape_rank, + element_type, encoding_type, + IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL | IREE_HAL_MEMORY_TYPE_HOST_VISIBLE, + IREE_HAL_BUFFER_USAGE_DISPATCH | IREE_HAL_BUFFER_USAGE_TRANSFER | + IREE_HAL_BUFFER_USAGE_MAPPING, + iree_const_byte_span_empty(), &buffer_view)); } - iree_device_size_t allocation_size = 0; - IREE_RETURN_IF_ERROR(iree_hal_buffer_compute_view_size( - shape, shape_rank, element_type, IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, - &allocation_size)); - - iree_hal_buffer_t* buffer = NULL; - IREE_RETURN_IF_ERROR(iree_hal_allocator_allocate_buffer( - iree_hal_device_allocator(replay->device), - IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL | IREE_HAL_MEMORY_TYPE_HOST_VISIBLE, - IREE_HAL_BUFFER_USAGE_ALL, allocation_size, &buffer)); - iree_status_t status = iree_trace_replay_generate_hal_buffer( - replay, document, contents_generator_node, element_type, buffer, shape, - shape_rank); - if (!iree_status_is_ok(status)) { - iree_hal_buffer_release(buffer); - return status; - } - status = iree_trace_replay_parse_hal_buffer(replay, document, contents_node, - element_type, buffer); - if (!iree_status_is_ok(status)) { - iree_hal_buffer_release(buffer); - return status; - } - - iree_hal_buffer_view_t* buffer_view = NULL; - status = iree_hal_buffer_view_create(buffer, shape, shape_rank, element_type, - encoding_type, replay->host_allocator, - &buffer_view); - iree_hal_buffer_release(buffer); - IREE_RETURN_IF_ERROR(status); - iree_vm_ref_t buffer_view_ref = iree_hal_buffer_view_move_ref(buffer_view); - status = iree_vm_list_push_ref_move(target_list, &buffer_view_ref); + iree_status_t status = + iree_vm_list_push_ref_move(target_list, &buffer_view_ref); iree_vm_ref_release(&buffer_view_ref); return status; } +// Parses a !hal.buffer in tensor form and appends it to |target_list|. +// The tensor form is used to size and initialize the buffer but then the +// metadata is thrown away. +// +// ```yaml +// !!hal.buffer 4xf32=[0 1 2 3] +// ``` +static iree_status_t iree_trace_replay_parse_inline_hal_buffer( + iree_trace_replay_t* replay, yaml_document_t* document, + yaml_node_t* value_node, iree_vm_list_t* target_list) { + iree_hal_buffer_view_t* buffer_view = NULL; + IREE_RETURN_IF_ERROR(iree_hal_buffer_view_parse( + iree_yaml_node_as_string(value_node), + iree_hal_device_allocator(replay->device), &buffer_view)); + iree_vm_ref_t buffer_ref = + iree_hal_buffer_retain_ref(iree_hal_buffer_view_buffer(buffer_view)); + iree_status_t status = iree_vm_list_push_ref_move(target_list, &buffer_ref); + iree_hal_buffer_view_release(buffer_view); + return status; +} + // Parses a !hal.buffer_view in tensor form and appends it to |target_list|. // // ```yaml -// !tensor 4xf32=[0 1 2 3] +// !hal.buffer_view 4xf32=[0 1 2 3] // ``` static iree_status_t iree_trace_replay_parse_inline_hal_buffer_view( iree_trace_replay_t* replay, yaml_document_t* document, @@ -756,10 +834,7 @@ iree_yaml_node_as_string(value_node), iree_hal_device_allocator(replay->device), &buffer_view)); iree_vm_ref_t buffer_view_ref = iree_hal_buffer_view_move_ref(buffer_view); - iree_status_t status = - iree_vm_list_push_ref_move(target_list, &buffer_view_ref); - iree_vm_ref_release(&buffer_view_ref); - return status; + return iree_vm_list_push_ref_move(target_list, &buffer_view_ref); } // Parses a typed item from |value_node| and appends it to |target_list|. @@ -778,7 +853,10 @@ yaml_document_t* document, yaml_node_t* value_node, iree_vm_list_t* target_list) { - if (strcmp(value_node->tag, "!hal.buffer_view") == 0) { + if (strcmp(value_node->tag, "!hal.buffer") == 0) { + return iree_trace_replay_parse_inline_hal_buffer(replay, document, + value_node, target_list); + } else if (strcmp(value_node->tag, "!hal.buffer_view") == 0) { return iree_trace_replay_parse_inline_hal_buffer_view( replay, document, value_node, target_list); } @@ -796,6 +874,10 @@ } else if (iree_string_view_equal(type, iree_make_cstring_view("vm.list"))) { return iree_trace_replay_parse_vm_list(replay, document, value_node, target_list); + } else if (iree_string_view_equal(type, + iree_make_cstring_view("hal.buffer"))) { + return iree_trace_replay_parse_hal_buffer(replay, document, value_node, + target_list); } else if (iree_string_view_equal( type, iree_make_cstring_view("hal.buffer_view"))) { return iree_trace_replay_parse_hal_buffer_view(replay, document, value_node, @@ -819,7 +901,8 @@ return iree_ok_status(); } -static iree_status_t iree_trace_replay_print_item(iree_vm_variant_t* value); +static iree_status_t iree_trace_replay_print_item( + iree_vm_variant_t* value, iree_allocator_t host_allocator); static iree_status_t iree_trace_replay_print_scalar(iree_vm_variant_t* value) { switch (value->type.value_type) { @@ -848,35 +931,34 @@ return iree_ok_status(); } -static iree_status_t iree_trace_replay_print_vm_list(iree_vm_list_t* list) { +static iree_status_t iree_trace_replay_print_vm_list( + iree_vm_list_t* list, iree_allocator_t host_allocator) { for (iree_host_size_t i = 0; i < iree_vm_list_size(list); ++i) { iree_vm_variant_t variant = iree_vm_variant_empty(); IREE_RETURN_IF_ERROR(iree_vm_list_get_variant(list, i, &variant), "variant %zu not present", i); - IREE_RETURN_IF_ERROR(iree_trace_replay_print_item(&variant)); + IREE_RETURN_IF_ERROR( + iree_trace_replay_print_item(&variant, host_allocator)); fprintf(stdout, "\n"); } return iree_ok_status(); } -static iree_status_t iree_trace_replay_print_hal_buffer_view( - iree_hal_buffer_view_t* buffer_view) { - return iree_hal_buffer_view_fprint(stdout, buffer_view, - /*max_element_count=*/1024); -} - -static iree_status_t iree_trace_replay_print_item(iree_vm_variant_t* value) { +static iree_status_t iree_trace_replay_print_item( + iree_vm_variant_t* value, iree_allocator_t host_allocator) { if (iree_vm_variant_is_value(*value)) { IREE_RETURN_IF_ERROR(iree_trace_replay_print_scalar(value)); } else if (iree_vm_variant_is_ref(*value)) { if (iree_hal_buffer_view_isa(value->ref)) { iree_hal_buffer_view_t* buffer_view = iree_hal_buffer_view_deref(value->ref); - IREE_RETURN_IF_ERROR( - iree_trace_replay_print_hal_buffer_view(buffer_view)); + IREE_RETURN_IF_ERROR(iree_hal_buffer_view_fprint( + stdout, buffer_view, + /*max_element_count=*/1024, host_allocator)); } else if (iree_vm_list_isa(value->ref)) { iree_vm_list_t* list = iree_vm_list_deref(value->ref); - IREE_RETURN_IF_ERROR(iree_trace_replay_print_vm_list(list)); + IREE_RETURN_IF_ERROR( + iree_trace_replay_print_vm_list(list, host_allocator)); } else { // TODO(benvanik): a way for ref types to describe themselves. fprintf(stdout, "(no printer)"); @@ -985,7 +1067,8 @@ // Print the outputs. if (iree_status_is_ok(status)) { - status = iree_trace_replay_print_vm_list(output_list); + status = + iree_trace_replay_print_vm_list(output_list, replay->host_allocator); } iree_vm_list_release(output_list);
diff --git a/iree/tools/utils/vm_util.cc b/iree/tools/utils/vm_util.cc index f668aee..39e1331 100644 --- a/iree/tools/utils/vm_util.cc +++ b/iree/tools/utils/vm_util.cc
@@ -84,14 +84,27 @@ bool has_x = iree_string_view_find_char(input_view, 'x', 0) != IREE_STRING_VIEW_NPOS; if (has_equal || has_x) { - // Buffer view (either just a shape or a shape=value). + // Buffer view (either just a shape or a shape=value) or buffer. + bool is_storage_reference = iree_string_view_consume_prefix( + &input_view, iree_make_cstring_view("&")); iree_hal_buffer_view_t* buffer_view = nullptr; IREE_RETURN_IF_ERROR( iree_hal_buffer_view_parse(input_view, allocator, &buffer_view), "parsing value '%.*s'", (int)input_view.size, input_view.data); - auto buffer_view_ref = iree_hal_buffer_view_move_ref(buffer_view); - IREE_RETURN_IF_ERROR( - iree_vm_list_push_ref_move(variant_list.get(), &buffer_view_ref)); + if (is_storage_reference) { + // Storage buffer reference; just take the storage for the buffer view - + // it'll still have whatever contents were specified (or 0) but we'll + // discard the metadata. + auto buffer_ref = iree_hal_buffer_retain_ref( + iree_hal_buffer_view_buffer(buffer_view)); + iree_hal_buffer_view_release(buffer_view); + IREE_RETURN_IF_ERROR( + iree_vm_list_push_ref_move(variant_list.get(), &buffer_ref)); + } else { + auto buffer_view_ref = iree_hal_buffer_view_move_ref(buffer_view); + IREE_RETURN_IF_ERROR( + iree_vm_list_push_ref_move(variant_list.get(), &buffer_view_ref)); + } } else { // Scalar. bool has_dot = iree_string_view_find_char(input_view, '.', 0) !=
diff --git a/iree/tools/utils/vm_util_test.cc b/iree/tools/utils/vm_util_test.cc index 29df93b..d0477be 100644 --- a/iree/tools/utils/vm_util_test.cc +++ b/iree/tools/utils/vm_util_test.cc
@@ -38,6 +38,18 @@ }; TEST_F(VmUtilTest, ParsePrintBuffer) { + std::string buf_string = "&2x2xi32=[42 43][44 45]"; + vm::ref<iree_vm_list_t> variant_list; + IREE_ASSERT_OK(ParseToVariantList( + allocator_, std::vector<std::string>{buf_string}, &variant_list)); + std::stringstream os; + IREE_ASSERT_OK(PrintVariantList(variant_list.get(), &os)); + // TODO(benvanik): add a !hal.buffer printer. + EXPECT_EQ(os.str(), + std::string("result[0]: hal.buffer\n") + "(no printer)" + "\n"); +} + +TEST_F(VmUtilTest, ParsePrintBufferView) { std::string buf_string = "2x2xi32=[42 43][44 45]"; vm::ref<iree_vm_list_t> variant_list; IREE_ASSERT_OK(ParseToVariantList( @@ -58,7 +70,7 @@ EXPECT_EQ(os.str(), std::string("result[0]: i32=") + input_string + "\n"); } -TEST_F(VmUtilTest, ParsePrintRank0Buffer) { +TEST_F(VmUtilTest, ParsePrintRank0BufferView) { std::string buf_string = "i32=42"; vm::ref<iree_vm_list_t> variant_list; IREE_ASSERT_OK(ParseToVariantList( @@ -69,7 +81,7 @@ std::string("result[0]: hal.buffer_view\n") + buf_string + "\n"); } -TEST_F(VmUtilTest, ParsePrintMultipleBuffers) { +TEST_F(VmUtilTest, ParsePrintMultipleBufferViews) { std::string buf_string1 = "2x2xi32=[42 43][44 45]"; std::string buf_string2 = "2x3xf64=[1 2 3][4 5 6]"; vm::ref<iree_vm_list_t> variant_list;
diff --git a/iree/vm/bytecode_module.c b/iree/vm/bytecode_module.c index af15527..f64ac0f 100644 --- a/iree/vm/bytecode_module.c +++ b/iree/vm/bytecode_module.c
@@ -788,6 +788,11 @@ return iree_ok_status(); } +static iree_status_t IREE_API_PTR iree_vm_bytecode_module_notify( + void* self, iree_vm_module_state_t* module_state, iree_vm_signal_t signal) { + return iree_ok_status(); +} + static iree_status_t iree_vm_bytecode_module_begin_call( void* self, iree_vm_stack_t* stack, const iree_vm_function_call_t* call, iree_vm_execution_result_t* out_result) { @@ -914,6 +919,7 @@ module->interface.alloc_state = iree_vm_bytecode_module_alloc_state; module->interface.free_state = iree_vm_bytecode_module_free_state; module->interface.resolve_import = iree_vm_bytecode_module_resolve_import; + module->interface.notify = iree_vm_bytecode_module_notify; module->interface.begin_call = iree_vm_bytecode_module_begin_call; module->interface.get_function_reflection_attr = iree_vm_bytecode_module_get_function_reflection_attr;
diff --git a/iree/vm/context.c b/iree/vm/context.c index 4fc7938..e060dd6 100644 --- a/iree/vm/context.c +++ b/iree/vm/context.c
@@ -487,3 +487,127 @@ (int)module_name.size, module_name.data, (int)full_name.size, full_name.data); } + +// Calls the '__notify(i32)' function in |module|, if present. +static iree_status_t iree_vm_context_call_module_notify( + iree_vm_stack_t* stack, iree_vm_module_t* module, + iree_vm_module_state_t* module_state, iree_vm_signal_t signal) { + // Single i32 argument with the signal number. + uint32_t signal_arg = (uint32_t)signal; + iree_vm_function_call_t call; + memset(&call, 0, sizeof(call)); + call.arguments = iree_make_byte_span(&signal_arg, sizeof(signal_arg)); + + // Try to find the function. Modules are not required to export it. + iree_status_t status = iree_vm_module_lookup_function_by_name( + module, IREE_VM_FUNCTION_LINKAGE_EXPORT, + iree_make_cstring_view("__notify"), &call.function); + if (iree_status_is_not_found(status)) { + // Function doesn't exist; that's ok as this was an optional call. + return iree_status_ignore(status); + } else if (!iree_status_is_ok(status)) { + // Failed during trim. + return status; + } + + // Call the resolved function. + iree_vm_execution_result_t result; + status = module->begin_call(module->self, stack, &call, &result); + if (!iree_status_is_ok(status)) { + status = IREE_VM_STACK_ANNOTATE_BACKTRACE_IF_ENABLED(stack, status); + } + + // TODO(benvanik): ensure completed synchronously. + + return status; +} + +// Calls the module notify methods in registration order. +static iree_status_t iree_vm_context_notify_forward(iree_vm_stack_t* stack, + iree_vm_context_t* context, + iree_vm_signal_t signal) { + IREE_TRACE_ZONE_BEGIN(z0); + iree_status_t status = iree_ok_status(); + for (iree_host_size_t i = 0; i < context->list.count; ++i) { + iree_vm_module_t* module = context->list.modules[i]; + iree_vm_module_state_t* module_state = context->list.module_states[i]; + + // Call the module internal interface notify method. + // This handles the resources owned by the module implementation itself + // such as JITed binaries or other module infrastructure. + status = module->notify(module->self, module_state, signal); + if (!iree_status_is_ok(status)) break; + + // Call the user-level notify method. + // This may new use the reallocated resources from the module internal + // implementation above. + status = + iree_vm_context_call_module_notify(stack, module, module_state, signal); + if (!iree_status_is_ok(status)) break; + } + IREE_TRACE_ZONE_END(z0); + return status; +} + +// Calls the module notify methods in reverse registration order. +static iree_status_t iree_vm_context_notify_reverse(iree_vm_stack_t* stack, + iree_vm_context_t* context, + iree_vm_signal_t signal) { + IREE_TRACE_ZONE_BEGIN(z0); + iree_status_t status = iree_ok_status(); + for (int i = (int)context->list.count - 1; i >= 0; --i) { + iree_vm_module_t* module = context->list.modules[i]; + iree_vm_module_state_t* module_state = context->list.module_states[i]; + + // Call the user-level notify method first. + // This allows users to drop any state that they can rematerialize and + // return the resources to pools/caches to be trimmed below. + status = + iree_vm_context_call_module_notify(stack, module, module_state, signal); + if (!iree_status_is_ok(status)) break; + + // Call the module internal interface notify method. + // This handles the resources owned by the module implementation itself + // such as JITed binaries or other module infrastructure. Since we've + // already called the user-level function we likely have all of the + // resources that could be returned to pools there for this to reclaim. + status = module->notify(module->self, module_state, signal); + if (!iree_status_is_ok(status)) break; + } + IREE_TRACE_ZONE_END(z0); + return status; +} + +IREE_API_EXPORT iree_status_t iree_vm_context_notify(iree_vm_context_t* context, + iree_vm_signal_t signal) { + IREE_TRACE_ZONE_BEGIN(z0); + IREE_TRACE_ZONE_APPEND_VALUE(z0, (uint64_t)signal); + + // VM stack used to call into module __init methods. + IREE_VM_INLINE_STACK_INITIALIZE( + stack, + context->flags & IREE_VM_CONTEXT_FLAG_TRACE_EXECUTION + ? IREE_VM_INVOCATION_FLAG_TRACE_EXECUTION + : IREE_VM_INVOCATION_FLAG_NONE, + iree_vm_context_state_resolver(context), context->allocator); + + // Resumes are walked forward while suspends are walked backward. + // This follows the expected construction/destruction pattern where for + // example on suspend one would walk user modules to release resources back + // to system module pools before the system modules then clean up the pools. + iree_status_t status = iree_ok_status(); + switch (signal) { + default: + case IREE_VM_SIGNAL_RESUME: + status = iree_vm_context_notify_forward(stack, context, signal); + break; + case IREE_VM_SIGNAL_SUSPEND: + case IREE_VM_SIGNAL_LOW_MEMORY: + status = iree_vm_context_notify_reverse(stack, context, signal); + break; + } + + iree_vm_stack_deinitialize(stack); + IREE_TRACE_ZONE_END(z0); + return status; +}
diff --git a/iree/vm/context.h b/iree/vm/context.h index 71c36d3..b58bca6 100644 --- a/iree/vm/context.h +++ b/iree/vm/context.h
@@ -106,6 +106,10 @@ const iree_vm_context_t* context, iree_string_view_t full_name, iree_vm_function_t* out_function); +// Notifies all modules in the context of a system signal. +IREE_API_EXPORT iree_status_t iree_vm_context_notify(iree_vm_context_t* context, + iree_vm_signal_t signal); + #ifdef __cplusplus } // extern "C" #endif // __cplusplus
diff --git a/iree/vm/list.c b/iree/vm/list.c index a373d74..f1046cc 100644 --- a/iree/vm/list.c +++ b/iree/vm/list.c
@@ -11,6 +11,8 @@ #include <stdint.h> #include <string.h> +#include "iree/base/tracing.h" + // Size of each iree_vm_value_type_t in bytes. static const iree_host_size_t kValueTypeSizes[7] = { 0, // IREE_VM_VALUE_TYPE_NONE @@ -114,6 +116,8 @@ IREE_API_EXPORT iree_status_t iree_vm_list_initialize( iree_byte_span_t storage, const iree_vm_type_def_t* element_type, iree_host_size_t capacity, iree_vm_list_t** out_list) { + IREE_TRACE_ZONE_BEGIN(z0); + iree_vm_list_storage_mode_t storage_mode = IREE_VM_LIST_STORAGE_MODE_VARIANT; iree_host_size_t element_size = sizeof(iree_vm_variant_t); if (element_type) { @@ -151,22 +155,30 @@ list->storage = storage.data + storage_offset; *out_list = list; + IREE_TRACE_ZONE_END(z0); return iree_ok_status(); } IREE_API_EXPORT void iree_vm_list_deinitialize(iree_vm_list_t* list) { IREE_ASSERT_ARGUMENT(list); + IREE_TRACE_ZONE_BEGIN(z0); + iree_atomic_ref_count_abort_if_uses(&list->ref_object.counter); iree_vm_list_reset_range(list, 0, list->count); list->count = 0; + + IREE_TRACE_ZONE_END(z0); } IREE_API_EXPORT iree_status_t iree_vm_list_create( const iree_vm_type_def_t* element_type, iree_host_size_t initial_capacity, iree_allocator_t allocator, iree_vm_list_t** out_list) { + IREE_ASSERT_ARGUMENT(out_list); + IREE_TRACE_ZONE_BEGIN(z0); + iree_vm_list_t* list = NULL; - IREE_RETURN_IF_ERROR( - iree_allocator_malloc(allocator, sizeof(*list), (void**)&list)); + IREE_RETURN_AND_END_ZONE_IF_ERROR( + z0, iree_allocator_malloc(allocator, sizeof(*list), (void**)&list)); memset(list, 0, sizeof(*list)); iree_atomic_ref_count_init(&list->ref_object.counter); list->allocator = allocator; @@ -186,20 +198,25 @@ } iree_status_t status = iree_vm_list_reserve(list, initial_capacity); - if (!iree_status_is_ok(status)) { - iree_allocator_free(allocator, list); - return status; - } - *out_list = list; - return iree_ok_status(); + if (iree_status_is_ok(status)) { + *out_list = list; + } else { + iree_allocator_free(allocator, list); + } + IREE_TRACE_ZONE_END(z0); + return status; } static void iree_vm_list_destroy(void* ptr) { + IREE_TRACE_ZONE_BEGIN(z0); + iree_vm_list_t* list = (iree_vm_list_t*)ptr; iree_vm_list_reset_range(list, 0, list->count); iree_allocator_free(list->allocator, list->storage); iree_allocator_free(list->allocator, list); + + IREE_TRACE_ZONE_END(z0); } IREE_API_EXPORT void iree_vm_list_retain(iree_vm_list_t* list) {
diff --git a/iree/vm/module.h b/iree/vm/module.h index f164527..832d5aa 100644 --- a/iree/vm/module.h +++ b/iree/vm/module.h
@@ -24,11 +24,9 @@ typedef struct iree_vm_stack_t iree_vm_stack_t; typedef struct iree_vm_stack_frame_t iree_vm_stack_frame_t; -// An opaque offset into a source map that a source resolver can calculate. -// Do not assume that iree_vm_source_offset_t+1 means the next byte offset as -// backends are free to treat these as everything from pointers to machine code -// to hash codes. -typedef int64_t iree_vm_source_offset_t; +//===----------------------------------------------------------------------===// +// Module / function reflection +//===----------------------------------------------------------------------===// // A key-value pair of module/function reflection information. typedef struct iree_vm_reflection_attr_t { @@ -36,24 +34,6 @@ iree_string_view_t value; } iree_vm_reflection_attr_t; -// A variable-length list of registers. -// -// This structure is an overlay for the bytecode that is serialized in a -// matching format, though it can be stack allocated as needed. -// -// TODO(benvanik): this should be made private to the bytecode module, but is -// used for toll-free variadic argument lists here. We could just define an -// identical structure (and static_assert) to at least rename it to something -// sensible (iree_vm_segment_size_list_t). -typedef struct iree_vm_register_list_t { - uint16_t size; - uint16_t registers[]; -} iree_vm_register_list_t; -static_assert(iree_alignof(iree_vm_register_list_t) == 2, - "expecting byte alignment (to avoid padding)"); -static_assert(offsetof(iree_vm_register_list_t, registers) == 2, - "expect no padding in the struct"); - // Describes the type of a function reference. typedef enum iree_vm_function_linkage_e { // Function is internal to the module and may not be reflectable. @@ -136,6 +116,28 @@ // VM functions and accessing this state. typedef struct iree_vm_module_state_t iree_vm_module_state_t; +//===----------------------------------------------------------------------===// +// Function calls and coroutines +//===----------------------------------------------------------------------===// + +// A variable-length list of registers. +// +// This structure is an overlay for the bytecode that is serialized in a +// matching format, though it can be stack allocated as needed. +// +// TODO(benvanik): this should be made private to the bytecode module, but is +// used for toll-free variadic argument lists here. We could just define an +// identical structure (and static_assert) to at least rename it to something +// sensible (iree_vm_segment_size_list_t). +typedef struct iree_vm_register_list_t { + uint16_t size; + uint16_t registers[]; +} iree_vm_register_list_t; +static_assert(iree_alignof(iree_vm_register_list_t) == 2, + "expecting byte alignment (to avoid padding)"); +static_assert(offsetof(iree_vm_register_list_t, registers) == 2, + "expect no padding in the struct"); + // Function call data. // // Arguments and results are encoded following a standard format shared across @@ -261,6 +263,16 @@ int reserved; } iree_vm_execution_result_t; +//===----------------------------------------------------------------------===// +// Source locations +//===----------------------------------------------------------------------===// + +// An opaque offset into a source map that a source resolver can calculate. +// Do not assume that iree_vm_source_offset_t+1 means the next byte offset as +// backends are free to treat these as everything from pointers to machine code +// to hash codes. +typedef int64_t iree_vm_source_offset_t; + // Controls how source locations are formatted into strings. enum iree_vm_source_location_format_flag_bits_e { IREE_VM_SOURCE_LOCATION_FORMAT_FLAG_NONE = 0u, @@ -290,6 +302,34 @@ iree_vm_source_location_format_flags_t flags, iree_string_builder_t* builder); +//===----------------------------------------------------------------------===// +// iree_vm_module_t +//===----------------------------------------------------------------------===// + +// Indicates an event that can be signaled in modules from the hosting program. +typedef enum iree_vm_signal_e { + // Program is resuming from a suspended state. + // Modules may reallocate memory for pools and caches. + // + // Modules are walked in registration order (A->B->C). + IREE_VM_SIGNAL_RESUME = 0, + + // Program is entering a suspended state. + // Modules should drop any transient memory that is possible to reallocate + // upon resume. + // + // Modules are walked in reverse registration order (C->B->A). + IREE_VM_SIGNAL_SUSPEND = 1, + + // Program has received a low memory alert. + // Modules must aggressively drop all possible memory even if expensive to + // rematerialize it. On some platforms this is sent as a threat that if + // sufficient memory is not unwired/freed ASAP the process will be killed. + // + // Modules are walked in reverse registration order (C->B->A). + IREE_VM_SIGNAL_LOW_MEMORY = 2, +} iree_vm_signal_t; + // Defines an interface that can be used to reflect and execute functions on a // module. // @@ -349,6 +389,11 @@ iree_host_size_t ordinal, const iree_vm_function_t* function, const iree_vm_function_signature_t* signature); + // Notifies the module of a system signal. + iree_status_t(IREE_API_PTR* notify)(void* self, + iree_vm_module_state_t* module_state, + iree_vm_signal_t signal); + // Begins a function call with the given |call| arguments. // Execution may yield in the case of asynchronous code and require one or // more calls to the resume method to complete.
diff --git a/iree/vm/native_module.c b/iree/vm/native_module.c index e941089..c45a75c 100644 --- a/iree/vm/native_module.c +++ b/iree/vm/native_module.c
@@ -270,6 +270,15 @@ "native module does not support imports"); } +static iree_status_t IREE_API_PTR iree_vm_native_module_notify( + void* self, iree_vm_module_state_t* module_state, iree_vm_signal_t signal) { + iree_vm_native_module_t* module = (iree_vm_native_module_t*)self; + if (module->user_interface.notify) { + return module->user_interface.notify(module->self, module_state, signal); + } + return iree_ok_status(); +} + static iree_status_t IREE_API_PTR iree_vm_native_module_begin_call( void* self, iree_vm_stack_t* stack, const iree_vm_function_call_t* call, iree_vm_execution_result_t* out_result) { @@ -428,6 +437,7 @@ module->base_interface.alloc_state = iree_vm_native_module_alloc_state; module->base_interface.free_state = iree_vm_native_module_free_state; module->base_interface.resolve_import = iree_vm_native_module_resolve_import; + module->base_interface.notify = iree_vm_native_module_notify; module->base_interface.begin_call = iree_vm_native_module_begin_call; module->base_interface.resume_call = iree_vm_native_module_resume_call;
diff --git a/iree/vm/native_module_cc.h b/iree/vm/native_module_cc.h index df1e3b9..015fdb8 100644 --- a/iree/vm/native_module_cc.h +++ b/iree/vm/native_module_cc.h
@@ -78,6 +78,7 @@ interface_.alloc_state = NativeModule::ModuleAllocState; interface_.free_state = NativeModule::ModuleFreeState; interface_.resolve_import = NativeModule::ModuleResolveImport; + interface_.notify = NativeModule::ModuleNotify; interface_.begin_call = NativeModule::ModuleBeginCall; } @@ -91,6 +92,11 @@ virtual StatusOr<std::unique_ptr<State>> CreateState( iree_allocator_t allocator) = 0; + // Notifies the module a signal has been raised. + virtual Status Notify(State* state, iree_vm_signal_t signal) { + return OkStatus(); + } + private: static NativeModule* FromModulePointer(void* self) { return reinterpret_cast<NativeModule*>(self); @@ -201,6 +207,13 @@ "C++ API does not support imports"); } + static iree_status_t ModuleNotify(void* self, + iree_vm_module_state_t* module_state, + iree_vm_signal_t signal) { + auto* module = FromModulePointer(self); + return module->Notify(FromStatePointer(module_state), signal); + } + static iree_status_t ModuleBeginCall(void* self, iree_vm_stack_t* stack, const iree_vm_function_call_t* call, iree_vm_execution_result_t* out_result) {
diff --git a/llvm-external-projects/iree-compiler-api/BUILD.bazel b/llvm-external-projects/iree-compiler-api/BUILD.bazel index 94afcee..f103ced 100644 --- a/llvm-external-projects/iree-compiler-api/BUILD.bazel +++ b/llvm-external-projects/iree-compiler-api/BUILD.bazel
@@ -69,6 +69,7 @@ "//iree/compiler/InputConversion/MHLO", "//iree/compiler/InputConversion/TOSA", "//iree/compiler/Translation:IREEVM", + "//iree/compiler/Utils", "//iree/tools:init_targets", "//iree/tools:iree_translate_lib", "@llvm-project//lld:COFF",
diff --git a/llvm-external-projects/iree-compiler-api/build_tools/smoketest.py b/llvm-external-projects/iree-compiler-api/build_tools/smoketest.py index d5e407f..385451c 100644 --- a/llvm-external-projects/iree-compiler-api/build_tools/smoketest.py +++ b/llvm-external-projects/iree-compiler-api/build_tools/smoketest.py
@@ -12,7 +12,7 @@ from iree.compiler.dialects import arith from iree.compiler.dialects import chlo from iree.compiler.dialects import mhlo -from iree.compiler.dialects import iree as iree_dialect +from iree.compiler.dialects import iree_input from iree.compiler.dialects import builtin from iree.compiler.dialects import std from iree.compiler.dialects import linalg @@ -24,13 +24,13 @@ from iree.compiler.dialects import tosa from iree.compiler.dialects import vector -from iree.compiler.api import driver +from iree.compiler.transforms import ireec # Test the compiler API. with ir.Context() as ctx: chlo.register_chlo_dialect(ctx) mhlo.register_mhlo_dialect(ctx) - iree_dialect.register_dialect(ctx) + iree_input.register_dialect(ctx) input_module = ir.Module.parse(r""" builtin.module { @@ -42,16 +42,16 @@ } """) - options = driver.CompilerOptions() - options.set_input_dialect_mhlo() - options.add_target_backend("cpu") + options = ireec.CompilerOptions("--iree-input-type=mhlo", + "--iree-hal-target-backends=cpu") + print(options) pm = passmanager.PassManager() - driver.build_iree_vm_pass_pipeline(options, pm) + ireec.build_iree_vm_pass_pipeline(options, pm) pm.run(input_module) print(input_module) bytecode_io = io.BytesIO() - driver.translate_module_to_vm_bytecode(options, input_module, bytecode_io) + ireec.translate_module_to_vm_bytecode(options, input_module, bytecode_io) print(f"Bytecode module len = {len(bytecode_io.getbuffer())}") # Check console scripts.
diff --git a/llvm-external-projects/iree-compiler-api/include/iree-compiler-c/Compiler.h b/llvm-external-projects/iree-compiler-api/include/iree-compiler-c/Compiler.h index c1fd471..e4c0d3d 100644 --- a/llvm-external-projects/iree-compiler-api/include/iree-compiler-c/Compiler.h +++ b/llvm-external-projects/iree-compiler-api/include/iree-compiler-c/Compiler.h
@@ -37,14 +37,15 @@ MLIR_CAPI_EXPORTED IreeCompilerOptions ireeCompilerOptionsCreate(); MLIR_CAPI_EXPORTED void ireeCompilerOptionsDestroy(IreeCompilerOptions options); -MLIR_CAPI_EXPORTED void ireeCompilerOptionsSetInputDialectMHLO( - IreeCompilerOptions options); -MLIR_CAPI_EXPORTED void ireeCompilerOptionsSetInputDialectTOSA( - IreeCompilerOptions options); -MLIR_CAPI_EXPORTED void ireeCompilerOptionsSetInputDialectXLA( - IreeCompilerOptions options); -MLIR_CAPI_EXPORTED void ireeCompilerOptionsAddTargetBackend( - IreeCompilerOptions options, const char *targetBackend); +// Parses argv style arguments into a compiler options structure. +MLIR_CAPI_EXPORTED MlirLogicalResult ireeCompilerOptionsSetFlags( + IreeCompilerOptions options, int argc, const char *const *argv, + void (*onError)(MlirStringRef, void *), void *userData); + +// Enumerates any non default flags and invokes the callback. +MLIR_CAPI_EXPORTED void ireeCompilerOptionsGetFlags( + IreeCompilerOptions options, bool nonDefaultOnly, + void (*onFlag)(MlirStringRef, void *), void *userData); //===----------------------------------------------------------------------===// // Compiler stages.
diff --git a/llvm-external-projects/iree-compiler-api/lib/CAPI/Compiler.cpp b/llvm-external-projects/iree-compiler-api/lib/CAPI/Compiler.cpp index bc7a7a6..27b658d 100644 --- a/llvm-external-projects/iree-compiler-api/lib/CAPI/Compiler.cpp +++ b/llvm-external-projects/iree-compiler-api/lib/CAPI/Compiler.cpp
@@ -11,6 +11,7 @@ #include "iree/compiler/InputConversion/MHLO/Passes.h" #include "iree/compiler/InputConversion/TOSA/Passes.h" #include "iree/compiler/Translation/IREEVM.h" +#include "iree/compiler/Utils/OptionUtils.h" #include "iree/tools/init_targets.h" #include "mlir/CAPI/IR.h" #include "mlir/CAPI/Pass.h" @@ -39,6 +40,17 @@ HALTargetOptions executableOptions; VMTargetOptions vmTargetOptions; VMBytecodeTargetOptions vmBytecodeTargetOptions; + + OptionsBinder binder; + + CompilerOptions() : binder(OptionsBinder::local()) { + bindingOptions.bindOptions(binder); + inputDialectOptions.bindOptions(binder); + highLevelOptimizationOptions.bindOptions(binder); + executableOptions.bindOptions(binder); + vmTargetOptions.bindOptions(binder); + vmBytecodeTargetOptions.bindOptions(binder); + } }; } // namespace @@ -53,6 +65,31 @@ return wrap(options); } +MlirLogicalResult ireeCompilerOptionsSetFlags( + IreeCompilerOptions options, int argc, const char *const *argv, + void (*onError)(MlirStringRef, void *), void *userData) { + CompilerOptions *optionsCpp = unwrap(options); + auto callback = [&](llvm::StringRef message) { + if (onError) { + onError(wrap(message), userData); + } + }; + if (failed(optionsCpp->binder.parseArguments(argc, argv, callback))) { + return mlirLogicalResultFailure(); + } + return mlirLogicalResultSuccess(); +} + +void ireeCompilerOptionsGetFlags(IreeCompilerOptions options, + bool nonDefaultOnly, + void (*onFlag)(MlirStringRef, void *), + void *userData) { + auto flagVector = unwrap(options)->binder.printArguments(nonDefaultOnly); + for (std::string &value : flagVector) { + onFlag(wrap(llvm::StringRef(value)), userData); + } +} + void ireeCompilerOptionsDestroy(IreeCompilerOptions options) { delete unwrap(options); }
diff --git a/llvm-external-projects/iree-compiler-api/python/CMakeLists.txt b/llvm-external-projects/iree-compiler-api/python/CMakeLists.txt index cede62e..e1beabc 100644 --- a/llvm-external-projects/iree-compiler-api/python/CMakeLists.txt +++ b/llvm-external-projects/iree-compiler-api/python/CMakeLists.txt
@@ -11,7 +11,7 @@ declare_mlir_python_sources(IREECompilerAPIPythonSources ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/iree/compiler" SOURCES - api/driver.py + transforms/ireec.py version.py ) declare_mlir_python_sources(IREECompilerAPIPythonExtensions) @@ -31,11 +31,11 @@ # Extensions ################################################################################ -declare_mlir_python_extension(IREECompilerAPIPythonExtensions.CompilerDriver - MODULE_NAME _ireeCompilerDriver +declare_mlir_python_extension(IREECompilerAPIPythonExtensions.IREECTransforms + MODULE_NAME _ireecTransforms ADD_TO_PARENT IREECompilerAPIPythonExtensions SOURCES - CompilerModule.cpp + IREECTransforms.cpp EMBED_CAPI_LINK_LIBS IREECompilerAPICompilerCAPI PRIVATE_LINK_LIBS
diff --git a/llvm-external-projects/iree-compiler-api/python/CompilerModule.cpp b/llvm-external-projects/iree-compiler-api/python/IREECTransforms.cpp similarity index 71% rename from llvm-external-projects/iree-compiler-api/python/CompilerModule.cpp rename to llvm-external-projects/iree-compiler-api/python/IREECTransforms.cpp index 37256ce..c7c1d7a 100644 --- a/llvm-external-projects/iree-compiler-api/python/CompilerModule.cpp +++ b/llvm-external-projects/iree-compiler-api/python/IREECTransforms.cpp
@@ -55,6 +55,40 @@ bool binary; }; +void setOptionsFromArgs(IreeCompilerOptions &options, py::args &args) { + std::vector<std::string> allocedArgs; + std::vector<const char *> cArgs; + for (auto &argObject : args) { + allocedArgs.push_back(py::cast<std::string>(argObject)); + cArgs.push_back(allocedArgs.back().c_str()); + } + + std::string errorMessage; + auto callback = +[](MlirStringRef msg, void *userData) { + std::string *innerErrorMessage = static_cast<std::string *>(userData); + if (!innerErrorMessage->empty()) innerErrorMessage->append("\n"); + innerErrorMessage->append(msg.data, msg.length); + }; + if (mlirLogicalResultIsFailure(ireeCompilerOptionsSetFlags( + options, cArgs.size(), cArgs.data(), callback, + static_cast<void *>(&errorMessage)))) { + throw std::invalid_argument(std::move(errorMessage)); + } +} + +py::list getFlagsFromOptions(IreeCompilerOptions &options, + bool nonDefaultOnly) { + py::list flags; + auto callback = +[](MlirStringRef flag, void *userData) { + py::list *innerFlags = static_cast<py::list *>(userData); + py::str s(flag.data, flag.length); + innerFlags->append(std::move(s)); + }; + ireeCompilerOptionsGetFlags(options, nonDefaultOnly, callback, + static_cast<void *>(&flags)); + return flags; +} + } // namespace static const char BUILD_MHLO_IMPORT_PASS_PIPELINE_DOCSTRING[] = @@ -101,40 +135,39 @@ binary data written to it. )"; -PYBIND11_MODULE(_ireeCompilerDriver, m) { - m.doc() = "iree-compiler driver api"; +PYBIND11_MODULE(_ireecTransforms, m) { + m.doc() = "ireec transforms API"; ireeCompilerRegisterTargetBackends(); py::class_<PyCompilerOptions>(m, "CompilerOptions", "Options for the IREE backend compiler.") - .def(py::init<>()) + .def(py::init([](py::args args) { + PyCompilerOptions self; + if (!args.empty()) { + setOptionsFromArgs(self.options, args); + } + return self; + })) .def( - "set_input_dialect_mhlo", - [](PyCompilerOptions &self) { - ireeCompilerOptionsSetInputDialectMHLO(self.options); + "set", + [](PyCompilerOptions &self, py::args args) { + setOptionsFromArgs(self.options, args); }, - "Sets the input type to the 'mhlo' dialect") + "Sets options from flag values in the usual form for command line " + "options") .def( - "set_input_dialect_tosa", - [](PyCompilerOptions &self) { - ireeCompilerOptionsSetInputDialectTOSA(self.options); + "get", + [](PyCompilerOptions &self, bool nonDefaultOnly) { + return getFlagsFromOptions(self.options, nonDefaultOnly); }, - "Sets the input type to the 'tosa' dialect") - .def( - "set_input_dialect_xla", - [](PyCompilerOptions &self) { - ireeCompilerOptionsSetInputDialectTOSA(self.options); - }, - "Sets the input type to the 'mhlo' dialect with XLA compatibility " - "cleanups") - .def( - "add_target_backend", - [](PyCompilerOptions &self, const std::string &targetBackend) { - ireeCompilerOptionsAddTargetBackend(self.options, - targetBackend.c_str()); - }, - py::arg("target_backend"), - "Adds a target backend (i.e. 'cpu', 'vulkan-spirv', etc)"); + py::arg("non_default_only") = true, + "Gets a list of flag values for the options (by default only those " + "that have changed).") + .def("__repr__", [](PyCompilerOptions &self) { + py::list flags = + getFlagsFromOptions(self.options, /*nonDefaultOnly=*/true); + return py::str("<CompilerOptions:") + py::repr(flags) + py::str(">"); + }); m.def( "build_mhlo_import_pass_pipeline", [](MlirPassManager passManager) {
diff --git a/llvm-external-projects/iree-compiler-api/python/iree/compiler/tools/core.py b/llvm-external-projects/iree-compiler-api/python/iree/compiler/tools/core.py index 2281ad7..b675b07 100644 --- a/llvm-external-projects/iree-compiler-api/python/iree/compiler/tools/core.py +++ b/llvm-external-projects/iree-compiler-api/python/iree/compiler/tools/core.py
@@ -184,10 +184,8 @@ cl.append(f"--iree-hal-target-backends={target_backend}") # Output file. - output_file = tfs.alloc_optional("core-output.bin", - export_as=options.output_file) - if output_file: - cl.append(f"-o={output_file}") + if options.output_file: + cl.append(f"-o={options.output_file}") # Translation to perform. cl.append("--iree-mlir-to-vm-bytecode-module") @@ -237,10 +235,25 @@ """ with TempFileSaver.implicit() as tfs: options = CompilerOptions(**kwargs) + retained_output_file = tfs.alloc_optional("core-output.bin", + export_as=options.output_file) + if options.output_file: + options.output_file = retained_output_file cl = build_compile_command_line(input_file, tfs, options) + + # Save a temp file with the command line. + retained_cl = tfs.alloc_optional("core-command-line.txt") + if retained_cl: + with open(retained_cl, "wt") as f: + f.write(" ".join(cl)) + result = invoke_immediate(cl) if options.output_file: return None + # Output as string needs to write to the retained output file itself. + if retained_output_file: + with open(retained_output_file, "wb") as f: + f.write(result) return result @@ -255,11 +268,32 @@ was specified in the options. """ with TempFileSaver.implicit() as tfs: + retained_input_file = tfs.alloc_optional("core-input.mlir") + if retained_input_file: + with open(retained_input_file, + "wt" if isinstance(input_str, str) else "wb") as f: + f.write(input_str) options = CompilerOptions(**kwargs) + retained_output_file = tfs.alloc_optional("core-output.bin", + export_as=options.output_file) + if options.output_file: + options.output_file = retained_output_file cl = build_compile_command_line("-", tfs, options) input_bytes = input_str.encode("utf-8") if isinstance(input_str, str) else input_str + + # Save a temp file with the command line. + retained_cl = tfs.alloc_optional("core-command-line.txt") + if retained_cl: + with open(retained_cl, "wt") as f: + f.write(" ".join(cl)) + result = invoke_immediate(cl, immediate_input=input_bytes) if options.output_file: return None + + # Output as string needs to write to the retained output file itself. + if retained_output_file: + with open(retained_output_file, "wb") as f: + f.write(result) return result
diff --git a/llvm-external-projects/iree-compiler-api/python/iree/compiler/api/driver.py b/llvm-external-projects/iree-compiler-api/python/iree/compiler/transforms/ireec.py similarity index 81% rename from llvm-external-projects/iree-compiler-api/python/iree/compiler/api/driver.py rename to llvm-external-projects/iree-compiler-api/python/iree/compiler/transforms/ireec.py index e8bcd96..571e6db 100644 --- a/llvm-external-projects/iree-compiler-api/python/iree/compiler/api/driver.py +++ b/llvm-external-projects/iree-compiler-api/python/iree/compiler/transforms/ireec.py
@@ -4,4 +4,4 @@ # See https://llvm.org/LICENSE.txt for license information. # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -from .._mlir_libs._ireeCompilerDriver import * +from .._mlir_libs._ireecTransforms import *
diff --git a/llvm-external-projects/iree-compiler-api/setup.py b/llvm-external-projects/iree-compiler-api/setup.py index 2612810..3b8b497 100644 --- a/llvm-external-projects/iree-compiler-api/setup.py +++ b/llvm-external-projects/iree-compiler-api/setup.py
@@ -138,6 +138,7 @@ ext_modules=[ CMakeExtension("iree.compiler._mlir_libs._mlir"), CMakeExtension("iree.compiler._mlir_libs._ireeDialects"), + CMakeExtension("iree.compiler._mlir_libs._ireecTransforms"), CMakeExtension("iree.compiler._mlir_libs._mlirHlo"), CMakeExtension("iree.compiler._mlir_libs._mlirLinalgPasses"), ],
diff --git a/llvm-external-projects/iree-compiler-api/unittests/CMakeLists.txt b/llvm-external-projects/iree-compiler-api/unittests/CMakeLists.txt index f7c6b56..1750864 100644 --- a/llvm-external-projects/iree-compiler-api/unittests/CMakeLists.txt +++ b/llvm-external-projects/iree-compiler-api/unittests/CMakeLists.txt
@@ -1 +1,21 @@ +function(iree_compiler_api_py_test) + cmake_parse_arguments( + ARG + "" + "NAME;MAIN" + "" + ${ARGN} + ) + set(TEST_NAME "iree-compiler-api-${ARG_NAME}") + add_test( + NAME + ${TEST_NAME} + WORKING_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}" + COMMAND "${Python3_EXECUTABLE}" "${CMAKE_CURRENT_SOURCE_DIR}/${ARG_MAIN}" + ) + set_tests_properties(${TEST_NAME} PROPERTIES + ENVIRONMENT PYTHONPATH=${IREE_COMPILER_API_BINARY_DIR}/python_package) +endfunction() + add_subdirectory(tools) +add_subdirectory(transforms/ireec)
diff --git a/llvm-external-projects/iree-compiler-api/unittests/tools/CMakeLists.txt b/llvm-external-projects/iree-compiler-api/unittests/tools/CMakeLists.txt index 9233c3f..e276d13 100644 --- a/llvm-external-projects/iree-compiler-api/unittests/tools/CMakeLists.txt +++ b/llvm-external-projects/iree-compiler-api/unittests/tools/CMakeLists.txt
@@ -4,25 +4,6 @@ # See https://llvm.org/LICENSE.txt for license information. # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -function(iree_compiler_api_py_test) - cmake_parse_arguments( - ARG - "" - "NAME;MAIN" - "" - ${ARGN} - ) - set(TEST_NAME "iree-compiler-api-${ARG_NAME}") - add_test( - NAME - ${TEST_NAME} - WORKING_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}" - COMMAND "${Python3_EXECUTABLE}" "${CMAKE_CURRENT_SOURCE_DIR}/${ARG_MAIN}" - ) - set_tests_properties(${TEST_NAME} PROPERTIES - ENVIRONMENT PYTHONPATH=${IREE_COMPILER_API_BINARY_DIR}/python_package) -endfunction() - iree_compiler_api_py_test( NAME compiler_core_test
diff --git a/llvm-external-projects/iree-compiler-api/unittests/tools/compiler_core_test.py b/llvm-external-projects/iree-compiler-api/unittests/tools/compiler_core_test.py index 996fc29..f0accd6 100644 --- a/llvm-external-projects/iree-compiler-api/unittests/tools/compiler_core_test.py +++ b/llvm-external-projects/iree-compiler-api/unittests/tools/compiler_core_test.py
@@ -199,6 +199,54 @@ self.assertEqual(temp_contents, output_contents) temp_dir.cleanup() + def testExplicitTempFileSaverCompileToStrTextInput(self): + temp_dir = tempfile.TemporaryDirectory() + with iree.compiler.tools.TempFileSaver(temp_dir.name): + output = iree.compiler.tools.compile_str( + SIMPLE_MUL_ASM, + input_type="mhlo", + target_backends=iree.compiler.tools.DEFAULT_TESTING_BACKENDS) + self.assertIsNotNone(output) + self.assertGreater(len(output), 0) + + # There should be a core-input.mlir and core-output.bin in the temp dir. + expected_temp_file = os.path.join(temp_dir.name, "core-output.bin") + self.assertTrue(os.path.exists(expected_temp_file)) + with open(expected_temp_file, "rb") as f: + temp_output = f.read() + self.assertEqual(output, temp_output) + + expected_temp_file = os.path.join(temp_dir.name, "core-input.mlir") + self.assertTrue(os.path.exists(expected_temp_file)) + with open(expected_temp_file, "rt") as f: + input_contents = f.read() + self.assertEqual(SIMPLE_MUL_ASM, input_contents) + temp_dir.cleanup() + + def testExplicitTempFileSaverBinaryInput(self): + temp_dir = tempfile.TemporaryDirectory() + with iree.compiler.tools.TempFileSaver(temp_dir.name): + output = iree.compiler.tools.compile_str( + SIMPLE_MUL_ASM, + input_type="mhlo", + target_backends=iree.compiler.tools.DEFAULT_TESTING_BACKENDS) + self.assertIsNotNone(output) + self.assertGreater(len(output), 0) + + # There should be a core-input.mlir and core-output.bin in the temp dir. + expected_temp_file = os.path.join(temp_dir.name, "core-output.bin") + self.assertTrue(os.path.exists(expected_temp_file)) + with open(expected_temp_file, "rb") as f: + temp_output = f.read() + self.assertEqual(output, temp_output) + + expected_temp_file = os.path.join(temp_dir.name, "core-input.mlir") + self.assertTrue(os.path.exists(expected_temp_file)) + with open(expected_temp_file, "rt") as f: + input_contents = f.read() + self.assertEqual(SIMPLE_MUL_ASM, input_contents) + temp_dir.cleanup() + def testEnvTempFileSaver(self): temp_dir = tempfile.TemporaryDirectory() os.environ["IREE_SAVE_TEMPS"] = temp_dir.name
diff --git a/llvm-external-projects/iree-compiler-api/unittests/transforms/ireec/CMakeLists.txt b/llvm-external-projects/iree-compiler-api/unittests/transforms/ireec/CMakeLists.txt new file mode 100644 index 0000000..e0a903d --- /dev/null +++ b/llvm-external-projects/iree-compiler-api/unittests/transforms/ireec/CMakeLists.txt
@@ -0,0 +1,12 @@ +# Copyright 2022 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 + +iree_compiler_api_py_test( + NAME + compiler_options_test + MAIN + "compiler_options_test.py" +)
diff --git a/llvm-external-projects/iree-compiler-api/unittests/transforms/ireec/compiler_options_test.py b/llvm-external-projects/iree-compiler-api/unittests/transforms/ireec/compiler_options_test.py new file mode 100644 index 0000000..945db55 --- /dev/null +++ b/llvm-external-projects/iree-compiler-api/unittests/transforms/ireec/compiler_options_test.py
@@ -0,0 +1,50 @@ +# Copyright 2022 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 + +import logging +import unittest + +from iree.compiler.transforms import ireec + + +class CompilerTest(unittest.TestCase): + + def testDefaultOptions(self): + options = ireec.CompilerOptions() + self.assertEqual(repr(options), "<CompilerOptions:[]>") + + def testOptionsBadArg(self): + with self.assertRaisesRegex(ValueError, "option not found: foobar"): + options = ireec.CompilerOptions("--foobar") + + def testOptionsBoolArgImplicit(self): + options = ireec.CompilerOptions("--iree-tflite-bindings-support") + self.assertEqual( + repr(options), + "<CompilerOptions:['--iree-tflite-bindings-support=true']>") + + def testOptionsBoolArgExplicit(self): + options = ireec.CompilerOptions("--iree-tflite-bindings-support=true") + self.assertEqual( + repr(options), + "<CompilerOptions:['--iree-tflite-bindings-support=true']>") + + def testOptionsEnumArg(self): + options = ireec.CompilerOptions("--iree-input-type=mhlo") + self.assertEqual(repr(options), + "<CompilerOptions:['--iree-input-type=mhlo']>") + + def testListOption(self): + options = ireec.CompilerOptions("--iree-hal-target-backends=cpu,vmvx") + self.assertEqual( + repr(options), + "<CompilerOptions:['--iree-hal-target-backends=cpu,vmvx']>") + print(options) + + +if __name__ == "__main__": + logging.basicConfig(level=logging.DEBUG) + unittest.main()