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;
  // TODO(@sunggg): Fix dependency issue in the header file and correct the types
  // e.g., relax::trace, relax::database in tvm/relax/tuning_api.h
  mutable Array<ObjectRef> trace_stack;
  Optional<Map<String, Bool>> make_traceable;
  mutable int num_evals{0};
  Optional<ObjectRef> tuning_api_database;
  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);
    v->Visit("trace_stack", &trace_stack);
    v->Visit("make_traceable", &make_traceable);
    v->Visit("num_evals", &num_evals);
    v->Visit("tuning_api_daatabase", &tuning_api_database);
  }
 
  Array<ObjectRef> GetTraceStack() { return trace_stack; }
  void PushTrace(ObjectRef new_trace) { trace_stack.push_back(new_trace); }
  void PopTrace() {
    ICHECK(GetTraceStackSize()) << "Trace stack is currently empty. Please double check.";
    trace_stack.pop_back();
  }
  int GetTraceStackSize() { return trace_stack.size(); }
  ObjectRef GetCurrentTrace() {
    ICHECK(GetTraceStackSize()) << "Trace stack is currently empty. Please double check.";
    return trace_stack.back();
  }
  void SetNumEvals(int _num_evals) { num_evals = _num_evals; }
  void IncNumEvals(int _num_evals) { num_evals += _num_evals; }
 
  Optional<ObjectRef> GetTuningAPIDatabase() { return tuning_api_database; }
 
  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();
    auto type_key = runtime::Object::TypeIndex2Key(tindex);
 
    auto* reflection = ReflectionVTable::Global();
 
    auto legalization = [=](ObjectRef obj) -> ObjectRef {
      if (obj->IsInstance<Map<String, ObjectRef>::ContainerType>()) {
        return reflection->CreateObject(type_key, Downcast<Map<String, ObjectRef>>(obj));
      } else {
        // Backwards compatibility for config options defined prior to
        // https://github.com/apache/tvm/pull/16183.  This commit
        // changed the default FFI conversion of python integers from
        // `tvm::IntImm` to `runtime::Int`.
        //
        // This backwards compatibility fix can be removed when all
        // options registered with TVM_REGISTER_PASS_CONFIG_OPTION are
        // updated to use `runtime::Int` and `runtime::Bool`.
        TVMRetValue ret;
        ret = obj;
        try {
          ValueType legalized = ret;
          return legalized;
        } catch (Error& err) {
          LOG(FATAL) << "AttributeError: expect config " << key << " to have type " << type_key
                     << ", but received error when converting to this type.\n"
                     << err.what();
        }
      }
    };
 
    RegisterConfigOption(key, tindex, legalization);
    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,
                                           std::function<ObjectRef(ObjectRef)> 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;
 
  String name;
 
  bool traceable;
 
  Array<String> required;
 
  PassInfoNode() = default;
 
  void VisitAttrs(AttrVisitor* v) {
    v->Visit("opt_level", &opt_level);
    v->Visit("name", &name);
    v->Visit("required", &required);
    v->Visit("traceable", &traceable);
  }
 
  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, bool traceable);
 
  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, 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, String func_name_regex,
                                 bool error_if_no_function_matches_regex = false);
 
TVM_DLL Pass PrintIR(String header = "", bool show_meta_data = false);
 
}  // namespace transform
}  // namespace tvm
 
#endif  // TVM_IR_TRANSFORM_H_