def test_mangle_mod_name():
assert utils.mangle_module_name("default") == "tvmgen_default"
assert utils.mangle_module_name("ccompiler") == "tvmgen_ccompiler"
assert utils.mangle_module_name("1234"), "tvmgen_1234"
assert utils.mangle_module_name(""), "tvmgen"
assert utils.mangle_module_name(None), "tvmgen"
with pytest.raises(ValueError):
utils.mangle_module_name("\u018e")
utils.mangle_module_name("\xf1")
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
def mangle_name(mod_name, name):
mod_name = mangle_module_name(mod_name)
return mod_name + "_" + name