samples/custom_modules/module.cc - 3p/openxla/iree - Git at Google

 // 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 "samples/custom_modules/module.h"

 #include <atomic>
 #include <cstdint>
 #include <cstdio>
 #include <cstring>
 #include <memory>
 #include <string>
 #include <type_traits>
 #include <utility>

 #include "iree/base/api.h"
 #include "iree/base/status_cc.h"
 #include "iree/hal/api.h"
 #include "iree/modules/hal/module.h"
 #include "iree/vm/native_module_cc.h"
 #include "iree/vm/ref_cc.h"

 //===----------------------------------------------------------------------===//
 // !custom.message type
 //===----------------------------------------------------------------------===//

 // Runtime type descriptor for the !custom.message describing how to manage it
 // and destroy it. The type ID is allocated at runtime and does not need to
 // match the compiler ID.
 static iree_vm_ref_type_descriptor_t iree_custom_message_descriptor = {0};

 // The "message" type we use to store string messages to print.
 // This could be arbitrarily complex or simply wrap another user-defined type.
 // The descriptor that is registered at startup defines how to manage the
 // lifetime of the type (such as which destruction function is called, if any).
 // See ref.h for more information and additional utilities.
 typedef struct iree_custom_message_t {
   // Ideally first; used to track the reference count of the object.
   iree_vm_ref_object_t ref_object;
   // Allocator the message was created from.
   // Ideally pools and nested allocators would be used to avoid needing to store
   // the allocator with every object.
   iree_allocator_t allocator;
   // String message value.
   iree_string_view_t value;
 } iree_custom_message_t;

 IREE_VM_DEFINE_TYPE_ADAPTERS(iree_custom_message, iree_custom_message_t);

 iree_status_t iree_custom_message_create(iree_string_view_t value,
                                          iree_allocator_t allocator,
                                          iree_custom_message_t** out_message) {
   IREE_ASSERT_ARGUMENT(out_message);
   // Note that we allocate the message and the string value together.
   iree_custom_message_t* message = NULL;
   IREE_RETURN_IF_ERROR(iree_allocator_malloc(
       allocator, sizeof(*message) + value.size, (void**)&message));
   message->ref_object.counter = IREE_ATOMIC_VAR_INIT(1);
   message->allocator = allocator;
   message->value.data = ((const char*)message) + sizeof(iree_custom_message_t);
   message->value.size = value.size;
   memcpy((void*)message->value.data, value.data, message->value.size);
   *out_message = message;
   return iree_ok_status();
 }

 iree_status_t iree_custom_message_wrap(iree_string_view_t value,
                                        iree_allocator_t allocator,
                                        iree_custom_message_t** out_message) {
   IREE_ASSERT_ARGUMENT(out_message);
   iree_custom_message_t* message = NULL;
   IREE_RETURN_IF_ERROR(
       iree_allocator_malloc(allocator, sizeof(*message), (void**)&message));
   message->ref_object.counter = IREE_ATOMIC_VAR_INIT(1);
   message->allocator = allocator;
   message->value = value;  // Unowned.
   *out_message = message;
   return iree_ok_status();
 }

 void iree_custom_message_destroy(void* ptr) {
   iree_custom_message_t* message = (iree_custom_message_t*)ptr;
   iree_allocator_free(message->allocator, ptr);
 }

 iree_status_t iree_custom_message_read_value(iree_custom_message_t* message,
                                              char* buffer,
                                              size_t buffer_capacity) {
   IREE_ASSERT_ARGUMENT(message);
   IREE_ASSERT_ARGUMENT(buffer);
   if (buffer_capacity < message->value.size + 1) {
     // Not an error; just a size query.
     return iree_status_from_code(IREE_STATUS_OUT_OF_RANGE);
   }
   memcpy(buffer, message->value.data, message->value.size);
   buffer[message->value.size] = 0;
   return iree_ok_status();
 }

 iree_status_t iree_custom_native_module_register_types() {
   if (iree_custom_message_descriptor.type) {
     return iree_ok_status();  // Already registered.
   }
   iree_custom_message_descriptor.type_name =
       iree_make_cstring_view("custom.message");
   iree_custom_message_descriptor.offsetof_counter =
       offsetof(iree_custom_message_t, ref_object.counter);
   iree_custom_message_descriptor.destroy = iree_custom_message_destroy;
   return iree_vm_ref_register_type(&iree_custom_message_descriptor);
 }

 //===----------------------------------------------------------------------===//
 // VM module interface implementation
 //===----------------------------------------------------------------------===//

 namespace iree {
 namespace samples {
 namespace {

 // Per-context module state.
 // This can contain "globals" and other arbitrary state.
 //
 // Thread-compatible; the runtime will not issue multiple calls at the same
 // time using the same state. If the implementation uses external threads then
 // it must synchronize itself.
 class CustomModuleState final {
  public:
   explicit CustomModuleState(iree_allocator_t host_allocator,
                              vm::ref<iree_hal_allocator_t> device_allocator)
       : host_allocator_(host_allocator),
         device_allocator_(std::move(device_allocator)) {}
   ~CustomModuleState() = default;

   Status Initialize(int32_t unique_id) {
     // Allocate a unique ID to demonstrate per-context state.
     auto str_buffer = "ctx_" + std::to_string(unique_id);
     IREE_RETURN_IF_ERROR(
         iree_custom_message_create(iree_make_cstring_view(str_buffer.c_str()),
                                    host_allocator_, &unique_message_));

     return OkStatus();
   }

   // custom.buffer_to_message(%buffer_view) -> %result
   StatusOr<vm::ref<iree_custom_message_t>> BufferToMessage(
       vm::ref<iree_hal_buffer_view_t> buffer_view) {
     // Convert the buffer view to a [shape]x[type]=[contents] string.
     std::string string_value(4096, '\0');
     iree_status_t status;
     do {
       iree_host_size_t actual_length = 0;
       status = iree_hal_buffer_view_format(
           buffer_view.get(), /*max_element_count=*/1024,
           string_value.size() + 1, &string_value[0], &actual_length);
       string_value.resize(actual_length);
     } while (iree_status_is_out_of_range(status));
     IREE_RETURN_IF_ERROR(std::move(status));

     // Pack the string contents into a message.
     vm::ref<iree_custom_message_t> message;
     IREE_RETURN_IF_ERROR(iree_custom_message_create(
         iree_string_view_t{string_value.data(), string_value.size()},
         host_allocator_, &message));
     return std::move(message);
   }

   // custom.message_to_buffer(%message) -> %buffer_view
   StatusOr<vm::ref<iree_hal_buffer_view_t>> MessageToBuffer(
       vm::ref<iree_custom_message_t> message) {
     // Convert the [shape]x[type]=[contents] string to a buffer view.
     vm::ref<iree_hal_buffer_view_t> buffer_view;
     IREE_RETURN_IF_ERROR(
         iree_hal_buffer_view_parse(message->value, device_allocator_.get(),
                                    &buffer_view),
         "parsing value '%.*s'", (int)message->value.size, message->value.data);
     return std::move(buffer_view);
   }

   // custom.print(%message, %count)
   Status Print(vm::ref<iree_custom_message_t> message, int32_t count) {
     for (int i = 0; i < count; ++i) {
       fwrite(message->value.data, 1, message->value.size, stdout);
       fputc('\n', stdout);
     }
     fflush(stdout);
     return OkStatus();
   }

   // custom.reverse(%message) -> %result
   StatusOr<vm::ref<iree_custom_message_t>> Reverse(
       vm::ref<iree_custom_message_t> message) {
     vm::ref<iree_custom_message_t> reversed_message;
     IREE_RETURN_IF_ERROR(iree_custom_message_create(
         message->value, message->allocator, &reversed_message));
     char* str_ptr = const_cast<char*>(reversed_message->value.data);
     for (int low = 0, high = reversed_message->value.size - 1; low < high;
          ++low, --high) {
       char temp = str_ptr[low];
       str_ptr[low] = str_ptr[high];
       str_ptr[high] = temp;
     }
     return std::move(reversed_message);
   }

   // custom.get_unique_message() -> %result
   StatusOr<vm::ref<iree_custom_message_t>> GetUniqueMessage() {
     return vm::retain_ref(unique_message_);
   }

  private:
   // Allocator that the caller requested we use for any allocations we need to
   // perform during operation.
   iree_allocator_t host_allocator_ = iree_allocator_system();

   // HAL buffer allocator that uses host-local memory. This is just for this
   // test as we don't actually use a HAL device and don't have a real device
   // allocator to use.
   vm::ref<iree_hal_allocator_t> device_allocator_;

   // A unique message owned by the state and returned to the VM.
   // This demonstrates any arbitrary per-context state one may want to store.
   vm::ref<iree_custom_message_t> unique_message_;
 };

 // Function table mapping imported function names to their implementation.
 // The signature of the target function is expected to match that in the
 // custom.imports.mlir file.
 static const vm::NativeFunction<CustomModuleState> kCustomModuleFunctions[] = {
     vm::MakeNativeFunction("buffer_to_message",
                            &CustomModuleState::BufferToMessage),
     vm::MakeNativeFunction("message_to_buffer",
                            &CustomModuleState::MessageToBuffer),
     vm::MakeNativeFunction("print", &CustomModuleState::Print),
     vm::MakeNativeFunction("reverse", &CustomModuleState::Reverse),
     vm::MakeNativeFunction("get_unique_message",
                            &CustomModuleState::GetUniqueMessage),
 };

 // The module instance that will be allocated and reused across contexts.
 // Any context-specific state must be stored in a state structure such as
 // CustomModuleState below.
 //
 // Assumed thread-safe (by construction here, as it's immutable), though if more
 // state is stored here it will need to be synchronized by the implementation.
 class CustomModule final : public vm::NativeModule<CustomModuleState> {
  public:
   using vm::NativeModule<CustomModuleState>::NativeModule;

   // Example of global initialization (shared across all contexts), such as
   // loading native libraries or creating shared pools.
   Status Initialize(iree_hal_allocator_t* device_allocator) {
     device_allocator_ = vm::retain_ref(device_allocator);
     next_unique_id_ = 0;
     return OkStatus();
   }

   // Creates per-context state when the module is added to a new context.
   // May be called from any thread.
   StatusOr<std::unique_ptr<CustomModuleState>> CreateState(
       iree_allocator_t host_allocator) override {
     auto state = std::make_unique<CustomModuleState>(
         host_allocator, vm::retain_ref(device_allocator_));
     IREE_RETURN_IF_ERROR(state->Initialize(next_unique_id_++));
     return state;
   }

  private:
   // Used for any HAL buffer allocations we need to make.
   vm::ref<iree_hal_allocator_t> device_allocator_;

   // The next ID to allocate for states that will be used to form the
   // per-context unique message. This shows state at the module level. Note that
   // this must be thread-safe.
   std::atomic<int32_t> next_unique_id_;
 };

 }  // namespace

 // Note that while we are using C++ bindings internally we still expose the
 // module as a C instance. This hides the details of our implementation.
 extern "C" iree_status_t iree_custom_native_module_create(
     iree_allocator_t host_allocator, iree_hal_allocator_t* device_allocator,
     iree_vm_module_t** out_module) {
   IREE_ASSERT_ARGUMENT(out_module);
   *out_module = NULL;
   auto module = std::make_unique<CustomModule>(
       "custom", host_allocator,
       iree::span<const vm::NativeFunction<CustomModuleState>>(
           kCustomModuleFunctions));
   IREE_RETURN_IF_ERROR(module->Initialize(device_allocator));
   *out_module = module.release()->interface();
   return iree_ok_status();
 }

 }  // namespace samples
 }  // namespace iree
	// 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 "samples/custom_modules/module.h"

	#include <atomic>
	#include <cstdint>
	#include <cstdio>
	#include <cstring>
	#include <memory>
	#include <string>
	#include <type_traits>
	#include <utility>

	#include "iree/base/api.h"
	#include "iree/base/status_cc.h"
	#include "iree/hal/api.h"
	#include "iree/modules/hal/module.h"
	#include "iree/vm/native_module_cc.h"
	#include "iree/vm/ref_cc.h"

	//===----------------------------------------------------------------------===//
	// !custom.message type
	//===----------------------------------------------------------------------===//

	// Runtime type descriptor for the !custom.message describing how to manage it
	// and destroy it. The type ID is allocated at runtime and does not need to
	// match the compiler ID.
	static iree_vm_ref_type_descriptor_t iree_custom_message_descriptor = {0};

	// The "message" type we use to store string messages to print.
	// This could be arbitrarily complex or simply wrap another user-defined type.
	// The descriptor that is registered at startup defines how to manage the
	// lifetime of the type (such as which destruction function is called, if any).
	// See ref.h for more information and additional utilities.
	typedef struct iree_custom_message_t {
	// Ideally first; used to track the reference count of the object.
	iree_vm_ref_object_t ref_object;
	// Allocator the message was created from.
	// Ideally pools and nested allocators would be used to avoid needing to store
	// the allocator with every object.
	iree_allocator_t allocator;
	// String message value.
	iree_string_view_t value;
	} iree_custom_message_t;

	IREE_VM_DEFINE_TYPE_ADAPTERS(iree_custom_message, iree_custom_message_t);

	iree_status_t iree_custom_message_create(iree_string_view_t value,
	iree_allocator_t allocator,
	iree_custom_message_t** out_message) {
	IREE_ASSERT_ARGUMENT(out_message);
	// Note that we allocate the message and the string value together.
	iree_custom_message_t* message = NULL;
	IREE_RETURN_IF_ERROR(iree_allocator_malloc(
	allocator, sizeof(message) + value.size, (void*)&message));
	message->ref_object.counter = IREE_ATOMIC_VAR_INIT(1);
	message->allocator = allocator;
	message->value.data = ((const char*)message) + sizeof(iree_custom_message_t);
	message->value.size = value.size;
	memcpy((void*)message->value.data, value.data, message->value.size);
	*out_message = message;
	return iree_ok_status();
	}

	iree_status_t iree_custom_message_wrap(iree_string_view_t value,
	iree_allocator_t allocator,
	iree_custom_message_t** out_message) {
	IREE_ASSERT_ARGUMENT(out_message);
	iree_custom_message_t* message = NULL;
	IREE_RETURN_IF_ERROR(
	iree_allocator_malloc(allocator, sizeof(message), (void*)&message));
	message->ref_object.counter = IREE_ATOMIC_VAR_INIT(1);
	message->allocator = allocator;
	message->value = value; // Unowned.
	*out_message = message;
	return iree_ok_status();
	}

	void iree_custom_message_destroy(void* ptr) {
	iree_custom_message_t* message = (iree_custom_message_t*)ptr;
	iree_allocator_free(message->allocator, ptr);
	}

	iree_status_t iree_custom_message_read_value(iree_custom_message_t* message,
	char* buffer,
	size_t buffer_capacity) {
	IREE_ASSERT_ARGUMENT(message);
	IREE_ASSERT_ARGUMENT(buffer);
	if (buffer_capacity < message->value.size + 1) {
	// Not an error; just a size query.
	return iree_status_from_code(IREE_STATUS_OUT_OF_RANGE);
	}
	memcpy(buffer, message->value.data, message->value.size);
	buffer[message->value.size] = 0;
	return iree_ok_status();
	}

	iree_status_t iree_custom_native_module_register_types() {
	if (iree_custom_message_descriptor.type) {
	return iree_ok_status(); // Already registered.
	}
	iree_custom_message_descriptor.type_name =
	iree_make_cstring_view("custom.message");
	iree_custom_message_descriptor.offsetof_counter =
	offsetof(iree_custom_message_t, ref_object.counter);
	iree_custom_message_descriptor.destroy = iree_custom_message_destroy;
	return iree_vm_ref_register_type(&iree_custom_message_descriptor);
	}

	//===----------------------------------------------------------------------===//
	// VM module interface implementation
	//===----------------------------------------------------------------------===//

	namespace iree {
	namespace samples {
	namespace {

	// Per-context module state.
	// This can contain "globals" and other arbitrary state.
	//
	// Thread-compatible; the runtime will not issue multiple calls at the same
	// time using the same state. If the implementation uses external threads then
	// it must synchronize itself.
	class CustomModuleState final {
	public:
	explicit CustomModuleState(iree_allocator_t host_allocator,
	vm::ref<iree_hal_allocator_t> device_allocator)
	: host_allocator_(host_allocator),
	device_allocator_(std::move(device_allocator)) {}
	~CustomModuleState() = default;

	Status Initialize(int32_t unique_id) {
	// Allocate a unique ID to demonstrate per-context state.
	auto str_buffer = "ctx_" + std::to_string(unique_id);
	IREE_RETURN_IF_ERROR(
	iree_custom_message_create(iree_make_cstring_view(str_buffer.c_str()),
	host_allocator_, &unique_message_));

	return OkStatus();
	}

	// custom.buffer_to_message(%buffer_view) -> %result
	StatusOr<vm::ref<iree_custom_message_t>> BufferToMessage(
	vm::ref<iree_hal_buffer_view_t> buffer_view) {
	// Convert the buffer view to a [shape]x[type]=[contents] string.
	std::string string_value(4096, '\0');
	iree_status_t status;
	do {
	iree_host_size_t actual_length = 0;
	status = iree_hal_buffer_view_format(
	buffer_view.get(), /max_element_count=/1024,
	string_value.size() + 1, &string_value[0], &actual_length);
	string_value.resize(actual_length);
	} while (iree_status_is_out_of_range(status));
	IREE_RETURN_IF_ERROR(std::move(status));

	// Pack the string contents into a message.
	vm::ref<iree_custom_message_t> message;
	IREE_RETURN_IF_ERROR(iree_custom_message_create(
	iree_string_view_t{string_value.data(), string_value.size()},
	host_allocator_, &message));
	return std::move(message);
	}

	// custom.message_to_buffer(%message) -> %buffer_view
	StatusOr<vm::ref<iree_hal_buffer_view_t>> MessageToBuffer(
	vm::ref<iree_custom_message_t> message) {
	// Convert the [shape]x[type]=[contents] string to a buffer view.
	vm::ref<iree_hal_buffer_view_t> buffer_view;
	IREE_RETURN_IF_ERROR(
	iree_hal_buffer_view_parse(message->value, device_allocator_.get(),
	&buffer_view),
	"parsing value '%.*s'", (int)message->value.size, message->value.data);
	return std::move(buffer_view);
	}

	// custom.print(%message, %count)
	Status Print(vm::ref<iree_custom_message_t> message, int32_t count) {
	for (int i = 0; i < count; ++i) {
	fwrite(message->value.data, 1, message->value.size, stdout);
	fputc('\n', stdout);
	}
	fflush(stdout);
	return OkStatus();
	}

	// custom.reverse(%message) -> %result
	StatusOr<vm::ref<iree_custom_message_t>> Reverse(
	vm::ref<iree_custom_message_t> message) {
	vm::ref<iree_custom_message_t> reversed_message;
	IREE_RETURN_IF_ERROR(iree_custom_message_create(
	message->value, message->allocator, &reversed_message));
	char* str_ptr = const_cast<char*>(reversed_message->value.data);
	for (int low = 0, high = reversed_message->value.size - 1; low < high;
	++low, --high) {
	char temp = str_ptr[low];
	str_ptr[low] = str_ptr[high];
	str_ptr[high] = temp;
	}
	return std::move(reversed_message);
	}

	// custom.get_unique_message() -> %result
	StatusOr<vm::ref<iree_custom_message_t>> GetUniqueMessage() {
	return vm::retain_ref(unique_message_);
	}

	private:
	// Allocator that the caller requested we use for any allocations we need to
	// perform during operation.
	iree_allocator_t host_allocator_ = iree_allocator_system();

	// HAL buffer allocator that uses host-local memory. This is just for this
	// test as we don't actually use a HAL device and don't have a real device
	// allocator to use.
	vm::ref<iree_hal_allocator_t> device_allocator_;

	// A unique message owned by the state and returned to the VM.
	// This demonstrates any arbitrary per-context state one may want to store.
	vm::ref<iree_custom_message_t> unique_message_;
	};

	// Function table mapping imported function names to their implementation.
	// The signature of the target function is expected to match that in the
	// custom.imports.mlir file.
	static const vm::NativeFunction<CustomModuleState> kCustomModuleFunctions[] = {
	vm::MakeNativeFunction("buffer_to_message",
	&CustomModuleState::BufferToMessage),
	vm::MakeNativeFunction("message_to_buffer",
	&CustomModuleState::MessageToBuffer),
	vm::MakeNativeFunction("print", &CustomModuleState::Print),
	vm::MakeNativeFunction("reverse", &CustomModuleState::Reverse),
	vm::MakeNativeFunction("get_unique_message",
	&CustomModuleState::GetUniqueMessage),
	};

	// The module instance that will be allocated and reused across contexts.
	// Any context-specific state must be stored in a state structure such as
	// CustomModuleState below.
	//
	// Assumed thread-safe (by construction here, as it's immutable), though if more
	// state is stored here it will need to be synchronized by the implementation.
	class CustomModule final : public vm::NativeModule<CustomModuleState> {
	public:
	using vm::NativeModule<CustomModuleState>::NativeModule;

	// Example of global initialization (shared across all contexts), such as
	// loading native libraries or creating shared pools.
	Status Initialize(iree_hal_allocator_t* device_allocator) {
	device_allocator_ = vm::retain_ref(device_allocator);
	next_unique_id_ = 0;
	return OkStatus();
	}

	// Creates per-context state when the module is added to a new context.
	// May be called from any thread.
	StatusOr<std::unique_ptr<CustomModuleState>> CreateState(
	iree_allocator_t host_allocator) override {
	auto state = std::make_unique<CustomModuleState>(
	host_allocator, vm::retain_ref(device_allocator_));
	IREE_RETURN_IF_ERROR(state->Initialize(next_unique_id_++));
	return state;
	}

	private:
	// Used for any HAL buffer allocations we need to make.
	vm::ref<iree_hal_allocator_t> device_allocator_;

	// The next ID to allocate for states that will be used to form the
	// per-context unique message. This shows state at the module level. Note that
	// this must be thread-safe.
	std::atomic<int32_t> next_unique_id_;
	};

	} // namespace

	// Note that while we are using C++ bindings internally we still expose the
	// module as a C instance. This hides the details of our implementation.
	extern "C" iree_status_t iree_custom_native_module_create(
	iree_allocator_t host_allocator, iree_hal_allocator_t* device_allocator,
	iree_vm_module_t** out_module) {
	IREE_ASSERT_ARGUMENT(out_module);
	*out_module = NULL;
	auto module = std::make_unique<CustomModule>(
	"custom", host_allocator,
	iree::span<const vm::NativeFunction<CustomModuleState>>(
	kCustomModuleFunctions));
	IREE_RETURN_IF_ERROR(module->Initialize(device_allocator));
	*out_module = module.release()->interface();
	return iree_ok_status();
	}

	} // namespace samples
	} // namespace iree