mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-25 19:48:11 +00:00
Eager Mode generator: add comments, rename functions (#12385)
* eager generator: add comments, rename functions * lint
This commit is contained in:
msftlincoln
GitHub
parent
f77ab4fea6
commit
62922f4c3c
1 changed files with 70 additions and 76 deletions
|
|
@ -239,12 +239,14 @@ class ORTGen:
|
|||
writer.writeline(");")
|
||||
writer.pop_indent()
|
||||
|
||||
def _write_function_body_debug_header(self, writer, func_parameters):
|
||||
# Debug Logging
|
||||
# Generates a log line for method entry, with the function parameters.
|
||||
def _write_function_body_entry_logging(self, writer, func_parameters):
|
||||
log_params = ", ".join([p.member.identifier.value for p in func_parameters if p.member.identifier])
|
||||
writer.writeline(f"ORT_LOG_FN({log_params});")
|
||||
writer.writeline()
|
||||
|
||||
# Generates code to resize a passed in output tensor to self.size().
|
||||
# TODO: allow resizing to other sizes.
|
||||
def _write_function_body_resize_output(self, writer):
|
||||
writer.writeline("// resize the output and then create output ort value to be updated.")
|
||||
writer.writeline(
|
||||
|
|
@ -253,6 +255,7 @@ class ORTGen:
|
|||
writer.writeline("auto ort_input_out = create_ort_value(invoker, out);")
|
||||
writer.writeline()
|
||||
|
||||
# Generates code to do type promotion via casting for a single input for an onnx op.
|
||||
def _write_function_body_onnx_op_input_type_promotion(self, writer, cpp_param, onnx_op_index, op_input):
|
||||
type_func_str = (
|
||||
"type()" if cpp_param.parameter_type.desugar().identifier_tokens[0].value == "Scalar" else "scalar_type()"
|
||||
|
|
@ -266,6 +269,7 @@ class ORTGen:
|
|||
writer.pop_indent()
|
||||
writer.writeline("}")
|
||||
|
||||
# Generates code to declare and populate node attributes for one onnx operator.
|
||||
def _write_function_body_onnx_op_node_attributes(self, writer, onnx_op, attrs, attrs_arg):
|
||||
writer.writeline()
|
||||
writer.writeline(f"NodeAttributes {attrs_arg}({len(attrs)});")
|
||||
|
|
@ -289,7 +293,8 @@ class ORTGen:
|
|||
writer.writeline(");")
|
||||
writer.pop_indent()
|
||||
|
||||
def _write_function_body_return_info_outputs(
|
||||
# Generates code which assigns an inplace (ort wrapped) input parameter to this onnx op's corresponding output
|
||||
def _write_function_body_assign_onnx_op_outputs_to_inplace_params(
|
||||
self, writer, in_place_params, cpp_func, return_info, onnx_op, onnx_op_index
|
||||
):
|
||||
for input_index, op_input in enumerate(onnx_op.inputs):
|
||||
|
|
@ -304,9 +309,7 @@ class ORTGen:
|
|||
for torch_p in cpp_param.torch_param:
|
||||
if isinstance(return_info, ast.TupleType):
|
||||
for output_index, output_param in enumerate(return_info.elements):
|
||||
assert isinstance(
|
||||
output_param.member, ast.TupleMemberType
|
||||
), "output_param.member must be of TupleMemberType"
|
||||
assert isinstance(output_param.member, ast.TupleMemberType)
|
||||
if self._is_inplace(output_param.member.element_type, torch_p):
|
||||
writer.writeline(
|
||||
f"{onnx_op.outputs}[{output_index}] = ort_input_{onnx_op_index}_{onnx_op.inputs[input_index]};"
|
||||
|
|
@ -329,6 +332,7 @@ class ORTGen:
|
|||
in_place_params[0] = cpp_param.identifier.value
|
||||
break
|
||||
|
||||
# Generates Onnx 'Invoke' call for one onnx op, with parameters, node attributes and output.
|
||||
def _write_function_body_onnx_op_invocation(self, writer, onnx_op, onnx_op_index, cpp_func, attrs_arg_ptr):
|
||||
# Perform the invocation
|
||||
writer.writeline()
|
||||
|
|
@ -354,25 +358,24 @@ class ORTGen:
|
|||
|
||||
return onnx_op.outputs
|
||||
|
||||
def _write_function_body_set_out(
|
||||
# Generates code to assign the aten "out" parameter (now in ort_input_out) to
|
||||
# the Onnx Operator output so that it will populate when the onnx op is invoked.
|
||||
# In the case that type casting is needed, _write_function_body_return_no_inplace
|
||||
# will cast the operator output and assign the out param after invocation.
|
||||
def _write_function_body_assign_onnx_op_output_to_out_param(
|
||||
self,
|
||||
writer,
|
||||
in_place_params,
|
||||
num_in_place_params,
|
||||
set_out_tensor,
|
||||
onnx_op_index,
|
||||
ctx,
|
||||
need_type_promotion,
|
||||
cast_op_found,
|
||||
is_last_onnx_op,
|
||||
need_type_promotion_without_cast,
|
||||
onnx_op,
|
||||
return_info,
|
||||
cpp_func,
|
||||
):
|
||||
# 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):
|
||||
# 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:
|
||||
if num_in_place_params == 0 and set_out_tensor and is_last_onnx_op:
|
||||
if need_type_promotion_without_cast:
|
||||
writer.writeline("if (*promoted_type == out.scalar_type()) {")
|
||||
writer.push_indent()
|
||||
writer.writeline(f"{onnx_op.outputs}[0] = ort_input_out;")
|
||||
|
|
@ -381,39 +384,21 @@ class ORTGen:
|
|||
else:
|
||||
writer.writeline(f"{onnx_op.outputs}[0] = ort_input_out;")
|
||||
|
||||
if len(in_place_params) != 0 and len(in_place_params) != (
|
||||
if num_in_place_params != 0 and num_in_place_params != (
|
||||
len(return_info.elements) if isinstance(return_info, ast.TupleType) else 1
|
||||
):
|
||||
raise Exception(
|
||||
f"Cannot mix and match inplace with non-inplace parameters - function: {cpp_func.identifier.value} "
|
||||
+ f"in_place_params={in_place_params}, return_elements={return_info.elements}"
|
||||
)
|
||||
raise Exception("Cannot mix and match inplace with non-inplace parameters.")
|
||||
|
||||
# Generates a non-ref return (the value returned is not one of the function's parameters).
|
||||
def _write_function_body_return_no_inplace(
|
||||
self,
|
||||
writer,
|
||||
set_out_tensor,
|
||||
need_type_promotion,
|
||||
cast_op_found,
|
||||
onnx_op_outputs,
|
||||
last_param,
|
||||
first_param,
|
||||
mapped_func,
|
||||
cpp_func,
|
||||
return_outputs,
|
||||
):
|
||||
# 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
|
||||
|
||||
# TODO: revisit the hardcoded use of TensorList.
|
||||
writer.write(f"at::TensorOptions tensor_options = {first_param.identifier.value}")
|
||||
if first_param.parameter_type.desugar().identifier_tokens[0].value == "TensorList":
|
||||
|
|
@ -441,6 +426,7 @@ class ORTGen:
|
|||
writer.writeline("tensor_options);")
|
||||
writer.pop_indent()
|
||||
|
||||
# Generates the return statement when a tuple is returned.
|
||||
def _write_function_body_return_multiple(self, writer, cpp_func, in_place_params):
|
||||
if not (
|
||||
isinstance(cpp_func.return_type, ast.TemplateType)
|
||||
|
|
@ -456,7 +442,10 @@ class ORTGen:
|
|||
writer.write(in_place_params[key])
|
||||
writer.writeline(");")
|
||||
|
||||
def _need_set_out_tensor(self, first_param, last_param, return_info):
|
||||
# determines if the "out" param exists and needs to be set.
|
||||
# "out", while similar to other modifiable (e.g. 'in_place') params, is treated
|
||||
# differently because the intent is to put result data into it, not modify data
|
||||
def _should_set_out_tensor(self, first_param, last_param, return_info):
|
||||
# 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. Record this need in set_out_tensor.
|
||||
# TODO: make this more general to handle cases where the first param is not self such as
|
||||
|
|
@ -479,13 +468,14 @@ class ORTGen:
|
|||
return True
|
||||
return False
|
||||
|
||||
# Evals the outer ONNX op to produce a topologically ordered list of ops (in case of nested ops).
|
||||
def _get_onnx_ops_eval_context(self, op):
|
||||
# may have nested operations - eval the outer ONNX op to produce a topologically ordered list of ops
|
||||
ctx = ONNXOpEvalContext()
|
||||
op.eval(ctx)
|
||||
ctx.prepare_outputs()
|
||||
return ctx
|
||||
|
||||
# Generates an assert to verify the first param is a tensor of size > 0.
|
||||
def _write_function_body_first_param_assert(self, writer, first_param):
|
||||
# Check if the first parameter is tensorlist and if yes it's size should be > 0
|
||||
if first_param.parameter_type.desugar().identifier_tokens[0].value == "TensorList":
|
||||
|
|
@ -498,8 +488,9 @@ class ORTGen:
|
|||
):
|
||||
raise FunctionGenerationError(cpp_func, "First parameter must be an at::Tensor")
|
||||
|
||||
def _write_function_body_invoker(self, writer, first_param):
|
||||
# Fetch the ORT invoker from an at::Tensor.device()
|
||||
# Generates code to get an ORT Invoker for the device from the first param.
|
||||
# The Invoker will be used and reused in _write_function_body_onnx_op_invocation.
|
||||
def _write_function_body_get_invoker(self, writer, first_param):
|
||||
writer.write("auto& invoker = GetORTInvoker(")
|
||||
writer.write(first_param.identifier.value)
|
||||
if first_param.parameter_type.desugar().identifier_tokens[0].value == "TensorList":
|
||||
|
|
@ -507,14 +498,14 @@ class ORTGen:
|
|||
writer.writeline(".device());")
|
||||
writer.writeline()
|
||||
|
||||
# Generates code to declare a vector which receives the full output for one onnx op.
|
||||
def _write_function_body_onnx_op_output_vector(self, writer, onnx_op):
|
||||
# Outputs vector
|
||||
writer.writeline()
|
||||
writer.write(f"std::vector<OrtValue> {onnx_op.outputs}")
|
||||
writer.writeline(f"({onnx_op.outputs.count});")
|
||||
|
||||
# Generates code to create ORT values from each torch function parameter that needs to be passed to this onnx op.
|
||||
def _write_function_body_onnx_op_inputs(self, writer, onnx_op, onnx_op_index, need_type_promotion, cpp_func):
|
||||
# Torch -> ORT inputs
|
||||
for op_input in onnx_op.inputs:
|
||||
if isinstance(op_input, Outputs):
|
||||
continue
|
||||
|
|
@ -524,6 +515,7 @@ class ORTGen:
|
|||
if need_type_promotion:
|
||||
self._write_function_body_onnx_op_input_type_promotion(writer, cpp_param, onnx_op_index, op_input)
|
||||
|
||||
# Ends the function by writing a return statement (or not for void).
|
||||
def _write_function_body_return(
|
||||
self,
|
||||
writer,
|
||||
|
|
@ -531,7 +523,7 @@ class ORTGen:
|
|||
in_place_params,
|
||||
set_out_tensor,
|
||||
need_type_promotion,
|
||||
cast_op_found,
|
||||
impl_uses_cast,
|
||||
onnx_op_outputs,
|
||||
last_param,
|
||||
first_param,
|
||||
|
|
@ -541,23 +533,27 @@ class ORTGen:
|
|||
if cpp_func.return_type.desugar().identifier_tokens[0].value == "void":
|
||||
pass
|
||||
elif len(in_place_params) == 0:
|
||||
self._write_function_body_return_no_inplace(
|
||||
writer,
|
||||
set_out_tensor,
|
||||
need_type_promotion,
|
||||
cast_op_found,
|
||||
onnx_op_outputs,
|
||||
last_param,
|
||||
first_param,
|
||||
mapped_func,
|
||||
cpp_func,
|
||||
return_outputs,
|
||||
)
|
||||
if set_out_tensor:
|
||||
if need_type_promotion and not impl_uses_cast:
|
||||
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};")
|
||||
else:
|
||||
self._write_function_body_return_no_inplace(
|
||||
writer, need_type_promotion, first_param, mapped_func, cpp_func, return_outputs
|
||||
)
|
||||
elif len(in_place_params) == 1:
|
||||
writer.writeline(f"return {in_place_params[0]};")
|
||||
else:
|
||||
self._write_function_body_return_multiple(writer, cpp_func, in_place_params)
|
||||
|
||||
# Generates all of the code for a single onnx op call, including mapping inputs, outputs and attributes.
|
||||
def _write_function_body_onnx_op(
|
||||
self,
|
||||
writer,
|
||||
|
|
@ -568,7 +564,7 @@ class ORTGen:
|
|||
return_info,
|
||||
set_out_tensor,
|
||||
ctx,
|
||||
cast_op_found,
|
||||
impl_uses_cast,
|
||||
):
|
||||
self._write_function_body_onnx_op_inputs(writer, onnx_op, onnx_op_index, need_type_promotion, cpp_func)
|
||||
|
||||
|
|
@ -585,20 +581,17 @@ class ORTGen:
|
|||
in_place_params = {}
|
||||
|
||||
if return_info:
|
||||
self._write_function_body_return_info_outputs(
|
||||
self._write_function_body_assign_onnx_op_outputs_to_inplace_params(
|
||||
writer, in_place_params, cpp_func, return_info, onnx_op, onnx_op_index
|
||||
)
|
||||
self._write_function_body_set_out(
|
||||
self._write_function_body_assign_onnx_op_output_to_out_param(
|
||||
writer,
|
||||
in_place_params,
|
||||
len(in_place_params),
|
||||
set_out_tensor,
|
||||
onnx_op_index,
|
||||
ctx,
|
||||
need_type_promotion,
|
||||
cast_op_found,
|
||||
onnx_op_index == (len(ctx.ops) - 1),
|
||||
need_type_promotion and not impl_uses_cast,
|
||||
onnx_op,
|
||||
return_info,
|
||||
cpp_func,
|
||||
)
|
||||
|
||||
# We'll potentially return back to Torch from this op
|
||||
|
|
@ -608,6 +601,7 @@ class ORTGen:
|
|||
|
||||
return in_place_params, return_outputs
|
||||
|
||||
# Generates code for the entire body of the function (everything between { and }.)
|
||||
# TODO: Pick the right "out" Torch parameter; do not assume the first one
|
||||
# TODO: Handle multiple results
|
||||
# TODO: Assert return type
|
||||
|
|
@ -616,7 +610,7 @@ class ORTGen:
|
|||
full_onnx_op, cpp_func = mapped_func.onnx_op, mapped_func.cpp_func
|
||||
assert len(cpp_func.parameters) > 0
|
||||
|
||||
self._write_function_body_debug_header(writer, cpp_func.parameters)
|
||||
self._write_function_body_entry_logging(writer, cpp_func.parameters)
|
||||
|
||||
if mapped_func.make_torch_fallback:
|
||||
return self._write_cpu_fall_back(writer, mapped_func)
|
||||
|
|
@ -628,10 +622,10 @@ class ORTGen:
|
|||
need_type_promotion = self._write_type_promotion(writer, mapped_func, cpp_func, ctx)
|
||||
return_info = cpp_func.torch_func.return_type if cpp_func.torch_func else None
|
||||
last_param = cpp_func.parameters[-1].member
|
||||
set_out_tensor = self._need_set_out_tensor(first_param, last_param, return_info)
|
||||
cast_op_found = self._write_type_check(writer, mapped_func, cpp_func, ctx, need_type_promotion, set_out_tensor)
|
||||
set_out_tensor = self._should_set_out_tensor(first_param, last_param, return_info)
|
||||
impl_uses_cast = self._write_type_check(writer, mapped_func, cpp_func, ctx, need_type_promotion, set_out_tensor)
|
||||
|
||||
self._write_function_body_invoker(writer, first_param)
|
||||
self._write_function_body_get_invoker(writer, first_param)
|
||||
|
||||
if set_out_tensor:
|
||||
self._write_function_body_resize_output(writer)
|
||||
|
|
@ -646,7 +640,7 @@ class ORTGen:
|
|||
return_info,
|
||||
set_out_tensor,
|
||||
ctx,
|
||||
cast_op_found,
|
||||
impl_uses_cast,
|
||||
)
|
||||
|
||||
self._write_function_body_return(
|
||||
|
|
@ -655,7 +649,7 @@ class ORTGen:
|
|||
in_place_params,
|
||||
set_out_tensor,
|
||||
need_type_promotion,
|
||||
cast_op_found,
|
||||
impl_uses_cast,
|
||||
full_onnx_op.outputs,
|
||||
last_param,
|
||||
first_param,
|
||||
|
|
@ -664,7 +658,7 @@ class ORTGen:
|
|||
)
|
||||
|
||||
def _write_type_check(self, writer, mapped_func, cpp_func, ctx, need_type_promotion, set_out_tensor):
|
||||
cast_op_found = False
|
||||
impl_uses_cast = False
|
||||
need_type_check = False
|
||||
if not self._custom_ops:
|
||||
for onnx_op in ctx.ops:
|
||||
|
|
@ -678,7 +672,7 @@ class ORTGen:
|
|||
for onnx_op in ctx.ops:
|
||||
# track is the CAST op was explicitly used
|
||||
if onnx_op.name == "Cast":
|
||||
cast_op_found = True
|
||||
impl_uses_cast = True
|
||||
for idx, op_input in enumerate(onnx_op.inputs):
|
||||
if isinstance(op_input, Outputs):
|
||||
continue
|
||||
|
|
@ -691,14 +685,14 @@ class ORTGen:
|
|||
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:
|
||||
if need_type_promotion and set_out_tensor and not impl_uses_cast:
|
||||
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("}")
|
||||
return cast_op_found
|
||||
return impl_uses_cast
|
||||
|
||||
def _write_type_promotion(self, writer, mapped_func, cpp_func, ctx):
|
||||
need_type_promotion = False
|
||||
|
|
|
|||
Loading…
Reference in a new issue