frame.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_RELAX_SCRIPT_IR_BUILDER_FRAME_H_
#define TVM_RELAX_SCRIPT_IR_BUILDER_FRAME_H_
 
#include <tvm/ffi/reflection/registry.h>
#include <tvm/relax/block_builder.h>
#include <tvm/relax/expr.h>
#include <tvm/script/ir_builder/base.h>
#include <tvm/script/ir_builder/frame.h>
#include <tvm/script/ir_builder/ir.h>
 
#include <utility>
 
namespace tvm {
namespace script {
namespace ir_builder {
namespace relax {
 
class RelaxFrameObj : public IRBuilderFrameObj {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<RelaxFrameObj>();
  }
  TVM_FFI_DECLARE_OBJECT_INFO("script.ir_builder.relax.RelaxFrame", RelaxFrameObj,
                              IRBuilderFrameObj);
};
 
class RelaxFrame : public IRBuilderFrame {
 public:
  explicit RelaxFrame(ffi::ObjectPtr<RelaxFrameObj> data) : IRBuilderFrame(ffi::UnsafeInit{}) {
    TVM_FFI_ICHECK(data != nullptr);
    data_ = std::move(data);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(RelaxFrame, IRBuilderFrame, RelaxFrameObj);
 
 protected:
  RelaxFrame() = default;
};
 
class SeqExprFrameObj : public RelaxFrameObj {
 public:
  ffi::Array<tvm::relax::BindingBlock> binding_blocks;
  ffi::Optional<tvm::relax::Expr> output;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<SeqExprFrameObj>()
        .def_ro("binding_blocks", &SeqExprFrameObj::binding_blocks)
        .def_ro("output", &SeqExprFrameObj::output);
  }
  TVM_FFI_DECLARE_OBJECT_INFO("script.ir_builder.relax.SeqExprFrame", SeqExprFrameObj,
                              RelaxFrameObj);
 
 public:
  void EnterWithScope() override;
  void ExitWithScope() override;
};
 
class SeqExprFrame : public RelaxFrame {
 public:
  explicit SeqExprFrame(ffi::ObjectPtr<SeqExprFrameObj> data) : RelaxFrame(data) {
    TVM_FFI_ICHECK(data != nullptr);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SeqExprFrame, RelaxFrame, SeqExprFrameObj);
};
 
class FunctionFrameObj : public SeqExprFrameObj {
 public:
  explicit FunctionFrameObj(tvm::relax::BlockBuilder builder) : block_builder(std::move(builder)) {}
  explicit FunctionFrameObj(ffi::UnsafeInit) : block_builder(ffi::UnsafeInit{}) {}
  ffi::Optional<ffi::String> name;
  ffi::Array<tvm::Var> params;
  ffi::Optional<tvm::Type> ret_ty;
  ffi::Optional<bool> is_pure;
  ffi::Optional<bool> is_private;
  ffi::Map<ffi::String, Any> attrs;
  tvm::relax::BlockBuilder block_builder;
  bool declaration = false;
  bool local = false;
  ffi::Optional<tvm::Var> local_var;
  ffi::Optional<tvm::relax::Function> function;
  ffi::Optional<tvm::GlobalVar> global_var;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<FunctionFrameObj>()
        .def_ro("name", &FunctionFrameObj::name)
        .def_ro("params", &FunctionFrameObj::params)
        .def_ro("ret_ty", &FunctionFrameObj::ret_ty)
        .def_ro("is_pure", &FunctionFrameObj::is_pure)
        .def_ro("attrs", &FunctionFrameObj::attrs)
        .def_ro("function", &FunctionFrameObj::function)
        .def_ro("global_var", &FunctionFrameObj::global_var)
        .def_ro("local_var", &FunctionFrameObj::local_var);
    // `binding_blocks` and `output` are inherited from SeqExprFrameObj.
    // `block_builder` is not registered as it's not visited.
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.relax.FunctionFrame", FunctionFrameObj,
                                    SeqExprFrameObj);
 
