apache--tvm/include/tvm/tirx/script/builder/frame.h

/*
 * 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_IR_BUILDER_TIR_FRAME_H_
#define TVM_SCRIPT_IR_BUILDER_TIR_FRAME_H_

#include <tvm/script/ir_builder/base.h>
#include <tvm/script/ir_builder/ir/frame.h>
#include <tvm/tirx/exec_scope.h>
#include <tvm/tirx/stmt.h>

#include <utility>

namespace tvm {
namespace script {
namespace ir_builder {
namespace tirx {

/*!
 * \brief A base frame that represents the TIR fame with body of statements.
 *
 * \sa TIRFrame
 */
class TIRFrameNode : public IRBuilderFrameNode {
 public:
  /*! \brief The Stmt within in this frame. */
  ffi::Array<tvm::tirx::Stmt> stmts;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<TIRFrameNode>().def_ro("stmts", &TIRFrameNode::stmts);
  }
  TVM_FFI_DECLARE_OBJECT_INFO("script.ir_builder.tirx.TIRFrame", TIRFrameNode, IRBuilderFrameNode);
};

/*!
 * \brief Managed reference to TIRFrameNode.
 *
 * \sa TIRFrameNode
 */
class TIRFrame : public IRBuilderFrame {
 public:
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TIRFrame, IRBuilderFrame, TIRFrameNode);

 protected:
  TIRFrame() = default;
  explicit TIRFrame(ffi::ObjectPtr<TIRFrameNode> data) : IRBuilderFrame(data) {}
};

/*!
 * \brief A frame that represents the PrimFunc containing TIR statements.
 *
 * \sa PrimFuncFrame
 */
class PrimFuncFrameNode : public TIRFrameNode {
 public:
  /*! \brief The name of the block. */
  ffi::Optional<ffi::String> name;
  /*! \brief Function parameters. */
  ffi::Array<tvm::tirx::Var> args;
  /*! \brief Whether the PrimFunc is annotated as private. */
  bool is_private;
  /*! \brief The return type of the function. */
  ffi::Optional<Type> ret_type;
  /*! \brief Maps some parameters to specific Buffer data structures. */
  ffi::Map<tvm::tirx::Var, tvm::tirx::Buffer> buffer_map;
  /*! \brief Additional attributes storing the meta-data */
  ffi::Map<ffi::String, Any> attrs;
  /*! \brief The variable map bound to thread env. */
  ffi::Map<tvm::tirx::Var, tvm::tirx::IterVar> env_threads;
  /*! \brief The buffer allocated in root block. */
  ffi::Array<tvm::tirx::Buffer> root_alloc_buffers;

  // TIR utils
  /*! \brief Whether this PrimFunc uses s_tir semantics (root SBlock wrap,
   *  parser layout default = None). Default (false) = tirx semantics. */
  bool s_tir;
  /*! \brief Whether it is a persistent kernel. */
  bool persistent;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<PrimFuncFrameNode>()
        .def_ro("name", &PrimFuncFrameNode::name)
        .def_ro("args", &PrimFuncFrameNode::args)
        .def_ro("is_private", &PrimFuncFrameNode::is_private)
        .def_ro("ret_type", &PrimFuncFrameNode::ret_type)
        .def_ro("buffer_map", &PrimFuncFrameNode::buffer_map)
        .def_ro("attrs", &PrimFuncFrameNode::attrs)
        .def_ro("env_threads", &PrimFuncFrameNode::env_threads)
        .def_ro("root_alloc_buffers", &PrimFuncFrameNode::root_alloc_buffers)
        .def_ro("s_tir", &PrimFuncFrameNode::s_tir)
        .def_ro("persistent", &PrimFuncFrameNode::persistent);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.PrimFuncFrame", PrimFuncFrameNode,
                                    TIRFrameNode);

 public:
  /*!
   * \brief The method called when exiting RAII scope.
   * \sa tvm::support::With
   */
  void ExitWithScope() final;
};

/*!
 * \brief Managed reference to PrimFuncFrameNode.
 *
 * \sa PrimFuncFrameNode
 */
class PrimFuncFrame : public TIRFrame {
 public:
  explicit PrimFuncFrame(ffi::ObjectPtr<PrimFuncFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
    TVM_FFI_ICHECK(data != nullptr);
    data_ = std::move(data);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PrimFuncFrame, TIRFrame, PrimFuncFrameNode);
};

/*!
 * \brief A frame that represents the block.
 *
 * \sa SBlockFrame
 */
