diff --git a/orttraining/orttraining/eager/opgen/opgen/atenops.py b/orttraining/orttraining/eager/opgen/opgen/atenops.py index 9e28c147a6..0c83e0e223 100644 --- a/orttraining/orttraining/eager/opgen/opgen/atenops.py +++ b/orttraining/orttraining/eager/opgen/opgen/atenops.py @@ -46,6 +46,57 @@ class GeluGrad(ONNXOp): ops = {} type_promotion_ops = [] +# the following op list is for ops that have a .out version. Often this is the only op needing to be implemented +# and the regular and inplace(_) version derive from the .out. +# Example: abs.out is the elementary op and abs and abs_ derive from that. +# However, that is not alway the case. Example: nonzero has .out and a normal version which is not derived. +# In this case nonzero will also be in the regular list. +unary_ops_with_out = [ + "abs", + "acos", + "acosh", + "asin", + "asinh", + "atan", + "atanh", + "ceil", + "cos", + "cosh", + "erf", + "exp", + "floor", + "hardsigmoid", + "log", + "neg", + "nonzero", + "reciprocal", + "round", + "sigmoid", + "sign", + "sin", + "sinh", + "sqrt", + "tan", + "tanh", +] + +# the following ops have explicit inplace(_) needing to be implemented. +unary_ops_with_inplace = [ + "relu", + "selu", +] + +# the following ops have the regular version needing to be implemented. +# Note det, isinf, selu - are composite aten ops but have direct onnx ops. +unary_ops = [ + "det", + "isinf", + "isnan", + "nonzero", + "relu", + "selu", +] + for binary_op, onnx_op in { "add": Add("self", Mul("alpha", "other")), "sub": Sub("self", Mul("alpha", "other")), @@ -59,44 +110,14 @@ for binary_op, onnx_op in { ops[f"aten::{binary_op}{variant}.{dtype}"] = deepcopy(onnx_op) type_promotion_ops.append(f"aten::{binary_op}{variant}.{dtype}") -for unary_op in [ - "acos", - "acosh", - "asinh", - "atanh", - "asin", - "atan", - "ceil", - "cos", - "cosh", - "erf", - "exp", - "floor", - "isnan", - "log", - "reciprocal", - "neg", - "round", - "relu", - "selu", - "sigmoid", - "sin", - "sinh", - "sqrt", - "tan", - "tanh", - "nonzero", - "sign", - "hardsigmoid", - "isinf", - "det", -]: - aten_name = f"aten::{unary_op}" - onnx_op = onnx_ops[unary_op]("self") - ops[aten_name] = onnx_op - # produce the in-place variant as well for ops that support it - if unary_op not in ["isnan", "nonzero", "min", "max", "isinf", "det"]: - ops[f"{aten_name}_"] = onnx_op +for unary_op in unary_ops_with_out: + ops[f"aten::{unary_op}.out"] = onnx_ops[unary_op]("self") + +for unary_op in unary_ops_with_inplace: + ops[f"aten::{unary_op}_"] = onnx_ops[unary_op]("self") + +for unary_op in unary_ops: + ops[f"aten::{unary_op}"] = onnx_ops[unary_op]("self") # Notes on Onnx op mapping # @@ -107,7 +128,6 @@ for unary_op in [ # --------------------------- hand_implemented = { - "aten::abs.out": Abs("self"), "aten::empty.memory_format": SignatureOnly(), "aten::empty_strided": SignatureOnly(), "aten::zero_": SignatureOnly(), diff --git a/orttraining/orttraining/eager/opgen/opgen/generator.py b/orttraining/orttraining/eager/opgen/opgen/generator.py index 1c32446c7f..77a67bc1f0 100644 --- a/orttraining/orttraining/eager/opgen/opgen/generator.py +++ b/orttraining/orttraining/eager/opgen/opgen/generator.py @@ -334,6 +334,9 @@ class ORTGen: # 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. + # 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 ( @@ -510,6 +513,7 @@ class ORTGen: 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]")