Class ImmutableTensor::DevValueCache

src/opr/impl/io.cpp:432–559

}

class ImmutableTensor::DevValueCache final : public UserDataContainer::UserData {
    MGB_TYPEINFO_OBJ_DECL;
    CompNode m_comp_node;

    class TensorKey {
        struct Trait {
            size_t hash = 0, size_bytes = 0;
            TensorLayout layout;
        };
        Trait m_trait;
        std::vector<dt_byte> m_val;
        HostTensorND m_val_ref;

        const dt_byte* val_ptr() const {
            mgb_assert(m_trait.size_bytes);
            return m_val.empty() ? m_val_ref.raw_ptr() : m_val.data();
        }

    public:
        TensorKey() = default;
        TensorKey(const HostTensorND& v) : m_val_ref{v} {
            mgb_assert(v.layout().is_contiguous() || v.layout().is_empty());
            m_trait.size_bytes = v.layout().span().high_byte;

            auto&& layout = m_trait.layout;
            // zero to enable byte-comparison
            memset(&layout, 0, sizeof(layout));
            layout.ndim = v.layout().ndim;
            layout.dtype = v.layout().dtype;
            for (size_t i = 0; i < layout.ndim; ++i) {
                layout.shape[i] = v.layout().shape[i];
                layout.stride[i] = v.layout().stride[i];
            }
            XXHash hasher;
            if (!v.empty()) {
                hasher.update(v.raw_ptr(), m_trait.size_bytes);
            }
            hasher.update(&m_trait.layout, sizeof(m_trait.layout));
            m_trait.hash = hasher.digest();
        }

        bool operator==(const TensorKey& rhs) const {
            return !memcmp(&m_trait, &rhs.m_trait, sizeof(Trait)) &&
                   ((m_trait.size_bytes == 0 && rhs.m_trait.size_bytes == 0) ||
                    !memcmp(val_ptr(), rhs.val_ptr(), m_trait.size_bytes));
        }

        size_t hash() const { return m_trait.hash; }

        //! copy from m_val_ref to m_val, to avoid refed value being
        //! modified
        void copy_val_permanent() {
            if (m_trait.size_bytes == 0)
                return;
            mgb_assert(m_val.empty());
            m_val.resize(m_trait.size_bytes);
            memcpy(m_val.data(), m_val_ref.raw_ptr(), m_trait.size_bytes);
            m_val_ref = {};