def process_ir_type( typ: Type, properties: "LazyIrProperties" ) -> Union[BaseCType, VectorCType, OptionalCType, ListCType]: """ This function takes a type from NativeFunctions and converts it for use with lazy tensor codegen. Type conversion for lazy currently consists of (1) changing at::Tensors into lazy::Values (2) wrapping everything in a BaseCType (3) making cpp-reference types into cpp-value types (e.g. vector instead of IntArrayRef) (1) converts at::Tensors to lazy::Values (which wrap lazy::Nodes, with which Lazy IR represents tensors.) There is special handling for Optional[Tensor] or List[Tensor], etc- hence 'tensor-like' This is incomplete- there are assertions in places that it's expected to need to add more types as the codegen is used with more operators. """ if isinstance(typ, BaseType): if typ.name == BaseTy.Tensor: return BaseCType(getValueT()) elif typ.name == BaseTy.Scalar: if properties.TreatScalarsAsConstants: return BaseCType(scalarT) # at::scalar has special handling, # and is wrapped in an lazy::Value just like at::tensor return BaseCType(getValueT()) elif typ.name == BaseTy.ScalarType: return BaseCType(scalarTypeT) elif typ.name == BaseTy.int: return BaseCType(longT) elif typ.name == BaseTy.SymInt: return BaseCType(getValueT()) elif typ.name == BaseTy.bool: return BaseCType(boolT) elif typ.name == BaseTy.float: return BaseCType(doubleT) elif typ.name == BaseTy.str: return BaseCType(stringT) elif typ.name == BaseTy.Device: return BaseCType(deviceT) elif typ.name == BaseTy.Layout: return BaseCType(layoutT) elif typ.name == BaseTy.MemoryFormat: return BaseCType(memoryFormatT) else: raise AssertionError(f"TODO add support for type {repr(typ)}") elif isinstance(typ, OptionalType): return OptionalCType(process_ir_type(typ.elem, properties)) elif isinstance(typ, ListType): if str(typ.elem) == "Tensor?": # TODO(whc) is this actually correct? or should it use a Vector like above return ListCType(OptionalCType(BaseCType(getValueT()))) elif str(typ.elem) == "Tensor": # this is a TensorList which comes in from GetTensorList as a Value return BaseCType(tensorListValueT) else: return VectorCType(process_ir_type(typ.elem, properties)) else: raise AssertionError(f"unrecognized type {repr(typ)}")
def argumenttype_type(t: Type, *, mutable: bool, binds: ArgName, remove_non_owning_ref_types: bool = False) -> NamedCType: # If it's a value type, do the value type translation r = valuetype_type(t, binds=binds, remove_non_owning_ref_types=remove_non_owning_ref_types) if r is not None: return r if isinstance(t, BaseType): if t.name == BaseTy.Tensor: if mutable and not local.use_const_ref_for_mutable_tensors(): return NamedCType(binds, MutRefCType(BaseCType(tensorT))) else: return NamedCType(binds, ConstRefCType(BaseCType(tensorT))) elif t.name == BaseTy.Scalar: return NamedCType(binds, ConstRefCType(BaseCType(scalarT))) else: raise AssertionError(f"base type should have been value type {t}") elif isinstance(t, OptionalType): if str(t.elem) == "Tensor": if mutable and not local.use_const_ref_for_mutable_tensors(): return NamedCType(binds, MutRefCType( BaseCType(tensorT))) # TODO: fix this discrepancy else: return NamedCType( binds, ConstRefCType(OptionalCType(BaseCType(tensorT)))) elif str(t.elem) == "Scalar": return NamedCType(binds, ConstRefCType(OptionalCType(BaseCType(scalarT)))) elif isinstance(t.elem, ListType) and str(t.elem.elem) == "int": return NamedCType(binds, BaseCType(optionalIntArrayRefT)) elem = argumenttype_type(t.elem, mutable=mutable, binds=binds) return NamedCType(binds, OptionalCType(elem.type)) elif isinstance(t, ListType): # TODO: remove these special cases, ArrayRef fallthrough works fine if str(t.elem) == "int": if remove_non_owning_ref_types: return NamedCType(binds, VectorCType(BaseCType(longT))) else: return NamedCType(binds, BaseCType(intArrayRefT)) elif str(t.elem) == "Tensor": return NamedCType(binds, BaseCType(tensorListT)) elif str(t.elem) == "Scalar": return NamedCType(binds, ArrayRefCType(BaseCType(scalarT))) elif str(t.elem) == "SymInt": return NamedCType(binds, BaseCType(symIntArrayRefT)) elif str(t.elem) == "Dimname": return NamedCType(binds, BaseCType(dimnameListT)) elif str(t.elem) == "Tensor?": return NamedCType( binds, ConstRefCType(ListCType(OptionalCType(BaseCType(tensorT))))) elem = argumenttype_type(t.elem, mutable=mutable, binds=binds) return NamedCType(binds, ArrayRefCType(elem.type)) else: raise AssertionError(f"unrecognized type {repr(t)}")
def argumenttype_type(t: Type, *, mutable: bool, binds: ArgName) -> NamedCType: if str(t) == "Tensor?": tensor_type: OptionalCType = OptionalCType(BaseCType(tensorT)) if mutable and not local.use_const_ref_for_mutable_tensors(): return NamedCType(binds, MutRefCType(tensor_type)) else: return NamedCType(binds, ConstRefCType(tensor_type)) elif str(t) == "Tensor?[]": return NamedCType( binds, ConstRefCType(ListCType(OptionalCType(BaseCType(tensorT))))) elif str(t) == "Scalar": return NamedCType(binds, ConstRefCType(BaseCType(scalarT))) elif str(t) == "Scalar?": return NamedCType(binds, ConstRefCType(OptionalCType(BaseCType(scalarT)))) return cpp.argumenttype_type(t, mutable=mutable, binds=binds)
def translate( bindings: Sequence[Union[Expr, Binding]], goals: Sequence[Union[NamedCType, Binding]], *, method: bool = False, allow_expensive_conversions: bool = False, ) -> List[Expr]: binding_exprs: List[Expr] = [] for b in bindings: if isinstance(b, Binding): binding_exprs.append(Expr( expr=b.name, type=b.nctype, )) else: binding_exprs.append(b) goal_ctypes: List[NamedCType] = [] for g in goals: if isinstance(g, Binding): goal_ctypes.append(g.nctype) else: goal_ctypes.append(g) # Add all the bindings to the context ctx: Dict[NamedCType, str] = {} for b in binding_exprs: ctx[b.type] = b.expr # While we're at it, do some simple forward inference, looking through # constructors. # # NB: When should you do forward inference versus backward inference? # The general idea: # # - Backward inference WHEN the goal gets smaller # - Forward inference WHEN the hypothesis gets smaller # # This helps ensure termination: backward inference starts with a goal # and tries to make it simpler and simpler until it's trivial; if the # goal can grow in size, we blow up to a really huge goal size. # Similarly, with forward inference we take hypotheses and decompose # them into simpler hypotheses; if hypotheses could expand in size, # we also have potential nontermination. (In the code below, forward # inference is only ever carried out at a single step, but you could # imagine repeated application of forward inference being profitable.) # # A good starting point in the literature for exploring more about proof # search are these lecture notes # https://www.cs.cmu.edu/~fp/courses/oregon-m10/04-focusing.pdf # # TODO: My kingdom for a pattern matcher # https://www.python.org/dev/peps/pep-0634/ # # TODO: This could get us in recomputation trouble if b.expr is nontrivial. # Fix this by implementing some sort of sharing so that if multiple # goals share the same expression, we only compute it once. This seems # to matter in practice as compiler is often unwilling to CSE nontrivial # expressions like scalar.to<scalar_t>() t = b.type if (isinstance(t, ConstRefCType) and isinstance(t.elem, OptionalCType) and isinstance(t.elem.elem, BaseCType) and str(t.elem.elem.type) == "at::Tensor"): ctx[NamedCType( t.elem.elem.name, ConstRefCType(BaseCType(tensorT)) )] = f"({b.expr}.has_value() ? *{b.expr} : at::Tensor())" if t.type == ConstRefCType(OptionalCType(BaseCType(tensorT))): ctx[NamedCType( t.name, BaseCType(optionalTensorRefT) )] = f"(({b.expr}.has_value() && (*{b.expr}).defined()) ? at::OptionalTensorRef(*{b.expr}) : at::OptionalTensorRef())" if t.type == ConstRefCType(BaseCType(scalarT)): ctx[NamedCType(t.name, BaseCType(opmath_t))] = f"({b.expr}).to<opmath_t>()" if t.type == ConstRefCType(OptionalCType(BaseCType(scalarT))): ctx[NamedCType( t.name, BaseCType(optionalScalarRefT) )] = f"({b.expr}.has_value() ? at::OptionalScalarRef(&({b.expr}.value())) : at::OptionalScalarRef())" if t.type == BaseCType(scalar_t): ctx[NamedCType( t.name, BaseCType(opmath_t))] = f"static_cast<opmath_t>({b.expr})" # [Note: ITensorListRef] if t.type == BaseCType(tensorListT): ctx[NamedCType( t.name, BaseCType(iTensorListRefT))] = f"at::ITensorListRef({b.expr})" # [Note: IOptTensorListRef] if t.type == ConstRefCType(ListCType(OptionalCType( BaseCType(tensorT)))): ctx[NamedCType(t.name, BaseCType( iOptTensorListRefT))] = f"at::IOptTensorListRef({b.expr})" # Add implicit bindings if the generated code is inside a Tensor method if method: ctx[NamedCType("self", MutRefCType( BaseCType(tensorT)))] = "const_cast<Tensor&>(*this)" ctx[NamedCType("self", ConstRefCType( BaseCType(tensorT)))] = "const_cast<Tensor&>(*this)" # This is better! Byte-for-byte compat # ctx[NamedCType("self", ConstRefCType(BaseCType(tensorT)))] = "*this" def unsat(goal: NamedCType) -> NoReturn: ctx_desc = "\n".join(f" {t.cpp_type()} {t.name}; // {e}" for t, e in ctx.items()) raise UnsatError(f""" Failed to synthesize the expression "{goal.cpp_type()} {goal.name}". When I failed, the following bindings were available in the context: {ctx_desc} This probably means there is a missing rule in the rules of torchgen.api.translate. Check this module for more information. """) # A shitty backtracking search implementation. It's shitty because it # does backtracking via stack (bad idea!) and for the most part tries to # avoid backtracking. In particular, if # direct=True, we won't try to do any fancy synthesis, just trivial # conversions (e.g., "T a" is OK for "const T& a"). So all of the # existing rules in this function simply try to solve immediately, # and bail if things don't work out. def solve(goal: NamedCType, *, direct: bool) -> str: def direct_solve(goal: NamedCType) -> str: return solve(goal, direct=True) if goal in ctx: # Trivial return ctx[goal] # const & is satisfied with mutable & if isinstance(goal.type, ConstRefCType): try: # WARNING: not strictly decreasing; be careful not # to add a direct conversion that goes satisfies # mutable& with const& return solve(NamedCType(goal.name, MutRefCType(goal.type.elem)), direct=direct) except UnsatError: pass # mutable & is satisfied with value if isinstance(goal.type, MutRefCType): try: return solve(NamedCType(goal.name, goal.type.elem), direct=direct) except UnsatError: pass if direct: unsat(goal) # For now, all of these rules are mutually exclusive. if goal == NamedCType("memory_format", OptionalCType(BaseCType(memoryFormatT))): memory_format = direct_solve( NamedCType( SpecialArgName.possibly_redundant_memory_format, OptionalCType(BaseCType(memoryFormatT)), )) # No need to join "memory_format" and "options" if the target API takes "options" directly. # Otherwise it will cause the redundant memory_format error. if options_ctype in goal_ctypes: return memory_format try: options = direct_solve(options_ctype) return f"c10::impl::check_tensor_options_and_extract_memory_format({options}, {memory_format})" except UnsatError: return memory_format elif goal == NamedCType("options", BaseCType(tensorOptionsT)): dtype = direct_solve( NamedCType("dtype", OptionalCType(BaseCType(scalarTypeT)))) pin_memory = direct_solve( NamedCType("pin_memory", OptionalCType(BaseCType(boolT)))) device = direct_solve( NamedCType("device", OptionalCType(BaseCType(deviceT)))) layout = direct_solve( NamedCType("layout", OptionalCType(BaseCType(layoutT)))) return f"TensorOptions().dtype({dtype}).layout({layout}).device({device}).pinned_memory({pin_memory})" elif goal == NamedCType("dtype", OptionalCType(BaseCType(scalarTypeT))): try: options = direct_solve(options_ctype) return f"optTypeMetaToScalarType({options}.dtype_opt())" except UnsatError: out_tensor = direct_solve(out_tensor_ctype) return f"{out_tensor}.scalar_type()" elif goal == NamedCType("layout", OptionalCType(BaseCType(layoutT))): try: options = direct_solve(options_ctype) return f"{options}.layout_opt()" except UnsatError: out_tensor = direct_solve(out_tensor_ctype) return f"{out_tensor}.layout()" elif goal == NamedCType("device", OptionalCType(BaseCType(deviceT))): try: options = direct_solve(options_ctype) return f"{options}.device_opt()" except UnsatError: out_tensor = direct_solve(out_tensor_ctype) return f"{out_tensor}.device()" elif goal == NamedCType("pin_memory", OptionalCType(BaseCType(boolT))): try: options = direct_solve(options_ctype) return f"{options}.pinned_memory_opt()" except UnsatError: # If we're calling a factory op from its out= variant, # We don't actually care about the value of pin_memory. out_tensor = direct_solve(out_tensor_ctype) return "c10::nullopt" # We can always do translations from value types to reference types, like vector<int> -> IntArrayRef elif goal.type == BaseCType(intArrayRefT): try: return direct_solve(NamedCType(goal.name, longVec_ctype)) except UnsatError: # We can also go SymIntArrayRef -> IntArrayRef symIntArrayRef_type = direct_solve( NamedCType(goal.name, BaseCType(symIntArrayRefT))) return f"c10::asIntArrayRefSlow({symIntArrayRef_type})" elif goal.type == BaseCType(symIntArrayRefT): return direct_solve(NamedCType(goal.name, longSymVec_ctype)) elif goal.type == BaseCType(longT): symInt_type = direct_solve( NamedCType(goal.name, BaseCType(SymIntT))) return f"{symInt_type}.expectInt()" elif goal.type == BaseCType(optionalIntArrayRefT): return direct_solve(NamedCType(goal.name, optionalLongVec_ctype)) elif goal.type == BaseCType(optionalScalarRefT): return direct_solve(NamedCType(goal.name, optionalScalar_ctype)) elif goal.type == BaseCType(optionalTensorRefT): return direct_solve(NamedCType(goal.name, optionalTensor_ctype)) # Note [translation from C++ reference to value types] # The below cases are all for when we have an argument with a reference type, # and a corresponding goal with a value type. # These are needed when we populate the inputs to a lambda capture and we need # to guarantee the lifetime of each captured argument. # We guard it with an explicit kwarg because converting to a value type is expensive # (O(n)) to convert from IntArrayRef to vector<int>), # so the caller of translate() should be explicit that they need it. if allow_expensive_conversions: if goal.type == VectorCType(BaseCType(longT)): intArrayRef_ctype = NamedCType(goal.name, BaseCType(intArrayRefT)) argname = direct_solve(intArrayRef_ctype) return f"{argname}.vec()" if goal.type == VectorCType(BaseCType(SymIntT)): symIntArrayRef_ctype = NamedCType(goal.name, BaseCType(symIntArrayRefT)) argname = direct_solve(symIntArrayRef_ctype) return f"{argname}.vec()" elif goal.type == OptionalCType(VectorCType(BaseCType(longT))): optionalIntArrayRef_ctype = NamedCType( goal.name, BaseCType(optionalIntArrayRefT)) argname = direct_solve(optionalIntArrayRef_ctype) return f"{argname}.has_value() ? c10::make_optional({argname}->vec()) : c10::nullopt" elif goal.type == OptionalCType(BaseCType(scalarT)): optionalScalarRef_ctype = NamedCType( goal.name, BaseCType(optionalScalarRefT)) argname = direct_solve(optionalScalarRef_ctype) return f"{argname}.has_value() ? c10::make_optional({argname}) : c10::nullopt" elif goal.type == OptionalCType(BaseCType(scalarT)): optionalTensorRef_ctype = NamedCType( goal.name, BaseCType(optionalTensorRefT)) argname = direct_solve(optionalTensorRef_ctype) return f"{argname}.has_value() ? c10::make_optional({argname}) : c10::nullopt" # Technically, we also need to handle cases of C++ containers holding reference types. # But there currently aren't any ops that require lambda capture codegen # With arguments like std::vector<IntArrayRef>. # If that changes, we'll have to add the translation here. # We allow const casting on tensors, since const-correctness is a bit broken for at::Tensor. # We could probably generalize this to non-tensor types too. if goal.type == MutRefCType(BaseCType(tensorT)): const_ref_tensor_ctype = NamedCType( goal.name, ConstRefCType(BaseCType(tensorT))) argname = direct_solve(const_ref_tensor_ctype) return f"const_cast<Tensor&>({argname})" unsat(goal) return [Expr(solve(g, direct=False), g) for g in goal_ctypes]
def save_var(var: SavedAttribute, is_output: bool) -> None: name = var.nctype.name type = var.nctype.type should_append_getsetdef = True should_append_raw_getsetdef = False if ( type == BaseCType(tensorT) or type == OptionalCType(BaseCType(tensorT)) or type == MutRefCType(OptionalCType(BaseCType(tensorT))) or (type == BaseCType(scalarT) and is_output) ): saved_variables.append(f"SavedVariable {name}_;") release_variables.append(f"{name}_.reset_data();") ptr = "shared_from_this()" if is_output else "" unpack.append(f"auto {name} = {name}_.unpack({ptr});") getter_definitions.append( GETTER_DEFINITION_SAVEDVAR.substitute( op=info.op, name=name, body=GETTER_BODY_SAVEDVAR ) ) getter_definitions.append( GETTER_DEFINITION_RAW_SAVEDVAR.substitute( op=info.op, name=name, body=GETTER_BODY_RAW_SAVEDVAR ) ) should_append_raw_getsetdef = True elif type == BaseCType(tensorListT): saved_variables.append(f"std::vector<SavedVariable> {name}_;") saved_variables.append(f"bool {name}_released_ = false;") # Just clear() is sufficient, we don't need to loop and clear each variable. # Because the SavedVariable owns a tensor and a grad_fn, removing the SavedVariable makes them go away as well. release_variables.append(f"{name}_.clear();") release_variables.append(f"{name}_released_ = true;") unpack.append(f"auto {name} = unpack_list({name}_);") asserts.append(f"TORCH_CHECK(!{name}_released_, ERR_BACKWARD_TWICE);") getter_definitions.append( GETTER_DEFINITION_VEC_SAVEDVAR.substitute( op=info.op, name=name, body=GETTER_BODY_VEC_SAVEDVAR ) ) getter_definitions.append( GETTER_DEFINITION_RAW_VEC_SAVEDVAR.substitute( op=info.op, name=name, body=GETTER_BODY_RAW_VEC_SAVEDVAR ) ) should_append_raw_getsetdef = True elif type == ListCType(OptionalCType(BaseCType(tensorT))): saved_variables.append(f"std::vector<SavedVariable> {name}_;") saved_variables.append(f"bool {name}_released_ = false;") # Just clear() is sufficient, we don't need to loop and clear each variable. # Because the SavedVariable owns a tensor and a grad_fn, removing the SavedVariable makes them go away as well. release_variables.append(f"{name}_.clear();") release_variables.append(f"{name}_released_ = true;") unpack.append(f"auto {name} = unpack_opt_list({name}_);") asserts.append(f"TORCH_CHECK(!{name}_released_, ERR_BACKWARD_TWICE);") getter_definitions.append( GETTER_DEFINITION_VEC_SAVEDVAR.substitute( op=info.op, name=name, body=GETTER_BODY_VEC_SAVEDVAR ) ) getter_definitions.append( GETTER_DEFINITION_RAW_VEC_SAVEDVAR.substitute( op=info.op, name=name, body=GETTER_BODY_RAW_VEC_SAVEDVAR ) ) should_append_raw_getsetdef = True elif type == BaseCType(intArrayRefT): saved_variables.append(f"std::vector<int64_t> {name};") getter_definitions.append( GETTER_DEFINITION.substitute( op=info.op, name=name, body=GETTER_BODY_ARRAYREF_LONG ) ) elif type == BaseCType(optionalIntArrayRefT): saved_variables.append(f"c10::OptionalArray<int64_t> {name};") getter_definitions.append( GETTER_DEFINITION_OPT_ARRAYREF.substitute( op=info.op, name=name, body=GETTER_BODY_ARRAYREF_LONG ) ) elif type == OptionalCType(BaseCType(intArrayRefT)): saved_variables.append(f"c10::OptionalArray<int64_t> {name};") getter_definitions.append( GETTER_DEFINITION_OPT_ARRAYREF.substitute( op=info.op, name=name, body=GETTER_BODY_ARRAYREF_LONG ) ) elif type == OptionalCType(ArrayRefCType(BaseCType(doubleT))): saved_variables.append(f"c10::OptionalArray<double> {name};") getter_definitions.append( GETTER_DEFINITION_OPT_ARRAYREF.substitute( op=info.op, name=name, body=GETTER_BODY_ARRAYREF_DOUBLE ) ) elif type == BaseCType(longT): saved_variables.append(f"{type.cpp_type()} {name} = 0;") getter_definitions.append( GETTER_DEFINITION.substitute( op=info.op, name=name, body=GETTER_BODY_INT64_T ) ) elif type == BaseCType(stringT): saved_variables.append(f"std::string {name};") getter_definitions.append( GETTER_DEFINITION.substitute( op=info.op, name=name, body=GETTER_BODY_STRING ) ) elif type == OptionalCType(BaseCType(stringT)): saved_variables.append(f"c10::optional<std::string> {name};") getter_definitions.append( GETTER_DEFINITION_OPT.substitute( op=info.op, name=name, body=GETTER_BODY_STRING ) ) else: saved_variables.append(f"{type.cpp_type()} {name};") if type in MISC_GETTER_DEFS: getter_def, body = MISC_GETTER_DEFS[type] getter_definitions.append( getter_def.substitute(op=info.op, name=name, body=body) ) else: # Types we don't expose python bindings to yet: # TypeAndSize, at::ScalarType, TensorOptions, TensorGeometry, # std::vector<std::vector<int64_t>>, std::vector<at::ScalarType> should_append_getsetdef = False if should_append_getsetdef: py_getsetdef_structs.append( PY_GETSETDEF_STRUCT.substitute(op=info.op, name=name) ) if should_append_raw_getsetdef: py_getsetdef_structs.append( PY_RAW_GETSETDEF_STRUCT.substitute(op=info.op, name=name) )