scalar flags added, and used to dispatch when there is a scalar variant of a function. broadcast annotations are used to figure out when a scalar s + A should also be converted.

Author: zdevito

Formatting.cpp

@@ -231,9 +231,8 @@ std::ostream& print(std::ostream& stream, const Tensor & tensor_, int64_t linesi
     Tensor tensor = tensor_.toType(getType(kCPU,kDouble)).contiguous();
     if(tensor.ndimension() == 0) {
       stream << std::defaultfloat << tensor.data<double>()[0] << std::endl;
-//      stream << "[Tensor<" << typedesc() << "," << devicedesc() << "> (value)]";
-    }
-    else if(tensor.ndimension() == 1) {
+      stream << "[" << tensor_.pImpl->toString() << " of size {}]";
+    } else if(tensor.ndimension() == 1) {
       double scale;
       int64_t sz;
       std::tie(scale, sz) =  __printFormat(stream, tensor);

THLongStorageView.h

@@ -5,24 +5,33 @@
 namespace at {
 
 // make a fake storage out of a size, pointer pair...
+// used as an argument where THSize and THStride are passed into TH
 class THLongStorageView {
 public:
-  static THLongStorageView make(ArrayRef<int64_t> ref) {
-    return THLongStorageView(ref);
+  static THLongStorageView make(ArrayRef<int64_t> ref, bool zero_dim_to_one = false) {
+    return THLongStorageView(ref,zero_dim_to_one);
   }
   operator THLongStorage*() {
     return &storage;
   }
 private:
-  THLongStorageView(ArrayRef<int64_t> ref) {
-    storage.data = (long*)(ref.data());
-    storage.size = ref.size();
+  THLongStorageView(ArrayRef<int64_t> ref, bool zero_dim_to_one) {
+    if(zero_dim_to_one && ref.size() == 0) {
+      // make storage of size 0 actually a 1-length storage with 1 element
+      // so that our 0-dim tensors get allocated as 1-dim inside TH
+      one = 1;
+      storage.data = &one;
+      storage.size = 1;
+    } else {
+      storage.data = (long*)(ref.data());
+      storage.size = ref.size();
+    }
     storage.refcount = 0;
     storage.flag = 0;
     storage.allocator = nullptr;
     storage.allocatorContext = nullptr;
   }
-
+  long one;
   THLongStorage storage;
 };
 

TensorImpl.h

@@ -10,14 +10,15 @@ namespace at {
 class Type;
 struct TensorImpl {
   TensorImpl(Type * type)
-  : type_(type), refcount(1) {}
+  : type_(type), refcount(1), is_scalar(false) {}
   Type & type() const {
     return *type_;
   }
-
   virtual const char * toString() const = 0;
   virtual IntList sizes() = 0;
   virtual IntList strides() = 0;
+  virtual int64_t dim() = 0;
+  virtual Scalar localScalar() = 0;
   void retain() {
     ++refcount;
   }
@@ -27,10 +28,23 @@ struct TensorImpl {
     }
   }
   virtual ~TensorImpl() {}
-
   friend class Type;
+
+  // 0-dim patchup of TH requires us to have a flag marking
+  // if a Tensor should be treated as 0-dim.
+  // the generated wrapper manipulates this flag.
+  // the setter should never be exposed in Tensor's public API
+  // because eventually we would like isScalar() to just be dim() == 0;
+  bool isScalar() const {
+    return is_scalar;
+  }
+  void setScalar(bool s) {
+    is_scalar = s;
+  }
+
 private:
   std::atomic<int> refcount;
+  bool is_scalar;
   Type * type_;
 };
 

copy_wrapper.py

@@ -32,6 +32,7 @@
       throw std::runtime_error("unsupported type in copy");
       break;
   }
+  dst.pImpl->setScalar(src.pImpl->isScalar());
 }
 """)
 

function_wrapper.py

@@ -45,6 +45,18 @@
 }
 """)
 
+ZERO_DIM_CHECK = CodeTemplate("""\
+if(${check_name}.dim() == 0) {
+    return ${method_prefix}${api_name}(${zero_dim_actuals});
+}""")
+
+SCALAR_EXPAND = CodeTemplate("""\
+Tensor ${name}__;
+if(${name}_->isScalar()) {
+    ${name}__ = ${name}.expand(${other}.sizes());
+    ${name}_ = static_cast<${Tensor}*>(${name}__.pImpl);
+}
+""")
 
 class NYIError(Exception):
     """Indicates we don't support this declaration yet"""
@@ -83,8 +95,8 @@ def __init__(self, reason):
     'THIntegerTensor*': CodeTemplate('checked_cast<${Backend}IntTensor>(${arg_name}.pImpl,"${arg_name}",${arg_pos})'),
     'THStorage*': CodeTemplate('checked_cast<${Storage}>(&${arg_name},"${arg_name}",${arg_pos})'),
     'THGenerator*': CodeTemplate('check_generator(&${arg_name})'),
-    'THSize*': CodeTemplate('THLongStorageView::make(${arg_name})'),
-    'THStride*': CodeTemplate('THLongStorageView::make(${arg_name})'),
+    'THSize*': CodeTemplate('THLongStorageView::make(${arg_name},true)'),
+    'THStride*': CodeTemplate('THLongStorageView::make(${arg_name},true)'),
     'real': CodeTemplate('${arg_name}.to${ScalarName}()'),
     'accreal': CodeTemplate('${arg_name}.to${AccScalarName}()'),
 
@@ -290,16 +302,27 @@ def is_actual_return_long(ret):
         return ret['type'] == 'long' or (backend_type_env['ScalarName'] == 'Long' and
                                          ret['type'] == 'real' or ret['type'] == 'accreal')
 
+    def handle_zero_dim(env,option):
+        if 'zero_dim_dispatch_when_scalar' not in option:
+            return []
+        check_name = option['zero_dim_dispatch_when_scalar']
+        zero_dim_actuals = [ arg['name']
+                             if arg['name'] != check_name else arg['name']+'.scalar()'
+                             for arg in option['formals_list'] ]
+        return [ ZERO_DIM_CHECK.substitute(env,check_name = check_name, zero_dim_actuals=zero_dim_actuals) ]
+
     def emit_body(env, option):
         body = []
+        body += handle_zero_dim(env,option)
         # arguments are potentially duplicated because of one argument
         # referencing another
         seen_names = set()
-        # only generated checked casts the first time we see it
         count = 0
         for arg in option['arguments']:
             if is_real_argument_to_wrapper(arg):
                 count += 1
+
+            # only generated checked casts the first time we see it
             if not arg['name'] in seen_names and requires_checked_cast(arg):
                 seen_names.add(arg['name'])
                 if arg.get('allocate', False):
@@ -326,19 +349,25 @@ def emit_body(env, option):
                             arg['name'], ','.join(dims)))
                 if arg.get('cpu_zero', False):
                     body.append("{}.zero_();".format(arg['name']))
