Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[tflchef] Extract more methods #13661

Merged
merged 1 commit into from
Aug 13, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
156 changes: 156 additions & 0 deletions compiler/tflchef/core/src/ModelChef.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -187,6 +187,10 @@ class ModelChef
void cook(const ::tflchef::ModelRecipe &model_recipe);

private:
void prepare_initial_buffer(void);
void gather_operator_codes(const ::tflchef::ModelRecipe &model_recipe);
void gather_signature_defs(const ::tflchef::ModelRecipe &model_recipe);

template <typename T> void cook_operands(const T &graph);

template <typename T>
Expand Down Expand Up @@ -739,8 +743,153 @@ void ModelChef::cook_graph(const T &graph, std::map<std::string, int32_t> &symbo
_subgraph_vec.emplace_back(subgraph_builder.Finish());
}

void ModelChef::gather_operator_codes(const ::tflchef::ModelRecipe &model_recipe)
{
// Create OperatorCode with Builtin Operator
_builtin_code_map = gather_builtincode_map(model_recipe);
for (auto const &opcode : _builtin_code_map)
{
tflite::OperatorCodeBuilder code_builder{*_flatbuffer_builder};
// 127 is BuiltinOperator_PLACEHOLDER_FOR_GREATER_OP_CODES
// This is the way to handle deprecated builtin code
// See
// https://github.com/tensorflow/tensorflow/blob/a0afe8f9218be5eb9ed5dffc2dff652996da8c28/tensorflow/lite/schema/schema.fbs#L1061-L1077
if (opcode.first < 127)
{
code_builder.add_deprecated_builtin_code(opcode.first);
}
else
{
code_builder.add_deprecated_builtin_code(
::tflite::BuiltinOperator_PLACEHOLDER_FOR_GREATER_OP_CODES);
}
code_builder.add_version(opcode.second);
code_builder.add_builtin_code(opcode.first);
auto code = code_builder.Finish();
// Update OperatorCode vector
_code_vec.emplace_back(code);
}

// Create OperatorCode with Custom Operator
{
std::set<std::string> custom_code_set = gather_customcode_set(model_recipe);
std::vector<std::string> custom_code_vec{custom_code_set.begin(), custom_code_set.end()};
_custom_code_vec = custom_code_vec;
}

for (auto opcode : _custom_code_vec)
{
auto custom_code = _flatbuffer_builder->CreateString(opcode);
tflite::OperatorCodeBuilder code_builder{*_flatbuffer_builder};
code_builder.add_deprecated_builtin_code(tflite::BuiltinOperator_CUSTOM);
code_builder.add_custom_code(custom_code);
code_builder.add_builtin_code(tflite::BuiltinOperator_CUSTOM);
auto code = code_builder.Finish();
// Update OperatorCode vector
_code_vec.emplace_back(code);
}
}

void ModelChef::prepare_initial_buffer(void)
{
// Create an Empty Buffer
//
// Buffer 0 SHOULD be an empty buffer in TensorFlow Lite model file
// (Please refer to the comment for Tensor.buffer field in schema)
tflite::BufferBuilder buffer_builder{*_flatbuffer_builder};
_buffer_vec.emplace_back(buffer_builder.Finish());
}

