1/*
2 * Licensed to the Apache Software Foundation (ASF) under one
3 * or more contributor license agreements. See the NOTICE file
4 * distributed with this work for additional information
5 * regarding copyright ownership. The ASF licenses this file
6 * to you under the Apache License, Version 2.0 (the
7 * "License"); you may not use this file except in compliance
8 * with the License. You may obtain a copy of the License at
9 *
10 * http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing,
13 * software distributed under the License is distributed on an
14 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15 * KIND, either express or implied. See the License for the
16 * specific language governing permissions and limitations
17 * under the License.
18 */
19
20/*!
21 * \file tvm/runtime/metadata.h
22 * \brief Defines types which can be used in Metadata.
23 */
24#ifndef TVM_RUNTIME_METADATA_H_
25#define TVM_RUNTIME_METADATA_H_
26
27#include <dmlc/memory_io.h>
28#include <tvm/runtime/c_runtime_api.h>
29#include <tvm/runtime/metadata_base.h>
30#include <tvm/runtime/metadata_types.h>
31#include <tvm/runtime/object.h>
32#include <tvm/support/span.h>
33
34#include <memory>
35#include <string>
36#include <vector>
37
38// Version number recorded in emitted artifacts for runtime checking.
39#define TVM_METADATA_VERSION 1
40
41namespace tvm {
42namespace runtime {
43namespace metadata {
44/*!
45 * \brief Version of metadata emitted and understood by this compiler/runtime.
46 * Should be populated into the `version` field of all TVMMetadata.
47 */
48static const constexpr int64_t kMetadataVersion = TVM_METADATA_VERSION;
49
50class Metadata;
51class TensorInfo;
52class ConstantInfoMetadata;
53
54class MetadataNode : public MetadataBaseNode {
55 public:
56 explicit MetadataNode(const struct ::TVMMetadata* data) : data_{data} {}
57 static constexpr const char* _type_key = "metadata.MetadataNode";
58 const char* get_c_struct_name() const override;
59 inline int64_t version() const { return int64_t(data_->version); }
60 inline int64_t num_inputs() const { return data_->num_inputs; }
61 ArrayAccessor<struct TVMTensorInfo, TensorInfo> inputs();
62 inline int64_t num_outputs() const { return data_->num_outputs; }
63 ArrayAccessor<struct TVMTensorInfo, TensorInfo> outputs();
64 inline int64_t num_workspace_pools() const { return data_->num_workspace_pools; }
65 ArrayAccessor<struct TVMTensorInfo, TensorInfo> workspace_pools();
66 inline ::tvm::runtime::String mod_name() const { return ::tvm::runtime::String(data_->mod_name); }
67 const struct ::TVMMetadata* data() const { return data_; }
68 ArrayAccessor<struct TVMConstantInfo, ConstantInfoMetadata> constant_pools();
69 inline int64_t num_constant_pools() const { return data_->num_constant_pools; }
70 TVM_DECLARE_FINAL_OBJECT_INFO(MetadataNode, MetadataBaseNode);
71
72 private:
73 const struct ::TVMMetadata* data_;
74};
75
76class Metadata : public MetadataBase {
77 public:
78 explicit Metadata(const struct ::TVMMetadata* data);
79 TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(Metadata, MetadataBase, MetadataNode);
80};
81
82class TensorInfoNode : public MetadataBaseNode {
83 public:
84 explicit TensorInfoNode(const struct ::TVMTensorInfo* data) : data_{data} {}
85 static constexpr const char* _type_key = "metadata.TensorInfoNode";
86 const char* get_c_struct_name() const override;
87 inline ::tvm::runtime::String name() const { return ::tvm::runtime::String(data_->name); }
88 inline int64_t num_shape() const { return data_->num_shape; }
89 inline ::tvm::support::Span<const int64_t, int64_t> shape() const {
90 return ::tvm::support::Span<const int64_t, int64_t>(data_->shape,
91 data_->shape + data_->num_shape);
92 }
93 inline ::tvm::runtime::DataType dtype() const { return ::tvm::runtime::DataType(data_->dtype); }
94 const struct ::TVMTensorInfo* data() const { return data_; }
95 TVM_DECLARE_FINAL_OBJECT_INFO(TensorInfoNode, MetadataBaseNode);
96
97 private:
98 const struct ::TVMTensorInfo* data_;
99};
100
101class TensorInfo : public MetadataBase {
102 public:
103 explicit TensorInfo(const struct ::TVMTensorInfo* data);
104 TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(TensorInfo, MetadataBase, TensorInfoNode);
105};
106
107class ConstantInfoMetadataNode : public MetadataBaseNode {
108 public:
109 explicit ConstantInfoMetadataNode(const struct ::TVMConstantInfo* data) : data_{data} {}
110 // This name should match TVMConstantInfo after processing
111 static constexpr const char* _type_key = "metadata.ConstantInfoNode";
112 const char* get_c_struct_name() const override;
113 inline ::tvm::runtime::String name_hint() const {
114 return ::tvm::runtime::String(data_->name_hint);
115 }
116 inline size_t byte_offset() const { return data_->byte_offset; }
117 inline ::tvm::runtime::NDArray data() const {
118 ::tvm::runtime::NDArray ndarray;
119 if (data_->data_len) {
120 dmlc::MemoryFixedSizeStream bytes(const_cast<void*>(data_->data_bytes), data_->data_len);
121 ndarray.Load(&bytes);
122 }
123 return ndarray;
124 }
125 TVM_DECLARE_FINAL_OBJECT_INFO(ConstantInfoMetadataNode, MetadataBaseNode);
126
127 protected:
128 const struct ::TVMConstantInfo* data_;
129};
130
131class ConstantInfoMetadata : public MetadataBase {
132 public:
133 explicit ConstantInfoMetadata(const struct ::TVMConstantInfo* data);
134 TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(ConstantInfoMetadata, MetadataBase,
135 ConstantInfoMetadataNode);
136};
137
138} // namespace metadata
139} // namespace runtime
140} // namespace tvm
141
142#endif // TVM_RUNTIME_METADATA_H_
143