-
-        option['actuals'] = get_arguments(option)
+                # handle scalars that occur on LHS of things like a - b
+                if 'broadcast' in arg and 'inplace' not in arg['broadcast']:
+                    other = arg['broadcast'].split(' ')[0].split(',')[0]
+                    body.append(SCALAR_EXPAND.substitute(env,
+                                                         name=arg['name'],
+                                                         other=other))
+
+        option['derived_actuals'] = get_arguments(option)
         is_cuda = backend_type_env['Backend'] == 'CUDA'
         is_nn = option['mode'] == 'NN'
         if is_cuda or is_nn:
-            option['actuals'] = ['context->thc_state'] + option['actuals']
+            option['derived_actuals'] = ['context->thc_state'] + option['derived_actuals']
 
         if is_nn:
             prefix = 'THNN_{}'.format(env['THType'])
         else:
             prefix = env['THTensor'] + '_'
 
-        call = prefix + CodeTemplate("${cname}(${actuals})").substitute(env)
+        call = prefix + CodeTemplate("${cname}(${derived_actuals})").substitute(env)
         ret = option['return']
         if ret['kind'] == 'arguments':
             body.append(call + ";")

preprocess_declarations.py

@@ -2,6 +2,7 @@
 from copy import deepcopy
 from function_wrapper import TYPE_FORMAL_GENERIC
 import common_with_cwrap
