#pragma once #include #include #include #include #include #include #include #include #include namespace torch { namespace jit { struct Function; namespace script { struct CompilationUnit; } } // namespace jit } // namespace torch namespace c10 { struct IValue; struct ClassType; struct TupleType; // For custom class __init__ registration, we need to pass in a function // that looks like this: [](IValue x, args...) // However, kernel_functor.h automatically sets the input types of the function // by introspecting the types of the functor (which is IValue in this case). // However, we need the type it binds to be Foo. // Instead, we pass in a lambda [](ivalue_holder x, args...) from // which getTypePtr can recover the original class pointer. template struct tagged_capsule { IValue ivalue; }; template c10::intrusive_ptr IValue::moveToIntrusivePtr() { auto t = c10::intrusive_ptr::reclaim(static_cast(payload.as_intrusive_ptr)); clearToNone(); return t; } template c10::intrusive_ptr IValue::toIntrusivePtr() const { auto r = c10::intrusive_ptr::reclaim(static_cast(payload.as_intrusive_ptr)); auto p = r; r.release(); return p; } template intrusive_ptr static_intrusive_pointer_cast(intrusive_ptr r) { return intrusive_ptr::reclaim(static_cast(r.release())); } inline c10::intrusive_ptr IValue::toFuture() && { AT_ASSERT(isFuture(), "Expected Future but got ", tagKind()); return moveToIntrusivePtr(); } inline c10::intrusive_ptr IValue::toFuture() const & { AT_ASSERT(isFuture(), "Expected Future but got ", tagKind()); return toIntrusivePtr(); } inline c10::intrusive_ptr IValue::toString() && { AT_ASSERT(isString(), "Expected String but got ", tagKind()); return moveToIntrusivePtr(); } inline c10::intrusive_ptr IValue::toString() const & { AT_ASSERT(isString(), "Expected String but got ", tagKind()); return toIntrusivePtr(); } inline c10::intrusive_ptr IValue::toObject() && { AT_ASSERT(isObject(), "Expected Object but got ", tagKind()); return toIntrusivePtr(); } inline c10::intrusive_ptr IValue::toObject() const & { AT_ASSERT(isObject(), "Expected Object but got ", tagKind()); return toIntrusivePtr(); } inline at::Tensor IValue::toTensor() && { AT_ASSERT(isTensor(), "Expected Tensor but got ", tagKind()); return at::Tensor(moveToIntrusivePtr()); } inline at::Tensor IValue::toTensor() const & { AT_ASSERT(isTensor(), "Expected Tensor but got ", tagKind()); return at::Tensor(toIntrusivePtr()); } inline c10::intrusive_ptr IValue::toBlob() && { AT_ASSERT(isBlob(), "Expected Blob but got ", tagKind()); return moveToIntrusivePtr(); } inline c10::intrusive_ptr IValue::toBlob() const & { AT_ASSERT(isBlob(), "Expected Blob but got ", tagKind()); return toIntrusivePtr();; } inline c10::intrusive_ptr IValue::toCapsule() && { TORCH_INTERNAL_ASSERT(isCapsule()); return moveToIntrusivePtr(); } inline c10::intrusive_ptr IValue::toCapsule() const & { TORCH_INTERNAL_ASSERT(isCapsule()); return toIntrusivePtr(); } namespace ivalue { template using Shared = c10::intrusive_ptr; // string struct CAFFE2_API ConstantString final : c10::intrusive_ptr_target { private: const std::string str_; public: ConstantString(std::string str) : str_(std::move(str)) {} static c10::intrusive_ptr create(std::string str_); const std::string & string() const { return str_; } operator const std::string & () const { return string(); } CAFFE2_API friend std::ostream& operator<<( std::ostream& out, const ConstantString& v); }; struct Future; struct CAFFE2_API Tuple : c10::intrusive_ptr_target { private: std::vector elements_; mutable std::shared_ptr type_; // lazily computed for unnamed tuples public: // named tuples have additional type information, so we // directly create them tagged static c10::intrusive_ptr createNamed( std::vector elements_, std::shared_ptr type_) { return c10::make_intrusive(std::move(elements_), type_); } static c10::intrusive_ptr create(std::vector elements_) { return c10::make_intrusive(std::move(elements_)); } const std::vector& elements() const & { return elements_; } operator const std::vector&() const { return elements(); } std::vector& elements() & { return elements_; } operator std::vector&() { return elements(); } std::vector&& elements() && { return std::move(elements_); } std::shared_ptr type() const; private: Tuple(std::vector elements, std::shared_ptr type = nullptr) : elements_(std::move(elements)), type_(std::move(type)) {} friend class c10::intrusive_ptr; }; struct Object; } // Future struct C10_EXPORT ivalue::Future final : c10::intrusive_ptr_target { private: c10::intrusive_ptr intrusive_from_this() { c10::raw::intrusive_ptr::incref(this); // we are creating a new pointer // from a raw `this` pointer // so we need to bump the refcount // to account for this ownership return c10::intrusive_ptr::reclaim(this); } public: Future(TypePtr type) : type_(type) {} struct CAFFE2_API FutureError final : public std::exception { FutureError(std::string&& error_msg_) : error_msg(std::move(error_msg_)) {} FutureError() = default; const char* what() const noexcept override { return error_msg.c_str(); } std::string error_msg; }; /** * Wait on the future until it completes. */ void wait() { std::unique_lock lock(mutex_); while (!completed_) { finished_cv_.wait(lock); } } /** * Explicitly mark the future as completed with the output value. */ void markCompleted(IValue value) { std::unique_lock lock(mutex_); AT_ASSERT(!completed()); completed_ = true; value_ = std::move(value); fireCallbacks(); finished_cv_.notify_all(); } void markCompleted() { markCompleted(IValue {}); } void markCompleted(FutureError&& error_) { std::unique_lock lock(mutex_); AT_ASSERT(!completed()); completed_ = true; has_error = true; error = std::move(error_); fireCallbacks(); finished_cv_.notify_all(); } // Get the result of the current future. IValue value() { std::unique_lock lock(mutex_); AT_ASSERT(completed()); if (has_error) { throw error; } return value_; } /** * Add a callback to the future. * The callbacks will be executed once the future completes. * If the future has already completed, * this function will execute the callback immediately. */ void addCallback(std::function callback) { std::unique_lock lock(mutex_); if (completed()) { lock.unlock(); callback(); return; } callbacks.push_back(callback); } // Check if the current future has completed bool completed() const{ return completed_; } CAFFE2_API friend std::ostream& operator<<( std::ostream& out, const Future& v); TypePtr type() const { return type_; } private: void fireCallbacks() { AT_ASSERT(completed()); // There is no need to protect callbacks with the lock. // Once completed_ is set to true, no one can add new callback to the list. for (auto& callback : callbacks) { callback(); } callbacks.clear(); } std::mutex mutex_; std::atomic_bool completed_ = {false}; // is this future complete std::condition_variable finished_cv_; IValue value_; // when finished the value TypePtr type_; std::vector> callbacks; bool has_error = false; FutureError error; }; // User-defined object. struct C10_EXPORT ivalue::Object final : c10::intrusive_ptr_target { public: Object(StrongTypePtr type, size_t numSlots) : type_(std::move(type)) { slots_.resize(numSlots); } static c10::intrusive_ptr create( StrongTypePtr type, size_t numSlots) { return c10::make_intrusive(std::move(type), numSlots); } /** * Slot API. * * Attributes are stored as a simple vector so that lookups are fast at * runtime. A "slot" is just an index into that vector, which can be computed * statically if you have access to the class type. Use this API if you are * writing compiler stuff. */ void setSlot(size_t slot, IValue v) { if (slot >= slots_.size()) { // for module types, it is possible that the members of the class have // expanded after the object was created. In this case, we expand // the slots to the right size resizeObject(slot); } slots_[slot] = v; } const IValue& getSlot(size_t slot) const { return slots_.at(slot); } /** * Attribute API. * * Wrappers around the slot stuff so that users can access attributes * directly. Use this API if you are a user. * * Note: Unlike in Python, TorchScript must make a distinction between * attributes (which are IValues) and methods (which are Methods). If you * want a method, use `obj.type()->getMethod()` */ IValue getAttr(const std::string& name) const; void setAttr(const std::string& name, IValue v); std::string name() const; const std::vector& slots() const { return slots_; } std::shared_ptr type() const { return type_.type_; } std::shared_ptr compilation_unit() { return type_.cu_; } private: void resizeObject(size_t slot); StrongTypePtr type_; std::vector slots_; }; std::vector> iterationOrder(const c10::Dict& dict); #undef TORCH_FORALL_TAGS namespace detail { struct _guarded_unsigned_long_unique_dummy final { _guarded_unsigned_long_unique_dummy(int64_t){}; }; using _guarded_unsigned_long = c10::guts::conditional_t< std::is_same::value || std::is_same::value, _guarded_unsigned_long_unique_dummy, unsigned long>; } // namespace detail inline const ivalue::Object& IValue::toObjectRef() const { AT_ASSERT(isObject(), "Expected Object but got ", tagKind()); return *static_cast(payload.as_intrusive_ptr); } // note: when adding a DEFINE_TO case here you should also add a // toX method to IValue. These named methods are much more discoverable // than the to templated function. #define DEFINE_TO(type, method_name) \ template<> \ inline type IValue::to() && { \ return std::move(*this).method_name(); \ } \ template<> \ inline type IValue::to() const & { \ return this->method_name(); \ } DEFINE_TO(at::Tensor, toTensor) DEFINE_TO(float, toDouble) DEFINE_TO(double, toDouble) DEFINE_TO(unsigned char, toInt) DEFINE_TO(signed char, toInt) DEFINE_TO(unsigned short, toInt) DEFINE_TO(short, toInt) DEFINE_TO(int, toInt) DEFINE_TO(uint32_t, toInt) DEFINE_TO(uint64_t, toInt) DEFINE_TO(detail::_guarded_unsigned_long, toInt) DEFINE_TO(int64_t, toInt) DEFINE_TO(bool, toBool) DEFINE_TO(c10::intrusive_ptr, toBlob); DEFINE_TO(c10::intrusive_ptr, toString) DEFINE_TO(c10::intrusive_ptr, toObject) DEFINE_TO(at::Scalar, toScalar) DEFINE_TO(c10::List, toIntList) DEFINE_TO(c10::List, toDoubleList) DEFINE_TO(c10::List, toBoolList) DEFINE_TO(c10::List, toTensorList) DEFINE_TO(c10::impl::GenericList, toGenericList) DEFINE_TO(c10::impl::GenericDict, toGenericDict) DEFINE_TO(c10::intrusive_ptr, toTuple) DEFINE_TO(std::string, toStringRef) DEFINE_TO(c10::intrusive_ptr, toFuture) DEFINE_TO(IValue, toIValue) DEFINE_TO(c10::Device, toDevice) DEFINE_TO(at::ScalarType, toScalarType) DEFINE_TO(at::Layout, toLayout) DEFINE_TO(at::MemoryFormat, toMemoryFormat) DEFINE_TO(at::QScheme, toQScheme) template struct _fake_type {}; // generic_to converts an IValue from a generic list or generic dict // to a concrete list/dict type likelike List, Dict<...> or optional. // Note that in the case of lists, this only works for IValue-based lists, // i.e. not for int64_t, double, ... // generic_to is an implementation detail of IValue::to and not // supposed to be called directly. // The _fake_type parameter allows us to overload // based on the return type. template C10_DEPRECATED_MESSAGE("IValues based on std::vector are potentially slow and deprecated. Please use c10::List instead.") std::vector generic_to( IValue ivalue, _fake_type>) { // We need to do a deep copy of the vector because there might be other // references to this same IValue that also use the list. We can't just // move the elements out. auto list = std::move(ivalue).to>(); std::vector result; result.reserve(list.size()); for (Elem v : list) { result.push_back(std::move(v)); } return result; } template T generic_to( IValue ivalue, _fake_type) { using ElemType = typename std::remove_pointer::type::element_type; auto obj = ivalue.toObject(); auto capsule = obj->getSlot(0); return c10::static_intrusive_pointer_cast(capsule.toCapsule()); } template tagged_capsule generic_to( IValue ivalue, _fake_type>) { return tagged_capsule{ivalue}; } template c10::List generic_to( IValue ivalue, _fake_type>) { return impl::toTypedList(std::move(ivalue).toGenericList()); } template c10::Dict generic_to( IValue ivalue, _fake_type>) { return impl::toTypedDict(std::move(ivalue).toGenericDict()); } template C10_DEPRECATED_MESSAGE("IValues based on std::unordered_map are slow and deprecated. Please use c10::Dict instead.") std::unordered_map generic_to( IValue ivalue, _fake_type>) { std::unordered_map specialized_dict; for (const auto& item : std::move(ivalue).toGenericDict()) { specialized_dict[item.key().to()] = item.value().to(); } return specialized_dict; } template c10::optional generic_to( IValue ivalue, _fake_type>) { if (ivalue.isNone()) { return c10::nullopt; } return std::move(ivalue).to(); } template inline T IValue::to() && { return generic_to(std::move(*this), _fake_type{}); } template inline T IValue::to() const& { return generic_to(*this, _fake_type{}); } inline c10::List IValue::toIntList() && { AT_ASSERT(isIntList(), "Expected IntList but got ", tagKind()); return c10::List(moveToIntrusivePtr>()); } inline c10::List IValue::toIntList() const & { AT_ASSERT(isIntList(), "Expected IntList but got ", tagKind()); return c10::List(toIntrusivePtr>()); } inline c10::ArrayRef IValue::toIntListRef() const { AT_ASSERT(isIntList(), "Expected IntList but got ", tagKind()); return static_cast*>(payload.as_intrusive_ptr)->list; } inline c10::List IValue::toDoubleList() && { AT_ASSERT(isDoubleList(), "Expected DoubleList but got ", tagKind()); return c10::List(moveToIntrusivePtr>()); } inline c10::List IValue::toDoubleList() const & { AT_ASSERT(isDoubleList(), "Expected DoubleList but got ", tagKind()); return c10::List(toIntrusivePtr>()); } inline c10::ArrayRef IValue::toDoubleListRef() const { AT_ASSERT(isDoubleList(), "Expected DoubleList but got ", tagKind()); return static_cast*>(payload.as_intrusive_ptr)->list; } inline c10::List IValue::toBoolList() && { AT_ASSERT(isBoolList(), "Expected BoolList but got ", tagKind()); return c10::List(moveToIntrusivePtr>()); } inline c10::List IValue::toBoolList() const & { AT_ASSERT(isBoolList(), "Expected BoolList but got ", tagKind()); return c10::List(toIntrusivePtr>()); } inline c10::List IValue::toTensorList() && { AT_ASSERT(isTensorList(), "Expected TensorList but got ", tagKind()); return c10::List(moveToIntrusivePtr>()); } inline c10::List IValue::toTensorList() const & { AT_ASSERT(isTensorList(), "Expected TensorList but got ", tagKind()); return c10::List(toIntrusivePtr>()); } inline c10::ArrayRef IValue::toTensorListRef() const { AT_ASSERT(isTensorList(), "Expected TensorList but got ", tagKind()); return static_cast*>(payload.as_intrusive_ptr)->list; } inline c10::List IValue::toGenericList() && { AT_ASSERT(isGenericList(), "Expected GenericList but got ", tagKind()); return c10::List(moveToIntrusivePtr>()); } inline c10::List IValue::toGenericList() const & { AT_ASSERT(isGenericList(), "Expected GenericList but got ", tagKind()); return c10::List(toIntrusivePtr>()); } inline c10::ArrayRef IValue::toGenericListRef() const { AT_ASSERT(isGenericList(), "Expected GenericList but got ", tagKind()); return static_cast*>(payload.as_intrusive_ptr)->list; } inline c10::Dict IValue::toGenericDict() && { AT_ASSERT(isGenericDict(), "Expected GenericDict but got ", tagKind()); return c10::Dict(moveToIntrusivePtr()); } inline c10::Dict IValue::toGenericDict() const & { AT_ASSERT(isGenericDict(), "Expected GenericDict but got ", tagKind()); return c10::Dict(toIntrusivePtr()); } inline c10::intrusive_ptr IValue::toTuple() && { AT_ASSERT(isTuple(), "Expected Tuple but got ", tagKind()); return moveToIntrusivePtr(); } inline c10::intrusive_ptr IValue::toTuple() const & { AT_ASSERT(isTuple(), "Expected Tuple but got ", tagKind()); return toIntrusivePtr(); } inline IValue::IValue(c10::intrusive_ptr v) : tag(Tag::Tuple), is_intrusive_ptr(true) { payload.as_intrusive_ptr = v.release(); } inline IValue::IValue(c10::List v) : tag(Tag::IntList), is_intrusive_ptr(true) { payload.as_intrusive_ptr = v.impl_.release(); } inline IValue::IValue(std::vector v) : IValue(c10::impl::toList(v)) {} inline IValue::IValue(c10::ArrayRef v) : IValue(c10::List(v)) {} inline IValue::IValue(c10::intrusive_ptr v) : tag(Tag::String), is_intrusive_ptr(true) { payload.as_intrusive_ptr = v.release(); } inline IValue::IValue(std::string v) : IValue(ivalue::ConstantString::create(std::move(v))) {} inline IValue::IValue(c10::List v) : tag(Tag::DoubleList), is_intrusive_ptr(true) { payload.as_intrusive_ptr = v.impl_.release(); } inline IValue::IValue(std::vector v) : IValue(c10::impl::toList(std::move(v))) {} inline IValue::IValue(c10::List v) : tag(Tag::BoolList), is_intrusive_ptr(true) { payload.as_intrusive_ptr = v.impl_.release(); } inline IValue::IValue(std::vector v) : IValue(c10::impl::toList(std::move(v))) {} inline IValue::IValue(c10::List v) : tag(Tag::TensorList), is_intrusive_ptr(true) { payload.as_intrusive_ptr = v.impl_.release(); } inline IValue::IValue(std::vector v) : IValue(c10::impl::toList(std::move(v))) {} inline IValue::IValue(c10::impl::GenericList v) : tag(Tag::GenericList), is_intrusive_ptr(true) { payload.as_intrusive_ptr = v.impl_.release(); } template inline IValue::IValue(c10::List v) : IValue(impl::toGenericList(std::move(v))) { static_assert(std::is_same::StorageT>::value, "Can only use this constructor for generic list types"); } template inline IValue::IValue(std::vector v) : IValue(c10::List()) { static_assert(std::is_same::StorageT>::value, "Can only use this constructor for generic list types"); auto list = to>(); list.reserve(v.size()); for (auto& e : v) { list.push_back(std::move(e)); } } inline IValue::IValue(c10::impl::GenericDict v) : tag(Tag::GenericDict), is_intrusive_ptr(true) { payload.as_intrusive_ptr = v.impl_.release(); } template inline IValue::IValue(c10::Dict v) : IValue(impl::toGenericDict(std::move(v))) {} template inline IValue::IValue(std::unordered_map v) : IValue(Dict()) { auto dict = to>(); dict.reserve(v.size()); for (auto& e : v) { dict.insert(std::move(e.first), std::move(e.second)); } } template inline IValue::IValue(c10::optional v): IValue() { if (v.has_value()) { *this = IValue(std::move(*v)); } } inline IValue::IValue(c10::nullopt_t): IValue() {} inline IValue::IValue(c10::intrusive_ptr v) : tag(Tag::Object), is_intrusive_ptr(true) { payload.as_intrusive_ptr = v.release(); } inline IValue::IValue(c10::intrusive_ptr v) : tag(Tag::Capsule), is_intrusive_ptr(true) { payload.as_intrusive_ptr = v.release(); } inline IValue::IValue(c10::intrusive_ptr v) : tag(Tag::Future), is_intrusive_ptr(true) { payload.as_intrusive_ptr = v.release(); } inline const std::string& IValue::toStringRef() const { return toString()->string(); } template inline optional IValue::toOptional() { if (this->isNone()) { return nullopt; } return this->to(); } inline bool IValue::isSameIdentity(const IValue& rhs) const { // We choose to not use memcmp for payload check due to potential random padding characters on union type // Semantics: // 1. None is None, False is False, and True is True are all true // 2. If it is a tensor type, we need to take undefined tensor into account // 3. Undefined_tensor is None and vice versa should be true // 4. If it is a reference type (i.e. is_intrusive_ptr), then is is True when the pointed-to object is the same. // 5. False for all other comparisons. if (this->isNone() && rhs.isNone()) { return true; } else if (this->isBool() && rhs.isBool()) { // for bool type, do equality check return this->toBool() == rhs.toBool(); } else if (this->isTensor() && rhs.isTensor()) { // for tensor type, just check the as_intrusive_ptr since is_intrusive_ptr is false for undefined tensor return this->payload.as_intrusive_ptr == rhs.payload.as_intrusive_ptr; } else if (this->isTensor() && rhs.isNone()) { // special case: undefined tensor and None are the same identity return !this->is_intrusive_ptr; } else if (this->isNone() && rhs.isTensor()) { // special case: undefined tensor and None are the same identity return !rhs.is_intrusive_ptr; } else { // for objects holding in IValue, do shallow compare on pointer address to testify the identity return this->is_intrusive_ptr && rhs.is_intrusive_ptr && this->payload.as_intrusive_ptr == rhs.payload.as_intrusive_ptr; } } namespace ivalue { namespace detail { // This code allows us to template on a function based on whether IValue has a // constructor for it. Specifically, has_constructor{} inherits from std::true_type if // IValue(T) compiles, and inherits from std::false_type if IValue(T) doesn't. // We use it for calling the IValue constructor for `from` if it exists, and otherwise // attempt to use our custom class code. template struct type_sink { typedef void type; }; template using type_sink_t = typename type_sink::type; template struct has_constructor : std::false_type {}; \ template struct has_constructor< T, type_sink_t< decltype( IValue(std::declval())) > >: std::true_type {}; template IValue from_(T x, std::true_type) { return IValue(x); } template IValue from_(c10::intrusive_ptr x, std::false_type) { using inputType = c10::intrusive_ptr; if (!isCustomClassRegistered()) { throw c10::Error("Trying to return a class that we don't support and isn't a registered custom class.", ""); } auto res = getCustomClassType(); auto retObject = ivalue::Object::create(res->second, 1); auto objPtr = c10::static_intrusive_pointer_cast(x); retObject->setSlot(0, IValue(objPtr)); auto resIVal = IValue(std::move(retObject)); return resIVal; } template IValue from_(T x, std::false_type) { static_assert(guts::false_t::value, "You are calling from with a type that it doesn't support, and isn't a potential custom class (ie: is an intrusive_ptr)"); return IValue(); } } template IValue from(T x) { return detail::from_(x, detail::has_constructor{}); } } } // namespace c10