Cutlass CUDA Templates for Linear Algebra Subroutines and Solvers

gemm_stream_pair.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/convert.h"
31 #include "cutlass/matrix_traits.h"
32 #include "cutlass/reshape_tile.h"
33 #include "cutlass/tile_allocation.h"
34 #include "cutlass/tile_iterator.h"
35
36 #include "cutlass/gemm/clear_accumulators.h"
37 #include "cutlass/gemm/gemm_config.h"
38 #include "cutlass/gemm/gemm_global_stream.h"
39 #include "cutlass/gemm/gemm_operand.h"
40 #include "cutlass/gemm/gemm_shared_stream.h"
41 #include "cutlass/gemm/threadblock_swizzle.h"
42
43 namespace cutlass {
44 namespace gemm {
45
47
49 template <typename StreamA_, typename StreamB_, bool kResidueInProlog_>
50 struct GlobalLoadStreamPair {
51 //
52 // Type definitions
53 //
54
56 typedef StreamA_ StreamA;
57
59 typedef StreamB_ StreamB;
60
62 struct Params {
64 typename StreamA::Params stream_a;
65
67 typename StreamB::Params stream_b;
68
70 CUTLASS_HOST_DEVICE
71 Params() {}
72
74 CUTLASS_HOST_DEVICE
75 Params(typename StreamA::Params const &_params_A, typename StreamB::Params const &_params_B)
76 : stream_a(_params_A), stream_b(_params_B) {}
77 };
78
80 typedef typename StreamA::Index Index;
81
83 typedef ZipTileAllocation<typename StreamA::ThreadblockTileStorage,
84 typename StreamB::ThreadblockTileStorage>
85 ThreadblockTileStorage;
86
88 typedef typename ThreadblockTileStorage::TensorRef ThreadblockTileRef;
89
91 struct SharedStorage {
92 typename StreamA::SharedStorage stream_a;
93 typename StreamB::SharedStorage stream_b;
94 };
95
96 //
97 // Data members
98 //
99
101 StreamA stream_a;
102
104 StreamB stream_b;
105
106 //
107 // Methods
108 //
109
111 CUTLASS_DEVICE GlobalLoadStreamPair(Params const &params,
112 SharedStorage &shared_storage,
113 ThreadblockTileRef const &threadblock_tile_ref,
114 Coord<3> const bounds,
115 Coord<3> const &block_offset = make_Coord(0, 0, 0))
116 : stream_a(params.stream_a,
117 shared_storage.stream_a,
118 threadblock_tile_ref.first,
119 bounds,
120 block_offset),
121 stream_b(params.stream_b,
122 shared_storage.stream_b,
123 threadblock_tile_ref.second,
124 bounds,
125 block_offset) {}
126
127 CUTLASS_DEVICE
128 GlobalLoadStreamPair & operator+=(Coord<3> const offset) {
129 stream_a += offset;
130 stream_b += offset;
131 return *this;
132 }
133
134 CUTLASS_DEVICE
135 GlobalLoadStreamPair & add_batch_offset(int batch_id) {
136 stream_a.add_batch_offset(batch_id);
137 stream_b.add_batch_offset(batch_id);
138 return *this;
139 }
140
142 CUTLASS_DEVICE void copy() {
143 stream_a.copy();
144 stream_b.copy();
145 }
146
148 CUTLASS_DEVICE void commit() {
149 stream_a.commit();
150 stream_b.commit();
151 }
152
154 CUTLASS_DEVICE void residue(Index k, bool skip_clear = false) {
155 stream_a.residue(k, skip_clear);
156 stream_b.residue(k, skip_clear);
157 }
158
160 CUTLASS_DEVICE void move_to_residue(Index k, Index kTileK) {
161 if (kResidueInProlog_) {
162 stream_a.move_to_residue(k, kTileK);
163 stream_b.move_to_residue(k, kTileK);
164 } else if (k < kTileK) {
165 residue(k, true);
166 }
167 }
168
170 CUTLASS_DEVICE void rollback(bool kRollback) {
171 if (kResidueInProlog_ && kRollback) {
172 stream_a.rollback();
173 stream_b.rollback();
174 }
175 }
176 };
177
179 template <typename StreamA_, typename StreamB_>
180 struct SharedStreamPair {
181 //
182 // Type definitions
183 //
184
186 typedef StreamA_ StreamA;
187
189 typedef StreamB_ StreamB;
190
192 struct Params {
194 typename StreamA::Params stream_a;
195
197 typename StreamB::Params stream_b;
198 };
199
201 typedef ZipTensorRef<typename StreamA::TensorRef,
202 typename StreamB::TensorRef >
203 ThreadblockTileRef;
204
205 //
206 // Data members
207 //
208
210 StreamA stream_a;
211
213 StreamB stream_b;
214
215 //
216 // Methods
217 //
218
220 CUTLASS_DEVICE SharedStreamPair(Params const &params, ThreadblockTileRef const &threadblock_tile_ref)
221 : stream_a(params.stream_a, threadblock_tile_ref.first),
222 stream_b(params.stream_b, threadblock_tile_ref.second) {}
223
225 CUTLASS_DEVICE void copy(int step) {
226 stream_a.copy(step);
227 stream_b.copy(step);
228 }
229
231 CUTLASS_DEVICE void commit(int step) {
232 stream_a.commit(step);
233 stream_b.commit(step);
234 }
235
237 CUTLASS_DEVICE
238 typename StreamA::TransformedFragment const &fragment_a(int step) const {
239 return stream_a.fragment(step);
240 }
241
243 CUTLASS_DEVICE
244 typename StreamB::TransformedFragment const &fragment_b(int step) const {
245 return stream_b.fragment(step);
246 }
247
249 CUTLASS_DEVICE void inc_stage() {
250 stream_a.inc_stage();
251 stream_b.inc_stage();
252 }
253 };
254
256
257 } // namespace gemm
258 } // namespace cutlass
cutlass::gemm::GlobalLoadStreamPair::Params::Params
CUTLASS_HOST_DEVICE Params(typename StreamA::Params const &_params_A, typename StreamB::Params const &_params_B)
Constructs a global load stream pair Params object.
Definition: gemm_stream_pair.h:75
cutlass::gemm::GlobalLoadStreamPair::operator+=
CUTLASS_DEVICE GlobalLoadStreamPair & operator+=(Coord< 3 > const offset)
Definition: gemm_stream_pair.h:128
cutlass::gemm::SharedStreamPair::StreamA
StreamA_ StreamA
Stream for A multiplicand.
Definition: gemm_stream_pair.h:186
cutlass::gemm::GlobalLoadStreamPair::move_to_residue
CUTLASS_DEVICE void move_to_residue(Index k, Index kTileK)
Move to residue.
Definition: gemm_stream_pair.h:160
cutlass
Definition: convert.h:33
cutlass::gemm::GlobalLoadStreamPair::SharedStorage::stream_b
StreamB::SharedStorage stream_b
Definition: gemm_stream_pair.h:93
cutlass::gemm::GlobalLoadStreamPair::add_batch_offset
CUTLASS_DEVICE GlobalLoadStreamPair & add_batch_offset(int batch_id)
Definition: gemm_stream_pair.h:135
cutlass::ZipTensorRef
Definition: zip_tensor_ref.h:38
tile_iterator.h
Defines the Tile Traits concept and iterators for loading and storing to tiles efficiently.
cutlass::gemm::SharedStreamPair::Params::stream_a
StreamA::Params stream_a
Definition: gemm_stream_pair.h:194
cutlass::gemm::GlobalLoadStreamPair::GlobalLoadStreamPair
CUTLASS_DEVICE GlobalLoadStreamPair(Params const &params, SharedStorage &shared_storage, ThreadblockTileRef const &threadblock_tile_ref, Coord< 3 > const bounds, Coord< 3 > const &block_offset=make_Coord(0, 0, 0))
Ctor.
Definition: gemm_stream_pair.h:111
threadblock_swizzle.h
Defies functors for mapping blockIdx to partitions of the GEMM computation.
cutlass::gemm::GlobalLoadStreamPair::SharedStorage
Defines a structure containing shared storage for each pair.
Definition: gemm_stream_pair.h:91
cutlass::make_Coord
CUTLASS_HOST_DEVICE Coord< 1 > make_Coord(int _0)
Helper to make a 2-element coordinate.
Definition: coord.h:368
cutlass::gemm::GlobalLoadStreamPair::ThreadblockTileStorage
ZipTileAllocation< typename StreamA::ThreadblockTileStorage, typename StreamB::ThreadblockTileStorage > ThreadblockTileStorage
Shared memory allocation for threadblock-scoped GEMM tile.
Definition: gemm_stream_pair.h:85
cutlass::gemm::GlobalLoadStreamPair::residue
CUTLASS_DEVICE void residue(Index k, bool skip_clear=false)
Execute the residue code.
Definition: gemm_stream_pair.h:154
cutlass::gemm::GlobalLoadStreamPair::ThreadblockTileRef
ThreadblockTileStorage::TensorRef ThreadblockTileRef
ZipTensorRef to threadblock tiles.
Definition: gemm_stream_pair.h:88
cutlass::gemm::SharedStreamPair::SharedStreamPair
CUTLASS_DEVICE SharedStreamPair(Params const &params, ThreadblockTileRef const &threadblock_tile_ref)
Construct with the composable structure.
Definition: gemm_stream_pair.h:220
cutlass::gemm::GlobalLoadStreamPair::StreamB
StreamB_ StreamB
Stream for B multiplicand.
Definition: gemm_stream_pair.h:59
cutlass::gemm::SharedStreamPair
Collect the global load streams for multiplicands.
Definition: gemm_stream_pair.h:180
cutlass::gemm::SharedStreamPair::fragment_b
CUTLASS_DEVICE StreamB::TransformedFragment const & fragment_b(int step) const
The fragment B.
Definition: gemm_stream_pair.h:244
cutlass::gemm::GlobalLoadStreamPair::Params::stream_b
StreamB::Params stream_b
Parameters object for StreamB.
Definition: gemm_stream_pair.h:67
cutlass::gemm::GlobalLoadStreamPair::rollback
CUTLASS_DEVICE void rollback(bool kRollback)
Rollback to beginning of first tile.
Definition: gemm_stream_pair.h:170
cutlass::gemm::SharedStreamPair::stream_a
StreamA stream_a
The stream for A.
Definition: gemm_stream_pair.h:210
cutlass::gemm::SharedStreamPair::fragment_a
CUTLASS_DEVICE StreamA::TransformedFragment const & fragment_a(int step) const
The fragment A.
Definition: gemm_stream_pair.h:238
cutlass::gemm::GlobalLoadStreamPair
Collect the global load streams for multiplicands.
Definition: gemm_stream_pair.h:50
cutlass::gemm::SharedStreamPair::copy
CUTLASS_DEVICE void copy(int step)
Trigger the copies from shared memory to registers.
Definition: gemm_stream_pair.h:225
tile_allocation.h
Defines a fragment based on a Shape<> template.
cutlass::gemm::GlobalLoadStreamPair::Params
Parameters object.
Definition: gemm_stream_pair.h:62
reshape_tile.h
Defines a type for restructuring a tile.
gemm_operand.h
Defines constant expressions for mapping GEMM problem size and strides onto pitch-linear memory...
clear_accumulators.h
Defines abstractions for efficiently clearing accumulator tiles.
cutlass::gemm::GlobalLoadStreamPair::StreamA
StreamA_ StreamA
Stream for A multiplicand.
Definition: gemm_stream_pair.h:56
cutlass::gemm::GlobalLoadStreamPair::stream_a
StreamA stream_a
Stream for A multiplicand.
Definition: gemm_stream_pair.h:101
cutlass::gemm::SharedStreamPair::ThreadblockTileRef
ZipTensorRef< typename StreamA::TensorRef, typename StreamB::TensorRef > ThreadblockTileRef
Shared memory allocation for threadblock-scoped GEMM tile.
Definition: gemm_stream_pair.h:203
cutlass::gemm::GlobalLoadStreamPair::Index
StreamA::Index Index
Assumes the A stream defines the index type.
Definition: gemm_stream_pair.h:80
cutlass::ZipTileAllocation
Manages a pair of tile allocations as if they are one allocation.
Definition: tile_allocation.h:125
CUTLASS_HOST_DEVICE
#define CUTLASS_HOST_DEVICE
Definition: cutlass.h:46
gemm_config.h
Defines properties of GEMM computation that impose some constraints on caller.
cutlass::gemm::GlobalLoadStreamPair::Params::Params
CUTLASS_HOST_DEVICE Params()
Default constructor.
Definition: gemm_stream_pair.h:71
cutlass::gemm::SharedStreamPair::StreamB
StreamB_ StreamB
Stream for B multiplicand.
Definition: gemm_stream_pair.h:189
cutlass::gemm::GlobalLoadStreamPair::SharedStorage::stream_a
StreamA::SharedStorage stream_a
Definition: gemm_stream_pair.h:92
cutlass::gemm::SharedStreamPair::commit
CUTLASS_DEVICE void commit(int step)
Commit the data.
Definition: gemm_stream_pair.h:231
cutlass::Coord< 3 >
cutlass::gemm::GlobalLoadStreamPair::commit
CUTLASS_DEVICE void commit()
Commit the data.
Definition: gemm_stream_pair.h:148
gemm_global_stream.h
Implements efficient loading of the thread block-level tile from global memory and storing to shared ...
cutlass::gemm::SharedStreamPair::stream_b
StreamB stream_b
The stream for B.
Definition: gemm_stream_pair.h:213
cutlass::gemm::SharedStreamPair::inc_stage
CUTLASS_DEVICE void inc_stage()
Increment the stage.
Definition: gemm_stream_pair.h:249
cutlass::gemm::SharedStreamPair::Params
Parameters object passed to load iterators.
Definition: gemm_stream_pair.h:192
cutlass::gemm::SharedStreamPair::Params::stream_b
StreamB::Params stream_b
Definition: gemm_stream_pair.h:197
matrix_traits.h
Defines properties of matrices used to denote layout and operands to GEMM kernels.
cutlass::gemm::GlobalLoadStreamPair::copy
CUTLASS_DEVICE void copy()
Trigger the copies from shared memory to registers.
Definition: gemm_stream_pair.h:142
cutlass::gemm::GlobalLoadStreamPair::Params::stream_a
StreamA::Params stream_a
Parameters object for StreamA.
Definition: gemm_stream_pair.h:64
gemm_shared_stream.h
Defines abstractions for managing loading and storing fragments to shared memory in the efficient GEM...
cutlass::gemm::GlobalLoadStreamPair::stream_b
StreamB stream_b
Stream for B multiplicand.
Definition: gemm_stream_pair.h:104
convert.h
Defines conversion operations among Fragments of different base type.