Skip to content

Commit addd67d

Browse files
BlealtanBlealtan
committed
Add stubgen fields; first commit to build TileLang IR
1 parent be5076f commit addd67d

File tree

7 files changed

+430
-108
lines changed

7 files changed

+430
-108
lines changed

‎rust/tvm-ffi-macros/src/object_macros.rs‎

Lines changed: 196 additions & 106 deletions
Original file line numberDiff line numberDiff line change
@@ -124,131 +124,221 @@ pub fn derive_object_ref(input: proc_macro::TokenStream) -> TokenStream {
124124
}
125125
.expect("First field must be `data: ObjectArc<T>`");
126126

127-
let mut expanded = quote! {
128-
unsafe impl #tvm_ffi_crate::object::ObjectRefCore for #struct_name {
129-
type ContainerType = <#data_ty as std::ops::Deref>::Target;
130-
#[inline]
131-
fn data(this: &Self) -> &ObjectArc<Self::ContainerType> {
132-
&this.data
133-
}
134-
#[inline]
135-
fn into_data(this: Self) -> ObjectArc<Self::ContainerType> {
136-
this.data
137-
}
138-
#[inline]
139-
fn from_data(data: ObjectArc<Self::ContainerType>) -> Self {
140-
Self { data}
141-
}
142-
}
127+
let is_object_ref = struct_name == syn::Ident::new("ObjectRef", struct_name.span());
143128

144-
// implement AnyCompatible for #struct_name
145-
unsafe impl #tvm_ffi_crate::type_traits::AnyCompatible for #struct_name {
146-
fn type_str() -> String {
147-
type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore>
148-
::ContainerType;
149-
<ContainerType as #tvm_ffi_crate::object::ObjectCore>::TYPE_KEY.into()
150-
}
129+
let any_compatible_tokens = if is_object_ref {
130+
quote! {
131+
// implement AnyCompatible for #struct_name
132+
unsafe impl #tvm_ffi_crate::type_traits::AnyCompatible for #struct_name {
133+
fn type_str() -> String {
134+
type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore>
135+
::ContainerType;
136+
<ContainerType as #tvm_ffi_crate::object::ObjectCore>::TYPE_KEY.into()
137+
}
151138

152-
unsafe fn copy_to_any_view(
153-
src: &Self,
154-
data: &mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny
155-
) {
156-
type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore>
157-
::ContainerType;
158-
let type_index =
159-
<ContainerType as #tvm_ffi_crate::object::ObjectCore>::type_index();
160-
data.type_index = type_index as i32;
161-
data.small_str_len = 0;
162-
let data_ptr = #tvm_ffi_crate::object::ObjectArc::<ContainerType>::as_raw(
163-
&src.data
164-
);
165-
data.data_union.v_obj =
166-
data_ptr as *mut ContainerType as *mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIObject;
167-
}
139+
unsafe fn copy_to_any_view(
140+
src: &Self,
141+
data: &mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny
142+
) {
143+
type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore>
144+
::ContainerType;
145+
let data_ptr = #tvm_ffi_crate::object::ObjectArc::<ContainerType>::as_raw(
146+
&src.data
147+
) as *mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIObject;
148+
data.type_index = (*data_ptr).type_index;
149+
data.small_str_len = 0;
150+
data.data_union.v_obj = data_ptr;
151+
}
168152

