junrushao commented on code in PR #454: URL: https://github.com/apache/tvm-ffi/pull/454#discussion_r2820949850
########## src/ffi/extra/repr_print.cc: ########## @@ -0,0 +1,478 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +/* + * \file src/ffi/extra/repr_print.cc + * + * \brief Reflection-based repr printing with BFS-based cycle/DAG handling. + */ +#include <tvm/ffi/any.h> +#include <tvm/ffi/container/array.h> +#include <tvm/ffi/container/list.h> +#include <tvm/ffi/container/map.h> +#include <tvm/ffi/container/shape.h> +#include <tvm/ffi/container/tensor.h> +#include <tvm/ffi/dtype.h> +#include <tvm/ffi/error.h> +#include <tvm/ffi/reflection/accessor.h> +#include <tvm/ffi/reflection/registry.h> + +#include <iomanip> +#include <sstream> +#include <string> +#include <unordered_map> +#include <unordered_set> +#include <vector> + +namespace tvm { +namespace ffi { + +namespace { + +/*! + * \brief Convert a DLDeviceType to a short name string. + */ +const char* DeviceTypeName(int device_type) { + switch (device_type) { + case kDLCPU: + return "cpu"; + case kDLCUDA: + return "cuda"; + case kDLCUDAHost: + return "cuda_host"; + case kDLOpenCL: + return "opencl"; + case kDLVulkan: + return "vulkan"; + case kDLMetal: + return "metal"; + case kDLVPI: + return "vpi"; + case kDLROCM: + return "rocm"; + case kDLROCMHost: + return "rocm_host"; + case kDLExtDev: + return "ext_dev"; + case kDLCUDAManaged: + return "cuda_managed"; + case kDLOneAPI: + return "oneapi"; + case kDLWebGPU: + return "webgpu"; + case kDLHexagon: + return "hexagon"; + default: + return "unknown"; + } +} + +/*! + * \brief Format a DLDevice as "device_name:device_id". + */ +std::string DeviceToString(DLDevice device) { + std::ostringstream os; + os << DeviceTypeName(device.device_type) << ":" << device.device_id; + return os.str(); +} + +/*! + * \brief Format raw bytes as a Python-style bytes literal: b"...". + */ +std::string FormatBytes(const char* data, size_t size) { + std::ostringstream os; + os << "b\""; + for (size_t i = 0; i < size; ++i) { + unsigned char c = static_cast<unsigned char>(data[i]); + if (c >= 32 && c < 127 && c != '\"' && c != '\\') { + os << static_cast<char>(c); + } else { + os << "\\x" << std::hex << std::setw(2) << std::setfill('0') << static_cast<int>(c); + } + } + os << "\""; + return os.str(); +} + +/*! + * \brief Format an object address as a hex string. + */ +std::string AddressStr(const Object* obj) { + std::ostringstream os; + os << "0x" << std::hex << reinterpret_cast<uintptr_t>(obj); + return os.str(); +} + +/*! + * \brief Get the type key of an object as a std::string. + */ +std::string GetTypeKeyStr(const Object* obj) { + const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(obj->type_index()); + return std::string(type_info->type_key.data, type_info->type_key.size); +} + +/*! + * \brief Lazily initialize and return the __ffi_repr__ TypeAttrColumn. + * + * Returns nullptr if the column does not exist (e.g., before registration). + */ +const TVMFFITypeAttrColumn* GetReprColumn() { + static const TVMFFITypeAttrColumn* column = []() -> const TVMFFITypeAttrColumn* { + TVMFFIByteArray name = {reflection::type_attr::kRepr, + std::char_traits<char>::length(reflection::type_attr::kRepr)}; + return TVMFFIGetTypeAttrColumn(&name); + }(); + return column; +} + +/*! + * \brief Look up a type attribute from a column by type_index. + * \return AnyView of the attribute, or a None AnyView if not found. + */ +AnyView LookupTypeAttr(const TVMFFITypeAttrColumn* column, int32_t type_index) { + if (column == nullptr) return AnyView(); + size_t tindex = static_cast<size_t>(type_index); + if (tindex >= column->size) return AnyView(); + const AnyView* data = reinterpret_cast<const AnyView*>(column->data); + return data[tindex]; +} + +/*! + * \brief BFS-based repr printer. + * + * Algorithm: + * 1. BFS collect all objects reachable from root (tracking visit count). + * 2. Process in reverse BFS order (leaves first), building repr strings. + * 3. Objects encountered more than once use short-form on second+ occurrence. + * + * Modeled after ObjectDeepCopier in deep_copy.cc. + */ +class ReprPrinter { + public: + String Run(const Any& value) { + // Phase 1: BFS collection + CollectAny(value); + // Phase 2: Reverse-BFS processing + for (int64_t i = static_cast<int64_t>(bfs_queue_.size()) - 1; i >= 0; --i) { + ProcessNode(bfs_queue_[i]); + } + // Phase 3: Return repr of root + return String(ReprOfAny(value)); + } + + private: + // ---------- Phase 1: BFS Collection ---------- + + void CollectAny(const Any& value) { + int32_t ti = value.type_index(); + if (ti < TypeIndex::kTVMFFIStaticObjectBegin) return; + const Object* obj = static_cast<const Object*>(value.as<Object>()); + if (obj == nullptr) return; + // Track visit count + auto [it, inserted] = visit_count_.emplace(obj, 1); + if (!inserted) { + it->second++; + return; // Already visited + } + bfs_queue_.push_back(obj); + CollectChildren(obj, ti); + } + + void CollectChildren(const Object* obj, int32_t ti) { + switch (ti) { + case TypeIndex::kTVMFFIStr: + case TypeIndex::kTVMFFIBytes: + case TypeIndex::kTVMFFIShape: + case TypeIndex::kTVMFFITensor: + case TypeIndex::kTVMFFIFunction: + case TypeIndex::kTVMFFIError: + case TypeIndex::kTVMFFIOpaquePyObject: + // Leaf types: no children + break; + case TypeIndex::kTVMFFIArray: { + const ArrayObj* arr = static_cast<const ArrayObj*>(obj); + for (const Any& elem : *arr) CollectAny(elem); + break; + } + case TypeIndex::kTVMFFIList: { + const ListObj* lst = static_cast<const ListObj*>(obj); + for (const Any& elem : *lst) CollectAny(elem); + break; + } + case TypeIndex::kTVMFFIMap: { + const MapObj* map = static_cast<const MapObj*>(obj); + for (const auto& [k, v] : *map) { + CollectAny(k); + CollectAny(v); + } + break; + } + default: + // User-defined object: collect via reflection fields + CollectFieldChildren(obj); + break; + } + } + + void CollectFieldChildren(const Object* obj) { + const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(obj->type_index()); + if (type_info == nullptr) return; + reflection::ForEachFieldInfo(type_info, [&](const TVMFFIFieldInfo* finfo) { + if (finfo->flags & kTVMFFIFieldFlagBitMaskReprOff) return; + reflection::FieldGetter getter(finfo); + Any fv = getter(obj); + if (fv.type_index() >= TypeIndex::kTVMFFIStaticObjectBegin) { + CollectAny(fv); + } + }); + } + + // ---------- Phase 2: Processing ---------- + + void ProcessNode(const Object* obj) { + if (repr_map_.count(obj)) return; // Already processed + + int32_t ti = obj->type_index(); + AnyView custom_repr = LookupTypeAttr(GetReprColumn(), ti); + + if (custom_repr != nullptr) { + // Custom __ffi_repr__: call it with fn_repr callback + Function repr_fn = custom_repr.cast<Function>(); + Function fn_repr = CreateFnRepr(); + String result = repr_fn(obj, fn_repr).cast<String>(); + repr_map_[obj] = std::string(result.data(), result.size()); + } else { + // Generic reflection-based repr + repr_map_[obj] = GenericRepr(obj); + } + } + + Function CreateFnRepr() { + return Function::FromTyped( + [this](AnyView value) -> String { return String(ReprOfAny(Any(value))); }); + } + + // ---------- Repr Helpers ---------- + + std::string ReprOfAny(const Any& value) { + int32_t ti = value.type_index(); + switch (ti) { + case TypeIndex::kTVMFFINone: + return "None"; + case TypeIndex::kTVMFFIBool: + return value.cast<bool>() ? "True" : "False"; + case TypeIndex::kTVMFFIInt: + return std::to_string(value.cast<int64_t>()); + case TypeIndex::kTVMFFIFloat: { + std::ostringstream os; + os << value.cast<double>(); + return os.str(); + } + case TypeIndex::kTVMFFIDataType: { + String s = DLDataTypeToString(value.cast<DLDataType>()); + return std::string(s.data(), s.size()); + } + case TypeIndex::kTVMFFIDevice: { + return DeviceToString(value.cast<DLDevice>()); + } + default: + break; + } + if (ti == TypeIndex::kTVMFFISmallStr) { + String s = value.cast<String>(); + return "\"" + std::string(s.data(), s.size()) + "\""; + } + if (ti == TypeIndex::kTVMFFISmallBytes) { + Bytes b = value.cast<Bytes>(); + return FormatBytes(b.data(), b.size()); + } + if (ti < TypeIndex::kTVMFFIStaticObjectBegin) { + // Other POD types + return value.GetTypeKey(); + } + // Object type + const Object* obj = static_cast<const Object*>(value.as<Object>()); + if (obj == nullptr) return "None"; + auto it = repr_map_.find(obj); + if (it != repr_map_.end()) { + auto [_, first_use] = rendered_.insert(obj); + if (!first_use) { + // Already rendered once — use short form for duplicate references + return ShortRepr(obj); + } + return it->second; + } + // Not yet processed -- use short form + return ShortRepr(obj); + } + + std::string ShortRepr(const Object* obj) { return GetTypeKeyStr(obj) + "@" + AddressStr(obj); } + + std::string GenericRepr(const Object* obj) { + const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(obj->type_index()); + if (type_info == nullptr) { + return ShortRepr(obj); + } + std::string type_key(type_info->type_key.data, type_info->type_key.size); + std::ostringstream os; + os << type_key << "@" << AddressStr(obj) << "("; + + bool first = true; + bool has_fields = false; + reflection::ForEachFieldInfo(type_info, [&](const TVMFFIFieldInfo* finfo) { + if (finfo->flags & kTVMFFIFieldFlagBitMaskReprOff) return; + has_fields = true; + if (!first) os << ", "; + first = false; + os << std::string_view(finfo->name.data, finfo->name.size) << "="; + reflection::FieldGetter getter(finfo); + Any fv = getter(obj); + os << ReprOfAny(fv); + }); + if (!has_fields) { + return type_key + "@" + AddressStr(obj); + } + os << ")"; + return os.str(); + } + + // ---------- Data members ---------- + std::vector<const Object*> bfs_queue_; + std::unordered_map<const Object*, int> visit_count_; + std::unordered_map<const Object*, std::string> repr_map_; + std::unordered_set<const Object*> rendered_; +}; + +// ---------- Built-in __ffi_repr__ functions ---------- + +String ReprString(const details::StringObj* obj, const Function& fn_repr) { + std::ostringstream os; + os << "\"" << std::string_view(obj->data, obj->size) << "\""; + return String(os.str()); +} + +String ReprBytes(const details::BytesObj* obj, const Function& fn_repr) { + return String(FormatBytes(obj->data, obj->size)); +} + +String ReprTensor(const TensorObj* obj, const Function& fn_repr) { Review Comment: I updated the syntax. `@0x{addr}` shows up only when `TVM_FFI_REPR_WITH_ADDR` is set to `1`. It means there's no point of confusion -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
