traced_object_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_TRACED_OBJECT_FUNCTOR_H_
#define TVM_SCRIPT_PRINTER_TRACED_OBJECT_FUNCTOR_H_

#include <tvm/node/node.h>
#include <tvm/runtime/logging.h>
#include <tvm/runtime/packed_func.h>
#include <tvm/script/printer/traced_object.h>

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

namespace tvm {
namespace script {
namespace printer {

/*
 * 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>>;

const runtime::PackedFunc& GetDispatchFunction(const DispatchTable& dispatch_table,
                                               const String& token, uint32_t type_index);

void SetDispatchFunction(DispatchTable* dispatch_table, const String& token, uint32_t type_index,
                         runtime::PackedFunc f);

void RemoveDispatchFunction(DispatchTable* dispatch_table, const String& token,
                            uint32_t type_index);

constexpr const char* kDefaultDispatchToken = "";

template <typename R, typename... Args>
class TracedObjectFunctor {
 private:
  using TSelf = TracedObjectFunctor<R, Args...>;

  template <class TObjectRef, class TCallable>
  using IsDispatchFunction =
      typename std::is_convertible<TCallable, std::function<R(TracedObject<TObjectRef>, Args...)>>;

 public:
  template <class TObjectRef>
  R operator()(const String& token, TracedObject<TObjectRef> traced_object, Args... args) const {
    const runtime::PackedFunc& dispatch_function =
        GetDispatchFunction(dispatch_table_, token, traced_object.Get()->type_index());
    return dispatch_function(traced_object.Get(), traced_object.GetPath(), args...);
  }

  TSelf& set_dispatch(String token, uint32_t type_index, runtime::PackedFunc f) {
    SetDispatchFunction(&dispatch_table_, token, type_index, std::move(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, ObjectPath, Args...)>(
            [f = std::move(f)](TObjectRef object, ObjectPath path, Args... args) -> R {
              return f(MakeTraced(object, path), args...);
            }));
  }
  template <typename TObjectRef, typename TCallable,
            typename = std::enable_if_t<IsDispatchFunction<TObjectRef, TCallable>::value>>
  TSelf& set_dispatch(TCallable&& f) {
    return set_dispatch<TObjectRef>(kDefaultDispatchToken, std::forward<TCallable>(f));
  }

  void remove_dispatch(String token, uint32_t type_index) {
    RemoveDispatchFunction(&dispatch_table_, token, type_index);
  }

 private:
  DispatchTable dispatch_table_;
};

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