mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
ORT Eager Mode Generator - make smaller functions (#12371)
These changes resulted in no change to the generated outputs ort_aten.g.cpp and ort_customops.g.cpp.
This commit is contained in:
msftlincoln
GitHub
parent
f1a04078e9
commit
9559d25da9
1 changed files with 374 additions and 244 deletions
|
|
@ -157,8 +157,6 @@ class ORTGen:
|
|||
writer.push_namespace(ns)
|
||||
|
||||
writer.writeline()
|
||||
if mapped_func.cpp_func.torch_func:
|
||||
writer.writeline(f"// {mapped_func.cpp_func.torch_func.torch_schema}")
|
||||
|
||||
self._write_function_signature(writer, mapped_func.cpp_func)
|
||||
if mapped_func.signature_only:
|
||||
|
|
@ -211,6 +209,8 @@ class ORTGen:
|
|||
writer.pop_namespaces()
|
||||
|
||||
def _write_function_signature(self, writer: opgenwriter.SourceWriter, cpp_func: ast.FunctionDecl):
|
||||
if cpp_func.torch_func:
|
||||
writer.writeline(f"// {cpp_func.torch_func.torch_schema}")
|
||||
cpp_func.return_type.write(writer)
|
||||
writer.write(f" {cpp_func.identifier.value}(")
|
||||
writer.push_indent()
|
||||
|
|
@ -239,48 +239,230 @@ class ORTGen:
|
|||
writer.writeline(");")
|
||||
writer.pop_indent()
|
||||
|
||||
def _write_function_body(self, writer: opgenwriter.SourceWriter, mapped_func: MappedOpFunction):
|
||||
onnx_op, cpp_func = mapped_func.onnx_op, mapped_func.cpp_func
|
||||
|
||||
assert len(cpp_func.parameters) > 0
|
||||
|
||||
def _write_function_body_debug_header(self, writer, func_parameters):
|
||||
# Debug Logging
|
||||
log_params = ", ".join([p.member.identifier.value for p in cpp_func.parameters if p.member.identifier])
|
||||
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()
|
||||
|
||||
if mapped_func.make_torch_fallback:
|
||||
return self._write_cpu_fall_back(writer, mapped_func)
|
||||
def _write_function_body_resize_output(self, writer):
|
||||
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()
|
||||
|
||||
# Eval the outer ONNX op to produce a topologically ordered list of ops
|
||||
ctx = ONNXOpEvalContext()
|
||||
onnx_op.eval(ctx)
|
||||
ctx.prepare_outputs()
|
||||
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()"
|
||||
)
|
||||
writer.write(f"if ({op_input}.{type_func_str} != *promoted_type)")
|
||||
writer.writeline("{")
|
||||
writer.push_indent()
|
||||
writer.writeline(
|
||||
f"ort_input_{onnx_op_index}_{op_input} = CastToType(invoker, ort_input_{onnx_op_index}_{op_input}, *promoted_type);"
|
||||
)
|
||||
writer.pop_indent()
|
||||
writer.writeline("}")
|
||||
|
||||
# Fetch the ORT invoker from an at::Tensor.device()
|
||||
# FIXME: find the first at::Tensor param anywhere in the signature
|
||||
# instead of simply the first parameter?
|
||||
first_param = cpp_func.parameters[0].member
|
||||
# Check if the first parameter is tensorlist and if yes it's size should be > 0
|
||||
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)});")
|
||||
|
||||
for attr_name, attr in attrs.items():
|
||||
writer.write(f'{attrs_arg}["{attr_name}"] = ')
|
||||
writer.writeline("create_ort_attribute(")
|
||||
writer.push_indent()
|
||||
writer.write(f'"{attr_name}", {attr.value}')
|
||||
if attr.type.startswith("at::ScalarType::"):
|
||||
writer.write(f", {attr.type}")
|
||||
elif attr.type == AttrType.TENSOR:
|
||||
writer.write(f", true")
|
||||
elif attr.type != AttrType.STRING:
|
||||
raise FunctionGenerationError(
|
||||
cpp_func,
|
||||
f'Unsure how how to map ONNX op "{onnx_op.name}" attribute '
|
||||
+ f'"{attr_name}" of type "{attr.type}" to a call to '
|
||||
+ "create_ort_attribute. Please teach generator.py.",
|
||||
)
|
||||
writer.writeline(");")
|
||||
writer.pop_indent()
|
||||
|
||||
def _write_function_body_return_info_outputs(
|
||||
self, writer, in_place_params, cpp_func, return_info, onnx_op, onnx_op_index
|
||||
):
|
||||
for input_index, op_input in enumerate(onnx_op.inputs):
|
||||
if isinstance(op_input, Outputs):
|
||||
continue
|
||||
|
||||
# See if this input is aliased as an in-place tensor
|
||||
cpp_param = cpp_func.get_parameter(op_input)
|
||||
if not cpp_param:
|
||||
continue
|
||||
|
||||
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"
|
||||
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]};"
|
||||
)
|
||||
in_place_params[output_index] = cpp_param.identifier.value
|
||||
break
|
||||
elif isinstance(return_info, ast.ArrayType):
|
||||
if self._is_inplace(return_info, torch_p):
|
||||
writer.writeline(f"for (int i = 0; i < {onnx_op.outputs.count}; i++) {{")
|
||||
writer.push_indent()
|
||||
writer.writeline(
|
||||
f"{onnx_op.outputs}[i] = ort_input_{onnx_op_index}_{onnx_op.inputs[input_index]}[i];"
|
||||
)
|
||||
writer.pop_indent()
|
||||
writer.writeline("}")
|
||||
in_place_params[0] = cpp_param.identifier.value
|
||||
break
|
||||
elif self._is_inplace(return_info, torch_p):
|
||||
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
|
||||
|
||||
def _write_function_body_onnx_op_invocation(self, writer, onnx_op, onnx_op_index, cpp_func, attrs_arg_ptr):
|
||||
# Perform the invocation
|
||||
writer.writeline()
|
||||
if onnx_op_index == 0:
|
||||
writer.write("auto ")
|
||||
writer.writeline(f'status = invoker.Invoke("{onnx_op.name}", {{')
|
||||
writer.push_indent()
|
||||
for op_input in onnx_op.inputs:
|
||||
if isinstance(op_input, Outputs):
|
||||
if op_input.count != 1:
|
||||
raise FunctionGenerationError(cpp_func, "multiple outputs not supported")
|
||||
op_input = f"{op_input}[0]"
|
||||
else:
|
||||
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_ptr}")
|
||||
if onnx_op.domain:
|
||||
writer.write(f", {onnx_op.domain}")
|
||||
writer.writeline(");")
|
||||
writer.writeline("CHECK_STATUS(status);")
|
||||
writer.writeline()
|
||||
|
||||
return onnx_op.outputs
|
||||
|
||||
def _write_function_body_set_out(
|
||||
self,
|
||||
writer,
|
||||
in_place_params,
|
||||
set_out_tensor,
|
||||
onnx_op_index,
|
||||
ctx,
|
||||
need_type_promotion,
|
||||
cast_op_found,
|
||||
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:
|
||||
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
|
||||
):
|
||||
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}"
|
||||
)
|
||||
|
||||
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":
|
||||
writer.write("assert(")
|
||||
writer.write(first_param.identifier.value)
|
||||
writer.writeline(".size()>0);")
|
||||
writer.write("[0]")
|
||||
writer.write(".options()")
|
||||
if need_type_promotion:
|
||||
writer.write(".dtype(*promoted_type)")
|
||||
|
||||
# check whether need type promotion, if we do we will use this later to confirm out cast is supported.
|
||||
need_type_promotion = self._write_type_promotion(writer, mapped_func, cpp_func, ctx)
|
||||
# do we need to set type on the returned value
|
||||
if mapped_func.mapped_op_name in self.aten_output_type:
|
||||
writer.write(f".dtype({self.aten_output_type[mapped_func.mapped_op_name]})")
|
||||
|
||||
# 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
|
||||
writer.writeline(";")
|
||||
|
||||
writer.writeline("return aten_tensor_from_ort(")
|
||||
writer.push_indent()
|
||||
if (
|
||||
isinstance(cpp_func.return_type, ast.TemplateType)
|
||||
and cpp_func.return_type.identifier_tokens[-1].value == "std::vector"
|
||||
):
|
||||
writer.writeline(f"{return_outputs},")
|
||||
writer.writeline("tensor_options);")
|
||||
else:
|
||||
writer.writeline(f"std::move({return_outputs}[0]),")
|
||||
writer.writeline("tensor_options);")
|
||||
writer.pop_indent()
|
||||
|
||||
def _write_function_body_return_multiple(self, writer, cpp_func, in_place_params):
|
||||
if not (
|
||||
isinstance(cpp_func.return_type, ast.TemplateType)
|
||||
and cpp_func.return_type.identifier_tokens[-1].value == "std::tuple"
|
||||
):
|
||||
raise Exception(f"")
|
||||
tensorRef = "Tensor&," * len(in_place_params)
|
||||
tensorRef = tensorRef[: len(tensorRef) - 1]
|
||||
writer.write(f"return std::tuple<{tensorRef}>(")
|
||||
for index, key in enumerate(sorted(in_place_params)):
|
||||
if index > 0:
|
||||
writer.write(", ")
|
||||
writer.write(in_place_params[key])
|
||||
writer.writeline(");")
|
||||
|
||||
def _need_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
|
||||
# - cat.out(Tensor[] tensors, int dim=0, *, Tensor(a!) out) -> Tensor(a!)
|
||||
# - complex.out(Tensor real, Tensor imag, *, Tensor(a!) out) -> Tensor(a!)
|
||||
set_out_tensor = False
|
||||
last_param = cpp_func.parameters[-1].member
|
||||
|
||||
if (
|
||||
return_info
|
||||
and last_param
|
||||
|
|
@ -294,17 +476,30 @@ class ORTGen:
|
|||
# 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
|
||||
return True
|
||||
return False
|
||||
|
||||
# generate the type check
|
||||
cast_op_found = self._write_type_check(writer, mapped_func, cpp_func, ctx, need_type_promotion, set_out_tensor)
|
||||
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
|
||||
|
||||
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":
|
||||
writer.write("assert(")
|
||||
writer.write(first_param.identifier.value)
|
||||
writer.writeline(".size()>0);")
|
||||
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")
|
||||
|
||||
def _write_function_body_invoker(self, writer, first_param):
|
||||
# Fetch the ORT invoker from an at::Tensor.device()
|
||||
writer.write("auto& invoker = GetORTInvoker(")
|
||||
writer.write(first_param.identifier.value)
|
||||
if first_param.parameter_type.desugar().identifier_tokens[0].value == "TensorList":
|
||||
|
|
@ -312,226 +507,161 @@ class ORTGen:
|
|||
writer.writeline(".device());")
|
||||
writer.writeline()
|
||||
|
||||
# FIXME: warn if we have not consumed all torch parameters (either as
|
||||
# an ORT input or ORT attribute).
|
||||
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});")
|
||||
|
||||
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<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:
|
||||
if isinstance(op_input, Outputs):
|
||||
continue
|
||||
cpp_param = cpp_func.get_parameter(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 = (
|
||||
"type()"
|
||||
if cpp_param.parameter_type.desugar().identifier_tokens[0].value == "Scalar"
|
||||
else "scalar_type()"
|
||||
)
|
||||
writer.write(f"if ({op_input}.{type_func_str} != *promoted_type)")
|
||||
writer.writeline("{")
|
||||
writer.push_indent()
|
||||
writer.writeline(
|
||||
f"ort_input_{onnx_op_index}_{op_input} = CastToType(invoker, ort_input_{onnx_op_index}_{op_input}, *promoted_type);"
|
||||
)
|
||||
writer.pop_indent()
|
||||
writer.writeline("}")
|
||||
|
||||
# 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 = f"attrs_{onnx_op_index}"
|
||||
writer.writeline()
|
||||
writer.writeline(f"NodeAttributes {attrs_arg}({len(attrs)});")
|
||||
|
||||
for attr_name, attr in attrs.items():
|
||||
writer.write(f'{attrs_arg}["{attr_name}"] = ')
|
||||
writer.writeline("create_ort_attribute(")
|
||||
writer.push_indent()
|
||||
writer.write(f'"{attr_name}", {attr.value}')
|
||||
if attr.type.startswith("at::ScalarType::"):
|
||||
writer.write(f", {attr.type}")
|
||||
elif attr.type == AttrType.TENSOR:
|
||||
writer.write(f", true")
|
||||
elif attr.type != AttrType.STRING:
|
||||
raise FunctionGenerationError(
|
||||
cpp_func,
|
||||
f'Unsure how how to map ONNX op "{onnx_op.name}" attribute '
|
||||
+ f'"{attr_name}" of type "{attr.type}" to a call to '
|
||||
+ "create_ort_attribute. Please teach generator.py.",
|
||||
)
|
||||
writer.writeline(");")
|
||||
writer.pop_indent()
|
||||
attrs_arg = f"&{attrs_arg}"
|
||||
else:
|
||||
attrs_arg = "nullptr"
|
||||
|
||||
# Outputs vector
|
||||
writer.writeline()
|
||||
writer.write(f"std::vector<OrtValue> {onnx_op.outputs}")
|
||||
writer.writeline(f"({onnx_op.outputs.count});")
|
||||
|
||||
in_place_params = {}
|
||||
|
||||
if return_info:
|
||||
for input_index, op_input in enumerate(onnx_op.inputs):
|
||||
if isinstance(op_input, Outputs):
|
||||
continue
|
||||
|
||||
# See if this input is aliased as an in-place tensor
|
||||
cpp_param = cpp_func.get_parameter(op_input)
|
||||
if cpp_param:
|
||||
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"
|
||||
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]};"
|
||||
)
|
||||
in_place_params[output_index] = cpp_param.identifier.value
|
||||
break
|
||||
elif isinstance(return_info, ast.ArrayType):
|
||||
if self._is_inplace(return_info, torch_p):
|
||||
writer.writeline(f"for (int i = 0; i < {onnx_op.outputs.count}; i++) {{")
|
||||
writer.push_indent()
|
||||
writer.writeline(
|
||||
f"{onnx_op.outputs}[i] = ort_input_{onnx_op_index}_{onnx_op.inputs[input_index]}[i];"
|
||||
)
|
||||
writer.pop_indent()
|
||||
writer.writeline("}")
|
||||
in_place_params[0] = cpp_param.identifier.value
|
||||
break
|
||||
else:
|
||||
if self._is_inplace(return_info, torch_p):
|
||||
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
|
||||
|
||||
# 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:
|
||||
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
|
||||
):
|
||||
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}"
|
||||
)
|
||||
|
||||
# Perform the invocation
|
||||
writer.writeline()
|
||||
if onnx_op_index == 0:
|
||||
writer.write("auto ")
|
||||
writer.writeline(f'status = invoker.Invoke("{onnx_op.name}", {{')
|
||||
writer.push_indent()
|
||||
for op_input in onnx_op.inputs:
|
||||
if isinstance(op_input, Outputs):
|
||||
if op_input.count != 1:
|
||||
raise FunctionGenerationError(cpp_func, "multiple outputs not supported")
|
||||
op_input = f"{op_input}[0]"
|
||||
else:
|
||||
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("CHECK_STATUS(status);")
|
||||
writer.writeline()
|
||||
|
||||
# We'll potentially return back to Torch from this op
|
||||
return_outputs = onnx_op.outputs
|
||||
|
||||
# TODO: Pick the right "out" Torch parameter; do not assume the first one
|
||||
# TODO: Handle multiple results
|
||||
# TODO: Assert return type
|
||||
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
|
||||
cpp_param = cpp_func.get_parameter(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:
|
||||
self._write_function_body_onnx_op_input_type_promotion(writer, cpp_param, onnx_op_index, op_input)
|
||||
|
||||
def _write_function_body_return(
|
||||
self,
|
||||
writer,
|
||||
cpp_func,
|
||||
in_place_params,
|
||||
set_out_tensor,
|
||||
need_type_promotion,
|
||||
cast_op_found,
|
||||
onnx_op_outputs,
|
||||
last_param,
|
||||
first_param,
|
||||
mapped_func,
|
||||
return_outputs,
|
||||
):
|
||||
if cpp_func.return_type.desugar().identifier_tokens[0].value == "void":
|
||||
pass
|
||||
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
|
||||
|
||||
# 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":
|
||||
writer.write("[0]")
|
||||
writer.write(".options()")
|
||||
if need_type_promotion:
|
||||
writer.write(".dtype(*promoted_type)")
|
||||
|
||||
# do we need to set type on the returned value
|
||||
if mapped_func.mapped_op_name in self.aten_output_type:
|
||||
writer.write(f".dtype({self.aten_output_type[mapped_func.mapped_op_name]})")
|
||||
|
||||
writer.writeline(";")
|
||||
|
||||
writer.writeline("return aten_tensor_from_ort(")
|
||||
writer.push_indent()
|
||||
if (
|
||||
isinstance(cpp_func.return_type, ast.TemplateType)
|
||||
and cpp_func.return_type.identifier_tokens[-1].value == "std::vector"
|
||||
):
|
||||
writer.writeline(f"{return_outputs},")
|
||||
writer.writeline("tensor_options);")
|
||||
else:
|
||||
writer.writeline(f"std::move({return_outputs}[0]),")
|
||||
writer.writeline("tensor_options);")
|
||||
writer.pop_indent()
|
||||
return
|
||||
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,
|
||||
)
|
||||
elif len(in_place_params) == 1:
|
||||
writer.writeline(f"return {in_place_params[0]};")
|
||||
else:
|
||||
if not (
|
||||
isinstance(cpp_func.return_type, ast.TemplateType)
|
||||
and cpp_func.return_type.identifier_tokens[-1].value == "std::tuple"
|
||||
):
|
||||
raise Exception(f"")
|
||||
tensorRef = "Tensor&," * len(in_place_params)
|
||||
tensorRef = tensorRef[: len(tensorRef) - 1]
|
||||
writer.write(f"return std::tuple<{tensorRef}>(")
|
||||
for index, key in enumerate(sorted(in_place_params)):
|
||||
if index > 0:
|
||||
writer.write(", ")
|
||||
writer.write(in_place_params[key])
|
||||
writer.writeline(");")
|
||||
self._write_function_body_return_multiple(writer, cpp_func, in_place_params)
|
||||
|
||||
def _write_function_body_onnx_op(
|
||||
self,
|
||||
writer,
|
||||
onnx_op,
|
||||
onnx_op_index,
|
||||
need_type_promotion,
|
||||
cpp_func,
|
||||
return_info,
|
||||
set_out_tensor,
|
||||
ctx,
|
||||
cast_op_found,
|
||||
):
|
||||
self._write_function_body_onnx_op_inputs(writer, onnx_op, onnx_op_index, need_type_promotion, cpp_func)
|
||||
|
||||
# Torch kwargs -> ORT attributes
|
||||
attrs = {k: v for k, v in onnx_op.attributes.items() if v and v.value is not None}
|
||||
attrs_arg_ptr = "nullptr"
|
||||
if len(attrs) > 0:
|
||||
attrs_arg = f"attrs_{onnx_op_index}"
|
||||
attrs_arg_ptr = f"&{attrs_arg}"
|
||||
self._write_function_body_onnx_op_node_attributes(writer, onnx_op, attrs, attrs_arg)
|
||||
|
||||
self._write_function_body_onnx_op_output_vector(writer, onnx_op)
|
||||
|
||||
in_place_params = {}
|
||||
|
||||
if return_info:
|
||||
self._write_function_body_return_info_outputs(
|
||||
writer, in_place_params, cpp_func, return_info, onnx_op, onnx_op_index
|
||||
)
|
||||
self._write_function_body_set_out(
|
||||
writer,
|
||||
in_place_params,
|
||||
set_out_tensor,
|
||||
onnx_op_index,
|
||||
ctx,
|
||||
need_type_promotion,
|
||||
cast_op_found,
|
||||
onnx_op,
|
||||
return_info,
|
||||
cpp_func,
|
||||
)
|
||||
|
||||
# We'll potentially return back to Torch from this op
|
||||
return_outputs = self._write_function_body_onnx_op_invocation(
|
||||
writer, onnx_op, onnx_op_index, cpp_func, attrs_arg_ptr
|
||||
)
|
||||
|
||||
return in_place_params, return_outputs
|
||||
|
||||
# TODO: Pick the right "out" Torch parameter; do not assume the first one
|
||||
# TODO: Handle multiple results
|
||||
# TODO: Assert return type
|
||||
# TODO: warn if we have not consumed all torch parameters (either as an ORT input or ORT attribute).
|
||||
def _write_function_body(self, writer: opgenwriter.SourceWriter, mapped_func: MappedOpFunction):
|
||||
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)
|
||||
|
||||
if mapped_func.make_torch_fallback:
|
||||
return self._write_cpu_fall_back(writer, mapped_func)
|
||||
|
||||
first_param = cpp_func.parameters[0].member # FIXME: Find first at::Tensor param instead of first param only?
|
||||
self._write_function_body_first_param_assert(writer, first_param)
|
||||
|
||||
ctx = self._get_onnx_ops_eval_context(full_onnx_op)
|
||||
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)
|
||||
|
||||
self._write_function_body_invoker(writer, first_param)
|
||||
|
||||
if set_out_tensor:
|
||||
self._write_function_body_resize_output(writer)
|
||||
|
||||
for onnx_op_index, onnx_op in enumerate(ctx.ops):
|
||||
in_place_params, return_outputs = self._write_function_body_onnx_op(
|
||||
writer,
|
||||
onnx_op,
|
||||
onnx_op_index,
|
||||
need_type_promotion,
|
||||
cpp_func,
|
||||
return_info,
|
||||
set_out_tensor,
|
||||
ctx,
|
||||
cast_op_found,
|
||||
)
|
||||
|
||||
self._write_function_body_return(
|
||||
writer,
|
||||
cpp_func,
|
||||
in_place_params,
|
||||
set_out_tensor,
|
||||
need_type_promotion,
|
||||
cast_op_found,
|
||||
full_onnx_op.outputs,
|
||||
last_param,
|
||||
first_param,
|
||||
mapped_func,
|
||||
return_outputs,
|
||||
)
|
||||
|
||||
def _write_type_check(self, writer, mapped_func, cpp_func, ctx, need_type_promotion, set_out_tensor):
|
||||
cast_op_found = False
|
||||
|
|
|
|||
Loading…
Reference in a new issue