Generalize scatter/gather ops (#413)

doc/guide/myelin.md

@@ -21,10 +21,10 @@ Build system: Bazel<br>
 
 ## Using Myelin in Python
 
-Myelin represents a computation graph using a _flow_. The graph is divivded into
+Myelin represents a computation graph using a _flow_. The graph is divided into
 _functions_ which can be computed independently. A function is a set of
 _operations_ with tensor inputs and outputs. The tensor inputs and outputs are
-_variables_ in the flow. Variables can either be global constant tensor, e.g.
+_variables_ in the flow. Variables can either be global constant tensors, e.g.
 learned weights in a neural network, or parameter tensors, which are local to
 the function.
 

python/myelin/builder.py

@@ -35,7 +35,7 @@
   "f": DT_FLOAT32,
   "d": DT_FLOAT64,
   "i": DT_INT32,
-  "l": DT_INT32,
+  "l": DT_INT64,
   "B": DT_INT8,
   "h": DT_INT16,
   "b": DT_INT8,
@@ -592,8 +592,8 @@ def assign(self, x, y, name=None):
     op.add_input(x)
     op.add_input(y)
 
-  def scatter_add(self, m, f, v, ref=False, name=None):
-    op = self.rawop("ScatterAdd", name)
+  def assign_add_scatter(self, m, f, v, ref=False, name=None):
+    op = self.rawop("AssignAddScatter", name)
     op.add_input(m)
     op.add_input(f)
     op.add_input(v)

python/myelin/simulator.py

@@ -229,7 +229,13 @@ def compute(flow, f, data):
       for k in range(len(splits)): v[o[k]] = splits[k]
     elif op.type == "Gather":
       v[o[0]] = gather(v[i[0]], v[i[1]])
-    elif op.type == "ScatterAdd":
+    elif op.type == "GatherSum":
+      v[o[0]] = np.sum(gather(v[i[0]], v[i[1]]), axis=1)
+    elif op.type == "GatherMax":
+      v[o[0]] = np.max(gather(v[i[0]], v[i[1]]), axis=1)
+    elif op.type == "GatherAvg":
+      v[o[0]] = np.sum(gather(v[i[0]], v[i[1]]), axis=1) / v[i[1]].shape[1]
+    elif op.type == "AssignAddScatter":
       m = v[i[0]]
       f = v[i[1]]
       x = v[i[2]]

sling/myelin/builder.h

@@ -336,6 +336,10 @@ class FlowBuilder : public Scope {
     int n = f->rank() == 0 ? 1 : f->dim(1);
     return Op("Gather", {M, f}, M->type, {n, M->dim(1)});
   }
+  Variable *Gather(Variable *M, Variable *f, Variable *oov) {
+    int n = f->rank() == 0 ? 1 : f->dim(1);
+    return Op("Gather", {M, f, oov}, M->type, {n, M->dim(1)});
+  }
   Variable *GatherSum(Variable *M, Variable *f) {
     return Op("GatherSum", {M, f}, M->type, {1, M->dim(1)});
   }
@@ -350,6 +354,9 @@ class FlowBuilder : public Scope {
   Variable *Scatter(Variable *f, Variable *v, int size) {
     return Op("Scatter", {f, v}, v->type, {size, v->dim(1)});
   }
+  Variable *Scatter(Variable *f, Variable *v, int size, Variable *oov) {
+    return Op("Scatter", {f, v, oov}, v->type, {size, v->dim(1)});
+  }
 
   // Assignment.
   Operation *Assign(Variable *var, Variable *value) {
@@ -364,8 +371,8 @@ class FlowBuilder : public Scope {
     return Op("Assign", {var, value})->set_ref();
   }
 
-  Operation *ScatterAdd(Variable *M, Variable *f, Variable *v) {
-    return RawOp("ScatterAdd", {M, f, v});
+  Operation *AssignAddScatter(Variable *M, Variable *f, Variable *v) {
+    return RawOp("AssignAddScatter", {M, f, v});
   }
 
   // Concatenation.

sling/myelin/flow.cc

@@ -1790,6 +1790,7 @@ Flow::Operation *Flow::AddOperation(Function *func,
                                     const string &type) {
   Operation *op = AddOperation(name, type);
   func->AddOperation(op);
+  if (op->name.empty()) op->name = OpName(func->name + "/" + type);
   return op;
 }
 
@@ -1798,8 +1799,7 @@ Flow::Operation *Flow::AddOperation(Function *func,
                                     const string &type,
                                     const std::vector<Variable *> &inputs,
                                     const std::vector<Variable *> &outputs) {
-  Operation *op = AddOperation(name, type);
-  func->AddOperation(op);
+  Operation *op = AddOperation(func, name, type);
   for (auto *input : inputs) op->AddInput(input);
   for (auto *output : outputs) op->AddOutput(output);
   return op;
@@ -2148,15 +2148,21 @@ Flow::Blob *Flow::DataBlock(const string &name) {
 string Flow::VarName(const string &prefix) {
   for (int n = 0;; ++n) {
     string name = prefix;
-    if (n > 0) name.append(std::to_string(n));
+    if (n > 0) {
+      name.push_back('_');
+      name.append(std::to_string(n));
+    }
     if (Var(name) == nullptr) return name;
   }
 }
 
 string Flow::OpName(const string &prefix) {
   for (int n = 0;; ++n) {
     string name = prefix;
-    if (n > 0) name.append(std::to_string(n));
+    if (n > 0) {
+      name.push_back('_');
+      name.append(std::to_string(n));
+    }
     if (Op(name) == nullptr) return name;
   }
 }