aten add_ op supports bf16 (#10084)

* hand implemented add_

orttraining/orttraining/eager/opgen/opgen/atenops.py

@@ -25,8 +25,32 @@ class GeluGrad(ONNXOp):
     super().__init__('GeluGrad', 1, [{'at::kHalf', 'at::kFloat', 'at::kBFloat16'}, {'at::kHalf', 'at::kFloat', 'at::kBFloat16'}], dY, X)
     self.domain = kMSDomain
 
-ops = {
-  # Hand-Implemented Ops
+ops = {}
+
+for binary_op, onnx_op in {
+  'add': Add('self', Mul('alpha', 'other')),
+  'sub': Sub('self', Mul('alpha', 'other')),
+  'mul': Mul('self', 'other'),
+  'div': Div('self', 'other')}.items():
+  for dtype in ['Tensor', 'Scalar']:
+    for variant in ['', '_']:
+      name = f'aten::{binary_op}{variant}.{dtype}'
+      if name not in ops:
+        ops[f'aten::{binary_op}{variant}.{dtype}'] = deepcopy(onnx_op)
+
+for unary_op in [
+  'abs','acos','acosh', 'asinh', 'atanh', 'asin', 'atan', 'ceil', 'cos',
+  'cosh', 'erf', 'exp', 'floor', 'isnan', 'log', 'reciprocal', 'neg', 'round',
+  'relu', 'selu', 'sigmoid', 'sin', 'sinh', 'sqrt', 'tan', 'tanh', 'nonzero',
+  'sign', 'hardsigmoid', 'isinf', 'det']:
+  aten_name = f'aten::{unary_op}'
+  onnx_op = onnx_ops[unary_op]('self')
+  ops[aten_name] = onnx_op
+  # produce the in-place variant as well for ops that support it
+  if unary_op not in ['isnan', 'nonzero', 'min', 'max', 'isinf', 'det']:
+    ops[f'{aten_name}_'] = onnx_op
+
+hand_implemented = {
   'aten::empty.memory_format': SignatureOnly(),
   'aten::empty_strided': SignatureOnly(),
   'aten::zero_': SignatureOnly(),
@@ -35,6 +59,7 @@ ops = {
   'aten::view': SignatureOnly(),
   'aten::_copy_from_and_resize' : SignatureOnly(),
   'aten::addmm': Gemm('mat1', 'mat2', 'self', alpha='alpha', beta='beta'),
+  'aten::add_.Tensor': SignatureOnly(),
   'aten::t': Transpose('self'),
   'aten::mm': MatMul('self', 'mat2'),
   'aten::zeros_like': ConstantOfShape(Shape('self')), #the default constant is 0, so don't need to speicify attribute
@@ -65,23 +90,4 @@ ops = {
   'aten::gt.Scalar_out' : MakeTorchFallback(),
 }
 
-for binary_op, onnx_op in {
-  'add': Add('self', Mul('alpha', 'other')),
-  'sub': Sub('self', Mul('alpha', 'other')),
-  'mul': Mul('self', 'other'),
-  'div': Div('self', 'other')}.items():
-  for dtype in ['Tensor', 'Scalar']:
-    for variant in ['', '_']:
-      ops[f'aten::{binary_op}{variant}.{dtype}'] = deepcopy(onnx_op)
-
-for unary_op in [
-  'abs','acos','acosh', 'asinh', 'atanh', 'asin', 'atan', 'ceil', 'cos',
-  'cosh', 'erf', 'exp', 'floor', 'isnan', 'log', 'reciprocal', 'neg', 'round',
-  'relu', 'selu', 'sigmoid', 'sin', 'sinh', 'sqrt', 'tan', 'tanh', 'nonzero',
-  'sign', 'hardsigmoid', 'isinf', 'det']:
-  aten_name = f'aten::{unary_op}'
-  onnx_op = onnx_ops[unary_op]('self')
-  ops[aten_name] = onnx_op
-  # produce the in-place variant as well for ops that support it
-  if unary_op not in ['isnan', 'nonzero', 'min', 'max', 'isinf', 'det']:
-    ops[f'{aten_name}_'] = onnx_op
+ops = {**ops, **hand_implemented} 

orttraining/orttraining/eager/ort_aten.cpp

@@ -80,16 +80,38 @@ onnxruntime::MLDataType ort_scalar_type_from_aten(
 OrtValue create_ort_value(
   onnxruntime::ORTInvoker& invoker,
   const at::Scalar& scalar) {
-  // TODO: support more types
+  return create_ort_value(invoker, scalar, at::kFloat);
+}
+
+OrtValue create_ort_value(
+  onnxruntime::ORTInvoker& invoker,
+  const at::Scalar& scalar,
+  at::ScalarType type) {
   float val = scalar.toFloat();
   OrtValue ort_val;
   CreateMLValue(
     invoker.GetCurrentExecutionProvider().GetAllocator(0, OrtMemTypeDefault),
-    ort_scalar_type_from_aten(at::kFloat),
+    ort_scalar_type_from_aten(type),
     {},
     &ort_val);
   auto* ort_tensor = ort_val.GetMutable<onnxruntime::Tensor>();
-  CopyVectorToTensor<float>(invoker, &val, 1, *ort_tensor);
+  switch (type) {
+    case at::ScalarType::Float:
+      CopyVectorToTensor<float>(invoker, &val, 1, *ort_tensor);
+      break;
+    case at::ScalarType::BFloat16: {
+      at::BFloat16 valBFloat16 = scalar.toBFloat16();
+      Ort::BFloat16_t *valOrtBFloat16 = reinterpret_cast<Ort::BFloat16_t *>(&valBFloat16);
+      CopyVectorToTensor<Ort::BFloat16_t>(invoker, valOrtBFloat16, 1, *ort_tensor);
+      break;
+    }      
+    default:
+      // TODO: support more types
+      // For most at::ScalarType, it should be safe to just call value.to<>
+      // on it, but for now we want to explicitly know when we've encountered
+      // a new scalar type while bringing up ORT eager mode.
+      ORT_THROW("Unsupported: at::ScalarType::", scalar.type());
+  }
   return ort_val;
 }
 
@@ -387,6 +409,55 @@ at::Tensor& zero_(at::Tensor& self){
   return self;
 }
 
+// TODO: enhance opgen.py to support inplace binary operations.
+// aten::add_.Tensor(Tensor(a!) self, Tensor other, *, Scalar alpha=1) -> Tensor(a!)
+at::Tensor& add__Tensor(
+  at::Tensor& self, 
+  const at::Tensor& other, 
+  const at::Scalar& alpha) {
+  ORT_LOG_FN(self, other, alpha);
+  
+  if (
+    !IsSupportedType(alpha, {at::kDouble,at::kLong,at::kHalf,at::kShort,at::kInt,at::kByte,at::kFloat,at::kBFloat16}) || 
+    !IsSupportedType(other, {at::kDouble,at::kLong,at::kHalf,at::kShort,at::kInt,at::kByte,at::kFloat,at::kBFloat16}) || 
+    !IsSupportedType(self, {at::kDouble,at::kLong,at::kHalf,at::kShort,at::kInt,at::kByte,at::kFloat,at::kBFloat16})) {
+    return at::native::call_fallback_fn<
+      &at::native::cpu_fallback,
+      ATEN_OP(add__Tensor)>::call(self, other, alpha);
+  }
+  auto& invoker = GetORTInvoker(self.device());
+  
+  auto ort_input_alpha = create_ort_value(invoker, alpha, other.scalar_type());
+  auto ort_input_other = create_ort_value(invoker, other);
+  
+  std::vector<OrtValue> ort_outputs_0_Mul(1);
+  
+  auto status = invoker.Invoke("Mul", {
+    std::move(ort_input_alpha),
+    std::move(ort_input_other),
+  }, ort_outputs_0_Mul, nullptr);
+  
+  if (!status.IsOK())
+    throw std::runtime_error(
+      "ORT return failure status:" + status.ErrorMessage());
+  
+  auto ort_input_self = create_ort_value(invoker, self);
+  
+  std::vector<OrtValue> ort_outputs_1_Add(1);
+  ort_outputs_1_Add[0] = ort_input_self;
+  
+  status = invoker.Invoke("Add", {
+    std::move(ort_input_self),
+    std::move(ort_outputs_0_Mul[0]),
+  }, ort_outputs_1_Add, nullptr);
+  
+  if (!status.IsOK())
+    throw std::runtime_error(
+      "ORT return failure status:" + status.ErrorMessage());
+  
+  return self;
+}
+
 } // namespace aten
 
 //#pragma endregion

orttraining/orttraining/eager/ort_aten.h

@@ -28,6 +28,11 @@ OrtValue create_ort_value(
   onnxruntime::ORTInvoker& invoker,
   const at::Scalar& scalar);
 
+OrtValue create_ort_value(
+  onnxruntime::ORTInvoker& invoker,
+  const at::Scalar& scalar,
+  at::ScalarType type);
+
 OrtValue create_ort_value(
   onnxruntime::ORTInvoker& invoker,
   const at::Tensor& tensor);

orttraining/orttraining/eager/ort_ops.cpp

@@ -30,8 +30,11 @@ void createInplaceOutputValue(OrtValue& input, V<int64_t> shape, OrtValue* p_mlv
                 element_type, new_shape, p_mlvalue);
 }
 
+template <typename T> 
+using Vector = std::vector<T, std::allocator<T>>;
+
 template <>
-void createInplaceOutputValue<std::vector>(OrtValue& input, std::vector<int64_t> shape, OrtValue* p_mlvalue){
+void createInplaceOutputValue<Vector>(OrtValue& input, Vector<int64_t> shape, OrtValue* p_mlvalue){
   auto* input_ort_tensor = input.GetMutable<onnxruntime::Tensor>();
   auto element_type = onnxruntime::DataTypeImpl::GetType<int64_t>();
   CreateMLValue(input_ort_tensor->MutableDataRaw(),

orttraining/orttraining/eager/test/ort_ops.py

@@ -34,6 +34,7 @@ class OrtOpTests(unittest.TestCase):
       torch.mul(cpu_ones, cpu_ones),
       torch.mul(ort_ones, ort_ones).cpu())
 
+  # TODO: Add BFloat16 test coverage
   def test_add_(self):
     device = self.get_device()
     cpu_ones = torch.Tensor([[1, 1, 1], [1, 1, 1], [1, 1, 1]])