1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
|
/***************************************************************************************************
* Copyright (c) 2017-2018, NVIDIA CORPORATION. All rights reserved.
*
* Redistribution and use in source and binary forms, with or without modification, are permitted
* provided that the following conditions are met:
* * Redistributions of source code must retain the above copyright notice, this list of
* conditions and the following disclaimer.
* * Redistributions in binary form must reproduce the above copyright notice, this list of
* conditions and the following disclaimer in the documentation and/or other materials
* provided with the distribution.
* * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
* to endorse or promote products derived from this software without specific prior written
* permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
* IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
* FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
* BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
* OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
* STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
/*! \file
\brief Defines tile iterator traits for loading thread block-level tile from global memory.
*/
#pragma once
#include <cutlass/gemm/gemm_global_tile.h>
namespace cutlass {
namespace gemm {
////////////////////////////////////////////////////////////////////////////////////////////////////
template <typename Scalar_, typename Tile_, typename Threads_, int kAccessSize_>
struct WmmaGemmGlobalIteratorCdTraits : public GemmGlobalTileTraits<GemmOperand::kC,
MatrixLayout::kColumnMajor,
Scalar_,
Tile_,
Threads_,
kAccessSize_> {
/// The base class.
typedef GemmGlobalTileTraits<GemmOperand::kC,
MatrixLayout::kColumnMajor,
Scalar_,
Tile_,
Threads_,
kAccessSize_>
Base;
/// Override the strides in each dimension between different loads/stores.
typedef Shape<0, 0, Base::Delta::kW, Base::Delta::kC> Delta;
/// Computes the thread offset in (H, W) based on thread ID
struct ThreadOffset {
CUTLASS_HOST_DEVICE
Coord<4> operator()() const {
int thread_offset_h = threadIdx.x / Base::Threads::kW;
int thread_offset_w = threadIdx.x % Base::Threads::kW * Base::ThreadsDelta::kW;
return make_Coord(0, thread_offset_h, thread_offset_w, 0);
}
};
};
////////////////////////////////////////////////////////////////////////////////////////////////////
template <typename TileTraits_, typename Index_ = int>
struct WmmaGemmGlobalIteratorCd : public TileIteratorBase<TileTraits_,
typename TileTraits_::Scalar,
IteratorAdvance::kH,
MemorySpace::kGlobal,
Index_> {
/// This class.
typedef WmmaGemmGlobalIteratorCd<TileTraits_, Index_> This_;
/// The traits.
typedef TileTraits_ Traits;
/// The base class.
typedef TileIteratorBase<Traits,
typename TileTraits_::Scalar,
IteratorAdvance::kH,
MemorySpace::kGlobal,
Index_>
Base;
/// Override the strides in each dimension between different loads/stores.
typedef Shape<0, 0, Base::Delta::kW, Base::Delta::kC> ImmediateOffsetStrides;
/// The layout.
static MatrixLayout::Kind const kLayout = TileTraits_::kLayout;
/// The scalar.
typedef typename TileTraits_::Scalar Scalar;
/// The pointer.
typedef typename TileTraits_::Pointer Pointer;
/// The threads.
typedef typename TileTraits_::Threads Threads;
/// The index.
typedef Index_ Index;
/// The thread offset functor.
typedef typename TileTraits_::ThreadOffset ThreadOffset;
/// The params.
struct Params {
/// The pointer.
Pointer pointer;
/// The stride in the H dimension to setup the thread in the block.
Index stride_h;
/// The strides to increment the pointer.
Index inc_h, inc_advance;
/// The column offset to compute the predicate for the columns.
Index predicate_offset;
/// The strides to increment the predicate offset.
Index predicate_inc_h, predicate_inc_advance;
/// Setup the params.
CUTLASS_HOST_DEVICE int initialize(
Pointer pointer, Index ld, Index n, Index epilogue_stride_w, Index epilogue_delta_w) {
// The pointer.
this->pointer = pointer;
// Setup the base stride. One "group of threads" per column.
stride_h = ld;
// Each thread output 1 column per iteration. .
inc_h = ld * TileTraits_::Threads::kH;
inc_advance = inc_h + epilogue_stride_w;
predicate_offset = n;
predicate_inc_h = TileTraits_::Threads::kH;
predicate_inc_advance = predicate_inc_h + epilogue_delta_w;
// It worked.
return 0;
}
};
Params params;
Coord<4> thread_offset;
/// Ctor.
CUTLASS_DEVICE WmmaGemmGlobalIteratorCd() {}
/// Ctor.
CUTLASS_DEVICE WmmaGemmGlobalIteratorCd(Params const& params,
const Coord<3>& bounds,
const Coord<3>& block,
int const pointer_offset = 0,
int const pred_offset = 0,
ThreadOffset thread_offset_func = ThreadOffset())
: params(params) {
thread_offset = thread_offset_func();
// Each warp works on a different column of the tile.
int const h = thread_offset[1] + block[1];
// Each lane writes a different element.
int const w = thread_offset[2] + block[2];
// Setup the pointer.
this->params.pointer += ((h * params.stride_h + w) + pointer_offset);
// Prepare the vector of predicates.
for (int i = 0; i < Base::Iterations::kW; ++i) {
predicates.set(i, w + i * Base::Delta::kW < bounds[2]);
}
this->params.predicate_offset -= (h + pred_offset);
}
/// The accessor.
CUTLASS_DEVICE void get(typename Base::AccessType& value, int d, int h, int w, int c) const {
int const imm =
ComputeOffsetFromStrides<typename Base::ImmediateOffsetStrides>::get(0, 0, w, c);
Load<Scalar, TileTraits_::kAccessSize, MemorySpace::kGlobal>::load(value, params.pointer, imm);
}
/// Increment the pointer in the C dimension.
CUTLASS_DEVICE void inc_c() {}
/// Increment the pointer in the W dimension.
CUTLASS_DEVICE void inc_w() {}
/// Increment the pointer in the H dimension.
CUTLASS_DEVICE void inc_h() {
params.pointer += params.inc_h;
params.predicate_offset -= params.predicate_inc_h;
}
/// Increment the pointer in the D dimension.
CUTLASS_DEVICE void inc_d() {}
/// Increment the pointer to move to the next iteration.
CUTLASS_DEVICE void inc_advance() {
params.pointer += params.inc_advance;
params.predicate_offset -= params.predicate_inc_advance;
}
/// The accessor.
CUTLASS_DEVICE void set(typename Base::AccessType const& value, int d, int h, int w, int c) {
int const imm =
ComputeOffsetFromStrides<typename Base::ImmediateOffsetStrides>::get(d, h, w, 0);
Store<Scalar, TileTraits_::kAccessSize, MemorySpace::kGlobal>::store(
value, params.pointer, imm);
}
/// Test the predicate.
CUTLASS_DEVICE bool valid(int d, int h, int w, int c) const {
return predicates.at(w) && params.predicate_offset > 0;
}
/// The predicates for the row.
cutlass::PredicateVector<Base::Iterations::kW> predicates;
};
////////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace gemm
} // namespace cutlass
|