class SBlockFrameNode : public TIRFrameNode {
 public:
  /*! \brief The name of the block. */
  ffi::String name;
  /*! \brief The variables of the block. */
  ffi::Array<tvm::tirx::IterVar> iter_vars;
  /*! \brief The read buffer regions of the block. */
  ffi::Optional<ffi::Array<tvm::tirx::BufferRegion>> reads;
  /*! \brief The write buffer regions of the block. */
  ffi::Optional<ffi::Array<tvm::tirx::BufferRegion>> writes;
  /*! \brief The init statement of the bolck. */
  ffi::Optional<tvm::tirx::Stmt> init;
  /*! \brief The buffer allocated in the block. */
  ffi::Array<tvm::tirx::Buffer> alloc_buffers;
  /*! \brief The match buffer regions. */
  ffi::Array<tvm::tirx::MatchBufferRegion> match_buffers;
  /*! \brief The annotation of the block. */
  ffi::Optional<ffi::Map<ffi::String, Any>> annotations;
  /*! \brief The corresponding values of the iter vars. */
  ffi::Array<PrimExpr> iter_values;
  /*!
   * \brief The predicate of the block realization, the block will only be executed when the
   * predicate is true.
   */
  ffi::Optional<PrimExpr> predicate;
  /*! \brief The flag whether to construct BlockRealize or Block. */
  bool no_realize;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<SBlockFrameNode>()
        .def_ro("name", &SBlockFrameNode::name)
        .def_ro("iter_vars", &SBlockFrameNode::iter_vars)
        .def_ro("reads", &SBlockFrameNode::reads)
        .def_ro("writes", &SBlockFrameNode::writes)
        .def_ro("init", &SBlockFrameNode::init)
        .def_ro("alloc_buffers", &SBlockFrameNode::alloc_buffers)
        .def_ro("match_buffers", &SBlockFrameNode::match_buffers)
        .def_ro("annotations", &SBlockFrameNode::annotations)
        .def_ro("iter_values", &SBlockFrameNode::iter_values)
        .def_ro("predicate", &SBlockFrameNode::predicate)
        .def_ro("no_realize", &SBlockFrameNode::no_realize);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.SSBlockFrame", SBlockFrameNode,
                                    TIRFrameNode);

 public:
  /*!
   * \brief The method called when exiting RAII scope.
   * \sa tvm::support::With
   */
  void ExitWithScope() final;
};

/*!
 * \brief Managed reference to SBlockFrameNode.
 *
 * \sa SBlockFrameNode
 */

class SBlockFrame : public TIRFrame {
 public:
  explicit SBlockFrame(ffi::ObjectPtr<SBlockFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
    TVM_FFI_ICHECK(data != nullptr);
    data_ = std::move(data);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SBlockFrame, TIRFrame, SBlockFrameNode);
};

/*!
 * \brief A frame that represents the block initialization statment.
 *
 * \sa BlockInitFrame
 */
class BlockInitFrameNode : public TIRFrameNode {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<BlockInitFrameNode>();
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.SBlockInitFrame", BlockInitFrameNode,
                                    TIRFrameNode);

 public:
  /*!
   * \brief The method called when entering RAII scope.
   * \sa tvm::support::With
   */
  void EnterWithScope() final;
  /*!
   * \brief The method called when exiting RAII scope.
   * \sa tvm::support::With
   */
  void ExitWithScope() final;
};

/*!
 * \brief Managed reference to BlockInitFrameNode.
 *
 * \sa BlockInitFrameNode
 */
class BlockInitFrame : public TIRFrame {
 public:
  explicit BlockInitFrame(ffi::ObjectPtr<BlockInitFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
    TVM_FFI_ICHECK(data != nullptr);
    data_ = std::move(data);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BlockInitFrame, TIRFrame, BlockInitFrameNode);
};

/*!
 * \brief A frame that represents the for loop.
 *
 * \sa ForFrame
 */
