diff --git a/orttraining/orttraining/eager/opgen/opgen/atenops.py b/orttraining/orttraining/eager/opgen/opgen/atenops.py index 1bd426958a..9e28c147a6 100644 --- a/orttraining/orttraining/eager/opgen/opgen/atenops.py +++ b/orttraining/orttraining/eager/opgen/opgen/atenops.py @@ -98,8 +98,16 @@ for unary_op in [ if unary_op not in ["isnan", "nonzero", "min", "max", "isinf", "det"]: ops[f"{aten_name}_"] = onnx_op +# Notes on Onnx op mapping +# +# Equal - Onnx spec has the return as a bool tensor, but aten will keep the tensor +# return type matching that of the "out" tensor if one is passed. To support this behavior +# we will CAST the Equal result to match the "out" as seen in eq and ne below. +# +# --------------------------- + hand_implemented = { - "aten::abs.out": SignatureOnly(), + "aten::abs.out": Abs("self"), "aten::empty.memory_format": SignatureOnly(), "aten::empty_strided": SignatureOnly(), "aten::zero_": SignatureOnly(), @@ -129,11 +137,10 @@ hand_implemented = { "aten::min": ReduceMin("self", keepdims=0), "aten::_cat": Concat("tensors", "dim"), "aten::fill_.Scalar": ConstantOfShape("self", value="value"), - "aten::ne.Scalar": MakeTorchFallback(), - "aten::ne.Scalar_out": MakeTorchFallback(), - "aten::ne.Tensor_out": MakeTorchFallback(), - "aten::eq.Tensor": MakeTorchFallback(), - "aten::eq.Tensor_out": MakeTorchFallback(), + "aten::ne.Scalar_out": Cast(Not(Equal("self", "other")), to="GetONNXTensorProtoDataType(out.scalar_type())"), + "aten::ne.Tensor_out": Cast(Not(Equal("self", "other")), to="GetONNXTensorProtoDataType(out.scalar_type())"), + "aten::eq.Tensor_out": Cast(Equal("self", "other"), to="GetONNXTensorProtoDataType(out.scalar_type())"), + "aten::eq.Scalar_out": Cast(Equal("self", "other"), to="GetONNXTensorProtoDataType(out.scalar_type())"), "aten::bitwise_and.Tensor_out": MakeTorchFallback(), "aten::masked_select": MakeTorchFallback(), "aten::_local_scalar_dense": MakeTorchFallback(), diff --git a/orttraining/orttraining/eager/opgen/opgen/generator.py b/orttraining/orttraining/eager/opgen/opgen/generator.py index 7189fe74bb..1c32446c7f 100644 --- a/orttraining/orttraining/eager/opgen/opgen/generator.py +++ b/orttraining/orttraining/eager/opgen/opgen/generator.py @@ -328,6 +328,36 @@ class ORTGen: ) writer.writeline() + # Gather return info on the torch func, such as Tensor(a! -> a) + # which implies the return value is writeable tensor (Tensor&) + 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. + set_out_tensor = False + last_param = cpp_func.parameters[-1].member + if ( + return_info + and last_param + 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) + + for torch_p in last_param.torch_param: + # Confirm we have an output alias and that it is writable (!) + # 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() + for onnx_op_index, onnx_op in enumerate(ctx.ops): # Torch -> ORT inputs for op_input in onnx_op.inputs: @@ -385,7 +415,6 @@ class ORTGen: writer.write(f"std::vector {onnx_op.outputs}") writer.writeline(f"({onnx_op.outputs.count});") - return_info = cpp_func.torch_func.return_type if cpp_func.torch_func else None in_place_params = {} if return_info: @@ -424,6 +453,11 @@ class ORTGen: in_place_params[0] = cpp_param.identifier.value break + # 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 len(in_place_params) != 0 and len(in_place_params) != ( len(return_info.elements) if isinstance(return_info, ast.TupleType) else 1 ): @@ -472,6 +506,10 @@ class ORTGen: if len(in_place_params) == 0: # tensor options + if set_out_tensor: + writer.writeline(f"return {last_param.identifier.value};") + return + writer.write(f"at::TensorOptions tensor_options = {first_param.identifier.value}") if first_param.parameter_type.desugar().identifier_tokens[0].value == "TensorList": writer.write("[0]") diff --git a/orttraining/orttraining/eager/ort_aten.cpp b/orttraining/orttraining/eager/ort_aten.cpp index d26dfd71ec..5de3da912d 100644 --- a/orttraining/orttraining/eager/ort_aten.cpp +++ b/orttraining/orttraining/eager/ort_aten.cpp @@ -901,42 +901,6 @@ const at::Tensor& resize_( return self; } -// aten::abs.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) -at::Tensor& abs_out( - const at::Tensor& self, - // *, - at::Tensor& out) { - ORT_LOG_FN(self, out); - - if ( - !IsSupportedType(self, {at::kHalf,at::kByte,at::kInt,at::kBFloat16,at::kFloat,at::kDouble,at::kShort,at::kLong})) { - return at::native::call_fallback_fn< - &at::native::cpu_fallback, - ATEN_OP(abs_out)>::call(self, out); - } - auto& invoker = GetORTInvoker(self.device()); - - auto ort_input_self = create_ort_value(invoker, self); - - resize_output(invoker, - dynamic_cast(out.unsafeGetTensorImpl()), - self.sizes()); - - auto ort_out = create_ort_value(invoker, out); - std::vector ort_outputs_0_Abs{ort_out}; - - auto status = invoker.Invoke("Abs", { - std::move(ort_input_self), - }, ort_outputs_0_Abs, nullptr); - - if (!status.IsOK()) { - throw std::runtime_error( - "ORT return failure status:" + status.ErrorMessage()); - } - - return out; -} - } // namespace aten //#pragma endregion diff --git a/orttraining/orttraining/eager/test/ort_ops.py b/orttraining/orttraining/eager/test/ort_ops.py index 7f6b412aa9..4224e472d0 100644 --- a/orttraining/orttraining/eager/test/ort_ops.py +++ b/orttraining/orttraining/eager/test/ort_ops.py @@ -289,6 +289,61 @@ class OrtOpTests(unittest.TestCase): assert torch.equal(cpu_out_tensor, ort_out_tensor.cpu()) assert torch.equal(ort_result.cpu(), ort_out_tensor.cpu()) + def test_eq_tensor(self): + device = self.get_device() + cpu_a = torch.Tensor([1.0, 1.5, 2.0]) + ort_a = cpu_a.to(device) + cpu_b = torch.Tensor([1.0, 1.5, 2.1]) + ort_b = cpu_b.to(device) + + for tensor_type in {torch.float, torch.bool}: + for func in {"eq", "ne"}: + print(f"Testing {func} with type {tensor_type}") + cpu_out_tensor = torch.tensor([], dtype=tensor_type) + ort_out_tensor = cpu_out_tensor.to(device) + cpu_a_b_eq_result = eval("torch." + func + "(cpu_a, cpu_b, out=cpu_out_tensor)") + ort_a_b_eq_result = eval("torch." + func + "(ort_a, ort_b, out=ort_out_tensor)") + assert torch.equal(cpu_a_b_eq_result.to(device), ort_a_b_eq_result) + assert torch.equal(cpu_out_tensor, ort_out_tensor.to("cpu")) + assert ort_out_tensor.dtype == tensor_type + + def test_eq_scalar(self): + device = self.get_device() + cpu_tensor_int = torch.tensor([1, 1], dtype=torch.int32) + cpu_scalar_int = torch.scalar_tensor(1, dtype=torch.int) + cpu_scalar_int_not = torch.scalar_tensor(2, dtype=torch.int) + cpu_tensor_float = torch.tensor([1.1, 1.1], dtype=torch.float32) + cpu_scalar_float = torch.scalar_tensor(1.1, dtype=torch.float32) + cpu_scalar_float_not = torch.scalar_tensor(1.0, dtype=torch.float32) + + ort_tensor_int = cpu_tensor_int.to(device) + ort_scalar_int = cpu_scalar_int.to(device) + ort_scalar_int_not = cpu_scalar_int_not.to(device) + ort_tensor_float = cpu_tensor_float.to(device) + ort_scalar_float = cpu_scalar_float.to(device) + ort_scalar_float_not = cpu_scalar_float_not.to(device) + + # compare int to int, float to float - ort only supports same type at the moment + cpu_out_tensor = torch.tensor([], dtype=torch.bool) + ort_out_tensor = cpu_out_tensor.to(device) + + for func in {"eq", "ne"}: + cpu_int_int_result = eval("torch." + func + "(cpu_tensor_int, cpu_scalar_int, out=cpu_out_tensor)") + cpu_int_int_not_result = eval("torch." + func + "(cpu_tensor_int, cpu_scalar_int_not)") + cpu_float_float_result = eval("torch." + func + "(cpu_tensor_float, cpu_scalar_float)") + cpu_float_float_not_result = eval("torch." + func + "(cpu_tensor_float, cpu_scalar_float_not)") + + ort_int_int_result = eval("torch." + func + "(ort_tensor_int, ort_scalar_int, out=ort_out_tensor)") + ort_int_int_not_result = eval("torch." + func + "(ort_tensor_int, ort_scalar_int_not)") + ort_float_float_result = eval("torch." + func + "(ort_tensor_float, ort_scalar_float)") + ort_float_float_not_result = eval("torch." + func + "(ort_tensor_float, ort_scalar_float_not)") + + assert torch.equal(cpu_out_tensor, ort_out_tensor.to("cpu")) + assert torch.equal(cpu_int_int_result, ort_int_int_result.to("cpu")) + assert torch.equal(cpu_int_int_not_result, ort_int_int_not_result.to("cpu")) + assert torch.equal(cpu_float_float_result, ort_float_float_result.to("cpu")) + assert torch.equal(cpu_float_float_not_result, ort_float_float_not_result.to("cpu")) + if __name__ == "__main__": unittest.main()