Add generic THStorage_newFromFile

Author: lantiga

include/torch2c.h

@@ -0,0 +1,9 @@
+#ifndef TORCH2C_H
+#define TORCH2C_H
+
+#include <TH.h>
+#include "torch2c_generic.h"
+#include <THGenerateAllTypes.h>
+
+#endif
+

include/torch2c_generic.h

@@ -0,0 +1,30 @@
+#ifndef TH_GENERIC_FILE
+#define TH_GENERIC_FILE "torch2c_generic.h"
+#else
+
+// TODO: this only works with little endian for now, we need to make more general
+
+TH_API THStorage *THStorage_(newFromFile)(const char *filename)
+{
+  FILE *f = fopen(filename,"rb");
+
+  if (!f) {
+    THError("cannot open file %s for reading");
+    return NULL;
+  }
+
+  long size;
+  size_t result = fread(&size,sizeof(long),1,f);
+
+  THStorage *out = THStorage_(newWithSize)(size);
+  char *bytes = (char *) out->data;
+
+  uint64_t remaining = sizeof(real) * out->size;
+  result = fread(bytes,sizeof(real),out->size,f);
+
+  fclose(f);
+
+  return out;
+}
+
+#endif

scripts/CMakeLists.txt

@@ -4,6 +4,7 @@ project(torch2ctest)
 
 include_directories(${INSTALL_DIR}/include/TH)
 include_directories(${INSTALL_DIR}/include/THNN)
+include_directories(${INSTALL_DIR}/include/torch2c)
 
 link_directories(${INSTALL_DIR}/lib)
 

scripts/run_test.sh

@@ -11,6 +11,9 @@ INSTALL_DIR=$BASE_DIR/tmp/install
 OUT_BASE_DIR=$BASE_DIR/out
 OUT_BUILD_DIR=$BASE_DIR/tmp/out-build
 
+mkdir -p $INSTALL_DIR/include/torch2c
+cp $BASE_DIR/include/*.h $INSTALL_DIR/include/torch2c
+
 rm -rf $OUT_BASE_DIR
 rm -rf $OUT_BUILD_DIR
 python3 test/$1.py

torch2c/__init__.py

@@ -52,7 +52,7 @@ def _generate_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'
+    includes = '#include "TH.h"\n#include "THNN.h"\n#include "torch2c.h"'
     fndecl = 'void %s(%s)' % (fnname, 
               ', '.join([el.generate_decl() for el in var_nodes + [out_node]]))
     calls = [el.generate_call(out_path,'data') for el in nodes]
@@ -103,6 +103,7 @@ def _generate_test(nodes, out, fnname, filename, out_path):
 
 
 def compile(node, fnname, out_path, compile_test=False):
+    includedir = os.path.join(os.path.dirname(__file__),'..','include')
     nodes = _traverse_graph(node)
     if not os.path.isdir(out_path):
         os.mkdir(out_path)

torch2c/emitters.py

@@ -180,35 +180,6 @@ def persist_tensor(tensor, name, out_path, datadir, size_name='size_$id', stride
     return os.path.join(datadir,filename), meta, meta_free
 
 
-# TODO: add this function to an auxiliary file
-# call it something like TH${T}Storage_newFromFile(filename);
-def read_storage(storage_name,filepath,numtype):
-    subs = {
-      'filepath': filepath,
-      'storage_name': storage_name,
-      'real': type_map[numtype],
-      'T': numtype
-    }
-    # TODO: extend past little endian
-    tpl = '''
-      TH${T}Storage *${storage_name};
-      {
-        FILE *f = fopen("${filepath}","rb");
-        if (!f) {
-          THError("cannot open file ${filepath} for reading");
-        }
-        long size;
-        size_t result = fread(&size,sizeof(long),1,f);
-        ${storage_name} = TH${T}Storage_newWithSize(size);
-        char *bytes = (char *) ${storage_name}->data;
-        uint64_t remaining = sizeof(${real}) * ${storage_name}->size;
-        result = fread(bytes,sizeof(${real}),${storage_name}->size,f);
-        fclose(f);
-      }
-    '''
-    return Template(tpl).substitute(subs)
-
-
 class PersistedVariable(Variable):
 
     def __init__(self, obj, prevfns):
@@ -219,7 +190,7 @@ def call_tpl(self):
                                                    self.out_path,self.datadir)
 
         return '\n'.join([
-            read_storage('storage_$id',filepath,self.numtype), 
+            'TH${T}Storage *storage_$id = TH${T}Storage_newFromFile("%s");' % filepath,
             meta,
             'TH${T}Tensor *$id = TH${T}Tensor_newWithStorage(storage_$id,0,size_$id,stride_$id);',
             meta_free