From 0e0c0ef89e4884019fde4b6ed770511304196738 Mon Sep 17 00:00:00 2001 From: Zachary DeVito Date: Fri, 2 Jun 2017 23:39:23 +0000 Subject: [PATCH] add storage to generator --- aten/src/aten/CMakeLists.txt | 33 +++++++++ aten/src/aten/CodeTemplate.py | 19 ++++- aten/src/aten/HalfConvert.h | 13 ++++ aten/src/aten/Storage.h | 3 +- aten/src/aten/TensorLib.h | 17 +++++ aten/src/aten/gen.py | 64 +++++++++++++++-- aten/src/aten/scalar_test.cc | 8 ++- aten/src/aten/templates/StorageDerived.cpp | 82 ++++++++++++++++++++++ aten/src/aten/templates/StorageDerived.h | 40 +++++++++++ aten/src/aten/templates/TensorImpl.cpp | 36 ++++++++++ 10 files changed, 306 insertions(+), 9 deletions(-) create mode 100644 aten/src/aten/HalfConvert.h diff --git a/aten/src/aten/CMakeLists.txt b/aten/src/aten/CMakeLists.txt index 36a82e90e3a..9c9bab50a80 100644 --- a/aten/src/aten/CMakeLists.txt +++ b/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/*") diff --git a/aten/src/aten/CodeTemplate.py b/aten/src/aten/CodeTemplate.py index 7d2f9385f12..4f597ee6284 100644 --- a/aten/src/aten/CodeTemplate.py +++ b/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=[])) diff --git a/aten/src/aten/HalfConvert.h b/aten/src/aten/HalfConvert.h new file mode 100644 index 00000000000..9b3553d3903 --- /dev/null +++ b/aten/src/aten/HalfConvert.h @@ -0,0 +1,13 @@ +#pragma once + +#include "Scalar.h" +#include "TH/TH.h" + +namespace tlib { + + template + static inline To HalfFix(From h) { + return To { h.x }; + } + +} diff --git a/aten/src/aten/Storage.h b/aten/src/aten/Storage.h index f80ff16c909..b7d70a7fd88 100644 --- a/aten/src/aten/Storage.h +++ b/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; diff --git a/aten/src/aten/TensorLib.h b/aten/src/aten/TensorLib.h index ca6bd4eef3f..c24a1fad002 100644 --- a/aten/src/aten/TensorLib.h +++ b/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" diff --git a/aten/src/aten/gen.py b/aten/src/aten/gen.py index b0c45281c0e..f504c3c3551 100644 --- a/aten/src/aten/gen.py +++ b/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' + else: + env['to_th_half'] = 'HalfFix' + env['to_tlib_half'] = 'HalfFix' + 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) diff --git a/aten/src/aten/scalar_test.cc b/aten/src/aten/scalar_test.cc index cb99e506302..f91693724dd 100644 --- a/aten/src/aten/scalar_test.cc +++ b/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"; } diff --git a/aten/src/aten/templates/StorageDerived.cpp b/aten/src/aten/templates/StorageDerived.cpp index e69de29bb2d..61256fa4030 100644 --- a/aten/src/aten/templates/StorageDerived.cpp +++ b/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()"); +} + +} diff --git a/aten/src/aten/templates/StorageDerived.h b/aten/src/aten/templates/StorageDerived.h index e69de29bb2d..649b8b0a417 100644 --- a/aten/src/aten/templates/StorageDerived.h +++ b/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 diff --git a/aten/src/aten/templates/TensorImpl.cpp b/aten/src/aten/templates/TensorImpl.cpp index e69de29bb2d..7c2987a90e0 100644 --- a/aten/src/aten/templates/TensorImpl.cpp +++ b/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); +} + + +} + +*/