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"
34 #include "cutlass/tile_iterator.h"
35 
42 
43 namespace cutlass {
44 namespace gemm {
45 
47 
49 template <typename StreamA_, typename StreamB_, bool kResidueInProlog_>
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 
71  Params() {}
72 
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>
86 
89 
91  struct SharedStorage {
92  typename StreamA::SharedStorage stream_a;
93  typename StreamB::SharedStorage stream_b;
94  };
95 
96  //
97  // Data members
98  //
99 
102 
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
129  stream_a += offset;
130  stream_b += offset;
131  return *this;
132  }
133 
134  CUTLASS_DEVICE
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_>
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 >
204 
205  //
206  // Data members
207  //
208 
211 
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_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_DEVICE GlobalLoadStreamPair & operator+=(Coord< 3 > const offset)
Definition: gemm_stream_pair.h:128
StreamA_ StreamA
Stream for A multiplicand.
Definition: gemm_stream_pair.h:186
CUTLASS_DEVICE void move_to_residue(Index k, Index kTileK)
Move to residue.
Definition: gemm_stream_pair.h:160
Definition: convert.h:33
StreamB::SharedStorage stream_b
Definition: gemm_stream_pair.h:93
CUTLASS_DEVICE GlobalLoadStreamPair & add_batch_offset(int batch_id)
Definition: gemm_stream_pair.h:135
Definition: zip_tensor_ref.h:38
Defines the Tile Traits concept and iterators for loading and storing to tiles efficiently.
StreamA::Params stream_a
Definition: gemm_stream_pair.h:194
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
Defies functors for mapping blockIdx to partitions of the GEMM computation.
Defines a structure containing shared storage for each pair.
Definition: gemm_stream_pair.h:91
CUTLASS_HOST_DEVICE Coord< 1 > make_Coord(int _0)
Helper to make a 2-element coordinate.
Definition: coord.h:368
ZipTileAllocation< typename StreamA::ThreadblockTileStorage, typename StreamB::ThreadblockTileStorage > ThreadblockTileStorage
Shared memory allocation for threadblock-scoped GEMM tile.
Definition: gemm_stream_pair.h:85
CUTLASS_DEVICE void residue(Index k, bool skip_clear=false)
Execute the residue code.
Definition: gemm_stream_pair.h:154
ThreadblockTileStorage::TensorRef ThreadblockTileRef
ZipTensorRef to threadblock tiles.
Definition: gemm_stream_pair.h:88
CUTLASS_DEVICE SharedStreamPair(Params const &params, ThreadblockTileRef const &threadblock_tile_ref)
Construct with the composable structure.
Definition: gemm_stream_pair.h:220
StreamB_ StreamB
Stream for B multiplicand.
Definition: gemm_stream_pair.h:59
Collect the global load streams for multiplicands.
Definition: gemm_stream_pair.h:180
CUTLASS_DEVICE StreamB::TransformedFragment const & fragment_b(int step) const
The fragment B.
Definition: gemm_stream_pair.h:244
StreamB::Params stream_b
Parameters object for StreamB.
Definition: gemm_stream_pair.h:67
CUTLASS_DEVICE void rollback(bool kRollback)
Rollback to beginning of first tile.
Definition: gemm_stream_pair.h:170
StreamA stream_a
The stream for A.
Definition: gemm_stream_pair.h:210
CUTLASS_DEVICE StreamA::TransformedFragment const & fragment_a(int step) const
The fragment A.
Definition: gemm_stream_pair.h:238
Collect the global load streams for multiplicands.
Definition: gemm_stream_pair.h:50
CUTLASS_DEVICE void copy(int step)
Trigger the copies from shared memory to registers.
Definition: gemm_stream_pair.h:225
Defines a fragment based on a Shape<> template.
Parameters object.
Definition: gemm_stream_pair.h:62
Defines a type for restructuring a tile.
Defines constant expressions for mapping GEMM problem size and strides onto pitch-linear memory...
Defines abstractions for efficiently clearing accumulator tiles.
StreamA_ StreamA
Stream for A multiplicand.
Definition: gemm_stream_pair.h:56
StreamA stream_a
Stream for A multiplicand.
Definition: gemm_stream_pair.h:101
ZipTensorRef< typename StreamA::TensorRef, typename StreamB::TensorRef > ThreadblockTileRef
Shared memory allocation for threadblock-scoped GEMM tile.
Definition: gemm_stream_pair.h:203
StreamA::Index Index
Assumes the A stream defines the index type.
Definition: gemm_stream_pair.h:80
Manages a pair of tile allocations as if they are one allocation.
Definition: tile_allocation.h:125
#define CUTLASS_HOST_DEVICE
Definition: cutlass.h:46
Defines properties of GEMM computation that impose some constraints on caller.
CUTLASS_HOST_DEVICE Params()
Default constructor.
Definition: gemm_stream_pair.h:71
StreamB_ StreamB
Stream for B multiplicand.
Definition: gemm_stream_pair.h:189
StreamA::SharedStorage stream_a
Definition: gemm_stream_pair.h:92
CUTLASS_DEVICE void commit(int step)
Commit the data.
Definition: gemm_stream_pair.h:231
CUTLASS_DEVICE void commit()
Commit the data.
Definition: gemm_stream_pair.h:148
Implements efficient loading of the thread block-level tile from global memory and storing to shared ...
StreamB stream_b
The stream for B.
Definition: gemm_stream_pair.h:213
CUTLASS_DEVICE void inc_stage()
Increment the stage.
Definition: gemm_stream_pair.h:249
Parameters object passed to load iterators.
Definition: gemm_stream_pair.h:192
StreamB::Params stream_b
Definition: gemm_stream_pair.h:197
Defines properties of matrices used to denote layout and operands to GEMM kernels.
CUTLASS_DEVICE void copy()
Trigger the copies from shared memory to registers.
Definition: gemm_stream_pair.h:142
StreamA::Params stream_a
Parameters object for StreamA.
Definition: gemm_stream_pair.h:64
Defines abstractions for managing loading and storing fragments to shared memory in the efficient GEM...
StreamB stream_b
Stream for B multiplicand.
Definition: gemm_stream_pair.h:104
Defines conversion operations among Fragments of different base type.