transform.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_TRANSFORM_H_
#define TVM_IR_TRANSFORM_H_
 
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/reflection/creator.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include <tvm/ir/diagnostic.h>
#include <tvm/ir/instrument.h>
#include <tvm/ir/module.h>
#include <tvm/support/with.h>
 
#include <string>
#include <utility>
 
namespace tvm {
namespace transform {
 
class PassContextNode : public Object {
 public:
  int opt_level{2};
 
  ffi::Array<ffi::String> required_pass;
  ffi::Array<ffi::String> disabled_pass;
  mutable ffi::Optional<DiagnosticContext> diag_ctx;
  ffi::Map<ffi::String, Any> config;
 
  ffi::Array<instrument::PassInstrument> instruments;
 
  PassContextNode() = default;
 
  template <typename TObjectRef>
  ffi::Optional<TObjectRef> GetConfig(
      const std::string& key,
      ffi::Optional<TObjectRef> default_value = ffi::Optional<TObjectRef>(std::nullopt)) const {
    if (!config.defined()) return default_value;
    auto it = config.find(key);
    if (it != config.end()) {
      return Downcast<ffi::Optional<TObjectRef>>((*it).second);
    } else {
      return default_value;
    }
  }
  // variant that uses TObjectRef to enable implicit conversion to default value.
  template <typename TObjectRef>
  ffi::Optional<TObjectRef> GetConfig(const std::string& key, TObjectRef default_value) const {
    return GetConfig<TObjectRef>(key, ffi::Optional<TObjectRef>(default_value));
  }
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<PassContextNode>()
        .def_ro("opt_level", &PassContextNode::opt_level)
        .def_ro("required_pass", &PassContextNode::required_pass)
        .def_ro("disabled_pass", &PassContextNode::disabled_pass)
        .def_ro("instruments", &PassContextNode::instruments)
        .def_ro("config", &PassContextNode::config)
        .def_ro("diag_ctx", &PassContextNode::diag_ctx);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("transform.PassContext", PassContextNode, Object);
};
 
class PassContext : public ObjectRef {
 public:
  PassContext() {}
  explicit PassContext(ffi::UnsafeInit tag) : ObjectRef(tag) {}
  explicit PassContext(ObjectPtr<PassContextNode> n) : ObjectRef(n) {}
  const PassContextNode* operator->() const {
    ICHECK(get() != nullptr);
    return static_cast<const PassContextNode*>(get());
  }
  PassContextNode* operator->() {
    ICHECK(get() != nullptr);
    return static_cast<PassContextNode*>(get_mutable());
  }
 
  TVM_DLL static PassContext Create();
  TVM_DLL static PassContext Current();
 
  TVM_DLL static ffi::Map<ffi::String, ffi::Map<ffi::String, ffi::String>> ListConfigs();
 
  TVM_DLL void InstrumentEnterPassContext();
 
  TVM_DLL void InstrumentExitPassContext();
 
  TVM_DLL bool InstrumentBeforePass(const IRModule& mod, const PassInfo& info) const;
 
  TVM_DLL void InstrumentAfterPass(const IRModule& mod, const PassInfo& info) const;
 
  TVM_DLL bool PassEnabled(const PassInfo& info) const;
 
  template <typename ValueType>
  static int32_t RegisterConfigOption(const char* key) {
    // NOTE: we could further update the function later.
    if constexpr (std::is_base_of_v<ObjectRef, ValueType>) {
      int32_t tindex = ffi::TypeToRuntimeTypeIndex<ValueType>::v();
      auto type_key = ffi::TypeIndexToTypeKey(tindex);
      auto legalization = [=](ffi::Any value) -> ffi::Any {
        if (auto opt_map = value.try_cast<ffi::Map<ffi::String, ffi::Any>>()) {
          return ffi::reflection::ObjectCreator(type_key)(opt_map.value());
        } else {
          auto opt_val = value.try_cast<ValueType>();
          if (!opt_val.has_value()) {
            TVM_FFI_THROW(AttributeError)
                << "Expect config " << key << " to have type " << type_key << ", but instead get "
                << ffi::details::AnyUnsafe::GetMismatchTypeInfo<ValueType>(value);
          }
          return *opt_val;
        }
      };
      RegisterConfigOption(key, type_key, legalization);
    } else {
      // non-object type, do not support implicit conversion from map
      std::string type_str = ffi::TypeTraits<ValueType>::TypeStr();
      auto legalization = [=](ffi::Any value) -> ffi::Any {
        auto opt_val = value.try_cast<ValueType>();
        if (!opt_val.has_value()) {
          TVM_FFI_THROW(AttributeError)
              << "Expect config " << key << " to have type " << type_str << ", but instead get "
              << ffi::details::AnyUnsafe::GetMismatchTypeInfo<ValueType>(value);
        } else {
          return *opt_val;
        }
      };
      RegisterConfigOption(key, type_str, legalization);
    }
    return 0;
  }
 
  // accessor.
  using ContainerType = PassContextNode;
  class Internal;
 
