25 #ifndef TVM_TIR_DATA_LAYOUT_H_
26 #define TVM_TIR_DATA_LAYOUT_H_
28 #include <tvm/ffi/reflection/registry.h>
53 inline bool IsPrimal()
const {
return name_ >=
'A' && name_ <=
'Z'; }
54 inline std::string
name()
const {
return std::string(1, name_); }
59 if (name_ >=
'A' && name_ <=
'Z') {
113 refl::ObjectDef<LayoutNode>()
126 explicit Layout(
const ffi::Array<tir::IterVar>& axes);
182 if (!defined())
return 0;
188 if (!defined())
return 0;
190 for (
auto x : operator->()->axes) {
206 std::string new_src_layout_str =
"";
207 for (
auto dst_axis : dst_layout->
axes) {
210 new_src_layout_str += dst_axis->var->name_hint;
215 new_src_layout_str += this->
name();
216 new_src_layout =
Layout(new_src_layout_str);
217 return new_src_layout;
228 if (!this->defined())
return -1;
230 for (
size_t i = 0; i < axes.size(); ++i) {
251 if (!defined())
return false;
261 ICHECK(defined()) <<
"Try to access axis from an undefined layout.";
262 int32_t index = i < 0 ? static_cast<int32_t>(
ndim() + i) : i;
263 ICHECK(index >= 0 &&
static_cast<size_t>(index) <
ndim()) <<
"Invalid index " << i;
269 inline std::string
name()
const {
270 if (!defined())
return "__undef__";
316 refl::ObjectDef<BijectiveLayoutNode>()
345 TVM_DLL ffi::Array<PrimExpr>
BackwardShape(
const ffi::Array<PrimExpr>& dst_shape)
const;
347 TVM_DLL ffi::Array<PrimExpr>
ForwardIndex(
const ffi::Array<PrimExpr>& index)
const;
349 TVM_DLL ffi::Array<PrimExpr>
BackwardIndex(
const ffi::Array<PrimExpr>& dst_index)
const;
Runtime primitive data type.
Definition: data_type.h:47
static DataType Int(int bits, int lanes=1)
Construct an int type.
Definition: data_type.h:274
Definition: data_layout.h:296
Layout src_layout
The source layout.
Definition: data_layout.h:310
ffi::Array< PrimExpr > shape_forward_rule
Describes how source shapes can be mapped to the destination shapes.
Definition: data_layout.h:305
Layout dst_layout
The destination layout.
Definition: data_layout.h:312
ffi::Array< PrimExpr > index_forward_rule
Describes how source axes can be mapped to the destination axes, e.g., [i0 / 16, i1,...
Definition: data_layout.h:301
ffi::Array< PrimExpr > shape_backward_rule
Describes how destination shapes can be mapped to the source shapes.
Definition: data_layout.h:307
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tir.BijectiveLayout", BijectiveLayoutNode, Object)
static void RegisterReflection()
Definition: data_layout.h:314
ffi::Array< PrimExpr > index_backward_rule
Describes how destination axes can be mapped to the source axes.
Definition: data_layout.h:303
Bijective function mapping for data layout transformation. Given two Layout, BijectiveLayout build an...
Definition: data_layout.h:333
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(BijectiveLayout, ObjectRef, BijectiveLayoutNode)
BijectiveLayout(Layout src_layout, Layout dst_layout)
The constructor.
ffi::Array< PrimExpr > BackwardIndex(const ffi::Array< PrimExpr > &dst_index) const
ffi::Array< PrimExpr > BackwardShape(const ffi::Array< PrimExpr > &dst_shape) const
ffi::Array< PrimExpr > ForwardIndex(const ffi::Array< PrimExpr > &index) const
ffi::Array< PrimExpr > ForwardShape(const ffi::Array< PrimExpr > &shape) const
Iteration Variable, represents an iteration over an integer interval.
Definition: var.h:297
Definition: data_layout.h:43
std::string name() const
Definition: data_layout.h:54
static const LayoutAxis & Get(const char name)
const LayoutAxis & ToPrimal() const
Definition: data_layout.h:67
bool IsPrimal() const
Definition: data_layout.h:53
static const LayoutAxis & Get(const tir::IterVar &itvar)
const LayoutAxis & ToSubordinate() const
Definition: data_layout.h:70
const LayoutAxis & ToDual() const
Definition: data_layout.h:58
friend std::ostream & operator<<(std::ostream &os, const LayoutAxis &l)
Definition: data_layout.h:74
static const LayoutAxis & Get(const std::string &name)
bool operator==(const LayoutAxis &rhs) const
Definition: data_layout.h:72
Layout is to describe how data is organized within an N-dimention tensor. It is composed of upper cas...
Definition: data_layout.h:99
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tir.Layout", LayoutNode, Object)
ffi::Array< tir::IterVar > axes
specify each axis of the layout, in which the variable name is the name of the axis....
Definition: data_layout.h:109
static void RegisterReflection()
Definition: data_layout.h:111
ffi::String name
string representation of layout, "" for scalar.
Definition: data_layout.h:102
Managed reference to LayoutNode.
Definition: data_layout.h:124
Layout(const char *name)
construct from a string
Definition: data_layout.h:132
static const Layout & Undef()
Return an undefined layout.
Definition: data_layout.h:156
LayoutNode * operator->()
access the internal node container
Definition: data_layout.h:150
size_t ndim_primal() const
Definition: data_layout.h:187
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Layout, ObjectRef, LayoutNode)
bool Equals(const Layout &rhs) const
Whether the two layouts are equal.
Definition: data_layout.h:279
friend std::ostream & operator<<(std::ostream &os, const Layout &l)
allow output string of layout to ostream
Definition: data_layout.h:287
Layout SubLayout(size_t pos, size_t len) const
Returns a sub-layout which is the portion of the object that starts at dimension pos and spans len di...
int32_t IndexOf(const LayoutAxis &axis) const
return the index of the input axis. If it is not found in the layout or the layout is undefined,...
Definition: data_layout.h:227
Layout(const tvm::ffi::String &name)
construct from a string
Definition: data_layout.h:129
size_t ndim() const
Definition: data_layout.h:181
Layout ExpandPrimal(const Layout &dst_layout)
Returns a new layout where the dims have been expanded to match the primal dimensions.
Definition: data_layout.h:203
std::string name() const
Definition: data_layout.h:269
bool Contains(const LayoutAxis &axis) const
Whether the layout contains an axis.
Definition: data_layout.h:250
const LayoutAxis & operator[](int32_t i) const
Definition: data_layout.h:260
int32_t FactorOf(const LayoutAxis &axis) const
Get the factor size of the subordinate axis.
Layout(const std::string &name, DataType dtype=DataType::Int(32))
construct from a string.
Layout(const ffi::Array< tir::IterVar > &axes)
Layout Split(const LayoutAxis &axis, size_t target_pos, int32_t factor) const
Split axis by size and put the sub-axis to position target_pos.
ffi::String name_hint
The hint to the variable name.
Definition: var.h:54
Definition: repr_printer.h:91
Var var(std::string name_hint, DataType t=DataType::Int(32))
Construct a new Var expression.
const Op & undef()
Returns an initialized but arbitrary value.
Tensor shape(const Tensor &src, DataType dtype, const std::string name="T_shape", const std::string tag=kInjective)
Get the shape of input tensor.
Definition: transform.h:1960
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
Common operators defined for Expr.