diff --git a/bazel/BUILD b/bazel/BUILD index 90a0762ec..8e8cc859a 100644 --- a/bazel/BUILD +++ b/bazel/BUILD @@ -33,6 +33,7 @@ cc_library( "ngraph_bridge/ngraph_encapsulate_op.h", "ngraph_bridge/ngraph_freshness_tracker.h", "ngraph_bridge/ngraph_mark_for_clustering.h", + "ngraph_bridge/ngraph_partial_shapes.h", "ngraph_bridge/ngraph_pipelined_tensors.h", "ngraph_bridge/ngraph_rewrite_for_tracking.h", "ngraph_bridge/ngraph_timer.h", @@ -70,6 +71,7 @@ cc_binary( "ngraph_bridge/ngraph_encapsulate_op.cc", "ngraph_bridge/ngraph_freshness_tracker.cc", "ngraph_bridge/ngraph_mark_for_clustering.cc", + "ngraph_bridge/ngraph_partial_shapes.cc", "ngraph_bridge/ngraph_pipelined_tensors.cc", "ngraph_bridge/ngraph_rewrite_for_tracking.cc", "ngraph_bridge/ngraph_tracked_variable.cc", diff --git a/ngraph_bridge/CMakeLists.txt b/ngraph_bridge/CMakeLists.txt index 0ad8d4161..0f8312bab 100644 --- a/ngraph_bridge/CMakeLists.txt +++ b/ngraph_bridge/CMakeLists.txt @@ -48,6 +48,7 @@ set(SRC ngraph_encapsulate_op.cc ngraph_freshness_tracker.cc ngraph_mark_for_clustering.cc + ngraph_partial_shapes.cc ngraph_rewrite_for_tracking.cc ngraph_rewrite_pass.cc ngraph_tracked_variable.cc diff --git a/ngraph_bridge/enable_variable_ops/ngraph_rewrite_pass.cc b/ngraph_bridge/enable_variable_ops/ngraph_rewrite_pass.cc index f9e57e040..fbe2e3fa4 100644 --- a/ngraph_bridge/enable_variable_ops/ngraph_rewrite_pass.cc +++ b/ngraph_bridge/enable_variable_ops/ngraph_rewrite_pass.cc @@ -308,7 +308,7 @@ class NGraphEncapsulationPass : public NGraphRewritePass { // 4. Encapsulate clusters then, if requested, dump the graphs. FunctionDefLibrary* fdeflib_new = new FunctionDefLibrary(); TF_RETURN_IF_ERROR(EncapsulateClusters(options.graph->get(), idx, - fdeflib_new, config_map)); + fdeflib_new, config_map, {0, {}})); // TODO: not using fdeflib_new in this path. Only grappler path uses it free(fdeflib_new); if (DumpEncapsulatedGraphs()) { diff --git a/ngraph_bridge/grappler/ngraph_optimizer.cc b/ngraph_bridge/grappler/ngraph_optimizer.cc index 669f58ca8..2304409f2 100644 --- a/ngraph_bridge/grappler/ngraph_optimizer.cc +++ b/ngraph_bridge/grappler/ngraph_optimizer.cc @@ -53,14 +53,42 @@ Status NgraphOptimizer::Init( config_backend_name = params.at("ngraph_backend").s(); config_device_id = params.at("device_id").s(); NGRAPH_VLOG(3) << "Backend name from config: " << config_backend_name; + std::set shape_hints; + // typedef std::map> ShapeHintMap; for (auto i : params) { if (i.first != "ngraph_backend") { - config_map[(i.first == "device_id" ? "" : "_") + std::string("ngraph_") + - i.first] = i.second.s(); - NGRAPH_VLOG(3) << "Attribute: " << i.first - << " Value: " << config_map["_ngraph_" + i.first]; + // TODO: slightly hacky. The bridge reserves the right to use optional + // attributes whose names start with shape_hint + if (i.first.rfind("shape_hint", 0) != 0) { + config_map[(i.first == "device_id" ? "" : "_") + + std::string("ngraph_") + i.first] = i.second.s(); + NGRAPH_VLOG(3) << "Attribute: " << i.first + << " Value: " << config_map["_ngraph_" + i.first]; + } else { + ShapeHintMap hint; + for (auto k : i.second.func().attr().at("hint_body").func().attr()) { + vector full_or_partial_shape; + for (auto dim : k.second.tensor().int_val()) { + full_or_partial_shape.push_back(dim); + } + hint[k.first] = full_or_partial_shape; + } + shape_hints.insert(hint); + } } } + auto itr = params.find("aot_requested"); + bool do_aot = false; + if (itr != params.end()) { + do_aot = itr->second.s() == "1"; + } + if (!do_aot && shape_hints.size() > 0) { + return errors::Internal( + "Did not requested AOT, but passed shape hints. Please request to use " + "shape hints (by using --precompile in tf2ngraph.py), or if AOT is not " + "desired then do not pass shape hints"); + } + aot_info = make_pair(do_aot, shape_hints); return Status::OK(); } @@ -230,7 +258,9 @@ Status NgraphOptimizer::Optimize(tensorflow::grappler::Cluster* cluster, // 4. Encapsulate clusters then, if requested, dump the graphs. FunctionDefLibrary* fdeflib_new = new FunctionDefLibrary(); - TF_RETURN_IF_ERROR(EncapsulateClusters(&graph, idx, fdeflib_new, config_map)); + TF_RETURN_IF_ERROR( + // TODO: right now _ngraph_aot_requested is passed along in config_map. + EncapsulateClusters(&graph, idx, fdeflib_new, config_map, aot_info)); if (DumpEncapsulatedGraphs()) { DumpGraphs(graph, idx, "encapsulated", "Graph with Clusters Encapsulated"); } diff --git a/ngraph_bridge/grappler/ngraph_optimizer.h b/ngraph_bridge/grappler/ngraph_optimizer.h index 64fb6d194..e6914ae5c 100644 --- a/ngraph_bridge/grappler/ngraph_optimizer.h +++ b/ngraph_bridge/grappler/ngraph_optimizer.h @@ -83,6 +83,7 @@ class NgraphOptimizer : public tensorflow::grappler::CustomGraphOptimizer { static int s_serial_counter GUARDED_BY(s_serial_counter_mutex); static mutex s_serial_counter_mutex; + AOTInfo aot_info; }; int NgraphOptimizer::s_serial_counter = 0; diff --git a/ngraph_bridge/ngraph_encapsulate_clusters.cc b/ngraph_bridge/ngraph_encapsulate_clusters.cc index 0d6f1eb41..fcc076c65 100644 --- a/ngraph_bridge/ngraph_encapsulate_clusters.cc +++ b/ngraph_bridge/ngraph_encapsulate_clusters.cc @@ -40,11 +40,13 @@ #include "ngraph_bridge/ngraph_api.h" #include "ngraph_bridge/ngraph_assign_clusters.h" #include "ngraph_bridge/ngraph_bridge_registry.h" +#include "ngraph_bridge/ngraph_builder.h" #include "ngraph_bridge/ngraph_cluster_manager.h" #include "ngraph_bridge/ngraph_encapsulate_clusters.h" #include "ngraph_bridge/ngraph_mark_for_clustering.h" -#include "ngraph_bridge/ngraph_mark_for_clustering.h" +#include "ngraph_bridge/ngraph_partial_shapes.h" #include "ngraph_bridge/ngraph_utils.h" +#include "ngraph_bridge/version.h" using namespace std; @@ -74,7 +76,8 @@ static void AddInput(NodeDef* dst, StringPiece src_name, int src_slot) { Status EncapsulateClusters( Graph* graph, int graph_id, FunctionDefLibrary* fdeflib, - std::unordered_map device_config) { + std::unordered_map device_config, + AOTInfo aot_info) { // A map from cluster indices to the expected device name for nodes // in that cluster. std::map device_name_map; @@ -160,7 +163,6 @@ Status EncapsulateClusters( std::map arg_index_count; int count_arg = 0, count_retval = 0, count_both_arg_retval = 0, count_free = 0, count_encapsulated = 0, count_tot = 0; - for (auto edge : graph->edges()) { count_tot++; // TODO(amprocte): should actually keep of these. During clustering we @@ -527,7 +529,346 @@ Status EncapsulateClusters( fdef)); } - // Pass 8 (optional, only run if environment variable + // Pass 8: + bool aot_requested; + set performed_aot_on_enc; + std::set>> node_shapes_hints_sets; + std::tie(aot_requested, node_shapes_hints_sets) = aot_info; + if (aot_requested) { + NGRAPH_VLOG(3) << "AOT requested"; + if (!ngraph_tf_is_grappler_enabled()) { + return errors::Internal( + "AOT requested for non grappler build. Please use grappler build if " + "AOT is required"); + } + string input_node_type = "Placeholder"; + // In case of grappler, we have Placeholder, which might contain shape info, + // so it is possible we can aot without any provided shapes + // in normal pass its args. unless shapes are provided there is no chance of + // reading shapes from args. + + auto get_shape_for_node_from_shape_hint = [](Node* node, + ShapeHintMap single_hint) { + auto find_itr = single_hint.find(node->name()); + return find_itr == single_hint.end() ? PartialShape() + : PartialShape(find_itr->second); + }; + + auto hint_as_string = [](ShapeHintMap single_hint) { + string hint_str; + for (auto itr_node : single_hint) { + hint_str += + ((itr_node.first) + ":[" + ng::join(itr_node.second) + "],"); + } + return hint_str; + }; + + std::map> inputs_node_shapes_for_compilation; + // map between node name and the PartialShape it contains + std::map node_partial_shape_map; + // This is a map of placeholder names and the shapes we can infer from them + std::map> shape_from_placeholders_as_hints; + for (auto node : graph->op_nodes()) { + if (node->type_string() == input_node_type) { + NGRAPH_VLOG(5) << "Checking input for AOT: " << node->name() << "(" + << node->type_string() + << "): " << node->attrs().SummarizeNode(); + // TODO: need to confirm if its _output_shapes or shape + auto shape_field = node->attrs().Find("_output_shapes"); + if (shape_field == nullptr) { + shape_field = node->attrs().Find("shape"); + } + // It seems that _output_shapes is not found and hence the shape is + // inferred only from the hints. however if "shape" is present, it is + // empty, and in that case the empty shape and the rank!=0 hint fuse + // to give an invalid shape according to our current logic. have to + // modify that + PartialShape partial_shape_from_node; + if (shape_field != nullptr) { + // Get shape from the node + partial_shape_from_node = PartialShape(shape_field->shape()); + } + NGRAPH_VLOG(5) << "For node " << node->name() + << " got shape from nose: " + << partial_shape_from_node.to_string(); + node_partial_shape_map.insert({node->name(), partial_shape_from_node}); + shape_from_placeholders_as_hints.insert( + {node->name(), partial_shape_from_node.get_shape_vector()}); + } + } + + // If no shape hints are provided but the placeholders contain complete + // shape, then we still need to enter the for loop below to compute AOT. + // Hence adding the shapes from placeholders as hints. + if (node_shapes_hints_sets.size() == 0) { + NGRAPH_VLOG(5) << "Using shapes from placeholders as hint"; + node_shapes_hints_sets.insert(shape_from_placeholders_as_hints); + } + // TODO: .....CHECK ABOVE IF + + // Iterate over each shape hint and see if they can be used + for (ShapeHintMap single_hint : node_shapes_hints_sets) { + // A boolean to determine if we can AOT for this single_hint + bool can_aot = true; + + for (auto itr_single_hint : single_hint) { + if (shape_from_placeholders_as_hints.find(itr_single_hint.first) == + shape_from_placeholders_as_hints.end()) { + return errors::Internal("Passed hint for node ", + itr_single_hint.first, + " but there is no input with that name"); + } + } + + for (auto node : graph->op_nodes()) { + if (node->type_string() == input_node_type) { + PartialShape partial_shape_from_node = + node_partial_shape_map.at(node->name()); + + PartialShape shape_hint_for_node = + get_shape_for_node_from_shape_hint(node, single_hint); + + // If a shape has been found in the input node, match with + // shape_hints if they exist + PartialShape combined_shape_info; + if (shape_hint_for_node.is_valid()) { + NGRAPH_VLOG(5) << "For node " << node->name() << " shape hint (", + hint_as_string(single_hint), + ") for node is valid and is: " + + shape_hint_for_node.to_string(); + if (partial_shape_from_node.is_valid()) { + NGRAPH_VLOG(5) << "Partial shape from node is also valid. So " + "will attempt to concretize if possible"; + if (partial_shape_from_node.size() == 0) { + // TODO: revisit this if-else + NGRAPH_VLOG(5) << "Partial shape from node is empty, so will " + "use shape from hint"; + combined_shape_info = shape_hint_for_node; + } else { + NGRAPH_VLOG(5) << "Concretizing shape " + + partial_shape_from_node.to_string() + + "from node with hint for node, " + + shape_hint_for_node.to_string(); + partial_shape_from_node.concretize(shape_hint_for_node); + combined_shape_info = partial_shape_from_node; + } + } else { + NGRAPH_VLOG(5) << "Partial shape from node is invalid. So using " + "hint for the node as shape"; + combined_shape_info = shape_hint_for_node; + } + } else { + NGRAPH_VLOG(5) << "For node " << node->name() + << " shape hint (" + hint_as_string(single_hint) + + ") for node is invalid"; + if (partial_shape_from_node.is_valid()) { + // No shape hints found. But the node itself has some shape info + NGRAPH_VLOG(5) << "Partial shape from node is valid and is: " + + partial_shape_from_node.to_string(); + combined_shape_info = partial_shape_from_node; + } else { + NGRAPH_VLOG(5) << "Partial shape from node is invalid"; + combined_shape_info = PartialShape(); + } + } + + can_aot = combined_shape_info.is_valid() && + combined_shape_info.is_concrete(); + if (can_aot) { + inputs_node_shapes_for_compilation[node->name()] = + combined_shape_info.get_shape_vector(); + } else { + // TODO: necessarily break? Maybe some things can be AOT, others + // maybe not + string fail_reason = + (combined_shape_info.is_valid() + ? (node->name() + " could not be concretized") + : "it is invalid for " + node->name()); + return errors::Internal("Cannot AOT using this hint (", + hint_as_string(single_hint), ") as ", + fail_reason); + break; + } + } // end of if (node->type_string() == input_node_type) + } // End of for loop that goes through all nodes + + // Did we manage to concretize all input shapes? + for (auto itr : node_partial_shape_map) { // iterate over all inputs + if (inputs_node_shapes_for_compilation.find(itr.first) == + inputs_node_shapes_for_compilation.end()) { + can_aot = false; + // TODO: print "this" hint + return errors::Internal("Cannot AOT using this hint (", + hint_as_string(single_hint), ") for ", + (itr.first), " was not concretized"); + } + } + + if (!can_aot) { + return errors::Internal("AOT requested, but could not perform AOT"); + } + for (auto node : graph->op_nodes()) { + if (node->type_string() == "NGraphEncapsulate") { + // Check inputs of the encapsulates. They can only be fed by fully + // concrete shapes (after going through the shape hints) or consts + std::vector st_inputs; + GetStaticInputs(node, &st_inputs); + // Current assumption is that only encapsulates without static + // inputs are AOT + if (st_inputs.size() != 0) { + return errors::Internal( + "AOT requested. Found an encapsulate with static inputs, but " + "that is not supported"); + } + + std::vector input_shapes; + std::stringstream signature_ss; + for (auto in_node : node->in_nodes()) { + if (!in_node->IsSource()) { + auto itr_shape = + inputs_node_shapes_for_compilation.find(in_node->name()); + if (itr_shape == inputs_node_shapes_for_compilation.end()) { + // TODO: this error could potentially happen due to 2 reasons: + // 1. Enough valid shape hints were not passed + // 2. It is an encapsulate that has atleast 1 input fed by a + // non-placeholder (like another TF node or another + // encapsulate) + // Later provide more explicit debug message (reason 1 or 2 or + // anything else) + return errors::Internal( + "AOT requested. Found an encapsulate that has a " + "non-concrete input"); + } else { + std::vector converted_to_int64(itr_shape->second.begin(), + itr_shape->second.end()); + input_shapes.push_back(TensorShape(converted_to_int64)); + for (auto itr1 : itr_shape->second) { + signature_ss << itr1 << ","; + } + signature_ss << ";"; + } + } + } + + signature_ss << "/"; + string signature = signature_ss.str(); + NGRAPH_VLOG(3) << "Performing AOT for " << node->name() + << " for signature = " << signature << "\n"; + + std::vector static_input_map; + std::shared_ptr ng_function; + int cluster_idx; + TF_RETURN_IF_ERROR( + GetNodeAttr(node->attrs(), "ngraph_cluster", &cluster_idx)); + GraphDef* gdef_for_current_encapsulate; + gdef_for_current_encapsulate = + NGraphClusterManager::GetClusterGraph(cluster_idx); + GraphConstructorOptions opts; + opts.allow_internal_ops = true; + Graph graph_for_current_encapsulate(OpRegistry::Global()); + TF_RETURN_IF_ERROR( + ConvertGraphDefToGraph(opts, *gdef_for_current_encapsulate, + &graph_for_current_encapsulate)); + TF_RETURN_IF_ERROR(Builder::TranslateGraph( + input_shapes, static_input_map, &graph_for_current_encapsulate, + ng_function)); + string serialized_ngfunc(ngraph::serialize(ng_function, 4)); + + // get backend. + // TODO: Note that this is code duplication of some stuff present + // in NGraphEncapsulateOp + // Once NGraphEncapsulateOp is refactored, this code should be + // removed and a common function should be used + + // TODO: these sections can be hoisted out of the main loop + std::string backend_name; + TF_RETURN_IF_ERROR( + GetNodeAttr(node->attrs(), "ngraph_backend", &backend_name)); + std::string device_id; + TF_RETURN_IF_ERROR( + GetNodeAttr(node->attrs(), "ngraph_device_id", &device_id)); + + string op_backend_name; + try { + op_backend_name = BackendManager::GetBackendCreationString( + backend_name, device_id); + } catch (const std::exception& exp) { + return errors::Internal( + "Caught exception while creating backend string ", exp.what(), + "\n"); + } + BackendManager::CreateBackend( + op_backend_name); // Created a backend here. must free it + std::unordered_map additional_attribute_map; + for (auto itr : node->attrs()) { + // Find the optional attributes to be sent to the backend. + // The optional attributes have '_ngraph_' appended to the start + // so we need to get rid of that and only send the remaining + // string + // since the backend will only look for that. + // '_ngraph_' is only appended for the bridge. + // For e.g. _ngraph_ice_cores --> ice_cores + if (itr.first.find("_ngraph_") != std::string::npos) { + // leave out _ngraph_aot_requested + if (itr.first.find("_ngraph_aot_requested") != + std::string::npos) { + additional_attribute_map.insert( + {itr.first.substr(strlen("_ngraph_")), itr.second.s()}); + } + } + } + BackendManager::SetConfig(op_backend_name, additional_attribute_map); + ng::runtime::Backend* op_backend = nullptr; + try { + op_backend = BackendManager::GetBackend(op_backend_name); + } catch (const std::out_of_range& e) { + BackendManager::ReleaseBackend(op_backend_name); + throw; + } + BackendManager::LockBackend(op_backend_name); + std::shared_ptr ng_exec; + try { + ng_exec = op_backend->compile(ng_function); + } catch (...) { + BackendManager::UnlockBackend(op_backend_name); + NgraphSerialize("tf_function_error_aot.json", ng_function); + BackendManager::ReleaseBackend(op_backend_name); + return errors::Internal("Failed to compile ng_function for AOT"); + } + BackendManager::UnlockBackend(op_backend_name); + BackendManager::ReleaseBackend(op_backend_name); + + stringstream exec_dump; + ng_exec->save(exec_dump); + // ng function attached as debugging information + node->AddAttr("_ngraph_aot_ngfunction_" + signature, + serialized_ngfunc); + // Compute will use this ngexec + node->AddAttr("_ngraph_aot_ngexec_" + signature, exec_dump.str()); + // We do not need to add "_ngraph_aot_requested" attribute since it + // already is already present in device_config and inserted into the + // currently created NGraphEncapsulate + // TODO: create a separate namespace of node attributes for backend + // and for bridge + performed_aot_on_enc.insert(node->name()); + NGRAPH_VLOG(5) << "Performed AOT on " << node->name(); + } + } + } // end of for (ShapeHintMap single_hint : node_shapes_hints_sets) + + // In the end assert that all encapsulates have performed AOT + for (auto node : graph->op_nodes()) { + if (node->type_string() == "NGraphEncapsulate") { + if (performed_aot_on_enc.find(node->name()) == + performed_aot_on_enc.end()) { + return errors::Internal("Requested AOT, but did not perform AOT on ", + node->name()); + } + } + } + } // end of if (aot_requested) + + // Pass 9 (optional, only run if environment variable // NGRAPH_TF_DUMP_CLUSTERS is set): validate the graph def, and // make sure we can construct a graph from it. if (std::getenv("NGRAPH_TF_DUMP_CLUSTERS")) { diff --git a/ngraph_bridge/ngraph_encapsulate_clusters.h b/ngraph_bridge/ngraph_encapsulate_clusters.h index 26ab583f1..c2aea6111 100644 --- a/ngraph_bridge/ngraph_encapsulate_clusters.h +++ b/ngraph_bridge/ngraph_encapsulate_clusters.h @@ -18,15 +18,27 @@ #define NGRAPH_TF_BRIDGE_ENCAPSULATE_CLUSTERS_H_ #pragma once +#include +#include +#include +#include + +#include #include "tensorflow/core/graph/graph.h" namespace tensorflow { namespace ngraph_bridge { +typedef std::map> ShapeHintMap; + +// the integer represent AOT level requested. +typedef std::pair> AOTInfo; + Status EncapsulateClusters( Graph* graph, int graph_id, FunctionDefLibrary* fdeflib, - std::unordered_map device_config); + std::unordered_map device_config, + AOTInfo aot_info); } // namespace ngraph_bridge } // namespace tensorflow diff --git a/ngraph_bridge/ngraph_encapsulate_impl.cc b/ngraph_bridge/ngraph_encapsulate_impl.cc index f7a8e9486..f7a767b73 100644 --- a/ngraph_bridge/ngraph_encapsulate_impl.cc +++ b/ngraph_bridge/ngraph_encapsulate_impl.cc @@ -134,9 +134,19 @@ Status NGraphEncapsulateImpl::GetNgExecutable( MemoryProfile(vm0, rss0); NGRAPH_VLOG(1) << "Compilation cache miss: " << m_name; - TF_RETURN_IF_ERROR(Builder::TranslateGraph(input_shapes, static_input_map, - &m_graph, ng_function)); - ng_function->set_friendly_name(m_name); + if (!m_do_aot) { + TF_RETURN_IF_ERROR(Builder::TranslateGraph(input_shapes, static_input_map, + &m_graph, ng_function)); + ng_function->set_friendly_name(m_name); + } else { + auto itr = m_aot_functions.find(signature); + if (itr == m_aot_functions.end()) { + return errors::Internal( + "Expected to find AOT precompiled ng function of signature: ", + signature); + } + ng_function = ng::deserialize(itr->second); + } auto function_size = ng_function->get_graph_size() / 1024; // kb unit @@ -192,11 +202,24 @@ Status NGraphEncapsulateImpl::GetNgExecutable( << output_tensors_bytes_free / (1024 * 1024) << " MB"; } // cache eviction if cache size greater than cache depth - BackendManager::LockBackend(m_op_backend_name); - ngraph::Event event_compile("Compile nGraph", m_name, ""); + BackendManager::LockBackend(m_op_backend_name); try { - ng_exec = op_backend->compile(ng_function); + if (m_do_aot) { + auto itr = m_aot_execs.find(signature); + if (itr == m_aot_execs.end()) { + BackendManager::UnlockBackend(m_op_backend_name); + return errors::Internal( + "Requested AOT, but could not find string with the " + "signature: ", + signature); + } + stringstream serialized_exec_read; + serialized_exec_read << (itr->second); + ng_exec = op_backend->load(serialized_exec_read); + } else { + ng_exec = op_backend->compile(ng_function); + } } catch (const std::exception& exp) { BackendManager::UnlockBackend(m_op_backend_name); NgraphSerialize("tf_function_error_" + m_name + ".json", ng_function); @@ -450,6 +473,65 @@ std::shared_ptr NGraphEncapsulateImpl::GetCurrentNgTensor( return current_ng_tensor; } +Status NGraphEncapsulateImpl::ParseNodeAttributes( + const google::protobuf::Map& additional_attributes, + std::unordered_map* additional_attribute_map) { + for (auto itx : additional_attributes) { + // Find the optional attributes to be sent to the backend. + // The optional attributes have '_ngraph_' appended to the start + // so we need to get rid of that and only send the remaining string + // since the backend will only look for that. + // '_ngraph_' is only appended for the bridge. + // For e.g. _ngraph_ice_cores --> ice_cores + if (itx.first.find("_ngraph_") != std::string::npos) { + // TODO: decide what the node attributes should be. + // right now _ngraph_aot_ is used by aot, _ngraph_ is used for optional + // attributes + auto attr_name = itx.first; + auto attr_value = itx.second.s(); + if (attr_name.find("_ngraph_aot_") != std::string::npos) { + // The string is in the format: _ngraph_aot_ngexec_signature or + // _ngraph_aot_ngfunction_signature or _ngraph_aot_requested + // TODO: do not pass these 3 attributes to set_config of backend + if (attr_name.find("_ngraph_aot_ngexec_") != std::string::npos) { + m_aot_execs[ng::split(attr_name, '_')[4]] = attr_value; + } else if (attr_name.find("_ngraph_aot_ngfunction_") != + std::string::npos) { + // The other option is _ngraph_aot_ngfunction_ + // No need to save or do anything with _ngraph_aot_ngfunction_. They + // are there for debugging only + m_aot_functions[ng::split(attr_name, '_')[4]] = attr_value; + } else if (attr_name.find("_ngraph_aot_requested") != + std::string::npos) { + m_do_aot = (attr_value == "1"); + if (m_do_aot) { + NGRAPH_VLOG(1) << "Using AOT for encapsulate " + + to_string(m_ngraph_cluster); + } + } else { + return errors::Internal( + "Ngraph attribues beginning with _ngraph_aot_ " + "must be _ngraph_aot_ngexec_ or " + "_ngraph_aot_ngfunction_. But got " + "attribute named: ", + itx.first); + } + } else { + NGRAPH_VLOG(4) << "Attribute: " << attr_name.substr(strlen("_ngraph_")) + << " Value: " << attr_value; + additional_attribute_map->insert( + {attr_name.substr(strlen("_ngraph_")), attr_value}); + } + } + } + if (((m_aot_functions.size() > 0) || (m_aot_execs.size() > 0)) && !m_do_aot) { + return errors::Internal("The encapsulate ", m_name, + " has ngraph functions or executables embedded " + "in it, even though AOT was not requested."); + } + return Status::OK(); +} + Status NGraphEncapsulateImpl::CachePipelinedTensorIfNeeded( std::shared_ptr ng_exec) { if (!m_executable_can_create_tensor) { diff --git a/ngraph_bridge/ngraph_encapsulate_impl.h b/ngraph_bridge/ngraph_encapsulate_impl.h index dcd758976..cb5fa9688 100644 --- a/ngraph_bridge/ngraph_encapsulate_impl.h +++ b/ngraph_bridge/ngraph_encapsulate_impl.h @@ -189,6 +189,9 @@ class NGraphEncapsulateImpl { void SetName(string name) { m_name = name; } + Status ParseNodeAttributes( + const google::protobuf::Map& additional_attributes, + std::unordered_map* additional_attribute_map); void SetExecCanCreateTensor(bool b) { m_executable_can_create_tensor = b; } bool GetExecCanCreateTensor() { return m_executable_can_create_tensor; } @@ -226,6 +229,9 @@ class NGraphEncapsulateImpl { std::vector m_input_is_static; std::list m_lru; static int s_instance_count; + bool m_do_aot = false; + map m_aot_functions; + map m_aot_execs; // ng_function, ng_executable, Output and Input Cache maps std::unordered_map> diff --git a/ngraph_bridge/ngraph_encapsulate_op.cc b/ngraph_bridge/ngraph_encapsulate_op.cc index 640e1cd3b..597fab0b1 100644 --- a/ngraph_bridge/ngraph_encapsulate_op.cc +++ b/ngraph_bridge/ngraph_encapsulate_op.cc @@ -172,21 +172,8 @@ NGraphEncapsulateOp::NGraphEncapsulateOp(OpKernelConstruction* ctx) // Get the optional attributes std::unordered_map additional_attribute_map; auto node_def = ctx->def(); - auto additional_attributes = node_def.attr(); - for (auto itx : additional_attributes) { - // Find the optional attributes to be sent to the backend. - // The optional attributes have '_ngraph_' appended to the start - // so we need to get rid of that and only send the remaining string - // since the backend will only look for that. - // '_ngraph_' is only appended for the bridge. - // For e.g. _ngraph_ice_cores --> ice_cores - if (itx.first.find("_ngraph_") != std::string::npos) { - NGRAPH_VLOG(4) << "Attribute: " << itx.first.substr(strlen("_ngraph_")) - << " Value: " << itx.second.s(); - additional_attribute_map.insert( - {itx.first.substr(strlen("_ngraph_")), itx.second.s()}); - } - } + OP_REQUIRES_OK(ctx, ng_encap_impl.ParseNodeAttributes( + node_def.attr(), &additional_attribute_map)); // Concatenate the backend_name:device_id try { @@ -736,4 +723,4 @@ int NGraphEncapsulateImpl::s_instance_count = 0; REGISTER_KERNEL_BUILDER(Name("NGraphEncapsulate").Device(DEVICE_CPU), ngraph_bridge::NGraphEncapsulateOp); -} // namespace tensorflow +} // namespace tensorflow \ No newline at end of file diff --git a/ngraph_bridge/ngraph_partial_shapes.cc b/ngraph_bridge/ngraph_partial_shapes.cc new file mode 100644 index 000000000..90de172c6 --- /dev/null +++ b/ngraph_bridge/ngraph_partial_shapes.cc @@ -0,0 +1,114 @@ +/******************************************************************************* + * Copyright 2019 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + *******************************************************************************/ + +#include "ngraph_bridge/ngraph_partial_shapes.h" + +namespace tensorflow { + +namespace ngraph_bridge { + +PartialShape::PartialShape(std::vector shape) + : m_shape(shape), m_valid(true) { + for (auto i : m_shape) { + if (i < -1) { + invalidate(); + } + } +} +PartialShape::PartialShape() : m_valid(false) {} + +PartialShape::PartialShape( + const tensorflow::TensorShapeProto& tensor_shape_proto) { + try { + m_shape.resize(tensor_shape_proto.dim_size()); + for (uint shape_idx = 0; shape_idx < tensor_shape_proto.dim_size(); + shape_idx++) { + auto num_elems_in_this_dim = tensor_shape_proto.dim(shape_idx).size(); + m_shape[shape_idx] = num_elems_in_this_dim; + // -1 means not specified + } + m_valid = true; + } catch (...) { + invalidate(); + } +} + +bool PartialShape::is_concrete() const { + check_valid(); + return std::all_of(m_shape.begin(), m_shape.end(), + [](int i) { return i >= 0; }); +} + +int PartialShape::size() const { + check_valid(); + return m_shape.size(); +} + +int PartialShape::operator[](int idx) const { + check_valid(); + return m_shape[idx]; +} + +std::vector PartialShape::get_shape_vector() const { + check_valid(); + return m_shape; +} + +string PartialShape::to_string() const { + std::string st = m_valid ? "valid:" : "invalid:"; + for (auto i : m_shape) { + st += (std::to_string(i) + ","); + } + return st; +} + +void PartialShape::concretize(const PartialShape& shape_hint) { + // Both PartialShapes are expected to be valid + check_valid(); + uint base_rank = m_shape.size(); + if (base_rank != shape_hint.size()) { // different ranks + invalidate(); + return; + } else { + for (int i = 0; i < base_rank; i++) { + if (m_shape[i] != shape_hint[i]) { + if (m_shape[i] == -1) { + m_shape[i] = shape_hint[i]; + } else { + if (shape_hint[i] != -1) { + invalidate(); + return; + } + } + } + } + return; + } +} + +void PartialShape::check_valid() const { + if (!m_valid) { + throw std::runtime_error( + string("Attempted to use an invalid PartialShape")); + } +} + +void PartialShape::invalidate() { + m_shape.clear(); + m_valid = false; +} +} +} \ No newline at end of file diff --git a/ngraph_bridge/ngraph_partial_shapes.h b/ngraph_bridge/ngraph_partial_shapes.h new file mode 100644 index 000000000..1d1678886 --- /dev/null +++ b/ngraph_bridge/ngraph_partial_shapes.h @@ -0,0 +1,88 @@ +/******************************************************************************* + * Copyright 2019 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + *******************************************************************************/ + +#ifndef NGRAPH_TF_BRIDGE_PARTIAL_SHAPES_H_ +#define NGRAPH_TF_BRIDGE_PARTIAL_SHAPES_H_ +#pragma once + +#include +#include +#include +#include + +#include "tensorflow/core/graph/graph.h" + +namespace tensorflow { + +namespace ngraph_bridge { + +// TODO unit test the class and move to a different file +class PartialShape { + // This is a simple class that can represent full or partial shapes + // a full shape has all dimensions >= 0. A partial shape has atleast one + // dimension < 0 + // Its built like an optional/maybe, so please use is_valid before accessing + // other functions + + // A sample lattice that this class operates on is shown for rank = 2 below + // Invalid + // / \ + // (1,2) (2,2) + // / \ / \ + // (-1,2)(1,-1) (2,-1)(-1,2) + // \ \ / / + // ---- (-1,-1) ---- + // + // The PartialShape allows only 2 functionsthat modify object state + // Copy assignment (=) and concretize() + // Subsequent calls to concretize() can only make one move up the lattice. + // Copy assignment can reset the state of teh object arbitrarily + + // The class is not thread safe + + // Any scalar is represented by {True, {}} by this class + + public: + PartialShape(std::vector shape); + PartialShape(); + + PartialShape(const tensorflow::TensorShapeProto& tensor_shape_proto); + + bool is_concrete() const; + + int size() const; + + int operator[](int idx) const; + + std::vector get_shape_vector() const; + + bool is_valid() const { return m_valid; } + + string to_string() const; + + void concretize(const PartialShape& shape_hint); + + private: + std::vector m_shape; + bool m_valid; + + void check_valid() const; + + void invalidate(); +}; +} +} +#endif // NGRAPH_TF_BRIDGE_PARTIAL_SHAPES_H_ \ No newline at end of file diff --git a/ngraph_bridge/ngraph_rewrite_pass.cc b/ngraph_bridge/ngraph_rewrite_pass.cc index 1621a1eb1..1be64f10d 100644 --- a/ngraph_bridge/ngraph_rewrite_pass.cc +++ b/ngraph_bridge/ngraph_rewrite_pass.cc @@ -261,7 +261,7 @@ class NGraphEncapsulationPass : public NGraphRewritePass { // 4. Encapsulate clusters then, if requested, dump the graphs. FunctionDefLibrary* fdeflib_new = new FunctionDefLibrary(); TF_RETURN_IF_ERROR(EncapsulateClusters(options.graph->get(), idx, - fdeflib_new, config_map)); + fdeflib_new, config_map, {0, {}})); // TODO: not using fdeflib_new in this path. Only grappler path uses it free(fdeflib_new); if (DumpEncapsulatedGraphs()) { diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index c4b5c80a0..e6cc3f2a1 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -67,8 +67,9 @@ if(NGRAPH_TF_ENABLE_VARIABLES_AND_OPTIMIZERS) list(APPEND SRC graph_rewrites/remove_ngraphassigns.cc) endif() -if(NGRAPH_TF_USE_GRAPPLER_OPTIMIZER) - list(APPEND SRC graph_rewrites/config_for_grappler_test.cc) +if(NGRAPH_TF_USE_GRAPPLER_OPTIMIZER) + list(APPEND SRC graph_rewrites/config_for_grappler_test.cc) + list(APPEND SRC graph_rewrites/partial_shapes_test.cc) endif() set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11") diff --git a/test/graph_rewrites/encapsulate_clusters_test.cc b/test/graph_rewrites/encapsulate_clusters_test.cc index e94914fad..898feec6b 100644 --- a/test/graph_rewrites/encapsulate_clusters_test.cc +++ b/test/graph_rewrites/encapsulate_clusters_test.cc @@ -17,6 +17,7 @@ #include "gtest/gtest.h" #include "tensorflow/core/graph/node_builder.h" +#include "version.h" #include "ngraph_bridge/ngraph_cluster_manager.h" #include "ngraph_bridge/ngraph_encapsulate_clusters.h" @@ -80,9 +81,10 @@ TEST(EncapsulateClusters, PopulateLibrary) { g.AddEdge(node3, Graph::kControlSlot, sink, Graph::kControlSlot); FunctionDefLibrary* fdeflib_new = new FunctionDefLibrary(); + std::unordered_map config_map; config_map["ngraph_device_id"] = ""; - ASSERT_OK(EncapsulateClusters(&g, 0, fdeflib_new, config_map)); + ASSERT_OK(EncapsulateClusters(&g, 0, fdeflib_new, config_map, {0, {}})); int num_encapsulates = 0; int num_tf_nodes = 0; @@ -117,6 +119,534 @@ TEST(EncapsulateClusters, PopulateLibrary) { ASSERT_EQ(present, expected); free(fdeflib_new); } + +// Placeholder-->Add(0)--->IdN +// ^ +// | +// Placeholder or Const(0) +TEST(EncapsulateClusters, AOT0) { + if (!ngraph_tf_is_grappler_enabled()) return; + + vector fed_by_placeholder{true, false}; + for (auto using_placeholder : fed_by_placeholder) { + NGraphClusterManager::EvictAllClusters(); + Graph g(OpRegistry::Global()); + + int cluster_idx = NGraphClusterManager::NewCluster(); + + Node* node1; + Node* node2; + + ASSERT_OK(NodeBuilder("node1", "Placeholder") + .Attr("dtype", DT_FLOAT) + .Finalize(&g, &node1)); + + if (using_placeholder) { + ASSERT_OK(NodeBuilder("node2", "Placeholder") + .Attr("dtype", DT_FLOAT) + .Finalize(&g, &node2)); + } else { + Tensor t_shape(DT_INT32, TensorShape{2}); + t_shape.flat().data()[0] = 2; + t_shape.flat().data()[1] = 2; + ASSERT_OK(NodeBuilder("node1", "Const") + .Attr("dtype", DT_FLOAT) + .Attr("value", t_shape) + .Attr("_ngraph_marked_for_clustering", true) + .Attr("_ngraph_cluster", cluster_idx) + .Attr("_ngraph_backend", "INTERPRETER") + .Finalize(&g, &node2)); + } + + Node* node3; + ASSERT_OK(NodeBuilder("node3", "Add") + .Input(node1, 0) + .Input(node2, 0) + .Attr("T", DT_FLOAT) + .Attr("_ngraph_marked_for_clustering", true) + .Attr("_ngraph_cluster", cluster_idx) + .Attr("_ngraph_backend", "INTERPRETER") + .Finalize(&g, &node3)); + Node* node4; + std::vector inputs; + std::vector input_types; + inputs.push_back(NodeBuilder::NodeOut(node3, 0)); + input_types.push_back(node3->output_type(0)); + ASSERT_OK(NodeBuilder("node4", "IdentityN") + .Input(inputs) + .Attr("T", input_types) + .Finalize(&g, &node4)); + + Node* source = g.source_node(); + Node* sink = g.sink_node(); + g.AddEdge(source, Graph::kControlSlot, node1, Graph::kControlSlot); + g.AddEdge(source, Graph::kControlSlot, node2, Graph::kControlSlot); + g.AddEdge(node4, Graph::kControlSlot, sink, Graph::kControlSlot); + + FunctionDefLibrary* fdeflib_new = new FunctionDefLibrary(); + + std::vector> node_shapes_hints_vect = { + {}, {{{"node1", {2, 2}}, {"node2", {2, 2}}}}}; + std::vector did_aot = {true, true}; + // Interesting case when shape hints = {}. + // The placeholder has shape = []. + // Which means, given no hints it is indeed possible to compile that, since + // its fully specified (no -1s) + int num_cases = node_shapes_hints_vect.size(); + for (int i = 0; i < num_cases; i++) { + std::set hint; + for (auto itr : node_shapes_hints_vect[i]) { + if (using_placeholder) { + hint.insert(itr); + } else { + ShapeHintMap temp; + for (auto it : itr) { + if (it.first != "node2") { + temp.insert({it.first, it.second}); + } + } + hint.insert(temp); + } + } + auto status = + EncapsulateClusters(&g, 0, fdeflib_new, {{"ngraph_device_id", ""}}, + make_pair(true, hint)); + if (did_aot[i]) { + ASSERT_OK(status); + } else { + ASSERT_NOT_OK(status); + continue; + } + + int num_encapsulates = 0; + int num_tf_nodes = 0; + for (auto itr : g.nodes()) { + auto node_type = itr->type_string(); + num_encapsulates += (node_type == "NGraphEncapsulate" ? 1 : 0); + num_tf_nodes += ((node_type == "Add" || node_type == "Const") ? 1 : 0); + } + + ASSERT_EQ(num_encapsulates, fdeflib_new->function_size()); + ASSERT_EQ(num_encapsulates, 1); + + // No Add or Const nodes left in the graph + ASSERT_EQ(num_tf_nodes, 0); + + auto get_attr_name_for_aot = [&i, &using_placeholder](bool is_exec) { + string attrname = "_ngraph_aot_"; + attrname += (string(is_exec ? "ngexec" : "ngfunction") + "_"); + string signature_portion = (string(i == 0 ? "" : "2,2,") + ";"); + attrname += signature_portion; + attrname += (using_placeholder ? signature_portion : ""); + attrname += "/"; + return attrname; + }; + + for (auto itr : g.nodes()) { + if (itr->type_string() == "NGraphEncapsulate") { + string aot_info; + bool found_exec = + GetNodeAttr(itr->attrs(), get_attr_name_for_aot(true), + &aot_info) == tensorflow::Status::OK(); + bool found_function = + GetNodeAttr(itr->attrs(), get_attr_name_for_aot(false), + &aot_info) == tensorflow::Status::OK(); + ASSERT_TRUE(found_exec == did_aot[i]); + ASSERT_TRUE(found_function == did_aot[i]); + } + } + } + + free(fdeflib_new); + } +} + +// Placeholder-->Add(0)--->Abs(1)-->IdN +// ^ +// | +// Placeholder +// 2 Encapsulates connected in serial. Cannot AOT here +TEST(EncapsulateClusters, AOT1) { + if (!ngraph_tf_is_grappler_enabled()) return; // GTEST_SKIP() did not compile + NGraphClusterManager::EvictAllClusters(); + Graph g(OpRegistry::Global()); + + int cluster_idx = NGraphClusterManager::NewCluster(); + + Node* node1; + Node* node2; + + ASSERT_OK(NodeBuilder("node1", "Placeholder") + .Attr("dtype", DT_FLOAT) + .Finalize(&g, &node1)); + ASSERT_OK(NodeBuilder("node2", "Placeholder") + .Attr("dtype", DT_FLOAT) + .Finalize(&g, &node2)); + + Node* node3; + ASSERT_OK(NodeBuilder("node3", "Add") + .Input(node1, 0) + .Input(node2, 0) + .Attr("T", DT_FLOAT) + .Attr("_ngraph_marked_for_clustering", true) + .Attr("_ngraph_cluster", cluster_idx) + .Attr("_ngraph_backend", "INTERPRETER") + .Finalize(&g, &node3)); + + Node* node4; + cluster_idx = NGraphClusterManager::NewCluster(); + ASSERT_OK(NodeBuilder("node4", "Abs") + .Input(node3, 0) + .Attr("T", DT_FLOAT) + .Attr("_ngraph_marked_for_clustering", true) + .Attr("_ngraph_cluster", cluster_idx) + .Attr("_ngraph_backend", "INTERPRETER") + .Finalize(&g, &node4)); + + Node* node5; + std::vector inputs; + std::vector input_types; + inputs.push_back(NodeBuilder::NodeOut(node4, 0)); + input_types.push_back(node4->output_type(0)); + ASSERT_OK(NodeBuilder("node5", "IdentityN") + .Input(inputs) + .Attr("T", input_types) + .Finalize(&g, &node5)); + + Node* source = g.source_node(); + Node* sink = g.sink_node(); + g.AddEdge(source, Graph::kControlSlot, node1, Graph::kControlSlot); + g.AddEdge(source, Graph::kControlSlot, node2, Graph::kControlSlot); + g.AddEdge(node5, Graph::kControlSlot, sink, Graph::kControlSlot); + + FunctionDefLibrary* fdeflib_new = new FunctionDefLibrary(); + + std::vector> node_shapes_hints_vect = { + {}, {{{"node1", {2, 2}}, {"node2", {2, 2}}}}}; + int num_cases = node_shapes_hints_vect.size(); + for (int i = 0; i < num_cases; i++) { + ASSERT_NOT_OK( + EncapsulateClusters(&g, 0, fdeflib_new, {{"ngraph_device_id", ""}}, + make_pair(true, node_shapes_hints_vect[i]))); + } + + free(fdeflib_new); +} + +// Placeholder-->Abs(0)-->IdN +// +// Placeholder-->Abs(1)-->IdN +// 2 Encapsulates connected in parallel +TEST(EncapsulateClusters, AOT2) { + if (!ngraph_tf_is_grappler_enabled()) return; + NGraphClusterManager::EvictAllClusters(); + Graph g(OpRegistry::Global()); + + int cluster_idx = NGraphClusterManager::NewCluster(); + + Node* node1; + Node* node2; + + ASSERT_OK(NodeBuilder("node1", "Placeholder") + .Attr("dtype", DT_FLOAT) + .Finalize(&g, &node1)); + ASSERT_OK(NodeBuilder("node2", "Placeholder") + .Attr("dtype", DT_FLOAT) + .Finalize(&g, &node2)); + Node* node3; + ASSERT_OK(NodeBuilder("node3", "Abs") + .Input(node1, 0) + .Attr("T", DT_FLOAT) + .Attr("_ngraph_marked_for_clustering", true) + .Attr("_ngraph_cluster", cluster_idx) + .Attr("_ngraph_backend", "INTERPRETER") + .Finalize(&g, &node3)); + + Node* node4; + cluster_idx = NGraphClusterManager::NewCluster(); + ASSERT_OK(NodeBuilder("node4", "Abs") + .Input(node2, 0) + .Attr("T", DT_FLOAT) + .Attr("_ngraph_marked_for_clustering", true) + .Attr("_ngraph_cluster", cluster_idx) + .Attr("_ngraph_backend", "INTERPRETER") + .Finalize(&g, &node4)); + + Node* node5; + std::vector inputs; + std::vector input_types; + inputs.push_back(NodeBuilder::NodeOut(node3, 0)); + input_types.push_back(node3->output_type(0)); + ASSERT_OK(NodeBuilder("node5", "IdentityN") + .Input(inputs) + .Attr("T", input_types) + .Finalize(&g, &node5)); + Node* node6; + inputs.clear(); + input_types.clear(); + inputs.push_back(NodeBuilder::NodeOut(node4, 0)); + input_types.push_back(node4->output_type(0)); + ASSERT_OK(NodeBuilder("node6", "IdentityN") + .Input(inputs) + .Attr("T", input_types) + .Finalize(&g, &node6)); + + Node* source = g.source_node(); + Node* sink = g.sink_node(); + g.AddEdge(source, Graph::kControlSlot, node1, Graph::kControlSlot); + g.AddEdge(source, Graph::kControlSlot, node2, Graph::kControlSlot); + g.AddEdge(node5, Graph::kControlSlot, sink, Graph::kControlSlot); + g.AddEdge(node6, Graph::kControlSlot, sink, Graph::kControlSlot); + + FunctionDefLibrary* fdeflib_new = new FunctionDefLibrary(); + + std::vector> node_shapes_hints_vect = { + {}, + {{{"node1", {2, 2}}, {"node2", {2, 2}}}, + {{"node1", {2, 3}}, {"node2", {2, 3}}}}}; + std::vector did_aot = {true, true}; + int num_cases = node_shapes_hints_vect.size(); + for (int i = 0; i < num_cases; i++) { + auto status = + EncapsulateClusters(&g, 0, fdeflib_new, {{"ngraph_device_id", ""}}, + make_pair(true, node_shapes_hints_vect[i])); + if (did_aot[i]) { + ASSERT_OK(status); + } else { + ASSERT_NOT_OK(status); + continue; + } + + int num_encapsulates = 0; + int num_tf_nodes = 0; + for (auto itr : g.nodes()) { + auto node_type = itr->type_string(); + num_encapsulates += (node_type == "NGraphEncapsulate" ? 1 : 0); + num_tf_nodes += ((node_type == "Abs") ? 1 : 0); + } + + ASSERT_EQ(num_encapsulates, fdeflib_new->function_size()); + ASSERT_EQ(num_encapsulates, 2); + + // No Abs nodes left in the graph + ASSERT_EQ(num_tf_nodes, 0); + + for (auto itr : g.nodes()) { + if (itr->type_string() == "NGraphEncapsulate") { + string aot_info; + bool found_exec = + GetNodeAttr(itr->attrs(), string("_ngraph_aot_ngexec_") + + (i == 0 ? "" : "2,2,") + ";/", + &aot_info) == tensorflow::Status::OK(); + bool found_function = + GetNodeAttr(itr->attrs(), string("_ngraph_aot_ngfunction_") + + (i == 0 ? "" : "2,2,") + ";/", + &aot_info) == tensorflow::Status::OK(); + if (i == 1) { + bool found_second_exec = + GetNodeAttr(itr->attrs(), "_ngraph_aot_ngexec_2,3,;/", + &aot_info) == tensorflow::Status::OK(); + found_exec = found_exec && found_second_exec; + bool found_second_function = + GetNodeAttr(itr->attrs(), "_ngraph_aot_ngfunction_2,3,;/", + &aot_info) == tensorflow::Status::OK(); + found_function = found_function && found_second_function; + } + ASSERT_TRUE(found_exec == did_aot[i]); + ASSERT_TRUE(found_function == did_aot[i]); + } + } + } + + free(fdeflib_new); +} + +// Placeholder-->Add(0)--->IdN +// ^ +// | +// Placeholder +// Passing shape hints that will cause TranslateGraph to fail +TEST(EncapsulateClusters, AOT3) { + if (!ngraph_tf_is_grappler_enabled()) return; + + NGraphClusterManager::EvictAllClusters(); + Graph g(OpRegistry::Global()); + + int cluster_idx = NGraphClusterManager::NewCluster(); + + Node* node1; + Node* node2; + + ASSERT_OK(NodeBuilder("node1", "Placeholder") + .Attr("dtype", DT_FLOAT) + .Finalize(&g, &node1)); + + ASSERT_OK(NodeBuilder("node2", "Placeholder") + .Attr("dtype", DT_FLOAT) + .Finalize(&g, &node2)); + + Node* node3; + ASSERT_OK(NodeBuilder("node3", "Add") + .Input(node1, 0) + .Input(node2, 0) + .Attr("T", DT_FLOAT) + .Attr("_ngraph_marked_for_clustering", true) + .Attr("_ngraph_cluster", cluster_idx) + .Attr("_ngraph_backend", "INTERPRETER") + .Finalize(&g, &node3)); + Node* node4; + std::vector inputs; + std::vector input_types; + inputs.push_back(NodeBuilder::NodeOut(node3, 0)); + input_types.push_back(node3->output_type(0)); + ASSERT_OK(NodeBuilder("node4", "IdentityN") + .Input(inputs) + .Attr("T", input_types) + .Finalize(&g, &node4)); + + Node* source = g.source_node(); + Node* sink = g.sink_node(); + g.AddEdge(source, Graph::kControlSlot, node1, Graph::kControlSlot); + g.AddEdge(source, Graph::kControlSlot, node2, Graph::kControlSlot); + g.AddEdge(node4, Graph::kControlSlot, sink, Graph::kControlSlot); + + FunctionDefLibrary* fdeflib_new = new FunctionDefLibrary(); + + std::vector> node_shapes_hints_vect = { + {{{"node1", {2, 4}}, {"node2", {2, 2}}}}, + {{{"node1", {3, 2}}, {"node2", {2, 2}}}}}; + int num_cases = node_shapes_hints_vect.size(); + for (int i = 0; i < num_cases; i++) { + ASSERT_NOT_OK( + EncapsulateClusters(&g, 0, fdeflib_new, {{"ngraph_device_id", ""}}, + make_pair(true, node_shapes_hints_vect[i]))); + } + + free(fdeflib_new); +} + +// Placeholder-->Add(0)--->IdN +// ^ +// | +// Placeholder +// Placeholders contain full shape information. So AOT can happen even hints are +// empty +TEST(EncapsulateClusters, AOT4) { + if (!ngraph_tf_is_grappler_enabled()) return; + + NGraphClusterManager::EvictAllClusters(); + Graph g(OpRegistry::Global()); + + int cluster_idx = NGraphClusterManager::NewCluster(); + + Node* node1; + Node* node2; + + ASSERT_OK(NodeBuilder("node1", "Placeholder") + .Attr("dtype", DT_FLOAT) + .Attr("shape", TensorShape{2, 3}) + .Finalize(&g, &node1)); + + ASSERT_OK(NodeBuilder("node2", "Placeholder") + .Attr("dtype", DT_FLOAT) + .Attr("shape", TensorShape{2, 3}) + .Finalize(&g, &node2)); + + Node* node3; + ASSERT_OK(NodeBuilder("node3", "Add") + .Input(node1, 0) + .Input(node2, 0) + .Attr("T", DT_FLOAT) + .Attr("_ngraph_marked_for_clustering", true) + .Attr("_ngraph_cluster", cluster_idx) + .Attr("_ngraph_backend", "INTERPRETER") + .Finalize(&g, &node3)); + Node* node4; + std::vector inputs; + std::vector input_types; + inputs.push_back(NodeBuilder::NodeOut(node3, 0)); + input_types.push_back(node3->output_type(0)); + ASSERT_OK(NodeBuilder("node4", "IdentityN") + .Input(inputs) + .Attr("T", input_types) + .Finalize(&g, &node4)); + + Node* source = g.source_node(); + Node* sink = g.sink_node(); + g.AddEdge(source, Graph::kControlSlot, node1, Graph::kControlSlot); + g.AddEdge(source, Graph::kControlSlot, node2, Graph::kControlSlot); + g.AddEdge(node4, Graph::kControlSlot, sink, Graph::kControlSlot); + + FunctionDefLibrary* fdeflib_new = new FunctionDefLibrary(); + + // The first hint is empty (but placeholders contain complete information so + // can AOT) + // The second hint contains info that matches info in placeholders + // The third hint contains hints that do not match. So Encapsulation fails + std::vector> node_shapes_hints_vect = { + {}, + {{{"node1", {2, 3}}, {"node2", {2, 3}}}}, + {{{"node1", {5, 10}}, {"node2", {15, 20}}}}}; + std::vector did_aot = {true, true, false}; + int num_cases = node_shapes_hints_vect.size(); + for (int i = 0; i < num_cases; i++) { + auto encapsulate_status = + EncapsulateClusters(&g, 0, fdeflib_new, {{"ngraph_device_id", ""}}, + make_pair(true, node_shapes_hints_vect[i])); + if (did_aot[i]) { + ASSERT_OK(encapsulate_status); + } else { + ASSERT_NOT_OK(encapsulate_status); + continue; + } + + int num_encapsulates = 0; + int num_tf_nodes = 0; + for (auto itr : g.nodes()) { + auto node_type = itr->type_string(); + num_encapsulates += (node_type == "NGraphEncapsulate" ? 1 : 0); + num_tf_nodes += ((node_type == "Add" || node_type == "Const") ? 1 : 0); + } + + ASSERT_EQ(num_encapsulates, fdeflib_new->function_size()); + ASSERT_EQ(num_encapsulates, 1); + + // No Add or Const nodes left in the graph + ASSERT_EQ(num_tf_nodes, 0); + + for (auto itr : g.nodes()) { + if (itr->type_string() == "NGraphEncapsulate") { + string aot_info; + bool found_exec = + GetNodeAttr(itr->attrs(), "_ngraph_aot_ngexec_2,3,;2,3,;/", + &aot_info) == tensorflow::Status::OK(); + bool found_function = + GetNodeAttr(itr->attrs(), "_ngraph_aot_ngfunction_2,3,;2,3,;/", + &aot_info) == tensorflow::Status::OK(); + ASSERT_TRUE(found_exec == did_aot[i]); + ASSERT_TRUE(found_function == did_aot[i]); + } + } + } + + free(fdeflib_new); +} + +// Test cases for AOT: +// TODO: 1. what of scalar inputs. placeholder shape is {}? +// 2. Shape hints that cause errors in TranslateGraph?. eg trying to add [2,2] +// with [2,4]? (done, AOT3) +// 3. Encapsulate being fed by another enc (wont work) (done, AOT1) +// 4. 2 encapsulates, but both are attached to inputs, so we can AOT (done, +// AOT2) +// 5. Have a test where enc is fed by a const and a placeholder (done, AOT0) +// 6. Placeholders contain full shape. Then even with no shape hints, AOT can +// happen (done, AOT4) +// 7. Fail on bad hints (partly done, AOT4) +// 8. EncapsulateClusters compiles 2 executables due to 2 different shape hints +// (done, AOT2) +// TODO: 9. Same hint passed twice or functionally same hint passed twice (it +// should create 1 exec only). } } } \ No newline at end of file diff --git a/test/graph_rewrites/enter_in_catalog_test.cc b/test/graph_rewrites/enter_in_catalog_test.cc index c31872e5d..51321aac6 100644 --- a/test/graph_rewrites/enter_in_catalog_test.cc +++ b/test/graph_rewrites/enter_in_catalog_test.cc @@ -78,7 +78,7 @@ TEST(CatalogTest, SmallGraph1) { FunctionDefLibrary* fdeflib_new = new FunctionDefLibrary(); std::unordered_map config_map; config_map["ngraph_device_id"] = ""; - ASSERT_OK(EncapsulateClusters(&graph, 0, fdeflib_new, config_map)); + ASSERT_OK(EncapsulateClusters(&graph, 0, fdeflib_new, config_map, {0, {}})); ASSERT_OK(EnterInCatalog(&graph, 0)); bool remove = false; @@ -122,7 +122,7 @@ TEST(CatalogTest, SmallGraph2) { FunctionDefLibrary* fdeflib_new = new FunctionDefLibrary(); std::unordered_map config_map; config_map["ngraph_device_id"] = ""; - ASSERT_OK(EncapsulateClusters(&graph, 0, fdeflib_new, config_map)); + ASSERT_OK(EncapsulateClusters(&graph, 0, fdeflib_new, config_map, {0, {}})); ASSERT_OK(EnterInCatalog(&graph, 0)); bool remove = false; @@ -193,7 +193,7 @@ TEST(CatalogTest, SmallGraph3) { FunctionDefLibrary* fdeflib_new = new FunctionDefLibrary(); std::unordered_map config_map; config_map["ngraph_device_id"] = ""; - ASSERT_OK(EncapsulateClusters(&graph, 0, fdeflib_new, config_map)); + ASSERT_OK(EncapsulateClusters(&graph, 0, fdeflib_new, config_map, {0, {}})); ASSERT_OK(EnterInCatalog(&graph, 0)); // check if the _ngraph_remove attribute is added/not-added as expected @@ -259,7 +259,7 @@ TEST(CatalogTest, SmallGraph4) { FunctionDefLibrary* fdeflib_new = new FunctionDefLibrary(); std::unordered_map config_map; config_map["ngraph_device_id"] = ""; - ASSERT_OK(EncapsulateClusters(&graph, 0, fdeflib_new, config_map)); + ASSERT_OK(EncapsulateClusters(&graph, 0, fdeflib_new, config_map, {0, {}})); ASSERT_OK(EnterInCatalog(&graph, 0)); string key; diff --git a/test/graph_rewrites/partial_shapes_test.cc b/test/graph_rewrites/partial_shapes_test.cc new file mode 100644 index 000000000..3d0254e39 --- /dev/null +++ b/test/graph_rewrites/partial_shapes_test.cc @@ -0,0 +1,139 @@ +/******************************************************************************* + * Copyright 2019 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + *******************************************************************************/ +#include "gtest/gtest.h" + +#include "ngraph_bridge/ngraph_partial_shapes.h" +#include "ngraph_bridge/ngraph_utils.h" +#include "test/test_utilities.h" + +using namespace std; +namespace ng = ngraph; + +namespace tensorflow { + +namespace ngraph_bridge { + +namespace testing { + +#define ASSERT_OK(x) ASSERT_EQ((x), ::tensorflow::Status::OK()); +#define ASSERT_NOT_OK(x) ASSERT_NE((x), ::tensorflow::Status::OK()); + +// The result of concretize would be {2, 3} +TEST(PartialShapes, ValidConcretize1) { + PartialShape p1({2, -1}); + PartialShape p2({-1, 3}); + p1.concretize(p2); + ASSERT_EQ(p1.is_valid(), true); +} + +// The result of concretize would be {-1, -1} +TEST(PartialShapes, ValidConcretize2) { + PartialShape p1({-1, -1}); + PartialShape p2({-1, -1}); + p1.concretize(p2); + ASSERT_EQ(p1.is_valid(), true); + ASSERT_EQ(p1.to_string(), "valid:-1,-1,"); +} + +// The result of concretize would be {} +TEST(PartialShapes, ValidConcretize3) { + PartialShape p1(vector{}); + PartialShape p2(vector{}); + p1.concretize(p2); + ASSERT_EQ(p1.is_valid(), true); + ASSERT_EQ(p1.to_string(), "valid:"); +} + +// This would result in an invalid case and should fail because +// p1[0] = 2 does not match p2[0] = 3 +TEST(PartialShapes, InvalidConcretize1) { + PartialShape p1({2, -1}); + PartialShape p2({3, 3}); + p1.concretize(p2); + ASSERT_EQ(p1.is_valid(), false); +} + +// This would result in an invalid case and should fail because +// the ranks are not same +TEST(PartialShapes, InvalidConcretize2) { + PartialShape p1({2, -1, -1}); + PartialShape p2({3, 3}); + p1.concretize(p2); + ASSERT_EQ(p1.is_valid(), false); +} + +// The result of concretize would be {3, 3}, which is a concrete shape +TEST(PartialShapes, IsConcrete1) { + PartialShape p1({-1, -1}); + PartialShape p2({3, 3}); + p1.concretize(p2); + ASSERT_EQ(p1.is_concrete(), true); + ASSERT_EQ(p2.is_concrete(), true); + ASSERT_EQ(p1.to_string(), "valid:3,3,"); +} + +// The result of concretize would be {3, -1}, which is not a concrete shape +TEST(PartialShapes, IsConcrete2) { + PartialShape p1({-1, -1}); + PartialShape p2({3, -1}); + p1.concretize(p2); + ASSERT_EQ(p1.is_concrete(), false); + ASSERT_EQ(p1.is_valid(), true); + ASSERT_EQ(p2.is_valid(), true); +} + +// The result of concretize would be {}, which is a concrete shape +TEST(PartialShapes, IsConcrete3) { + PartialShape p1(vector{}); + PartialShape p2(vector{}); + p1.concretize(p2); + ASSERT_EQ(p1.is_concrete(), true); +} + +// Test default constructor +TEST(PartialShapes, DefaultConstructor) { + PartialShape p1; + ASSERT_EQ(p1.is_valid(), false); +} + +TEST(PartialShapes, Constructor) { + PartialShape p1({-2, 1}); + ASSERT_EQ(p1.is_valid(), false); +} + +// Concretizing a shape with invalid shape should fail +TEST(PartialShapes, ConcretizeWithInvalid) { + PartialShape p1({2, -1}); + PartialShape p2({3, -1}); + p1.concretize(p2); // p1 is now invalid + ASSERT_EQ(p1.is_valid(), false); + ASSERT_EQ(p2.is_valid(), true); + + PartialShape p3({-1, -1}); + bool test_passed = true; + try { + p3.concretize(p1); + } catch (...) { + test_passed = false; + } + ASSERT_EQ(test_passed, false); +} + +} // namespace testing + +} // namespace ngraph_bridge + +} // namespace tensorflow diff --git a/test/graph_rewrites/remove_ngraphassigns.cc b/test/graph_rewrites/remove_ngraphassigns.cc index 3d501080a..521c203a4 100644 --- a/test/graph_rewrites/remove_ngraphassigns.cc +++ b/test/graph_rewrites/remove_ngraphassigns.cc @@ -72,7 +72,7 @@ TEST(RemoveNGraphAssigns, Graph1) { FunctionDefLibrary* fdeflib_new = new FunctionDefLibrary(); std::unordered_map config_map; config_map["ngraph_device_id"] = ""; - ASSERT_OK(EncapsulateClusters(&graph, 0, fdeflib_new, config_map)); + ASSERT_OK(EncapsulateClusters(&graph, 0, fdeflib_new, config_map, {0, {}})); // Get all the nodes in map [utility] map node_map; @@ -130,7 +130,7 @@ TEST(RemoveNGraphAssigns, Graph2) { FunctionDefLibrary* fdeflib_new = new FunctionDefLibrary(); std::unordered_map config_map; config_map["ngraph_device_id"] = ""; - ASSERT_OK(EncapsulateClusters(&graph, 0, fdeflib_new, config_map)); + ASSERT_OK(EncapsulateClusters(&graph, 0, fdeflib_new, config_map, {0, {}})); // clean up config::ngraph_set_disabled_ops(""); @@ -231,7 +231,7 @@ TEST(RemoveNGraphAssigns, Graph3) { FunctionDefLibrary* fdeflib_new = new FunctionDefLibrary(); std::unordered_map config_map; config_map["ngraph_device_id"] = ""; - ASSERT_OK(EncapsulateClusters(&graph, 0, fdeflib_new, config_map)); + ASSERT_OK(EncapsulateClusters(&graph, 0, fdeflib_new, config_map, {0, {}})); // Get all the nodes in map [utility] map node_map; @@ -337,7 +337,7 @@ TEST(RemoveNGraphAssigns, Graph4) { FunctionDefLibrary* fdeflib_new = new FunctionDefLibrary(); std::unordered_map config_map; config_map["ngraph_device_id"] = ""; - ASSERT_OK(EncapsulateClusters(&graph, 0, fdeflib_new, config_map)); + ASSERT_OK(EncapsulateClusters(&graph, 0, fdeflib_new, config_map, {0, {}})); // clean up config::ngraph_set_disabled_ops(""); @@ -442,7 +442,7 @@ TEST(RemoveNGraphAssigns, Graph5) { FunctionDefLibrary* fdeflib_new = new FunctionDefLibrary(); std::unordered_map config_map; config_map["ngraph_device_id"] = ""; - ASSERT_OK(EncapsulateClusters(&graph, 0, fdeflib_new, config_map)); + ASSERT_OK(EncapsulateClusters(&graph, 0, fdeflib_new, config_map, {0, {}})); // Get all the nodes in map [utility] map node_map; diff --git a/test/python/test_tf2ngraph_functions.py b/test/python/test_tf2ngraph_functions.py index 102cb98aa..a36f6f4d5 100644 --- a/test/python/test_tf2ngraph_functions.py +++ b/test/python/test_tf2ngraph_functions.py @@ -28,7 +28,7 @@ import ngraph_bridge from tools.build_utils import command_executor -from tools.tf2ngraph import parse_extra_params_string, update_config_to_include_custom_config +from tools.tf2ngraph import update_config_to_include_custom_config from common import NgraphTest @@ -40,53 +40,28 @@ def test_config_updater_api(self): '0', { 'abc': '1', 'def': '2' - }) + }, [{ + 'x': [1] + }], True) assert config.HasField('graph_options') assert config.graph_options.HasField('rewrite_options') custom_opts = config.graph_options.rewrite_options.custom_optimizers assert len(custom_opts) == 1 assert custom_opts[0].name == 'ngraph-optimizer' assert set(custom_opts[0].parameter_map.keys()) == { - 'abc', 'ngraph_backend', 'def', 'device_id' + 'abc', 'ngraph_backend', 'def', 'device_id', 'aot_requested', + 'shape_hint_0' } retrieved_dict = {} for key, val in custom_opts[0].parameter_map.items(): - retrieved_dict[key] = val.ListFields()[0][1].decode() + # For everything other than shape_hint_0, the values are simple strings + # shape_hint_0 contains a complex data structure representing shape hint + if (key != "shape_hint_0"): + retrieved_dict[key] = val.ListFields()[0][1].decode() assert retrieved_dict == { 'abc': '1', 'def': '2', 'ngraph_backend': 'CPU', - 'device_id': '0' + 'device_id': '0', + 'aot_requested': '1' } - - # In this test we test that the spaces are handled correctly - @pytest.mark.parametrize(('ng_device',), - (('{abc:1,def:2}',), ('{abc:1, def:2}',), - ('{abc :1,def: 2}',), ('{abc:1,def: 2} ',))) - def parse_extra_params_string(self, raw_string): - assert (parse_extra_params_string(raw_string) == { - 'abc': '1', - 'def': '2' - }) - - # This test checks the parsing of the empty dictionary {} - def parse_extra_params_empty(self): - assert (parse_extra_params_string('{}') == {}) - - # In this test we pass badly formatted strings, - # and expect parse_extra_params_string to fail - @pytest.mark.parametrize(('ng_device',), ( - ('abc:1,def:2}',), - ('{abc:1, def:2',), - ('{abc :1:2 ,def: 2}',), - )) - def parse_extra_params_string(self, raw_string): - function_failed = False - try: - parse_extra_params_string(raw_string) - except: - function_failed = True - assert function_failed - - # TODO: write a test for parse_extra_params_string for incorrect parsings - # where parse_extra_params_string is expected to fail diff --git a/test/python/test_tf2ngraph_script.py b/test/python/test_tf2ngraph_script.py index f87d89fbc..8b92ceab5 100644 --- a/test/python/test_tf2ngraph_script.py +++ b/test/python/test_tf2ngraph_script.py @@ -28,7 +28,7 @@ import ngraph_bridge from tools.build_utils import command_executor -from tools.tf2ngraph import convert, get_gdef +from tools.tf2ngraph import convert, get_gdef, Tf2ngraphJson from common import NgraphTest @@ -58,9 +58,12 @@ def format_and_loc_match(format, loc): ('pb',), ('savedmodel',), )) - @pytest.mark.parametrize(('ng_device',), (('CPU',), ('INTERPRETER',))) + @pytest.mark.parametrize(('ng_device', 'shape_hints', 'precompile'), + (('CPU', [], False), ('INTERPRETER', [{}], True), + ('INTERPRETER', [], False))) + # In sample_graph.pbtxt, the input shape is fully specified, so we don't need to pass shape hints for precompile def test_command_line_api(self, inp_format, inp_loc, out_format, - commandline, ng_device): + commandline, ng_device, shape_hints, precompile): # Only run this test when grappler is enabled if not ngraph_bridge.is_grappler_enabled(): return @@ -73,31 +76,42 @@ def test_command_line_api(self, inp_format, inp_loc, out_format, pass conversion_successful = False try: - extra_params = { - 'CPU': '{device_config:0}', - 'INTERPRETER': '{test_echo:1}' + optional_backend_params = { + 'CPU': { + 'device_config': '0' + }, + 'INTERPRETER': { + 'test_echo': '1' + } }[ng_device] + config_file_name = 'temp_config_file.json' + Tf2ngraphJson.dump_json(config_file_name, optional_backend_params, + shape_hints) if commandline: # In CI this test is expected to be run out of artifacts/test/python - command_executor('python ../../tools/tf2ngraph.py --input_' + - inp_format + ' ' + inp_loc + - ' --output_nodes out_node --output_' + - out_format + ' ' + out_loc + ' --ng_backend ' + - ng_device + ' --extra_params ' + extra_params) + command_executor( + 'python ../../tools/tf2ngraph.py --input_' + inp_format + + ' ' + inp_loc + ' --output_nodes out_node --output_' + + out_format + ' ' + out_loc + ' --ng_backend ' + ng_device + + ' --config_file ' + config_file_name + + ("", " --precompile ")[precompile]) else: convert(inp_format, inp_loc, out_format, out_loc, ['out_node'], - ng_device, extra_params) + ng_device, optional_backend_params, shape_hints, + precompile) conversion_successful = True finally: if not conversion_successful: try: (shutil.rmtree, os.remove)[os.path.isfile(out_loc)](out_loc) + os.remove(config_file_name) except: pass assert conversion_successful gdef = get_gdef(out_format, out_loc) (shutil.rmtree, os.remove)[os.path.isfile(out_loc)](out_loc) + os.remove(config_file_name) with tf.Graph().as_default() as g: tf.import_graph_def(gdef, name='') diff --git a/test/python/test_tf2ngraph_shape_hints.py b/test/python/test_tf2ngraph_shape_hints.py new file mode 100644 index 000000000..c7dc9d274 --- /dev/null +++ b/test/python/test_tf2ngraph_shape_hints.py @@ -0,0 +1,172 @@ +# ============================================================================== +# Copyright 2019 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""nGraph TensorFlow bridge test for tf2ngraph script for precompilation with shape hints + +""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import pytest +import os +import numpy as np +import shutil +import tensorflow as tf +import ngraph_bridge +import json + +from google.protobuf import text_format +from tensorflow.core.framework import graph_pb2 +from tools.build_utils import command_executor +from tools.tf2ngraph import convert, get_gdef, Tf2ngraphJson + +from common import NgraphTest + + +def get_pbtxt_name(tag, p0_shape, p1_shape): + return tag + ','.join(map(lambda x: str(x), p0_shape)) + '__' + ','.join( + map(lambda x: str(x), p1_shape)) + '.pbtxt' + + +def create_graph(p0_shape, p1_shape): + temp_pbtxt_name = get_pbtxt_name('temp_graph_in_', p0_shape, p1_shape) + with tf.Session() as sess: + x = tf.placeholder(tf.float32, shape=p0_shape, name='x') + y = tf.placeholder(tf.float32, shape=p1_shape, name='y') + z = tf.add(tf.abs(x), tf.abs(y), name="z") + tf.io.write_graph(sess.graph, '.', temp_pbtxt_name, as_text=True) + return x, y, z, temp_pbtxt_name + + +def get_inputs(p_shape): + return np.random.rand(*p_shape) + + +def run_pbtxt(pbtxt_filename, inp0, inp1): + tf.reset_default_graph() + gdef = graph_pb2.GraphDef() + with open(pbtxt_filename, 'r') as f: + raw_contents = f.read() + text_format.Parse(raw_contents, gdef) + with tf.Session() as sess: + tf.import_graph_def(gdef, name='') + x = tf.get_default_graph().get_tensor_by_name("x:0") + y = tf.get_default_graph().get_tensor_by_name("y:0") + z = tf.get_default_graph().get_tensor_by_name("z:0") + return sess.run(z, feed_dict={x: inp0, y: inp1}) + + +def check_pbtxt_has_exec(pbtxt_filename, num_expected_execs): + with open(pbtxt_filename, 'r') as f: + contents = '\n'.join(f.readlines()) + assert contents.count('_ngraph_aot_requested') == 1 + assert contents.count('_ngraph_aot_ngexec_') == num_expected_execs + assert contents.count('_ngraph_aot_ngfunction_') == num_expected_execs + + +def helper(p0_shape, p1_shape, p0_actual_shape, p1_actual_shape, shapehints): + inp0 = get_inputs(p0_actual_shape) + inp1 = get_inputs(p1_actual_shape) + x, y, z, temp_in_pbtxt_name = create_graph(p0_shape, p1_shape) + temp_out_pbtxt_name = get_pbtxt_name('temp_graph_out_', p0_shape, p1_shape) + json_name = 'temp_config_file.json' + # shapehints is a list of dictionaries (keys are node names, vals are lists (of shapes)) + Tf2ngraphJson.dump_json(json_name, None, shapehints) + + command_executor('python ../../tools/tf2ngraph.py --input_pbtxt ' + + temp_in_pbtxt_name + ' --output_nodes z --output_pbtxt ' + + temp_out_pbtxt_name + ' --ng_backend INTERPRETER ' + + ' --config_file ' + json_name + ' --precompile') + + num_expected_execs = (len(shapehints), 1)[len(shapehints) == 0] + check_pbtxt_has_exec(temp_out_pbtxt_name, num_expected_execs) + + tf_out_val = run_pbtxt(temp_in_pbtxt_name, inp0, inp1) + ng_out_vals = run_pbtxt(temp_out_pbtxt_name, inp0, inp1) + assert ((tf_out_val == ng_out_vals).all()) + + os.remove(temp_in_pbtxt_name) + os.remove(temp_out_pbtxt_name) + os.remove(json_name) + + +# TODO: Add more test cases +class Testtf2ngraphShapehints(NgraphTest): + + @pytest.mark.parametrize( + ('p0_shape', 'p1_shape', 'p0_actual_shape', 'p1_actual_shape', + 'shapehints'), + ( + ([2, 2], [2, 2], [2, 2], [2, 2], [{} + ]), # np input needs shape hints + ([2, 2], [None, 2], [2, 2], [2, 2], [{ + 'y': [2, -1] + }]), # only 1 input needs shape hints + ( + [2, None], + [None, 3], + [2, 3], + [2, 3], + [{ + 'y': [2, -1], + 'x': [2, 3] # both inputs need shape hints + }]), + ([None, None], [None, None], [5, 1], [5, 1], [{ + 'y': [2, 3], + 'x': [2, 3] + }, { + 'y': [5, 1], + 'x': [5, 1] + }]), # 2 executables are compiled + )) + @pytest.mark.skipif( + not ngraph_bridge.is_grappler_enabled(), + reason="Requires grappler build for tf2ngraph and AOT") + def test_tf2ngraph_with_shape_hints_0(self, p0_shape, p1_shape, + p0_actual_shape, p1_actual_shape, + shapehints): + helper(p0_shape, p1_shape, p0_actual_shape, p1_actual_shape, shapehints) + + @pytest.mark.parametrize( + ('p0_shape', 'p1_shape', 'p0_actual_shape', 'p1_actual_shape', + 'shapehints'), + ( + ([2, 2], [None, 2], [2, 2], [2, 2], [{ + 'y': [2, 3] + }]), # conflicting shape hint + ([2, 2], [None, 2], [2, 2], [2, 2], [{ + 'y': [2] + }]), # shape hint is of conflicting rank + ([2, 2], [None, 2], [2, 5], [2, 5], [{ + 'y': [2, 2] + }]), # During run time bad shapes are passed + ([2, 2], [None, 2], [2, 2], [2, 2], [{ + 'x': [2, -1] + }]), # Input y does not have enough hints to concretize it + ([2, 2], [None, 2], [2, 2], [2, 2], [{ + 'y': [2, -1], + 'bogus': [1, 2] + }]), # passing a bogus node name + )) + @pytest.mark.skipif( + not ngraph_bridge.is_grappler_enabled(), + reason="Requires grappler build for tf2ngraph and AOT") + def test_tf2ngraph_with_shape_hints_1(self, p0_shape, p1_shape, + p0_actual_shape, p1_actual_shape, + shapehints): + with pytest.raises(Exception): + helper(p0_shape, p1_shape, p0_actual_shape, p1_actual_shape, + shapehints) diff --git a/tools/run_tf2ngraph_model.py b/tools/run_tf2ngraph_model.py new file mode 100644 index 000000000..a8f07f66a --- /dev/null +++ b/tools/run_tf2ngraph_model.py @@ -0,0 +1,33 @@ +import tensorflow as tf +import ngraph_bridge +from google.protobuf import text_format +import pdb + + +def get_graphdef(pb_filename): + graph_def = tf.GraphDef() + if pb_filename.endswith("pbtxt"): + with open(pb_filename, "r") as f: + text_format.Merge(f.read(), graph_def) + else: + with open(pb_filename, "rb") as f: + graph_def.ParseFromString(f.read()) + return graph_def + + +inputs = ['import/x:0', 'import/y:0'] +with tf.Graph().as_default() as graph: + tf.import_graph_def(get_graphdef('axpy_ngraph.pbtxt')) +config = tf.ConfigProto( + inter_op_parallelism_threads=1, allow_soft_placement=True) +sess = tf.Session(graph=graph, config=config) + +xval = [[1, 1, 1], [2, 2, 2]] +yval = xval +out = sess.run( + graph.get_tensor_by_name('import/add:0'), + feed_dict={ + graph.get_tensor_by_name('import/x:0'): xval, + graph.get_tensor_by_name('import/y:0'): yval + }) +print(out) diff --git a/tools/sample_optional_params_and_shape_hints.json b/tools/sample_optional_params_and_shape_hints.json new file mode 100644 index 000000000..2e8f649cd --- /dev/null +++ b/tools/sample_optional_params_and_shape_hints.json @@ -0,0 +1,14 @@ +{ + "shape_hints":[ + { + "x":[2,3] + }, + { + "x":[2,-1] + } + ], + "backend_optional_params":{ + "a" : "1", + "b" : "2" + } +} diff --git a/tools/tf2ngraph.py b/tools/tf2ngraph.py index c475d15ac..bb00dc213 100644 --- a/tools/tf2ngraph.py +++ b/tools/tf2ngraph.py @@ -19,39 +19,84 @@ from google.protobuf import text_format from tensorflow.core.protobuf import meta_graph_pb2 from tensorflow.core.protobuf import rewriter_config_pb2 +from tensorflow.core.framework import attr_value_pb2 +from tensorflow.core.framework import tensor_pb2 from tensorflow.python.grappler import tf_optimizer import ngraph_bridge import os import sys +import json from functools import partial -def parse_extra_params_string(raw_extra_params): - raw_extra_params = raw_extra_params.strip(' ') - assert raw_extra_params[0] == '{' and raw_extra_params[ - -1] == '}', "Expected extra_params string to be a dictionary beginning with { and ending with }" - raw_extra_params_contents = raw_extra_params[1:-1].strip(' ') - extra_params_dict = {} - if len(raw_extra_params_contents) == 0: - return extra_params_dict - # could have used eval(extra_params_string), but then the string would have to be the cumbersome {\"abc\":1} - # and not {"abc":1} or {abc:1}. Hence explicity parsing the string without using eval - for key_val in raw_extra_params_contents.split(','): - key_val = key_val.strip(' ') - try: - key, val = key_val.split(':') - extra_params_dict[key.strip(' ')] = val.strip(' ') - except Exception as e: - raise type( - e - )(e.message + - 'Got an entry in extra_params, that is an invalid entry for a python dictionary: ' - + key_val) - return extra_params_dict +class Tf2ngraphJson(object): + + @staticmethod + def allowed_fields(): + return set(["shape_hints", "backend_optional_params"]) + + @staticmethod + def assert_type(obj, expected_type, tag): + assert type(obj) == expected_type, "Expected " + tag + " to be " + str( + expected_type) + " but got " + str(type(obj)) + + @staticmethod + def check_shape_hints(shape_hints): + Tf2ngraphJson.assert_type(shape_hints, type([]), 'shape_hints') + for item in shape_hints: + Tf2ngraphJson.assert_type(item, type({}), + 'each element of the shape_hints list') + for k in item: + Tf2ngraphJson.assert_type( + k, type(""), 'the keys of dictionaries in shape_hints list') + Tf2ngraphJson.assert_type( + item[k], type([]), + 'the values of dictionaries in shape_hints list') + + @staticmethod + def check_optional_params(opt_params): + for optional_attr in opt_params: + Tf2ngraphJson.assert_type(opt_params[optional_attr], type(""), + 'keys of backend_optional_params') + Tf2ngraphJson.assert_type(optional_attr, type(""), + 'values of backend_optional_params') + + @staticmethod + def parse_json(json_name): + if json_name == '': + return {}, [] + optional_backend_params = {} + shape_hints = [] + with open(json_name) as f: + dct = json.load(f) + for k in dct: + if k == 'shape_hints': + Tf2ngraphJson.check_shape_hints(dct[k]) + shape_hints = dct[k] + elif k == 'backend_optional_params': + Tf2ngraphJson.check_optional_params(dct[k]) + optional_backend_params = dct[k] + else: + assert False, "Expected keys to be only in " + str( + allowed_fields()) + return optional_backend_params, shape_hints + + @staticmethod + def dump_json(json_name, optional_params=None, shape_hints=None): + optional_params = {} if optional_params is None else optional_params + shape_hints = [] if shape_hints is None else shape_hints + Tf2ngraphJson.check_optional_params(optional_params) + Tf2ngraphJson.check_shape_hints(shape_hints) + dict_to_dump = {"backend_optional_params": optional_params} + if len(shape_hints) > 0: + dict_to_dump["shape_hints"] = shape_hints + with open(json_name, 'w') as fp: + json.dump(dict_to_dump, fp) def update_config_to_include_custom_config(config, backend, device_id, - extra_params): + backend_optional_params, shape_hints, + do_aot): rewriter_options = rewriter_config_pb2.RewriterConfig() rewriter_options.meta_optimizer_iterations = ( rewriter_config_pb2.RewriterConfig.ONE) @@ -60,8 +105,24 @@ def update_config_to_include_custom_config(config, backend, device_id, ngraph_optimizer.name = "ngraph-optimizer" ngraph_optimizer.parameter_map["ngraph_backend"].s = backend.encode() ngraph_optimizer.parameter_map["device_id"].s = device_id.encode() - for k in extra_params: - ngraph_optimizer.parameter_map[k].s = extra_params[k].encode() + for k in backend_optional_params: + ngraph_optimizer.parameter_map[k].s = backend_optional_params[k].encode( + ) + # Attach shape hints + for hint_id, shape_hint in enumerate(shape_hints): + shape_hint_name = "shape_hint_" + str(hint_id) + ngraph_optimizer.parameter_map[ + shape_hint_name].func.name = shape_hint_name.encode() + ngraph_optimizer.parameter_map[shape_hint_name].func.attr.get_or_create( + 'hint_body').func.name = b'hint_body' + for node_name in shape_hint: # TODO: verify that node names passed in shape hints are valid node names present in the graph + ngraph_optimizer.parameter_map[ + shape_hint_name].func.attr.get_or_create( + 'hint_body').func.attr.get_or_create( + node_name).tensor.int_val.extend(shape_hint[node_name]) + # Attach aot request + ngraph_optimizer.parameter_map["aot_requested"].s = str( + ("0", "1")[do_aot]).encode() config.MergeFrom( tf.ConfigProto( graph_options=tf.GraphOptions(rewrite_options=rewriter_options))) @@ -69,7 +130,8 @@ def update_config_to_include_custom_config(config, backend, device_id, def run_ngraph_grappler_optimizer(input_gdef, output_nodes, ng_backend, - device_id, extra_params): + device_id, backend_optional_params, + shape_hints, do_aot): graph = tf.Graph() with graph.as_default(): tf.import_graph_def(input_gdef, name="") @@ -89,10 +151,11 @@ def run_ngraph_grappler_optimizer(input_gdef, output_nodes, ng_backend, output_collection) session_config = tf.ConfigProto() - # Pass backend and extra backend params to grappler through rewriter config by updating the config + # Pass backend and backend_optional_params to grappler through rewriter config by updating the config # TODO: move update_config_to_include_custom_config to ngraph_bridge session_config = update_config_to_include_custom_config( - session_config, ng_backend, device_id, extra_params) + session_config, ng_backend, device_id, backend_optional_params, + shape_hints, do_aot) output_gdef = tf_optimizer.OptimizeGraph( session_config, grappler_meta_graph_def, graph_id=b"tf_graph") return output_gdef @@ -147,6 +210,7 @@ def prepare_argparser(formats): python tf2ngraph.py --input_pbtxt mobilenet.pbtxt --output_nodes out_node --output_pbtxt mobilenet_ngraph.pbtxt python tf2ngraph.py --input_pb inception_v3_2016_08_28_frozen.pb --output_nodes InceptionV3/Predictions/Reshape_1 --output_pb inception_v3_2016_08_28_frozen_ngraph.pb python tf2ngraph.py --input_pbtxt ../test/test_axpy.pbtxt --output_nodes add --output_pbtxt axpy_ngraph.pbtxt --ng_backend CPU + python tf2ngraph.py --input_pbtxt ../test/test_axpy.pbtxt --output_nodes add --output_pbtxt axpy_ngraph.pbtxt --ng_backend INTERPRETER --config_file sample_optional_params_and_shape_hints.json --precompile ''') in_out_groups = [ parser.add_argument_group(i, j) for i, j in zip( @@ -170,10 +234,16 @@ def prepare_argparser(formats): "--ng_backend", default='CPU', help="Ngraph backend. Eg, NNPI") parser.add_argument("--device_id", default='', help="Device id. Eg, 0") parser.add_argument( - "--extra_params", - default='{}', + "--config_file", + default='', help= - "Other params that the backend needs in the form of a dictionary. Eg, {max_cores: 4}." + "Json file that contains optional backend configuration settings and shape hints" + ) + parser.add_argument( + "--precompile", + action='store_true', + help= + "Perform precompilation to embed the ngraph executable in the dumped TF graph" ) if len(sys.argv) == 1: parser.print_help(sys.stderr) @@ -237,7 +307,7 @@ def attach_device(gdef): def convert(inp_format, inp_loc, out_format, out_loc, output_nodes, ng_backend, - device_id, extra_params): + device_id, backend_optional_params, shape_hints, do_aot): """Functional api for converting TF models by inserting ngraph nodes. Sample usage: from tf2ngraph import convert @@ -259,7 +329,8 @@ def convert(inp_format, inp_loc, out_format, out_loc, output_nodes, ng_backend, input_gdef = get_gdef(inp_format, inp_loc) attach_device(input_gdef) output_gdef = run_ngraph_grappler_optimizer( - input_gdef, output_nodes, ng_backend, device_id, extra_params) + input_gdef, output_nodes, ng_backend, device_id, + backend_optional_params, shape_hints, do_aot) save_model(output_gdef, out_format, out_loc) @@ -273,11 +344,19 @@ def main(): inp_format, inp_loc = filter_dict("input", args.__dict__) out_format, out_loc = filter_dict("output", args.__dict__) output_nodes = args.output_nodes.split(',') - extra_params = parse_extra_params_string(args.extra_params) + backend_optional_params, shape_hints = Tf2ngraphJson.parse_json( + args.config_file) convert(inp_format, inp_loc, out_format, out_loc, output_nodes, - args.ng_backend, args.device_id, extra_params) + args.ng_backend, args.device_id, backend_optional_params, + shape_hints, args.precompile) print('Converted the model. Exiting now') if __name__ == '__main__': main() + + # TODO remove these lines + # python tf2ngraph.py --input_pbtxt ../test/test_axpy.pbtxt --output_nodes add --output_pbtxt axpy_ngraph.pbtxt --ng_backend INTERPRETER --config_file sample_optional_params_and_shape_hints.json --precompile + # python run_tf2ngraph_model.py + + # TODO what happens if same shape is passed twice