node_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_IR_NODE_FUNCTOR_H_
#define TVM_IR_NODE_FUNCTOR_H_
 
#include <tvm/ffi/error.h>
 
#include <cstring>
#include <type_traits>
#include <utility>
#include <vector>
 
namespace tvm {
 
template <typename FType>
class NodeFunctor;
 
template <typename R, typename... Args>
class NodeFunctor<R(const ffi::ObjectRef& n, Args...)> {
 private:
  typedef R (*FPointer)(const ffi::ObjectRef& n, Args...);
  using TSelf = NodeFunctor<R(const ffi::ObjectRef& n, Args...)>;
  std::vector<FPointer> func_;
  uint32_t begin_type_index_{0};
 
 public:
  using result_type = R;
  bool can_dispatch(const ffi::ObjectRef& n) const {
    uint32_t type_index = n->type_index();
    if (type_index < begin_type_index_) return false;
    type_index -= begin_type_index_;
    return type_index < func_.size() && func_[type_index] != nullptr;
  }
  R operator()(const ffi::ObjectRef& n, Args... args) const {
    uint32_t type_index = n->type_index();
    if (type_index >= begin_type_index_) {
      uint32_t index = type_index - begin_type_index_;
      if (index < func_.size() && func_[index] != nullptr) {
        return (*func_[index])(n, std::forward<Args>(args)...);
      }
    }
 
    const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(type_index);
    for (int32_t i = type_info->type_depth - 1; i >= 0; --i) {
      type_index = type_info->type_ancestors[i]->type_index;
      if (type_index >= begin_type_index_) {
        uint32_t index = type_index - begin_type_index_;
        if (index < func_.size() && func_[index] != nullptr) {
          return (*func_[index])(n, std::forward<Args>(args)...);
        }
      }
    }
    TVM_FFI_THROW(InternalError) << "NodeFunctor calls un-registered function on type "
                                 << n->GetTypeKey();
    throw;
  }
  template <typename TNode>
  TSelf& set_dispatch(FPointer f) {  // NOLINT(*)
    uint32_t tindex = TNode::RuntimeTypeIndex();
    if (func_.size() <= tindex) {
      func_.resize(tindex + 1, nullptr);
    }
    TVM_FFI_ICHECK(func_[tindex] == nullptr)
        << "Dispatch for " << TNode::_type_key << " is already set";
    TVM_FFI_ICHECK_EQ(begin_type_index_, 0) << " Cannot call set_dispatch after calling Finalize";
    func_[tindex] = f;
    return *this;
  }
  template <typename TNode>
  TSelf& clear_dispatch() {  // NOLINT(*)
    uint32_t tindex = TNode::RuntimeTypeIndex();
    TVM_FFI_ICHECK_LT(tindex, func_.size()) << "clear_dispatch: index out of range";
    TVM_FFI_ICHECK_EQ(begin_type_index_, 0) << " Cannot call clear_dispatch after calling Finalize";
    func_[tindex] = nullptr;
    return *this;
  }
  void Finalize() {
    TVM_FFI_ICHECK_EQ(begin_type_index_, 0) << "Can only call Finalize once";
    while (begin_type_index_ < func_.size() && func_[begin_type_index_] == nullptr) {
      ++begin_type_index_;
    }
    // shift up the function value
    size_t new_ftable_size = func_.size() - begin_type_index_;
    if (begin_type_index_ != 0) {
      std::memmove(func_.data(), func_.data() + begin_type_index_,
                   new_ftable_size * sizeof(FPointer));
    }
    func_.resize(new_ftable_size);
    func_.shrink_to_fit();
  }
};
 
#define TVM_REG_FUNC_VAR_DEF(ClsName) [[maybe_unused]] static auto& __make_functor##_##ClsName
 
#define TVM_STATIC_IR_FUNCTOR(ClsName, FField) \
  TVM_FFI_STR_CONCAT(TVM_REG_FUNC_VAR_DEF(ClsName), __COUNTER__) = ClsName::FField()
}  // namespace tvm
#endif  // TVM_IR_NODE_FUNCTOR_H_