Introduces new operator to get the lemma logits for factored vocabulary models for GPU inference

.gitmodules

@@ -20,3 +20,6 @@
 [submodule "src/3rd_party/simple-websocket-server"]
 	path = src/3rd_party/simple-websocket-server
 	url = https://github.com/marian-nmt/Simple-WebSocket-Server
+[submodule "src/3rd_party/cub"]
+	path = src/3rd_party/cub
+	url = https://github.com/NVIDIA/cub

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 ## [Unreleased]
 
 ### Added
+- Add new operator to compute path score for lemmas for a factored vocabulary
 - Local/global sharding with MPI training via `--sharding local`
 - fp16 support for factors.
 - Correct training with fp16 via `--fp16`. 

src/CMakeLists.txt

@@ -1,3 +1,5 @@
+add_definitions(-DCUB_IGNORE_DEPRECATED_CPP_DIALECT=1)
+add_definitions(-DTHRUST_IGNORE_DEPRECATED_CPP_DIALECT=1)
 add_subdirectory(3rd_party)
 
 include_directories(.)

src/data/factored_vocab.h

@@ -3,6 +3,11 @@
 // and via dynamic_cast to FactoredVocab for factored-specific things used by
 // the Embedding and Output layers.
 
+/* Part of this file was contributed by NVIDIA under license:
+ *   Copyright (C) 2020 NVIDIA Corporation
+ *   SPDX-License-Identifier: MIT
+ */
+
 #pragma once
 
 #include "common/definitions.h"
