Add Softmax and Add. Fix infer and don't duplicate Variables for diverging paths.

Author: lantiga

torch2c/__init__.py

@@ -5,30 +5,44 @@
 
 
 def _wrap(obj, prevfns=[]):
+    if obj.__class__ not in emitters._class_map:
+        raise Exception('%s does not have an Emitter' % obj.__class__)
     return emitters._class_map[obj.__class__](obj,prevfns)
 
 
-def _traverse_graph_recursive(out, el):
+def _traverse_graph_recursive(out, id_set, el):
     if isinstance(el, Variable):
         var = _wrap(el)
     else:
         prevfns = []
         if hasattr(el,'previous_functions'):
             prevfns = [f[0] for f in el.previous_functions]
         var = _wrap(el,prevfns)
-    out.append(var)
+    var_name = var.id_var_name()
+    if var_name not in id_set:
+        out.append(var)
+        id_set.add(var_name)
     if hasattr(el, 'previous_functions'):
         for u in el.previous_functions:
-            _traverse_graph_recursive(out,u[0])
+            _traverse_graph_recursive(out,id_set,u[0])
 
 
 def _traverse_graph(node):
     nodes = []
-    _traverse_graph_recursive(nodes,node.creator)
+    id_set = set()
+    _traverse_graph_recursive(nodes,id_set,node.creator)
     nodes.reverse()
     var_dict = dict([(el.id,el) for el in nodes])
-    for el in nodes:
-        el.infer_type(var_dict)
+    prev_none_count = 0
+    while True:
+        for el in nodes:
+            el.infer_type(var_dict)
+        none_count = len([el for el in nodes if el.numtype == None])
+        if none_count == 0:
+            break
+        if none_count == prev_none_count:
+            raise Exception('Cannot infer types for all nodes in the graphs')
+        prev_none_count = none_count
     return nodes
 
 
@@ -52,12 +66,10 @@ def _emit_c(nodes, out, fnname, out_path):
     last_node = nodes[-1]
     ifndef = '#ifndef __%s__\n#define __%s__\n' % (2*(fnname.upper(),))
     endif = '#endif'
-    includes = '#include "TH.h"\n#include "THNN.h"\n#include "torch2c.h"'
+    includes = '#include "TH.h"\n#include "THNN.h"\n#include "torch2c.h"\n'
     fndecl = 'void %s(%s)' % (fnname, 
               ', '.join([el.emit_decl() for el in var_nodes + [out_node]]))
-    print(fndecl)
     calls = [el.emit_call(out_path,'data') for el in nodes]
-    print('\n'.join(calls))
     copy_out = out_node.emit_copy(last_node.id_var_name())
     # TODO: be smarter re: frees
     # analyze calls backwards and free right after last use

torch2c/emitters.py

@@ -54,6 +54,9 @@ def id_var_name(self):
     def var_names(self):
         return {k: self.var_name(v) for k,v in self.vars.items()}
 
+    def persisted_vars(self):
+        return []
+
     def call_tpl(self):
         return ''
 
@@ -241,6 +244,51 @@ def free_tpl(self):
 register(Linear, torch.nn._functions.linear.Linear)
 
 
+class Add(Emitter):
+
+    def __init__(self, obj, prevfns):
+        Emitter.__init__(self, obj, prevfns)
+        self.def_vars({'input0': id(prevfns[0]),
+                       'input1': id(prevfns[1])})
+        self.infer_type_var = 'input0'
+
+    def call_tpl(self):
+        return '''
+            TH${T}Tensor *$id = TH${T}Tensor_new();
+            TH${T}Tensor_cadd($id,$input0,1.0,$input1);
+            '''
+
+    def free_tpl(self):
+        return '''
+            TH${T}Tensor_free($id);
+            '''
+
+register(Add, torch.autograd._functions.basic_ops.Add)
+
+
+
+
+class Softmax(Emitter):
+
+    def __init__(self, obj, prevfns):
+        Emitter.__init__(self, obj, prevfns)
+        self.def_vars({'input': id(prevfns[0])})
+        self.infer_type_var = 'input'
+
+    def call_tpl(self):
+        return '''
+            TH${T}Tensor *$id = TH${T}Tensor_new();
+            THNN_${T}SoftMax_updateOutput(NULL,$input,$id);
+            '''
+
+    def free_tpl(self):
+        return '''
+            TH${T}Tensor_free($id);
+            '''
+
+register(Softmax, torch.nn._functions.thnn.auto.Softmax)
+
+
 class LogSoftmax(Emitter):
 
     def __init__(self, obj, prevfns):