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_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 {
 
class TIRFrameNode : public IRBuilderFrameNode {
 public:
  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);
};
 
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) {}
};
 
class PrimFuncFrameNode : public TIRFrameNode {
 public:
  ffi::Optional<ffi::String> name;
  ffi::Array<tvm::tirx::Var> args;
  bool is_private;
  ffi::Optional<Type> ret_type;
  ffi::Map<tvm::tirx::Var, tvm::tirx::BufferVar> buffer_map;
  ffi::Map<ffi::String, Any> attrs;
  ffi::Map<tvm::tirx::Var, tvm::tirx::IterVar> env_threads;
  ffi::Array<tvm::tirx::BufferVar> root_alloc_buffers;
 
  // TIR utils
  bool s_tir;
  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:
  void ExitWithScope() final;
};
 
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);
};
 
class SBlockFrameNode : public TIRFrameNode {
 public:
  ffi::String name;
  ffi::Array<tvm::tirx::IterVar> iter_vars;
  ffi::Optional<ffi::Array<tvm::tirx::BufferRegion>> reads;
  ffi::Optional<ffi::Array<tvm::tirx::BufferRegion>> writes;
  ffi::Optional<tvm::tirx::Stmt> init;
  ffi::Array<tvm::tirx::BufferVar> alloc_buffers;
  ffi::Array<tvm::tirx::MatchBufferRegion> match_buffers;
  ffi::Optional<ffi::Map<ffi::String, Any>> annotations;
  ffi::Array<PrimExpr> iter_values;
  ffi::Optional<PrimExpr> predicate;
  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:
  void ExitWithScope() final;
};
 
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);
};
 
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:
  void EnterWithScope() final;
  void ExitWithScope() final;
};
 
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);
};
 
class ForFrameNode : public TIRFrameNode {
 public:
  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)>;
  ffi::Array<tvm::tirx::Var> vars;
  ffi::Array<Range> doms;
  ffi::Array<ffi::Optional<PrimExpr>> steps;
  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:
  void ExitWithScope() final;
};
 
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);
};
 
class AssertFrameNode : public TIRFrameNode {
 public:
  PrimExpr condition;
  tvm::prim::StringImm error_kind;
  ffi::Array<tvm::prim::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:
  void ExitWithScope() final;
};
 
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);
};
 
class LaunchThreadFrameNode : public TIRFrameNode {
 public:
  PrimExpr extent;
  ffi::String attr_key;
  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:
  void ExitWithScope() final;
};
 
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);
};
 
class AttrFrameNode : public TIRFrameNode {
 public:
  Any node;
  ffi::String attr_key;
  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:
  void ExitWithScope() final;
};
 
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);
};
 
class WhileFrameNode : public TIRFrameNode {
 public:
  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:
  void ExitWithScope() final;
};
 
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);
};
 
class IfFrameNode : public TIRFrameNode {
 public:
  PrimExpr condition;
  ffi::Optional<ffi::Array<tvm::tirx::Stmt>> then_stmts;
  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:
  void ExitWithScope() final;
};
 
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);
};
 
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:
  void EnterWithScope() final;
  void ExitWithScope() final;
};
 
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);
};
 
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:
  void EnterWithScope() final;
  void ExitWithScope() final;
};
 
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:
  tvm::tirx::BufferVar buffer;
  Expr data;
  bool allocated;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<DeclBufferFrameNode>()
        .def_ro("buffer", &DeclBufferFrameNode::buffer)
        .def_ro("data", &DeclBufferFrameNode::data)
        .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:
  ffi::Map<ffi::String, tvm::tirx::BufferVar> workspace;
  ffi::Map<ffi::String, ffi::Any> config;
  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:
  tvm::tirx::BufferVar 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);
};
 
class HintFrameNode : public TIRFrameNode {
 public:
  ffi::String message;
  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;
};
 
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_