add storage to generator

aten/src/aten/CMakeLists.txt

@@ -74,6 +74,39 @@ EXCLUDE_DIR(base_cpp ".*/templates/.*$")
 SET(generated_cpp
   CPUGenerator.h
   CUDAGenerator.h
+  CPUByteStorage.h
+  CPUCharStorage.h
+  CPUDoubleStorage.h
+  CPUFloatStorage.h
+  CPUIntStorage.h
+  CPULongStorage.h
+  CPUShortStorage.h
+  CPUHalfStorage.h
+  CUDAByteStorage.h
+  CUDACharStorage.h
+  CUDADoubleStorage.h
+  CUDAFloatStorage.h
+  CUDAIntStorage.h
+  CUDALongStorage.h
+  CUDAShortStorage.h
+  CUDAHalfStorage.h
+  
+  CPUByteStorage.cpp
+  CPUCharStorage.cpp
+  CPUDoubleStorage.cpp
+  CPUFloatStorage.cpp
+  CPUIntStorage.cpp
+  CPULongStorage.cpp
+  CPUShortStorage.cpp
+  CPUHalfStorage.cpp
+  CUDAByteStorage.cpp
+  CUDACharStorage.cpp
+  CUDADoubleStorage.cpp
+  CUDAFloatStorage.cpp
+  CUDAIntStorage.cpp
+  CUDALongStorage.cpp
+  CUDAShortStorage.cpp
+  CUDAHalfStorage.cpp
 )
 
 FILE(GLOB_RECURSE all_templates "templates/*")

aten/src/aten/CodeTemplate.py

@@ -5,7 +5,7 @@ import re
 # block subsitution by identing to that depth and replacing
 
 class CodeTemplate(object):
-    subtitution = re.compile('(^[^\n\S]*)?\$([^\d\W]\w*|\{[^\d\W]\w*\})',re.MULTILINE)
+    subtitution = re.compile('(^[^\n\S]*)?\$([^\d\W]\w*|\{,?[^\d\W]\w*\,?})',re.MULTILINE)
     def from_file(filename):
         with open(filename,'r') as f:
             return CodeTemplate(f.read())
@@ -19,13 +19,24 @@ class CodeTemplate(object):
         def replace(match):
             indent = match.group(1)
             key = match.group(2)
+            comma_before=''
+            comma_after=''
             if key[0] == "{":
                 key = key[1:-1]
+                if key[0] == ",":
+                    comma_before = ', '
+                    key = key[1:]
+                if key[-1] == ',':
+                    comma_after = ', '
+                    key = key[:-1]
             v = lookup(key)
             if indent is not None and isinstance(v,list):
                 return indent_lines(indent,v)
             elif isinstance(v,list):
-                return ', '.join([str(x) for x in v])
+                middle = ', '.join([str(x) for x in v])
+                if len(v) == 0:
+                    return middle
+                return comma_before+middle+comma_after
             else:
                 return (indent or '') + str(v)
         return self.subtitution.sub(replace,self.pattern)
@@ -38,5 +49,7 @@ if __name__ == "__main__":
             $bar
         $a+$b
     }