class ForFrameNode : public TIRFrameNode {
 public:
  /*!
   * \brief Functions that generate loop nests.
   * \param loop_vars The loop variables, from outer to inner
   * \param loop_extents The loop extents that correspond to loop variables
   * \param loop_body The loop body
   * \return A stmt, the loop nest
   */
  using FMakeForLoop = ffi::TypedFunction<tvm::tirx::Stmt(
      ffi::Array<tvm::tirx::Var> loop_vars, ffi::Array<Range> loop_extents,
      ffi::Array<ffi::Optional<PrimExpr>> loop_steps, tvm::tirx::Stmt loop_body)>;
  /*! \brief The loop variable. */
  ffi::Array<tvm::tirx::Var> vars;
  /*! \brief The domains of iteration. */
  ffi::Array<Range> doms;
  /*! \brief The optional steps of iteration. */
  ffi::Array<ffi::Optional<PrimExpr>> steps;
  /*! \brief The for loop generating function. */
  FMakeForLoop f_make_for_loop;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ForFrameNode>()
        .def_ro("vars", &ForFrameNode::vars)
        .def_ro("doms", &ForFrameNode::doms);
    // `f_make_for_loop` is not registered as it's not visited.
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.ForFrame", ForFrameNode, TIRFrameNode);

 public:
  /*!
   * \brief The method called when exiting RAII scope.
   * \sa tvm::support::With
   */
  void ExitWithScope() final;
};

/*!
 * \brief Managed reference to ForFrameNode.
 *
 * \sa ForFrameNode
 */
class ForFrame : public TIRFrame {
 public:
  explicit ForFrame(ffi::ObjectPtr<ForFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
    TVM_FFI_ICHECK(data != nullptr);
    data_ = std::move(data);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ForFrame, TIRFrame, ForFrameNode);
};

/*!
 * \brief A frame that represents the assert statement. Proceeds if the condition is true,
 * otherwise aborts with the message.
 *
 * \sa AssertFrame
 */
class AssertFrameNode : public TIRFrameNode {
 public:
  /*! \brief The PrimExpr to test. */
  PrimExpr condition;
  /*! \brief The error kind, e.g. "RuntimeError", "TypeError", "ValueError". */
  tvm::tirx::StringImm error_kind;
  /*! \brief Error message fragments, concatenated at runtime when assertion fails. */
  ffi::Array<tvm::tirx::StringImm> message_parts;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<AssertFrameNode>()
        .def_ro("condition", &AssertFrameNode::condition)
        .def_ro("error_kind", &AssertFrameNode::error_kind)
        .def_ro("message_parts", &AssertFrameNode::message_parts);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.AssertFrame", AssertFrameNode,
                                    TIRFrameNode);

 public:
  /*!
   * \brief The method called when exiting RAII scope.
   * \sa tvm::support::With
   */
  void ExitWithScope() final;
};

/*!
 * \brief Managed reference to AssertFrameNode.
 *
 * \sa AssertFrameNode
 */
class AssertFrame : public TIRFrame {
 public:
  explicit AssertFrame(ffi::ObjectPtr<AssertFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
    TVM_FFI_ICHECK(data != nullptr);
    data_ = std::move(data);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AssertFrame, TIRFrame, AssertFrameNode);
};

/*!
 * \brief The LaunchThreadFrameNode.
 * \note It is used only inside a PrimFunc.
 */
class LaunchThreadFrameNode : public TIRFrameNode {
 public:
  /*! \brief The extent of environment thread. */
  PrimExpr extent;
  /*! \brief The attribute key, could be either virtual_thread or thread_extent. */
  ffi::String attr_key;
  /*! \brief The iteration variable. */
  tvm::tirx::IterVar iter_var;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<LaunchThreadFrameNode>()
        .def_ro("extent", &LaunchThreadFrameNode::extent)
        .def_ro("attr_key", &LaunchThreadFrameNode::attr_key)
        .def_ro("iter_var", &LaunchThreadFrameNode::iter_var);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.LaunchThreadFrame",
                                    LaunchThreadFrameNode, TIRFrameNode);

 public:
  /*!
   * \brief The method called when exiting RAII scope.
   * \sa tvm::support::With
   */
  void ExitWithScope() final;
};

/*!
 * \brief Managed reference to LaunchThreadFrameNode.
 *
 * \sa LaunchThreadFrameNode
 */
class LaunchThreadFrame : public TIRFrame {
 public:
  explicit LaunchThreadFrame(ffi::ObjectPtr<LaunchThreadFrameNode> data)
      : TIRFrame(ffi::UnsafeInit{}) {
    TVM_FFI_ICHECK(data != nullptr);
    data_ = std::move(data);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(LaunchThreadFrame, TIRFrame, LaunchThreadFrameNode);
};

/*!
 * \brief A frame that represents attribute node.
 *
 * \sa AttrFrame
 */
class AttrFrameNode : public TIRFrameNode {
 public:
  /*! \brief The node to annotate the attribute. */
  Any node;
  /*! \brief Attribute type key. */
  ffi::String attr_key;
  /*! \brief The value of the attribute. */
  PrimExpr value;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<AttrFrameNode>()
        .def_ro("node", &AttrFrameNode::node)
        .def_ro("attr_key", &AttrFrameNode::attr_key)
        .def_ro("value", &AttrFrameNode::value);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.AttrFrame", AttrFrameNode,
                                    TIRFrameNode);

