From 62922f4c3cee5363ba55dbbf0c9515ef05a5d900 Mon Sep 17 00:00:00 2001 From: msftlincoln <107071614+msftlincoln@users.noreply.github.com> Date: Mon, 1 Aug 2022 15:52:47 -0400 Subject: [PATCH] Eager Mode generator: add comments, rename functions (#12385) * eager generator: add comments, rename functions * lint --- .../eager/opgen/opgen/generator.py | 146 +++++++++--------- 1 file changed, 70 insertions(+), 76 deletions(-) diff --git a/orttraining/orttraining/eager/opgen/opgen/generator.py b/orttraining/orttraining/eager/opgen/opgen/generator.py index bfdac107df..8813ad15ba 100644 --- a/orttraining/orttraining/eager/opgen/opgen/generator.py +++ b/orttraining/orttraining/eager/opgen/opgen/generator.py @@ -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 {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