// Copyright (C) 2018-2026 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//

#include "dnnl_memory_desc.h"

#include <oneapi/dnnl/dnnl_types.h>

#include <common/memory_desc.hpp>
#include <common/memory_desc_wrapper.hpp>
#include <cstddef>
#include <memory>
#include <oneapi/dnnl/dnnl.hpp>
#include <string>

#include "cpu_types.h"
#include "dnnl_extension_utils.h"
#include "memory_desc/cpu_memory_desc.h"
#include "openvino/core/except.hpp"
#include "openvino/core/type/element_type.hpp"

namespace ov::intel_cpu {

DnnlMemoryDesc::DnnlMemoryDesc(const dnnl::memory::desc& desc) : DnnlMemoryDesc(desc.get()) {}

DnnlMemoryDesc::DnnlMemoryDesc(const_dnnl_memory_desc_t cdesc)
    : MemoryDesc(Shape(DnnlExtensionUtils::convertToVectorDims(cdesc->dims, cdesc->ndims)), Dnnl),
      desc(DnnlExtensionUtils::clone_desc(cdesc)) {
    OPENVINO_ASSERT(getFormatKind() != dnnl::memory::format_kind::any, "Unexpected: Memory format any is prohibited!");
}

ov::element::Type DnnlMemoryDesc::getPrecision() const {
    return DnnlExtensionUtils::DataTypeToElementType(getDataType());
}

MemoryDescPtr DnnlMemoryDesc::clone() const {
    return std::make_shared<DnnlMemoryDesc>(*this);
}

MemoryDescPtr DnnlMemoryDesc::cloneWithNewPrecision(const ov::element::Type prec) const {
    auto newDesc = std::make_shared<DnnlMemoryDesc>(*this);
    newDesc->setPrecision(prec);
    return newDesc;
}

bool DnnlMemoryDesc::isCompatible(const MemoryDesc& rhs) const {
    if (MemoryDescType::Dnnl & rhs.getType()) {
        const auto* dnnMemDesc = rhs.as<DnnlMemoryDesc>();
        return isCompatible(*dnnMemDesc);
    }
    return false;
}

bool DnnlMemoryDesc::isCompatible(const DnnlMemoryDesc& rhs) const {
    return this->desc == rhs.desc;
}

std::string DnnlMemoryDesc::serializeFormat() const {
    dnnl::impl::memory_desc_wrapper wrapped(desc.get());
    if (wrapped.is_wino_desc()) {
        switch (desc.get()->format_desc.wino_desc.wino_format) {
        case dnnl::impl::wino_memory_format_t::wino_wei_aaOio:
            return "wino_aaOio";
        case dnnl::impl::wino_memory_format_t::wino_wei_aaOBiOo:
            return "wino_aaOBiOo";
        case dnnl::impl::wino_memory_format_t::wino_wei_OBaaIBOIio:
            return "wino_OBaaIBOIio";
        default:
            return "wino_undef";
        }
    } else if (wrapped.is_rnn_packed_desc()) {
        switch (desc.get()->format_desc.rnn_packed_desc.format) {
        case dnnl::impl::rnn_packed_format::ldigo_p:
            return "packed_ldigo";
        case dnnl::impl::rnn_packed_format::ldgoi_p:
            return "packed_ldgoi";
        case dnnl::impl::rnn_packed_format::ldio_p:
            return "packed_ldio";
        default:
            return "packed_undef";
        }
    }
    return "undef";
}

size_t DnnlMemoryDesc::getMaxMemSize() const {
    OPENVINO_ASSERT(!shape.isDynamic(), "Can't compute max mem size for DnnlMemoryDesc with dynamic shape");

    return getCurrentMemSize();
}

dnnl::memory::data_type DnnlMemoryDesc::getDataType() const {
    return desc.get_data_type();
}

dnnl::memory::format_kind DnnlMemoryDesc::getFormatKind() const {
    return desc.get_format_kind();
}

bool DnnlMemoryDesc::hasEmptyExtraData() const {
    dnnl::impl::memory_desc_wrapper wrapped(desc.get());
    return wrapped.extra().flags == dnnl_memory_extra_flag_none;
}

bool DnnlMemoryDesc::canComputeMemSizeZeroDims() const {
    if (!getShape().hasZeroDims()) {
        return false;
    }

    dnnl::impl::memory_desc_wrapper wrapped(desc.get());
    return getShape().hasZeroDims() && wrapped.offset0() != DNNL_RUNTIME_DIM_VAL;
}

size_t DnnlMemoryDesc::getCurrentMemSizeImp() const {
    return DnnlExtensionUtils::getMemSizeForDnnlDesc(desc);
}

size_t DnnlMemoryDesc::getElementOffset(size_t elemNumber) const {
    dnnl::impl::memory_desc_wrapper wrapped(desc.get());
    return wrapped.off_l(elemNumber);
}

bool DnnlMemoryDesc::isDefinedImp() const {
    dnnl::impl::memory_desc_wrapper wrappedThis(desc.get());

    if (wrappedThis.has_runtime_dims_or_strides()) {
        return false;
    }

    return wrappedThis.offset0() != DNNL_RUNTIME_DIM_VAL;
}

MemoryDescPtr DnnlMemoryDesc::cloneWithNewDimsImp([[maybe_unused]] const VectorDims& dims) const {
    OPENVINO_THROW("Unexpected: Cannot clone non blocked oneDNN desc with new dims");
}

size_t DnnlMemoryDesc::getOffsetPadding() const {
    dnnl::impl::memory_desc_wrapper wrap(desc.get());
    return DnnlExtensionUtils::convertToDim(wrap.offset0());
}

}  // namespace ov::intel_cpu
