Grant Watson
06a0d63c43
* Sanitize names of input tensors in interface header Change-Id: I7f02a993887bf84316262cd2586a734a9079c338 * Update tensor name sanitizer tests to parameterize them. Change-Id: I157d8d8d607de2904285e403893f146e97b510d5 * Only test unpacked, C interface API, AOT case Change-Id: I9082ae32079a1a3924c06c7f26c757aafa46dec2
97 lines
3.8 KiB
Python
97 lines
3.8 KiB
Python
# 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.
|
|
|
|
"""Defines functions for generating a C interface header"""
|
|
|
|
# TODO: Currently the Interface API header is generated in Python but the source it references
|
|
# is generated in C++. These should be consolidated to generate both header and source in C++
|
|
# and avoid re-implementing logic, such as name sanitising, in the two different languages.
|
|
# See https://github.com/apache/tvm/issues/8792 .
|
|
|
|
import os
|
|
import re
|
|
|
|
from tvm.relay.backend.utils import mangle_module_name
|
|
|
|
|
|
def _emit_brief(header_file, module_name, description):
|
|
header_file.write("/*!\n")
|
|
header_file.write(f' * \\brief {description} for TVM module "{module_name}" \n')
|
|
header_file.write(" */\n")
|
|
|
|
|
|
def generate_c_interface_header(module_name, inputs, outputs, output_path):
|
|
"""Generates a C interface header for a given modules inputs and outputs
|
|
|
|
Parameters
|
|
----------
|
|
module_name : str
|
|
Name of the module to be used in defining structs and naming the header
|
|
inputs : list[str]
|
|
List of module input names to be placed in generated structs
|
|
outputs : list[str]
|
|
List of module output names to be placed in generated structs
|
|
output_path : str
|
|
Path to the output folder to generate the header into
|
|
|
|
Returns
|
|
-------
|
|
str :
|
|
Name of the generated file.
|
|
"""
|
|
mangled_name = mangle_module_name(module_name)
|
|
metadata_header = os.path.join(output_path, f"{mangled_name}.h")
|
|
with open(metadata_header, "w") as header_file:
|
|
header_file.write(
|
|
"#include <stdint.h>\n"
|
|
f"#ifndef {mangled_name.upper()}_H_\n"
|
|
f"#define {mangled_name.upper()}_H_\n"
|
|
)
|
|
|
|
_emit_brief(header_file, module_name, "Input tensor pointers")
|
|
header_file.write(f"struct {mangled_name}_inputs {{\n")
|
|
sanitized_names = []
|
|
for input_name in inputs:
|
|
sanitized_input_name = re.sub(r"\W", "_", input_name)
|
|
if sanitized_input_name in sanitized_names:
|
|
raise ValueError(f"Sanitized input tensor name clash: {sanitized_input_name}")
|
|
sanitized_names.append(sanitized_input_name)
|
|
header_file.write(f" void* {sanitized_input_name};\n")
|
|
header_file.write("};\n\n")
|
|
|
|
_emit_brief(header_file, module_name, "Output tensor pointers")
|
|
header_file.write(f"struct {mangled_name}_outputs {{\n")
|
|
for output_name in outputs:
|
|
header_file.write(f" void* {output_name};\n")
|
|
header_file.write("};\n\n")
|
|
|
|
header_file.write(
|
|
"/*!\n"
|
|
f' * \\brief entrypoint function for TVM module "{module_name}"\n'
|
|
" * \\param inputs Input tensors for the module \n"
|
|
" * \\param outputs Output tensors for the module \n"
|
|
" */\n"
|
|
f"int32_t {mangled_name}_run(\n"
|
|
f" struct {mangled_name}_inputs* inputs,\n"
|
|
f" struct {mangled_name}_outputs* outputs\n"
|
|
");\n"
|
|
)
|
|
|
|
header_file.write(f"#endif // {mangled_name.upper()}_H_\n")
|
|
|
|
return metadata_header
|