+    int commatest(int a${,stuff})
+    int notest(int a${,empty,})
     """)
-    print(c.substitute(args=["hi",8],bar=["what",7],a=3,b=4))
+    print(c.substitute(args=["hi",8],bar=["what",7],a=3,b=4,stuff=["things...","others"],empty=[]))

aten/src/aten/HalfConvert.h

@@ -0,0 +1,13 @@
+#pragma once
+
+#include "Scalar.h"
+#include "TH/TH.h"
+
+namespace tlib {
+
+  template<typename To, typename From>
+  static inline To HalfFix(From h) {
+    return To { h.x };
+  }
+
+}

aten/src/aten/Storage.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include "Scalar.h"
+#include "Type.h"
 
 namespace tlib {
 
@@ -22,7 +23,7 @@ struct Storage {
   virtual Type & type() const = 0;
   virtual int getDevice() const = 0;
 
-  virtual Scalar fill(Scalar value) = 0;
+  virtual Storage& fill(Scalar value) = 0;
   virtual Storage& set(std::size_t ind, Scalar value) = 0;
   virtual Storage& fast_set(std::size_t ind, Scalar value) = 0;
   virtual Scalar get(std::size_t ind) = 0;

aten/src/aten/TensorLib.h

@@ -7,3 +7,20 @@
 #include "CUDAGenerator.h"
 #include "Context.h"
 #include "Storage.h"
+
+#include "CPUByteStorage.h"
+#include "CPUCharStorage.h"
+#include "CPUDoubleStorage.h"
+#include "CPUFloatStorage.h"
+#include "CPUIntStorage.h"
+#include "CPULongStorage.h"
+#include "CPUShortStorage.h"
+#include "CPUHalfStorage.h"
+#include "CUDAByteStorage.h"
+#include "CUDACharStorage.h"
+#include "CUDADoubleStorage.h"
+#include "CUDAFloatStorage.h"
+#include "CUDAIntStorage.h"
+#include "CUDALongStorage.h"
+#include "CUDAShortStorage.h"
+#include "CUDAHalfStorage.h"

aten/src/aten/gen.py

@@ -7,6 +7,11 @@ parser.add_option('-s', '--source-path', help='path to source director for tenso
     action='store', default='.')
 options,args = parser.parse_args()
 
+TEMPLATE_PATH =  options.source_path+"/templates"
+GENERATOR_DERIVED = CodeTemplate.from_file(TEMPLATE_PATH+"/GeneratorDerived.h")
+STORAGE_DERIVED_CPP = CodeTemplate.from_file(TEMPLATE_PATH+"/StorageDerived.cpp")
+STORAGE_DERIVED_H = CodeTemplate.from_file(TEMPLATE_PATH+"/StorageDerived.h")
+
 generators = {
     'CPUGenerator.h' : {
         'name' : 'CPU',
@@ -20,12 +25,63 @@ generators = {
     },
 }
 
-TEMPLATE_PATH =  options.source_path+"/templates"
-GENERATOR_DERIVED = CodeTemplate.from_file(TEMPLATE_PATH+"/GeneratorDerived.h")
+processors = [ 'CPU', 'CUDA']
+scalar_types = [
+    ('Byte','uint8_t'),
+    ('Char','int8_t'),
+    ('Double','double'),
+    ('Float','float'),
+    ('Int','int'),
+    ('Long','int64_t'),
+    ('Short','int16_t'),
+    ('Half','Half'),
+]
 
-def write(f,s):
-    with open(fname,"w") as f:
+
+def write(filename,s):
+    with open(filename,"w") as f:
         f.write(s)
 
+def generate_storage(processor, scalar_type):
+    scalar_name, c_type = scalar_type
+    env = {}
+    env['ScalarName'] = scalar_name
+    env['ScalarType'] = c_type
+    env['Storage'] = "{}{}Storage".format(processor,scalar_name)
+    if processor == 'CUDA':
+        env['th_header'] = "THC/THC.h"
+        sname = '' if scalar_name == "Float" else scalar_name
+        env['THStorage'] = 'THCuda{}Storage'.format(sname)
+        env['state'] = ['context->thc_state']
+        env['isCUDA'] = 'true'
+        env['storage_device'] = 'return storage->device;'
+    else:
+        env['th_header'] = "TH/TH.h"
+        env['THStorage'] = "TH{}Storage".format(scalar_name)
+        env['state'] = []
+        env['isCUDA'] = 'false'
+        env['storage_device'] = 'throw std::runtime_error("CPU storage has no device");'
+
+    if scalar_name == "Half":
+        if processor == "CUDA":
+            env['to_th_half'] = 'HalfFix<__half,Half>'
+            env['to_tlib_half'] = 'HalfFix<Half,__half>'
+        else:
+            env['to_th_half'] = 'HalfFix<THHalf,Half>'
+            env['to_tlib_half'] = 'HalfFix<Half,THHalf>'
+    else:
+        env['to_th_half'] = ''
+        env['to_tlib_half'] = ''
+
+
+    write(env['Storage']+".cpp",STORAGE_DERIVED_CPP.substitute(env))
+    write(env['Storage']+".h",STORAGE_DERIVED_H.substitute(env))
+
+    print("#include '{}.h'".format(env['Storage']))
+
 for fname,env in generators.items():
     write(fname,GENERATOR_DERIVED.substitute(env))
+
+for processor in processors:
+    for scalar_type in scalar_types:
+        generate_storage(processor,scalar_type)

aten/src/aten/scalar_test.cc

@@ -12,7 +12,13 @@ int main() {
 
   cout << "H2: " << h2.toDouble() << " " << what.toFloat() << " " << bar.toDouble() << " " << what.isIntegral() <<  "\n";
   CUDAGenerator gen(tlib::globalContext());
-  cout << gen.seed();
+  cout << gen.seed() << "\n";
 
+  CPUFloatStorage s(tlib::globalContext());
+
+  s.resize(4);
+  s.fill(7);
+
+  cout << "GET " << s.get(3).toFloat() << "\n";
 
 }

aten/src/aten/templates/StorageDerived.cpp

@@ -0,0 +1,82 @@
+#include "${Storage}.h"
+#include "HalfConvert.h"
+
+namespace tlib {
+
+${Storage}::${Storage}(Context* context):
+    storage(${THStorage}_new(${state})), context(context) {}
+
+${Storage}::${Storage}(Context* context, ${THStorage}* storage):
+    storage(storage), context(context) {}
+
+${Storage}::${Storage}(Context* context, std::size_t storage_size)
+  : storage(${THStorage}_newWithSize(${state,} storage_size)), context(context) {}
+
+${Storage}::~${Storage}() {
+  ${THStorage}_free(${state,} storage);
+}
+
+std::size_t ${Storage}::elementSize() const {
+  return sizeof(${ScalarType});
+}
+
+std::size_t ${Storage}::size() const {
+  return storage->size;
+}
+
+void* ${Storage}::data() {
+  return storage->data;
+}
+
+const void* ${Storage}::data() const {
+  return storage->data;
+}
+
+auto ${Storage}::retain() -> ${Storage}& {
+  ${THStorage}_retain(${state,} storage);
+  return *this;
+}
+
+auto ${Storage}::free() -> ${Storage}& {
+  ${THStorage}_free(${state,} storage);
+  return *this;
+}
+
+auto ${Storage}::resize(long new_size) -> ${Storage}& {
+  ${THStorage}_resize(${state,} storage, new_size);
+  return *this;
+}
+
+auto ${Storage}::fill(Scalar value) -> ${Storage}& {
+  ${THStorage}_fill(${state,} storage, ${to_th_half}(value.to${ScalarName}()));
+  return *this;
+}
+
+auto ${Storage}::set(std::size_t ind, Scalar value) -> ${Storage}& {
+  ${THStorage}_set(${state,} storage, ind, ${to_th_half}(value.to${ScalarName}()));
+  return *this;
+}
+
+auto ${Storage}::fast_set(std::size_t ind, Scalar value) -> ${Storage}& {
+  throw std::runtime_error("unsupported operation 'fast_set'");
+}
+
+auto ${Storage}::get(std::size_t ind) -> Scalar {
+  return ${to_tlib_half}(${THStorage}_get(${state,} storage, ind));
+}
+
+auto ${Storage}::fast_get(std::size_t ind) -> Scalar {
+  if(${isCUDA})
+    throw std::runtime_error("unsupported operation 'fast_get'");
+  return ${to_tlib_half}(storage->data[ind]);
+}
+
+int ${Storage}::getDevice() const {
+  ${storage_device} //storage->device;
+}
+
+Type& ${Storage}::type() const {
+  throw std::runtime_error("NYI - Storage::type()");
+}
+
+}

aten/src/aten/templates/StorageDerived.h

@@ -0,0 +1,40 @@
+#pragma once
+
+#include <$th_header>
+
+#include "Storage.h"
+#include "Context.h"
+
+namespace tlib {
+
+struct ${Storage} : public Storage {
+public:
+  ${Storage}(Context* context);
+  ${Storage}(Context* context, ${THStorage} *wrapped);
+  ${Storage}(Context* context, std::size_t size);
+  virtual ~${Storage}();
+
+  virtual std::size_t elementSize() const override;
+  virtual std::size_t size() const override;
+  virtual void* data() override;
+  virtual const void* data() const override;
+  virtual ${Storage}& retain() override;
+  virtual ${Storage}& free() override;
+
+  virtual ${Storage}& resize(long new_size) override;
+  virtual ${Storage}& fill(Scalar value) override;
+  virtual ${Storage}& set(std::size_t ind, Scalar value) override;
+  virtual ${Storage}& fast_set(std::size_t ind, Scalar value) override;
+  virtual Scalar get(std::size_t ind) override;
+  virtual Scalar fast_get(std::size_t ind) override;
+
+  virtual Type& type() const override;
+  virtual int getDevice() const override;
+
+
+protected:
+  ${THStorage} *storage;
+  Context* context;
+};
+
+} // namespace thpp

aten/src/aten/templates/TensorImpl.cpp

@@ -0,0 +1,36 @@
+
+
+//sketch:
+
+/*
+
+class Type {
+  // ctor is configurable as well
+  Tensor new(...) {}
+  // all  methods and pure functions, virtually dispatched
+  virtual Tensor add(Tensor a, Tensor b);
+
+  // no Tensor arguments but virtually dispatched on type
+  virtual Tensor linspace(int a, int b);
+  //
+}
+
+// base class of derived tensors
+class TensorImpl {
+  Type * type_; //non-virtual so that we don't double indirect unnecessarily ...
+
+  // not virtual - Type * handles the virtual dispatch
+  Tensor add(Tensor b) {
+    type_->add(self,b);
+  }
+}
+
+// non-method style dispatch
+Tensor add(Tensor a, Tensor b) {
+  return a->type_->add(a,b);
+}
+
+
+}
+
+*/