blob: a8ba56e67d1f644ccf80058e4259275c11757131 [file] [log] [blame]
Olivier Deprezf4ef2d02021-04-20 13:36:24 +02001# 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.
7function(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
38endfunction()