From de57daaab055766f6a5fa0c4ef34318fe611174b Mon Sep 17 00:00:00 2001 From: Wil Brady <25513670+WilBrady@users.noreply.github.com> Date: Tue, 26 Jul 2022 09:14:57 -0400 Subject: [PATCH] Eager mode: binary ops more complete behavior and testing. (#12293) * Remove hand written add_.Tensor as it can now be generated. * Generate .out for tensor version of basic math ops. Add.out testing added too. * Remove sin tests as they are covered by parameterized tests. Also, moved all parameterized tests to the end in their own section. * Add binary ops tests for tensors. Scalar tests are calling the aten .out which is for tensor. * Add support for scalar input to add, div, mul, and sub. --- .../orttraining/eager/opgen/opgen/atenops.py | 28 +- .../eager/opgen/opgen/generator.py | 139 +++-- orttraining/orttraining/eager/ort_aten.cpp | 96 +--- orttraining/orttraining/eager/ort_aten.h | 1 + orttraining/orttraining/eager/test/ort_ops.py | 484 ++++++++++-------- 5 files changed, 403 insertions(+), 345 deletions(-) diff --git a/orttraining/orttraining/eager/opgen/opgen/atenops.py b/orttraining/orttraining/eager/opgen/opgen/atenops.py index 606b677ade..1215f63a35 100644 --- a/orttraining/orttraining/eager/opgen/opgen/atenops.py +++ b/orttraining/orttraining/eager/opgen/opgen/atenops.py @@ -93,19 +93,6 @@ unary_ops = [ "selu", ] -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) - type_promotion_ops.append(f"aten::{binary_op}{variant}.{dtype}") - for unary_op in unary_ops_with_out: ops[f"aten::{unary_op}.out"] = onnx_ops[unary_op]("self") @@ -115,6 +102,20 @@ for unary_op in unary_ops_with_inplace: for unary_op in unary_ops: ops[f"aten::{unary_op}"] = onnx_ops[unary_op]("self") +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 Tensor, binary_op.out is used by both binary_op and binary_op_, so we only generate .out + # from testing and call stacks, it also apears scalar ops fall back to the (Tensor) binary_op.out, + # so this is all we need. + name = f"aten::{binary_op}.out" + if name not in ops: + ops[f"aten::{binary_op}.out"] = deepcopy(onnx_op) + type_promotion_ops.append(f"aten::{binary_op}.out") + # Notes on Onnx op mapping # # Equal - Onnx spec has the return as a bool tensor, but aten will keep the tensor @@ -136,7 +137,6 @@ hand_implemented = { # manually implement Slice using stride and offset. "aten::slice.Tensor": SignatureOnly(), "aten::addmm": Gemm("mat1", "mat2", "self", alpha="alpha", beta="beta"), - "aten::add_.Tensor": SignatureOnly(), "aten::t": Transpose("self"), # MatMul("self", "mat2"), fails since it resizes based on self but should be based on result shape of the mult "aten::mm.out": SignatureOnly(), diff --git a/orttraining/orttraining/eager/opgen/opgen/generator.py b/orttraining/orttraining/eager/opgen/opgen/generator.py index 9b1d37f309..bea5494430 100644 --- a/orttraining/orttraining/eager/opgen/opgen/generator.py +++ b/orttraining/orttraining/eager/opgen/opgen/generator.py @@ -1,6 +1,8 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. +# pylint: disable=missing-docstring, too-many-public-methods, too-many-nested-blocks + import json import sys from typing import Dict, List, Optional, Union @@ -264,53 +266,7 @@ class ORTGen: writer.write(first_param.identifier.value) writer.writeline(".size()>0);") - # generate the type check - need_type_check = False - if not self._custom_ops: - for onnx_op_index, onnx_op in enumerate(ctx.ops): - for op_input in onnx_op.inputs: - if not isinstance(op_input, Outputs): - need_type_check = True - break - if need_type_check: - writer.write("if (") - i = 0 - for onnx_op_index, onnx_op in enumerate(ctx.ops): - for idx, op_input in enumerate(onnx_op.inputs): - if isinstance(op_input, Outputs): - continue - writer.writeline(" || " if i > 0 else "") - if i == 0: - writer.push_indent() - cpp_param = cpp_func.get_parameter(op_input) - supported_types = ",".join(sorted([type for type in onnx_op.input_types[idx]])) - writer.write(f"!IsSupportedType({cpp_param.identifier.value}, {{{supported_types}}})") - i += 1 - writer.writeline(") {") - self._write_cpu_fall_back(writer, mapped_func) - writer.pop_indent() - writer.writeline("}") - - if ( - not isinstance(first_param.parameter_type.desugar(), ast.ConcreteType) - or "Tensor" not in first_param.parameter_type.desugar().identifier_tokens[0].value - ): - raise FunctionGenerationError(cpp_func, "First parameter must be an at::Tensor") - - writer.write("auto& invoker = GetORTInvoker(") - writer.write(first_param.identifier.value) - if first_param.parameter_type.desugar().identifier_tokens[0].value == "TensorList": - writer.write("[0]") - writer.writeline(".device());") - writer.writeline() - - # FIXME: warn if we have not consumed all torch parameters (either as - # an ORT input or ORT attribute). - - # Perform kernel fission on the ATen op to yield a chain of ORT Invokes - # e.g. aten::add(x, y, α) -> onnx::Add(x, onnx::Mul(α, y)) - - # whether need type promotion + # check whether need type promotion, if we do we will use this later to confirm out cast is supported. need_type_promotion = False if mapped_func.mapped_op_name in self.type_promotion_ops: types_from_tensor = [] @@ -338,7 +294,7 @@ class ORTGen: return_info = cpp_func.torch_func.return_type if cpp_func.torch_func else None # if the torch func has a return ref tensor, out is the last param, and self param is the first input - # then we need to update and return out. + # then we need to update and return out. Record this need in set_out_tensor. # TODO: make this more general to handle cases where the first param is not self such as # - cat.out(Tensor[] tensors, int dim=0, *, Tensor(a!) out) -> Tensor(a!) # - complex.out(Tensor real, Tensor imag, *, Tensor(a!) out) -> Tensor(a!) @@ -350,7 +306,6 @@ class ORTGen: and last_param.identifier.value == "out" and first_param.identifier.value == "self" ): - # output_alias is how the return tensor is marked, normally (a! -> a) output_alias = self._get_alias_info(return_info) @@ -359,12 +314,66 @@ class ORTGen: # and the current param (torch_p/last_param) is marked with the output alias if output_alias and output_alias.is_writable and self._get_alias_info(torch_p) == output_alias: set_out_tensor = True - writer.writeline("// resize the output and then create output ort value to be updated.") - writer.writeline( - "resize_output(invoker, dynamic_cast(out.unsafeGetTensorImpl()), self.sizes());" - ) - writer.writeline("auto ort_input_out = create_ort_value(invoker, out);") - writer.writeline() + + # generate the type check + need_type_check = False + cast_op_found = False + if not self._custom_ops: + for onnx_op_index, onnx_op in enumerate(ctx.ops): + for op_input in onnx_op.inputs: + if not isinstance(op_input, Outputs): + need_type_check = True + break + if need_type_check: + writer.write("if (") + i = 0 + for onnx_op_index, onnx_op in enumerate(ctx.ops): + # track is the CAST op was explicitly used + if onnx_op.name == "Cast": + cast_op_found = True + for idx, op_input in enumerate(onnx_op.inputs): + if isinstance(op_input, Outputs): + continue + writer.writeline(" || " if i > 0 else "") + if i == 0: + writer.push_indent() + cpp_param = cpp_func.get_parameter(op_input) + supported_types = ",".join(sorted(list(onnx_op.input_types[idx]))) + writer.write(f"!IsSupportedType({cpp_param.identifier.value}, {{{supported_types}}})") + i += 1 + # if we have type promotion and need to set the out tensor and CAST op not explictily listed, + # then we confirm the promotion type is castable to the out type. + if need_type_promotion and set_out_tensor and not cast_op_found: + writer.writeline(" || ") + writer.write("!c10::canCast(*promoted_type, out.scalar_type())") + writer.writeline(") {") + self._write_cpu_fall_back(writer, mapped_func) + writer.pop_indent() + writer.writeline("}") + + if ( + not isinstance(first_param.parameter_type.desugar(), ast.ConcreteType) + or "Tensor" not in first_param.parameter_type.desugar().identifier_tokens[0].value + ): + raise FunctionGenerationError(cpp_func, "First parameter must be an at::Tensor") + + writer.write("auto& invoker = GetORTInvoker(") + writer.write(first_param.identifier.value) + if first_param.parameter_type.desugar().identifier_tokens[0].value == "TensorList": + writer.write("[0]") + writer.writeline(".device());") + writer.writeline() + + # FIXME: warn if we have not consumed all torch parameters (either as + # an ORT input or ORT attribute). + + if set_out_tensor: + writer.writeline("// resize the output and then create output ort value to be updated.") + writer.writeline( + "resize_output(invoker, dynamic_cast(out.unsafeGetTensorImpl()), self.sizes());" + ) + writer.writeline("auto ort_input_out = create_ort_value(invoker, out);") + writer.writeline() for onnx_op_index, onnx_op in enumerate(ctx.ops): # Torch -> ORT inputs @@ -467,7 +476,16 @@ class ORTGen: # if no in_place_params found and there is an out input to set # and this is the last onnx op, we set the out to be written to if len(in_place_params) == 0 and set_out_tensor and onnx_op_index == (len(ctx.ops) - 1): - writer.writeline(f"{onnx_op.outputs}[0] = ort_input_out;") + # if we have type promotion, need to set the out tensor and CAST op not explictily listed, + # check if we need to do a cast + if need_type_promotion and not cast_op_found: + writer.writeline("if (*promoted_type == out.scalar_type()) {") + writer.push_indent() + writer.writeline(f"{onnx_op.outputs}[0] = ort_input_out;") + writer.pop_indent() + writer.writeline("}") + else: + writer.writeline(f"{onnx_op.outputs}[0] = ort_input_out;") if len(in_place_params) != 0 and len(in_place_params) != ( len(return_info.elements) if isinstance(return_info, ast.TupleType) else 1 @@ -511,6 +529,15 @@ class ORTGen: elif len(in_place_params) == 0: # tensor options if set_out_tensor: + if need_type_promotion and not cast_op_found: + writer.writeline("if (*promoted_type != out.scalar_type()) {") + writer.push_indent() + writer.writeline( + f"CastToType_out(invoker, {onnx_op.outputs}[0], ort_input_out, out.scalar_type());" + ) + writer.pop_indent() + writer.writeline("}") + writer.writeline(f"return {last_param.identifier.value};") return diff --git a/orttraining/orttraining/eager/ort_aten.cpp b/orttraining/orttraining/eager/ort_aten.cpp index 4a18bc61f0..b8d8b87dc4 100644 --- a/orttraining/orttraining/eager/ort_aten.cpp +++ b/orttraining/orttraining/eager/ort_aten.cpp @@ -367,7 +367,7 @@ c10::optional PromoteScalarTypesWithCategory( OrtValue CastToType(onnxruntime::ORTInvoker& invoker, const OrtValue& input, at::ScalarType type) { std::vector output(1); - NodeAttributes attrs(2); + NodeAttributes attrs(1); attrs["to"] = create_ort_attribute( "to", GetONNXTensorProtoDataType(type), at::ScalarType::Long); @@ -375,12 +375,24 @@ OrtValue CastToType(onnxruntime::ORTInvoker& invoker, const OrtValue& input, at: {std::move(input)}, output, &attrs); - if (!status.IsOK()) - throw std::runtime_error( - "ORT return failure status:" + status.ErrorMessage()); + CHECK_STATUS(status); return output[0]; } +void CastToType_out(onnxruntime::ORTInvoker& invoker, const OrtValue& input, OrtValue& output, at::ScalarType type) { + std::vector output_result(1); + output_result[0] = output; + NodeAttributes attrs(1); + attrs["to"] = create_ort_attribute( + "to", GetONNXTensorProtoDataType(type), at::ScalarType::Long); + + auto status = invoker.Invoke("Cast", + {std::move(input)}, + output_result, &attrs); + + CHECK_STATUS(status); +} + /* * Utility method to calculate the resulting shape of tensor after a reduction operation. * @@ -662,9 +674,7 @@ at::Tensor& copy_( {std::move(ort_src)}, ort_cast_output, &attrs); - if (!status.IsOK()) - throw std::runtime_error( - "ORT return failure status:" + status.ErrorMessage()); + CHECK_STATUS(status); copy(invoker, ort_cast_output[0], ort_self); } else { @@ -709,57 +719,7 @@ at::Tensor& zero_(at::Tensor& self) { {std::move(ort_in_self), std::move(flag_val)}, ort_out, nullptr, onnxruntime::kMSDomain, 1); - if (!status.IsOK()) - throw std::runtime_error( - "ORT return failure status:" + status.ErrorMessage()); - - return self; -} - -// TODO(unknown): 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); - - auto st = {at::kDouble, at::kLong, at::kHalf, at::kShort, at::kInt, at::kByte, at::kFloat, at::kBFloat16}; - if ( - !IsSupportedType(alpha, st) || - !IsSupportedType(other, st) || - !IsSupportedType(self, st)) { - 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()); + CHECK_STATUS(status); return self; } @@ -864,9 +824,7 @@ at::Tensor& argmax_out( {std::move(ort_input_self)}, ort_outputs_0_ArgMax, &attrs); - if (!status.IsOK()) - throw std::runtime_error( - "ORT return failure status:" + status.ErrorMessage()); + CHECK_STATUS(status); return out; } @@ -909,9 +867,7 @@ bool equal( {std::move(ort_input_self), std::move(ort_input_other)}, ort_outputs_0_Equal, nullptr); - if (!equalStatus.IsOK()) - throw std::runtime_error( - "ORT Equal return failure status:" + equalStatus.ErrorMessage()); + CHECK_STATUS(equalStatus); // now reduce the resulting tensor of bool values to its minimum value (any false) NodeAttributes attrs(1); @@ -928,9 +884,7 @@ bool equal( {std::move(equalAsInt)}, ort_outputs_0_ReduceMin, &attrs); - if (!reduceStatus.IsOK()) - throw std::runtime_error( - "ORT ReduceMin return failure reduceStatus:" + reduceStatus.ErrorMessage()); + CHECK_STATUS(reduceStatus); auto* ort_tensor = ort_outputs_0_ReduceMin[0].GetMutable(); // the first (and only) value of the tensor will be 0 for false else true @@ -983,9 +937,7 @@ at::Tensor& fill__Scalar( {std::move(ort_input_self)}, ort_outputs_0_Shape, nullptr); - if (!status.IsOK()) - throw std::runtime_error( - "ORT return failure status:" + status.ErrorMessage()); + CHECK_STATUS(status); std::vector ort_outputs_1_ConstantOfShape(1); ort_outputs_1_ConstantOfShape[0] = ort_input_self; @@ -998,9 +950,7 @@ at::Tensor& fill__Scalar( {std::move(ort_outputs_0_Shape[0])}, ort_outputs_1_ConstantOfShape, &attrs); - if (!status.IsOK()) - throw std::runtime_error( - "ORT return failure status:" + status.ErrorMessage()); + CHECK_STATUS(status); return self; } diff --git a/orttraining/orttraining/eager/ort_aten.h b/orttraining/orttraining/eager/ort_aten.h index 968e57469b..abcec45c60 100644 --- a/orttraining/orttraining/eager/ort_aten.h +++ b/orttraining/orttraining/eager/ort_aten.h @@ -127,6 +127,7 @@ c10::optional PromoteScalarTypesWithCategory( ONNX_NAMESPACE::TensorProto_DataType GetONNXTensorProtoDataType(at::ScalarType dtype); OrtValue CastToType(onnxruntime::ORTInvoker& invoker, const OrtValue& input, at::ScalarType type); +void CastToType_out(onnxruntime::ORTInvoker& invoker, const OrtValue& input, OrtValue& output, at::ScalarType type); void resize_output( onnxruntime::ORTInvoker& invoker, diff --git a/orttraining/orttraining/eager/test/ort_ops.py b/orttraining/orttraining/eager/test/ort_ops.py index f1fd8c5d81..2692fabac7 100644 --- a/orttraining/orttraining/eager/test/ort_ops.py +++ b/orttraining/orttraining/eager/test/ort_ops.py @@ -1,7 +1,7 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. -# pylint: disable=missing-docstring +# pylint: disable=missing-docstring, too-many-public-methods, no-member import unittest @@ -10,49 +10,6 @@ import onnxruntime_pybind11_state as torch_ort import torch from parameterized import parameterized -# OPS - is a list of list of [test_operator, the tested_tensor]. -# The default value for tested_tensor is torch.rand (6)- size of 6 uniform distribution on the interval [0, 1). -# for floor and erf the ort- produce a roundoff error for floor(NaN), compare to cpu- stay NaN. thus, nan_to_num change nan to zero. -ops = [ - ["abs", torch.tensor([-1, -2, 3, -6, -7])], - ["acos"], - ["acosh"], - ["asinh"], - ["atanh"], - ["asin"], - ["atan"], - ["ceil"], - ["cos"], - ["cosh"], - ["erf", torch.nan_to_num(torch.rand(5))], - ["exp"], - ["floor", torch.nan_to_num(torch.rand(5))], - ["log"], - ["reciprocal"], - ["neg"], - ["round"], - ["relu"], - ["sigmoid"], - ["sin"], - ["sinh"], - ["sqrt"], - ["tan"], - ["tanh"], -] -# the following unary ops not been tested: -# ["isnan", torch.tensor([1, float('nan'), 2])]] -# ["selu", torch.randn(10)]] -# ["sign", ]] -# ["hardsigmoid", ], -# ["isinf", ], -# ["det" ]] - -math_sign_ops = ["eq", "ne", "lt", "gt"] - -# The function renames the test function: ops/math_sign_ops (e.g. abs)+ the test name(e.g. out), results in: test_abs_out -def rename_func(testcase_func, param_num, param): - return f"test_{parameterized.to_safe_name(str(param.args[0]))}{testcase_func.__name__[7:]}" - class OrtOpTests(unittest.TestCase): """test cases for supported eager ops""" @@ -60,37 +17,36 @@ class OrtOpTests(unittest.TestCase): def get_device(self): return torch_ort.device() - def test_add(self): - device = self.get_device() - cpu_ones = torch.Tensor([[1, 1, 1], [1, 1, 1], [1, 1, 1]]) - ort_ones = cpu_ones.to(device) - cpu_twos = cpu_ones + cpu_ones - ort_twos = ort_ones + ort_ones - assert torch.allclose(cpu_twos, ort_twos.cpu()) - - def test_type_promotion_add(self): - device = self.get_device() - x = torch.ones(2, 5, dtype=torch.int64) - y = torch.ones(2, 5, dtype=torch.float32) - ort_x = x.to(device) - ort_y = y.to(device) - ort_z = ort_x + ort_y - assert ort_z.dtype == torch.float32 - assert torch.allclose(ort_z.cpu(), (x + y)) - - def test_add_alpha(self): - device = self.get_device() - cpu_ones = torch.Tensor([[1, 1, 1], [1, 1, 1], [1, 1, 1]]) - ort_ones = cpu_ones.to(device) - assert torch.allclose(torch.add(cpu_ones, cpu_ones, alpha=2.5), torch.add(ort_ones, ort_ones, alpha=2.5).cpu()) - - # the onnx operator Mul does not support type bool. The following test verifies cpu fall back works. - def test_mul_bool(self): + def test_fallback_to_cpu(self): device = self.get_device() cpu_ones = torch.ones(3, 3, dtype=bool) ort_ones = cpu_ones.to(device) + # the onnx operator Mul does not support type bool so will fallback to cpu. assert torch.allclose(torch.mul(cpu_ones, cpu_ones), torch.mul(ort_ones, ort_ones).cpu()) + def test_type_promotion_add(self): + device = self.get_device() + cpu_ones_int64 = torch.ones(2, 5, dtype=torch.int64) + cpu_ones_float32 = torch.ones(2, 5, dtype=torch.float32) + ort_ones_int64 = cpu_ones_int64.to(device) + ort_ones_float32 = cpu_ones_float32.to(device) + cpu_result = cpu_ones_int64 + cpu_ones_float32 + ort_result = ort_ones_int64 + ort_ones_float32 + assert ort_result.dtype == torch.float32 + assert torch.allclose(cpu_result, ort_result.cpu()) + + # verify scalar addition promotion + cpu_result = cpu_ones_int64 + 1.1 + ort_result = ort_ones_int64 + 1.1 + assert ort_result.dtype == torch.float32 + assert torch.allclose(cpu_result, ort_result.cpu()) + + ## verify setting out to type int while inputs are float cause an error as casting float to int is not allowed. + cpu_out_tensor = torch.tensor([], dtype=torch.int) + ort_out_tensor = cpu_out_tensor.to(device) + with self.assertRaises(RuntimeError): + ort_result = torch.add(ort_ones_int64, ort_ones_float32, out=ort_out_tensor) + # TODO: Add BFloat16 test coverage def test_add_(self): device = self.get_device() @@ -102,24 +58,6 @@ class OrtOpTests(unittest.TestCase): ort_twos += ort_ones assert torch.allclose(cpu_twos, ort_twos.cpu()) - def test_sin_(self): - device = self.get_device() - cpu_sin_pi_ = torch.Tensor([np.pi]) - torch.sin_(cpu_sin_pi_) - ort_sin_pi_ = torch.Tensor([np.pi]).to(device) - torch.sin_(ort_sin_pi_) - cpu_sin_pi = torch.sin(torch.Tensor([np.pi])) - ort_sin_pi = torch.sin(torch.Tensor([np.pi]).to(device)) - assert torch.allclose(cpu_sin_pi, ort_sin_pi.cpu()) - assert torch.allclose(cpu_sin_pi_, ort_sin_pi_.cpu()) - assert torch.allclose(ort_sin_pi.cpu(), ort_sin_pi_.cpu()) - - def test_sin(self): - device = self.get_device() - cpu_sin_pi = torch.sin(torch.Tensor([np.pi])) - ort_sin_pi = torch.sin(torch.Tensor([np.pi]).to(device)) - assert torch.allclose(cpu_sin_pi, ort_sin_pi.cpu()) - def test_zero_like(self): device = self.get_device() ones = torch.ones((10, 10), dtype=torch.float32) @@ -373,50 +311,6 @@ class OrtOpTests(unittest.TestCase): ort_result = torch.bitwise_and(ort_a, ort_b) assert torch.equal(cpu_result, ort_result.cpu()) - # @parameterized.expand generate test methods for ops and using name_func we renaming the test to be test_{ops} - @parameterized.expand(ops, name_func=rename_func) - def test_op(self, test_name, tensor_test=torch.rand(6)): - # compile eval- creates a code object that evaluates the operator (for example torch.abs(tensor_test)) and returns its result. - cpu_result = eval(compile("torch." + test_name + "(tensor_test)", "", "eval")) - ort_result = eval(compile("torch." + test_name + "(tensor_test.to(self.get_device()))", "", "eval")) - assert torch.allclose(cpu_result, ort_result.cpu(), equal_nan=True) - - @parameterized.expand(ops, name_func=rename_func) - def test_op_(self, test_name, tensor_test=torch.rand(6)): - device = self.get_device() - - cpu_tensor = tensor_test - ort_tensor = cpu_tensor.to(device) - - eval(compile("torch." + test_name + "_(cpu_tensor)", "", "eval")) - eval(compile("torch." + test_name + "_(ort_tensor)", "", "eval")) - - assert torch.allclose(cpu_tensor, ort_tensor.cpu(), equal_nan=True) - - @parameterized.expand(ops, name_func=rename_func) - def test_op_out(self, test_name, tensor_test=torch.rand(6)): - ##relu -don't have output - if test_name == "relu": - self.skipTest(f"no {test_name}_output") - ### troubleshoot later: the following tests are Failing. - if test_name == "asin" or test_name == "log" or test_name == "atanh": - self.skipTest(f" {test_name}_output Fails - skipping for now") - device = self.get_device() - cpu_tensor = tensor_test - ort_tensor = cpu_tensor.to(device) - - cpu_out_tensor = torch.tensor([], dtype=tensor_test.dtype) - ort_out_tensor = cpu_out_tensor.to(device) - - st_cpu = f"torch.{test_name}(cpu_tensor, out=cpu_out_tensor)" - st_ort = f"torch.{test_name}(ort_tensor, out=ort_out_tensor)" - cpu_result = eval(compile(st_cpu, "", "eval")) - ort_result = eval(compile(st_ort, "", "eval")) - - assert torch.allclose(cpu_result, ort_result.cpu(), equal_nan=True) - assert torch.allclose(cpu_out_tensor, ort_out_tensor.cpu(), equal_nan=True) - assert torch.allclose(ort_result.cpu(), ort_out_tensor.cpu(), equal_nan=True) - def test_resize(self): device = self.get_device() @@ -475,6 +369,176 @@ class OrtOpTests(unittest.TestCase): self.assertEqual(cpu_tensor.size(), ort_tensor.size()) self.assertTrue(torch.allclose(cpu_tensor, ort_tensor.cpu())) + def test_fill(self): + device = self.get_device() + for torch_type in [torch.int, torch.float]: + cpu_tensor = torch.zeros(2, 2, dtype=torch_type) + ort_tensor = cpu_tensor.to(device) + for value in [True, 1.1, -1, 0]: + cpu_tensor.fill_(value) + ort_tensor.fill_(value) + assert cpu_tensor.dtype == ort_tensor.dtype + assert torch.equal(cpu_tensor, ort_tensor.to("cpu")) + + # tests both nonzero and nonzero.out + def test_nonzero(self): + device = self.get_device() + + for cpu_tensor in [ + torch.tensor([[[-1, 0, 1], [0, 1, 0]], [[0, 1, 0], [-1, 0, 1]]], dtype=torch.long), + torch.tensor([[[-1, 0, 1], [0, 1, 0]], [[0, 1, 0], [-1, 0, 1]]], dtype=torch.float), + ]: + ort_tensor = cpu_tensor.to(device) + + cpu_out_tensor = torch.tensor([], dtype=torch.long) + ort_out_tensor = cpu_out_tensor.to(device) + + # nonzero.out + cpu_result = torch.nonzero(cpu_tensor, out=cpu_out_tensor) + ort_result = torch.nonzero(ort_tensor, out=ort_out_tensor) + assert torch.equal(cpu_out_tensor, ort_out_tensor.to("cpu")) + assert torch.equal(cpu_result, ort_result.to("cpu")) + + # nonzero + cpu_result = torch.nonzero(cpu_tensor) + ort_result = torch.nonzero(ort_tensor) + assert torch.equal(cpu_result, ort_result.to("cpu")) + + # check result between nonzero.out and nonzero + assert torch.equal(ort_result.to("cpu"), ort_out_tensor.to("cpu")) + + def test_mm(self): + device = self.get_device() + + # out version test + cpu_mat1 = torch.rand(3, 2) + cpu_mat2 = torch.rand(2, 2) + ort_mat1 = cpu_mat1.to(device) + ort_mat2 = cpu_mat2.to(device) + cpu_out = torch.tensor([]) + ort_out = cpu_out.to(device) + cpu_result = torch.mm(cpu_mat1, cpu_mat2, out=cpu_out) + ort_result = torch.mm(ort_mat1, ort_mat2, out=ort_out) + assert torch.allclose(cpu_result, ort_result.cpu()) + assert torch.allclose(cpu_out, ort_out.cpu()) + assert torch.allclose(cpu_result, ort_out.cpu()) + + # non-out version with alternate dimension matrices + cpu_mat1 = torch.rand(7, 5) + cpu_mat2 = torch.rand(5, 4) + ort_mat1 = cpu_mat1.to(device) + ort_mat2 = cpu_mat2.to(device) + cpu_result = torch.mm(cpu_mat1, cpu_mat2) + ort_result = torch.mm(ort_mat1, ort_mat2) + assert torch.allclose(cpu_result, ort_result.cpu()) + + # check error cases + ort_mat1 = torch.rand(1, 1).to(device) + ort_bad_dim = torch.rand(2, 2).to(device) + ort_wrong_type = torch.ones(1, 1, dtype=torch.int).to(device) + ort_not_matrix = torch.ones(1).to(device) + with self.assertRaises(RuntimeError): + torch.mm(ort_mat1, ort_bad_dim) + with self.assertRaises(RuntimeError): + torch.mm(ort_mat1, ort_wrong_type) + with self.assertRaises(RuntimeError): + torch.mm(ort_mat1, ort_not_matrix) + + ################################ parameterized test follow ####################################### + # OPS - is a list of [test_operator, tested_tensor=torch.rand (6)]. + # The default value for tested_tensor is torch.rand (6)- size of 6 uniform distribution on the interval [0, 1). + # for floor and erf, the ort produces a roundoff error for NaN input, but cpu keeps it a NaN. + # Thus, we use nan_to_num to ensure actual numbers are passed in. + + # As many of the following use eval and make it appear to pylint that there are many unused variables, + # we disable those warnings + + # pylint: disable=eval-used, unused-argument, no-self-argument, reportSelfClsParameterName + + ops = [ + ["abs", torch.tensor([-1, -2, 3, -6, -7])], + ["acos"], + ["acosh"], + ["asinh"], + ["atanh"], + ["asin"], + ["atan"], + ["ceil"], + ["cos"], + ["cosh"], + ["erf", torch.nan_to_num(torch.rand(5))], + ["exp"], + ["floor", torch.nan_to_num(torch.rand(5))], + ["log"], + ["reciprocal"], + ["neg"], + ["round"], + ["relu"], + ["sigmoid"], + ["sin"], + ["sinh"], + ["sqrt"], + ["tan"], + ["tanh"], + ] + # the following unary ops not been tested: + # ["isnan", torch.tensor([1, float('nan'), 2])]] + # ["selu", torch.randn(10)]] + # ["sign", ]] + # ["hardsigmoid", ], + # ["isinf", ], + # ["det" ]] + + # The function renames the test function: ops/math_sign_ops (e.g. abs)+ the test name(e.g. out), results in: test_abs_out + def rename_func(testcase_func, param_num, param): + return f"test_{parameterized.to_safe_name(str(param.args[0]))}{testcase_func.__name__[7:]}" + + # @parameterized.expand generate test methods for ops and using name_func we renaming the test to be test_{ops} + @parameterized.expand(ops, name_func=rename_func) + def test_op(self, test_name, tensor_test=torch.rand(6)): + # compile eval- creates a code object that evaluates the operator (for example torch.abs(tensor_test)) and returns its result. + cpu_result = eval(compile("torch." + test_name + "(tensor_test)", "", "eval")) + ort_result = eval(compile("torch." + test_name + "(tensor_test.to(self.get_device()))", "", "eval")) + assert torch.allclose(cpu_result, ort_result.cpu(), equal_nan=True) + + @parameterized.expand(ops, name_func=rename_func) + def test_op_(self, test_name, tensor_test=torch.rand(6)): + device = self.get_device() + + cpu_tensor = tensor_test + ort_tensor = cpu_tensor.to(device) + + eval(compile("torch." + test_name + "_(cpu_tensor)", "", "eval")) + eval(compile("torch." + test_name + "_(ort_tensor)", "", "eval")) + + assert torch.allclose(cpu_tensor, ort_tensor.cpu(), equal_nan=True) + + @parameterized.expand(ops, name_func=rename_func) + def test_op_out(self, test_name, tensor_test=torch.rand(6)): + ##relu -don't have output + if test_name == "relu": + self.skipTest(f"no {test_name}_output") + ### troubleshoot later: the following tests are Failing. + if test_name in ("asin", "log", "atanh"): + self.skipTest(f" {test_name}_output Fails - skipping for now") + device = self.get_device() + cpu_tensor = tensor_test + ort_tensor = cpu_tensor.to(device) + + cpu_out_tensor = torch.tensor([], dtype=tensor_test.dtype) + ort_out_tensor = cpu_out_tensor.to(device) + + st_cpu = f"torch.{test_name}(cpu_tensor, out=cpu_out_tensor)" + st_ort = f"torch.{test_name}(ort_tensor, out=ort_out_tensor)" + cpu_result = eval(compile(st_cpu, "", "eval")) + ort_result = eval(compile(st_ort, "", "eval")) + + assert torch.allclose(cpu_result, ort_result.cpu(), equal_nan=True) + assert torch.allclose(cpu_out_tensor, ort_out_tensor.cpu(), equal_nan=True) + assert torch.allclose(ort_result.cpu(), ort_out_tensor.cpu(), equal_nan=True) + + math_sign_ops = ["eq", "ne", "lt", "gt"] + @parameterized.expand(math_sign_ops, name_func=rename_func) def test_op_tensor(self, math_sign_ops): device = self.get_device() @@ -551,81 +615,97 @@ class OrtOpTests(unittest.TestCase): assert torch.equal(cpu_float_float_lt_result, ort_float_float_lt_result.to("cpu")) assert torch.equal(cpu_float_float_gt_result, ort_float_float_gt_result.to("cpu")) - def test_fill(self): + binary_ops = [ # [op, op_sign, alpha_supported] + ["add", "+", True], + ["sub", "-", True], + ["mul", "*", False], + ["div", "/", False], + ] + + @parameterized.expand(binary_ops, name_func=rename_func) + def test_op_binary_tensor(self, binary_op, op_sign, alpha_supported): device = self.get_device() - for torch_type in [torch.int, torch.float]: - cpu_tensor = torch.zeros(2, 2, dtype=torch_type) - ort_tensor = cpu_tensor.to(device) - for value in [True, 1.1, -1, 0]: - cpu_tensor.fill_(value) - ort_tensor.fill_(value) - assert cpu_tensor.dtype == ort_tensor.dtype - assert torch.equal(cpu_tensor, ort_tensor.to("cpu")) + cpu_input = torch.rand(3, 3) + ort_input = cpu_input.to(device) + cpu_other = torch.rand(3, 3) + ort_other = cpu_other.to(device) - # tests both nonzero and nonzero.out - def test_nonzero(self): - device = self.get_device() - - for cpu_tensor in [ - torch.tensor([[[-1, 0, 1], [0, 1, 0]], [[0, 1, 0], [-1, 0, 1]]], dtype=torch.long), - torch.tensor([[[-1, 0, 1], [0, 1, 0]], [[0, 1, 0], [-1, 0, 1]]], dtype=torch.float), - ]: - ort_tensor = cpu_tensor.to(device) - - cpu_out_tensor = torch.tensor([], dtype=torch.long) - ort_out_tensor = cpu_out_tensor.to(device) - - # nonzero.out - cpu_result = torch.nonzero(cpu_tensor, out=cpu_out_tensor) - ort_result = torch.nonzero(ort_tensor, out=ort_out_tensor) - assert torch.equal(cpu_out_tensor, ort_out_tensor.to("cpu")) - assert torch.equal(cpu_result, ort_result.to("cpu")) - - # nonzero - cpu_result = torch.nonzero(cpu_tensor) - ort_result = torch.nonzero(ort_tensor) - assert torch.equal(cpu_result, ort_result.to("cpu")) - - # check result between nonzero.out and nonzero - assert torch.equal(ort_result.to("cpu"), ort_out_tensor.to("cpu")) - - def test_mm(self): - device = self.get_device() - - # out version test - cpu_mat1 = torch.rand(3, 2) - cpu_mat2 = torch.rand(2, 2) - ort_mat1 = cpu_mat1.to(device) - ort_mat2 = cpu_mat2.to(device) - cpu_out = torch.tensor([]) - ort_out = cpu_out.to(device) - cpu_result = torch.mm(cpu_mat1, cpu_mat2, out=cpu_out) - ort_result = torch.mm(ort_mat1, ort_mat2, out=ort_out) - assert torch.allclose(cpu_result, ort_result.cpu()) - assert torch.allclose(cpu_out, ort_out.cpu()) - assert torch.allclose(cpu_result, ort_out.cpu()) - - # non-out version with alternate dimension matrices - cpu_mat1 = torch.rand(7, 5) - cpu_mat2 = torch.rand(5, 4) - ort_mat1 = cpu_mat1.to(device) - ort_mat2 = cpu_mat2.to(device) - cpu_result = torch.mm(cpu_mat1, cpu_mat2) - ort_result = torch.mm(ort_mat1, ort_mat2) + # verify op_sign works + cpu_result = eval(compile("cpu_input " + op_sign + " cpu_other", "", "eval")) + ort_result = eval(compile("ort_input " + op_sign + " ort_other", "", "eval")) assert torch.allclose(cpu_result, ort_result.cpu()) - # check error cases - ort_mat1 = torch.rand(1, 1).to(device) - ort_bad_dim = torch.rand(2, 2).to(device) - ort_wrong_type = torch.ones(1, 1, dtype=torch.int).to(device) - ort_not_matrix = torch.ones(1).to(device) - with self.assertRaises(RuntimeError): - torch.mm(ort_mat1, ort_bad_dim) - with self.assertRaises(RuntimeError): - torch.mm(ort_mat1, ort_wrong_type) - with self.assertRaises(RuntimeError): - torch.mm(ort_mat1, ort_not_matrix) + # verify torch op with out param works + cpu_out_tensor = torch.tensor([]) + ort_out_tensor = cpu_out_tensor.to(device) + cpu_result = eval( + compile("torch." + binary_op + "(cpu_input, cpu_other, out=cpu_out_tensor)", "", "eval") + ) + ort_result = eval( + compile("torch." + binary_op + "(ort_input, ort_other, out=ort_out_tensor)", "", "eval") + ) + assert torch.allclose(cpu_result, ort_result.cpu()) + assert torch.allclose(cpu_out_tensor, ort_out_tensor.cpu()) + + if alpha_supported: + cpu_result = eval( + compile( + "torch." + binary_op + "(cpu_input, cpu_other, alpha=2.5, out=cpu_out_tensor)", "", "eval" + ) + ) + ort_result = eval( + compile( + "torch." + binary_op + "(ort_input, ort_other, alpha=2.5, out=ort_out_tensor)", "", "eval" + ) + ) + assert torch.allclose(cpu_result, ort_result.cpu()) + assert torch.allclose(cpu_out_tensor, ort_out_tensor.cpu()) + + @parameterized.expand(binary_ops, name_func=rename_func) + def test_op_binary_scalar(self, binary_op, op_sign, alpha_supported): + device = self.get_device() + cpu_input = torch.ones(3, 3) + ort_input = cpu_input.to(device) + cpu_other = 3.1 + ort_other = 3.1 + + # verify op_sign works + cpu_result = eval(compile("cpu_input " + op_sign + " cpu_other", "", "eval")) + ort_result = eval(compile("ort_input " + op_sign + " ort_other", "", "eval")) + assert torch.allclose(cpu_result, ort_result.cpu()) + + # verify torch op with out param works + cpu_out_tensor = torch.tensor([]) + ort_out_tensor = cpu_out_tensor.to(device) + cpu_result = eval( + compile("torch." + binary_op + "(cpu_input, cpu_other, out=cpu_out_tensor)", "", "eval") + ) + ort_result = eval( + compile("torch." + binary_op + "(ort_input, ort_other, out=ort_out_tensor)", "", "eval") + ) + assert torch.allclose(cpu_result, ort_result.cpu()) + assert torch.allclose(cpu_out_tensor, ort_out_tensor.cpu()) + + if alpha_supported: + cpu_result = eval( + compile( + "torch." + binary_op + "(cpu_input, cpu_other, alpha=2.5, out=cpu_out_tensor)", "", "eval" + ) + ) + ort_result = eval( + compile( + "torch." + binary_op + "(ort_input, ort_other, alpha=2.5, out=ort_out_tensor)", "", "eval" + ) + ) + assert torch.allclose(cpu_result, ort_result.cpu()) + assert torch.allclose(cpu_out_tensor, ort_out_tensor.cpu()) + + ################################################################ + # Please add new non-parameterized tests above the parameterized section. + ################################################################ if __name__ == "__main__": + # torch_ort.set_default_logger_severity(0) + # torch_ort.set_default_logger_verbosity(4) unittest.main()