Sarkars/nmsv4 (#489)

Author: sachinmuradi

src/ngraph_builder.cc

@@ -2404,6 +2404,80 @@ static Status TranslateMaxPoolGradOp(
   return Status::OK();
 }
 
+static Status TranslateNonMaxSuppressionV4Op(
+    const Node* op, const std::vector<const Tensor*>& static_input_map,
+    Builder::OpMap& ng_op_map) {
+  shared_ptr<ng::Node> ng_boxes, ng_scores;
+  TF_RETURN_IF_ERROR(GetInputNodes(ng_op_map, op, &ng_boxes, &ng_scores,
+                                   nullptr, nullptr, nullptr));
+
+  std::vector<int> max_output_size;
+  TF_RETURN_IF_ERROR(
+      GetStaticInputVector(op, 2, static_input_map, &max_output_size));
+  std::vector<float> iou_threshold;
+  TF_RETURN_IF_ERROR(
+      GetStaticInputVector(op, 3, static_input_map, &iou_threshold));
+
+  std::vector<float> score_threshold;
+  TF_RETURN_IF_ERROR(
+      GetStaticInputVector(op, 4, static_input_map, &score_threshold));
+
+  bool pad_to_max_output_size;
+  if (GetNodeAttr(op->attrs(), "pad_to_max_output_size",
+                  &pad_to_max_output_size) != Status::OK()) {
+    pad_to_max_output_size = false;
+  }
+  // max_output_size must be scalar
+  if (max_output_size.size() != 1) {
+    return errors::InvalidArgument(
+        "NonMaxSuppressionV4 Op: max_output_size of nms must be scalar ",
+        max_output_size.size());
+  }
+  // iou_threshold must be scalar
+  if (iou_threshold.size() != 1) {
+    return errors::InvalidArgument(
+        "NonMaxSuppressionV4 Op: iou_threshold of nms must be scalar ",
+        iou_threshold.size());
+  }
+
+  // score_threshold must be scalar
+  if (score_threshold.size() != 1) {
+    return errors::InvalidArgument(
+        "NonMaxSuppressionV4 Op: score_threshold of nms must be scalar ",
+        score_threshold.size());
+  }
+
+  std::string backend_name;
+  TF_RETURN_IF_ERROR(ngraph_bridge::GetNodeBackend(op, &backend_name));
+
+  if (backend_name != "NNPI") {
+    return errors::Internal("In translating NonMaxSuppressionV4 op ",
+                            op->name(), " found requested backend ",
+                            backend_name, " which is unsupported");
+  }
+
+  ng::runtime::Backend* backend = BackendManager::GetBackend(backend_name);
+
+  shared_ptr<ng::Node> ng_nmsv4 = backend->get_backend_op(
+      "NonMaxSuppressionV4", &ng_boxes, &ng_scores,
+      (size_t)(max_output_size[0]), (float)(iou_threshold[0]),
+      (float)score_threshold[0], (bool)pad_to_max_output_size);
+  if (ng_nmsv4 == nullptr) {
+    return errors::Internal("In translating NonMaxSuppressionV4 op ",
+                            op->name(),
+                            " backend could not return valid ngraph node");
+  }
+  shared_ptr<ngraph::Node> ng_selected_indices =
+      ConstructNgNode<ngraph::op::GetOutputElement>(op->name(), ng_nmsv4, 0);
+  shared_ptr<ngraph::Node> ng_valid_output =
+      ConstructNgNode<ngraph::op::GetOutputElement>(op->name(), ng_nmsv4, 1);
+
+  SaveNgOp(ng_op_map, op->name(), ng_selected_indices);
+  SaveNgOp(ng_op_map, op->name(), ng_valid_output);
+
+  return Status::OK();
+}
+
 static Status TranslateReduceOp(
     const Node* op, const std::vector<const Tensor*>& static_input_map,
     Builder::OpMap& ng_op_map,
@@ -4357,6 +4431,7 @@ const static std::map<
         {"MaxPool", TranslateMaxPoolOp},
         {"MaxPool3D", TranslateMaxPool3DOp},
         {"MaxPoolGrad", TranslateMaxPoolGradOp},
+        {"NonMaxSuppressionV4", TranslateNonMaxSuppressionV4Op},
         {"Mean", TranslateMeanOp},
         {"Min", TranslateDirectReduceOp<ng::op::Min>},
         {"Minimum", TranslateBinaryOp<ngraph::op::Minimum>},

src/ngraph_mark_for_clustering.cc

@@ -460,6 +460,8 @@ Status MarkForClustering(Graph* graph,
       type_constraint_map["Minimum"]["T"] = NGraphNumericDTypes();
       type_constraint_map["Mul"]["T"] = NGraphNumericDTypes();
       type_constraint_map["Neg"]["T"] = NGraphNumericDTypes();
+      type_constraint_map["NonMaxSuppressionV4"]["T"] = {
+          DT_FLOAT};  // TF allows half too
       type_constraint_map["OneHot"]["T"] = NGraphDTypes();
       type_constraint_map["Pack"]["T"] = NGraphDTypes();
       type_constraint_map["Pad"]["T"] = NGraphDTypes();
@@ -574,6 +576,7 @@ Status MarkForClustering(Graph* graph,
       set_attributes_map["Max"] = SetStaticInputs({1});
       set_attributes_map["Mean"] = SetStaticInputs({1});
       set_attributes_map["Min"] = SetStaticInputs({1});
+      set_attributes_map["NonMaxSuppressionV4"] = SetStaticInputs({2, 3, 4});
       set_attributes_map["OneHot"] = SetStaticInputs({1});
       set_attributes_map["Pad"] = SetStaticInputs({1});
       set_attributes_map["Prod"] = SetStaticInputs({1});
@@ -626,7 +629,6 @@ Status MarkForClustering(Graph* graph,
                               " is not supported");
     }
     current_backend = backend_env;
-    // TODO: set backend. Then don't use current_backend
   }
 
   // Right now it cannot be inside the if(!initialized) block, because it is
@@ -639,6 +641,12 @@ Status MarkForClustering(Graph* graph,
     return Status::OK();
   };
 
+  confirmation_function_map["NonMaxSuppressionV4"] = [&current_backend](
+      Node* n, bool* result) {
+    *result = (current_backend == "NNPI");
+    return Status::OK();
+  };
+
   std::unordered_map<string, int> no_support_histogram;
   std::unordered_map<string, int> fail_confirmation_histogram;
   std::unordered_map<string, int> fail_constraint_histogram;