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/tir/stmt.h>
 
#include <utility>
 
namespace tvm {
namespace script {
namespace ir_builder {
namespace tir {
 
class TIRFrameNode : public IRBuilderFrameNode {
 public:
  ffi::Array<tvm::tir::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.tir.TIRFrame", TIRFrameNode, IRBuilderFrameNode);
};
 
class TIRFrame : public IRBuilderFrame {
 public:
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TIRFrame, IRBuilderFrame, TIRFrameNode);
 
 protected:
  TIRFrame() = default;
  explicit TIRFrame(ObjectPtr<TIRFrameNode> data) : IRBuilderFrame(data) {}
};
 
class PrimFuncFrameNode : public TIRFrameNode {
 public:
  ffi::Optional<ffi::String> name;
  ffi::Array<tvm::tir::Var> args;
  bool is_private;
  ffi::Optional<Type> ret_type;
  ffi::Map<tvm::tir::Var, tvm::tir::Buffer> buffer_map;
  ffi::Map<ffi::String, Any> attrs;
  ffi::Map<tvm::tir::Var, tvm::tir::IterVar> env_threads;
  ffi::Array<tvm::tir::Buffer> root_alloc_buffers;
 
  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);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tir.PrimFuncFrame", PrimFuncFrameNode,
                                    TIRFrameNode);
 
 public:
  void ExitWithScope() final;
};
 
class PrimFuncFrame : public TIRFrame {
 public:
  explicit PrimFuncFrame(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 BlockFrameNode : public TIRFrameNode {
 public:
  ffi::String name;
  ffi::Array<tvm::tir::IterVar> iter_vars;
  ffi::Optional<ffi::Array<tvm::tir::BufferRegion>> reads;
  ffi::Optional<ffi::Array<tvm::tir::BufferRegion>> writes;
  ffi::Optional<tvm::tir::Stmt> init;
  ffi::Array<tvm::tir::Buffer> alloc_buffers;
  ffi::Array<tvm::tir::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<BlockFrameNode>()
        .def_ro("name", &BlockFrameNode::name)
        .def_ro("iter_vars", &BlockFrameNode::iter_vars)
        .def_ro("reads", &BlockFrameNode::reads)
        .def_ro("writes", &BlockFrameNode::writes)
        .def_ro("init", &BlockFrameNode::init)
        .def_ro("alloc_buffers", &BlockFrameNode::alloc_buffers)
        .def_ro("match_buffers", &BlockFrameNode::match_buffers)
        .def_ro("annotations", &BlockFrameNode::annotations)
        .def_ro("iter_values", &BlockFrameNode::iter_values)
        .def_ro("predicate", &BlockFrameNode::predicate)
        .def_ro("no_realize", &BlockFrameNode::no_realize);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tir.BlockFrame", BlockFrameNode,
                                    TIRFrameNode);
 
 public:
  void ExitWithScope() final;
};
 
class BlockFrame : public TIRFrame {
 public:
  explicit BlockFrame(ObjectPtr<BlockFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
    TVM_FFI_ICHECK(data != nullptr);
    data_ = std::move(data);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BlockFrame, TIRFrame, BlockFrameNode);
};
 
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.tir.BlockInitFrame", BlockInitFrameNode,
                                    TIRFrameNode);
 
 public:
  void EnterWithScope() final;
  void ExitWithScope() final;
};
 
class BlockInitFrame : public TIRFrame {
 public:
  explicit BlockInitFrame(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::tir::Stmt(ffi::Array<tvm::tir::Var> loop_vars,
                                        ffi::Array<Range> loop_extents, tvm::tir::Stmt loop_body)>;
  ffi::Array<tvm::tir::Var> vars;
  ffi::Array<Range> doms;
  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.tir.ForFrame", ForFrameNode, TIRFrameNode);
 
 public:
  void ExitWithScope() final;
};
 
class ForFrame : public TIRFrame {
 public:
  explicit ForFrame(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;
  PrimExpr message;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<AssertFrameNode>()
        .def_ro("condition", &AssertFrameNode::condition)
        .def_ro("message", &AssertFrameNode::message);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tir.AssertFrame", AssertFrameNode,
                                    TIRFrameNode);
 
