1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
|
/***************************************************************************************************
* Copyright (c) 2017-2018, NVIDIA CORPORATION. All rights reserved.
*
* Redistribution and use in source and binary forms, with or without modification, are permitted
* provided that the following conditions are met:
* * Redistributions of source code must retain the above copyright notice, this list of
* conditions and the following disclaimer.
* * Redistributions in binary form must reproduce the above copyright notice, this list of
* conditions and the following disclaimer in the documentation and/or other materials
* provided with the distribution.
* * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
* to endorse or promote products derived from this software without specific prior written
* permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
* IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
* FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
* BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
* OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
* STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
/*! \file
\brief Defines Shape implementing the Layout concept for representing a 4D hypercube of objects.
*/
#pragma once
#include <cutlass/cutlass.h>
namespace cutlass {
////////////////////////////////////////////////////////////////////////////////////////////////////
/*!@defgroup layout_concept Layout Concept
* @{
* @par Implementations of \ref layout_concept are used to describe a cube with DxHxW elements and C
scalars per element.
A HxW slice of a cube is called an image and a cube consists of D images.
*
* @par Notations
* Let Layout be an implementation of the \ref layout_concept.
*
* @par Valid Expressions
* - <b>Layout::D</b> specifies the depth of a cube
* - <b>Layout::H</b> specifies the height of a cube
* - <b>Layout::W</b> specifies the height of a cube
* - <b>Layout::C</b> specifies the number of channels of each element in a cube
* - <b>Layout::W_c</b> specifies the number of scalars of each row in one image of a cube.
* - <b>Layout::H_w</b> specifies the number of elements in an image slice.
* - <b>Layout::H_w_c</b>_specifies the number of scalars in an image slice.
* - <b>Layout::D_h_w</b> specifies the number of elements in a cube.
* - <b>Layout::D_h_w_c</b> specifies the number of scalars in a cube.
* - <b>Layout::Strides</b> is a \ref layout_concept specifying the strides.
* @}
*/
/**
* @brief A Shape implementing \ref layout_concept describing the dimensions of a cube.
* @concept{layout_concept}
*/
template <int kD_ = 1, int kH_ = 1, int kW_ = 1, int kC_ = 1>
struct Shape {
/// The depth of the cube.
static int const kD = kD_;
/// The height of the cube.
static int const kH = kH_;
/// The width of the cube.
static int const kW = kW_;
/// The number of scalars per element.
static int const kC = kC_;
};
/**
* @brief Compute derived counted of a \ref layout_concept based class
*/
template <typename Shape>
struct ShapeCount {
/// The number of elements per row.
static int const kWc = Shape::kW * Shape::kC;
/// The number of pixels per image.
static int const kHw = Shape::kH * Shape::kW;
/// The number of elements per image.
static int const kHwc = Shape::kH * kWc;
/// The number of pixels per cube.
static int const kDhw = Shape::kD * kHw;
/// The number of elements in the 4D space.
static int const kDhwc = Shape::kD * kHwc;
/// The number of elements in the 4D space.
static int const kCount = kDhwc;
};
////////////////////////////////////////////////////////////////////////////////////////////////////
template <typename A_, int kScale_>
struct ShapeScale {
typedef Shape<A_::kD * kScale_, A_::kH * kScale_, A_::kW * kScale_, A_::kC * kScale_> Shape;
};
////////////////////////////////////////////////////////////////////////////////////////////////////
template <typename A_, typename B_>
struct ShapeAdd {
typedef Shape<A_::kD + B_::kD, A_::kH + B_::kH, A_::kW + B_::kW, A_::kC + B_::kC> Shape;
};
////////////////////////////////////////////////////////////////////////////////////////////////////
template <typename A_, typename B_>
struct ShapeSub {
typedef Shape<A_::kD - B_::kD, A_::kH - B_::kH, A_::kW - B_::kW, A_::kC - B_::kC> Shape;
};
////////////////////////////////////////////////////////////////////////////////////////////////////
template <typename A_, typename B_>
struct ShapeMul {
typedef Shape<A_::kD * B_::kD, A_::kH * B_::kH, A_::kW * B_::kW, A_::kC * B_::kC> Shape;
};
////////////////////////////////////////////////////////////////////////////////////////////////////
template <typename A_, typename B_>
struct ShapeDiv {
typedef Shape<A_::kD / B_::kD, A_::kH / B_::kH, A_::kW / B_::kW, A_::kC / B_::kC> Shape;
};
////////////////////////////////////////////////////////////////////////////////////////////////////
template <typename A_, typename B_>
struct ShapeMax {
typedef Shape<(A_::kD > B_::kD ? A_::kD : B_::kD),
(A_::kH > B_::kH ? A_::kH : B_::kH),
(A_::kW > B_::kW ? A_::kW : B_::kW),
(A_::kC > B_::kC ? A_::kC : B_::kC)>
Shape;
};
////////////////////////////////////////////////////////////////////////////////////////////////////
template <typename A_, typename B_>
struct ShapeMin {
typedef Shape<(A_::kD < B_::kD ? A_::kD : B_::kD),
(A_::kH < B_::kH ? A_::kH : B_::kH),
(A_::kW < B_::kW ? A_::kW : B_::kW),
(A_::kC < B_::kC ? A_::kC : B_::kC)>
Shape;
};
////////////////////////////////////////////////////////////////////////////////////////////////////
template <typename Shape_, int kElementsPerAccess>
struct ShapeStrides {
typedef Shape<Shape_::kH * Shape_::kW * Shape_::kC,
Shape_::kW * Shape_::kC,
Shape_::kC,
kElementsPerAccess>
Shape;
};
////////////////////////////////////////////////////////////////////////////////////////////////////
/**
* @brief Compute the offset for the given coordinates in a cube
* @tparam A \ref layout_concept where each dimension of the cube specifies the corresponding stride.
*/
template <typename Shape_>
struct ComputeOffsetFromShape {
static CUTLASS_DEVICE int get(int d, int h, int w, int c) {
// clang-format off
return d * Shape_::kH * Shape_::kW * Shape_::kC +
h * Shape_::kW * Shape_::kC +
w * Shape_::kC +
c;
// clang-format on
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
/**
* @brief Compute the offset for the given coordinates in a cube with a depth of 1
* @tparam kSh Elements in the H dimension
* @tparam kSw Elements in the W dimension
* @tparam kSc Separation between two elements in "elements"
*/
template <int kSh_, int kSw_, int kSc_>
struct ComputeOffsetFromShape<Shape<1, kSh_, kSw_, kSc_> > {
static CUTLASS_DEVICE int get(int d, int h, int w, int c) {
return h * kSw_ * kSc_ + w * kSc_ + c;
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
/**
* @brief Compute the offset for the given coordinates in a cube with one channel and a depth of 1
* @tparam kSh Elements in the H dimension
* @tparam kSw Elements in the W dimension
*/
template <int kSh_, int kSw_>
struct ComputeOffsetFromShape<Shape<1, kSh_, kSw_, 1> > {
static CUTLASS_DEVICE int get(int d, int h, int w, int c) { return h * kSw_ + w; }
};
////////////////////////////////////////////////////////////////////////////////////////////////////
/**
* @brief Compute the offset for the given coordinates in a cube
* @tparam A \ref layout_concept where each dimension of the cube specifies the corresponding stride.
*/
template <typename Strides_>
struct ComputeOffsetFromStrides {
static CUTLASS_DEVICE int get(int d, int h, int w, int c) {
return d * Strides_::kD + h * Strides_::kH + w * Strides_::kW + c * Strides_::kC;
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
/**
* @brief Compute the offset for the given coordinates in a cube with a depth of 1
* @tparam S_h Stride in the H dimension in scalars
* @tparam S_w Stride in the W dimension in scalars
* @tparam S_c Stride between two scalars.
*/
template <int S_h_, int S_w_, int S_c_>
struct ComputeOffsetFromStrides<Shape<1, S_h_, S_w_, S_c_> > {
static CUTLASS_DEVICE int get(int d, int h, int w, int c) {
return h * S_h_ + w * S_w_ + c * S_c_;
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
/**
* @brief Compute the offset for the given coordinates in a cube with one channel and a depth of 1
* @tparam S_h Stride in the H dimension in scalars
* @tparam S_w Stride in the W dimension in scalars
*/
template <int S_h_, int S_w_>
struct ComputeOffsetFromStrides<Shape<1, S_h_, S_w_, 1> > {
static CUTLASS_DEVICE int get(int d, int h, int w, int c) { return h * S_h_ + w * S_w_; }
};
////////////////////////////////////////////////////////////////////////////////////////////////////
/**
* @brief Decompose threadId.x into coordinate of a cube whose dimensions are specified by Threads_.
* Afterwards compute the offset of those coordinates using Strides_
* @tparam Threads_ The dimension of the cube the threadIdx.x value is mapped on
* @tparam Strides_ The strides to use when compute the offsets based on the coordinates of the cube.
*/
template <typename Threads_, typename Strides_>
struct ComputeThreadOffsetFromStrides {
static CUTLASS_DEVICE int get() {
// Decompose the thread index.
int c = threadIdx.x % Threads_::kC;
int w = threadIdx.x / Threads_::kC % Threads_::kW;
int h = threadIdx.x / Threads_::kC / Threads_::kW % Threads_::kH;
int d = threadIdx.x / Threads_::kC / Threads_::kW / Threads_::kH;
// Compute the offset.
return d * Strides_::kD + h * Strides_::kH + w * Strides_::kW + c * Strides_::kC;
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
/**
*@brief Specialization for D=1
*/
template <int T_h_, int T_w_, int T_c_, int S_h_, int S_w_, int S_c_>
struct ComputeThreadOffsetFromStrides<Shape<1, T_h_, T_w_, T_c_>, Shape<1, S_h_, S_w_, S_c_> > {
static CUTLASS_DEVICE int get() {
// Decompose the thread index.
int c = threadIdx.x % T_c_;
int w = threadIdx.x / T_c_ % T_w_;
int h = threadIdx.x / T_c_ / T_w_ % T_h_;
// Compute the offset.
return h * S_h_ + w * S_w_ + c * S_c_;
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
/**
*@brief Specialization for D=1 and C=1
*/
template <int T_h_, int T_w_, int S_h_, int S_w_>
struct ComputeThreadOffsetFromStrides<Shape<1, T_h_, T_w_, 1>, Shape<1, S_h_, S_w_, 1> > {
static CUTLASS_DEVICE int get() {
// Decompose the thread index.
int w = threadIdx.x % T_w_;
int h = threadIdx.x / T_w_;
// Compute the offset.
return h * S_h_ + w * S_w_;
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass
|
