Initial commit of framework for TensorLib

aten/src/aten/CMakeLists.txt

@@ -0,0 +1,117 @@
+
+CMAKE_MINIMUM_REQUIRED(VERSION 2.8)
+SET(CMAKE_MODULE_PATH ${CMAKE_CURRENT_SOURCE_DIR}/cmake ${CMAKE_MODULE_PATH})
+
+if(${CMAKE_VERSION} VERSION_LESS "2.8.12")
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11")
+else(${CMAKE_VERSION} VERSION_LESS "2.8.12")
+  if(${CMAKE_VERSION} VERSION_LESS "3.1")
+    add_compile_options(-std=c++11) # CMake 2.8.12 to 3.1
+  endif(${CMAKE_VERSION} VERSION_LESS "3.1")
+endif(${CMAKE_VERSION} VERSION_LESS "2.8.12")
+
+################################################################################
+# Helper functions
+################################################################################
+
+FUNCTION(EXCLUDE_DIR list_name dir_name)
+  # A helper that excludes all files that contain dir_name in their file path
+  SET(local_list ${${list_name}})
+  FOREACH(source ${local_list})
+    IF(${source} MATCHES ${dir_name})
+      MESSAGE(STATUS "Excluding " ${source} " from the build")
+      LIST(REMOVE_ITEM local_list ${source})
+    ENDIF()
+  ENDFOREACH()
+  SET(${list_name} ${local_list} PARENT_SCOPE)
+ENDFUNCTION()
+
+IF(NOT Torch_FOUND)
+  FIND_PACKAGE(Torch REQUIRED)
+ENDIF()
+
+IF(NOT TH_LIBRARIES)
+  SET(TH_LIBRARIES "TH")
+ENDIF(NOT TH_LIBRARIES)
+MESSAGE(STATUS "TH_LIBRARIES: ${TH_LIBRARIES}")
+
+IF(NOT THS_LIBRARIES)
+  SET(THS_LIBRARIES "THS")
+ENDIF()
+
+IF(NO_CUDA)
+  MESSAGE(STATUS "ignoring CUDA")
+ELSE()
+  FIND_PACKAGE(CUDA 5.5)
+  IF(CUDA_FOUND)
+    INCLUDE_DIRECTORIES(${CUDA_INCLUDE_DIRS})
+    INCLUDE_DIRECTORIES("${CUDA_SDK_ROOT_DIR}/common/inc")
+
+    IF(NOT THC_LIBRARIES)
+      SET(THC_LIBRARIES "THC")
+    ENDIF(NOT THC_LIBRARIES)
+    MESSAGE(STATUS "THC_LIBRARIES: ${THC_LIBRARIES}")
+
+    IF(NOT THCS_LIBRARIES)
+      SET(THCS_LIBRARIES "THCS")
+    ENDIF(NOT THCS_LIBRARIES)
+    MESSAGE(STATUS "THCS_LIBRARIES: ${THCS_LIBRARIES}")
+  ENDIF()
+ENDIF()
+
+# Can be compiled standalone
+IF(NOT TENSOR_LIB_INSTALL_BIN_DIR OR NOT TENSOR_LIB_INSTALL_LIB_DIR OR NOT TENSOR_LIB_INSTALL_INCLUDE_DIR)
+  SET(TENSOR_LIB_INSTALL_BIN_DIR "bin" CACHE PATH "TENSOR_LIB install binary subdirectory")
+  SET(TENSOR_LIB_INSTALL_LIB_DIR "lib" CACHE PATH "TENSOR_LIB install library subdirectory")
+  SET(TENSOR_LIB_INSTALL_INCLUDE_DIR "include" CACHE PATH "TENSOR_LIB install include subdirectory")
+ENDIF()
+
+FILE(GLOB_RECURSE base_h RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} "*.h" "*.hpp")
+FILE(GLOB_RECURSE base_cpp "*.cpp")
+
+EXCLUDE_DIR(base_cpp ".*/templates/.*$")
+
+SET(generated_cpp
+  CPUGenerator.h
+  CUDAGenerator.h
+)
+
+FILE(GLOB_RECURSE all_templates "templates/*")
+
+ADD_CUSTOM_COMMAND(OUTPUT ${generated_cpp}
+COMMAND python ${CMAKE_CURRENT_SOURCE_DIR}/gen.py -s ${CMAKE_CURRENT_SOURCE_DIR}
+DEPENDS gen.py ${all_templates})
+
+SET(all_cpp ${base_cpp} ${generated_cpp})
+SET(all_h ${base_h})
+
+INCLUDE_DIRECTORIES(${CMAKE_CURRENT_SOURCE_DIR})
+# so the build can find the generated header files
+INCLUDE_DIRECTORIES(${CMAKE_CURRENT_BINARY_DIR})
+ADD_LIBRARY(TensorLib SHARED ${all_cpp})
+SET_TARGET_PROPERTIES(TensorLib PROPERTIES VERSION 1 SOVERSION 1)
+
+ADD_EXECUTABLE(scalar_test scalar_test.cc)
+
+TARGET_LINK_LIBRARIES(scalar_test TensorLib)
+TARGET_LINK_LIBRARIES(scalar_test ${CUDA_LIBRARIES})
+
+if(NOT ${CMAKE_VERSION} VERSION_LESS "3.1")
+    SET_PROPERTY(TARGET TensorLib PROPERTY CXX_STANDARD 11)
+endif(NOT ${CMAKE_VERSION} VERSION_LESS "3.1")
+
+TARGET_LINK_LIBRARIES(TensorLib LINK_PRIVATE ${TH_LIBRARIES} ${THS_LIBRARIES})
+IF(CUDA_FOUND)
+  TARGET_LINK_LIBRARIES(TensorLib LINK_PRIVATE ${THC_LIBRARIES} ${THCS_LIBRARIES})
+  TARGET_LINK_LIBRARIES(TensorLib ${CUDA_LIBRARIES})
+ENDIF()
+
+INSTALL(TARGETS TensorLib
+  RUNTIME DESTINATION "${TENSOR_LIB_INSTALL_BIN_DIR}"
+  LIBRARY DESTINATION "${TENSOR_LIB_INSTALL_LIB_DIR}"
+  ARCHIVE DESTINATION "${TENSOR_LIB_INSTALL_LIB_DIR}")
+
+FOREACH(HEADER ${all_h})
+  STRING(REGEX MATCH "(.*)[/\\]" DIR ${HEADER})
+  INSTALL(FILES ${HEADER} DESTINATION ${TENSOR_LIB_INSTALL_INCLUDE_DIR}/TensorLib/${DIR})
+ENDFOREACH()

