diff options
Diffstat (limited to 'lib/cmake/llvm/TensorFlowCompile.cmake')
| -rw-r--r-- | lib/cmake/llvm/TensorFlowCompile.cmake | 38 |
1 files changed, 38 insertions, 0 deletions
diff --git a/lib/cmake/llvm/TensorFlowCompile.cmake b/lib/cmake/llvm/TensorFlowCompile.cmake new file mode 100644 index 0000000..a8ba56e --- /dev/null +++ b/lib/cmake/llvm/TensorFlowCompile.cmake @@ -0,0 +1,38 @@ +# Run the tensorflow compiler (saved_model_cli) on the saved model in the +# ${model} directory, looking for the ${tag_set} tag set, and the SignatureDef +# ${signature_def_key}. +# Produce a pair of files called ${fname}.h and ${fname}.o in the +# ${CMAKE_CURRENT_BINARY_DIR}. The generated header will define a C++ class +# called ${cpp_class} - which may be a namespace-qualified class name. +function(tfcompile model tag_set signature_def_key fname cpp_class) + if (IS_ABSOLUTE ${model}) + set(LLVM_ML_MODELS_ABSOLUTE ${model}) + else() + set(LLVM_ML_MODELS_ABSOLUTE + ${CMAKE_CURRENT_SOURCE_DIR}/${model}) + endif() + + set(prefix ${CMAKE_CURRENT_BINARY_DIR}/${fname}) + set(obj_file ${prefix}.o) + set(hdr_file ${prefix}.h) + add_custom_command(OUTPUT ${obj_file} ${hdr_file} + COMMAND "XLA_FLAGS=\"--xla_cpu_multi_thread_eigen=false\"" ${TENSORFLOW_AOT_COMPILER} aot_compile_cpu + --dir ${LLVM_ML_MODELS_ABSOLUTE} + --tag_set ${tag_set} + --signature_def_key ${signature_def_key} + --output_prefix ${prefix} + --cpp_class ${cpp_class} + --target_triple ${LLVM_HOST_TRIPLE} + ) + + # Aggregate the objects so that results of different tfcompile calls may be + # grouped into one target. + set(GENERATED_OBJS ${GENERATED_OBJS} ${obj_file} PARENT_SCOPE) + set_source_files_properties(${obj_file} PROPERTIES + GENERATED 1 EXTERNAL_OBJECT 1) + + set(GENERATED_HEADERS ${GENERATED_HEADERS} ${hdr_file} PARENT_SCOPE) + set_source_files_properties(${hdr_file} PROPERTIES + GENERATED 1) + +endfunction() |