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 returntype_type(t: Type, *, mutable: bool) -> CType: # placeholder is ignored r = valuetype_type(t, binds="__placeholder__") if r is not None: return r.type if isinstance(t, BaseType): if t.name == BaseTy.Tensor: if mutable: if local.use_const_ref_for_mutable_tensors(): return ConstRefCType(BaseCType(tensorT)) else: return MutRefCType(BaseCType(tensorT)) else: # Note [Tensor Copy Returns] # Currently, we use "Argument.is_write" to determine # whether or not Tensor return types should be copies or references. # If that ever changes, take a look at other locations of this note! return BaseCType(tensorT) elif t.name == BaseTy.Scalar: return BaseCType(scalarT) elif isinstance(t, ListType): assert ( not mutable ), "Native functions should never return a mutable tensor list. They should return void." elem = returntype_type(t.elem, mutable=False) assert t.size is None, f"fixed size list returns not supported: {t}" return VectorCType(elem) raise AssertionError(f"unrecognized return type {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)