CUTLASS: gemm.h Source File

| CUTLASS CUDA Templates for Linear Algebra Subroutines and Solvers

- Main Page
- Modules
- Namespaces
- Classes
- Files

- File List
- File Members

include/cutlass/gemm/gemm.h
Go to the documentation of this file.
1 /***************************************************************************************************
2 * Copyright (c) 2017-2019, 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/cutlass.h"
31 #include "cutlass/coord.h"
32
33 namespace cutlass {
34 namespace gemm {
35
37
39 enum class Operand {
40 kA,
41 kB,
42 kC,
43 kD
44 };
45
47
49 template <
51 int M = 1,
53 int N = 1,
55 int K = 1
56 >
57 struct GemmShape {
58 static int const kM = M;
59 static int const kN = N;
60 static int const kK = K;
61
62 static int const kMN = M * N;
63 static int const kMK = M * K;
64 static int const kKN = N * K;
65 static int const kMNK = M * N * K;
66
67 static int const kCount = kMNK;
68
69
70 //
71 // Static member functions
72 //
73
75 CUTLASS_HOST_DEVICE
76 static Coord<3> toCoord() {
77 return make_Coord(kM, kN, kK);
78 }
79 };
80
82
84 template <
86 typename Shape
87 >
88 using GemmShapeTranspose = GemmShape<Shape::kN, Shape::kM, Shape::kK>;
89
91
94 struct GemmCoord : public Coord<3, int> {
95
97 typedef int Index;
98
100 typedef Coord<3, Index> Base;
101
103 static int const kM = 0;
104
106 static int const kN = 1;
107
109 static int const kK = 2;
110
111 //
112 // Methods
113 //
114
116 CUTLASS_HOST_DEVICE
117 GemmCoord() { }
118
120 CUTLASS_HOST_DEVICE
121 GemmCoord(Coord<3, Index> const &coord): Base(make_Coord(coord[0], coord[1], coord[2])) { }
122
124 CUTLASS_HOST_DEVICE
125 GemmCoord(Index m, Index n, Index k): Base(make_Coord(m, n, k)) { }
126
128 CUTLASS_HOST_DEVICE
129 Index const & m() const { return this->at(kM); }
130
132 CUTLASS_HOST_DEVICE
133 Index & m() { return this->at(kM); }
134
136 CUTLASS_HOST_DEVICE
137 Index const & n() const { return this->at(kN); }
138
140 CUTLASS_HOST_DEVICE
141 Index & n() { return this->at(kN); }
142
144 CUTLASS_HOST_DEVICE
145 Index const & k() const { return this->at(kK); }
146
148 CUTLASS_HOST_DEVICE
149 Index & k() { return this->at(kK); }
150
152 CUTLASS_HOST_DEVICE
153 Coord<3> mnk() const {
154 return make_Coord(m(), n(), k());
155 }
156
158 CUTLASS_HOST_DEVICE
159 Coord<3> knm() const {
160 return make_Coord(k(), n(), m());
161 }
162
164 CUTLASS_HOST_DEVICE
165 Coord<2> nm() const {
166 return make_Coord(n(), m());
167 }
168
170 CUTLASS_HOST_DEVICE
171 Coord<2> mn() const {
172 return make_Coord(m(), n());
173 }
174
176 CUTLASS_HOST_DEVICE
177 Coord<2> mk() const {
178 return make_Coord(m(), k());
179 }
180
182 CUTLASS_HOST_DEVICE
183 Coord<2> km() const {
184 return make_Coord(k(), m());
185 }
186
188 CUTLASS_HOST_DEVICE
189 Coord<2> nk() const {
190 return make_Coord(n(), k());
191 }
192
194 CUTLASS_HOST_DEVICE
195 Coord<2> kn() const {
196 return make_Coord(k(), n());
197 }
198
199 //
200 // Coord operators
201 //
202
204 CUTLASS_HOST_DEVICE
205 GemmCoord operator+(Base const& b) const {
206 return GemmCoord(Base::operator+(b));
207 }
208
210 CUTLASS_HOST_DEVICE
211 GemmCoord operator-(Base const& b) const {
212 return GemmCoord(Base::operator-(b));
213 }
214
216 CUTLASS_HOST_DEVICE
217 GemmCoord operator*(Base const& b) const {
218 return GemmCoord(Base::operator*(b));
219 }
220
222 CUTLASS_HOST_DEVICE
223 GemmCoord operator/(Base const& b) const {
224 return GemmCoord(Base::operator/(b));
225 }
226
228 CUTLASS_HOST_DEVICE
229 GemmCoord& operator+=(Base const& b) {
230 Base::operator+=(b);
231 return *this;
232 }
233
235 CUTLASS_HOST_DEVICE
236 GemmCoord& operator-=(Base const& b) {
237 Base::operator-=(b);
238 return *this;
239 }
240
242 CUTLASS_HOST_DEVICE
243 GemmCoord& operator*=(Base const& b) {
244 Base::operator*=(b);
245 return *this;
246 }
247
249 CUTLASS_HOST_DEVICE
250 GemmCoord& operator/=(Base const& b) {
251 Base::operator/=(b);
252 return *this;
253 }
254 };
255
257
260 struct BatchedGemmCoord : public Coord<4, int> {
261
263 typedef int Index;
264
266 typedef Coord<4, Index> Base;
267
269 static int const kM = 0;
270
272 static int const kN = 1;
273
275 static int const kK = 2;
276
278 static int const kBatch = 3;
279
280 //
281 // Methods
282 //
283
285 CUTLASS_HOST_DEVICE
286 BatchedGemmCoord() { }
287
289 CUTLASS_HOST_DEVICE
290 BatchedGemmCoord(Base const &coord): Base(coord) { }
291
293 CUTLASS_HOST_DEVICE
294 BatchedGemmCoord(Index m, Index n, Index k, Index b): Base(make_Coord(m, n, k, b)) { }
295
297 CUTLASS_HOST_DEVICE
298 Index const & m() const { return this->at(kM); }
299
301 CUTLASS_HOST_DEVICE
302 Index & m() { return this->at(kM); }
303
305 CUTLASS_HOST_DEVICE
306 Index const & n() const { return this->at(kN); }
307
309 CUTLASS_HOST_DEVICE
310 Index & n() { return this->at(kN); }
311
313 CUTLASS_HOST_DEVICE
314 Index const & k() const { return this->at(kK); }
315
317 CUTLASS_HOST_DEVICE
318 Index & k() { return this->at(kK); }
319
321 CUTLASS_HOST_DEVICE
322 Index const & batch() const { return this->at(kBatch); }
323
325 CUTLASS_HOST_DEVICE
326 Index & batch() { return this->at(kBatch); }
327
329 CUTLASS_HOST_DEVICE
330 GemmCoord mnk() const {
331 return GemmCoord(m(), n(), k());
332 }
333
335 CUTLASS_HOST_DEVICE
336 Coord<4> mnkb() const {
337 return make_Coord(m(), n(), k(), batch());
338 }
339
340 //
341 // Coord operators
342 //
343
345 CUTLASS_HOST_DEVICE
346 BatchedGemmCoord operator+(Base const& b) const {
347 return BatchedGemmCoord(Base::operator+(b));
348 }
349
351 CUTLASS_HOST_DEVICE
352 BatchedGemmCoord operator-(Base const& b) const {
353 return BatchedGemmCoord(Base::operator-(b));
354 }
355
357 CUTLASS_HOST_DEVICE
358 BatchedGemmCoord operator*(Base const& b) const {
359 return BatchedGemmCoord(Base::operator*(b));
360 }
361
363 CUTLASS_HOST_DEVICE
364 BatchedGemmCoord operator/(Base const& b) const {
365 return BatchedGemmCoord(Base::operator/(b));
366 }
367
369 CUTLASS_HOST_DEVICE
370 BatchedGemmCoord& operator+=(Base const& b) {
371 Base::operator+=(b);
372 return *this;
373 }
374
376 CUTLASS_HOST_DEVICE
377 BatchedGemmCoord& operator-=(Base const& b) {
378 Base::operator-=(b);
379 return *this;
380 }
381
383 CUTLASS_HOST_DEVICE
384 BatchedGemmCoord& operator*=(Base const& b) {
385 Base::operator*=(b);
386 return *this;
387 }
388
390 CUTLASS_HOST_DEVICE
391 BatchedGemmCoord& operator/=(Base const& b) {
392 Base::operator/=(b);
393 return *this;
394 }
395 };
396
398
399 } // namespace gemm
400 } // namespace cutlass
cutlass::gemm::BatchedGemmCoord::mnkb
CUTLASS_HOST_DEVICE Coord< 4 > mnkb() const
Obtains a Coord<4> from BatchedGemmCoord.
Definition: include/cutlass/gemm/gemm.h:336
cutlass::gemm::BatchedGemmCoord::Base
Coord< 4, Index > Base
Base type is a Coord of rank=4.
Definition: include/cutlass/gemm/gemm.h:266
cutlass::gemm::BatchedGemmCoord::m
CUTLASS_HOST_DEVICE Index & m()
Returns reference to the GEMM M coordinate.
Definition: include/cutlass/gemm/gemm.h:302
cutlass
Definition: aligned_buffer.h:35
cutlass::gemm::BatchedGemmCoord::operator/
CUTLASS_HOST_DEVICE BatchedGemmCoord operator/(Base const &b) const
Element-wise division.
Definition: include/cutlass/gemm/gemm.h:364
cutlass::gemm::GemmCoord::operator/=
CUTLASS_HOST_DEVICE GemmCoord & operator/=(Base const &b)
In-place division.
Definition: include/cutlass/gemm/gemm.h:250
cutlass::gemm::GemmCoord::Index
int Index
Integer-valued index.
Definition: include/cutlass/gemm/gemm.h:97
cutlass::gemm::GemmCoord::GemmCoord
CUTLASS_HOST_DEVICE GemmCoord(Coord< 3, Index > const &coord)
Constructs from Coord<3> and a batch.
Definition: include/cutlass/gemm/gemm.h:121
cutlass::gemm::BatchedGemmCoord::mnk
CUTLASS_HOST_DEVICE GemmCoord mnk() const
Obtains a GemmCoord from BatchedGemmCoord.
Definition: include/cutlass/gemm/gemm.h:330
cutlass::gemm::GemmCoord::operator+
CUTLASS_HOST_DEVICE GemmCoord operator+(Base const &b) const
Element-wise addition.
Definition: include/cutlass/gemm/gemm.h:205
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:387
cutlass::gemm::GemmCoord::m
CUTLASS_HOST_DEVICE Index & m()
Returns reference to the GEMM M coordinate.
Definition: include/cutlass/gemm/gemm.h:133
cutlass::gemm::Operand
Operand
GEMM operand enumeration: D = A * B + C.
Definition: include/cutlass/gemm/gemm.h:39
cutlass::gemm::GemmCoord
Definition: include/cutlass/gemm/gemm.h:94
cutlass::gemm::GemmCoord::mn
CUTLASS_HOST_DEVICE Coord< 2 > mn() const
Obtains a Coord<2> from GemmCoord.
Definition: include/cutlass/gemm/gemm.h:171
cutlass::operator/=
CUTLASS_HOST_DEVICE half_t & operator/=(half_t &lhs, half_t const &rhs)
Definition: half.h:684
cutlass::gemm::BatchedGemmCoord::operator-
CUTLASS_HOST_DEVICE BatchedGemmCoord operator-(Base const &b) const
Element-wise subtraction.
Definition: include/cutlass/gemm/gemm.h:352
cutlass::gemm::BatchedGemmCoord::batch
CUTLASS_HOST_DEVICE Index & batch()
Returns reference to the GEMM batch coordinate.
Definition: include/cutlass/gemm/gemm.h:326
cutlass::gemm::GemmCoord::n
CUTLASS_HOST_DEVICE Index const & n() const
Returns the GEMM N coordinate.
Definition: include/cutlass/gemm/gemm.h:137
cutlass::gemm::GemmCoord::nm
CUTLASS_HOST_DEVICE Coord< 2 > nm() const
Obtains a Coord<2> from GemmCoord.
Definition: include/cutlass/gemm/gemm.h:165
cutlass::gemm::BatchedGemmCoord::operator+
CUTLASS_HOST_DEVICE BatchedGemmCoord operator+(Base const &b) const
Element-wise addition.
Definition: include/cutlass/gemm/gemm.h:346
cutlass::gemm::GemmCoord::k
CUTLASS_HOST_DEVICE Index & k()
Returns reference to the GEMM K coordinate.
Definition: include/cutlass/gemm/gemm.h:149
cutlass::gemm::BatchedGemmCoord::n
CUTLASS_HOST_DEVICE Index & n()
Returns reference to the GEMM N coordinate.
Definition: include/cutlass/gemm/gemm.h:310
cutlass::gemm::GemmCoord::operator/
CUTLASS_HOST_DEVICE GemmCoord operator/(Base const &b) const
Element-wise division.
Definition: include/cutlass/gemm/gemm.h:223
cutlass::gemm::GemmCoord::GemmCoord
CUTLASS_HOST_DEVICE GemmCoord(Index m, Index n, Index k)
Helper to construct from a K, N, M, batch variables.
Definition: include/cutlass/gemm/gemm.h:125
cutlass::gemm::BatchedGemmCoord::Index
int Index
Integer-valued index.
Definition: include/cutlass/gemm/gemm.h:263
cutlass::operator+=
CUTLASS_HOST_DEVICE half_t & operator+=(half_t &lhs, half_t const &rhs)
Definition: half.h:654
cutlass::gemm::GemmCoord::nk
CUTLASS_HOST_DEVICE Coord< 2 > nk() const
Obtains a Coord<2> from GemmCoord.
Definition: include/cutlass/gemm/gemm.h:189
cutlass::gemm::BatchedGemmCoord::BatchedGemmCoord
CUTLASS_HOST_DEVICE BatchedGemmCoord(Base const &coord)
Constructs from Coord<4>
Definition: include/cutlass/gemm/gemm.h:290
cutlass::gemm::GemmCoord::k
CUTLASS_HOST_DEVICE Index const & k() const
Returns the GEMM K coordinate.
Definition: include/cutlass/gemm/gemm.h:145
cutlass::gemm::Operand::kC
B multiplicand.
cutlass::operator-=
CUTLASS_HOST_DEVICE half_t & operator-=(half_t &lhs, half_t const &rhs)
Definition: half.h:664
cutlass::gemm::GemmCoord::Base
Coord< 3, Index > Base
Base type is a Coord of rank=4.
Definition: include/cutlass/gemm/gemm.h:100
cutlass::gemm::Operand::kA
cutlass::gemm::GemmShape::toCoord
static CUTLASS_HOST_DEVICE Coord< 3 > toCoord()
Returns a Coord object.
Definition: include/cutlass/gemm/gemm.h:76
cutlass::gemm::GemmCoord::km
CUTLASS_HOST_DEVICE Coord< 2 > km() const
Obtains a Coord<2> from GemmCoord.
Definition: include/cutlass/gemm/gemm.h:183
cutlass::gemm::BatchedGemmCoord
Definition: include/cutlass/gemm/gemm.h:260
cutlass::gemm::GemmCoord::mnk
CUTLASS_HOST_DEVICE Coord< 3 > mnk() const
Obtains a Coord<3> from GemmCoord.
Definition: include/cutlass/gemm/gemm.h:153
cutlass::gemm::GemmCoord::operator-
CUTLASS_HOST_DEVICE GemmCoord operator-(Base const &b) const
Element-wise subtraction.
Definition: include/cutlass/gemm/gemm.h:211
cutlass::gemm::BatchedGemmCoord::batch
CUTLASS_HOST_DEVICE Index const & batch() const
Returns the GEMM batch coordinate.
Definition: include/cutlass/gemm/gemm.h:322
cutlass::gemm::BatchedGemmCoord::operator*=
CUTLASS_HOST_DEVICE BatchedGemmCoord & operator*=(Base const &b)
In-place multiplication.
Definition: include/cutlass/gemm/gemm.h:384
cutlass::gemm::BatchedGemmCoord::operator*
CUTLASS_HOST_DEVICE BatchedGemmCoord operator*(Base const &b) const
Element-wise multiplication.
Definition: include/cutlass/gemm/gemm.h:358
CUTLASS_HOST_DEVICE
#define CUTLASS_HOST_DEVICE
Definition: cutlass.h:89
cutlass::gemm::BatchedGemmCoord::k
CUTLASS_HOST_DEVICE Index const & k() const
Returns the GEMM K coordinate.
Definition: include/cutlass/gemm/gemm.h:314
cutlass::gemm::GemmShape
Shape of a matrix multiply-add operation.
Definition: include/cutlass/gemm/gemm.h:57
cutlass::gemm::GemmCoord::mk
CUTLASS_HOST_DEVICE Coord< 2 > mk() const
Obtains a Coord<2> from GemmCoord.
Definition: include/cutlass/gemm/gemm.h:177
cutlass::gemm::BatchedGemmCoord::BatchedGemmCoord
CUTLASS_HOST_DEVICE BatchedGemmCoord()
Default ctor.
Definition: include/cutlass/gemm/gemm.h:286
cutlass::gemm::BatchedGemmCoord::k
CUTLASS_HOST_DEVICE Index & k()
Returns reference to the GEMM K coordinate.
Definition: include/cutlass/gemm/gemm.h:318
cutlass::gemm::Operand::kD
Source accumulator.
cutlass::gemm::GemmCoord::knm
CUTLASS_HOST_DEVICE Coord< 3 > knm() const
Obtains a Coord<3> from GemmCoord.
Definition: include/cutlass/gemm/gemm.h:159
cutlass::gemm::GemmCoord::operator+=
CUTLASS_HOST_DEVICE GemmCoord & operator+=(Base const &b)
In-place addition.
Definition: include/cutlass/gemm/gemm.h:229
cutlass::operator*=
CUTLASS_HOST_DEVICE half_t & operator*=(half_t &lhs, half_t const &rhs)
Definition: half.h:674
cutlass::Coord< 3 >
cutlass::gemm::GemmCoord::GemmCoord
CUTLASS_HOST_DEVICE GemmCoord()
Default ctor.
Definition: include/cutlass/gemm/gemm.h:117
cutlass::gemm::BatchedGemmCoord::operator/=
CUTLASS_HOST_DEVICE BatchedGemmCoord & operator/=(Base const &b)
In-place division.
Definition: include/cutlass/gemm/gemm.h:391
cutlass::gemm::GemmCoord::operator-=
CUTLASS_HOST_DEVICE GemmCoord & operator-=(Base const &b)
In-place subtraction.
Definition: include/cutlass/gemm/gemm.h:236
cutlass::gemm::BatchedGemmCoord::operator+=
CUTLASS_HOST_DEVICE BatchedGemmCoord & operator+=(Base const &b)
In-place addition.
Definition: include/cutlass/gemm/gemm.h:370
cutlass::gemm::GemmCoord::kn
CUTLASS_HOST_DEVICE Coord< 2 > kn() const
Obtains a Coord<2> from GemmCoord.
Definition: include/cutlass/gemm/gemm.h:195
cutlass::gemm::BatchedGemmCoord::n
CUTLASS_HOST_DEVICE Index const & n() const
Returns the GEMM N coordinate.
Definition: include/cutlass/gemm/gemm.h:306
cutlass::gemm::BatchedGemmCoord::BatchedGemmCoord
CUTLASS_HOST_DEVICE BatchedGemmCoord(Index m, Index n, Index k, Index b)
Helper to construct from a K, N, M, and batch variables.
Definition: include/cutlass/gemm/gemm.h:294
cutlass::gemm::GemmCoord::m
CUTLASS_HOST_DEVICE Index const & m() const
Returns the GEMM M coordinate.
Definition: include/cutlass/gemm/gemm.h:129
cutlass::gemm::GemmCoord::operator*=
CUTLASS_HOST_DEVICE GemmCoord & operator*=(Base const &b)
In-place multiplication.
Definition: include/cutlass/gemm/gemm.h:243
cutlass::gemm::Operand::kB
A multiplicand.
cutlass::gemm::BatchedGemmCoord::operator-=
CUTLASS_HOST_DEVICE BatchedGemmCoord & operator-=(Base const &b)
In-place subtraction.
Definition: include/cutlass/gemm/gemm.h:377
cutlass.h
Basic include for CUTLASS.
cutlass::gemm::GemmCoord::n
CUTLASS_HOST_DEVICE Index & n()
Returns reference to the GEMM N coordinate.
Definition: include/cutlass/gemm/gemm.h:141
cutlass::gemm::GemmCoord::operator*
CUTLASS_HOST_DEVICE GemmCoord operator*(Base const &b) const
Element-wise multiplication.
Definition: include/cutlass/gemm/gemm.h:217
cutlass::gemm::BatchedGemmCoord::m
CUTLASS_HOST_DEVICE Index const & m() const
Returns the GEMM M coordinate.
Definition: include/cutlass/gemm/gemm.h:298