llvm/cmake/modules/TensorFlowCompile.cmake

   1 function(tf_get_absolute_path path base final_path)
   2   if (IS_ABSOLUTE ${path})
   3     set(${final_path} ${path} PARENT_SCOPE)
   4   else()
   5     set(${final_path} ${base}/${path} PARENT_SCOPE)
   6   endif()
   7 endfunction()
   8
   9 function(tf_get_model model final_path)
  10   string(FIND ${model} "http:" pos_http)
  11   string(FIND ${model} "https:" pos_https)
  12   if (${pos_http} EQUAL 0 OR ${pos_https} EQUAL 0)
  13     message("Downloading model " ${model})
  14     string(FIND ${model} "/" fname_start REVERSE)
  15     math(EXPR fname_start "${fname_start}+1")
  16     string(SUBSTRING ${model} ${fname_start}+1 -1 fname)
  17     message("Model archive: " ${fname})
  18     file(DOWNLOAD ${model} ${CMAKE_CURRENT_BINARY_DIR}/${fname})
  19     file(ARCHIVE_EXTRACT INPUT
  20       ${CMAKE_CURRENT_BINARY_DIR}/${fname}
  21       DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/${fname}_model)
  22     set(${final_path} ${CMAKE_CURRENT_BINARY_DIR}/${fname}_model/model PARENT_SCOPE)
  23   else()
  24     tf_get_absolute_path(${model} ${CMAKE_CURRENT_BINARY_DIR} model_path)
  25     set(${final_path} ${model_path} PARENT_SCOPE)
  26   endif()
  27 endfunction()
  28
  29 # Generate a mock model for tests.
  30 function(generate_mock_model model generate_mock_model_py config)
  31   tf_get_absolute_path(${model} ${CMAKE_CURRENT_BINARY_DIR} LLVM_ML_MODELS_ABSOLUTE)
  32   tf_get_absolute_path(${generate_mock_model_py} ${CMAKE_CURRENT_SOURCE_DIR} GENERATE_INLINING_MODEL_ABSOLUTE)
  33   tf_get_absolute_path(${config} ${CMAKE_CURRENT_SOURCE_DIR} LLVM_ML_MODEL_CONFIG_ABSOLUTE)
  34   message(WARNING "Autogenerated mock models should not be used in production builds.")
  35   execute_process(COMMAND python3
  36     ${GENERATE_INLINING_MODEL_ABSOLUTE}
  37     ${LLVM_ML_MODEL_CONFIG_ABSOLUTE}
  38     ${LLVM_ML_MODELS_ABSOLUTE}-autogenerated
  39     WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
  40   )
  41 endfunction()
  42
  43 # Run the tensorflow compiler (saved_model_cli) on the saved model in the
  44 # ${model} directory, looking for the ${tag_set} tag set, and the SignatureDef
  45 # ${signature_def_key}.
  46 # Produce a pair of files called ${fname}.h and  ${fname}.o in the
  47 # ${CMAKE_CURRENT_BINARY_DIR}. The generated header will define a C++ class
  48 # called ${cpp_class} - which may be a namespace-qualified class name.
  49 function(tfcompile model tag_set signature_def_key fname cpp_class)
  50   set(prefix ${CMAKE_CURRENT_BINARY_DIR}/${fname})
  51   set(obj_file ${prefix}.o)
  52   set(hdr_file ${prefix}.h)
  53   string(TOUPPER ${fname} fname_allcaps)
  54   set(override_header ${LLVM_OVERRIDE_MODEL_HEADER_${fname_allcaps}})
  55   set(override_object ${LLVM_OVERRIDE_MODEL_OBJECT_${fname_allcaps}})
  56   if (EXISTS "${override_header}" AND EXISTS "${override_object}")
  57     configure_file(${override_header} ${hdr_file} COPYONLY)
  58     configure_file(${override_object} ${obj_file} COPYONLY)
  59     message("Using provided header "
  60       ${hdr_file} " and object "   ${obj_file}
  61       " files for model " ${model})
  62   else()
  63     tf_get_absolute_path(${model} ${CMAKE_CURRENT_BINARY_DIR} LLVM_ML_MODELS_ABSOLUTE)
  64     message("Using model at " ${LLVM_ML_MODELS_ABSOLUTE})
  65     add_custom_command(OUTPUT ${obj_file} ${hdr_file}
  66       COMMAND ${TENSORFLOW_AOT_COMPILER} aot_compile_cpu
  67             --multithreading false
  68             --dir ${LLVM_ML_MODELS_ABSOLUTE}
  69             --tag_set ${tag_set}
  70             --signature_def_key ${signature_def_key}
  71             --output_prefix ${prefix}
  72             --cpp_class ${cpp_class}
  73             --target_triple ${LLVM_HOST_TRIPLE}
  74     )
  75   endif()
  76
  77   # Aggregate the objects so that results of different tfcompile calls may be
  78   # grouped into one target.
  79   set(GENERATED_OBJS ${GENERATED_OBJS} ${obj_file} PARENT_SCOPE)
  80   set_source_files_properties(${obj_file} PROPERTIES
  81     GENERATED 1 EXTERNAL_OBJECT 1)
  82
  83   set(GENERATED_HEADERS ${GENERATED_HEADERS} ${hdr_file} PARENT_SCOPE)
  84   set_source_files_properties(${hdr_file} PROPERTIES
  85     GENERATED 1)
  86
  87 endfunction()
  88
  89 function(tf_find_and_compile model default_url default_path generation_config tag_set signature_def_key fname cpp_class)
  90   if ("${model}" STREQUAL "download")
  91     # Crash if the user wants to download a model but a URL is set to "TO_BE_UPDATED"
  92     if ("${LLVM_INLINER_MODEL_CURRENT_URL}" STREQUAL "TO_BE_UPDATED")
  93         message(FATAL_ERROR "LLVM_INLINER_MODEL_PATH was set to 'download' but there is no model url currently specified in cmake - likely, the model interface recently changed, and so there is not a released model available.")
  94     endif()
  95
  96     set(model ${default_url})
  97   endif()
  98
  99   if ("${model}" STREQUAL "autogenerate")
 100     generate_mock_model(${default_path} models/generate_mock_model.py ${generation_config})
 101     set(model ${default_path}-autogenerated)
 102   endif()
 103
 104   tf_get_model(${model} LLVM_ML_MODELS_ABSOLUTE)
 105   tfcompile(${LLVM_ML_MODELS_ABSOLUTE} ${tag_set} ${signature_def_key} ${fname} ${cpp_class})
 106
 107   set(GENERATED_OBJS ${GENERATED_OBJS} ${obj_file} PARENT_SCOPE)
 108   set_source_files_properties(${obj_file} PROPERTIES
 109     GENERATED 1 EXTERNAL_OBJECT 1)
 110
 111   set(GENERATED_HEADERS ${GENERATED_HEADERS} ${hdr_file} PARENT_SCOPE)
 112   set_source_files_properties(${hdr_file} PROPERTIES
 113     GENERATED 1)
 114
 115   set(GeneratedMLSources ${GeneratedMLSources} ${GENERATED_HEADERS} PARENT_SCOPE)
 116   set(MLDeps ${MLDeps} tf_xla_runtime PARENT_SCOPE)
 117   set(MLLinkDeps ${MLLinkDeps} tf_xla_runtime ${GENERATED_OBJS} PARENT_SCOPE)
 118
 119 endfunction()