Cutlass CUDA Templates for Linear Algebra Subroutines and Solvers

gemm_fragment_stream.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/fragment_stream.h>
31 #include <cutlass/matrix_traits.h>
32
33 #include <cutlass/gemm/gemm_operand.h>
34 #include <cutlass/gemm/gemm_traits.h>
35
36 namespace cutlass {
37 namespace gemm {
38
40
42 template <GemmOperand::Kind Usage, typename Scalar_, MatrixLayout::Kind Layout,
43 typename ThreadBlockTile_, int Threads, int ScalarsPerInst, typename Index_ = int,
44 typename DestinationSkew_ = Shape<0, 0, 0, 0> >
45 struct GemmFragmentStreamTraits {
47 static GemmOperand::Kind const kUsage = Usage;
48
50 typedef Scalar_ Scalar;
51
53 static MatrixLayout::Kind const kLayout = Layout;
54
56 typedef ThreadBlockTile_ ThreadBlockTile;
57
59 static int const kThreads = Threads;
60
62 static int const kAccessSize = ScalarsPerInst;
63
65 typedef Index_ Index;
66
68 typedef typename ShapeDiv<DestinationSkew_, Shape<ScalarsPerInst, ScalarsPerInst, ScalarsPerInst,
69 1> >::Shape DestinationSkew;
70
72 typedef GemmMultiplicandTraits<ThreadBlockTile, kUsage, kLayout> MultiplicandTraits;
73
75 typedef typename MultiplicandTraits::Shape ScalarTile;
76
78 typedef typename ReshapeTile<ScalarTile, kAccessSize>::Tile VectorizedTile;
79
81 typedef TileTraitsDefault<VectorizedTile, kThreads> TileTraits;
82
84 typedef FragmentStream<
85 TileTraits,
86 TileLoadIterator<TileTraits, Scalar,
87 MultiplicandTraits::kKstrided ? IteratorAdvance::kH : IteratorAdvance::kW,
88 MemorySpace::kGlobal, Index>,
89 TileStoreIterator<TileTraits, Scalar,
90 MultiplicandTraits::kKstrided ? IteratorAdvance::kH : IteratorAdvance::kW,
91 MemorySpace::kShared, Index, Scalar, IteratorFragment::kScalar,
92 DestinationSkew> >
93 FragmentStream;
94 };
95
97 template <typename Traits_>
98 struct GemmFragmentStream
99 : public FragmentStream<
100 typename Traits_::TileTraits,
101 TileLoadIterator<typename Traits_::TileTraits, typename Traits_::Scalar,
102 Traits_::MultiplicandTraits::kKstrided ? IteratorAdvance::kH
103 : IteratorAdvance::kW,
104 MemorySpace::kGlobal, typename Traits_::Index>,
105 TileStoreIterator<typename Traits_::TileTraits, typename Traits_::Scalar,
106 Traits_::MultiplicandTraits::kKstrided ? IteratorAdvance::kH
107 : IteratorAdvance::kW,
108 MemorySpace::kShared, typename Traits_::Index, typename Traits_::Scalar,
109 IteratorFragment::kScalar, typename Traits_::DestinationSkew> > {
111 typedef Traits_ Traits;
112
114 typedef typename Traits::FragmentStream Base;
115
116 //
117 // FragmentStream concept
118 //
119
121 typedef typename Traits::Scalar Scalar;
122
124 typedef typename Base::LoadIterator LoadIterator;
125
127 typedef typename Base::StoreIterator StoreIterator;
128
130 typedef typename Base::Convert Convert;
131
133 typedef typename Base::Fragment Fragment;
134
136 typedef typename Base::StoreFragment StoreFragment;
137
139 typedef typename Base::Storage Storage;
140
141 // Parameters type
142 // typedef typename Base::Params BaseParams;
143
145 typedef typename Traits::Index Index;
146
147 //
148 // Nested class definitions
149 //
150
152 typedef typename Traits::Scalar const *Pointer;
153
155 struct Params : public Base::Params {
156 //
157 // Methods
158 //
159
161 template <typename GemmDesc_>
162 CUTLASS_HOST_DEVICE int initialize(GemmDesc_ const &desc,
163 typename Traits::Scalar const *pointer, Index ldm) {
164 return this->load_params.initialize(pointer, ldm * Traits::MultiplicandTraits::Shape::kH, ldm,
165 Traits::kAccessSize);
166 }
167 };
168
169 //
170 // Static member functions
171 //
172
174 static CUTLASS_DEVICE void shared_store_fence() { Base::shared_store_fence(); }
175
176 //
177 // Methods
178 //
179
180 CUTLASS_DEVICE
181 GemmFragmentStream() {}
182
184 CUTLASS_DEVICE
185 GemmFragmentStream(Params const &params, Coord<3> const &bounds,
186 Coord<3> const &block_offset = make_Coord(0, 0, 0))
187 : Base(params, ProjectOperand<Traits::kUsage, Traits::MultiplicandTraits::kKstrided>::project(
188 bounds) +
189 make_Coord(1, 0, 0),
190 ProjectOperand<Traits::kUsage, Traits::MultiplicandTraits::kKstrided>::project(
191 block_offset)) {}
192
194 CUTLASS_DEVICE
195 void load() { Base::load(); }
196
198 CUTLASS_DEVICE
199 void commit() { Base::commit(); }
200
202 CUTLASS_DEVICE
203 void residue(Coord<3> const &bounds, Coord<3> const &block_offset) {
204 this->initialize_predicates(bounds, block_offset);
205
206 this->fetch.clear();
207 }
208
210 CUTLASS_DEVICE
211 void initialize_predicates(Coord<3> const &bounds, Coord<3> const &block_offset) {
212 Base::initialize_predicates(
213 ProjectOperand<Traits::kUsage, Traits::MultiplicandTraits::kKstrided>::project(bounds) +
214 make_Coord(1, 0, 0),
215 ProjectOperand<Traits::kUsage, Traits::MultiplicandTraits::kKstrided>::project(
216 block_offset));
217 }
218 };
219
221 }
222 }
cutlass::gemm::GemmMultiplicandTraits::Shape
nv_std::conditional< kKstrided, Shape< 1, ThreadBlockTile::kD, GetExtent< Usage, ThreadBlockTile >::kExtent >, Shape< 1, GetExtent< Usage, ThreadBlockTile >::kExtent, ThreadBlockTile::kD > >::type Shape
Map the ThreadBlockShape onto (kH, kW) dimensions for A and B operand.
Definition: gemm_operand.h:86
cutlass::gemm::GemmMultiplicandTraits::kKstrided
static bool const kKstrided
Definition: gemm_operand.h:81
cutlass::gemm::GemmFragmentStreamTraits::Scalar
Scalar_ Scalar
Scalar data type.
Definition: gemm_fragment_stream.h:50
cutlass::gemm::GemmFragmentStreamTraits::MultiplicandTraits
GemmMultiplicandTraits< ThreadBlockTile, kUsage, kLayout > MultiplicandTraits
Traits of multiplicand.
Definition: gemm_fragment_stream.h:72
cutlass::gemm::GemmFragmentStreamTraits::kAccessSize
static int const kAccessSize
Scalars per instruction.
Definition: gemm_fragment_stream.h:62
cutlass::MemorySpace::kShared
Definition: load_store.h:42
cutlass
Definition: convert.h:34
cutlass::gemm::GemmFragmentStream::StoreIterator
Base::StoreIterator StoreIterator
Defines the store iterator.
Definition: gemm_fragment_stream.h:127
cutlass::gemm::GemmFragmentStream::Params::initialize
CUTLASS_HOST_DEVICE int initialize(GemmDesc_ const &desc, typename Traits::Scalar const *pointer, Index ldm)
Initializes parameters.
Definition: gemm_fragment_stream.h:162
gemm_traits.h
Defines structural properties of complete GEMM computation.
cutlass::gemm::GemmFragmentStream::Base
Traits::FragmentStream Base
Base class.
Definition: gemm_fragment_stream.h:114
fragment_stream.h
An abstraction for implementing a stream loading a tile and storing a tile using a pair of tile itera...
cutlass::MemorySpace::kGlobal
Definition: load_store.h:43
cutlass::make_Coord
CUTLASS_HOST_DEVICE Coord< 1 > make_Coord(int _0)
Helper to make a 2-element coordinate.
Definition: coord.h:241
cutlass::gemm::GemmFragmentStreamTraits::VectorizedTile
ReshapeTile< ScalarTile, kAccessSize >::Tile VectorizedTile
Reshape for vectorized access.
Definition: gemm_fragment_stream.h:78
cutlass::gemm::GemmFragmentStream::Index
Traits::Index Index
Index type.
Definition: gemm_fragment_stream.h:145
cutlass::gemm::GemmFragmentStreamTraits::FragmentStream
FragmentStream< TileTraits, TileLoadIterator< TileTraits, Scalar, MultiplicandTraits::kKstrided ? IteratorAdvance::kH :IteratorAdvance::kW, MemorySpace::kGlobal, Index >, TileStoreIterator< TileTraits, Scalar, MultiplicandTraits::kKstrided ? IteratorAdvance::kH :IteratorAdvance::kW, MemorySpace::kShared, Index, Scalar, IteratorFragment::kScalar, DestinationSkew > > FragmentStream
Define the tile stream.
Definition: gemm_fragment_stream.h:93
cutlass::gemm::GemmFragmentStream::Traits
Traits_ Traits
Traits.
Definition: gemm_fragment_stream.h:111
cutlass::FragmentStream< Traits_::TileTraits, TileLoadIterator< Traits_::TileTraits, Traits_::Scalar, Traits_::MultiplicandTraits::kKstrided ? IteratorAdvance::kH :IteratorAdvance::kW, MemorySpace::kGlobal, Traits_::Index >, TileStoreIterator< Traits_::TileTraits, Traits_::Scalar, Traits_::MultiplicandTraits::kKstrided ? IteratorAdvance::kH :IteratorAdvance::kW, MemorySpace::kShared, Traits_::Index, Traits_::Scalar, IteratorFragment::kScalar, Traits_::DestinationSkew > >::fetch
Fragment fetch
Fragment fetched by load iterator.
Definition: fragment_stream.h:135
cutlass::IteratorAdvance::kH
Definition: tile_iterator.h:97
cutlass::gemm::GemmFragmentStream::GemmFragmentStream
CUTLASS_DEVICE GemmFragmentStream()
Definition: gemm_fragment_stream.h:181
cutlass::gemm::GemmFragmentStreamTraits::TileTraits
TileTraitsDefault< VectorizedTile, kThreads > TileTraits
Define structure of stripmined tile.
Definition: gemm_fragment_stream.h:81
cutlass::gemm::GemmFragmentStreamTraits::ScalarTile
MultiplicandTraits::Shape ScalarTile
Scalar tile shape.
Definition: gemm_fragment_stream.h:75
cutlass::gemm::GemmFragmentStream::shared_store_fence
static CUTLASS_DEVICE void shared_store_fence()
The memory fence for shared stores.
Definition: gemm_fragment_stream.h:174
cutlass::gemm::GemmFragmentStreamTraits
Defines a FragmentStream by mapping GEMM dimensions onto contiguous and strided dimensions.
Definition: gemm_fragment_stream.h:45
cutlass::gemm::GemmFragmentStream::GemmFragmentStream
CUTLASS_DEVICE GemmFragmentStream(Params const &params, Coord< 3 > const &bounds, Coord< 3 > const &block_offset=make_Coord(0, 0, 0))
Constructor - bounds and block offset are aligned to GEMM coordinates (K, N, M)
Definition: gemm_fragment_stream.h:185
cutlass::gemm::GemmFragmentStream::Fragment
Base::Fragment Fragment
Loaded fragment type.
Definition: gemm_fragment_stream.h:133
cutlass::gemm::GemmFragmentStream
GEMM Fragment Stream.
Definition: gemm_fragment_stream.h:98
cutlass::gemm::GemmFragmentStream::Pointer
Traits::Scalar const * Pointer
The pointer.
Definition: gemm_fragment_stream.h:152
cutlass::TileLoadIterator
An iterator implementing Tile Load Iterator Concept for loading a tile from memory.
Definition: tile_iterator.h:308
cutlass::gemm::GemmFragmentStream::commit
CUTLASS_DEVICE void commit()
Commits the fragment.
Definition: gemm_fragment_stream.h:199
cutlass::gemm::GemmFragmentStream::Storage
Base::Storage Storage
Destination storage.
Definition: gemm_fragment_stream.h:139
gemm_operand.h
Defines constant expressions for mapping GEMM problem size and strides onto pitch-linear memory...
cutlass::FragmentStream
Manages a pair of iterators to stream data from global memory to shared.
Definition: fragment_stream.h:50
cutlass::gemm::GemmMultiplicandTraits
Definition: gemm_operand.h:66
cutlass::gemm::GemmFragmentStreamTraits::kLayout
static MatrixLayout::Kind const kLayout
Layout of the operand.
Definition: gemm_fragment_stream.h:53
cutlass::gemm::GemmFragmentStream::Scalar
Traits::Scalar Scalar
Scalar type.
Definition: gemm_fragment_stream.h:121
cutlass::gemm::GemmFragmentStream::Params
Parameters object.
Definition: gemm_fragment_stream.h:155
CUTLASS_HOST_DEVICE
#define CUTLASS_HOST_DEVICE
Definition: cutlass.h:37
cutlass::gemm::GemmFragmentStreamTraits::Index
Index_ Index
Index type.
Definition: gemm_fragment_stream.h:65
cutlass::ShapeDiv
Definition: shape.h:124
cutlass::Shape
A Shape implementing Layout Concept describing the dimensions of a cube.
Definition: shape.h:63
cutlass::gemm::GemmFragmentStream::Convert
Base::Convert Convert
Converts between tiles.
Definition: gemm_fragment_stream.h:130
cutlass::gemm::ProjectOperand
Definition: gemm_operand.h:94
cutlass::Coord
Statically-sized array specifying Coords within a tensor.
Definition: coord.h:48
cutlass::gemm::GemmFragmentStreamTraits::kThreads
static int const kThreads
Number of threads.
Definition: gemm_fragment_stream.h:59
cutlass::MatrixLayout::Kind
Kind
Definition: matrix_traits.h:36
cutlass::gemm::GemmFragmentStream::StoreFragment
Base::StoreFragment StoreFragment
Stored fragment type.
Definition: gemm_fragment_stream.h:136
cutlass::gemm::GemmFragmentStream::LoadIterator
Base::LoadIterator LoadIterator
Defines the load iterator.
Definition: gemm_fragment_stream.h:124
cutlass::gemm::GemmFragmentStreamTraits::kUsage
static GemmOperand::Kind const kUsage
Indicates identity of multiplicand.
Definition: gemm_fragment_stream.h:47
cutlass::ReshapeTile::Tile
Tile_ Tile
Definition: tile.h:43
cutlass::IteratorAdvance::kW
Definition: tile_iterator.h:97
cutlass::gemm::GemmFragmentStreamTraits::DestinationSkew
ShapeDiv< DestinationSkew_, Shape< ScalarsPerInst, ScalarsPerInst, ScalarsPerInst, 1 > >::Shape DestinationSkew
Skew added to shared memory tile.
Definition: gemm_fragment_stream.h:69
cutlass::GemmOperand::Kind
Kind
Definition: matrix_traits.h:43
cutlass::gemm::GemmFragmentStream::residue
CUTLASS_DEVICE void residue(Coord< 3 > const &bounds, Coord< 3 > const &block_offset)
TODO - Recomputes predicates and clears fetch registers.
Definition: gemm_fragment_stream.h:203
cutlass::gemm::GemmFragmentStreamTraits::ThreadBlockTile
ThreadBlockTile_ ThreadBlockTile
Shape of the thread block tile (K, N, M)
Definition: gemm_fragment_stream.h:56
matrix_traits.h
Defines properties of matrices used to denote layout and operands to GEMM kernels.
cutlass::gemm::GemmFragmentStream::initialize_predicates
CUTLASS_DEVICE void initialize_predicates(Coord< 3 > const &bounds, Coord< 3 > const &block_offset)
Recomputes predicates aligned to GEMM coordinates (K, N, M)
Definition: gemm_fragment_stream.h:211
cutlass::IteratorFragment::kScalar
Definition: tile_iterator.h:102
cutlass::gemm::GemmFragmentStream::load
CUTLASS_DEVICE void load()
Loads the fragment.
Definition: gemm_fragment_stream.h:195
cutlass::TileStoreIterator
An iterator implementing Tile Store Iterator Concept for storing a tile to memory.
Definition: tile_iterator.h:556