 private:
  // The entry of a pass context scope.
  TVM_DLL void EnterWithScope();
  // The exit of a pass context scope.
  TVM_DLL void ExitWithScope();
  // Register configuration key value type.
  TVM_DLL static void RegisterConfigOption(const char* key, ffi::String value_type_str,
                                           std::function<ffi::Any(ffi::Any)> legalization);
 
  // Classes to get the Python `with` like syntax.
  friend class Internal;
  friend class With<PassContext>;
};
 
#define TVM_PASS_CTX_CONFIG_VAR_DEF static TVM_ATTRIBUTE_UNUSED uint32_t __make_PassContext_tid
 
#define TVM_REGISTER_PASS_CONFIG_OPTION(Key, ValueType)      \
  TVM_STR_CONCAT(TVM_PASS_CTX_CONFIG_VAR_DEF, __COUNTER__) = \
      ::tvm::transform::PassContext::RegisterConfigOption<ValueType>(Key)
 
class PassInfoNode : public Object {
 public:
  int opt_level;
 
  ffi::String name;
 
  bool traceable;
 
  ffi::Array<ffi::String> required;
 
  PassInfoNode() = default;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<PassInfoNode>()
        .def_ro("opt_level", &PassInfoNode::opt_level)
        .def_ro("name", &PassInfoNode::name)
        .def_ro("required", &PassInfoNode::required)
        .def_ro("traceable", &PassInfoNode::traceable);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("transform.PassInfo", PassInfoNode, Object);
};
 
class PassInfo : public ObjectRef {
 public:
  TVM_DLL PassInfo(int opt_level, ffi::String name, ffi::Array<ffi::String> required,
                   bool traceable);
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PassInfo, ObjectRef, PassInfoNode);
};
 
class PassNode : public Object {
 public:
  virtual ~PassNode() {}
  virtual PassInfo Info() const = 0;
 
  IRModule operator()(IRModule mod) const {
    return this->operator()(std::move(mod), PassContext::Current());
  }
 
  virtual IRModule operator()(IRModule mod, const PassContext& pass_ctx) const = 0;
  TVM_FFI_DECLARE_OBJECT_INFO("transform.Pass", PassNode, Object);
};
 
class Pass : public ObjectRef {
 public:
  IRModule operator()(IRModule mod) const;
 
  IRModule operator()(IRModule mod, const PassContext& pass_ctx) const;
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Pass, ObjectRef, PassNode);
 
 private:
  IRModule static AssertImmutableModule(const IRModule& mod, const PassNode* node,
                                        const PassContext& pass_ctx);
};
 
class SequentialNode : public PassNode {
 public:
  /* \brief The pass meta data.*/
  PassInfo pass_info;
 
  tvm::ffi::Array<Pass> passes;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<SequentialNode>()
        .def_ro("pass_info", &SequentialNode::pass_info)
        .def_ro("passes", &SequentialNode::passes);
  }
 
  PassInfo Info() const override { return pass_info; }
 
  void ResolveDependency(const IRModule& mod);
 
  IRModule operator()(IRModule mod, const PassContext& pass_ctx) const final;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("transform.Sequential", SequentialNode, PassNode);
};
 
class Sequential : public Pass {
 public:
  TVM_DLL Sequential(ffi::Array<Pass> passes, PassInfo pass_info);
 
  TVM_DLL Sequential(ffi::Array<Pass> passes, ffi::String name = "sequential");
 
  Sequential() = default;
  explicit Sequential(ObjectPtr<SequentialNode> n) : Pass(n) {}
 
  const SequentialNode* operator->() const;
  using ContainerType = SequentialNode;
};
 
/*
 * \brief Create a module pass.
 *
 * \param pass_func The packed function that contains the optimization.
 * \param opt_level The optimization level of the module pass.
 * \param name The name of the module pass.
 * \param required The list of the passes that the module pass is dependent on.
 *
 * \return The created module pass.
 */
TVM_DLL Pass CreateModulePass(std::function<IRModule(IRModule, PassContext)> pass_func,
                              int opt_level, ffi::String name, ffi::Array<ffi::String> required,
                              bool traceable = false);
 
/*
 * \brief Utility to apply a pass to specific functions in an IRModule
 *
 * TVM uses IRModule to IRModule transformations at all stages of
 * lowering.  These transformations may be useful when hand-writing an
 * optimized model, or to perform optimizations on specific kernels
 * within an IRModule.  This utility allows a pass to be applied to a
 * specified function, without altering other functions in the module.
 *
 * \param pass The IRModule to IRModule pass to be applied.
 *
 * \param func_name_regex A regex used to select the functions to be
 * updated.  The pass will be applied to all functions whose name
 * matches the regex.
 *
 * \param error_if_no_function_matches_regex Specifies the behavior if
 *     an IRModule does not contain any function matching the provided
 *     regex.  If true, an error will be raised.  If false (default),
 *     the IRModule will be returned unmodified.
 *
 * \return The modified IRModule to IRModule pass.
 */
TVM_DLL Pass ApplyPassToFunction(Pass pass, ffi::String func_name_regex,
                                 bool error_if_no_function_matches_regex = false);
 
TVM_DLL Pass PrintIR(ffi::String header = "", bool show_meta_data = false);
 
}  // namespace transform
}  // namespace tvm
 
#endif  // TVM_IR_TRANSFORM_H_