169-
unsafe fn check_any_strict(data: & #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny) -> bool {
170-
type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore>
171-
::ContainerType;
172-
let target_index =
173-
<ContainerType as #tvm_ffi_crate::object::ObjectCore>::type_index();
174-
if data.type_index == target_index as i32 {
175-
return true;
153+
unsafe fn check_any_strict(data: & #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny) -> bool {
154+
data.type_index >= #tvm_ffi_crate::TypeIndex::kTVMFFIStaticObjectBegin as i32
176155
}
177-
let info = #tvm_ffi_crate::tvm_ffi_sys::TVMFFIGetTypeInfo(data.type_index);
178-
if info.is_null() {
179-
return false;
156+
157+
unsafe fn copy_from_any_view_after_check(
158+
data: & #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny
159+
) -> Self {
160+
type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore>
161+
::ContainerType;
162+
let data_ptr = data.data_union.v_obj;
163+
// need to increase ref because original weak ptr
164+
// do not own the code
165+
#tvm_ffi_crate::object::unsafe_::inc_ref(
166+
data_ptr as *mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIObject
167+
);
168+
Self {
169+
data : #tvm_ffi_crate::object::ObjectArc::from_raw(
170+
data_ptr as *mut ContainerType
171+
)
172+
}
180173
}
181-
let info = &*info;
182-
let ancestors = info.type_acenstors;
183-
if ancestors.is_null() {
184-
return false;
174+
175+
unsafe fn move_to_any(
176+
src: Self,
177+
data: &mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny
178+
) {
179+
type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore>
180+
::ContainerType;
181+
let data_ptr = #tvm_ffi_crate::object::ObjectArc::into_raw(
182+
src.data
183+
) as *mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIObject;
184+
data.type_index = (*data_ptr).type_index;
185+
data.small_str_len = 0;
186+
data.data_union.v_obj = data_ptr;
185187
}
186-
for depth in 0..info.type_depth {
187-
let ancestor = *ancestors.add(depth as usize);
188-
if !ancestor.is_null() && (*ancestor).type_index == target_index {
189-
return true;
188+
189+
unsafe fn move_from_any_after_check(
190+
data: &mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny
191+
) -> Self {
192+
type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore>
193+
::ContainerType;
194+
let data_ptr = data.data_union.v_obj as *mut ContainerType;
195+
Self {
196+
data : #tvm_ffi_crate::object::ObjectArc::<ContainerType>::from_raw(data_ptr)
190197
}
191198
}
192-
false
193-
}
194199

195-
unsafe fn copy_from_any_view_after_check(
196-
data: & #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny
197-
) -> Self {
198-
type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore>
199-
::ContainerType;
200-
let data_ptr = data.data_union.v_obj;
201-
// need to increase ref because original weak ptr
202-
// do not own the code
203-
#tvm_ffi_crate::object::unsafe_::inc_ref(
204-
data_ptr as *mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIObject
205-
);
206-
Self {
207-
data : #tvm_ffi_crate::object::ObjectArc::from_raw(
208-
data_ptr as *mut ContainerType
209-
)
200+
unsafe fn try_cast_from_any_view(
201+
data: & #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny
202+
) -> Result<Self, ()> {
203+
if Self::check_any_strict(data) {
204+
Ok(Self::copy_from_any_view_after_check(data))
205+
} else {
206+
Err(())
207+
}
210208
}
211209
}
210+
}
211+
} else {
212+
quote! {
213+
// implement AnyCompatible for #struct_name
214+
unsafe impl #tvm_ffi_crate::type_traits::AnyCompatible for #struct_name {
215+
fn type_str() -> String {
216+
type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore>
217+
::ContainerType;
218+
<ContainerType as #tvm_ffi_crate::object::ObjectCore>::TYPE_KEY.into()
219+
}
212220

213-
unsafe fn move_to_any(
214-
src: Self,
215-
data: &mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny
216-
) {
217-
type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore>
218-
::ContainerType;
219-
let type_index =
220-
<ContainerType as #tvm_ffi_crate::object::ObjectCore>::type_index();
221-
data.type_index = type_index as i32;
222-
data.small_str_len = 0;
223-
let data_ptr = #tvm_ffi_crate::object::ObjectArc::into_raw(
224-
src.data
225-
);
226-
data.data_union.v_obj =
227-
data_ptr as *mut ContainerType as *mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIObject;
228-
}
221+
unsafe fn copy_to_any_view(
222+
src: &Self,
223+
data: &mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny
224+
) {
225+
type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore>
226+
::ContainerType;
227+
let type_index =
228+
<ContainerType as #tvm_ffi_crate::object::ObjectCore>::type_index();
229+
data.type_index = type_index as i32;
230+
data.small_str_len = 0;
231+
let data_ptr = #tvm_ffi_crate::object::ObjectArc::<ContainerType>::as_raw(
232+
&src.data
233+
);
234+
data.data_union.v_obj =
235+
data_ptr as *mut ContainerType as *mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIObject;
236+
}
237+
238+
unsafe fn check_any_strict(data: & #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny) -> bool {
239+
type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore>
240+
::ContainerType;
241+
let target_index =
242+
<ContainerType as #tvm_ffi_crate::object::ObjectCore>::type_index();
243+
if data.type_index == target_index as i32 {
244+
return true;
245+
}
246+
let info = #tvm_ffi_crate::tvm_ffi_sys::TVMFFIGetTypeInfo(data.type_index);
247+
if info.is_null() {
248+
return false;
249+
}
250+
let info = &*info;
251+
let ancestors = info.type_acenstors;
252+
if ancestors.is_null() {
253+
return false;
254+
}
255+
for depth in 0..info.type_depth {
256+
let ancestor = *ancestors.add(depth as usize);
257+
if !ancestor.is_null() && (*ancestor).type_index == target_index {
258+
return true;
259+
}
260+
}
261+
false
262+
}
229263

