Preserve output PH names in C2 loader (#3798)

Summary:
It's important to preserve the name of each external output in outside models, but since they must be unique with Node names we'll need to claim them before naming any nodes. This fixes the C2 Model Loader to do that.

Documentation: Fixes #3774
Pull Request resolved: #3798

Test Plan: ninja test, image-classifier with -dump-ir to make sure weights were named right.

Differential Revision: D18630019

Pulled By: nickgg

fbshipit-source-id: 0b0e771c0e6489384f7e40a772178bf45ff64b0e

Author: nickgg

examples/bundles/lenet_mnist/main.cpp

@@ -222,7 +222,7 @@ uint8_t activations[LENET_MNIST_ACTIVATIONS_MEM_SIZE];
 uint8_t *inputAddr = GLOW_GET_ADDR(mutableWeight, LENET_MNIST_data);
 
 /// Bundle output data absolute address.
-uint8_t *outputAddr = GLOW_GET_ADDR(mutableWeight, LENET_MNIST_softmax__1);
+uint8_t *outputAddr = GLOW_GET_ADDR(mutableWeight, LENET_MNIST_softmax);
 
 /// Copy the pre-processed images into the mutable region of the bundle.
 static void initInputImages() {

examples/bundles/resnet50/main.cpp

@@ -300,7 +300,7 @@ static uint8_t *allocateMutableWeightVars(const BundleConfig &config) {
 static void dumpInferenceResults(const BundleConfig &config,
                                  uint8_t *mutableWeightVars) {
   const SymbolTableEntry &outputWeights =
-      getMutableWeightVar(config, "gpu_0_softmax__1");
+      getMutableWeightVar(config, "gpu_0_softmax");
   int maxIdx = 0;
   float maxValue = 0;
   float *results = (float *)(mutableWeightVars + outputWeights.offset);

examples/resnet-runtime.cpp

@@ -101,7 +101,7 @@ void dispatchClassify(unsigned int id, HostManager *hostManager,
         EXIT_ON_ERR(std::move(err));
         auto *bindings = context->getPlaceholderBindings();
         size_t maxIdx =
-            bindings->get(bindings->getPlaceholderByName("gpu_0_softmax__1"))
+            bindings->get(bindings->getPlaceholderByName("gpu_0_softmax"))
                 ->getHandle()
                 .minMaxArg()
                 .second;

include/glow/Graph/Graph.h

@@ -109,6 +109,12 @@ class Module final {
     usedNodeNames_.insert(name);
   }
 
+  /// Registers a name as used by a Storage node (Constant or Placeholder) in
+  /// this module.
+  void registerStorageName(llvm::StringRef name) {
+    usedStorageNames_.insert(name);
+  }
+
   /// Return a pointer to a uniqued type \p T.
   TypeRef uniqueType(const Type &T);
 

lib/Importer/Caffe2ModelLoader.cpp

@@ -1444,6 +1444,12 @@ Error Caffe2ModelLoader::loadInputs(const caffe2::NetDef &net,
 }
 
 Error Caffe2ModelLoader::loadNetwork(caffe2::NetDef &net) {
+  // Make a claim on the unique name of all output Placeholders.
+  for (int i = 0; i < net.external_output_size(); i++) {
+    auto &outputName = net.external_output(i);
+    G_.getParent()->registerStorageName(legalizeName(outputName));
+  }
+
   /// Load the network operators:
   for (int i = 0; i < net.op_size(); i++) {
     auto &op = net.op(i);
@@ -1464,8 +1470,13 @@ Error Caffe2ModelLoader::loadNetwork(caffe2::NetDef &net) {
     auto &outputName = net.external_output(i);
     NodeValue r;
     ASSIGN_VALUE_OR_RETURN_ERR(r, getNodeValueByName(outputName));
-    auto *PH =
-        G_.getParent()->createPlaceholder(r.getType(), outputName, false);
+
+    PlaceholderList &PHList = G_.getParent()->getPlaceholders();
+    // Create a Placeholder with the previously claimed name.
+    auto *PH = new Placeholder(legalizeName(outputName),
+                               G_.getParent()->uniqueType(*r.getType()), false,
+                               ANY_LAYOUT);
+    PHList.push_back(PH);
     auto *SN = G_.createSave("save_" + outputName, r, PH);
     outputVarsByName_[outputName] = SN->getPlaceholder();
   }

tests/unittests/Caffe2ImporterTest.cpp

@@ -2551,7 +2551,8 @@ TEST(caffe2, importNames) {
   Tensor input(ElemKind::FloatTy, {6});
   Caffe2ModelLoader caffe2LD(NetDescFilename, NetWeightFilename,
                              {"sigmoid_test_input"}, {&input.getType()}, *F);
-  EXPECT_TRUE(F->getNodeByName("sigmoid_test_output"));
+  EXPECT_TRUE(mod.getPlaceholderByName("sigmoid_test_output"));
+  EXPECT_TRUE(F->getNodeByName("sigmoid_test_output__1"));
 }
 
 TEST(caffe2, importSqr) {