[ATen][Scalars] Remove Scalar from return types of functions. (pytorch#3557)

* Add direct C-type scalar conversions from Tensor, e.g. toCFloat() as an alias for Scalar(x).toFloat()

* Provide tensor overloads for fill_, masked_fill_, index_fill_.

* Everythign up to scalar overload.

* Fix pytorch build for aten scalar return type changes.

* Use valid expression instead of dangling else.

* Simplify code generation.

* Fix test_jit (why didn't this compile locally?)

Author: gchanan

aten/contrib/meter/APMeter.cc

@@ -82,7 +82,7 @@ void APMeter::value(Tensor& val) {
          if(targetbuffer_d[n] != 0.)
             val_d[k] += precision_d[n];
       }
-      auto norm = sum(targetbuffer).toDouble();
+      auto norm = sum(targetbuffer).toCDouble();
       if(norm > 0)
         val_d[k] /= norm;
    }

aten/src/ATen/Declarations.cwrap

@@ -126,11 +126,16 @@
 ]]
 [[
   name: fill_
-  cname: fill
   return: self
-  arguments:
-    - THTensor* self
-    - real value
+  cname: fill
+  options:
+    - arguments:
+      - THTensor* self
+      - real value
+    - zero_dim_tensor_only: True
+      arguments:
+      - THTensor* self
+      - THTensor* value
 ]]
 [[
   name: isContiguous
@@ -156,11 +161,18 @@
   cname: maskedFill
   python_name: masked_fill_
   return: self
-  arguments:
-    - arg: THTensor* self
-      broadcast: mask inplace fallback types:Byte
-    - THBoolTensor* mask
-    - real value
+  options:
+    - arguments:
+      - arg: THTensor* self
+        broadcast: mask inplace fallback types:Byte
+      - THBoolTensor* mask
+      - real value
+    - zero_dim_tensor_only: True
+      arguments:
+      - arg: THTensor* self
+        broadcast: mask inplace fallback types:Byte
+      - THBoolTensor* mask
+      - THTensor* value
 ]]
 [[
   name: maskedCopy_
@@ -364,12 +376,20 @@
   python_name: index_fill_
   cname: indexFill
   return: argument 0
-  arguments:
-    - THTensor* self
-    - arg: long dim
-      wrap_dim: self
-    - THIndexTensor* index
-    - real value
+  options:
+    - arguments:
+      - THTensor* self
+      - arg: long dim
+        wrap_dim: self
+      - THIndexTensor* index
+      - real value
+    - zero_dim_tensor_only: True
+      arguments:
+        - THTensor* self
+        - arg: long dim
+          wrap_dim: self
+        - THIndexTensor* index
+        - THTensor* value
 ]]
 [[
   name: narrow

aten/src/ATen/function_wrapper.py

@@ -74,10 +74,15 @@
 # implementation of ${api_name} if we have overloaded a function with
 # the same name (but different signature) already
 ZERO_DIM_CHECK = CodeTemplate("""\
-if(${check_name}.dim() == 0) {
+if (${check_name}.dim() == 0) {
     return static_cast<const Type*>(this)->${method_prefix}${api_name}(${zero_dim_actuals});
 }""")
 
+ZERO_DIM_ONLY = CodeTemplate("""\
+runtime_error("${api_name} only supports a 0-dimensional ${check_name} tensor, but got tensor "
+    "with %" PRId64 " dimension(s)", ${check_name}.dim());
+""")
+
 SPARSE_CHECK = CodeTemplate("""\
 if(${check_name}.type().isSparse()) {
     return static_cast<const Type*>(this)->${method_prefix}${api_name}(${sparse_actuals});
@@ -136,8 +141,8 @@ def __init__(self, reason):
     'THIndexTensor*': 'Tensor',
     'THBoolTensor*': 'Tensor',
     'THIntegerTensor*': 'Tensor',
-    'real': 'Scalar',
-    'accreal': 'Scalar',
+    'real': 'Tensor',
+    'accreal': 'Tensor',
     'long': 'int64_t',
 }
 
@@ -710,14 +715,24 @@ def is_actual_return_long(ret):
             return backend_type_env['AccScalarName'] == 'Long'
         return False
 
+    def get_zero_dim_dispatch_when_scalar(option):
+        return option.get('zero_dim_dispatch_when_scalar', False)
+
     def handle_zero_dim(env, option):
-        if 'zero_dim_dispatch_when_scalar' not in option:
+        zero_dim_dispatch = get_zero_dim_dispatch_when_scalar(option)
+        if not zero_dim_dispatch:
             return []
-        check_name = option['zero_dim_dispatch_when_scalar']
         zero_dim_actuals = [arg['name']
-                            if arg['name'] != check_name else "Scalar({})".format(arg['name'])
+                            if arg['name'] != zero_dim_dispatch else "Scalar({})".format(arg['name'])
                             for arg in option['formals_list']]
-        return [ZERO_DIM_CHECK.substitute(env, check_name=check_name, zero_dim_actuals=zero_dim_actuals)]
+        return [ZERO_DIM_CHECK.substitute(env, check_name=zero_dim_dispatch, zero_dim_actuals=zero_dim_actuals)]
+
+    def handle_only_zero_dim(env, option):
+        if option.get('zero_dim_tensor_only', False):
+            check_name = get_zero_dim_dispatch_when_scalar(option)
+            return [ZERO_DIM_ONLY.substitute(env, check_name=check_name)]
+        else:
+            return None
 
     def handle_sparse(env, option):
         if 'when_sparse_dispatch' not in option or 'Sparse' in backend_type_env['Backend']:
@@ -781,6 +796,12 @@ def emit_body(env, option):
         body = []
         body += handle_sparse(env, option)
         body += handle_zero_dim(env, option)
+        only_zero_dim_check = handle_only_zero_dim(env, option)
+        if only_zero_dim_check is not None:
+            #  code below only_zero_dim_check is unreachable so we do not need to generate the rest.
+            body += only_zero_dim_check
+            return body
+
         body += handle_buffers(env, option)
         # arguments are potentially duplicated because of one argument
         # referencing another
@@ -933,6 +954,10 @@ def emit_body(env, option):
                 return_tensor = "return Tensor((new ${Tensor}(context,${arg_name}))${maybe_scalar},false);"
                 body.append(CodeTemplate(return_tensor).substitute(
                     env, arg_name=call, maybe_scalar=maybe_scalar))
+            # return the same underlying Tensor type for both real and accreal; this ensures
+            # e.g. x.sum(0) and x.sum() return the same type.
+            elif ret['type'] == 'accreal' or ret['type'] == 'real':
+                body.append('return scalarTensor({});'.format(call))
             else:
                 # we using int64_t for long in the API, so correct it here...
                 if is_actual_return_long(ret):

aten/src/ATen/templates/Tensor.h

@@ -108,6 +108,11 @@ struct Tensor : public detail::TensorBase {
   AT_FORALL_SCALAR_TYPES(TO_TYPE_DATA)
   #undef TO_TYPE_DATA
 
+  #define TO_C_TYPE(T,name,_) \
+  T toC##name () const;
+  AT_FORALL_SCALAR_TYPES(TO_C_TYPE)
+  #undef TO_C_TYPE
+
   template<typename T, size_t N>
   TensorAccessor<T,N> accessor() {
     static_assert(N > 0, "accessor is used for indexing tensor, for scalars use *data<T>()");

aten/src/ATen/templates/TensorMethods.h

@@ -46,4 +46,10 @@ inline T* Tensor::to##name##Data() const { return data<T>(); }
 AT_FORALL_SCALAR_TYPES(DEFINE_CAST)
 #undef DEFINE_CAST
 
+#define DEFINE_TO_C_TYPE(T,name,_) \
+inline T Tensor::toC##name () const { return Scalar(*this).to##name (); }
+
+AT_FORALL_SCALAR_TYPES(DEFINE_TO_C_TYPE)
+#undef DEFINE_TO_C_TYPE
+
 } //namespace at

aten/src/ATen/test/basic.cpp

@@ -33,11 +33,11 @@ static void test(Type & type) {
     std::cout << "ones and dot:" << std::endl;
     Tensor b = type.ones({3, 4});
     std::cout << b << std::endl;
-    ASSERT(24 == (b+b).sum().toDouble());
+    ASSERT(24 == (b+b).sum().toCDouble());
     std::cout << b.numel() << std::endl;
     ASSERT(12 == b.numel());
     std::cout << b.dot(b) << std::endl;
-    ASSERT(b.dot(b).toDouble() == 12);
+    ASSERT(b.dot(b).toCDouble() == 12);
   }
 
   {
@@ -97,8 +97,8 @@ static void test(Type & type) {
     }
     auto end = std::chrono::high_resolution_clock::now();
     std::cout << std::dec << "   " << std::chrono::duration_cast<std::chrono::milliseconds>(end-begin).count() << " ms" << std::endl;
-    ASSERT(norm(100000*d).toDouble() == norm(r).toDouble());
-    std::cout << "   norm: " << norm(r).toDouble() << std::endl;
+    ASSERT(norm(100000*d).toCDouble() == norm(r).toCDouble());
+    std::cout << "   norm: " << norm(r).toCDouble() << std::endl;
   }
 
   {
@@ -111,8 +111,8 @@ static void test(Type & type) {
     }
     auto end = std::chrono::high_resolution_clock::now();
     std::cout << std::dec << "   " << std::chrono::duration_cast<std::chrono::milliseconds>(end-begin).count() << " ms" << std::endl;
-    ASSERT(norm(100000*d).toDouble() == norm(r).toDouble());
-    std::cout << "   norm: " << norm(r).toDouble() << std::endl;
+    ASSERT(norm(100000*d).toCDouble() == norm(r).toCDouble());
+    std::cout << "   norm: " << norm(r).toCDouble() << std::endl;
   }
 
   {
@@ -247,7 +247,7 @@ static void test(Type & type) {
       std::cout << c << std::endl;
 
       Tensor e = CPU(kFloat).rand({});
-      ASSERT(*e.data<float>()== e.sum().toFloat());
+      ASSERT(*e.data<float>()== e.sum().toCFloat());
   }
   {
     Tensor b = CPU(kFloat).ones({3,7})*.0000001f;

aten/src/ATen/test/scalar_tensor_test.cpp

@@ -172,6 +172,17 @@ int main() {
         ASSERT(false);
       } catch (std::runtime_error &e) {}
     }
+
+    // fill_
+    if (t.dim() > 0 && t.numel() != 0) {
+      try {
+        // can only fill_ 0-dim tensors
+        t.fill_(t.sum(0));
+        assert(t.dim() == 1);
+      } catch (std::runtime_error &e) {
+        assert(t.dim() != 1);
+      }
+    }
   }
 
   for (auto lhs_it = sizes.begin(); lhs_it != sizes.end(); ++lhs_it) {

aten/src/ATen/test/scalar_test.cpp

@@ -105,7 +105,7 @@ int main() {
   auto t = CPU(Float).ones({4,4});
 
   auto wha2 = CPU(Float).zeros({4,4}).add(t).sum();
-  cout << wha2.toDouble() << " <-ndim\n";
+  cout << wha2.toCDouble() << " <-ndim\n";
 
   cout << t.sizes() << " " << t.strides() << "\n";
 
@@ -143,6 +143,18 @@ int main() {
   ASSERT(Scalar(CPU(kFloat).ones({})).toTensor().type().scalarType() == kFloat);
 
   dispatch<Foo>(x.type(),x,prev_h);
+
+  // test direct C-scalar type conversions
+  try {
+    auto x = T.ones({1,2});
+    x.toCFloat();
+    ASSERT(false);
+  } catch (std::runtime_error &e) {}
+  auto float_one = T.ones({});
+  ASSERT(float_one.toCFloat() == 1);
+  ASSERT(float_one.toCInt() == 1);
+  ASSERT(float_one.toCHalf() == 1);
+
   return 0;
 
 }

tools/autograd/templates/Functions.cpp

@@ -76,7 +76,7 @@ Tensor norm_backward(Tensor grad, const Tensor & self, const Scalar & p_, Tensor
 
 Tensor reduce_to(const Tensor & grad, IntList sizes) {
   if (sizes.size() == 0) {
-    return grad.sum().toTensor();
+    return grad.sum();
   }
   Tensor result = grad;
   while (result.dim() > (int64_t)sizes.size()) {
@@ -306,9 +306,9 @@ Tensor glu_double_backward_grad_output(const Tensor & grad, const Tensor & input
 Tensor kl_div_double_backward_grad_output(const Tensor & grad, const Tensor & input, const Tensor & target, bool size_average, bool reduce) {
   auto result = kl_div_backward(grad, input, target, size_average, false);
   if (reduce && size_average) {
-    return result.mean().toTensor();
+    return result.mean();
   } else if (reduce) {
-    return result.sum().toTensor();
+    return result.sum();
   }
   return result;
 }
@@ -343,9 +343,9 @@ Tensor log_softmax_double_backward(const Tensor & grad, const Tensor & grad_outp
 Tensor l1_loss_double_backward_grad_output(const Tensor & grad, const Tensor & input, const Tensor & target, bool size_average, bool reduce) {
   auto output = l1_loss_backward(grad, input, target, size_average, false);
   if (reduce and size_average) {
-    return output.mean().toTensor();
+    return output.mean();
   } else if (reduce) {
-    return output.sum().toTensor();
+    return output.sum();
   }
   return output;
 }
@@ -364,7 +364,7 @@ Tensor smooth_l1_loss_double_backward_grad_output(const Tensor & grad, const Ten
     return smooth_l1_loss_backward(grad, input, target, size_average, reduce);
   }
   auto r = smooth_l1_loss_backward(ones_like(grad_output), input, target, size_average, true);
-  return (r * grad).sum().toTensor().view({1});
+  return (r * grad).sum().view({1});
 }
 
 Tensor max_pool2d_double_backward(const Tensor & grad, const Tensor & indices) {
@@ -389,7 +389,7 @@ Tensor mse_loss_double_backward_grad_output(const Tensor & grad, const Tensor &
     return mse_loss_backward(grad, input, target, size_average, reduce);
   }
   auto r = mse_loss_backward(ones_like(grad_output), input, target, size_average, true);
-  return (r * grad).sum().toTensor().view({1});
+  return (r * grad).sum().view({1});
 }
 
 Tensor soft_margin_loss_double_backward(const Tensor & grad, const Tensor & input, const Tensor & target, bool size_average) {

torch/csrc/jit/test_jit.cpp

@@ -98,7 +98,7 @@ static void fusionTests() {
     auto o = at::CUDA(at::kFloat).zeros({3,4});
     comp.debugLaunchGraph(graph, {a,b}, {o});
     auto o2 = a*b;
-    float max_diff = (o2 - o).abs().max().toDouble();
+    float max_diff = (o2 - o).abs().max().toCDouble();
     //std::cout << "max diff: " << max_diff << "\n";
     JIT_ASSERT(max_diff == 0);
   };
@@ -160,7 +160,7 @@ static void fusionTests() {
     //auto out0 = inputs[0]*inputs[1];
     comp.debugLaunchGraph(graph, inputs, outputs);
     JIT_ASSERT(out0.is_same_size(outputs.front()));
-    float max_diff = (outputs.front() - out0).abs().max().toDouble();
+    float max_diff = (outputs.front() - out0).abs().max().toCDouble();
     JIT_ASSERT(max_diff < 1e-6);
 
   };
@@ -191,9 +191,9 @@ static void fusionTests() {
     auto o2 = at::CUDA(at::kFloat).zeros(o2_r.sizes());
     comp.debugLaunchGraph(graph, {a,b}, {o, o2});
 
-    float max_diff = (o_r - o).abs().max().toDouble();
+    float max_diff = (o_r - o).abs().max().toCDouble();
     JIT_ASSERT(max_diff == 0);
-    float max_diff2 = (o2_r - o2).abs().max().toDouble();
+    float max_diff2 = (o2_r - o2).abs().max().toCDouble();
     JIT_ASSERT(max_diff2 == 0);
   };
   testConcat(0);