aten/src/aten/CPUGenerator.cpp

@@ -0,0 +1,36 @@
+#include "CPUGenerator.h"
+
+#define const_generator_cast(generator) \
+  dynamic_cast<const CPUGenerator&>(generator)
+
+namespace tlib {
+
+CPUGenerator::CPUGenerator(Context * context_)
+  : context(context_), generator(THGenerator_new())
+{}
+
+CPUGenerator::~CPUGenerator() {
+  if (generator)
+    THGenerator_free(generator);
+}
+
+CPUGenerator& CPUGenerator::copy(const Generator& from) {
+  THGenerator_copy(generator, const_generator_cast(from).generator);
+  return *this;
+}
+
+CPUGenerator& CPUGenerator::free() {
+  THGenerator_free(generator);
+  return *this;
+}
+
+unsigned long CPUGenerator::seed() {
+  return THRandom_seed(generator);
+}
+
+CPUGenerator& CPUGenerator::manualSeed(unsigned long seed) {
+  THRandom_manualSeed(generator, seed);
+  return *this;
+}
+
+} // namespace tlib

aten/src/aten/CUDAGenerator.cpp

@@ -0,0 +1,41 @@
+#include "CUDAGenerator.h"
+#include "Context.h"
+#include <stdexcept>
+
+#define const_generator_cast(generator) \
+  dynamic_cast<const CUDAGenerator&>(generator)
+
+namespace tlib {
+
+CUDAGenerator::CUDAGenerator(Context * context_)
+  : context(context_)
+{
+  int num_devices, current_device;
+  cudaGetDeviceCount(&num_devices);
+  cudaGetDevice(&current_device);
+  THCRandom_init(context->thc_state, num_devices, current_device);
+}
+
+CUDAGenerator::~CUDAGenerator() {
+  THCRandom_shutdown(context->thc_state);
+}
+
+CUDAGenerator& CUDAGenerator::copy(const Generator& from) {
+  throw std::runtime_error("CUDAGenerator::copy() not implemented");
+}
+
+CUDAGenerator& CUDAGenerator::free() {
+  THCRandom_shutdown(context->thc_state);
+  return *this;
+}
+
+unsigned long CUDAGenerator::seed() {
+  return THCRandom_initialSeed(context->thc_state);
+}
+
+CUDAGenerator& CUDAGenerator::manualSeed(unsigned long seed) {
+  THCRandom_manualSeed(context->thc_state, seed);
+  return *this;
+}
+
+} // namespace thpp