void ModelChef::gather_signature_defs(const ::tflchef::ModelRecipe &model_recipe)
{
for (int s = 0; s < model_recipe.signature_def_size(); ++s)
{
// load from recipe
const auto &rec_signature_def = model_recipe.signature_def(s);

std::vector<flatbuffers::Offset<::tflite::TensorMap>> tensormap_inputs;
std::vector<flatbuffers::Offset<::tflite::TensorMap>> tensormap_outputs;

// which subgraph index to cook
auto subgraph_index = 0;
if (rec_signature_def.has_subgraph_index())
{
subgraph_index = rec_signature_def.subgraph_index();
}
assert(subgraph_index < _symbol_tables.size());
auto &symbol_table = _symbol_tables[subgraph_index];

// cook for inputs
for (int si = 0; si < rec_signature_def.inputs_size(); ++si)
{
// recipe for input TensorMap
auto rec_tm_input = rec_signature_def.inputs(si);
auto name = _flatbuffer_builder->CreateString(rec_tm_input.name());
uint32_t tensor_index = 0;
// either tensor or tensor_index should exist
assert(rec_tm_input.has_tensor() || rec_tm_input.has_tensor_index());
if (rec_tm_input.has_tensor())
{
// we can get tensor_index from symbol_table
auto tensor = rec_tm_input.tensor();
tensor_index = symbol_table[tensor];
}
else
{
// or we can use tensor_index itself
tensor_index = rec_tm_input.tensor_index();
}

::tflite::TensorMapBuilder tensormap_builder{*_flatbuffer_builder};
tensormap_builder.add_name(name);
tensormap_builder.add_tensor_index(tensor_index);
tensormap_inputs.push_back(tensormap_builder.Finish());
}
// cook for outputs, same as inputs
for (int so = 0; so < rec_signature_def.outputs_size(); ++so)
{
auto rec_tm_output = rec_signature_def.outputs(so);
auto name = _flatbuffer_builder->CreateString(rec_tm_output.name());
uint32_t tensor_index = 0;
assert(rec_tm_output.has_tensor() || rec_tm_output.has_tensor_index());
if (rec_tm_output.has_tensor())
{
auto tensor = rec_tm_output.tensor();
tensor_index = symbol_table[tensor];
}
else
{
tensor_index = rec_tm_output.tensor_index();
}

::tflite::TensorMapBuilder tensormap_builder{*_flatbuffer_builder};
tensormap_builder.add_name(name);
tensormap_builder.add_tensor_index(tensor_index);
tensormap_outputs.push_back(tensormap_builder.Finish());
}

auto inputs = _flatbuffer_builder->CreateVector(tensormap_inputs);
auto outputs = _flatbuffer_builder->CreateVector(tensormap_outputs);
auto signature_key = _flatbuffer_builder->CreateString(rec_signature_def.signature_key());
// TODO add validation for signature_key

::tflite::SignatureDefBuilder signature_def_builder{*_flatbuffer_builder};
signature_def_builder.add_inputs(inputs);
signature_def_builder.add_outputs(outputs);
signature_def_builder.add_signature_key(signature_key);
signature_def_builder.add_subgraph_index(rec_signature_def.subgraph_index());

_signdef_vec.emplace_back(signature_def_builder.Finish());
}
}

void ModelChef::cook(const ::tflchef::ModelRecipe &model_recipe)
{
prepare_initial_buffer();

gather_operator_codes(model_recipe);

#if 0
// Create OperatorCode with Builtin Operator
_builtin_code_map = gather_builtincode_map(model_recipe);
for (auto const &opcode : _builtin_code_map)
Expand Down Expand Up @@ -781,7 +930,9 @@ void ModelChef::cook(const ::tflchef::ModelRecipe &model_recipe)
// Update OperatorCode vector
_code_vec.emplace_back(code);
}
#endif

#if 0
// Create an Empty Buffer
//
// Buffer 0 SHOULD be an empty buffer in TensorFlow Lite model file
Expand All @@ -790,6 +941,7 @@ void ModelChef::cook(const ::tflchef::ModelRecipe &model_recipe)
tflite::BufferBuilder buffer_builder{*_flatbuffer_builder};
_buffer_vec.emplace_back(buffer_builder.Finish());
}
#endif

//
// Create Main graph
Expand Down Expand Up @@ -820,6 +972,9 @@ void ModelChef::cook(const ::tflchef::ModelRecipe &model_recipe)
_symbol_tables.push_back(symbol_table);
}

gather_signature_defs(model_recipe);

#if 0
// Create Signature-Def
//
for (int s = 0; s < model_recipe.signature_def_size(); ++s)
Expand Down Expand Up @@ -901,6 +1056,7 @@ void ModelChef::cook(const ::tflchef::ModelRecipe &model_recipe)

_signdef_vec.emplace_back(signature_def_builder.Finish());
}
#endif

// Create "Model" arguments
auto buffers = _flatbuffer_builder->CreateVector(_buffer_vec);
Expand Down