From bd6c88af904f9271737ca6df092d4992d06f6b81 Mon Sep 17 00:00:00 2001 From: Chris Sidebottom Date: Thu, 5 Jan 2023 13:58:51 +0000 Subject: [PATCH 1/2] [Rust] Initial generator for Embedded Interface This is based on the older API in the RFC, we're landing this now so as to upstream it iteratively and allow others to see it in action :smile_cat: There'll be a follow up, as noted in #13705 to update this to the fully idiomatic Rust API Co-authored-by: Ashutosh Parkhi --- src/target/source/interface_rust.cc | 380 +++++++++ .../cpp/target/source/interface_rust_test.cc | 748 ++++++++++++++++++ 2 files changed, 1128 insertions(+) create mode 100644 src/target/source/interface_rust.cc create mode 100644 tests/cpp/target/source/interface_rust_test.cc diff --git a/src/target/source/interface_rust.cc b/src/target/source/interface_rust.cc new file mode 100644 index 000000000000..636ca38c3cff --- /dev/null +++ b/src/target/source/interface_rust.cc @@ -0,0 +1,380 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 + * + * http://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. + */ + +/*! + * \file interface_rust.cc + * \brief Generates a Rust interface header for a given modules inputs and outputs + * which works on top of the C interface API + */ + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "../../relay/backend/name_transforms.h" +#include "codegen_params.h" + +namespace tvm { +namespace codegen { + +using runtime::PackedFunc; +using namespace tvm::relay::backend; +using namespace tvm::runtime; + +class InterfaceRustNode : public runtime::ModuleNode { + public: + InterfaceRustNode(std::string module_name, Array input_names, Array output_names, + Array pools, + Map io_pool_allocations, + Array devices, int workspace_size, + Map> inputs, + Map> outputs) + : module_name_(module_name), + inputs_(inputs), + outputs_(outputs), + devices_(devices), + input_names_(input_names), + output_names_(output_names), + pools_(FilterExternalPools(pools)), + io_pool_allocations_(io_pool_allocations), + workspace_size_(workspace_size) { + ICHECK(io_pool_allocations_.empty()) << "Workspace Memory Pools IO unsupported"; + } + const char* type_key() const final { return "h"; } + + std::string GetSource(const std::string& format) final { + std::stringstream code; + + EmitBrief(code, "Input tensors"); + EmitDataStruct(code, "inputs", inputs_, input_names_); + EmitBrief(code, "Output tensors"); + EmitDataStruct(code, "outputs", outputs_, output_names_); + + if (!devices_.empty()) { + EmitBrief(code, "Device context pointers"); + EmitDeviceStruct(code, "devices", devices_); + } + + if (!pools_.empty()) { + EmitBrief(code, "Workspace pools"); + EmitWorkspacePoolsStruct(code); + } + + EmitRustRunFunction(code); + EmitCRunFunction(code); + + EmitIntegerValueConst(code, "Workspace size", "WORKSPACE_SIZE", workspace_size_); + EmitMemoryPools(code); + + return code.str(); + } + + PackedFunc GetFunction(const std::string& name, const ObjectPtr& sptr_to_self) final { + return PackedFunc(); + } + + private: + constexpr static const char* _macro_workspace_pool_size_postfix = "_WORKSPACE_POOL_SIZE"; + constexpr static const char* _macro_constant_pool_size_postfix = "_CONSTANT_POOL_SIZE"; + constexpr static const char* _macro_constant_pool_data_postfix = "_constant_pool_data"; + + void EmitBrief(std::stringstream& code_stream, const std::string& description) { + code_stream << "/// " << description << " for TVM module \"" << module_name_ << "\"\n"; + } + + void EmitMemoryPools(std::stringstream& code) { + for (const tir::usmp::AllocatedPoolInfo pool : pools_) { + String pool_name = pool->pool_info->pool_name; + Integer pool_size = pool->allocated_size; + if (const auto* pool_info = pool->pool_info.as()) { + EmitConstantPool(code, SanitizeName(pool_name) + " initialization data", pool_info); + } else { + EmitIntegerValueConst(code, SanitizeName(pool_name) + " size", + SanitizeName(pool_name) + _macro_workspace_pool_size_postfix, + pool_size->value); + } + } + } + + void EmitConstantPool(std::stringstream& code_, const std::string& brief_description, + const ConstantPoolInfoNode* pool_info) { + EmitBrief(code_, brief_description); + std::string const_name_prefixed = ToRustConstantStyle(SanitizeName(pool_info->pool_name)); + std::string macro_name_prefixed = ToRustMacroStyle(SanitizeName(pool_info->pool_name)); + + if (pool_info->constant_info_array.size() > 0) { + std::vector const_info_vec(pool_info->constant_info_array.begin(), + pool_info->constant_info_array.end()); + std::sort(const_info_vec.begin(), const_info_vec.end(), + [](const ConstantInfo& a, const ConstantInfo& b) { + return a->byte_offset->value < b->byte_offset->value; + }); + int64_t accumulated_pool_len = + const_info_vec.back()->byte_offset.IntValue() + + runtime::GetDataSize(*const_info_vec.back()->data.operator->()); + const auto& accumulated_pool = runtime::NDArray::Empty( + {accumulated_pool_len}, DataType::UInt(8), const_info_vec.back()->data->device); + for (const auto& const_info : const_info_vec) { + const auto& data = const_info->data; + const auto& offs = const_info->byte_offset; + data.CopyToBytes(static_cast(accumulated_pool->data) + offs.IntValue(), + runtime::GetDataSize(*data.operator->())); + } + + code_ << "pub const " << const_name_prefixed << _macro_constant_pool_size_postfix + << ": u32 = " << accumulated_pool_len << ";\n"; + code_ << "#[macro_export]\n" + << "macro_rules! " << macro_name_prefixed << _macro_constant_pool_data_postfix << " {\n" + << " () => {[\n"; + codegen::NDArrayDataToC(accumulated_pool, 8, code_, "\\\n"); + code_ << " ]};\n"; + code_ << "}\n"; + } else { + LOG(FATAL) << "No constant data in constant pool found " + << PrettyPrint(GetRef(pool_info)); + } + } + + void EmitStruct(std::stringstream& code_stream, const std::string& struct_name, + std::unordered_map> properties, + std::vector property_names_ordered) { + std::unordered_map> sanitized_properties; + std::vector sanitized_property_names_ordered; + + for (const std::string& property_name : property_names_ordered) { + std::string sanitized_property = SanitizeName(property_name); + std::pair property_values = properties.at(property_name); + ICHECK(std::find(sanitized_property_names_ordered.begin(), + sanitized_property_names_ordered.end(), + sanitized_property) == sanitized_property_names_ordered.end()) + << "Sanitized input tensor name clash" << sanitized_property; + + sanitized_properties.emplace(sanitized_property, property_values); + sanitized_property_names_ordered.push_back(sanitized_property); + } + std::reverse(sanitized_property_names_ordered.begin(), sanitized_property_names_ordered.end()); + + code_stream << "#[repr(C)]\n"; + code_stream << "pub struct " << ToRustStructStyle(struct_name) << " {\n"; + for (const std::string& property_name : sanitized_property_names_ordered) { + code_stream << " " << property_name << ": *mut ::std::os::raw::c_void,\n"; + } + code_stream << "}\n\n" + << "impl " << ToRustStructStyle(struct_name) << " {\n" + << " pub fn new <'a>(\n"; + for (const std::string& property_name : sanitized_property_names_ordered) { + std::string rust_data_type = sanitized_properties.at(property_name).first; + code_stream << " " << property_name << ": " << rust_data_type << ",\n"; + } + code_stream << " ) -> Self {\n" + << " Self {\n"; + for (const std::string& property_name : sanitized_property_names_ordered) { + std::string struct_conversion = sanitized_properties.at(property_name).second; + code_stream << " " << property_name << ": " << property_name << struct_conversion + << ",\n"; + } + code_stream << " }\n" + << " }\n" + << "}\n"; + } + + void EmitWorkspacePoolsStruct(std::stringstream& code_stream) { + std::unordered_map> struct_properties; + std::vector property_names_ordered; + for (const tir::usmp::AllocatedPoolInfo pool : pools_) { + int64_t allocated_size = pool->allocated_size.IntValue(); + std::string rust_type = "&mut [u8; " + std::to_string(allocated_size) + "]"; + struct_properties.emplace( + pool->pool_info->pool_name, + std::make_pair(rust_type, ".as_ptr() as *mut ::std::os::raw::c_void")); + property_names_ordered.push_back(pool->pool_info->pool_name); + } + std::reverse(property_names_ordered.begin(), property_names_ordered.end()); + + EmitStruct(code_stream, "workspace_pools", struct_properties, property_names_ordered); + } + + std::string DTypeToRust(std::string dtype) { + std::string width; + std::copy_if(dtype.begin(), dtype.end(), std::back_inserter(width), ::isdigit); + return dtype[0] + width; + } + + std::string NumElements(std::string dtype, int64_t size) { + std::string width; + std::copy_if(dtype.begin(), dtype.end(), std::back_inserter(width), ::isdigit); + return std::to_string(size * 8 / std::stoi(width)); + } + + void EmitDataStruct(std::stringstream& code_stream, const std::string& struct_name, + Map> properties, + Array property_names) { + std::unordered_map> struct_properties; + for (const auto& property : properties) { + Map values = property.second; + std::string dtype = Downcast(values.Get("dtype")); + int64_t size = Downcast(values.Get("size")).IntValue(); + std::string rust_dtype = DTypeToRust(dtype); + std::string num_elements = NumElements(dtype, size); + + struct_properties.emplace(property.first, + std::make_pair("&mut [" + rust_dtype + "; " + num_elements + "]", + ".as_ptr() as *mut ::std::os::raw::c_void")); + } + + std::vector property_names_ordered; + for (const String& property_name : property_names) { + property_names_ordered.push_back(property_name); + } + std::reverse(property_names_ordered.begin(), property_names_ordered.end()); + + EmitStruct(code_stream, struct_name, struct_properties, property_names_ordered); + } + + void EmitDeviceStruct(std::stringstream& code_stream, const std::string& struct_name, + Array devices) { + std::unordered_map> struct_properties; + std::vector property_names_ordered; + for (const auto& device : devices) { + struct_properties.emplace(device, std::make_pair("*mut ::std::os::raw::c_void", "")); + property_names_ordered.push_back(device); + } + std::reverse(property_names_ordered.begin(), property_names_ordered.end()); + + EmitStruct(code_stream, struct_name, struct_properties, property_names_ordered); + } + + void EmitIntegerValueConst(std::stringstream& code_stream, const std::string& brief_description, + const std::string& macro_name, int macro_value) { + EmitBrief(code_stream, brief_description); + std::string macro_name_prefixed = ToRustConstantStyle(macro_name); + code_stream << "pub const " << macro_name_prefixed << ": usize = " << macro_value << ";\n"; + } + + void EmitCRunFunction(std::stringstream& code_stream) { + std::string run_function = ToCVariableStyle(PrefixGeneratedName({module_name_, "run"})); + code_stream << "extern \"C\" {\n" + << " pub fn " << run_function << "(\n" + << " inputs: *mut Inputs,\n" + << " outputs: *mut Outputs,\n"; + if (!pools_.empty()) { + code_stream << " workspace_pools: *mut WorkspacePools,\n"; + } + if (!devices_.empty()) { + code_stream << " devices: *mut Devices,\n"; + } + code_stream << " ) -> i32;\n" + << "}\n"; + } + + void EmitRustRunFunction(std::stringstream& code_stream) { + std::string run_function = ToCVariableStyle(PrefixGeneratedName({module_name_, "run"})); + code_stream << "/// Entrypoint function for TVM module \"" << module_name_ << "\"\n" + << "/// # Arguments\n"; + if (io_pool_allocations_.empty()) { + code_stream << "/// * `inputs` - Input tensors for the module\n"; + code_stream << "/// * `outputs` - Output tensors for the module\n"; + } + + if (!pools_.empty()) { + code_stream << "/// * `workspace_pools` - Workspace memory pools for the module\n"; + } + if (!devices_.empty()) { + code_stream << "/// * `devices` - Device context pointers for the module\n"; + } + code_stream << "pub fn run(\n" + << " inputs: &mut Inputs,\n" + << " outputs: &mut Outputs,\n"; + if (!pools_.empty()) { + code_stream << " workspace_pools: &mut WorkspacePools,\n"; + } + if (!devices_.empty()) { + code_stream << " devices: &mut Devices,\n"; + } + code_stream << ") -> Result<(), ()> {\n" + << " unsafe {\n" + << " let ret = " << run_function << "(\n" + << " inputs,\n" + << " outputs,\n"; + if (!pools_.empty()) { + code_stream << " workspace_pools,\n"; + } + if (!devices_.empty()) { + code_stream << " devices,\n"; + } + code_stream << " );\n" + << " if ret == 0 {\n" + << " Ok(())\n" + << " } else {\n" + << " Err(())\n" + << " }\n" + << " }\n" + << "}\n\n"; + } + + void EmitRunFunction(std::stringstream& code_stream) { + EmitRustRunFunction(code_stream); + EmitCRunFunction(code_stream); + } + + Array FilterExternalPools( + const Array& pools) { + Array external_pools; + for (tir::usmp::AllocatedPoolInfo pool : pools) { + if (!pool->pool_info->is_internal) { + external_pools.push_back(pool); + } + } + return external_pools; + } + + std::string module_name_; + Map> inputs_; + Map> outputs_; + Array devices_; + Array input_names_; + Array output_names_; + Array pools_; + Map io_pool_allocations_; + int workspace_size_; +}; // namespace codegen + +runtime::Module InterfaceRustCreate(std::string module_name, + Map> inputs, + Map> outputs, + Array pools, + Map io_pool_allocations, + Array devices, Array input_names, + Array output_names, int workspace_size) { + auto n = + make_object(module_name, input_names, output_names, pools, + io_pool_allocations, devices, workspace_size, inputs, outputs); + return runtime::Module(n); +} + +TVM_REGISTER_GLOBAL("runtime.InterfaceRustCreate").set_body_typed(InterfaceRustCreate); + +} // namespace codegen +} // namespace tvm diff --git a/tests/cpp/target/source/interface_rust_test.cc b/tests/cpp/target/source/interface_rust_test.cc new file mode 100644 index 000000000000..841145b1eaff --- /dev/null +++ b/tests/cpp/target/source/interface_rust_test.cc @@ -0,0 +1,748 @@ + +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 + * + * http://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. + */ + +#include +#include +#include +#include +#include +#include + +using ::testing::ContainsRegex; +using ::testing::HasSubstr; + +namespace tvm { +namespace codegen { + +runtime::Module InterfaceRustCreate(std::string module_name, + Map> inputs, + Map> outputs, + Array pools, + Map io_pool_allocations, + Array devices, Array input_names, + Array output_names, int workspace_size); + +namespace { + +Map TestIO(String dtype, Integer size) { + return {{"dtype", dtype}, {"size", size}}; +} + +Map TestIO() { return TestIO("uint8", 100); } + +TEST(RustInterfaceAPI, ContainsRunFunction) { + std::stringstream run_function; + + run_function << "/// Entrypoint function for TVM module \"ultimate_cat_spotter\"\n" + << "/// # Arguments\n" + << "/// * `inputs` - Input tensors for the module\n" + << "/// * `outputs` - Output tensors for the module\n" + << "pub fn run(\n" + << " inputs: &mut Inputs,\n" + << " outputs: &mut Outputs,\n" + << ") -> Result<(), ()> {\n" + << " unsafe {\n" + << " let ret = tvmgen_ultimate_cat_spotter_run(\n" + << " inputs,\n" + << " outputs,\n" + << " );\n" + << " if ret == 0 {\n" + << " Ok(())\n" + << " } else {\n" + << " Err(())\n" + << " }\n" + << " }\n" + << "}\n" + << "\n" + << "extern \"C\" {\n" + << " pub fn tvmgen_ultimate_cat_spotter_run(\n" + << " inputs: *mut Inputs,\n" + << " outputs: *mut Outputs,\n" + << " ) -> i32;\n" + << "}\n"; + + runtime::Module test_module = + InterfaceRustCreate("ultimate_cat_spotter", {{"input", TestIO()}}, {{"output", TestIO()}}, {}, + {}, {}, {"input"}, {"output"}, 0); + std::string header_source = test_module->GetSource(); + ASSERT_THAT(header_source, HasSubstr(run_function.str())); +} + +TEST(RustInterfaceAPI, ContainsRunFunctionWithDevices) { + std::stringstream run_function; + + run_function << "/// Entrypoint function for TVM module \"ultimate_cat_spotter\"\n" + << "/// # Arguments\n" + << "/// * `inputs` - Input tensors for the module\n" + << "/// * `outputs` - Output tensors for the module\n" + << "/// * `devices` - Device context pointers for the module\n" + << "pub fn run(\n" + << " inputs: &mut Inputs,\n" + << " outputs: &mut Outputs,\n" + << " devices: &mut Devices,\n" + << ") -> Result<(), ()> {\n" + << " unsafe {\n" + << " let ret = tvmgen_ultimate_cat_spotter_run(\n" + << " inputs,\n" + << " outputs,\n" + << " devices,\n" + << " );\n" + << " if ret == 0 {\n" + << " Ok(())\n" + << " } else {\n" + << " Err(())\n" + << " }\n" + << " }\n" + << "}\n" + << "\n" + << "extern \"C\" {\n" + << " pub fn tvmgen_ultimate_cat_spotter_run(\n" + << " inputs: *mut Inputs,\n" + << " outputs: *mut Outputs,\n" + << " devices: *mut Devices,\n" + << " ) -> i32;\n" + << "}\n"; + + runtime::Module test_module = + InterfaceRustCreate("ultimate_cat_spotter", {{"input", TestIO()}}, {{"output", TestIO()}}, {}, + {}, {"device"}, {"input"}, {"output"}, 0); + std::string header_source = test_module->GetSource(); + + ASSERT_THAT(header_source, HasSubstr(run_function.str())); +} + +TEST(RustInterfaceAPI, ContainsRunFunctionWithWorkspacePools) { + std::stringstream run_function; + + run_function << "/// Entrypoint function for TVM module \"ultimate_cat_spotter\"\n" + << "/// # Arguments\n" + << "/// * `inputs` - Input tensors for the module\n" + << "/// * `outputs` - Output tensors for the module\n" + << "/// * `workspace_pools` - Workspace memory pools for the module\n" + << "pub fn run(\n" + << " inputs: &mut Inputs,\n" + << " outputs: &mut Outputs,\n" + << " workspace_pools: &mut WorkspacePools,\n" + << ") -> Result<(), ()> {\n" + << " unsafe {\n" + << " let ret = tvmgen_ultimate_cat_spotter_run(\n" + << " inputs,\n" + << " outputs,\n" + << " workspace_pools,\n" + << " );\n" + << " if ret == 0 {\n" + << " Ok(())\n" + << " } else {\n" + << " Err(())\n" + << " }\n" + << " }\n" + << "}\n" + << "\n" + << "extern \"C\" {\n" + << " pub fn tvmgen_ultimate_cat_spotter_run(\n" + << " inputs: *mut Inputs,\n" + << " outputs: *mut Outputs,\n" + << " workspace_pools: *mut WorkspacePools,\n" + << " ) -> i32;\n" + << "}\n"; + + PoolInfo pool_info = WorkspacePoolInfo("my_memory_pool", {}); + tir::usmp::AllocatedPoolInfo allocated_pool_info = + tir::usmp::AllocatedPoolInfo(pool_info, 100000); + runtime::Module test_module = + InterfaceRustCreate("ultimate_cat_spotter", {{"input", TestIO()}}, {{"output", TestIO()}}, + {allocated_pool_info}, {}, {}, {"input"}, {"output"}, 0); + std::string header_source = test_module->GetSource(); + + ASSERT_THAT(header_source, HasSubstr(run_function.str())); +} + +TEST(RustInterfaceAPI, ContainsRunFunctionWithWorkspaceAndConstantPools) { + std::stringstream run_function; + + run_function << "/// Entrypoint function for TVM module \"ultimate_cat_spotter\"\n" + << "/// # Arguments\n" + << "/// * `inputs` - Input tensors for the module\n" + << "/// * `outputs` - Output tensors for the module\n" + << "/// * `workspace_pools` - Workspace memory pools for the module\n" + << "pub fn run(\n" + << " inputs: &mut Inputs,\n" + << " outputs: &mut Outputs,\n" + << " workspace_pools: &mut WorkspacePools,\n" + << ") -> Result<(), ()> {\n"; + + PoolInfo pool_info = WorkspacePoolInfo("my_memory_pool", {}); + PoolInfo const_info = ConstantPoolInfo( + "my_constant_pool", {}, + {{"const1", 0, runtime::NDArray::Empty({1}, DataType::Int(32), {kDLCPU, 0})}, + {"const2", 16, runtime::NDArray::Empty({1}, DataType::Float(64), {kDLCPU, 0})}}); + tir::usmp::AllocatedPoolInfo allocated_pool_info = + tir::usmp::AllocatedPoolInfo(pool_info, 100000); + tir::usmp::AllocatedPoolInfo allocated_const_info = + tir::usmp::AllocatedPoolInfo(const_info, 100000); + runtime::Module test_module = InterfaceRustCreate( + "ultimate_cat_spotter", {{"input", TestIO()}}, {{"output", TestIO()}}, + {allocated_pool_info, allocated_const_info}, {}, {}, {"input"}, {"output"}, 0); + std::string header_source = test_module->GetSource(); + ASSERT_THAT(header_source, HasSubstr(run_function.str())); + ASSERT_THAT(header_source, HasSubstr("pub const MY_CONSTANT_POOL_CONSTANT_POOL_SIZE: u32 = 24;")); + ASSERT_THAT( + header_source, + ContainsRegex( + "#\\[macro_export\\]\\\n" + "macro_rules! my_constant_pool_constant_pool_data \\{\\\n" + " \\(\\) => \\{\\[\\\n" + " 0x\\w\\w, 0x\\w\\w, 0x\\w\\w, 0x\\w\\w, 0x\\w\\w, 0x\\w\\w, 0x\\w\\w, 0x\\w\\w, " + "0x\\w\\w, 0x\\w\\w, 0x\\w\\w, 0x\\w\\w, 0x\\w\\w, " + "0x\\w\\w, 0x\\w\\w, 0x\\w\\w, \\\\\\\n 0x\\w\\w, 0x\\w\\w, 0x\\w\\w, 0x\\w\\w, " + "0x\\w\\w, 0x\\w\\w, 0x\\w\\w, 0x\\w\\w\\\\\\\n" + " \\]\\};\\\n" + "}\\\n")); +} + +TEST(RustInterfaceAPI, ContainsRunFunctionWithWorkspacePoolsAndDevices) { + std::stringstream run_function; + + run_function << "/// Entrypoint function for TVM module \"ultimate_cat_spotter\"\n" + << "/// # Arguments\n" + << "/// * `inputs` - Input tensors for the module\n" + << "/// * `outputs` - Output tensors for the module\n" + << "/// * `workspace_pools` - Workspace memory pools for the module\n" + << "/// * `devices` - Device context pointers for the module\n" + << "pub fn run(\n" + << " inputs: &mut Inputs,\n" + << " outputs: &mut Outputs,\n" + << " workspace_pools: &mut WorkspacePools,\n" + << " devices: &mut Devices,\n" + << ") -> Result<(), ()> {\n" + << " unsafe {\n" + << " let ret = tvmgen_ultimate_cat_spotter_run(\n" + << " inputs,\n" + << " outputs,\n" + << " workspace_pools,\n" + << " devices,\n" + << " );\n" + << " if ret == 0 {\n" + << " Ok(())\n" + << " } else {\n" + << " Err(())\n" + << " }\n" + << " }\n" + << "}\n" + << "\n" + << "extern \"C\" {\n" + << " pub fn tvmgen_ultimate_cat_spotter_run(\n" + << " inputs: *mut Inputs,\n" + << " outputs: *mut Outputs,\n" + << " workspace_pools: *mut WorkspacePools,\n" + << " devices: *mut Devices,\n" + << " ) -> i32;\n" + << "}\n"; + + PoolInfo pool_info = WorkspacePoolInfo("my_memory_pool", {}); + tir::usmp::AllocatedPoolInfo allocated_pool_info = + tir::usmp::AllocatedPoolInfo(pool_info, 100000); + runtime::Module test_module = + InterfaceRustCreate("ultimate_cat_spotter", {{"input", TestIO()}}, {{"output", TestIO()}}, + {allocated_pool_info}, {}, {"device"}, {"input"}, {"output"}, 0); + std::string header_source = test_module->GetSource(); + + ASSERT_THAT(header_source, HasSubstr(run_function.str())); +} + +TEST(RustInterfaceAPI, ContainsRunFunctionWithWorkspaceIO_Unsupported) { + std::stringstream run_function_with_map_functions; + PoolInfo pool_info = WorkspacePoolInfo("my_memory_pool", {}); + tir::usmp::AllocatedPoolInfo allocated_pool_info = + tir::usmp::AllocatedPoolInfo(pool_info, 100000); + tir::usmp::PoolAllocation pool_allocation_input{pool_info, 1000}; + tir::usmp::PoolAllocation pool_allocation_output{pool_info, 2000}; + + ASSERT_THROW( + InterfaceRustCreate("ultimate_cat_spotter", {{"input", TestIO()}}, {{"output", TestIO()}}, + {allocated_pool_info}, + {{"input", pool_allocation_input}, {"output", pool_allocation_output}}, + {}, {"input"}, {"output"}, 0), + InternalError); +} + +TEST(RustInterfaceAPI, ContainsInputStructSingle) { + std::stringstream input_struct; + + input_struct << "/// Input tensors for TVM module \"ultimate_cat_spotter\"\n" + << "#[repr(C)]\n" + << "pub struct Inputs {\n" + << " input: *mut ::std::os::raw::c_void,\n" + << "}\n" + << "\n" + << "impl Inputs {\n" + << " pub fn new <'a>(\n" + << " input: &mut [u8; 100],\n" + << " ) -> Self {\n" + << " Self {\n" + << " input: input.as_ptr() as *mut ::std::os::raw::c_void,\n" + << " }\n" + << " }\n" + << "}\n"; + + runtime::Module test_module = + InterfaceRustCreate("ultimate_cat_spotter", {{"input", TestIO()}}, {{"output", TestIO()}}, {}, + {}, {}, {"input"}, {"output"}, 0); + std::string header_source = test_module->GetSource(); + + ASSERT_THAT(header_source, HasSubstr(input_struct.str())); +} + +TEST(RustInterfaceAPI, ContainsInputStructMany) { + std::stringstream input_struct; + + input_struct << "/// Input tensors for TVM module \"ultimate_cat_spotter\"\n" + << "#[repr(C)]\n" + << "pub struct Inputs {\n" + << " input1: *mut ::std::os::raw::c_void,\n" + << " input2: *mut ::std::os::raw::c_void,\n" + << " input3: *mut ::std::os::raw::c_void,\n" + << "}\n" + << "\n" + << "impl Inputs {\n" + << " pub fn new <'a>(\n" + << " input1: &mut [f32; 15],\n" + << " input2: &mut [u8; 120],\n" + << " input3: &mut [i32; 45],\n" + << " ) -> Self {\n" + << " Self {\n" + << " input1: input1.as_ptr() as *mut ::std::os::raw::c_void,\n" + << " input2: input2.as_ptr() as *mut ::std::os::raw::c_void,\n" + << " input3: input3.as_ptr() as *mut ::std::os::raw::c_void,\n" + << " }\n" + << " }\n" + << "}\n"; + + runtime::Module test_module = InterfaceRustCreate("ultimate_cat_spotter", + {{"input1", TestIO("float32", 60)}, + {"input2", TestIO("uint8", 120)}, + {"input3", TestIO("int32", 180)}}, + {{"output", TestIO()}}, {}, {}, {}, + {"input1", "input2", "input3"}, {"output"}, 0); + std::string header_source = test_module->GetSource(); + ASSERT_THAT(header_source, HasSubstr(input_struct.str())); +} + +TEST(RustInterfaceAPI, ContainsInputStructSanitised) { + std::stringstream input_struct; + + input_struct << "/// Input tensors for TVM module \"ultimate_cat_spotter\"\n" + << "#[repr(C)]\n" + << "pub struct Inputs {\n" + << " input_1: *mut ::std::os::raw::c_void,\n" + << " input_2: *mut ::std::os::raw::c_void,\n" + << "}\n" + << "\n" + << "impl Inputs {\n" + << " pub fn new <'a>(\n" + << " input_1: &mut [u8; 100],\n" + << " input_2: &mut [u8; 100],\n" + << " ) -> Self {\n" + << " Self {\n" + << " input_1: input_1.as_ptr() as *mut ::std::os::raw::c_void,\n" + << " input_2: input_2.as_ptr() as *mut ::std::os::raw::c_void,\n" + << " }\n" + << " }\n" + << "}\n"; + + runtime::Module test_module = InterfaceRustCreate( + "ultimate_cat_spotter", {{"input+1", TestIO()}, {"input+2", TestIO()}}, + {{"output", TestIO()}}, {}, {}, {}, {"input+1", "input+2"}, {"output"}, 0); + std::string header_source = test_module->GetSource(); + + ASSERT_THAT(header_source, HasSubstr(input_struct.str())); +} + +TEST(RustInterfaceAPI, ContainsInputStructClash) { + runtime::Module test_module = + InterfaceRustCreate("ultimate_cat_spotter", {{"input+", TestIO()}, {"input-", TestIO()}}, + {{"output", TestIO()}}, {}, {}, {}, {"input+", "input-"}, {"output"}, 0); + ASSERT_THROW(test_module->GetSource(), InternalError); +} + +TEST(RustInterfaceAPI, ContainsOutputStructSingle) { + std::stringstream output_struct; + + output_struct << "/// Output tensors for TVM module \"ultimate_cat_spotter\"\n" + << "#[repr(C)]\n" + << "pub struct Outputs {\n" + << " output: *mut ::std::os::raw::c_void,\n" + << "}\n" + << "\n" + << "impl Outputs {\n" + << " pub fn new <'a>(\n" + << " output: &mut [u8; 100],\n" + << " ) -> Self {\n" + << " Self {\n" + << " output: output.as_ptr() as *mut ::std::os::raw::c_void,\n" + << " }\n" + << " }\n" + << "}\n"; + + runtime::Module test_module = + InterfaceRustCreate("ultimate_cat_spotter", {{"input", TestIO()}}, {{"output", TestIO()}}, {}, + {}, {}, {"input"}, {"output"}, 0); + std::string header_source = test_module->GetSource(); + + ASSERT_THAT(header_source, HasSubstr(output_struct.str())); +} + +TEST(RustInterfaceAPI, ContainsOutputStructMany) { + std::stringstream output_struct; + output_struct << "/// Output tensors for TVM module \"ultimate_cat_spotter\"\n" + << "#[repr(C)]\n" + << "pub struct Outputs {\n" + << " output1: *mut ::std::os::raw::c_void,\n" + << " output2: *mut ::std::os::raw::c_void,\n" + << " output3: *mut ::std::os::raw::c_void,\n" + << "}\n" + << "\n" + << "impl Outputs {\n" + << " pub fn new <'a>(\n" + << " output1: &mut [i32; 25],\n" + << " output2: &mut [f32; 19],\n" + << " output3: &mut [u64; 2],\n" + << " ) -> Self {\n" + << " Self {\n" + << " output1: output1.as_ptr() as *mut ::std::os::raw::c_void,\n" + << " output2: output2.as_ptr() as *mut ::std::os::raw::c_void,\n" + << " output3: output3.as_ptr() as *mut ::std::os::raw::c_void,\n" + << " }\n" + << " }\n" + << "}\n"; + + runtime::Module test_module = + InterfaceRustCreate("ultimate_cat_spotter", {{"input", TestIO()}}, + {{"output1", TestIO("int32", 100)}, + {"output2", TestIO("float32", 76)}, + {"output3", TestIO("uint64", 16)}}, + {}, {}, {}, {"input"}, {"output1", "output2", "output3"}, 0); + std::string header_source = test_module->GetSource(); + + ASSERT_THAT(header_source, HasSubstr(output_struct.str())); +} + +TEST(RustInterfaceAPI, ContainsOutputStructSanitised) { + std::stringstream output_struct; + output_struct << "/// Output tensors for TVM module \"ultimate_cat_spotter\"\n" + << "#[repr(C)]\n" + << "pub struct Outputs {\n" + << " output_1: *mut ::std::os::raw::c_void,\n" + << " output_2: *mut ::std::os::raw::c_void,\n" + << "}\n" + << "\n" + << "impl Outputs {\n" + << " pub fn new <'a>(\n" + << " output_1: &mut [u8; 100],\n" + << " output_2: &mut [u8; 100],\n" + << " ) -> Self {\n" + << " Self {\n" + << " output_1: output_1.as_ptr() as *mut ::std::os::raw::c_void,\n" + << " output_2: output_2.as_ptr() as *mut ::std::os::raw::c_void,\n" + << " }\n" + << " }\n" + << "}\n"; + + runtime::Module test_module = + InterfaceRustCreate("ultimate_cat_spotter", {{"input", TestIO()}}, + {{"output+1", TestIO()}, {"output-2", TestIO()}}, {}, {}, {}, {"input"}, + {"output+1", "output-2"}, 0); + std::string header_source = test_module->GetSource(); + + ASSERT_THAT(header_source, HasSubstr(output_struct.str())); +} + +TEST(RustInterfaceAPI, ContainsOutputStructClash) { + runtime::Module test_module = InterfaceRustCreate( + "ultimate_cat_spotter", {{"input", TestIO()}}, {{"output+", TestIO()}, {"output-", TestIO()}}, + {}, {}, {}, {"input"}, {"output+", "output-"}, 0); + ASSERT_THROW(test_module->GetSource(), InternalError); +} + +TEST(RustInterfaceAPI, NoDeviceAPIStructIfNoDevices) { + std::stringstream device_struct; + device_struct << "/// Device context pointers for TVM module \"ultimate_cat_spotter\"\n" + << "#[repr(C)]\n" + << "pub struct Devices {\n" + << " device: *mut ::std::os::raw::c_void,\n" + << "}\n" + << "\n"; + + runtime::Module test_module = + InterfaceRustCreate("ultimate_cat_spotter", {{"input", TestIO()}}, {{"output", TestIO()}}, {}, + {}, {}, {"input"}, {"output"}, 0); + std::string header_source = test_module->GetSource(); + + ASSERT_THAT(header_source, Not(HasSubstr(device_struct.str()))); +} + +TEST(RustInterfaceAPI, ContainsDeviceStructSingle) { + std::stringstream device_struct; + + device_struct << "/// Device context pointers for TVM module \"ultimate_cat_spotter\"\n" + << "#[repr(C)]\n" + << "pub struct Devices {\n" + << " device: *mut ::std::os::raw::c_void,\n" + << "}\n" + << "\n" + << "impl Devices {\n" + << " pub fn new <'a>(\n" + << " device: *mut ::std::os::raw::c_void,\n" + << " ) -> Self {\n" + << " Self {\n" + << " device: device,\n" + << " }\n" + << " }\n" + << "}\n"; + + runtime::Module test_module = + InterfaceRustCreate("ultimate_cat_spotter", {{"input", TestIO()}}, {{"output", TestIO()}}, {}, + {}, {"device"}, {"input"}, {"output"}, 0); + std::string header_source = test_module->GetSource(); + + ASSERT_THAT(header_source, HasSubstr(device_struct.str())); +} + +TEST(RustInterfaceAPI, ContainsDeviceStructMany) { + std::stringstream device_struct; + + device_struct << "/// Device context pointers for TVM module \"ultimate_cat_spotter\"\n" + << "#[repr(C)]\n" + << "pub struct Devices {\n" + << " device1: *mut ::std::os::raw::c_void,\n" + << " device2: *mut ::std::os::raw::c_void,\n" + << "}\n" + << "\n" + << "impl Devices {\n" + << " pub fn new <'a>(\n" + << " device1: *mut ::std::os::raw::c_void,\n" + << " device2: *mut ::std::os::raw::c_void,\n" + << " ) -> Self {\n" + << " Self {\n" + << " device1: device1,\n" + << " device2: device2,\n" + << " }\n" + << " }\n" + << "}\n"; + + runtime::Module test_module = + InterfaceRustCreate("ultimate_cat_spotter", {{"input", TestIO()}}, {{"output", TestIO()}}, {}, + {}, {"device1", "device2"}, {"input"}, {"output"}, 0); + std::string header_source = test_module->GetSource(); + ASSERT_THAT(header_source, HasSubstr(device_struct.str())); +} + +TEST(RustInterfaceAPI, ContainsDeviceStructSanitised) { + std::stringstream device_struct; + + device_struct << "/// Device context pointers for TVM module \"ultimate_cat_spotter\"\n" + << "#[repr(C)]\n" + << "pub struct Devices {\n" + << " device_1: *mut ::std::os::raw::c_void,\n" + << " device_2: *mut ::std::os::raw::c_void,\n" + << "}\n" + << "\n" + << "impl Devices {\n" + << " pub fn new <'a>(\n" + << " device_1: *mut ::std::os::raw::c_void,\n" + << " device_2: *mut ::std::os::raw::c_void,\n" + << " ) -> Self {\n" + << " Self {\n" + << " device_1: device_1,\n" + << " device_2: device_2,\n" + << " }\n" + << " }\n" + << "}\n"; + + runtime::Module test_module = + InterfaceRustCreate("ultimate_cat_spotter", {{"input", TestIO()}}, {{"output", TestIO()}}, {}, + {}, {"device+1", "device+2"}, {"input"}, {"output"}, 0); + std::string header_source = test_module->GetSource(); + + ASSERT_THAT(header_source, HasSubstr(device_struct.str())); +} + +TEST(RustInterfaceAPI, ContainsDeviceStructClash) { + runtime::Module test_module = + InterfaceRustCreate("ultimate_cat_spotter", {{"input", TestIO()}}, {{"output", TestIO()}}, {}, + {}, {"device+", "device-"}, {"input"}, {"output"}, 0); + ASSERT_THROW(test_module->GetSource(), InternalError); +} + +TEST(RustInterfaceAPI, ContainsWorkspaceSize) { + runtime::Module test_module = + InterfaceRustCreate("ultimate_cat_spotter", {{"input", TestIO()}}, {{"output", TestIO()}}, {}, + {}, {}, {"input"}, {"output"}, 765432); + std::string header_source = test_module->GetSource(); + + ASSERT_THAT(header_source, + HasSubstr("/// Workspace size for TVM module \"ultimate_cat_spotter\"")); + + ASSERT_THAT(header_source, HasSubstr("pub const WORKSPACE_SIZE: usize = 765432;")); +} + +TEST(RustInterfaceAPI, ContainsWorkspacePoolStructSingle) { + PoolInfo pool_info = WorkspacePoolInfo("my_memory_pool", {}); + tir::usmp::AllocatedPoolInfo allocated_pool_info = + tir::usmp::AllocatedPoolInfo(pool_info, 100000); + + std::stringstream workspace_struct; + + workspace_struct + << "/// Workspace pools for TVM module \"ultimate_cat_spotter\"\n" + << "#[repr(C)]\n" + << "pub struct WorkspacePools {\n" + << " my_memory_pool: *mut ::std::os::raw::c_void,\n" + << "}\n" + << "\n" + << "impl WorkspacePools {\n" + << " pub fn new <'a>(\n" + << " my_memory_pool: &mut [u8; 100000],\n" + << " ) -> Self {\n" + << " Self {\n" + << " my_memory_pool: my_memory_pool.as_ptr() as *mut ::std::os::raw::c_void,\n" + << " }\n" + << " }\n" + << "}\n"; + + runtime::Module test_module = + InterfaceRustCreate("ultimate_cat_spotter", {{"input", TestIO()}}, {{"output", TestIO()}}, + {allocated_pool_info}, {}, {}, {"input"}, {"output"}, 0); + std::string header_source = test_module->GetSource(); + + ASSERT_THAT(header_source, HasSubstr(workspace_struct.str())); + + ASSERT_THAT(header_source, + HasSubstr("/// my_memory_pool size for TVM module \"ultimate_cat_spotter\"")); + + ASSERT_THAT(header_source, + HasSubstr("pub const MY_MEMORY_POOL_WORKSPACE_POOL_SIZE: usize = 100000;")); +} + +TEST(RustInterfaceAPI, ContainsWorkspacePoolStructMany) { + PoolInfo pool_info1 = WorkspacePoolInfo("my_memory_pool_1", {}); + tir::usmp::AllocatedPoolInfo allocated_pool_info1 = + tir::usmp::AllocatedPoolInfo(pool_info1, 100000); + PoolInfo pool_info2 = WorkspacePoolInfo("my_memory_pool_2", {}); + tir::usmp::AllocatedPoolInfo allocated_pool_info2 = + tir::usmp::AllocatedPoolInfo(pool_info2, 200000); + + std::stringstream workspace_struct; + + workspace_struct + << "/// Workspace pools for TVM module \"ultimate_cat_spotter\"\n" + << "#[repr(C)]\n" + << "pub struct WorkspacePools {\n" + << " my_memory_pool_1: *mut ::std::os::raw::c_void,\n" + << " my_memory_pool_2: *mut ::std::os::raw::c_void,\n" + << "}\n" + << "\n" + << "impl WorkspacePools {\n" + << " pub fn new <'a>(\n" + << " my_memory_pool_1: &mut [u8; 100000],\n" + << " my_memory_pool_2: &mut [u8; 200000],\n" + << " ) -> Self {\n" + << " Self {\n" + << " my_memory_pool_1: my_memory_pool_1.as_ptr() as *mut ::std::os::raw::c_void,\n" + << " my_memory_pool_2: my_memory_pool_2.as_ptr() as *mut ::std::os::raw::c_void,\n" + << " }\n" + << " }\n" + << "}\n"; + + runtime::Module test_module = InterfaceRustCreate( + "ultimate_cat_spotter", {{"input", TestIO()}}, {{"output", TestIO()}}, + {allocated_pool_info1, allocated_pool_info2}, {}, {}, {"input"}, {"output"}, 0); + std::string header_source = test_module->GetSource(); + + ASSERT_THAT(header_source, HasSubstr(workspace_struct.str())); + + ASSERT_THAT(header_source, + HasSubstr("/// my_memory_pool_1 size for TVM module \"ultimate_cat_spotter\"")); + + ASSERT_THAT(header_source, + HasSubstr("pub const MY_MEMORY_POOL_1_WORKSPACE_POOL_SIZE: usize = 100000;")); + + ASSERT_THAT(header_source, + HasSubstr("/// my_memory_pool_2 size for TVM module \"ultimate_cat_spotter\"")); + + ASSERT_THAT(header_source, + HasSubstr("pub const MY_MEMORY_POOL_2_WORKSPACE_POOL_SIZE: usize = 200000;")); +} + +TEST(RustInterfaceAPI, ContainsWorkspacePoolStructSanitized) { + PoolInfo pool_info = WorkspacePoolInfo("my_memory_pool+1", {}); + tir::usmp::AllocatedPoolInfo allocated_pool_info = + tir::usmp::AllocatedPoolInfo(pool_info, 100000); + + std::stringstream workspace_struct; + + workspace_struct + << "/// Workspace pools for TVM module \"ultimate_cat_spotter\"\n" + << "#[repr(C)]\n" + << "pub struct WorkspacePools {\n" + << " my_memory_pool_1: *mut ::std::os::raw::c_void,\n" + << "}\n" + << "\n" + << "impl WorkspacePools {\n" + << " pub fn new <'a>(\n" + << " my_memory_pool_1: &mut [u8; 100000],\n" + << " ) -> Self {\n" + << " Self {\n" + << " my_memory_pool_1: my_memory_pool_1.as_ptr() as *mut ::std::os::raw::c_void,\n" + << " }\n" + << " }\n" + << "}\n"; + + runtime::Module test_module = + InterfaceRustCreate("ultimate_cat_spotter", {{"input", TestIO()}}, {{"output", TestIO()}}, + {allocated_pool_info}, {}, {}, {"input"}, {"output"}, 0); + std::string header_source = test_module->GetSource(); + + ASSERT_THAT(header_source, HasSubstr(workspace_struct.str())); + + ASSERT_THAT(header_source, + HasSubstr("/// my_memory_pool_1 size for TVM module \"ultimate_cat_spotter\"")); + + ASSERT_THAT(header_source, + HasSubstr("pub const MY_MEMORY_POOL_1_WORKSPACE_POOL_SIZE: usize = 100000;")); +} + +TEST(RustInterfaceAPI, ContainsWorkspacePoolStructClash) { + PoolInfo pool_info1 = WorkspacePoolInfo("my_memory_pool+", {}); + tir::usmp::AllocatedPoolInfo allocated_pool_info1 = + tir::usmp::AllocatedPoolInfo(pool_info1, 100000); + PoolInfo pool_info2 = WorkspacePoolInfo("my_memory_pool-", {}); + tir::usmp::AllocatedPoolInfo allocated_pool_info2 = + tir::usmp::AllocatedPoolInfo(pool_info2, 200000); + + runtime::Module test_module = InterfaceRustCreate( + "ultimate_cat_spotter", {{"input", TestIO()}}, {{"output", TestIO()}}, + {allocated_pool_info1, allocated_pool_info2}, {}, {}, {"input"}, {"output"}, 0); + ASSERT_THROW(test_module->GetSource(), InternalError); +} + +} // namespace +} // namespace codegen +} // namespace tvm From b250c7b520c8868ee485486e3af9ecdfe61c1f07 Mon Sep 17 00:00:00 2001 From: Chris Sidebottom Date: Thu, 9 Feb 2023 09:30:23 +0000 Subject: [PATCH 2/2] Add missing header --- src/target/source/interface_rust.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/target/source/interface_rust.cc b/src/target/source/interface_rust.cc index 636ca38c3cff..50e0f21ae2cf 100644 --- a/src/target/source/interface_rust.cc +++ b/src/target/source/interface_rust.cc @@ -23,6 +23,7 @@ * which works on top of the C interface API */ +#include #include #include #include @@ -153,7 +154,7 @@ class InterfaceRustNode : public runtime::ModuleNode { code_ << "}\n"; } else { LOG(FATAL) << "No constant data in constant pool found " - << PrettyPrint(GetRef(pool_info)); + << tvm::relay::PrettyPrint(GetRef(pool_info)); } }