Olivier Deprez | f4ef2d0 | 2021-04-20 13:36:24 +0200 | [diff] [blame] | 1 | # Run the tensorflow compiler (saved_model_cli) on the saved model in the |
| 2 | # ${model} directory, looking for the ${tag_set} tag set, and the SignatureDef |
| 3 | # ${signature_def_key}. |
| 4 | # Produce a pair of files called ${fname}.h and ${fname}.o in the |
| 5 | # ${CMAKE_CURRENT_BINARY_DIR}. The generated header will define a C++ class |
| 6 | # called ${cpp_class} - which may be a namespace-qualified class name. |
| 7 | function(tfcompile model tag_set signature_def_key fname cpp_class) |
| 8 | if (IS_ABSOLUTE ${model}) |
| 9 | set(LLVM_ML_MODELS_ABSOLUTE ${model}) |
| 10 | else() |
| 11 | set(LLVM_ML_MODELS_ABSOLUTE |
| 12 | ${CMAKE_CURRENT_SOURCE_DIR}/${model}) |
| 13 | endif() |
| 14 | |
| 15 | set(prefix ${CMAKE_CURRENT_BINARY_DIR}/${fname}) |
| 16 | set(obj_file ${prefix}.o) |
| 17 | set(hdr_file ${prefix}.h) |
| 18 | add_custom_command(OUTPUT ${obj_file} ${hdr_file} |
| 19 | COMMAND "XLA_FLAGS=\"--xla_cpu_multi_thread_eigen=false\"" ${TENSORFLOW_AOT_COMPILER} aot_compile_cpu |
| 20 | --dir ${LLVM_ML_MODELS_ABSOLUTE} |
| 21 | --tag_set ${tag_set} |
| 22 | --signature_def_key ${signature_def_key} |
| 23 | --output_prefix ${prefix} |
| 24 | --cpp_class ${cpp_class} |
| 25 | --target_triple ${LLVM_HOST_TRIPLE} |
| 26 | ) |
| 27 | |
| 28 | # Aggregate the objects so that results of different tfcompile calls may be |
| 29 | # grouped into one target. |
| 30 | set(GENERATED_OBJS ${GENERATED_OBJS} ${obj_file} PARENT_SCOPE) |
| 31 | set_source_files_properties(${obj_file} PROPERTIES |
| 32 | GENERATED 1 EXTERNAL_OBJECT 1) |
| 33 | |
| 34 | set(GENERATED_HEADERS ${GENERATED_HEADERS} ${hdr_file} PARENT_SCOPE) |
| 35 | set_source_files_properties(${hdr_file} PROPERTIES |
| 36 | GENERATED 1) |
| 37 | |
| 38 | endfunction() |