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>()
128 explicit Layout(
const Array<tir::IterVar>& axes);
184 if (!defined())
return 0;
190 if (!defined())
return 0;
192 for (
auto x : operator->()->axes) {
208 std::string new_src_layout_str =
"";
209 for (
auto dst_axis : dst_layout->
axes) {
212 new_src_layout_str += dst_axis->var->name_hint;
217 new_src_layout_str += this->
name();
218 new_src_layout =
Layout(new_src_layout_str);
219 return new_src_layout;
230 if (!this->defined())
return -1;
232 for (
size_t i = 0; i < axes.size(); ++i) {
253 if (!defined())
return false;
263 ICHECK(defined()) <<
"Try to access axis from an undefined layout.";
264 int32_t index = i < 0 ? static_cast<int32_t>(
ndim() + i) : i;
265 ICHECK(index >= 0 &&
static_cast<size_t>(index) <
ndim()) <<
"Invalid index " << i;
271 inline std::string
name()
const {
272 if (!defined())
return "__undef__";
318 refl::ObjectDef<BijectiveLayoutNode>()
327 static constexpr
const char*
_type_key =
"tir.BijectiveLayout";
349 TVM_DLL Array<PrimExpr>
BackwardShape(
const Array<PrimExpr>& dst_shape)
const;
351 TVM_DLL Array<PrimExpr>
ForwardIndex(
const Array<PrimExpr>& index)
const;
353 TVM_DLL Array<PrimExpr>
BackwardIndex(
const 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:298
Array< PrimExpr > shape_backward_rule
Describes how destination shapes can be mapped to the source shapes.
Definition: data_layout.h:309
Layout src_layout
The source layout.
Definition: data_layout.h:312
Layout dst_layout
The destination layout.
Definition: data_layout.h:314
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:303
Array< PrimExpr > index_backward_rule
Describes how destination axes can be mapped to the source axes.
Definition: data_layout.h:305
Array< PrimExpr > shape_forward_rule
Describes how source shapes can be mapped to the destination shapes.
Definition: data_layout.h:307
static void RegisterReflection()
Definition: data_layout.h:316
TVM_DECLARE_FINAL_OBJECT_INFO(BijectiveLayoutNode, Object)
static constexpr const char * _type_key
Definition: data_layout.h:327
Bijective function mapping for data layout transformation. Given two Layout, BijectiveLayout build an...
Definition: data_layout.h:337
BijectiveLayout(Layout src_layout, Layout dst_layout)
The constructor.
Array< PrimExpr > ForwardShape(const Array< PrimExpr > &shape) const
TVM_DEFINE_OBJECT_REF_METHODS(BijectiveLayout, ObjectRef, BijectiveLayoutNode)
Array< PrimExpr > BackwardShape(const Array< PrimExpr > &dst_shape) const
Array< PrimExpr > BackwardIndex(const Array< PrimExpr > &dst_index) const
Array< PrimExpr > ForwardIndex(const Array< PrimExpr > &index) const
Iteration Variable, represents an iteration over an integer interval.
Definition: var.h:298
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
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
TVM_DECLARE_FINAL_OBJECT_INFO(LayoutNode, Object)
static constexpr const char * _type_key
Definition: data_layout.h:118
static void RegisterReflection()
Definition: data_layout.h:111
String name
string representation of layout, "" for scalar.
Definition: data_layout.h:102
Managed reference to LayoutNode.
Definition: data_layout.h:126
Layout(const char *name)
construct from a string
Definition: data_layout.h:134
static const Layout & Undef()
Return an undefined layout.
Definition: data_layout.h:158
LayoutNode * operator->()
access the internal node container
Definition: data_layout.h:152
size_t ndim_primal() const
Definition: data_layout.h:189
bool Equals(const Layout &rhs) const
Whether the two layouts are equal.
Definition: data_layout.h:281
TVM_DEFINE_OBJECT_REF_METHODS(Layout, ObjectRef, LayoutNode)
friend std::ostream & operator<<(std::ostream &os, const Layout &l)
allow output string of layout to ostream
Definition: data_layout.h:289
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:229
size_t ndim() const
Definition: data_layout.h:183
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:205
std::string name() const
Definition: data_layout.h:271
bool Contains(const LayoutAxis &axis) const
Whether the layout contains an axis.
Definition: data_layout.h:252
const LayoutAxis & operator[](int32_t i) const
Definition: data_layout.h:262
int32_t FactorOf(const LayoutAxis &axis) const
Get the factor size of the subordinate axis.
Layout(const tvm::String &name)
construct from a string
Definition: data_layout.h:131
Layout(const std::string &name, DataType dtype=DataType::Int(32))
construct from a string.
Layout(const 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.
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:1945
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
Common operators defined for Expr.