Various fixes (pytorch#3895)

Summary:
Pull Request resolved: pytorch#3895

Some fixes for models with EmbeddingBagByteRowwiseOffsets

Reviewed By: qizzzh

Differential Revision: D18940211

fbshipit-source-id: 0a1b9771602c6465e04d246c1a3b57a74366694e

torch_glow/src/GlowFuser.cpp

@@ -27,6 +27,7 @@
 #include <torch/csrc/jit/passes/alias_analysis.h>
 #include <torch/csrc/jit/passes/common_subexpression_elimination.h>
 #include <torch/csrc/jit/passes/dead_code_elimination.h>
+#include <torch/csrc/jit/passes/inliner.h>
 #include <torch/csrc/jit/passes/subgraph_rewrite.h>
 #include <torch/csrc/jit/passes/utils/subgraph_utils.h>
 
@@ -269,6 +270,8 @@ void glowCustomFuseImpl(std::shared_ptr<torch::jit::Graph> graph,
     return fn(ptNode);
   };
 
+  Inline(*graph);
+
   // Prepare the graph by fusing known patterns for the model loader.
   // TODO: this should be done only on Glow subgraphs to avoid modifying parts
   // of the graph that Glow will not be running.

torch_glow/src/PyTorchModelLoader.cpp

@@ -742,6 +742,7 @@ PyTorchModelLoader::getSymbolsMapping() {
       {{"aten::gelu"}, &PyTorchModelLoader::loadGelu, {}},
       {{"aten::tanh", "aten::tanh_"}, &PyTorchModelLoader::loadTanh, {}},
       {{"aten::t", "aten::t_"}, &PyTorchModelLoader::loadT, {}},
+      {{"aten::to"}, &PyTorchModelLoader::loadTo, {}},
       {{"aten::permute"}, &PyTorchModelLoader::loadPermute, {}},
       {{"aten::transpose", "aten::transpose_"},
        &PyTorchModelLoader::loadTranspose,
@@ -3060,6 +3061,18 @@ Error PyTorchModelLoader::loadPermute(const torch::jit::Node *ptNode) {
   return addValueMapping(outputs[0], output);
 }
 
+Error PyTorchModelLoader::loadTo(const torch::jit::Node *ptNode) {
+  auto inputs = ptNode->inputs();
+  auto outputs = ptNode->outputs();
+  RETURN_IF_ERR(checkInputAndOutputSizes(inputs, 5, outputs, 1));
+
+  // TODO: use ConvertTo
+  glow::NodeValue in;
+  ASSIGN_VALUE_OR_RETURN_ERR(in, getGlowNodeValueForValue(inputs[0]));
+
+  return addValueMapping(outputs[0], in);
+}
+
 Error PyTorchModelLoader::loadFlatten(const torch::jit::Node *ptNode) {
   auto inputs = ptNode->inputs();
   auto outputs = ptNode->outputs();
@@ -3284,6 +3297,19 @@ Error PyTorchModelLoader::loadEmbeddingBagByteRowwiseOffsets(
       "EmbeddingBagByteRowwiseOffsets.ones",
       glow::Type(weight.getElementType(), {indices.dims()[0]}), 1.0);
 
+  glow::Constant *weightConstant =
+      llvm::dyn_cast<glow::Constant>(weight.getNode());
+
+  RETURN_ERR_IF_NOT(weightConstant,
+                    strFormat("Expected Weight to be a Constant but found: %s",
+                              weight.getNode()->getKindName()));
+
+  TypeRef fusedTy = F_.getParent()->uniqueType(ElemKind::UInt8FusedQTy,
+                                               weight.dims(), 0.0, 0);
+
+  weightConstant->setType(Storage::OutputIdx, fusedTy);
+  weightConstant->setPayloadType(fusedTy);
+
   bool scaleGradByFreq;
   ASSIGN_VALUE_OR_RETURN_ERR(
       scaleGradByFreq,
@@ -3308,8 +3334,8 @@ Error PyTorchModelLoader::loadEmbeddingBagByteRowwiseOffsets(
   RETURN_ERR_IF_NOT(sparse == false, "Currently only support sparse='false'");
 
   auto *EB = F_.createEmbeddingBagByteRowwiseOffsets(
-      "EmbeddingBagByteRowwiseOffsets", weight, perSampleWeights, indices,
-      offsets);
+      "EmbeddingBagByteRowwiseOffsets", weightConstant->getOutput(),
+      perSampleWeights, indices, offsets);
 
   return addValueMapping(outputs[0], EB->getResult());
 }

torch_glow/src/PyTorchModelLoader.h

@@ -536,6 +536,10 @@ class PyTorchModelLoader {
   /// Load a PyTorch aten::permute node.
   /// \returns error on failure.
   Error loadPermute(const torch::jit::Node *ptNode);
+
+  /// Load a PyTorch aten::to node.
+  /// \returns error on failure.
+  Error loadTo(const torch::jit::Node *ptNode);
 };
 
 } // namespace glow