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>
 
namespace tvm {
namespace script {
namespace ir_builder {
namespace tir {
 
class TIRFrameNode : public IRBuilderFrameNode {
 public:
  Array<tvm::tir::Stmt> stmts;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<TIRFrameNode>().def_ro("stmts", &TIRFrameNode::stmts);
  }
 
  static constexpr const char* _type_key = "script.ir_builder.tir.TIRFrame";
  TVM_DECLARE_BASE_OBJECT_INFO(TIRFrameNode, IRBuilderFrameNode);
};
 
class TIRFrame : public IRBuilderFrame {
 public:
  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(TIRFrame, IRBuilderFrame, TIRFrameNode);
 
 protected:
  TIRFrame() = default;
};
 
class PrimFuncFrameNode : public TIRFrameNode {
 public:
  Optional<String> name;
  Array<tvm::tir::Var> args;
  bool is_private;
  Optional<Type> ret_type;
  Map<tvm::tir::Var, tvm::tir::Buffer> buffer_map;
  Map<String, Any> attrs;
  Map<tvm::tir::Var, tvm::tir::IterVar> env_threads;
  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);
  }
 
  static constexpr const char* _type_key = "script.ir_builder.tir.PrimFuncFrame";
  TVM_DECLARE_FINAL_OBJECT_INFO(PrimFuncFrameNode, TIRFrameNode);
 
 public:
  void ExitWithScope() final;
};
 
class PrimFuncFrame : public TIRFrame {
 public:
  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(PrimFuncFrame, TIRFrame, PrimFuncFrameNode);
};
 
class BlockFrameNode : public TIRFrameNode {
 public:
  String name;
  Array<tvm::tir::IterVar> iter_vars;
  Optional<Array<tvm::tir::BufferRegion>> reads;
  Optional<Array<tvm::tir::BufferRegion>> writes;
  Optional<tvm::tir::Stmt> init;
  Array<tvm::tir::Buffer> alloc_buffers;
  Array<tvm::tir::MatchBufferRegion> match_buffers;
  Optional<Map<String, Any>> annotations;
  Array<PrimExpr> iter_values;
  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);
  }
 
  static constexpr const char* _type_key = "script.ir_builder.tir.BlockFrame";
  TVM_DECLARE_FINAL_OBJECT_INFO(BlockFrameNode, TIRFrameNode);
 
 public:
  void ExitWithScope() final;
};
 
class BlockFrame : public TIRFrame {
 public:
  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(BlockFrame, TIRFrame, BlockFrameNode);
};
 
class BlockInitFrameNode : public TIRFrameNode {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<BlockInitFrameNode>();
  }
 
  static constexpr const char* _type_key = "script.ir_builder.tir.BlockInitFrame";
  TVM_DECLARE_FINAL_OBJECT_INFO(BlockInitFrameNode, TIRFrameNode);
 
 public:
  void EnterWithScope() final;
  void ExitWithScope() final;
};
 
class BlockInitFrame : public TIRFrame {
 public:
  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(BlockInitFrame, TIRFrame, BlockInitFrameNode);
};
 
class ForFrameNode : public TIRFrameNode {
 public:
  using FMakeForLoop = ffi::TypedFunction<tvm::tir::Stmt(
      Array<tvm::tir::Var> loop_vars, Array<Range> loop_extents, tvm::tir::Stmt loop_body)>;
  Array<tvm::tir::Var> vars;
  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.
  }
 
  static constexpr const char* _type_key = "script.ir_builder.tir.ForFrame";
  TVM_DECLARE_FINAL_OBJECT_INFO(ForFrameNode, TIRFrameNode);
 
 public:
  void ExitWithScope() final;
};
 
class ForFrame : public TIRFrame {
 public:
  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(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);
  }
 
  static constexpr const char* _type_key = "script.ir_builder.tir.AssertFrame";
  TVM_DECLARE_FINAL_OBJECT_INFO(AssertFrameNode, TIRFrameNode);
 
 public:
  void ExitWithScope() final;
};
 
class AssertFrame : public TIRFrame {
 public:
  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(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);
  }
 
  static constexpr const char* _type_key = "script.ir_builder.tir.LetFrame";
  TVM_DECLARE_FINAL_OBJECT_INFO(LetFrameNode, TIRFrameNode);
 
 public:
  void ExitWithScope() final;
};
 
class LetFrame : public TIRFrame {
 public:
  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(LetFrame, TIRFrame, LetFrameNode);
};
 
class LaunchThreadFrameNode : public TIRFrameNode {
 public:
  PrimExpr extent;
  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);
  }
 
  static constexpr const char* _type_key = "script.ir_builder.tir.LaunchThreadFrame";
  TVM_DECLARE_FINAL_OBJECT_INFO(LaunchThreadFrameNode, TIRFrameNode);
 
 public:
  void ExitWithScope() final;
};
 