@@ -69,6 +74,7 @@ class FactoredVocab : public IVocab {
   const std::string& getFactorGroupPrefix(size_t groupIndex) const { return groupPrefixes_[groupIndex]; } // for diagnostics only
   const std::string& getFactorName(size_t groupIndex, size_t factorIndex) const { return factorVocab_[(WordIndex)(factorIndex + groupRanges_[groupIndex].first)]; }
   std::string decodeForDiagnostics(const Words& sentence) const;
+  const std::vector<std::vector<bool>>& getLemmaHasFactorGroupVector() const {return lemmaHasFactorGroup_;};
 
   static constexpr size_t FACTOR_NOT_APPLICABLE = (SIZE_MAX - 1);
   static constexpr size_t FACTOR_NOT_SPECIFIED  = (SIZE_MAX - 2);

src/graph/expression_graph.h

@@ -1,3 +1,8 @@
+/* Part of this file was contributed by NVIDIA under license:
+ *   Copyright (C) 2020 NVIDIA Corporation
+ *   SPDX-License-Identifier: MIT
+ */
+
 #pragma once
 
 #include "common/config.h"
@@ -36,6 +41,7 @@ class Tensors {
   typedef std::unordered_map<size_t, std::vector<WExpr>> WeakMemory;
   typedef std::unordered_map<size_t, std::vector<Expr>> Memory;
 
+  std::map<std::string, Expr> memoizationMap_;
   Ptr<WeakMemory> shortterm_;  // holds all nodes for a graph
   Ptr<Memory> longterm_;  // holds memoized nodes
 
@@ -109,6 +115,19 @@ class Tensors {
     return nullptr;
   }
 
+  void rememberByName(const std::string& name, Expr e) {
+    ABORT_IF(e == nullptr, "Expression must be non-null");
+    ABORT_IF(e->type() == "param", "Not intended for graph parameters");
+    memoizationMap_[name] = e;
+    findOrRemember(e);
+  }
+
+  Expr findByName(const std::string&name) {
+    if(memoizationMap_.count(name))
+      return findOrRemember(memoizationMap_[name]);
+    return nullptr;
+  }
+
   void clear() {
     tensors_->clear();
     shortterm_->clear();
@@ -709,6 +728,23 @@ class ExpressionGraph : public std::enable_shared_from_this<ExpressionGraph> {
    */
   Ptr<TensorAllocator> getTensorAllocator() { return tensors_->getTensorAllocator(); }
 
+
+  /**
+    * @TODO Add comment
+    *
+    */
+  void rememberByName(const std::string& name, Expr e) {
+    tensors_->rememberByName(name, e);
+  }
+
+  /**
+    * @TODO Add comment
+    *
+    */
+  Expr findByName(const std::string&name) {
+    return tensors_->findByName(name);
+  }
+
   /** Clear everything apart from parameters and memoized nodes */
   void clear() {
     count_ = 0;

src/graph/expression_operators.cpp

@@ -1,3 +1,7 @@
+/* Part of this file was contributed by NVIDIA under license:
+ *   Copyright (C) 2020 NVIDIA Corporation
+ *   SPDX-License-Identifier: MIT
+ */ 
 #include "graph/expression_operators.h"
 #include "layers/constructors.h"
 
@@ -683,6 +687,22 @@ Expr unlikelihood(Expr logits, Expr indices) {
   return -log(gather(1.f - softmax(logits), /*axis=*/-1, indicesWithLayout));
 }
 
+Expr addFactorMaxes(Expr lemmaHasFactorGroup, std::vector<Expr> groupLosses, Expr hypIndices, size_t group0Start) {
+  if(groupLosses.size() == 1) {
+    return groupLosses[0];
+  }
+
+  int numLemmas = groupLosses[0]->shape()[-1];
+
+  std::vector<Expr> nodes({lemmaHasFactorGroup});
+  if (hypIndices) {
+    nodes.push_back(hypIndices);
+  }
+  nodes.insert(nodes.end(), groupLosses.begin(), groupLosses.end());
+  bool hasShortList = hypIndices != nullptr;
+  return Expression<AddFactorMaxesOp>(nodes, hasShortList, group0Start, numLemmas);
+}
+
 Expr plus(const std::vector<Expr>& nodes) {
   ABORT_IF(nodes.size() > 1, "Not implemented");
   return nodes[0];

src/graph/expression_operators.h

@@ -1,3 +1,8 @@
+/* Part of this file was contributed by NVIDIA under license:
+ *   Copyright (C) 2020 NVIDIA Corporation
+ *   SPDX-License-Identifier: MIT
+ */
+
 #pragma once
 #include "graph/expression_graph.h"
 #include "graph/node_initializers.h"
@@ -610,6 +615,15 @@ Expr atleast_3d(Expr a);
  */
 Expr atleast_4d(Expr a);
 
+/**
+  * @TODO Add Comment
+  *
+  */
+Expr addFactorMaxes(Expr lemmaHasFactorGroup, std::vector<Expr> groupLosses, Expr hypIndices, size_t group0Start);
+
+// create a constant of shape a->shape() and initialize with init
+// @TODO: add a && version, to avoid a ref count. NodeInitializers are typically temps.
+// @TODO: and/or make this a template on init
 /**
  * Converts input to an expression with a least n-dimension dimensions.
  * @param a     Expression

src/graph/node_operators_binary.h

@@ -1,3 +1,7 @@
+/* Part of this file was contributed by NVIDIA under license:
+ *   Copyright (C) 2020 NVIDIA Corporation
+ *   SPDX-License-Identifier: MIT
+ */
 #pragma once
 
 #include <thread>
@@ -1281,6 +1285,73 @@ struct LayerNormalizationOp : public NaryNodeOp {
   float eps_;
 };
 
+struct AddFactorMaxesOp : public NaryNodeOp {
+size_t groupStart_;
+size_t numLemmas_;
+bool hasShortlist_;
+public:
+  AddFactorMaxesOp(const std::vector<Expr>& nodes, bool hasShortlist, size_t groupStart, size_t numLemmas)
+      : NaryNodeOp(nodes, getShape(nodes, hasShortlist), commonType(std::vector<Expr>(nodes.begin() + 1 + (int)hasShortlist, nodes.end())) ) {
+    groupStart_ = groupStart;
+    numLemmas_ = numLemmas;
+    hasShortlist_ = hasShortlist;
+  }
+
+  Shape getShape(const std::vector<Expr>& nodes, bool hasShortlist) {
+    ABORT_IF(nodes.empty(), "No child nodes given");
+    int start = hasShortlist? 2 : 1;
+    return nodes[start]->shape();
+  }
+
+  NodeOps forwardOps() override {
+    int start = hasShortlist_? 2 : 1;
+    std::vector<Tensor> losses;
+    for(int i = start; i < children().size(); ++i) {
+      losses.push_back(child(i)->val());
+    }
+    return {NodeOp(
+                    AddFactorMaxes(val_,
+                                   graph()->allocator(),
+                                   child(0)->val(), // lemmaHasFactorGroupTensor
+                                   hasShortlist_? child(1)->val() : nullptr, // indices
+                                   losses,
+                                   groupStart_, numLemmas_))};
+  }
+
+  NodeOps backwardOps() override {
+    ABORT("Not Implemented for Training");
+  }
+
+  const std::string type() override { return "AddFactorMaxesOp"; }
+
+  virtual size_t hash() override {
+    size_t seed = NaryNodeOp::hash();
+    util::hash_combine(seed, hasShortlist_);
+    util::hash_combine(seed, groupStart_);
+    util::hash_combine(seed, numLemmas_);
+    return seed;
+  }
+
+  virtual bool equal(Expr node) override {
+    if(!NaryNodeOp::equal(node))
+      return false;
+    auto cnode = std::dynamic_pointer_cast<AddFactorMaxesOp>(node);
+    if(!cnode)
+      return false;
+    if(hasShortlist_ != cnode->hasShortlist_)
+      return false;
+    if(groupStart_ != cnode->groupStart_)
+      return false;
+    if(numLemmas_ != cnode->numLemmas_)
+      return false;
+    return true;
+  }
+
+private:
+  friend class SerializationHelpers; // @TODO: use the same name for this as SqrtNodeOp
+};
+
+
 struct HighwayNodeOp : public NaryNodeOp {
   HighwayNodeOp(const std::vector<Expr>& nodes) : NaryNodeOp(nodes) {}
 

src/layers/generic.cpp

@@ -1,3 +1,9 @@
+/* Part of this file was contributed by NVIDIA under license:
+ *   Copyright (C) 2020 NVIDIA Corporation
+ *   SPDX-License-Identifier: MIT
+ */
+
+#include <algorithm>
 #include "marian.h"
 
 #include "layers/generic.h"
@@ -80,15 +86,25 @@ namespace marian {
     //  - lemma: add all maxes of applicable factors
     if (groupIndex > 0) {
       sel = sel - max(sel, -1);
-    }
-    else {
+    } else {
       auto numGroups = getNumFactorGroups();
-      for (size_t g = 1; g < numGroups; g++) {
-        auto factorMaxima = max(logits_[g]->loss(), -1); // we cast since loss is likely ce-loss which has type float32
-        auto factorMasks = constant(getFactorMasks(g, shortlist ? shortlist->indices() : std::vector<WordIndex>()));
-        sel = sel + cast(factorMaxima, sel->value_type()) * cast(factorMasks, sel->value_type()); // those lemmas that don't have a factor get multiplied with 0
+      if(numGroups > 1 && graph()->isInference() && graph()->getBackend()->getDeviceId().type == DeviceType::gpu) {
+        Expr shortlistIndices = shortlist? indices(shortlist->indices()) : nullptr;
+        Expr lemmaHasFactorGroupTensor = getLemmaHasFactorGroupTensor();
+        std::vector<Expr> groupLosses(logits_.size());
+        std::transform(logits_.begin(), logits_.end(), groupLosses.begin(), [](const Ptr<RationalLoss>& loss) -> Expr {return loss->loss();});
+
+        sel = addFactorMaxes(lemmaHasFactorGroupTensor, groupLosses, 
+                             shortlistIndices, factoredVocab_->getGroupRange(0).first);
+      } else {
+        for (size_t g = 1; g < numGroups; g++) {
+          auto factorMaxima = max(logits_[g]->loss(), -1);
+          Expr factorMasks = constant(getFactorMasks(g, shortlist ? shortlist->indices() : std::vector<WordIndex>()));
+          sel = sel + factorMaxima * factorMasks; // those lemmas that don't have a factor get multiplied with 0
+        }
       }
     }
+
 
     // if selIdx are given, then we must reshuffle accordingly
     if (!hypIndices.empty()) // use the same function that shuffles decoder state
@@ -190,6 +206,43 @@ namespace marian {
     return res;
   }
 
+  Expr Logits::getLemmaHasFactorGroupTensor() const {
+    auto g = graph();
+    const std::string name = "lemmaHasFactorGroup";
+    auto lemmaHasFactorGroupTensor = g->findByName(name);
+
+    if(!lemmaHasFactorGroupTensor) {
+      // We want to make this a graph param so the GPU can use this to implement lemmaHasFactorGroup to avoid unneeded memcpys.
+      const auto lemmaHasFactorGroup = factoredVocab_->getLemmaHasFactorGroupVector();
+      int dimLemma = (int)lemmaHasFactorGroup.size();
+      int dimFactorGroup = (int)lemmaHasFactorGroup[0].size();
+
+      for(const auto&  lemma : lemmaHasFactorGroup) {
+        // Paranoid check - Instead of aborting I think we can pad here and copy the array
+        ABORT_IF(lemma.size() != dimFactorGroup, "All groups must be the same size");
+      }
+
+      auto initFunc = [lemmaHasFactorGroup, dimLemma, dimFactorGroup] (Tensor t) {
+        std::vector<int8_t> flattened(dimLemma * dimFactorGroup);
+        int row = 0;
+        for(const auto& v : lemmaHasFactorGroup) {
+          int offset = row * dimFactorGroup;
+          for(int i = 0; i < v.size(); ++i) {
+            flattened[offset + i] = static_cast<int8_t>(v[i]);
+          }
+          ++row;
+        }
+        t->set(flattened);
+      };
+
+      lemmaHasFactorGroupTensor = graph()->constant({dimLemma, dimFactorGroup}, inits::fromLambda(initFunc), Type::int8);
+      lemmaHasFactorGroupTensor->setMemoize(true);
+      g->rememberByName(name, lemmaHasFactorGroupTensor);
+    }
+    ABORT_IF(!lemmaHasFactorGroupTensor, "Lemma has factor group tensor undefined?");
+    return lemmaHasFactorGroupTensor;
+  }
+
   Logits Logits::applyUnaryFunction(const std::function<Expr(Expr)>& f) const { // clone this but apply f to all loss values
     std::vector<Ptr<RationalLoss>> newLogits;
     for (const auto& l : logits_)

src/layers/generic.h

@@ -1,3 +1,8 @@
+/* Part of this file was contributed by NVIDIA under license:
+ *   Copyright (C) 2020 NVIDIA Corporation
+ *   SPDX-License-Identifier: MIT
+ */
+
 #pragma once
 
 #include "marian.h"
@@ -142,6 +147,8 @@ class Logits {
     template<typename T> Expr constant(const std::vector<T>& data) const { return constant(Shape{(int)data.size()}, data); } // same as constant() but assuming vector
     Expr indices(const std::vector<uint32_t>& data) const { return graph()->indices(data); } // actually the same as constant(data) for this data type
     std::vector<float> getFactorMasks(size_t factorGroup, const std::vector<WordIndex>& indices) const;
+    Expr getLemmaHasFactorGroupTensor() const;
+    Expr addFactorMaxesHelper(Expr indices=nullptr) const;
 private:
     // members
     // @TODO: we don't use the RationalLoss component anymore, can be removed again, and replaced just by the Expr

src/tensors/cpu/tensor_operators.cpp

@@ -3,6 +3,11 @@
  *   SPDX-License-Identifier: MIT
  */
 
+/* Part of this file was contributed by NVIDIA under license:
+ *   Copyright (C) 2020 NVIDIA Corporation
+ *   SPDX-License-Identifier: MIT
+ */
+
 #include "tensors/tensor_operators.h"
 #include "tensors/cpu/backend.h"
 #include "tensors/allocator.h"
@@ -24,6 +29,16 @@ void IsNaN(const Tensor /*in*/, Ptr<Allocator> /*allocator*/, bool& /*isNaN*/, b
   ABORT("Not implemented");
 }
 
+void AddFactorMaxes(Tensor /*out*/,
+                    Ptr<Allocator> /*allocator*/,
+                    const Tensor /*lemmaHasFactorGroupTensor*/,
+                    const Tensor /*indices*/,
+                    const std::vector<marian::Tensor>& /*groupLosses*/,
+                    size_t /*groupStart*/,
+                    size_t /*numLemmas*/) {
+  ABORT("AddFactorMaxes not implemented on CPU");
+}
+
 template <bool add, typename To, typename From>
 void CopyCastTo(To* out, const From* in, int length) {
   for(int i = 0; i < length; ++i)