230-
unsafe fn move_from_any_after_check(
231-
data: &mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny
232-
) -> Self {
233-
type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore>
234-
::ContainerType;
235-
let data_ptr = data.data_union.v_obj as *mut ContainerType;
236-
Self {
237-
data : #tvm_ffi_crate::object::ObjectArc::<ContainerType>::from_raw(data_ptr)
264+
unsafe fn copy_from_any_view_after_check(
265+
data: & #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny
266+
) -> Self {
267+
type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore>
268+
::ContainerType;
269+
let data_ptr = data.data_union.v_obj;
270+
// need to increase ref because original weak ptr
271+
// do not own the code
272+
#tvm_ffi_crate::object::unsafe_::inc_ref(
273+
data_ptr as *mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIObject
274+
);
275+
Self {
276+
data : #tvm_ffi_crate::object::ObjectArc::from_raw(
277+
data_ptr as *mut ContainerType
278+
)
279+
}
238280
}
239-
}
240281

241-
unsafe fn try_cast_from_any_view(
242-
data: & #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny
243-
) -> Result<Self, ()> {
244-
if Self::check_any_strict(data) {
245-
Ok(Self::copy_from_any_view_after_check(data))
246-
} else {
247-
Err(())
282+
unsafe fn move_to_any(
283+
src: Self,
284+
data: &mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny
285+
) {
286+
type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore>
287+
::ContainerType;
288+
let type_index =
289+
<ContainerType as #tvm_ffi_crate::object::ObjectCore>::type_index();
290+
data.type_index = type_index as i32;
291+
data.small_str_len = 0;
292+
let data_ptr = #tvm_ffi_crate::object::ObjectArc::into_raw(
293+
src.data
294+
);
295+
data.data_union.v_obj =
296+
data_ptr as *mut ContainerType as *mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIObject;
297+
}
298+
299+
unsafe fn move_from_any_after_check(
300+
data: &mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny
301+
) -> Self {
302+
type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore>
303+
::ContainerType;
304+
let data_ptr = data.data_union.v_obj as *mut ContainerType;
305+
Self {
306+
data : #tvm_ffi_crate::object::ObjectArc::<ContainerType>::from_raw(data_ptr)
307+
}
308+
}
309+
310+
unsafe fn try_cast_from_any_view(
311+
data: & #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny
312+
) -> Result<Self, ()> {
313+
if Self::check_any_strict(data) {
314+
Ok(Self::copy_from_any_view_after_check(data))
315+
} else {
316+
Err(())
317+
}
248318
}
249319
}
250320
}
251321
};
322+
323+
let mut expanded = quote! {
324+
unsafe impl #tvm_ffi_crate::object::ObjectRefCore for #struct_name {
325+
type ContainerType = <#data_ty as std::ops::Deref>::Target;
326+
#[inline]
327+
fn data(this: &Self) -> &ObjectArc<Self::ContainerType> {
328+
&this.data
329+
}
330+
#[inline]
331+
fn into_data(this: Self) -> ObjectArc<Self::ContainerType> {
332+
this.data
333+
}
334+
#[inline]
335+
fn from_data(data: ObjectArc<Self::ContainerType>) -> Self {
336+
Self { data}
337+
}
338+
}
339+
340+
#any_compatible_tokens
341+
};
252342
// skip ObjectRef since it can create circular dependency with any.rs
253343
if struct_name != "ObjectRef" {
254344
expanded.extend(quote! {

0 commit comments

Comments
 (0)