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

Eager mode generator improvements for multiple onnx operators and extra test cases (#12111)

Browse source 
* test case for masked_select

* isolate variables per onnx_op, include line numbers for ORT errors

* format errors

* correct masked_select impl, broadcast test

* node attrs naming fixed
This commit is contained in:
msftlincoln 2022-07-12 16:05:09 -04:00 committed by GitHub
parent 6e051016c1
commit a6fd1a3b85
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 36 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -162,7 +162,7 @@ hand_implemented = {
"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::masked_select": GatherND("self", Transpose(NonZero(Expand("mask", Shape("self"))))),
"aten::_local_scalar_dense": MakeTorchFallback(),
"aten::gt.Scalar_out": MakeTorchFallback(),
"aten::lt.Scalar_out": MakeTorchFallback(),

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

@ -193,6 +193,10 @@ class ORTGen:
writer.writeline('#include "ort_aten.h"')
writer.writeline('#include "ort_log.h"')
writer.writeline()
writer.writeline(
'#define CHECK_STATUS(status) if(!status.IsOK()) { std::stringstream err; err << "ORT return failure (line " << __LINE__ << "): " << status.ErrorMessage(); throw std::runtime_error(err.str()); }'
)
writer.writeline()
writer.push_namespace("torch_ort")
writer.push_namespace("eager")
writer.writeline()
@ -367,7 +371,7 @@ class ORTGen:
if isinstance(op_input, Outputs):
continue
cpp_param = cpp_func.get_parameter(op_input)
writer.write(f"auto ort_input_{op_input} = ")
writer.write(f"auto ort_input_{onnx_op_index}_{op_input} = ")
writer.writeline(f"create_ort_value(invoker, {op_input});")
if need_type_promotion:
type_func_str = (
@ -379,7 +383,7 @@ class ORTGen:
writer.writeline("{")
writer.push_indent()
writer.writeline(
f"ort_input_{op_input} = CastToType(invoker, ort_input_{op_input}, *promoted_type);"
f"ort_input_{onnx_op_index}_{op_input} = CastToType(invoker, ort_input_{onnx_op_index}_{op_input}, *promoted_type);"
)
writer.pop_indent()
writer.writeline("}")
@ -387,7 +391,7 @@ class ORTGen:
# Torch kwargs -> ORT attributes
attrs = {k: v for k, v in onnx_op.attributes.items() if v and v.value is not None}
if len(attrs) > 0:
attrs_arg = "attrs"
attrs_arg = f"attrs_{onnx_op_index}"
writer.writeline()
writer.writeline(f"NodeAttributes {attrs_arg}({len(attrs)});")
@ -441,7 +445,7 @@ class ORTGen:
and output_alias.is_writable
):
writer.writeline(
f"{onnx_op.outputs}[{output_index}] = ort_input_{onnx_op.inputs[input_index]};"
f"{onnx_op.outputs}[{output_index}] = ort_input_{onnx_op_index}_{onnx_op.inputs[input_index]};"
)
in_place_params[output_index] = cpp_param.identifier.value
break
@ -452,7 +456,9 @@ class ORTGen:
and self._get_alias_info(torch_p) == output_alias
and output_alias.is_writable
):
writer.writeline(f"{onnx_op.outputs}[0] = ort_input_{onnx_op.inputs[input_index]};")
writer.writeline(
f"{onnx_op.outputs}[0] = ort_input_{onnx_op_index}_{onnx_op.inputs[input_index]};"
)
in_place_params[0] = cpp_param.identifier.value
break
@ -481,23 +487,14 @@ class ORTGen:
raise FunctionGenerationError(cpp_func, "multiple outputs not supported")
op_input = f"{op_input}[0]"
else:
op_input = f"ort_input_{op_input}"
op_input = f"ort_input_{onnx_op_index}_{op_input}"
writer.writeline(f"std::move({op_input}),")
writer.pop_indent()
writer.write(f"}}, {onnx_op.outputs}, {attrs_arg}")
if onnx_op.domain:
writer.write(f", {onnx_op.domain}")
writer.writeline(");")
writer.writeline()
# Assert invocation
writer.writeline("if (!status.IsOK())")
writer.push_indent()
writer.writeline("throw std::runtime_error(")
writer.push_indent()
writer.writeline('"ORT return failure status:" + status.ErrorMessage());')
writer.pop_indent()
writer.pop_indent()
writer.writeline("CHECK_STATUS(status);")
writer.writeline()
# We'll potentially return back to Torch from this op

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

@ -251,6 +251,28 @@ class OrtOpTests(unittest.TestCase):
assert torch.allclose(cpu_result, ort_result.cpu())
assert cpu_result.dim() == ort_result.dim()
def test_masked_select(self):
device = self.get_device()
cpu_tensor = torch.randn(3, 4)
cpu_mask = cpu_tensor.ge(0.5)
ort_tensor = cpu_tensor.to(device)
ort_mask = cpu_mask.to(device)
cpu_result = cpu_tensor.masked_select(cpu_mask)
ort_result = ort_tensor.masked_select(ort_mask)
assert torch.allclose(cpu_result, ort_result.cpu())
assert cpu_result.dim() == ort_result.dim()
def test_masked_select_broadcast(self):
device = self.get_device()
cpu_tensor = torch.randn(3, 4)
cpu_mask = torch.tensor([[0], [1], [1]], dtype=bool)
ort_tensor = cpu_tensor.to(device)
ort_mask = cpu_mask.to(device)
cpu_result = cpu_tensor.masked_select(cpu_mask)
ort_result = ort_tensor.masked_select(ort_mask)
assert torch.allclose(cpu_result, ort_result.cpu())
assert cpu_result.dim() == ort_result.dim()
# @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_to_op)
def test_op(self, test_name, tensor_test=torch.rand(6)):