aten/src/aten/CodeTemplate.py

@@ -0,0 +1,42 @@
+import re
+
+# match $idenifier or ${identifier}, if the first group matches, then
+# this identifier is at the beginning of whitespace on a line and should be treated as
+# 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)
+    def from_file(filename):
+        with open(filename,'r') as f:
+            return CodeTemplate(f.read())
+    def __init__(self,pattern):
+        self.pattern = pattern
+    def substitute(self,env={},**kwargs):
+        def lookup(v):
+             return kwargs[v] if v in kwargs else env[v]
+        def indent_lines(indent,v):
+            return "".join([indent+l+"\n" for e in v for l in str(e).splitlines()]).rstrip()
+        def replace(match):
+            indent = match.group(1)
+            key = match.group(2)
+            if key[0] == "{":
+                key = key[1:-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])
+            else:
+                return (indent or '') + str(v)
+        return self.subtitution.sub(replace,self.pattern)
+
+if __name__ == "__main__":
+    c = CodeTemplate("""\
+    int foo($args) {
+
+        $bar
+            $bar
+        $a+$b
+    }
+    """)
+    print(c.substitute(args=["hi",8],bar=["what",7],a=3,b=4))

aten/src/aten/Context.cpp

@@ -0,0 +1,33 @@
+#include "Context.h"
+#include <thread>
+#include <mutex>
+
+namespace tlib {
+
+Context::Context() {
+  thc_state = THCState_alloc();
+  THCState_setDeviceAllocator(thc_state, THCCachingAllocator_get());
+  thc_state->cudaHostAllocator = &THCCachingHostAllocator;
+  THCudaInit(thc_state);
+  cuda_gen.reset(new CUDAGenerator(this));
+  cpu_gen.reset(new CPUGenerator(this));
+}
+
+Context::~Context() {
+  THCState_free(thc_state);
+}
+
+static std::mutex context_lock;
+static Context * globalContext_ = nullptr;
+
+Context * globalContext() {
+  if(!globalContext_) {
+    std::lock_guard<std::mutex> lock(context_lock);
+    if (!globalContext_) {
+      globalContext_ = new Context();
+    }
+  }
+  return globalContext_;
+}
+
+}

aten/src/aten/Context.h

@@ -0,0 +1,20 @@
+#pragma once
+
+#include "CPUGenerator.h"
+#include "CUDAGenerator.h"
+#include <memory>
+
+namespace tlib {
+
+class Context {
+public:
+  Context();
+  ~Context();
+  std::unique_ptr<CPUGenerator> cpu_gen;
+  std::unique_ptr<CUDAGenerator> cuda_gen;
+  THCState * thc_state;
+};
+
+Context * globalContext();
+
+}

aten/src/aten/Generator.h

@@ -0,0 +1,18 @@
+#pragma once
+
+namespace tlib {
+
+struct Generator {
+  Generator() {};
+  Generator(const Generator& other) = delete;
+  Generator(Generator&& other) = delete;
+  virtual ~Generator() {};
+
+  virtual Generator& copy(const Generator& other) = 0;
+  virtual Generator& free() = 0;
+
+  virtual unsigned long seed() = 0;
+  virtual Generator& manualSeed(unsigned long seed) = 0;
+};
+
+} // namespace tlib

aten/src/aten/Scalar.cpp

@@ -0,0 +1,24 @@
+#include <Scalar.h>
+#include <TH/TH.h>
+
+namespace tlib {
+
+template<> Half convert(double f) {
+  float t = static_cast<float>(f);
+  Half h;
+  TH_float2halfbits(&t,&h.x);
+  return h;
+}
+template<> double convert(Half f) {
+  float t;
+  TH_halfbits2float(&f.x,&t);
+  return t;
+}
+template<> Half convert(int64_t f) {
+  return convert<Half,double>(static_cast<double>(f));
+}
+template<> int64_t convert(Half f) {
+  return static_cast<int64_t>(convert<double,Half>(f));
+}
+
+}

aten/src/aten/Scalar.h

