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/ir/diagnostic.h>
#include <tvm/ir/instrument.h>
#include <tvm/ir/module.h>
#include <tvm/runtime/container/array.h>
#include <tvm/runtime/container/string.h>
#include <tvm/support/with.h>

#include <string>
#include <utility>

namespace tvm {
namespace transform {

class PassContextNode : public Object {
 public:
  int opt_level{2};

  Array<String> required_pass;
  Array<String> disabled_pass;
  mutable Optional<DiagnosticContext> diag_ctx;
  Map<String, ObjectRef> config;

  Array<instrument::PassInstrument> instruments;

  PassContextNode() = default;

  template <typename TObjectRef>
  Optional<TObjectRef> GetConfig(const std::string& key, Optional<TObjectRef> default_value =
                                                             Optional<TObjectRef>(nullptr)) const {
    static_assert(std::is_base_of<ObjectRef, TObjectRef>::value,
                  "Can only call GetAttr with ObjectRef types.");
    if (!config.defined()) return default_value;
    auto it = config.find(key);
    if (it != config.end()) {
      return Downcast<Optional<TObjectRef>>((*it).second);
    } else {
      return default_value;
    }
  }
  // variant that uses TObjectRef to enable implicit conversion to default value.
  template <typename TObjectRef>
  Optional<TObjectRef> GetConfig(const std::string& key, TObjectRef default_value) const {
    return GetConfig<TObjectRef>(key, Optional<TObjectRef>(default_value));
  }

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("opt_level", &opt_level);
    v->Visit("required_pass", &required_pass);
    v->Visit("disabled_pass", &disabled_pass);
    v->Visit("instruments", &instruments);
    v->Visit("config", &config);
    v->Visit("diag_ctx", &diag_ctx);
  }

  static constexpr const char* _type_key = "transform.PassContext";
  static constexpr bool _type_has_method_sequal_reduce = false;
  TVM_DECLARE_FINAL_OBJECT_INFO(PassContextNode, Object);
};

class PassContext : public ObjectRef {
 public:
  PassContext() {}
  explicit PassContext(ObjectPtr<Object> 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 Map<String, Map<String, 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 uint32_t RegisterConfigOption(const char* key) {
    using ValueNodeType = typename ValueType::ContainerType;
    // NOTE: we could further update the function later.
    uint32_t tindex = ValueNodeType::_GetOrAllocRuntimeTypeIndex();
    RegisterConfigOption(key, tindex);
    return tindex;
  }

  // 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, uint32_t value_type_index);

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

  String name;

  Array<String> required;

  PassInfoNode() = default;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("opt_level", &opt_level);
    v->Visit("name", &name);
    v->Visit("required", &required);
  }

  static constexpr const char* _type_key = "transform.PassInfo";
  static constexpr bool _type_has_method_sequal_reduce = false;
  TVM_DECLARE_FINAL_OBJECT_INFO(PassInfoNode, Object);
};

class PassInfo : public ObjectRef {
 public:
  TVM_DLL PassInfo(int opt_level, String name, Array<runtime::String> required);

  TVM_DEFINE_OBJECT_REF_METHODS(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;

  void VisitAttrs(AttrVisitor* v) {}

  static constexpr const char* _type_key = "transform.Pass";
  TVM_DECLARE_BASE_OBJECT_INFO(PassNode, Object);
};

class Pass : public ObjectRef {
 public:
  IRModule operator()(IRModule mod) const;

  IRModule operator()(IRModule mod, const PassContext& pass_ctx) const;

  TVM_DEFINE_OBJECT_REF_METHODS(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::Array<Pass> passes;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("pass_info", &pass_info);
    v->Visit("passes", &passes);
  }

  PassInfo Info() const override { return pass_info; }

  void ResolveDependency(const IRModule& mod);

  IRModule operator()(IRModule mod, const PassContext& pass_ctx) const final;

  static constexpr const char* _type_key = "transform.Sequential";
  TVM_DECLARE_FINAL_OBJECT_INFO(SequentialNode, PassNode);
};

class Sequential : public Pass {
 public:
  TVM_DLL Sequential(Array<Pass> passes, PassInfo pass_info);

  TVM_DLL Sequential(Array<Pass> passes, String name = "sequential");

  Sequential() = default;
  explicit Sequential(ObjectPtr<Object> 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(const runtime::TypedPackedFunc<IRModule(IRModule, PassContext)>& pass_func,
                 int opt_level, String name, Array<runtime::String> required);

TVM_DLL Pass PrintIR(String header = "", bool show_meta_data = false);

}  // namespace transform
}  // namespace tvm

#endif  // TVM_IR_TRANSFORM_H_