class LaunchThreadFrame : public TIRFrame {
 public:
  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(LaunchThreadFrame, TIRFrame,
                                                    LaunchThreadFrameNode);
};
 
class RealizeFrameNode : public TIRFrameNode {
 public:
  tvm::tir::BufferRegion buffer_slice;
  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);
  }
 
  static constexpr const char* _type_key = "script.ir_builder.tir.RealizeFrame";
  TVM_DECLARE_FINAL_OBJECT_INFO(RealizeFrameNode, TIRFrameNode);
 
 public:
  void ExitWithScope() final;
};
 
class RealizeFrame : public TIRFrame {
 public:
  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(RealizeFrame, TIRFrame, RealizeFrameNode);
};
 
class AllocateFrameNode : public TIRFrameNode {
 public:
  Array<PrimExpr> extents;
  DataType dtype;
  String storage_scope;
  PrimExpr condition;
  Map<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);
  }
 
  static constexpr const char* _type_key = "script.ir_builder.tir.AllocateFrame";
  TVM_DECLARE_FINAL_OBJECT_INFO(AllocateFrameNode, TIRFrameNode);
 
 public:
  void ExitWithScope() final;
};
 
class AllocateFrame : public TIRFrame {
 public:
  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(AllocateFrame, TIRFrame, AllocateFrameNode);
};
 
class AllocateConstFrameNode : public TIRFrameNode {
 public:
  DataType dtype;
  Array<PrimExpr> extents;
  tvm::runtime::NDArray data;
  tvm::tir::Var buffer_var;
  Map<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);
  }
 
  static constexpr const char* _type_key = "script.ir_builder.tir.AllocateConstFrame";
  TVM_DECLARE_FINAL_OBJECT_INFO(AllocateConstFrameNode, TIRFrameNode);
 
 public:
  void ExitWithScope() final;
};
 
class AllocateConstFrame : public TIRFrame {
 public:
  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(AllocateConstFrame, TIRFrame,
                                                    AllocateConstFrameNode);
};
class AttrFrameNode : public TIRFrameNode {
 public:
  Any node;
  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);
  }
 
  static constexpr const char* _type_key = "script.ir_builder.tir.AttrFrame";
  TVM_DECLARE_FINAL_OBJECT_INFO(AttrFrameNode, TIRFrameNode);
 
 public:
  void ExitWithScope() final;
};
 
class AttrFrame : public TIRFrame {
 public:
  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(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);
  }
 
  static constexpr const char* _type_key = "script.ir_builder.tir.WhileFrame";
  TVM_DECLARE_FINAL_OBJECT_INFO(WhileFrameNode, TIRFrameNode);
 
 public:
  void ExitWithScope() final;
};
 
class WhileFrame : public TIRFrame {
 public:
  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(WhileFrame, TIRFrame, WhileFrameNode);
};
 
class IfFrameNode : public TIRFrameNode {
 public:
  PrimExpr condition;
  Optional<Array<tvm::tir::Stmt>> then_stmts;
  Optional<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);
  }
 
  static constexpr const char* _type_key = "script.ir_builder.tir.IfFrame";
  TVM_DECLARE_FINAL_OBJECT_INFO(IfFrameNode, TIRFrameNode);
 
 public:
  void ExitWithScope() final;
};
 
class IfFrame : public TIRFrame {
 public:
  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(IfFrame, TIRFrame, IfFrameNode);
};
 
class ThenFrameNode : public TIRFrameNode {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ThenFrameNode>();
  }
 
  static constexpr const char* _type_key = "script.ir_builder.tir.ThenFrame";
  TVM_DECLARE_FINAL_OBJECT_INFO(ThenFrameNode, TIRFrameNode);
 
 public:
  void EnterWithScope() final;
  void ExitWithScope() final;
};
 
class ThenFrame : public TIRFrame {
 public:
  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(ThenFrame, TIRFrame, ThenFrameNode);
};
 
class ElseFrameNode : public TIRFrameNode {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ElseFrameNode>();
  }
 
  static constexpr const char* _type_key = "script.ir_builder.tir.ElseFrame";
  TVM_DECLARE_FINAL_OBJECT_INFO(ElseFrameNode, TIRFrameNode);
 
 public:
  void EnterWithScope() final;
  void ExitWithScope() final;
};
 
class ElseFrame : public TIRFrame {
 public:
  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(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);
  }
 
  static constexpr const char* _type_key = "script.ir_builder.tir.DeclBufferFrame";
  TVM_DECLARE_FINAL_OBJECT_INFO(DeclBufferFrameNode, TIRFrameNode);
 
 public:
  void ExitWithScope() final;
};
 
class DeclBufferFrame : public TIRFrame {
 public:
  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(DeclBufferFrame, TIRFrame, DeclBufferFrameNode);
};
 
}  // namespace tir
}  // namespace ir_builder
}  // namespace script
}  // namespace tvm
 
#endif  // TVM_SCRIPT_IR_BUILDER_TIR_FRAME_H_