 public:
  void EnterWithScope() final;
  void ExitWithScope() final;
};
 
class FunctionFrame : public SeqExprFrame {
 public:
  explicit FunctionFrame(ffi::ObjectPtr<FunctionFrameObj> data) : SeqExprFrame(data) {
    TVM_FFI_ICHECK(data != nullptr);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(FunctionFrame, SeqExprFrame, FunctionFrameObj);
};
 
class BindingBlockFrameObj : public RelaxFrameObj {
 public:
  bool is_dataflow;
  ffi::Array<tvm::Var> emitted_vars;
  bool block_ended;
  ffi::Array<tvm::Var> output_vars;
  ffi::Map<tvm::Var, Location> binding_locs;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<BindingBlockFrameObj>()
        .def_ro("is_dataflow", &BindingBlockFrameObj::is_dataflow)
        .def_ro("emitted_vars", &BindingBlockFrameObj::emitted_vars)
        .def_ro("output_vars", &BindingBlockFrameObj::output_vars);
    // `block_ended` is not registered as it's not visited.
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.relax.BindingBlockFrame",
                                    BindingBlockFrameObj, RelaxFrameObj);
 
 public:
  void EnterWithScope() final;
  void ExitWithScope() final;
};
 
class BindingBlockFrame : public RelaxFrame {
 public:
  explicit BindingBlockFrame(ffi::ObjectPtr<BindingBlockFrameObj> data) : RelaxFrame(data) {
    TVM_FFI_ICHECK(data != nullptr);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BindingBlockFrame, RelaxFrame,
                                                BindingBlockFrameObj);
};
 
class IfFrameObj : public RelaxFrameObj {
 public:
  explicit IfFrameObj(tvm::relax::Expr condition) : condition(std::move(condition)) {}
  explicit IfFrameObj(ffi::UnsafeInit) : condition(ffi::UnsafeInit{}) {}
 
  tvm::relax::Expr condition;
  ffi::Optional<tvm::relax::Expr> then_expr;
  ffi::Optional<tvm::relax::Expr> else_expr;
  ffi::Optional<tvm::Var> var;
  ffi::String var_name;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<IfFrameObj>()
        .def_ro("condition", &IfFrameObj::condition)
        .def_ro("then_expr", &IfFrameObj::then_expr)
        .def_ro("else_expr", &IfFrameObj::else_expr)
        .def_ro("var", &IfFrameObj::var)
        .def_ro("var_name", &IfFrameObj::var_name);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.relax.IfFrame", IfFrameObj, RelaxFrameObj);
 
 public:
  void EnterWithScope() final;
  void ExitWithScope() final;
};
 
class IfFrame : public RelaxFrame {
 public:
  explicit IfFrame(ffi::ObjectPtr<IfFrameObj> data) : RelaxFrame(data) {
    TVM_FFI_ICHECK(data != nullptr);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IfFrame, RelaxFrame, IfFrameObj);
};
 
class ThenFrameObj : public SeqExprFrameObj {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ThenFrameObj>();
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.relax.ThenFrame", ThenFrameObj,
                                    SeqExprFrameObj);
 
 public:
  void EnterWithScope() final;
  void ExitWithScope() final;
};
 
class ThenFrame : public SeqExprFrame {
 public:
  explicit ThenFrame(ffi::ObjectPtr<ThenFrameObj> data) : SeqExprFrame(data) {
    TVM_FFI_ICHECK(data != nullptr);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ThenFrame, SeqExprFrame, ThenFrameObj);
};
 
class ElseFrameObj : public SeqExprFrameObj {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ElseFrameObj>();
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.relax.ElseFrame", ElseFrameObj,
                                    SeqExprFrameObj);
 
 public:
  void EnterWithScope() final;
  void ExitWithScope() final;
};
 
class ElseFrame : public SeqExprFrame {
 public:
  explicit ElseFrame(ffi::ObjectPtr<ElseFrameObj> data) : SeqExprFrame(data) {
    TVM_FFI_ICHECK(data != nullptr);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ElseFrame, SeqExprFrame, ElseFrameObj);
};
 
}  // namespace relax
}  // namespace ir_builder
}  // namespace script
}  // namespace tvm
 
#endif  // TVM_RELAX_SCRIPT_IR_BUILDER_FRAME_H_