@@ -0,0 +1,87 @@
+#pragma once
+
+#include<stdint.h>
+#include <stdexcept>
+#include <string>
+
+#if defined(__GNUC__)
+#define TLIB_ALIGN(n) __attribute__((aligned(n)))
+#elif defined(_WIN32)
+#define TLIB_ALIGN(n) __declspec(align(n))
+#else
+#define TLIB_ALIGN(n)
+#endif
+
+
+
+namespace tlib {
+
+typedef struct TLIB_ALIGN(2){
+  unsigned short x;
+} Half;
+
+#define TLIB_SCALAR_TYPES(_) \
+_(uint8_t,Byte,i) \
+_(int8_t,Char,i) \
+_(double,Double,d) \
+_(float,Float,d) \
+_(int,Int,i) \
+_(int64_t,Long,i) \
+_(int16_t,Short,i) \
+_(Half,Half,d)
+
+template<typename To, typename From> To convert(From f) {
+  return static_cast<To>(f);
+}
+template<> Half convert(double f);
+template<> double convert(Half f);
+template<> Half convert(int64_t f);
+template<> int64_t convert(Half f);
+
+class Scalar {
+public:
+#define DEFINE_IMPLICIT_CTOR(type,name,member) \
+  Scalar(type v) \
+  : tag(Tag::HAS_##member) { \
+    member = convert<decltype(member),type>(v); \
+  } \
+
+  TLIB_SCALAR_TYPES(DEFINE_IMPLICIT_CTOR)
+#undef DEFINE_IMPLICIT_CTOR
+
+#define DEFINE_ACCESSOR(type,name,member) \
+  type to##name () { \
+    if (Tag::HAS_d == tag) { \
+      auto casted = convert<type,double>(d); \
+      if(convert<double,type>(casted) != d) { \
+        throw std::domain_error(std::string("value cannot be losslessly represented in type " #name ": ") + std::to_string(d) ); \
+      } \
+      return casted; \
+    } else { \
+      auto casted = convert<type,int64_t>(i); \
+      if(convert<int64_t,type>(casted) != i) { \
+        throw std::domain_error(std::string("value cannot be losslessly represented in type " #name ": ") + std::to_string(i)); \
+      } \
+      return casted; \
+    } \
+  }
+
+  TLIB_SCALAR_TYPES(DEFINE_ACCESSOR)
+
+#undef DEFINE_ACCESSOR
+  bool isFloatingPoint() {
+    return Tag::HAS_d == tag;
+  }
+  bool isIntegral() {
+    return Tag::HAS_i == tag;
+  }
+private:
+  enum class Tag { HAS_d, HAS_i };
+  Tag tag;
+  union {
+    double d;
+    int64_t i;
+  };
+};
+
+}

aten/src/aten/Storage.h

@@ -0,0 +1,33 @@
+#pragma once
+
+#include "Scalar.h"
+
+namespace tlib {
+
+struct Storage {
+  Storage() {};
+  Storage(const Storage& other) = delete;
+  Storage(Storage&& other) = delete;
+  virtual ~Storage() {};
+
+  virtual std::size_t elementSize() const = 0;
+  virtual std::size_t size() const = 0;
+  virtual void* data() = 0;
+  virtual const void* data() const = 0;
+  virtual Storage& retain() = 0;
+  virtual Storage& free() = 0;
+
+  virtual Storage& resize(long new_size) = 0;
+
+  virtual Type & type() const = 0;
+  virtual int getDevice() const = 0;
+
+  virtual Scalar 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;
+  virtual Scalar fast_get(std::size_t ind) = 0;
+
+};
+
+} // namespace tlib

aten/src/aten/TensorLib.h

@@ -0,0 +1,9 @@
+#pragma once
+
+#include "Scalar.h"
+#include "Type.h"
+#include "Generator.h"
+#include "CPUGenerator.h"
+#include "CUDAGenerator.h"
+#include "Context.h"
+#include "Storage.h"

aten/src/aten/Type.h

@@ -0,0 +1,26 @@
+#pragma once
+
+#include "Scalar.h"
+
+namespace tlib {
+
+enum class ScalarType {
+#define DEFINE_ENUM(_1,n,_2) \
+  n,
+  TLIB_SCALAR_TYPES(DEFINE_ENUM)
+#undef DEFINE_ENUM
+};
+
+enum class Processor {
+  CPU,
+  CUDA
+};
+
+class Type {
+  virtual ScalarType scalarType() = 0;
+  virtual Processor processor() = 0;
+  virtual bool isSparse() = 0;
+  virtual bool isDistributed() = 0;
+};
+
+}

aten/src/aten/gen.py

@@ -0,0 +1,31 @@
+import os
+from optparse import OptionParser
+from CodeTemplate import CodeTemplate
+import sys
+parser = OptionParser()
+parser.add_option('-s', '--source-path', help='path to source director for tensorlib',
+    action='store', default='.')
+options,args = parser.parse_args()
+
+generators = {
+    'CPUGenerator.h' : {
+        'name' : 'CPU',
+        'th_generator':'THGenerator * generator;',
+        'header': 'TH/TH.h',
+    },
+    'CUDAGenerator.h' : {
+        'name' : 'CUDA',
+        'th_generator': '',
+        'header': 'THC/THC.h'
+    },
+}
+
+TEMPLATE_PATH =  options.source_path+"/templates"
+GENERATOR_DERIVED = CodeTemplate.from_file(TEMPLATE_PATH+"/GeneratorDerived.h")
+
+def write(f,s):
+    with open(fname,"w") as f:
+        f.write(s)
+
+for fname,env in generators.items():
+    write(fname,GENERATOR_DERIVED.substitute(env))

aten/src/aten/scalar_test.cc

@@ -0,0 +1,18 @@
+#include <iostream>
+#include "TensorLib.h"
+
+using std::cout;
+using namespace tlib;
+
+int main() {
+  Scalar what = 257;
+  Scalar bar = 3.0;
+  Half h = bar.toHalf();
+  Scalar h2 = h;
+
+  cout << "H2: " << h2.toDouble() << " " << what.toFloat() << " " << bar.toDouble() << " " << what.isIntegral() <<  "\n";
+  CUDAGenerator gen(tlib::globalContext());
+  cout << gen.seed();
+
+
+}

aten/src/aten/templates/GeneratorDerived.h

@@ -0,0 +1,25 @@
+#pragma once
+
+#include <$header>
+
+#include "Generator.h"
+
+namespace tlib {
+
+class Context;
+struct ${name}Generator : public Generator {
+  ${name}Generator(Context * context);
+  virtual ~${name}Generator();
+
+  virtual ${name}Generator& copy(const Generator& from) override;
+  virtual ${name}Generator& free() override;
+
+  virtual unsigned long seed() override;
+  virtual ${name}Generator& manualSeed(unsigned long seed) override;
+
+protected:
+  Context * context;
+  ${th_generator}
+};
+
+}

aten/tools/subtree_dir_add.sh

@@ -0,0 +1,17 @@
+SRC_BRANCH=$1
+SRC_PATH=$2
+DST_PATH=$3
+
+CURRENT_BRANCH=$(git rev-parse --abbrev-ref HEAD)
+set -x
+git branch -q -D temporary-split-branch
+git checkout $SRC_BRANCH
+git subtree split -P $SRC_PATH -b temporary-split-branch
+git checkout $CURRENT_BRANCH
+git subtree add -P $DST_PATH temporary-split-branch ${@:5}
+git branch -D temporary-split-branch
+
+#./subtree_dir.sh add pytorch/master torch/lib/TH lib/TH
+#./subtree_dir.sh add pytorch/master torch/lib/THC lib/THC
+#./subtree_dir.sh add pytorch/master torch/lib/THNN lib/THNN
+#./subtree_dir.sh add pytorch/master torch/lib/THCUNN lib/THCUNN

aten/tools/subtree_dir_merge.sh

@@ -0,0 +1,17 @@
+SRC_BRANCH=$1
+SRC_PATH=$2
+DST_PATH=$3
+
+CURRENT_BRANCH=$(git rev-parse --abbrev-ref HEAD)
+set -x
+git branch -q -D temporary-split-branch
+git checkout $SRC_BRANCH
+git subtree split -P $SRC_PATH -b temporary-split-branch
+git checkout $CURRENT_BRANCH
+git subtree merge -P $DST_PATH temporary-split-branch -m "Merge commit '`git rev-parse temporary-split-branch`'"
+git branch -D temporary-split-branch
+
+#./subtree_dir.sh merge pytorch/master torch/lib/TH lib/TH
+#./subtree_dir.sh merge pytorch/master torch/lib/THC lib/THC
+#./subtree_dir.sh merge pytorch/master torch/lib/THNN lib/THNN
+#./subtree_dir.sh merge pytorch/master torch/lib/THCUNN lib/THCUNN