ir_docsifier_functor.h

Go to the documentation of this file.

/*
 * 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/node/node.h>
#include <tvm/runtime/logging.h>
#include <tvm/runtime/packed_func.h>

#include <string>
#include <type_traits>
#include <unordered_map>
#include <utility>
#include <vector>

namespace tvm {
namespace script {
namespace printer {

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:
  template <class TObjectRef>
  R operator()(const String& token, TObjectRef obj, Args... args) const {
    uint32_t type_index = obj.defined() ? obj->type_index() : 0;
    const runtime::PackedFunc* pf = nullptr;
    if ((pf = LookupDispatchTable(token, type_index)) != nullptr) {
      return (*pf)(obj, args...);
    }
    if ((pf = LookupDispatchTable("", type_index)) != nullptr) {
      return (*pf)(obj, args...);
    }
    LOG(WARNING) << "ObjectFunctor calls un-registered function on type: "
                 << runtime::Object::TypeIndex2Key(type_index) << " (token: " << token << ")"
                 << ". ObjectType: " << obj->GetTypeKey() << ". Object: " << obj;
    ICHECK(false) << "ObjectFunctor calls un-registered function on type: "
                  << runtime::Object::TypeIndex2Key(type_index) << " (token: " << token << ")"
                  << ". ObjectType: " << obj->GetTypeKey() << ". Object: " << obj;
  }

  TSelf& set_dispatch(String token, uint32_t type_index, runtime::PackedFunc f) {
    std::vector<runtime::PackedFunc>* table = &dispatch_table_[token];
    if (table->size() <= type_index) {
      table->resize(type_index + 1, nullptr);
    }
    runtime::PackedFunc& slot = (*table)[type_index];
    if (slot != nullptr) {
      ICHECK(false) << "Dispatch for type is already registered: "
                    << runtime::Object::TypeIndex2Key(type_index);
    }
    slot = f;
    return *this;
  }

  template <typename TObjectRef, typename TCallable,
            typename = std::enable_if_t<IsDispatchFunction<TObjectRef, TCallable>::value>>
  TSelf& set_dispatch(String token, TCallable f) {
    return set_dispatch(token, TObjectRef::ContainerType::RuntimeTypeIndex(),
                        runtime::TypedPackedFunc<R(TObjectRef, Args...)>(f));
  }

  void remove_dispatch(String token, uint32_t type_index) {
    std::vector<runtime::PackedFunc>* table = &dispatch_table_[token];
    if (table->size() <= type_index) {
      return;
    }
    (*table)[type_index] = nullptr;
  }

 private:
  const runtime::PackedFunc* LookupDispatchTable(const String& token, uint32_t type_index) const {
    auto it = dispatch_table_.find(token);
    if (it == dispatch_table_.end()) {
      return nullptr;
    }
    const std::vector<runtime::PackedFunc>& tab = it->second;
    if (type_index >= tab.size()) {
      return nullptr;
    }
    const PackedFunc* f = &tab[type_index];
    if (f->defined()) {
      return f;
    } 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<runtime::PackedFunc>>;
  DispatchTable dispatch_table_;
};

}  // namespace printer
}  // namespace script
}  // namespace tvm
#endif  // TVM_SCRIPT_PRINTER_IR_DOCSIFIER_FUNCTOR_H_