Cutlass CUDA Templates for Linear Algebra Subroutines and Solvers

matrix_traits.h
Go to the documentation of this file.
1 /***************************************************************************************************
2 * Copyright (c) 2017-2018, NVIDIA CORPORATION. All rights reserved.
3 *
4 * Redistribution and use in source and binary forms, with or without modification, are permitted
5 * provided that the following conditions are met:
6 * * Redistributions of source code must retain the above copyright notice, this list of
7 * conditions and the following disclaimer.
8 * * Redistributions in binary form must reproduce the above copyright notice, this list of
9 * conditions and the following disclaimer in the documentation and/or other materials
10 * provided with the distribution.
11 * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
12 * to endorse or promote products derived from this software without specific prior written
13 * permission.
14 *
15 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
16 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
17 * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
18 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
19 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
20 * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
21 * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
22 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
23 *
24 **************************************************************************************************/
28 #pragma once
29
30 #include "cutlass/coord.h"
31
32 namespace cutlass {
33
35
38 struct MatrixCoord : public Coord<2, int> {
39
41 typedef int Index;
42
44 typedef Coord<2, Index> Base;
45
47 static int const kRow = 0;
48
50 static int const kColumn = 1;
51
52 //
53 // Methods
54 //
55
57 CUTLASS_HOST_DEVICE
58 MatrixCoord() { }
59
61 CUTLASS_HOST_DEVICE
62 MatrixCoord(Coord<2, Index> const &coord): Base(coord) { }
63
65 CUTLASS_HOST_DEVICE
66 MatrixCoord(Index row, Index column): Base(make_Coord(row, column)) { }
67
69 CUTLASS_HOST_DEVICE
70 Index const & row() const { return this->at(kRow); }
71
73 CUTLASS_HOST_DEVICE
74 Index & row() { return this->at(kRow); }
75
77 CUTLASS_HOST_DEVICE
78 Index const & column() const { return this->at(kColumn); }
79
81 CUTLASS_HOST_DEVICE
82 Index & column() { return this->at(kColumn); }
83
84 //
85 // Coord operators
86 //
87
89 CUTLASS_HOST_DEVICE
90 MatrixCoord operator+(Base const& b) const {
91 return MatrixCoord(Base::operator+(b));
92 }
93
95 CUTLASS_HOST_DEVICE
96 MatrixCoord operator-(Base const& b) const {
97 return MatrixCoord(Base::operator-(b));
98 }
99
101 CUTLASS_HOST_DEVICE
102 MatrixCoord operator*(Base const& b) const {
103 return MatrixCoord(Base::operator*(b));
104 }
105
107 CUTLASS_HOST_DEVICE
108 MatrixCoord operator/(Base const& b) const {
109 return MatrixCoord(Base::operator/(b));
110 }
111
113 CUTLASS_HOST_DEVICE
114 MatrixCoord& operator+=(Base const& b) {
115 Base::operator+=(b);
116 return *this;
117 }
118
120 CUTLASS_HOST_DEVICE
121 MatrixCoord& operator-=(Base const& b) {
122 Base::operator-=(b);
123 return *this;
124 }
125
127 CUTLASS_HOST_DEVICE
128 MatrixCoord& operator*=(Base const& b) {
129 Base::operator*=(b);
130 return *this;
131 }
132
134 CUTLASS_HOST_DEVICE
135 MatrixCoord& operator/=(Base const& b) {
136 Base::operator/=(b);
137 return *this;
138 }
139 };
140
142
144 //
145 // The following define classes satisfying the TensorRefMapFunc concept. These must support the
146 // following operations, where func is an instance of type TensorRefMapFunc.
147 //
148 // Coord<TensorRefMapFunc::kStorageRank> = func(Coord<kRank>);
149 //
150 // Though not required to be usable by TensorRef, each of the following also define a helper
151 // function to map the "leading dimension" to an appropriate stride vector. Implementations
152 // following this convention should also implement the following static method:
153 //
154 // Coord<TensorRefMapFunc::kStorageRank> stride = TensorRefMapFunc::stride(leading_dim);
155 //
156 namespace MatrixLayout {
157
159 enum Kind { kRowMajor, kColumnMajor };
160
161 //
162 // TensorRefMapFunc definitions for common layouts
163 //
164
166 struct RowMajor {
167 static int const kStorageRank = 2;
169 CUTLASS_HOST_DEVICE
170 Coord<kStorageRank> operator()(MatrixCoord const &coord) const {
171 return coord;
172 }
173 };
174
176 struct ColumnMajor {
177 static int const kStorageRank = 2;
179 CUTLASS_HOST_DEVICE
180 Coord<kStorageRank> operator()(MatrixCoord const &coord) const {
181 return make_Coord(coord.column(), coord.row());
182 }
183 };
184
187 template <int Interleave>
188 struct RowMajorInterleaved {
189
191 static int const kStorageRank = 3;
192
194 static int const kInterleave = Interleave;
195
197 CUTLASS_HOST_DEVICE
198 Coord<kStorageRank> operator()(MatrixCoord const &coord) const {
199 return make_Coord(
200 coord.row() / kInterleave,
201 coord.column(),
202 coord.row() % kInterleave
203 );
204 }
205
207 CUTLASS_HOST_DEVICE
208 static Coord<kStorageRank> stride(int ldm) {
209 return make_Coord(
210 ldm * kInterleave,
211 kInterleave,
212 1
213 );
214 }
215 };
216
219 template <int Interleave>
220 struct ColumnMajorInterleaved {
221
223 static int const kStorageRank = 3;
224
226 static int const kInterleave = Interleave;
227
229 CUTLASS_HOST_DEVICE
230 Coord<kStorageRank> operator()(MatrixCoord const &coord) const {
231 return make_Coord(
232 coord.column() / kInterleave,
233 coord.row(),
234 coord.column() % kInterleave
235 );
236 }
237
239 CUTLASS_HOST_DEVICE
240 static Coord<kStorageRank> stride(int ldm) {
241 return make_Coord(
242 ldm * kInterleave,
243 kInterleave,
244 1
245 );
246 }
247 };
248
251 struct ContiguousLayout {
253 static int const kStorageRank = 3;
254
256 static int const kRow = 0;
257
259 static int const kColumn = 1;
260
263 CUTLASS_HOST_DEVICE
264 Coord<kStorageRank> operator()(MatrixCoord const &coord) const {
265 return make_Coord(coord.row(), coord.column(), 0);
266 }
267
269 CUTLASS_HOST_DEVICE
270 static Coord<kStorageRank> stride(MatrixLayout::Kind layout, int ldm) {
271 if (layout == MatrixLayout::kRowMajor) {
272 return make_Coord(ldm, 1, 1);
273 }
274 return make_Coord(1, ldm, 1);
275 }
276 };
277
280 template <int BlockRows, int BlockColumns>
281 struct ColumnMajorBlockLinear {
282
284 static int const kStorageRank = 4;
285
287 static int const kBlockRows = BlockRows;
288
290 static int const kBlockColumns = BlockColumns;
291
293 CUTLASS_HOST_DEVICE
294 Coord<kStorageRank> operator()(MatrixCoord const &coord) const {
295 return make_Coord(
296 coord.column() / kBlockColumns,
297 coord.row() / kBlockRows,
298 coord.column() % kBlockColumns,
299 coord.row() % kBlockRows
300 );
301 }
302
304 CUTLASS_HOST_DEVICE
305 static Coord<kStorageRank> stride(int ldm) {
306 return make_Coord(
307 ldm * kBlockRows * kBlockColumns,
308 kBlockRows * kBlockColumns,
309 kBlockRows,
310 1
311 );
312 }
313 };
314
317 template <int BlockRows, int BlockColumns>
318 struct RowMajorBlockLinear {
319
321 static int const kStorageRank = 4;
322
324 static int const kBlockRows = BlockRows;
325
327 static int const kBlockColumns = BlockColumns;
328
330 CUTLASS_HOST_DEVICE
331 Coord<kStorageRank> operator()(MatrixCoord const &coord) const {
332 return make_Coord(
333 coord.row() / kBlockRows,
334 coord.column() / kBlockColumns,
335 coord.row() % kBlockRows,
336 coord.column() % kBlockColumns
337 );
338 }
339
341 CUTLASS_HOST_DEVICE
342 static Coord<kStorageRank> stride(int ldm) {
343 return make_Coord(
344 ldm * kBlockRows * kBlockColumns,
345 kBlockRows * kBlockColumns,
346 kBlockColumns,
347 1
348 );
349 }
350 };
351 };
352
354
356 struct GemmOperand {
357 enum Kind { kA, kB, kC, kD };
358 };
359
361
363 struct MatrixTransform {
364 enum Kind {
365 kNone,
366 kConjugate,
367 };
368 };
369
371
372 } // namespace cutlass
cutlass::MatrixCoord::Index
int Index
Integer-valued index.
Definition: matrix_traits.h:41
cutlass::MatrixLayout::ColumnMajor
Mapping function for column-major matrices.
Definition: matrix_traits.h:176
cutlass::MatrixLayout::RowMajorBlockLinear::kBlockColumns
static int const kBlockColumns
Interleaving size in columns dimension.
Definition: matrix_traits.h:327
cutlass
Definition: convert.h:33
cutlass::MatrixLayout::ColumnMajorInterleaved::operator()
CUTLASS_HOST_DEVICE Coord< kStorageRank > operator()(MatrixCoord const &coord) const
Maps (row, col) to (col, row, col)
Definition: matrix_traits.h:230
cutlass::MatrixLayout::RowMajor::operator()
CUTLASS_HOST_DEVICE Coord< kStorageRank > operator()(MatrixCoord const &coord) const
Maps (i, j) to (i, j)
Definition: matrix_traits.h:170
cutlass::MatrixTransform
Transformation applied to matrix operands.
Definition: matrix_traits.h:363
cutlass::MatrixLayout::RowMajorInterleaved
Definition: matrix_traits.h:188
cutlass::MatrixLayout::ColumnMajorBlockLinear::kBlockColumns
static int const kBlockColumns
Interleaving size in columns dimension.
Definition: matrix_traits.h:290
cutlass::MatrixTransform::kNone
Definition: matrix_traits.h:365
cutlass::MatrixLayout::ColumnMajorBlockLinear
Definition: matrix_traits.h:281
cutlass::MatrixLayout::ColumnMajorInterleaved
Definition: matrix_traits.h:220
coord.h
A Coord is a coordinate of arbitrary rank into a tensor or matrix.
cutlass::make_Coord
CUTLASS_HOST_DEVICE Coord< 1 > make_Coord(int _0)
Helper to make a 2-element coordinate.
Definition: coord.h:368
cutlass::MatrixTransform::kConjugate
no operation
Definition: matrix_traits.h:366
cutlass::MatrixCoord::operator/=
CUTLASS_HOST_DEVICE MatrixCoord & operator/=(Base const &b)
In-place division.
Definition: matrix_traits.h:135
cutlass::MatrixLayout::RowMajor::kStorageRank
static int const kStorageRank
Definition: matrix_traits.h:167
cutlass::MatrixLayout::ContiguousLayout
Definition: matrix_traits.h:251
cutlass::MatrixLayout::ColumnMajor::operator()
CUTLASS_HOST_DEVICE Coord< kStorageRank > operator()(MatrixCoord const &coord) const
Maps (i, j) to (j, i)
Definition: matrix_traits.h:180
cutlass::MatrixLayout::Kind
Kind
Enumeration defining fundamental contiguous layouts.
Definition: matrix_traits.h:159
cutlass::MatrixTransform::Kind
Kind
Definition: matrix_traits.h:364
cutlass::MatrixCoord::Base
Coord< 2, Index > Base
Base type is a Coord of rank=2.
Definition: matrix_traits.h:44
cutlass::MatrixCoord::operator+
CUTLASS_HOST_DEVICE MatrixCoord operator+(Base const &b) const
Element-wise addition.
Definition: matrix_traits.h:90
cutlass::Coord::operator*=
CUTLASS_HOST_DEVICE Coord & operator*=(Coord const &b)
In-place multiplication.
Definition: coord.h:197
cutlass::GemmOperand::kC
Definition: matrix_traits.h:357
cutlass::MatrixLayout::ContiguousLayout::kRow
static int const kRow
Dimension of rows.
Definition: matrix_traits.h:256
cutlass::MatrixLayout::ColumnMajor::kStorageRank
static int const kStorageRank
Definition: matrix_traits.h:177
cutlass::MatrixLayout::ColumnMajorBlockLinear::kBlockRows
static int const kBlockRows
Interleaving size in rows dimension.
Definition: matrix_traits.h:287
cutlass::MatrixLayout::RowMajorInterleaved::kInterleave
static int const kInterleave
Interleaving size.
Definition: matrix_traits.h:194
cutlass::MatrixCoord::column
CUTLASS_HOST_DEVICE Index const & column() const
Returns the column of the coordinate.
Definition: matrix_traits.h:78
cutlass::MatrixCoord::MatrixCoord
CUTLASS_HOST_DEVICE MatrixCoord(Index row, Index column)
Helper to construct from a row and column.
Definition: matrix_traits.h:66
cutlass::MatrixLayout::RowMajorInterleaved::stride
static CUTLASS_HOST_DEVICE Coord< kStorageRank > stride(int ldm)
Helper to compute stride vector from leading dimension.
Definition: matrix_traits.h:208
cutlass::MatrixLayout::ContiguousLayout::kColumn
static int const kColumn
Dimension of columns.
Definition: matrix_traits.h:259
cutlass::MatrixLayout::RowMajorInterleaved::kStorageRank
static int const kStorageRank
Rank of storage n-D array.
Definition: matrix_traits.h:191
cutlass::Coord::operator-=
CUTLASS_HOST_DEVICE Coord & operator-=(Coord const &b)
In-place subtraction.
Definition: coord.h:188
cutlass::MatrixLayout::ContiguousLayout::kStorageRank
static int const kStorageRank
Arbitrary storage rank.
Definition: matrix_traits.h:253
cutlass::GemmOperand::kB
Definition: matrix_traits.h:357
cutlass::Coord::operator+=
CUTLASS_HOST_DEVICE Coord & operator+=(Coord const &b)
In-place addition.
Definition: coord.h:179
cutlass::MatrixLayout::RowMajorInterleaved::operator()
CUTLASS_HOST_DEVICE Coord< kStorageRank > operator()(MatrixCoord const &coord) const
Maps (row, col) to (row, col, row)
Definition: matrix_traits.h:198
CUTLASS_HOST_DEVICE
#define CUTLASS_HOST_DEVICE
Definition: cutlass.h:46
cutlass::MatrixLayout::RowMajorBlockLinear::kBlockRows
static int const kBlockRows
Interleaving size in rows dimension.
Definition: matrix_traits.h:324
cutlass::MatrixCoord::row
CUTLASS_HOST_DEVICE Index const & row() const
Returns the row of the coordinate.
Definition: matrix_traits.h:70
cutlass::Coord< 2, int >::at
CUTLASS_HOST_DEVICE Index & at()
Gets the index of a given Coord element.
Definition: coord.h:240
cutlass::Coord::operator/=
CUTLASS_HOST_DEVICE Coord & operator/=(Coord const &b)
In-place division.
Definition: coord.h:206
cutlass::MatrixLayout::kColumnMajor
Definition: matrix_traits.h:159
cutlass::MatrixCoord::operator-
CUTLASS_HOST_DEVICE MatrixCoord operator-(Base const &b) const
Element-wise subtraction.
Definition: matrix_traits.h:96
cutlass::MatrixCoord::MatrixCoord
CUTLASS_HOST_DEVICE MatrixCoord(Coord< 2, Index > const &coord)
Constructs from Coord<2>
Definition: matrix_traits.h:62
cutlass::MatrixLayout::RowMajorBlockLinear::kStorageRank
static int const kStorageRank
Rank of storage n-D array.
Definition: matrix_traits.h:321
cutlass::Coord
Statically-sized array specifying Coords within a tensor.
Definition: coord.h:49
cutlass::MatrixLayout::kRowMajor
Definition: matrix_traits.h:159
cutlass::GemmOperand
Gemm operand - D = A * B + C.
Definition: matrix_traits.h:356
cutlass::MatrixLayout::RowMajorBlockLinear::stride
static CUTLASS_HOST_DEVICE Coord< kStorageRank > stride(int ldm)
Helper to compute stride vector from leading dimension.
Definition: matrix_traits.h:342
cutlass::MatrixCoord::kRow
static int const kRow
Rows dimension.
Definition: matrix_traits.h:47
cutlass::MatrixCoord::operator-=
CUTLASS_HOST_DEVICE MatrixCoord & operator-=(Base const &b)
In-place subtraction.
Definition: matrix_traits.h:121
cutlass::MatrixCoord::operator*
CUTLASS_HOST_DEVICE MatrixCoord operator*(Base const &b) const
Element-wise multiplication.
Definition: matrix_traits.h:102
cutlass::MatrixLayout::ColumnMajorInterleaved::stride
static CUTLASS_HOST_DEVICE Coord< kStorageRank > stride(int ldm)
Helper to compute stride vector from leading dimension.
Definition: matrix_traits.h:240
cutlass::MatrixLayout::ContiguousLayout::operator()
CUTLASS_HOST_DEVICE Coord< kStorageRank > operator()(MatrixCoord const &coord) const
Definition: matrix_traits.h:264
cutlass::MatrixCoord::row
CUTLASS_HOST_DEVICE Index & row()
Returns the row of the coordinate.
Definition: matrix_traits.h:74
cutlass::MatrixLayout::ColumnMajorBlockLinear::kStorageRank
static int const kStorageRank
Rank of storage n-D array.
Definition: matrix_traits.h:284
cutlass::MatrixLayout::ColumnMajorInterleaved::kInterleave
static int const kInterleave
Interleaving size.
Definition: matrix_traits.h:226
cutlass::MatrixLayout::RowMajorBlockLinear::operator()
CUTLASS_HOST_DEVICE Coord< kStorageRank > operator()(MatrixCoord const &coord) const
Maps (row, col) to (row, col, row, col)
Definition: matrix_traits.h:331
cutlass::GemmOperand::Kind
Kind
Definition: matrix_traits.h:357
cutlass::GemmOperand::kA
Definition: matrix_traits.h:357
cutlass::MatrixCoord::column
CUTLASS_HOST_DEVICE Index & column()
Returns the column of the coordinate.
Definition: matrix_traits.h:82
cutlass::MatrixCoord::operator*=
CUTLASS_HOST_DEVICE MatrixCoord & operator*=(Base const &b)
In-place multiplication.
Definition: matrix_traits.h:128
cutlass::MatrixLayout::ColumnMajorInterleaved::kStorageRank
static int const kStorageRank
Rank of storage n-D array.
Definition: matrix_traits.h:223
cutlass::MatrixLayout::RowMajorBlockLinear
Definition: matrix_traits.h:318
cutlass::MatrixCoord::operator+=
CUTLASS_HOST_DEVICE MatrixCoord & operator+=(Base const &b)
In-place addition.
Definition: matrix_traits.h:114
cutlass::MatrixLayout::ColumnMajorBlockLinear::stride
static CUTLASS_HOST_DEVICE Coord< kStorageRank > stride(int ldm)
Helper to compute stride vector from leading dimension.
Definition: matrix_traits.h:305
cutlass::MatrixLayout::ColumnMajorBlockLinear::operator()
CUTLASS_HOST_DEVICE Coord< kStorageRank > operator()(MatrixCoord const &coord) const
Maps (row, col) to (col, row, col, row)
Definition: matrix_traits.h:294
cutlass::MatrixLayout::ContiguousLayout::stride
static CUTLASS_HOST_DEVICE Coord< kStorageRank > stride(MatrixLayout::Kind layout, int ldm)
Helper to construct a stride vector based on contiguous matrix layout and leading dimension...
Definition: matrix_traits.h:270
cutlass::MatrixCoord
Definition: matrix_traits.h:38
cutlass::MatrixCoord::operator/
CUTLASS_HOST_DEVICE MatrixCoord operator/(Base const &b) const
Element-wise division.
Definition: matrix_traits.h:108
cutlass::MatrixCoord::kColumn
static int const kColumn
Columns dimension.
Definition: matrix_traits.h:50
cutlass::MatrixCoord::MatrixCoord
CUTLASS_HOST_DEVICE MatrixCoord()
Default ctor.
Definition: matrix_traits.h:58
cutlass::GemmOperand::kD
Definition: matrix_traits.h:357
cutlass::MatrixLayout::RowMajor
Mapping function for row-major matrices.
Definition: matrix_traits.h:166