 public:
  void ExitWithScope() final;
};
 
class AssertFrame : public TIRFrame {
 public:
  explicit AssertFrame(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 LetFrameNode : public TIRFrameNode {
 public:
  tvm::tir::Var var;
  PrimExpr value;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<LetFrameNode>()
        .def_ro("var", &LetFrameNode::var)
        .def_ro("value", &LetFrameNode::value);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tir.LetFrame", LetFrameNode, TIRFrameNode);
 
 public:
  void ExitWithScope() final;
};
 
class LetFrame : public TIRFrame {
 public:
  explicit LetFrame(ObjectPtr<LetFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
    TVM_FFI_ICHECK(data != nullptr);
    data_ = std::move(data);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(LetFrame, TIRFrame, LetFrameNode);
};
 
class LaunchThreadFrameNode : public TIRFrameNode {
 public:
  PrimExpr extent;
  ffi::String attr_key;
  tvm::tir::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.tir.LaunchThreadFrame",
                                    LaunchThreadFrameNode, TIRFrameNode);
 
 public:
  void ExitWithScope() final;
};
 
class LaunchThreadFrame : public TIRFrame {
 public:
  explicit LaunchThreadFrame(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 RealizeFrameNode : public TIRFrameNode {
 public:
  tvm::tir::BufferRegion buffer_slice;
  ffi::String storage_scope;
  PrimExpr condition;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<RealizeFrameNode>()
        .def_ro("buffer_slice", &RealizeFrameNode::buffer_slice)
        .def_ro("storage_scope", &RealizeFrameNode::storage_scope)
        .def_ro("condition", &RealizeFrameNode::condition);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tir.RealizeFrame", RealizeFrameNode,
                                    TIRFrameNode);
 
 public:
  void ExitWithScope() final;
};
 
class RealizeFrame : public TIRFrame {
 public:
  explicit RealizeFrame(ObjectPtr<RealizeFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
    TVM_FFI_ICHECK(data != nullptr);
    data_ = std::move(data);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(RealizeFrame, TIRFrame, RealizeFrameNode);
};
 
class AllocateFrameNode : public TIRFrameNode {
 public:
  ffi::Array<PrimExpr> extents;
  DataType dtype;
  ffi::String storage_scope;
  PrimExpr condition;
  ffi::Map<ffi::String, Any> annotations;
  tvm::tir::Var buffer_var;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<AllocateFrameNode>()
        .def_ro("extents", &AllocateFrameNode::extents)
        .def_ro("dtype", &AllocateFrameNode::dtype)
        .def_ro("storage_scope", &AllocateFrameNode::storage_scope)
        .def_ro("condition", &AllocateFrameNode::condition)
        .def_ro("annotations", &AllocateFrameNode::annotations)
        .def_ro("buffer_var", &AllocateFrameNode::buffer_var);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tir.AllocateFrame", AllocateFrameNode,
                                    TIRFrameNode);
 
 public:
  void ExitWithScope() final;
};
 
class AllocateFrame : public TIRFrame {
 public:
  explicit AllocateFrame(ObjectPtr<AllocateFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
    TVM_FFI_ICHECK(data != nullptr);
    data_ = std::move(data);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AllocateFrame, TIRFrame, AllocateFrameNode);
};
 
class AllocateConstFrameNode : public TIRFrameNode {
 public:
  DataType dtype;
  ffi::Array<PrimExpr> extents;
  tvm::runtime::Tensor data;
  tvm::tir::Var buffer_var;
  ffi::Map<ffi::String, Any> annotations;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<AllocateConstFrameNode>()
        .def_ro("dtype", &AllocateConstFrameNode::dtype)
        .def_ro("extents", &AllocateConstFrameNode::extents)
        .def_ro("data", &AllocateConstFrameNode::data)
        .def_ro("buffer_var", &AllocateConstFrameNode::buffer_var)
        .def_ro("annotations", &AllocateConstFrameNode::annotations);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tir.AllocateConstFrame",
                                    AllocateConstFrameNode, TIRFrameNode);
 
 public:
  void ExitWithScope() final;
};
 
class AllocateConstFrame : public TIRFrame {
 public:
  explicit AllocateConstFrame(ObjectPtr<AllocateConstFrameNode> data)
      : TIRFrame(ffi::UnsafeInit{}) {
    TVM_FFI_ICHECK(data != nullptr);
    data_ = std::move(data);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AllocateConstFrame, TIRFrame,
                                                AllocateConstFrameNode);
};
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.tir.AttrFrame", AttrFrameNode, TIRFrameNode);
 
 public:
  void ExitWithScope() final;
};
 
class AttrFrame : public TIRFrame {
 public:
  explicit AttrFrame(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.tir.WhileFrame", WhileFrameNode,
                                    TIRFrameNode);
 
 public:
  void ExitWithScope() final;
};
 
class WhileFrame : public TIRFrame {
 public:
  explicit WhileFrame(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::tir::Stmt>> then_stmts;
  ffi::Optional<ffi::Array<tvm::tir::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.tir.IfFrame", IfFrameNode, TIRFrameNode);
 
 public:
  void ExitWithScope() final;
};
 
class IfFrame : public TIRFrame {
 public:
  explicit IfFrame(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.tir.ThenFrame", ThenFrameNode, TIRFrameNode);
 
 public:
  void EnterWithScope() final;
  void ExitWithScope() final;
};
 
class ThenFrame : public TIRFrame {
 public:
  explicit ThenFrame(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.tir.ElseFrame", ElseFrameNode, TIRFrameNode);
 
 public:
  void EnterWithScope() final;
  void ExitWithScope() final;
};
 
class ElseFrame : public TIRFrame {
 public:
  explicit ElseFrame(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::tir::Buffer buffer;
  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.tir.DeclBufferFrame", DeclBufferFrameNode,
                                    TIRFrameNode);
 
 public:
  void ExitWithScope() final;
};
 
class DeclBufferFrame : public TIRFrame {
 public:
  explicit DeclBufferFrame(ObjectPtr<DeclBufferFrameNode> data) : TIRFrame(data) {
    TVM_FFI_ICHECK(data != nullptr);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DeclBufferFrame, TIRFrame, DeclBufferFrameNode);
};
 
}  // namespace tir
}  // namespace ir_builder
}  // namespace script
}  // namespace tvm
 
#endif  // TVM_SCRIPT_IR_BUILDER_TIR_FRAME_H_