+import yaml
 
 type_map = {
     'floating_point': [
@@ -54,7 +55,7 @@ def expand(pair):
 
 
 def exclude(declaration):
-    return 'only_register' in declaration
+    return 'only_register' in declaration or declaration.get('python_name') == 'ndimension'
 
 
 def add_variants(option):
@@ -109,6 +110,43 @@ def sanitize_return(option):
 def set_mode(option):
     option['mode'] = option.get('mode', 'TH')
 
+# To enable 0-dim support in TH operations
+# we find all places where a single Scalar replaced with a Tensor
+# as an argument is still a valid function
+# we then mark the tensor variant with a key zero_dim_dispatch_when_scalar: name
+# where 'name' is the name of the argument that should be a scalar
+# during dispatch, if that argument is marked internally as holding a scalar
+# then the method will dispatch to that function.
+def discover_zero_dim_tensor_operations(declaration):
+    def exclude(arg):
+        return arg.get('ignore_check')
+
+    def signature(option,i=None,value=None):
+        elements = [TYPE_FORMAL_GENERIC.get(arg['type'],arg['type'])
+                    if i is None or j != i else value
+                    for j, arg in enumerate(option['arguments'])
+                    if not exclude(arg) ]
+        return '#'.join(elements)
+    signature_to_option = {signature(option): option
+                           for option in declaration['options']}
+
+    for option in declaration['options']:
+        for i,arg in enumerate(option['arguments']):
+            if arg['type'] == 'real':
+                signature_of_tensor_version = signature(option,i,'Tensor &')
+                if signature_of_tensor_version in signature_to_option:
+                    tensor_version = \
+                        signature_to_option[signature_of_tensor_version]
+                    names = [arg['name'] for arg in tensor_version['arguments']
+                             if not exclude(arg)]
+                    tensor_version['zero_dim_dispatch_when_scalar'] = names[i]
+                    print("FOUND "+str(i))
+                    print("Scalar Version ===== ")
+                    print(yaml.dump(option))
+                    print("Tensor Version ===== ")
+                    print(yaml.dump(tensor_version))
+                    print("SHARED "+names[i])
+
 
 def run(declarations):
     declarations = [d for d in declarations if not exclude(d)]
@@ -120,6 +158,8 @@ def run(declarations):
             type_to_signature=TYPE_FORMAL_GENERIC,
             remove_self=True)
         common_with_cwrap.sort_by_number_of_options(declaration)
+        discover_zero_dim_tensor_operations(declaration)
+
         new_options = []
         for option in declaration['options']:
             set_mode(option)

templates/Tensor.h

@@ -76,6 +76,15 @@ struct Tensor {
   IntList strides() const {
     return pImpl->strides();
   }
+  int64_t dim() const {
+    return pImpl->dim();
+  }
+  int64_t ndimension() const {
+    return dim();
+  }
+  Scalar scalar() const {
+    return pImpl->localScalar();
+  }
   Type & type() const {
     return pImpl->type();
   }
@@ -95,9 +104,7 @@ struct Tensor {
   Tensor toBackend(Backend b) {
     return toType(type().toBackend(b));
   }
-  int64_t dim() const {
-    return ndimension();
-  }
+  
   template<typename T>
   T * data() const;
 

templates/TensorDerived.cpp

@@ -54,11 +54,24 @@ const char * ${Tensor}::toString() const {
 }
 
 IntList ${Tensor}::sizes() {
+  if(isScalar())
+    return IntList();
   return IntList(reinterpret_cast<int64_t*>(tensor->size),tensor->nDimension);
 }
 IntList ${Tensor}::strides() {
+  if(isScalar())
+    return IntList();
   return IntList(reinterpret_cast<int64_t*>(tensor->stride),tensor->nDimension);
 }
+int64_t ${Tensor}::dim() {
+  if(isScalar())
+    return 0;
+  return ${THTensor}_nDimension(${state,}tensor);
+}
+Scalar ${Tensor}::localScalar() {
+  AT_ASSERT(isScalar(),"localScalar() called on Tensor with %d dims",sizes().size());
+  return Scalar(${to_at_half}(${THTensor}_get1d(${state,}tensor, 0)));
+}
 
 const char * ${Tensor}::typeString() {
   return "${Type}";

templates/TensorDerived.h

@@ -16,6 +16,8 @@ struct ${Tensor} : public TensorImpl {
   virtual const char * toString() const override;
   virtual IntList sizes() override;
   virtual IntList strides() override;
+  virtual int64_t dim() override;
+  virtual Scalar localScalar() override;
   static const char * typeString();
 
 //TODO(zach): sort of friend permissions later so this

test/basic.cpp

@@ -35,6 +35,7 @@ static void test(Type & type) {
   {
     std::cout << "sort:" << std::endl;
     Tensor b = type.rand({3, 4});
+
     std::cout << b << std::endl;
     auto z = b.sort(1);
     std::cout << std::get<0>(z) << std::endl;
@@ -167,6 +168,18 @@ static void test(Type & type) {
     //std::cout << select(select(a, 1, 3), 0, 2) << std::endl;
   }
 
+  {
+      std::cout << "zero-dim: " << std::endl;
+      Tensor a = type.rand({1});
+      // TODO
+      a.pImpl->setScalar(true);
+      std::cout << a << "dims: " << a.dim() << std::endl;
+      std::cout << a.scalar() << std::endl;
+      Tensor b = type.rand({3,4});
+      std::cout << b + a << std::endl;
+      std::cout << a + b << std::endl;
+  }
+
 }
 
 int main()