// Copyright (c) 2020, Huawei Technologies Co., Ltd
// All rights reserved.
//
// This source code is licensed under the BSD-style license found in the
// LICENSE file in the root directory of this source tree.

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vllm-ascend project
// Adapted from
// https://github.com/vllm-project/vllm-ascend/blob/main/csrc/aclnn_torch_adapter/NPUStorageImpl.h

#pragma once

#include <ATen/Tensor.h>
#include <c10/core/Allocator.h>
#include <c10/core/ScalarType.h>
#include <c10/core/StorageImpl.h>
#include <c10/util/order_preserving_flat_hash_map.h>
#include <c10/util/typeid.h>
#include "acl/acl_base.h"
#include "acl/acl_rt.h"

namespace vllm_ascend {

struct NPUStorageDesc {
public:
    struct use_byte_size_t {};

    c10::SmallVector<int64_t, 5> base_sizes_;
    c10::SmallVector<int64_t, 5> base_strides_;
    c10::SmallVector<int64_t, 5> storage_sizes_;
    int64_t base_offset_ = 0;
    use_byte_size_t base_dtype_ = {};
    aclFormat origin_format_ = ACL_FORMAT_UNDEFINED;
    aclFormat npu_format_ = ACL_FORMAT_ND;
    // used to make CANN GE tensor from storagImpl
    caffe2::TypeMeta data_type_ = caffe2::TypeMeta::Make<uint8_t>();
};

struct NPUStorageImpl : public c10::StorageImpl {
    explicit NPUStorageImpl(use_byte_size_t use_byte_size, size_t size_bytes, at::DataPtr data_ptr,
                            at::Allocator* allocator, bool resizable);
    ~NPUStorageImpl() override = default;

    void release_resources() override;

    NPUStorageDesc npu_desc_;

    NPUStorageDesc get_npu_desc() const { return npu_desc_; }
};

c10::intrusive_ptr<c10::StorageImpl> make_npu_storage_impl(c10::StorageImpl::use_byte_size_t,
                                                           c10::SymInt size_bytes,
                                                           c10::DataPtr data_ptr,
                                                           c10::Allocator* allocator,
                                                           bool resizable);

}  // namespace vllm_ascend
