saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-27 20:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add eager support for eq and ne ops. (#12031)

Browse source 
* Add eager support for aten::eq and aten:ne.

* Add generator support for resizing output param.
This commit is contained in:
Wil Brady 2022-07-06 12:39:04 -04:00 committed by GitHub
parent 07b0469a23
commit 1948b7c726
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 107 additions and 43 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
orttraining/orttraining/eager/opgen/opgen/atenops.py
View file

@ -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(),

40
orttraining/orttraining/eager/opgen/opgen/generator.py
View file

@ -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<ORTTensorImpl*>(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<OrtValue> {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]")

36
orttraining/orttraining/eager/ort_aten.cpp
View file

@ -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<ORTTensorImpl*>(out.unsafeGetTensorImpl()),
self.sizes());
auto ort_out = create_ort_value(invoker, out);
std::vector<OrtValue> 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

55
orttraining/orttraining/eager/test/ort_ops.py
View file

@ -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()