1# Ensure the ${model} is available at ${final_path}.
2#
3function(tfgetmodel model final_path)
4  if (IS_ABSOLUTE ${model})
5    set(${final_path} ${model} PARENT_SCOPE)
6  else()
7    set(${final_path}
8      ${CMAKE_CURRENT_SOURCE_DIR}/${model} PARENT_SCOPE)
9  endif()
10endfunction()
11
12# Run the tensorflow compiler (saved_model_cli) on the saved model in the
13# ${model} directory, looking for the ${tag_set} tag set, and the SignatureDef
14# ${signature_def_key}.
15# Produce a pair of files called ${fname}.h and  ${fname}.o in the
16# ${CMAKE_CURRENT_BINARY_DIR}. The generated header will define a C++ class
17# called ${cpp_class} - which may be a namespace-qualified class name.
18function(tfcompile model tag_set signature_def_key fname cpp_class)
19  set(prefix ${CMAKE_CURRENT_BINARY_DIR}/${fname})
20  set(obj_file ${prefix}.o)
21  set(hdr_file ${prefix}.h)
22  string(TOUPPER ${fname} fname_allcaps)
23  set(override_header ${LLVM_OVERRIDE_MODEL_HEADER_${fname_allcaps}})
24  set(override_object ${LLVM_OVERRIDE_MODEL_OBJECT_${fname_allcaps}})
25  if (EXISTS "${override_header}" AND EXISTS "${override_object}")
26    configure_file(${override_header} ${hdr_file} COPYONLY)
27    configure_file(${override_object} ${obj_file} COPYONLY)
28    message("Using provided header "
29      ${hdr_file} " and object "   ${obj_file}
30      " files for model " ${model})
31  else()
32    tfgetmodel(${model} LLVM_ML_MODELS_ABSOLUTE)
33    message("Using model at " ${LLVM_ML_MODELS_ABSOLUTE})
34    add_custom_command(OUTPUT ${obj_file} ${hdr_file}
35      COMMAND "XLA_FLAGS=\"--xla_cpu_multi_thread_eigen=false\"" ${TENSORFLOW_AOT_COMPILER} aot_compile_cpu
36            --dir ${LLVM_ML_MODELS_ABSOLUTE}
37            --tag_set ${tag_set}
38            --signature_def_key ${signature_def_key}
39            --output_prefix ${prefix}
40            --cpp_class ${cpp_class}
41            --target_triple ${LLVM_HOST_TRIPLE}
42    )
43  endif()
44
45  # Aggregate the objects so that results of different tfcompile calls may be
46  # grouped into one target.
47  set(GENERATED_OBJS ${GENERATED_OBJS} ${obj_file} PARENT_SCOPE)
48  set_source_files_properties(${obj_file} PROPERTIES
49    GENERATED 1 EXTERNAL_OBJECT 1)
50
51  set(GENERATED_HEADERS ${GENERATED_HEADERS} ${hdr_file} PARENT_SCOPE)
52  set_source_files_properties(${hdr_file} PROPERTIES
53    GENERATED 1)
54
55endfunction()
56