 public:
  /*!
   * \brief The method called when exiting RAII scope.
   * \sa tvm::support::With
   */
  void ExitWithScope() final;
};

/*!
 * \brief Managed reference to AttrFrameNode.
 *
 * \sa AttrFrameNode
 */
class AttrFrame : public TIRFrame {
 public:
  explicit AttrFrame(ffi::ObjectPtr<AttrFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
    TVM_FFI_ICHECK(data != nullptr);
    data_ = std::move(data);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AttrFrame, TIRFrame, AttrFrameNode);
};

/*!
 * \brief A frame that represents while loop.
 *
 * \sa WhileFrame
 */
class WhileFrameNode : public TIRFrameNode {
 public:
  /*! \brief The termination condition of while. */
  PrimExpr condition;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<WhileFrameNode>().def_ro("condition", &WhileFrameNode::condition);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.WhileFrame", WhileFrameNode,
                                    TIRFrameNode);

 public:
  /*!
   * \brief The method called when exiting RAII scope.
   * \sa tvm::support::With
   */
  void ExitWithScope() final;
};

/*!
 * \brief Managed reference to WhileFrameNode.
 *
 * \sa WhileFrameNode
 */
class WhileFrame : public TIRFrame {
 public:
  explicit WhileFrame(ffi::ObjectPtr<WhileFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
    TVM_FFI_ICHECK(data != nullptr);
    data_ = std::move(data);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(WhileFrame, TIRFrame, WhileFrameNode);
};

/*!
 * \brief A frame that represents if statement.
 *
 * \sa IfFrame
 */
class IfFrameNode : public TIRFrameNode {
 public:
  /*! \brief The condition of the if statement. */
  PrimExpr condition;
  /*! \brief The statements in the true branch. */
  ffi::Optional<ffi::Array<tvm::tirx::Stmt>> then_stmts;
  /*! \brief The stetements in the false branch. */
  ffi::Optional<ffi::Array<tvm::tirx::Stmt>> else_stmts;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<IfFrameNode>()
        .def_ro("condition", &IfFrameNode::condition)
        .def_ro("then_stmts", &IfFrameNode::then_stmts)
        .def_ro("else_stmts", &IfFrameNode::else_stmts);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.IfFrame", IfFrameNode, TIRFrameNode);

 public:
  /*!
   * \brief The method called when exiting RAII scope.
   * \sa tvm::support::With
   */
  void ExitWithScope() final;
};

/*!
 * \brief Managed reference to IfFrameNode.
 *
 * \sa IfFrameNode
 */
class IfFrame : public TIRFrame {
 public:
  explicit IfFrame(ffi::ObjectPtr<IfFrameNode> data) : TIRFrame(data) {
    TVM_FFI_ICHECK(data != nullptr);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IfFrame, TIRFrame, IfFrameNode);
};

/*!
 * \brief A frame that represents then.
 *
 * \sa ThenFrame
 */
class ThenFrameNode : public TIRFrameNode {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ThenFrameNode>();
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.ThenFrame", ThenFrameNode,
                                    TIRFrameNode);

 public:
  /*!
   * \brief The method called when entering RAII scope.
   * \sa tvm::support::With
   */
  void EnterWithScope() final;
  /*!
   * \brief The method called when exiting RAII scope.
   * \sa tvm::support::With
   */
  void ExitWithScope() final;
};

/*!
 * \brief Managed reference to ThenFrameNode.
 *
 * \sa ThenFrameNode
 */
class ThenFrame : public TIRFrame {
 public:
  explicit ThenFrame(ffi::ObjectPtr<ThenFrameNode> data) : TIRFrame(data) {
    TVM_FFI_ICHECK(data != nullptr);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ThenFrame, TIRFrame, ThenFrameNode);
};

/*!
 * \brief A frame that represents else.
 *
 * \sa ElseFrame
 */
class ElseFrameNode : public TIRFrameNode {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ElseFrameNode>();
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.ElseFrame", ElseFrameNode,
                                    TIRFrameNode);

 public:
  /*!
   * \brief The method called when entering RAII scope.
   * \sa tvm::support::With
   */
  void EnterWithScope() final;
  /*!
   * \brief The method called when exiting RAII scope.
   * \sa tvm::support::With
   */
  void ExitWithScope() final;
};

