layout.h 7.59 KB
Newer Older
1
2
3
4
5
6
7
8
9
/*!
 * \file Layout.h
 *
 */

#ifndef TVM_TL_LAYOUT_LAYOUT_H_
#define TVM_TL_LAYOUT_LAYOUT_H_

#include <tvm/arith/analyzer.h>
10
#include <tvm/arith/iter_affine_map.h>
11
#include <tvm/ffi/object.h>
12
#include <utility>
13

14
15
#include "../support/ffi_aliases.h"

16
17
18
19
20
21
22
23
24
namespace tvm {
namespace tl {

using namespace tir;

class Layout;
class Fragment;

class LayoutNode : public Object {
25
public:
26
27
28
29
30
31
32
33
34
35
36
37
38
  LayoutNode() = default;
  LayoutNode(Array<PrimExpr> input_size, Array<PrimExpr> forward_index);

  size_t InputDim() const { return input_size_.size(); }

  size_t OutputDim() const { return forward_index_.size(); }

  Array<PrimExpr> InputShape() const { return input_size_; }

  Array<PrimExpr> OutputShape() const;

  Array<PrimExpr> GetForwardIndex() const { return forward_index_; }

39
40
  virtual Array<PrimExpr> GetForwardVars() const;

41
  virtual Array<PrimExpr> Forward(const Array<PrimExpr> &vars) const;
42
43

  virtual Layout Inverse() const;
44
45
46
47

  virtual Layout Reshape(const Array<PrimExpr> &shape,
                         arith::Analyzer *analyzer) const;

48
  virtual std::pair<Layout, arith::IterMapLevel> InverseWithLevel() const;
49

50
51
52
  virtual std::string DebugOutput() const;

  virtual bool IsEqual(const LayoutNode *other, bool skip_index = false) const;
53

54
  static void RegisterReflection();
55
56
57
  TVM_FFI_DECLARE_OBJECT_INFO("tl.Layout", LayoutNode, Object);
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind =
      kTVMFFISEqHashKindTreeNode;
58

59
protected:
60
  virtual Map<Var, Range> getVarMap() const;
61
  void UpdateAnalyzer(arith::Analyzer *analyzer) const;
62
63
64
65
66
67
68
69
  Array<PrimExpr> forward_index_;
  Array<PrimExpr> input_size_;
};

/*!
 * \brief Layout reference class.
 */
class Layout : public ObjectRef {
70
public:
71
72
73
  TVM_DLL Layout(Array<IterVar> forward_var, Array<PrimExpr> forward_index);
  TVM_DLL Layout(Array<PrimExpr> input_size, Array<PrimExpr> forward_index);

74
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Layout, ObjectRef, LayoutNode);
75
76
77
};

class FragmentNode : public LayoutNode {
78
public:
79
  FragmentNode() = default;
80
81
  FragmentNode(Array<PrimExpr> input_size, Array<PrimExpr> forward_index,
               PrimExpr forward_thread, PrimExpr replicate_size);
82
83
84

  PrimExpr GetForwardThread() const { return forward_thread_; }

85
86
  Array<PrimExpr> GetForwardVars() const final;

87
  Layout Inverse() const final;
88
89
90

  Layout Reshape(const Array<PrimExpr> &shape, arith::Analyzer *analyzer) const;

91
  std::pair<Layout, arith::IterMapLevel> InverseWithLevel() const final;
92
93
94
95
96

  PrimExpr ThreadExtent() const;

  PrimExpr ReplicateExtent() const { return replicate_size_; };

97
98
  PrimExpr ForwardThread(const Array<PrimExpr> &vars,
                         const Optional<PrimExpr> &rep_var) const;
99

100
  Fragment Repeat(const Array<PrimExpr> &repeats, bool repeat_on_thread,
101
102
103
104
105
106
107
108
                  bool lower_dim_first = true) const;

  Fragment Replicate(int repeats) const;

  Fragment DeReplicate() const;

  Fragment CondenseReplicateVar() const;

109
110
  std::string DebugOutput() const final;

111
  Fragment BindThreadRange(Range thread_range) const;
112
113
114

