From 7b5464ed7bdd4eda5ea33977c008538b08da9025 Mon Sep 17 00:00:00 2001 From: ashari4 <70242157+ashari4@users.noreply.github.com> Date: Wed, 5 Jan 2022 09:33:28 -0800 Subject: [PATCH] aten add_ op supports bf16 (#10084) * hand implemented add_ --- .../orttraining/eager/opgen/opgen/atenops.py | 50 ++++++------ orttraining/orttraining/eager/ort_aten.cpp | 77 ++++++++++++++++++- orttraining/orttraining/eager/ort_aten.h | 5 ++ orttraining/orttraining/eager/ort_ops.cpp | 5 +- orttraining/orttraining/eager/test/ort_ops.py | 1 + 5 files changed, 112 insertions(+), 26 deletions(-) diff --git a/orttraining/orttraining/eager/opgen/opgen/atenops.py b/orttraining/orttraining/eager/opgen/opgen/atenops.py index f8cda84df4..d3fd008ef6 100644 --- a/orttraining/orttraining/eager/opgen/opgen/atenops.py +++ b/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} diff --git a/orttraining/orttraining/eager/ort_aten.cpp b/orttraining/orttraining/eager/ort_aten.cpp index d62e783805..d11db28333 100644 --- a/orttraining/orttraining/eager/ort_aten.cpp +++ b/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(); - CopyVectorToTensor(invoker, &val, 1, *ort_tensor); + switch (type) { + case at::ScalarType::Float: + CopyVectorToTensor(invoker, &val, 1, *ort_tensor); + break; + case at::ScalarType::BFloat16: { + at::BFloat16 valBFloat16 = scalar.toBFloat16(); + Ort::BFloat16_t *valOrtBFloat16 = reinterpret_cast(&valBFloat16); + CopyVectorToTensor(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 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 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 diff --git a/orttraining/orttraining/eager/ort_aten.h b/orttraining/orttraining/eager/ort_aten.h index fc560914e2..58bdf626d3 100644 --- a/orttraining/orttraining/eager/ort_aten.h +++ b/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); diff --git a/orttraining/orttraining/eager/ort_ops.cpp b/orttraining/orttraining/eager/ort_ops.cpp index 9b5f2e8f50..473b92c315 100644 --- a/orttraining/orttraining/eager/ort_ops.cpp +++ b/orttraining/orttraining/eager/ort_ops.cpp @@ -30,8 +30,11 @@ void createInplaceOutputValue(OrtValue& input, V shape, OrtValue* p_mlv element_type, new_shape, p_mlvalue); } +template +using Vector = std::vector>; + template <> -void createInplaceOutputValue(OrtValue& input, std::vector shape, OrtValue* p_mlvalue){ +void createInplaceOutputValue(OrtValue& input, Vector shape, OrtValue* p_mlvalue){ auto* input_ort_tensor = input.GetMutable(); auto element_type = onnxruntime::DataTypeImpl::GetType(); CreateMLValue(input_ort_tensor->MutableDataRaw(), diff --git a/orttraining/orttraining/eager/test/ort_ops.py b/orttraining/orttraining/eager/test/ort_ops.py index f8f27e163a..84be88d73a 100644 --- a/orttraining/orttraining/eager/test/ort_ops.py +++ b/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]])