iree/vm/module_abi_cc.h - 3p/openxla/iree - Git at Google

 // Copyright 2019 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 //      https://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.

 #ifndef IREE_VM_MODULE_ABI_CC_H_
 #define IREE_VM_MODULE_ABI_CC_H_

 #include <cstring>
 #include <memory>

 #include "absl/strings/str_cat.h"
 #include "absl/types/span.h"
 #include "iree/base/api.h"
 #include "iree/base/api_util.h"
 #include "iree/base/status.h"
 #include "iree/vm/module.h"
 #include "iree/vm/module_abi_packing.h"
 #include "iree/vm/stack.h"

 #ifndef __cplusplus
 #error "This header is meant for use with C++ module implementations."
 #endif  // __cplusplus

 namespace iree {
 namespace vm {

 // A native module as exported to the VM dynamic module linking API.
 // This allows easy wrapping of C++ module implementations and removes a
 // majority of the boilerplate required with marshaling args/results out/in of
 // the VM via the ABI.
 //
 // Functions are defined on the State type as member functions returning either
 // Status or StatusOr. Arguments are passed as primitive types (int32_t),
 // wrapped ref objects (vm::ref<my_type_t>&), or some nesting of std::array,
 // std::tuple, and absl::Span to match fixed-length arrays of the same type,
 // tuples of mixed types, or dynamic arrays (variadic arguments). Results may be
 // returned as either their type or an std::tuple/std::array of types.
 //
 // Usage:
 //   // Per-context module state that must only be thread-compatible.
 //   // Define
 //   struct MyState final {
 //     StatusOr<std::tuple<int32_t, int32_t>> MyMethod1(vm::ref<my_type_t> t);
 //   };
 //
 //   // Table of functions mapped to their name in the IR.
 //   static const vm::NativeFunction<MyState> kMyFunctions[] = {
 //     vm::MakeNativeFunction("my_method_1", &MyState::MyMethod1),
 //   };
 //
 //   // The outer module wrapper shared across contexts.
 //   // Must be thread-safe.
 //   struct MyModule : public NativeModule<MyState> {
 //     StatusOr<std::unique_ptr<MyState>> CreateState(iree_allocator_t) {
 //       // You could pass in thread-safe shared resources to MyState.
 //       return std::make_unique<MyState>();
 //     }
 //   };
 //
 //   // Creates the module and exposes it as a C interface.
 //   // Ownership transfers to the caller.
 //   iree_vm_module_t* create_my_module(iree_allocator_t allocator) {
 //     return std::make_unique<MyModule>("my_module", allocator,
 //         absl::MakeConstSpan(kCustomModuleFunctions)).release()->interface();
 //   }
 template <typename State>
 class NativeModule {
  public:
   NativeModule(const char* name, iree_allocator_t allocator,
                absl::Span<const NativeFunction<State>> dispatch_table)
       : name_(name), allocator_(allocator), dispatch_table_(dispatch_table) {
     CHECK_OK(FromApiStatus(iree_vm_module_init(&interface_, this), IREE_LOC));
     interface_.destroy = NativeModule::ModuleDestroy;
     interface_.name = NativeModule::ModuleName;
     interface_.signature = NativeModule::ModuleSignature;
     interface_.get_function = NativeModule::ModuleGetFunction;
     interface_.lookup_function = NativeModule::ModuleLookupFunction;
     interface_.alloc_state = NativeModule::ModuleAllocState;
     interface_.free_state = NativeModule::ModuleFreeState;
     interface_.resolve_import = NativeModule::ModuleResolveImport;
     interface_.execute = NativeModule::ModuleExecute;
   }

   virtual ~NativeModule() = default;

   // C API module interface bound to this NativeModule instance.
   iree_vm_module_t* interface() { return &interface_; }

  protected:
   // Creates a new per-context module State holder.
   virtual StatusOr<std::unique_ptr<State>> CreateState(
       iree_allocator_t allocator) = 0;

  private:
   static NativeModule* FromModulePointer(void* self) {
     return reinterpret_cast<NativeModule*>(self);
   }
   static State* FromStatePointer(void* self) {
     return reinterpret_cast<State*>(self);
   }

   static iree_status_t ModuleDestroy(void* self) {
     delete FromModulePointer(self);
     return IREE_STATUS_OK;
   }

   static iree_string_view_t ModuleName(void* self) {
     auto* module = FromModulePointer(self);
     return iree_make_cstring_view(module->name_);
   }

   static iree_vm_module_signature_t ModuleSignature(void* self) {
     auto* module = FromModulePointer(self);
     iree_vm_module_signature_t signature = {0};
     signature.import_function_count = 0;
     signature.export_function_count = module->dispatch_table_.size();
     signature.internal_function_count = 0;
     return signature;
   }

   static iree_status_t ModuleGetFunction(
       void* self, iree_vm_function_linkage_t linkage, int32_t ordinal,
       iree_vm_function_t* out_function, iree_string_view_t* out_name,
       iree_vm_function_signature_t* out_signature) {
     if (out_function) {
       std::memset(out_function, 0, sizeof(*out_function));
     }
     if (out_name) {
       out_name->data = nullptr;
       out_name->size = 0;
     }
     if (out_signature) {
       std::memset(out_signature, 0, sizeof(*out_signature));
     }
     auto* module = FromModulePointer(self);
     if (ordinal < 0 || ordinal > module->dispatch_table_.size()) {
       return IREE_STATUS_INVALID_ARGUMENT;
     }
     if (out_function) {
       out_function->module = module->interface();
       out_function->linkage = IREE_VM_FUNCTION_LINKAGE_EXPORT;
       out_function->ordinal = ordinal;
     }
     if (out_name) {
       const auto& dispatch_function = module->dispatch_table_[ordinal];
       *out_name = iree_make_cstring_view(dispatch_function.name);
     }
     return IREE_STATUS_OK;
   }

   static iree_status_t ModuleLookupFunction(void* self,
                                             iree_vm_function_linkage_t linkage,
                                             iree_string_view_t name,
                                             iree_vm_function_t* out_function) {
     if (!out_function) return IREE_STATUS_INVALID_ARGUMENT;
     std::memset(out_function, 0, sizeof(*out_function));
     if (!name.data || !name.size) return IREE_STATUS_INVALID_ARGUMENT;

     auto* module = FromModulePointer(self);
     out_function->module = module->interface();
     out_function->linkage = IREE_VM_FUNCTION_LINKAGE_EXPORT;
     for (int i = 0; i < module->dispatch_table_.size(); ++i) {
       if (iree_string_view_compare(
               name, iree_make_cstring_view(module->dispatch_table_[i].name)) ==
           0) {
         out_function->ordinal = i;
         return IREE_STATUS_OK;
       }
     }
     return IREE_STATUS_NOT_FOUND;
   }

   static iree_status_t ModuleAllocState(
       void* self, iree_allocator_t allocator,
       iree_vm_module_state_t** out_module_state) {
     if (!out_module_state) return IREE_STATUS_INVALID_ARGUMENT;
     *out_module_state = nullptr;

     auto* module = FromModulePointer(self);
     auto module_state_or = module->CreateState(allocator);
     if (!module_state_or.ok()) {
       return ToApiStatus(module_state_or.status());
     }
     auto module_state = std::move(module_state_or).ValueOrDie();

     *out_module_state =
         reinterpret_cast<iree_vm_module_state_t*>(module_state.release());
     return IREE_STATUS_OK;
   }

   static iree_status_t ModuleFreeState(void* self,
                                        iree_vm_module_state_t* module_state) {
     if (!module_state) return IREE_STATUS_INVALID_ARGUMENT;
     delete FromStatePointer(module_state);
     return IREE_STATUS_OK;
   }

   static iree_status_t ModuleResolveImport(void* self,
                                            iree_vm_module_state_t* module_state,
                                            int32_t ordinal,
                                            iree_vm_function_t function) {
     // C++ API does not yet support imports.
     return IREE_STATUS_FAILED_PRECONDITION;
   }

   static iree_status_t ModuleExecute(void* self, iree_vm_stack_t* stack,
                                      iree_vm_stack_frame_t* frame,
                                      iree_vm_execution_result_t* out_result) {
     if (!out_result) return IREE_STATUS_INVALID_ARGUMENT;
     std::memset(out_result, 0, sizeof(*out_result));
     if (!stack || !frame) return IREE_STATUS_INVALID_ARGUMENT;
     int32_t ordinal = frame->function.ordinal;
     auto* module = FromModulePointer(self);
     if (ordinal < 0 || ordinal > module->dispatch_table_.size()) {
       return IREE_STATUS_INVALID_ARGUMENT;
     }
     const auto& info = module->dispatch_table_[ordinal];
     auto* state = FromStatePointer(frame->module_state);
     auto status = info.call(info.ptr, state, stack, frame, out_result);
     if (!status.ok()) {
       status = iree::Annotate(
           status,
           absl::StrCat("while executing ", module->name_, ".", info.name));
       return ToApiStatus(status);
     }
     return IREE_STATUS_OK;
   }

   const char* name_;
   const iree_allocator_t allocator_;
   iree_vm_module_t interface_;

   const absl::Span<const NativeFunction<State>> dispatch_table_;
 };

 }  // namespace vm
 }  // namespace iree

 #endif  // IREE_VM_MODULE_ABI_CC_H_
	// Copyright 2019 Google LLC
	//
	// Licensed under the Apache License, Version 2.0 (the "License");
	// you may not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// https://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.

	#ifndef IREE_VM_MODULE_ABI_CC_H_
	#define IREE_VM_MODULE_ABI_CC_H_

	#include <cstring>
	#include <memory>

	#include "absl/strings/str_cat.h"
	#include "absl/types/span.h"
	#include "iree/base/api.h"
	#include "iree/base/api_util.h"
	#include "iree/base/status.h"
	#include "iree/vm/module.h"
	#include "iree/vm/module_abi_packing.h"
	#include "iree/vm/stack.h"

	#ifndef __cplusplus
	#error "This header is meant for use with C++ module implementations."
	#endif // __cplusplus

	namespace iree {
	namespace vm {

	// A native module as exported to the VM dynamic module linking API.
	// This allows easy wrapping of C++ module implementations and removes a
	// majority of the boilerplate required with marshaling args/results out/in of
	// the VM via the ABI.
	//
	// Functions are defined on the State type as member functions returning either
	// Status or StatusOr. Arguments are passed as primitive types (int32_t),
	// wrapped ref objects (vm::ref<my_type_t>&), or some nesting of std::array,
	// std::tuple, and absl::Span to match fixed-length arrays of the same type,
	// tuples of mixed types, or dynamic arrays (variadic arguments). Results may be
	// returned as either their type or an std::tuple/std::array of types.
	//
	// Usage:
	// // Per-context module state that must only be thread-compatible.
	// // Define
	// struct MyState final {
	// StatusOr<std::tuple<int32_t, int32_t>> MyMethod1(vm::ref<my_type_t> t);
	// };
	//
	// // Table of functions mapped to their name in the IR.
	// static const vm::NativeFunction<MyState> kMyFunctions[] = {
	// vm::MakeNativeFunction("my_method_1", &MyState::MyMethod1),
	// };
	//
	// // The outer module wrapper shared across contexts.
	// // Must be thread-safe.
	// struct MyModule : public NativeModule<MyState> {
	// StatusOr<std::unique_ptr<MyState>> CreateState(iree_allocator_t) {
	// // You could pass in thread-safe shared resources to MyState.
	// return std::make_unique<MyState>();
	// }
	// };
	//
	// // Creates the module and exposes it as a C interface.
	// // Ownership transfers to the caller.
	// iree_vm_module_t* create_my_module(iree_allocator_t allocator) {
	// return std::make_unique<MyModule>("my_module", allocator,
	// absl::MakeConstSpan(kCustomModuleFunctions)).release()->interface();
	// }
	template <typename State>
	class NativeModule {
	public:
	NativeModule(const char* name, iree_allocator_t allocator,
	absl::Span<const NativeFunction<State>> dispatch_table)
	: name_(name), allocator_(allocator), dispatch_table_(dispatch_table) {
	CHECK_OK(FromApiStatus(iree_vm_module_init(&interface_, this), IREE_LOC));
	interface_.destroy = NativeModule::ModuleDestroy;
	interface_.name = NativeModule::ModuleName;
	interface_.signature = NativeModule::ModuleSignature;
	interface_.get_function = NativeModule::ModuleGetFunction;
	interface_.lookup_function = NativeModule::ModuleLookupFunction;
	interface_.alloc_state = NativeModule::ModuleAllocState;
	interface_.free_state = NativeModule::ModuleFreeState;
	interface_.resolve_import = NativeModule::ModuleResolveImport;
	interface_.execute = NativeModule::ModuleExecute;
	}

	virtual ~NativeModule() = default;

	// C API module interface bound to this NativeModule instance.
	iree_vm_module_t* interface() { return &interface_; }

	protected:
	// Creates a new per-context module State holder.
	virtual StatusOr<std::unique_ptr<State>> CreateState(
	iree_allocator_t allocator) = 0;

	private:
	static NativeModule* FromModulePointer(void* self) {
	return reinterpret_cast<NativeModule*>(self);
	}
	static State* FromStatePointer(void* self) {
	return reinterpret_cast<State*>(self);
	}

	static iree_status_t ModuleDestroy(void* self) {
	delete FromModulePointer(self);
	return IREE_STATUS_OK;
	}

	static iree_string_view_t ModuleName(void* self) {
	auto* module = FromModulePointer(self);
	return iree_make_cstring_view(module->name_);
	}

	static iree_vm_module_signature_t ModuleSignature(void* self) {
	auto* module = FromModulePointer(self);
	iree_vm_module_signature_t signature = {0};
	signature.import_function_count = 0;
	signature.export_function_count = module->dispatch_table_.size();
	signature.internal_function_count = 0;
	return signature;
	}

	static iree_status_t ModuleGetFunction(
	void* self, iree_vm_function_linkage_t linkage, int32_t ordinal,
	iree_vm_function_t* out_function, iree_string_view_t* out_name,
	iree_vm_function_signature_t* out_signature) {
	if (out_function) {
	std::memset(out_function, 0, sizeof(*out_function));
	}
	if (out_name) {
	out_name->data = nullptr;
	out_name->size = 0;
	}
	if (out_signature) {
	std::memset(out_signature, 0, sizeof(*out_signature));
	}
	auto* module = FromModulePointer(self);
	if (ordinal < 0 \|\| ordinal > module->dispatch_table_.size()) {
	return IREE_STATUS_INVALID_ARGUMENT;
	}
	if (out_function) {
	out_function->module = module->interface();
	out_function->linkage = IREE_VM_FUNCTION_LINKAGE_EXPORT;
	out_function->ordinal = ordinal;
	}
	if (out_name) {
	const auto& dispatch_function = module->dispatch_table_[ordinal];
	*out_name = iree_make_cstring_view(dispatch_function.name);
	}
	return IREE_STATUS_OK;
	}

	static iree_status_t ModuleLookupFunction(void* self,
	iree_vm_function_linkage_t linkage,
	iree_string_view_t name,
	iree_vm_function_t* out_function) {
	if (!out_function) return IREE_STATUS_INVALID_ARGUMENT;
	std::memset(out_function, 0, sizeof(*out_function));
	if (!name.data \|\| !name.size) return IREE_STATUS_INVALID_ARGUMENT;

	auto* module = FromModulePointer(self);
	out_function->module = module->interface();
	out_function->linkage = IREE_VM_FUNCTION_LINKAGE_EXPORT;
	for (int i = 0; i < module->dispatch_table_.size(); ++i) {
	if (iree_string_view_compare(
	name, iree_make_cstring_view(module->dispatch_table_[i].name)) ==
	0) {
	out_function->ordinal = i;
	return IREE_STATUS_OK;
	}
	}
	return IREE_STATUS_NOT_FOUND;
	}

	static iree_status_t ModuleAllocState(
	void* self, iree_allocator_t allocator,
	iree_vm_module_state_t** out_module_state) {
	if (!out_module_state) return IREE_STATUS_INVALID_ARGUMENT;
	*out_module_state = nullptr;

	auto* module = FromModulePointer(self);
	auto module_state_or = module->CreateState(allocator);
	if (!module_state_or.ok()) {
	return ToApiStatus(module_state_or.status());
	}
	auto module_state = std::move(module_state_or).ValueOrDie();

	*out_module_state =
	reinterpret_cast<iree_vm_module_state_t*>(module_state.release());
	return IREE_STATUS_OK;
	}

	static iree_status_t ModuleFreeState(void* self,
	iree_vm_module_state_t* module_state) {
	if (!module_state) return IREE_STATUS_INVALID_ARGUMENT;
	delete FromStatePointer(module_state);
	return IREE_STATUS_OK;
	}

	static iree_status_t ModuleResolveImport(void* self,
	iree_vm_module_state_t* module_state,
	int32_t ordinal,
	iree_vm_function_t function) {
	// C++ API does not yet support imports.
	return IREE_STATUS_FAILED_PRECONDITION;
	}

	static iree_status_t ModuleExecute(void* self, iree_vm_stack_t* stack,
	iree_vm_stack_frame_t* frame,
	iree_vm_execution_result_t* out_result) {
	if (!out_result) return IREE_STATUS_INVALID_ARGUMENT;
	std::memset(out_result, 0, sizeof(*out_result));
	if (!stack \|\| !frame) return IREE_STATUS_INVALID_ARGUMENT;
	int32_t ordinal = frame->function.ordinal;
	auto* module = FromModulePointer(self);
	if (ordinal < 0 \|\| ordinal > module->dispatch_table_.size()) {
	return IREE_STATUS_INVALID_ARGUMENT;
	}
	const auto& info = module->dispatch_table_[ordinal];
	auto* state = FromStatePointer(frame->module_state);
	auto status = info.call(info.ptr, state, stack, frame, out_result);
	if (!status.ok()) {
	status = iree::Annotate(
	status,
	absl::StrCat("while executing ", module->name_, ".", info.name));
	return ToApiStatus(status);
	}
	return IREE_STATUS_OK;
	}

	const char* name_;
	const iree_allocator_t allocator_;
	iree_vm_module_t interface_;

	const absl::Span<const NativeFunction<State>> dispatch_table_;
	};

	} // namespace vm
	} // namespace iree

	#endif // IREE_VM_MODULE_ABI_CC_H_