wehub-resource-sync
200 lines
6.9 KiB
C++
200 lines
6.9 KiB
C++
/*
|
|
* 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.
|
|
*/
|
|
#ifndef TVM_SCRIPT_PRINTER_IR_DOCSIFIER_FUNCTOR_H_
|
|
#define TVM_SCRIPT_PRINTER_IR_DOCSIFIER_FUNCTOR_H_
|
|
|
|
#include <tvm/ffi/function.h>
|
|
#include <tvm/runtime/logging.h>
|
|
|
|
#include <optional>
|
|
#include <string>
|
|
#include <type_traits>
|
|
#include <unordered_map>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
namespace tvm {
|
|
namespace script {
|
|
namespace printer {
|
|
|
|
/*!
|
|
* \brief Dynamic dispatch functor based on AccessPath.
|
|
*
|
|
* This functor dispatches based on the type of object and the input dispatch token.
|
|
*/
|
|
template <typename R, typename... Args>
|
|
class IRDocsifierFunctor {
|
|
private:
|
|
using TSelf = IRDocsifierFunctor<R, Args...>;
|
|
|
|
template <class TObjectRef, class TCallable>
|
|
using IsDispatchFunction =
|
|
typename std::is_convertible<TCallable, std::function<R(TObjectRef, Args...)>>;
|
|
|
|
public:
|
|
/*!
|
|
* \brief Call the dispatch function.
|
|
* \param token The dispatch token.
|
|
* \param obj The object.
|
|
* \param args Other args.
|
|
*
|
|
* \return The return value of the dispatch function
|
|
*
|
|
* If the TObjectRef isn't registered with the token, it will try to find
|
|
* dispatch function for TObjectRef with the default dispatch token (empty string).
|
|
*/
|
|
template <class TObjectRef>
|
|
R operator()(const ffi::String& token, TObjectRef obj, Args... args) const {
|
|
uint32_t type_index = obj.defined() ? obj->type_index() : 0;
|
|
const ffi::Function* pf = nullptr;
|
|
if ((pf = LookupDispatchTable(token, type_index)) != nullptr) {
|
|
return (*pf)(obj, args...).template cast<R>();
|
|
}
|
|
if ((pf = LookupDispatchTable("", type_index)) != nullptr) {
|
|
return (*pf)(obj, args...).template cast<R>();
|
|
}
|
|
if ((pf = LookupFallback()) != nullptr) {
|
|
return (*pf)(obj, args...).template cast<R>();
|
|
}
|
|
|
|
LOG(WARNING) << "ObjectFunctor calls un-registered function on type: "
|
|
<< ffi::Object::TypeIndex2Key(type_index) << " (token: " << token << ")"
|
|
<< ". ObjectType: " << obj->GetTypeKey() << ". Object: " << obj;
|
|
TVM_FFI_ICHECK(false) << "ObjectFunctor calls un-registered function on type: "
|
|
<< ffi::Object::TypeIndex2Key(type_index) << " (token: " << token << ")"
|
|
<< ". ObjectType: " << obj->GetTypeKey() << ". Object: " << obj;
|
|
TVM_FFI_UNREACHABLE();
|
|
}
|
|
|
|
/*!
|
|
* \brief Set the dispatch function
|
|
* \param token The dispatch token.
|
|
* \param type_index The TVM object type index for this dispatch function.
|
|
* \param f The dispatch function.
|
|
*
|
|
* This takes a type-erased packed function as input. It should be used
|
|
* through FFI boundary, for example, registering dispatch function from Python.
|
|
*/
|
|
TSelf& set_dispatch(ffi::String token, uint32_t type_index, ffi::Function f) {
|
|
std::vector<ffi::Function>* table = &dispatch_table_[token];
|
|
if (table->size() <= type_index) {
|
|
table->resize(type_index + 1, nullptr);
|
|
}
|
|
ffi::Function& slot = (*table)[type_index];
|
|
if (slot != nullptr) {
|
|
TVM_FFI_ICHECK(false) << "Dispatch for type is already registered: "
|
|
<< ffi::Object::TypeIndex2Key(type_index);
|
|
}
|
|
slot = f;
|
|
return *this;
|
|
}
|
|
|
|
TSelf& set_fallback(ffi::Function f) {
|
|
TVM_FFI_ICHECK(!dispatch_fallback_.has_value()) << "Fallback is already defined";
|
|
dispatch_fallback_ = f;
|
|
return *this;
|
|
}
|
|
|
|
void remove_fallback() { dispatch_fallback_ = std::nullopt; }
|
|
|
|
/*!
|
|
* \brief Set the dispatch function
|
|
* \param token The dispatch token.
|
|
* \param f The dispatch function.
|
|
*/
|
|
template <typename TObjectRef, typename TCallable,
|
|
typename = std::enable_if_t<IsDispatchFunction<TObjectRef, TCallable>::value>>
|
|
TSelf& set_dispatch(ffi::String token, TCallable f) {
|
|
return set_dispatch(token, TObjectRef::ContainerType::RuntimeTypeIndex(),
|
|
ffi::TypedFunction<R(TObjectRef, Args...)>(f));
|
|
}
|
|
|
|
template <typename TCallable,
|
|
typename = std::enable_if_t<IsDispatchFunction<ffi::ObjectRef, TCallable>::value>>
|
|
TSelf& set_fallback(TCallable f) {
|
|
ffi::Function func = ffi::TypedFunction<R(ffi::ObjectRef, Args...)>(f);
|
|
return set_fallback(func);
|
|
}
|
|
|
|
/*!
|
|
* \brief Remove dispatch function
|
|
* \param token The dispatch token.
|
|
* \param type_index The TVM object type index for the dispatch function to be removed.
|
|
*
|
|
* This is useful when dispatch function comes from other language's runtime, and
|
|
* those function should be removed before that language runtime shuts down.
|
|
*/
|
|
void remove_dispatch(ffi::String token, uint32_t type_index) {
|
|
std::vector<ffi::Function>* table = &dispatch_table_[token];
|
|
if (table->size() <= type_index) {
|
|
return;
|
|
}
|
|
(*table)[type_index] = nullptr;
|
|
}
|
|
|
|
private:
|
|
/*!
|
|
* \brief Look up the dispatch table for the given token and type_index.
|
|
* \param token The dispatch token.
|
|
* \param type_index The TVM object type index.
|
|
* \return Returns the functor if the lookup succeeds, nullptr otherwise.
|
|
*/
|
|
const ffi::Function* LookupDispatchTable(const ffi::String& token, uint32_t type_index) const {
|
|
auto it = dispatch_table_.find(token);
|
|
if (it == dispatch_table_.end()) {
|
|
return nullptr;
|
|
}
|
|
const std::vector<ffi::Function>& tab = it->second;
|
|
if (type_index >= tab.size()) {
|
|
return nullptr;
|
|
}
|
|
const ffi::Function* f = &tab[type_index];
|
|
if (f->defined()) {
|
|
return f;
|
|
} else {
|
|
return nullptr;
|
|
}
|
|
}
|
|
|
|
/*!
|
|
* \brief Look up the fallback to be used if no handler is registered
|
|
*/
|
|
const ffi::Function* LookupFallback() const {
|
|
if (dispatch_fallback_.has_value()) {
|
|
return &*dispatch_fallback_;
|
|
} else {
|
|
return nullptr;
|
|
}
|
|
}
|
|
|
|
/*
|
|
* This type alias and the following free functions are created to reduce the binary bloat
|
|
* from template and also hide implementation details from this header
|
|
*/
|
|
using DispatchTable = std::unordered_map<std::string, std::vector<ffi::Function>>;
|
|
/*! \brief The dispatch table. */
|
|
DispatchTable dispatch_table_;
|
|
std::optional<ffi::Function> dispatch_fallback_;
|
|
};
|
|
|
|
} // namespace printer
|
|
} // namespace script
|
|
} // namespace tvm
|
|
#endif // TVM_SCRIPT_PRINTER_IR_DOCSIFIER_FUNCTOR_H_
|