diff --git a/src/target/source/interface_rust.cc b/src/target/source/interface_rust.cc new file mode 100644 index 000000000000..50e0f21ae2cf --- /dev/null +++ b/src/target/source/interface_rust.cc @@ -0,0 +1,381 @@ +/* + * 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 + +#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 " + << tvm::relay::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