  Range ThreadRange() const { return thread_range_; }

115
  bool IsEqual(const FragmentNode *other, bool skip_index = false) const;
116

117
118
  bool IsCompletedReplicated() const;

119
  static void RegisterReflection();
120

121
122
123
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tl.Fragment", FragmentNode, LayoutNode);
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind =
      kTVMFFISEqHashKindTreeNode;
124

125
protected:
126
  Map<Var, Range> getVarMap() const final;
127
  Range thread_range_;
128
129
130
131
132
133
134
135
  PrimExpr forward_thread_;
  PrimExpr replicate_size_;
};

/*!
 * \brief Fragment reference class.
 */
class Fragment : public Layout {
136
public:
137
138
139
140
  TVM_DLL Fragment(Array<IterVar> forward_var, Array<PrimExpr> forward_index,
                   PrimExpr forward_thread, IterVar thread_replicate);

  TVM_DLL Fragment(Array<PrimExpr> input_size, Array<PrimExpr> forward_index,
141
142
                   PrimExpr forward_thread, PrimExpr replicate_size,
                   Optional<Var> replicate_var);
143

144
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Fragment, Layout, FragmentNode);
145
146
147
148
};

Var InputPlaceholder(size_t idx);
Var ReplicationPlaceholder();
149
IterVar make_itervar(std::string name, PrimExpr dom);
150
151
152

Fragment makeGemmFragment8x8();
Fragment makeGemmFragment8x8Transposed();
153
154
Fragment makeGemmFragmentC(const int block_m, const int block_n,
                           const int warp_m, const int warp_n,
155
                           const int element_size);
156
157
158
Fragment makeGemmSparseFragmentC(const int block_m, const int block_n,
                                 const int warp_m, const int warp_n,
                                 const int element_size);
159
160
161
162
163
164
165
166
Fragment makeGemmFragmentCCDNA(const int block_m, const int block_n,
                               const int warp_m, const int warp_n,
                               const int element_size);
Fragment makeGemmFragmentCHopper(const int block_m, const int block_n,
                                 const int warp_m, const int warp_n,
                                 const int element_size);
Fragment makeGemmFragmentA(const int block_m, const int block_n,
                           const int block_k, const int warp_m,
167
168
                           const int warp_n, const int element_size,
                           bool transposed = false);
169
170
Fragment makeGemmFragmentB(const int block_m, const int block_n,
                           const int block_k, const int warp_m,
171
                           const int warp_n, bool transposed = false);
172
173
174

Fragment makeGemmFragmentACDNA(const int block_m, const int block_n,
                               const int block_k, const int warp_m,
175
                               const int warp_n, const int element_size,
176
                               const int k_pack, bool transposed = false);
177
178
179
180

// Default Memory Layout
Layout makeGemmLayoutLinear(int stride, int continuous);
Layout makeGemmABLayoutPadded(int stride, int continuous, int element_size);
181
Layout makeGemmABLayout(int mat_stride, int mat_continuous, int continuity,
182
                        int element_size, bool k_inner = true);
183
Layout makeGemmABLayoutHopper(int mat_stride, int mat_continuous,
184
185
                              int continuity, int element_size,
                              bool k_inner = true);
186
Layout makeGemmABLayoutSm100(int mat_stride, int mat_continuous, int continuity,
187
                             int element_size, bool k_inner = true);
188
Layout makeGemmABLayoutCDNA(int stride, int continuous, int element_size,
189
                            int kPack);
190
191
192
193
194
195
196
197

Fragment makeGemmVoltaFragmentC(const int block_m, const int block_n,
                                const int warp_m, const int warp_n,
                                const int element_size);
Fragment makeGemmVoltaFragmentA(const int block_m, const int block_n,
                                const int block_k, const int warp_m,
                                const int warp_n);
Layout makeGemmVoltaABLayout(int stride, int continuous, bool is_a,
198
                             bool k_inner = true);
199

200
201
202
203
204
Layout makeTensorOpMultiplicand(int mat_stride, int mat_continuous,
                                int elementsize, int crosswise);
Layout makeGemmSparseAmpereABLayout(int mat_stride, int mat_continuous,
                                    int elementsize);

205
206
Layout makeFullBankSwizzleLayout(int stride, int continuous, int element_size);
Layout makeHalfBankSwizzleLayout(int stride, int continuous, int element_size);
207
208
Layout makeQuarterBankSwizzleLayout(int stride, int continuous,
                                    int element_size);
209
210
211

namespace attr {
// BlockAttr, Containing the layout for all the buffers in the block
212
213
constexpr const char *kLayoutMap = "layout_map";
} // namespace attr
214

215
216
} // namespace tl
} // namespace tvm
217

218
#endif // TVM_TL_LAYOUT_LAYOUT_H_