/*!
 * \brief Managed reference to ElseFrameNode.
 *
 * \sa ElseFrameNode
 */
class ElseFrame : public TIRFrame {
 public:
  explicit ElseFrame(ffi::ObjectPtr<ElseFrameNode> data) : TIRFrame(data) {
    TVM_FFI_ICHECK(data != nullptr);
  }

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ElseFrame, TIRFrame, ElseFrameNode);
};

class DeclBufferFrameNode : public TIRFrameNode {
 public:
  /*! \brief The declared buffer. */
  tvm::tirx::Buffer buffer;
  /*! \brief The buffer allocated or not. */
  bool allocated;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<DeclBufferFrameNode>()
        .def_ro("buffer", &DeclBufferFrameNode::buffer)
        .def_ro("allocated", &DeclBufferFrameNode::allocated);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.DeclBufferFrame", DeclBufferFrameNode,
                                    TIRFrameNode);

 public:
  void ExitWithScope() final;
};

class DeclBufferFrame : public TIRFrame {
 public:
  explicit DeclBufferFrame(ffi::ObjectPtr<DeclBufferFrameNode> data) : TIRFrame(data) {
    TVM_FFI_ICHECK(data != nullptr);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DeclBufferFrame, TIRFrame, DeclBufferFrameNode);
};

class ComposeOpFrameNode : public TIRFrameNode {
 public:
  /*! \brief The workspace of the compose op. */
  ffi::Map<ffi::String, tvm::tirx::Buffer> workspace;
  /*! \brief The config of the compose op. */
  ffi::Map<ffi::String, ffi::Any> config;
  /*! \brief The optional dispatch variant name of the compose op. */
  ffi::Optional<ffi::String> dispatch{std::nullopt};

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ComposeOpFrameNode>()
        .def_ro("workspace", &ComposeOpFrameNode::workspace)
        .def_ro("config", &ComposeOpFrameNode::config)
        .def_ro("dispatch", &ComposeOpFrameNode::dispatch);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.ComposeOpFrame", ComposeOpFrameNode,
                                    TIRFrameNode);

 public:
  void ExitWithScope() final;
};

class ComposeOpFrame : public TIRFrame {
 public:
  explicit ComposeOpFrame(ffi::ObjectPtr<ComposeOpFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
    TVM_FFI_ICHECK(data != nullptr);
    data_ = std::move(data);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ComposeOpFrame, TIRFrame, ComposeOpFrameNode);
};
class AllocBufferFrameNode : public TIRFrameNode {
 public:
  /*! \brief The allocated buffer. */
  tvm::tirx::Buffer buffer;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<AllocBufferFrameNode>().def_ro("buffer", &AllocBufferFrameNode::buffer);
  }

  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.AllocBufferFrame", AllocBufferFrameNode,
                                    TIRFrameNode);

 public:
  void ExitWithScope() final;
};

class AllocBufferFrame : public TIRFrame {
 public:
  explicit AllocBufferFrame(ffi::ObjectPtr<AllocBufferFrameNode> data)
      : TIRFrame(ffi::UnsafeInit{}) {
    TVM_FFI_ICHECK(data != nullptr);
    data_ = std::move(data);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AllocBufferFrame, TIRFrame, AllocBufferFrameNode);
};

/*!
 * \brief A frame that represents a hint directive for the sketch language.
 *
 * \sa HintFrame
 */
class HintFrameNode : public TIRFrameNode {
 public:
  /*! \brief The free-form hint message string. */
  ffi::String message;
  /*! \brief Optional structured key-value attributes. */
  ffi::Map<ffi::String, ffi::Any> attrs;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<HintFrameNode>()
        .def_ro("message", &HintFrameNode::message)
        .def_ro("attrs", &HintFrameNode::attrs);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.HintFrame", HintFrameNode,
                                    TIRFrameNode);

 public:
  void ExitWithScope() final;
};

/*!
 * \brief Managed reference to HintFrameNode.
 *
 * \sa HintFrameNode
 */
class HintFrame : public TIRFrame {
 public:
  explicit HintFrame(ffi::ObjectPtr<HintFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
    TVM_FFI_ICHECK(data != nullptr);
    data_ = std::move(data);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(HintFrame, TIRFrame, HintFrameNode);
};

}  // namespace tirx
}  // namespace ir_builder
}  // namespace script
}  // namespace tvm

#endif  // TVM_TIRX_SCRIPT_BUILDER_FRAME_H_