Skip to content

Commit 1746930

Browse files
MasterJH5574MasterJH5574
authored
[C++] Update reflection block for C++ objects (#3364)
This PR updates the reflection block for C++ objects, following a recent upstream refactor, so that all objects will be allocated with a type index at runtime.
1 parent fa503e2 commit 1746930

31 files changed

+185
-5
lines changed

‎3rdparty/tvm‎

Submodule tvm updated 71 files

‎cpp/serve/config.cc‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,11 @@ namespace mlc {
2020
namespace llm {
2121
namespace serve {
2222

23+
TVM_FFI_STATIC_INIT_BLOCK() {
24+
GenerationConfigNode::RegisterReflection();
25+
EngineConfigNode::RegisterReflection();
26+
}
27+
2328
uint64_t TotalDetectGlobalMemory(DLDevice device) {
2429
// Get single-card GPU size.
2530
tvm::ffi::Any rv;

‎cpp/serve/config.h‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
#include <picojson.h>
99
#include <tvm/ffi/container/array.h>
10+
#include <tvm/ffi/reflection/registry.h>
1011
#include <tvm/ffi/string.h>
1112
#include <tvm/runtime/device_api.h>
1213
#include <tvm/runtime/int_tuple.h>
@@ -135,6 +136,11 @@ class GenerationConfigNode : public Object {
135136

136137
picojson::object AsJSON() const;
137138

139+
static void RegisterReflection() {
140+
namespace refl = tvm::ffi::reflection;
141+
refl::ObjectDef<GenerationConfigNode>();
142+
}
143+
138144
static constexpr const bool _type_has_method_sequal_reduce = false;
139145
static constexpr const bool _type_has_method_shash_reduce = false;
140146
TVM_FFI_DECLARE_OBJECT_INFO("mlc.serve.GenerationConfig", GenerationConfigNode, Object);
@@ -302,6 +308,11 @@ class EngineConfigNode : public Object {
302308

303309
String AsJSONString() const;
304310

311+
static void RegisterReflection() {
312+
namespace refl = tvm::ffi::reflection;
313+
refl::ObjectDef<EngineConfigNode>();
314+
}
315+
305316
static constexpr const bool _type_has_method_sequal_reduce = false;
306317
static constexpr const bool _type_has_method_shash_reduce = false;
307318
static constexpr const bool _type_mutable = true;

‎cpp/serve/data.cc‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,14 @@ namespace mlc {
1313
namespace llm {
1414
namespace serve {
1515

16+
TVM_FFI_STATIC_INIT_BLOCK() {
17+
DataNode::RegisterReflection();
18+
TextDataNode::RegisterReflection();
19+
TokenDataNode::RegisterReflection();
20+
ImageDataNode::RegisterReflection();
21+
RequestStreamOutputObj::RegisterReflection();
22+
}
23+
1624
/****************** Data ******************/
1725

1826
std::pair<Array<Data>, Array<Data>> SplitData(const Array<Data>& original_data, int total_length,

‎cpp/serve/data.h‎

Lines changed: 30 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
#include <tvm/ffi/container/array.h>
99
#include <tvm/ffi/container/shape.h>
1010
#include <tvm/ffi/optional.h>
11+
#include <tvm/ffi/reflection/registry.h>
1112
#include <tvm/ffi/string.h>
1213
#include <tvm/node/cast.h>
1314
#include <tvm/runtime/int_tuple.h>
@@ -50,8 +51,14 @@ class DataNode : public Object {
5051
*/
5152
virtual ObjectRef GetEmbedding(Model model, ObjectRef* dst = nullptr, int offset = 0) const = 0;
5253

54+
static void RegisterReflection() {
55+
namespace refl = tvm::ffi::reflection;
56+
refl::ObjectDef<DataNode>();
57+
}
58+
5359
static constexpr const bool _type_has_method_sequal_reduce = false;
5460
static constexpr const bool _type_has_method_shash_reduce = false;
61+
static constexpr const uint32_t _type_child_slots = 3;
5562
TVM_FFI_DECLARE_OBJECT_INFO("mlc.serve.Data", DataNode, Object);
5663
};
5764

@@ -75,7 +82,12 @@ class TextDataNode : public DataNode {
7582
int GetLength() const final;
7683
ObjectRef GetEmbedding(Model model, ObjectRef* dst = nullptr, int offset = 0) const final;
7784

78-
TVM_FFI_DECLARE_OBJECT_INFO("mlc.serve.TextData", TextDataNode, DataNode);
85+
static void RegisterReflection() {
86+
namespace refl = tvm::ffi::reflection;
87+
refl::ObjectDef<TextDataNode>();
88+
}
89+
90+
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("mlc.serve.TextData", TextDataNode, DataNode);
7991
};
8092

8193
class TextData : public Data {
@@ -96,7 +108,12 @@ class TokenDataNode : public DataNode {
96108
int GetLength() const final;
97109
ObjectRef GetEmbedding(Model model, ObjectRef* dst = nullptr, int offset = 0) const final;
98110

99-
TVM_FFI_DECLARE_OBJECT_INFO("mlc.serve.TokenData", TokenDataNode, DataNode);
111+
static void RegisterReflection() {
112+
namespace refl = tvm::ffi::reflection;
113+
refl::ObjectDef<TokenDataNode>();
114+
}
115+
116+
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("mlc.serve.TokenData", TokenDataNode, DataNode);
100117
};
101118

102119
class TokenData : public Data {
@@ -120,7 +137,12 @@ class ImageDataNode : public DataNode {
120137
int GetLength() const final;
121138
ObjectRef GetEmbedding(Model model, ObjectRef* dst = nullptr, int offset = 0) const final;
122139

123-
TVM_FFI_DECLARE_OBJECT_INFO("mlc.serve.ImageData", ImageDataNode, DataNode);
140+
static void RegisterReflection() {
141+
namespace refl = tvm::ffi::reflection;
142+
refl::ObjectDef<ImageDataNode>();
143+
}
144+
145+
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("mlc.serve.ImageData", ImageDataNode, DataNode);
124146
};
125147

126148
class ImageData : public Data {
@@ -197,6 +219,11 @@ class RequestStreamOutputObj : public Object {
197219

198220
std::atomic<bool> unpacked = false;
199221

222+
static void RegisterReflection() {
223+
namespace refl = tvm::ffi::reflection;
224+
refl::ObjectDef<RequestStreamOutputObj>();
225+
}
226+
200227
static constexpr const bool _type_has_method_sequal_reduce = false;
201228
static constexpr const bool _type_has_method_shash_reduce = false;
202229
static constexpr const bool _type_mutable = true;

‎cpp/serve/draft_token_workspace_manager.cc‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@ namespace mlc {
1111
namespace llm {
1212
namespace serve {
1313

14+
TVM_FFI_STATIC_INIT_BLOCK() { DraftTokenWorkspaceManagerObj::RegisterReflection(); }
15+
1416
DraftTokenWorkspaceManagerObj::DraftTokenWorkspaceManagerObj(int max_num_tokens, int vocab_size,
1517
int hidden_size,
1618
DLDataType hidden_states_dtype,

‎cpp/serve/draft_token_workspace_manager.h‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55

66
#ifndef MLC_LLM_SERVE_DRAFT_TOKEN_WORKSPACE_MANAGER_H_
77
#define MLC_LLM_SERVE_DRAFT_TOKEN_WORKSPACE_MANAGER_H_
8+
#include <tvm/ffi/reflection/registry.h>
89
#include <tvm/runtime/device_api.h>
910

1011
#include <numeric>
@@ -73,6 +74,11 @@ class DraftTokenWorkspaceManagerObj : public Object {
7374
*/
7475
void FreeSlots(const std::vector<int>& slots);
7576

77+
static void RegisterReflection() {
78+
namespace refl = tvm::ffi::reflection;
79+
refl::ObjectDef<DraftTokenWorkspaceManagerObj>();
80+
}
81+
7682
static constexpr const bool _type_has_method_sequal_reduce = false;
7783
static constexpr const bool _type_has_method_shash_reduce = false;
7884
static constexpr const bool _type_mutable = true;

‎cpp/serve/engine_actions/action.cc‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,10 @@
77

88
namespace mlc {
99
namespace llm {
10-
namespace serve {} // namespace serve
10+
namespace serve {
11+
12+
TVM_FFI_STATIC_INIT_BLOCK() { EngineActionObj::RegisterReflection(); }
13+
14+
} // namespace serve
1115
} // namespace llm
1216
} // namespace mlc

‎cpp/serve/engine_actions/action.h‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,11 @@ class EngineActionObj : public Object {
3838
*/
3939
virtual Array<Request> Step(EngineState estate) = 0;
4040

41+
static void RegisterReflection() {
42+
namespace refl = tvm::ffi::reflection;
43+
refl::ObjectDef<EngineActionObj>();
44+
}
45+
4146
static constexpr const bool _type_has_method_sequal_reduce = false;
4247
static constexpr const bool _type_has_method_shash_reduce = false;
4348
static constexpr const bool _type_mutable = true;

‎cpp/serve/engine_state.cc‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@ namespace mlc {
88
namespace llm {
99
namespace serve {
1010

11+
TVM_FFI_STATIC_INIT_BLOCK() { EngineStateObj::RegisterReflection(); }
12+
1113
EngineState::EngineState() { data_ = tvm::ffi::make_object<EngineStateObj>(); }
1214

1315
void EngineStateObj::Reset() {

0 commit comments

Comments
 (0)