Cutlass CUDA Templates for Linear Algebra Subroutines and Solvers

tensor_ref.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/coord.h"
31 #include "cutlass/cutlass.h"
32 #include "cutlass/vector.h"
33
34 namespace cutlass {
35
37
40 template <int Rank>
41 struct IdentityTensorMapFunc {
42 static int const kStorageRank = Rank;
43 CUTLASS_HOST_DEVICE
44 Coord<Rank> operator()(Coord<Rank> const &coord) const {
45 return coord;
46 }
47 };
48
50
51 /* \brief Structure modeling a pointer and stride into a tensor.
52
53 A tensor consists of an index space with Rank_ dimensions. It is stored in memory modeled
54 as an n-D array, where n = StorageRank_. A mapping function maps the logical coordinates of the
55 tensor's index space into the n-D array, and a stride vector maps the n-D array to linear memory.
56
57 CUTLASS requires the n-D array's least significant, "fastest changing" dimension to
58 be contiguous in memory. It therefore has a stride of 1 and is not stored. Construction is offered
59 from vectors of full StorageRank and of the 'compact' rank, though it is in error to construct
60 with the least significant stride != 1.
61
62 The requirement that the least significant dimension be consecutive enables numerous optimizations
63 and assumptions about vectorizing memory accesses throughout CUTLASS. It also matches various
64 BLAS conventions in which only the "leading dimension" or most significant stride of a rank=2
65 matrix is provided.
66
67 This does affect the ability of constructing arbitrary "sparse" 2-D matrices in memory where all
68 stride elements are > 1. This can be overcome by defining a custom mapping function and a
69 StorageRank of 3 or more.
70
71
72 Examples:
73
74 (These examples use helpers for matrix layouts defined in cutlass/matrix_traits.h)
75
76 1. Column-major matrix may be represented as a rank=2 tensor:
77
78 TensorRef<float, 2, MatrixLayout::ColumnMajor> A(ptr_A, make_Coord(ldm, 1));
79
80 2. Row-major matrix may be represented as a rank=2 tensor:
81
82 TensorRef<float, 2, MatrixLayout::RowMajor> B(ptr_A, ldm);
83
84 3. An interleaved matrix may be represented as a rank=2 tensor:
85
86 TensorRef<int8_t, 2, MatrixLayout::ColumnMajorInterleaved<32> > C;
87
88 4. Defining a sparse matrix with arbitrary strides in each dimension
89
90 struct ContiguousLayout {
91
93 static int const kStorageRank = 3;
94
96 CUTLASS_HOST_DEVICE
97 Coord<3> operator()(MatrixCoord const &coord) const {
98 return make_Coord(coord.row(), coord.column(), 0);
99 }
100 };
101
102 typedef TensorRef<float, 2, ContiguousLayout> ContiguousTensorRef;
103
104 // Construct the TensorRef object from a pair of stride values
105 ContiguousTensorRef D(ptr_D, make_Coord(row_stride, column_stride));
106
107
108 5. A helper exists to define a TensorRef for a contiguous matrix whose layout
109 is not known at compile time.
110
111 MatrixLayout::Kind layout; // Could be MatrixLayout::kRowMajor or MatrixLayout::kColumnMajor
112 int ldm; // leading dimension
113
114 ContiguousTensorRef E(ptr_E, ContiguousLayout::stride(layout, ldm));
115
116 */
117 template <
119 typename Storage_,
121 int Rank_,
123 typename MapFunc_ = IdentityTensorMapFunc<Rank_>,
125 int StorageRank_ = MapFunc_::kStorageRank,
127 typename Index_ = int,
129 typename LongIndex_ = long long
130 >
131 class TensorRef {
132 public:
134 typedef Storage_ Storage;
135
137 static int const kRank = Rank_;
138
140 typedef MapFunc_ MapFunc;
141
143 static int const kStorageRank = StorageRank_;
144
146 typedef Index_ Index;
147
149 typedef LongIndex_ LongIndex;
150
152 typedef Coord<kRank> TensorCoord;
153
155 typedef Coord<kStorageRank> StorageCoord;
156
159 typedef Coord<kStorageRank - 1> StrideVector;
160
162 typedef TensorRef<
163 typename platform::remove_const<Storage>::type const,
164 Rank_,
165 MapFunc_,
166 StorageRank_,
167 Index_,
168 LongIndex_> ConstTensorRef;
169
173 static_assert(kRank > 0, "Cannot define a zero-rank TensorRef");
174
175 //
176 // Definitions included for backwards compatibility - to be removed in next major release
177 //
178
180 typedef TensorCoord Coord_t;
181
183 static int const Rank = kRank;
184
185 private:
186
188 Storage* ptr_;
189
191 StrideVector stride_;
192
194 MapFunc coord_map_;
195
196 public:
197
198 //
199 // Methods
200 //
201
203 CUTLASS_HOST_DEVICE
204 TensorRef(Storage *ptr = nullptr): ptr_(ptr) {
205 for (int i = 0; i < kStorageRank - 1; ++i) {
206 stride_[i] = 1;
207 }
208 }
209
211 // Higher ranks are projected onto the fastest-changing rank.
212 CUTLASS_HOST_DEVICE
213 TensorRef(Storage* ptr, Index ldm) {
214 ptr_ = ptr;
215 for (int i = 0; i < kStorageRank - 1; ++i) {
216 stride_[i] = ldm;
217 }
218 }
219
221 CUTLASS_HOST_DEVICE
222 TensorRef(Storage* ptr, StrideVector const& stride) : ptr_(ptr), stride_(stride) {
223
224 }
225
228 CUTLASS_HOST_DEVICE
229 TensorRef(Storage* ptr, StorageCoord const& stride) {
230 // Fastest-changing stride must be one
231 if (stride.at(kStorageRank - 1) == 1) {
232 ptr_ = ptr;
233 for (int i = 0; i < kStorageRank - 1; ++i) {
234 stride_[i] = stride[i];
235 }
236 }
237 else {
238 // Fastest-chaning stride must be 1.
239 reset();
240 }
241 }
242
244 CUTLASS_HOST_DEVICE
245 TensorRef(
246 TensorRef<
247 typename platform::remove_const<Storage>::type,
248 kRank,
249 MapFunc,
250 kStorageRank,
251 Index,
252 LongIndex> const &ref
253 ):
254 ptr_(ref.data()) {
255 for (int i = 0; i < kStorageRank - 1; ++i) {
256 stride_[i] = ref.stride(i);
257 }
258 }
259
261 CUTLASS_HOST_DEVICE
262 ConstTensorRef const_ref() const {
263 return ConstTensorRef(*this);
264 }
265
267 CUTLASS_HOST_DEVICE
268 void reset(Storage* ptr = nullptr) {
269 ptr_ = ptr;
270 }
271
273 CUTLASS_HOST_DEVICE
274 void reset(Storage* ptr, StorageCoord const & stride) {
275 // Fastest-changing stride must be one
276 if (stride.at(kStorageRank - 1) == 1) {
277 ptr_ = ptr;
278 for (int i = 0; i < kStorageRank - 1; ++i) {
279 stride_[i] = stride[i];
280 }
281 }
282 else {
283 // Fastest-changing stride must be 1 - this is an error.
284 reset();
285 }
286 }
287
289 CUTLASS_HOST_DEVICE
290 bool good() const {
291 return ptr_ != nullptr;
292 }
293
295 CUTLASS_HOST_DEVICE
296 Storage * data() const { return ptr_; }
297
299 CUTLASS_HOST_DEVICE
300 StorageCoord stride() const {
301 StorageCoord ld;
302 for (int i = 0; i < kStorageRank - 1; ++i) {
303 ld[i] = stride_[i];
304 }
305 ld[kStorageRank - 1] = 1;
306 return ld;
307 }
308
310 CUTLASS_HOST_DEVICE
311 Index stride(int dim) const {
312 // fastest-changing stride assumbed to be 1
313 if (dim + 1 >= kStorageRank) {
314 return 1;
315 }
316 return stride_.at(dim);
317 }
318
320 CUTLASS_HOST_DEVICE
321 Index leading_dim(int idx = 0) const { return stride(idx); }
322
324 CUTLASS_HOST_DEVICE
325 StorageCoord map(TensorCoord const &coord) const {
326 return coord_map_(coord);
327 }
328
330 CUTLASS_HOST_DEVICE
331 LongIndex offset(TensorCoord const& coord) const {
332 return stride().template dot<LongIndex>(map(coord));
333 }
334
336 CUTLASS_HOST_DEVICE
337 Storage& at(TensorCoord const& coord) const {
338 return ptr_[offset(coord)];
339 }
340
342 CUTLASS_HOST_DEVICE
343 Storage& at(LongIndex idx) const { return ptr_[idx]; }
344
346 CUTLASS_HOST_DEVICE
347 Storage& operator[](TensorCoord const& coord) const {
348 return ptr_[offset(coord)];
349 }
350
352 CUTLASS_HOST_DEVICE
353 Storage& operator[](LongIndex idx) const { return ptr_[idx]; }
354
356 CUTLASS_HOST_DEVICE
357 TensorRef & add_pointer_offset(LongIndex delta) {
358 ptr_ += delta;
359 return *this;
360 }
361
363 CUTLASS_HOST_DEVICE
364 TensorRef operator+(TensorCoord const& b) const {
365 TensorRef result(*this);
366 result.add_pointer_offset(offset(b));
367 return result;
368 }
369
371 CUTLASS_HOST_DEVICE
372 TensorRef& operator+=(TensorCoord const& b) {
373 add_pointer_offset(offset(b));
374 return *this;
375 }
376
378 CUTLASS_HOST_DEVICE
379 TensorRef operator-(TensorCoord const& b) const {
380 TensorRef result(*this);
381 result.add_pointer_offset(-offset(b));
382 return result;
383 }
384
386 CUTLASS_HOST_DEVICE
387 TensorRef& operator-=(TensorCoord const& b) {
388 add_pointer_offset(-offset(b));
389 return *this;
390 }
391 };
392
394 //
395 // Partial specializations to handle degenerate cases.
396 //
398
400 template <
402 typename Storage_,
404 int Rank_,
406 typename MapFunc_,
408 typename Index_,
410 typename LongIndex_
411 >
412 class TensorRef<Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_> {
413 public:
415 typedef Storage_ Storage;
416
418 static int const kRank = Rank_;
419
421 typedef MapFunc_ MapFunc;
422
424 static int const kStorageRank = 1;
425
427 typedef Index_ Index;
428
430 typedef LongIndex_ LongIndex;
431
433 typedef Coord<kRank> TensorCoord;
434
436 typedef Coord<kStorageRank> StorageCoord;
437
440 struct StrideVector { };
441
443 typedef TensorRef<
444 typename platform::remove_const<Storage>::type const,
445 Rank_,
446 MapFunc_,
447 kStorageRank,
448 Index_,
449 LongIndex_> ConstTensorRef;
450
451 //
452 // Definitions included for backwards compatibility - to be removed in next major release
453 //
454
456 typedef TensorCoord Coord_t;
457
459 static int const Rank = kRank;
460
461 private:
462
464 Storage* ptr_;
465
467 MapFunc coord_map_;
468
469 public:
470
471 //
472 // Methods
473 //
474
476 CUTLASS_HOST_DEVICE
477 TensorRef(Storage *ptr = nullptr): ptr_(ptr) { }
478
480 CUTLASS_HOST_DEVICE
481 TensorRef(Storage* ptr, StrideVector const& stride) : ptr_(ptr) {
482
483 }
484
487 CUTLASS_HOST_DEVICE
488 TensorRef(Storage* ptr, StorageCoord const& stride) {
489 // Fastest-changing stride must be one
490 if (stride.at(kStorageRank - 1) == 1) {
491 ptr_ = ptr;
492 }
493 else {
494 // Fastest-chaning stride must be 1.
495 reset();
496 }
497 }
498
500 CUTLASS_HOST_DEVICE
501 TensorRef(
502 TensorRef<
503 typename platform::remove_const<Storage>::type,
504 kRank,
505 MapFunc,
506 kStorageRank,
507 Index,
508 LongIndex> const &ref
509 ):
510 ptr_(ref.data()) {
511 }
512
514 CUTLASS_HOST_DEVICE
515 ConstTensorRef const_ref() const {
516 return ConstTensorRef(*this);
517 }
518
520 CUTLASS_HOST_DEVICE
521 void reset(Storage* ptr = nullptr) {
522 ptr_ = ptr;
523 }
524
526 CUTLASS_HOST_DEVICE
527 void reset(Storage* ptr, StorageCoord const & stride) {
528 // Fastest-changing stride must be one
529 if (stride.at(kStorageRank - 1) == 1) {
530 ptr_ = ptr;
531 }
532 else {
533 // Fastest-changing stride must be 1 - this is an error.
534 reset();
535 }
536 }
537
539 CUTLASS_HOST_DEVICE
540 bool good() const {
541 return ptr_ != nullptr;
542 }
543
545 CUTLASS_HOST_DEVICE
546 Storage * data() const { return ptr_; }
547
549 CUTLASS_HOST_DEVICE
550 StorageCoord stride() const {
551 StorageCoord ld;
552 ld[kStorageRank - 1] = 1;
553 return ld;
554 }
555
557 CUTLASS_HOST_DEVICE
558 Index stride(int dim) const {
559 // fastest-changing stride assumbed to be 1
560 return 1;
561 }
562
564 CUTLASS_HOST_DEVICE
565 Index leading_dim(int idx = 0) const { return 1; }
566
568 CUTLASS_HOST_DEVICE
569 StorageCoord map(TensorCoord const &coord) const {
570 return coord_map_(coord);
571 }
572
574 CUTLASS_HOST_DEVICE
575 LongIndex offset(TensorCoord const& coord) const {
576 return stride().template dot<LongIndex>(map(coord));
577 }
578
580 CUTLASS_HOST_DEVICE
581 Storage& at(TensorCoord const& coord) const {
582 return ptr_[offset(coord)];
583 }
584
586 CUTLASS_HOST_DEVICE
587 Storage& at(LongIndex idx) const { return ptr_[idx]; }
588
590 CUTLASS_HOST_DEVICE
591 Storage& operator[](TensorCoord const& coord) const {
592 return ptr_[offset(coord)];
593 }
594
596 CUTLASS_HOST_DEVICE
597 Storage& operator[](LongIndex idx) const { return ptr_[idx]; }
598
600 CUTLASS_HOST_DEVICE
601 TensorRef & add_pointer_offset(LongIndex delta) {
602 ptr_ += delta;
603 return *this;
604 }
605
607 CUTLASS_HOST_DEVICE
608 TensorRef operator+(TensorCoord const& b) const {
609 TensorRef result(*this);
610 result.add_pointer_offset(offset(b));
611 return result;
612 }
613
615 CUTLASS_HOST_DEVICE
616 TensorRef& operator+=(TensorCoord const& b) {
617 add_pointer_offset(offset(b));
618 return *this;
619 }
620
622 CUTLASS_HOST_DEVICE
623 TensorRef operator-(TensorCoord const& b) const {
624 TensorRef result(*this);
625 result.add_pointer_offset(-offset(b));
626 return result;
627 }
628
630 CUTLASS_HOST_DEVICE
631 TensorRef& operator-=(TensorCoord const& b) {
632 add_pointer_offset(-offset(b));
633 return *this;
634 }
635 };
636
638
639 } // namespace cutlass
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::StorageCoord
Coord< kStorageRank > StorageCoord
Coordinate in storage n-D array.
Definition: tensor_ref.h:436
cutlass::TensorRef::operator[]
CUTLASS_HOST_DEVICE Storage & operator[](LongIndex idx) const
Returns a reference to the element at a given linear index.
Definition: tensor_ref.h:353
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::TensorRef
CUTLASS_HOST_DEVICE TensorRef(TensorRef< typename platform::remove_const< Storage >::type, kRank, MapFunc, kStorageRank, Index, LongIndex > const &ref)
Enables conversion from TensorRef of non-const type.
Definition: tensor_ref.h:501
cutlass::TensorRef::Coord_t
TensorCoord Coord_t
Coordinate in logical tensor space.
Definition: tensor_ref.h:173
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::add_pointer_offset
CUTLASS_HOST_DEVICE TensorRef & add_pointer_offset(LongIndex delta)
Adds an offset to each pointer.
Definition: tensor_ref.h:601
cutlass::TensorRef::data
CUTLASS_HOST_DEVICE Storage * data() const
Returns the pointer to referenced data.
Definition: tensor_ref.h:296
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::reset
CUTLASS_HOST_DEVICE void reset(Storage *ptr, StorageCoord const &stride)
Updates the pointer, stride, and location within a TensorRef.
Definition: tensor_ref.h:527
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::at
CUTLASS_HOST_DEVICE Storage & at(TensorCoord const &coord) const
Returns a reference to the element at a given Coord.
Definition: tensor_ref.h:581
cutlass
Definition: convert.h:33
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::TensorRef
CUTLASS_HOST_DEVICE TensorRef(Storage *ptr=nullptr)
Helper for 1-D memory. All higher ranks are projected onto the fastest changing rank.
Definition: tensor_ref.h:477
cutlass::platform::remove_const::type
T type
Definition: platform.h:377
cutlass::TensorRef::kRank
static int const kRank
Logical rank of tensor index space.
Definition: tensor_ref.h:137
coord.h
A Coord is a coordinate of arbitrary rank into a tensor or matrix.
cutlass::TensorRef::reset
CUTLASS_HOST_DEVICE void reset(Storage *ptr, StorageCoord const &stride)
Updates the pointer, stride, and location within a TensorRef.
Definition: tensor_ref.h:274
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::data
CUTLASS_HOST_DEVICE Storage * data() const
Returns the pointer to referenced data.
Definition: tensor_ref.h:546
cutlass::TensorRef::reset
CUTLASS_HOST_DEVICE void reset(Storage *ptr=nullptr)
Updates only the pointer.
Definition: tensor_ref.h:268
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::operator[]
CUTLASS_HOST_DEVICE Storage & operator[](LongIndex idx) const
Returns a reference to the element at a given linear index.
Definition: tensor_ref.h:597
cutlass::TensorRef::TensorRef
CUTLASS_HOST_DEVICE TensorRef(TensorRef< typename platform::remove_const< Storage >::type, kRank, MapFunc, kStorageRank, Index, LongIndex > const &ref)
Enables conversion from TensorRef of non-const type.
Definition: tensor_ref.h:245
cutlass::TensorRef::operator+
CUTLASS_HOST_DEVICE TensorRef operator+(TensorCoord const &b) const
Returns a TensorRef offset by a given amount.
Definition: tensor_ref.h:364
cutlass::TensorRef::MapFunc
MapFunc_ MapFunc
Mapping function from logical coordinate to internal n-D array.
Definition: tensor_ref.h:140
cutlass::TensorRef::Storage
Storage_ Storage
Data type of individual access.
Definition: tensor_ref.h:134
cutlass::TensorRef::stride
CUTLASS_HOST_DEVICE Index stride(int dim) const
Returns the stride of the tensor in the given dimension.
Definition: tensor_ref.h:311
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::stride
CUTLASS_HOST_DEVICE Index stride(int dim) const
Returns the stride of the tensor in the given dimension.
Definition: tensor_ref.h:558
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::TensorRef
CUTLASS_HOST_DEVICE TensorRef(Storage *ptr, StrideVector const &stride)
Constructs from a single pointer and stride vector.
Definition: tensor_ref.h:481
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::operator-=
CUTLASS_HOST_DEVICE TensorRef & operator-=(TensorCoord const &b)
Returns a TensorRef offset by a given amount.
Definition: tensor_ref.h:631
cutlass::TensorRef::stride
CUTLASS_HOST_DEVICE StorageCoord stride() const
Returns the stride of the tensor.
Definition: tensor_ref.h:300
cutlass::TensorRef::add_pointer_offset
CUTLASS_HOST_DEVICE TensorRef & add_pointer_offset(LongIndex delta)
Adds an offset to each pointer.
Definition: tensor_ref.h:357
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::reset
CUTLASS_HOST_DEVICE void reset(Storage *ptr=nullptr)
Updates only the pointer.
Definition: tensor_ref.h:521
cutlass::TensorRef::Index
Index_ Index
Index type.
Definition: tensor_ref.h:146
cutlass::TensorRef::operator[]
CUTLASS_HOST_DEVICE Storage & operator[](TensorCoord const &coord) const
Returns a reference to the element at a given Coord.
Definition: tensor_ref.h:347
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::ConstTensorRef
TensorRef< typename platform::remove_const< Storage >::type const, Rank_, MapFunc_, kStorageRank, Index_, LongIndex_ > ConstTensorRef
Tensor reference to of constant value.
Definition: tensor_ref.h:449
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::at
CUTLASS_HOST_DEVICE Storage & at(LongIndex idx) const
Returns a reference to the element at a given linear index.
Definition: tensor_ref.h:587
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::operator-
CUTLASS_HOST_DEVICE TensorRef operator-(TensorCoord const &b) const
Returns a TensorRef offset by a given amount.
Definition: tensor_ref.h:623
cutlass::TensorRef::TensorRef
CUTLASS_HOST_DEVICE TensorRef(Storage *ptr, Index ldm)
Helper to construct from a pointer and single stride element for 2-D pitch linear memory...
Definition: tensor_ref.h:213
cutlass::TensorRef::map
CUTLASS_HOST_DEVICE StorageCoord map(TensorCoord const &coord) const
Maps a logical coordinate to an n-D array in memory.
Definition: tensor_ref.h:325
cutlass::TensorRef::at
CUTLASS_HOST_DEVICE Storage & at(LongIndex idx) const
Returns a reference to the element at a given linear index.
Definition: tensor_ref.h:343
cutlass::TensorRef::LongIndex
LongIndex_ LongIndex
Typically, strides in memory can be very large.
Definition: tensor_ref.h:149
cutlass::TensorRef::StorageCoord
Coord< kStorageRank > StorageCoord
Coordinate in storage n-D array.
Definition: tensor_ref.h:155
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::Storage
Storage_ Storage
Data type of individual access.
Definition: tensor_ref.h:415
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::operator+
CUTLASS_HOST_DEVICE TensorRef operator+(TensorCoord const &b) const
Returns a TensorRef offset by a given amount.
Definition: tensor_ref.h:608
cutlass::TensorRef::TensorRef
CUTLASS_HOST_DEVICE TensorRef(Storage *ptr=nullptr)
Helper for 1-D memory. All higher ranks are projected onto the fastest changing rank.
Definition: tensor_ref.h:204
cutlass::TensorRef::operator+=
CUTLASS_HOST_DEVICE TensorRef & operator+=(TensorCoord const &b)
Returns a TensorRef offset by a given amount.
Definition: tensor_ref.h:372
cutlass::TensorRef::ConstTensorRef
TensorRef< typename platform::remove_const< Storage >::type const, Rank_, MapFunc_, StorageRank_, Index_, LongIndex_ > ConstTensorRef
Tensor reference to of constant value.
Definition: tensor_ref.h:168
cutlass::TensorRef
Definition: tensor_ref.h:131
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::stride
CUTLASS_HOST_DEVICE StorageCoord stride() const
Returns the stride of the tensor.
Definition: tensor_ref.h:550
cutlass::TensorRef::TensorCoord
Coord< kRank > TensorCoord
Coordinate in logical tensor space.
Definition: tensor_ref.h:152
cutlass::TensorRef::kStorageRank
static int const kStorageRank
Rank of internal storage.
Definition: tensor_ref.h:143
cutlass::TensorRef::Rank
static int const Rank
Logical rank of tensor index space.
Definition: tensor_ref.h:183
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::Coord_t
TensorCoord Coord_t
Coordinate in logical tensor space.
Definition: tensor_ref.h:456
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::TensorRef
CUTLASS_HOST_DEVICE TensorRef(Storage *ptr, StorageCoord const &stride)
Definition: tensor_ref.h:488
CUTLASS_HOST_DEVICE
#define CUTLASS_HOST_DEVICE
Definition: cutlass.h:46
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::good
CUTLASS_HOST_DEVICE bool good() const
Returns true if the TensorRef may be safely accessed.
Definition: tensor_ref.h:540
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::leading_dim
CUTLASS_HOST_DEVICE Index leading_dim(int idx=0) const
Returns the maximum stride element as the &#39;leading dimension&#39;.
Definition: tensor_ref.h:565
cutlass::Coord::at
CUTLASS_HOST_DEVICE Index & at()
Gets the index of a given Coord element.
Definition: coord.h:240
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::operator[]
CUTLASS_HOST_DEVICE Storage & operator[](TensorCoord const &coord) const
Returns a reference to the element at a given Coord.
Definition: tensor_ref.h:591
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::operator+=
CUTLASS_HOST_DEVICE TensorRef & operator+=(TensorCoord const &b)
Returns a TensorRef offset by a given amount.
Definition: tensor_ref.h:616
static_assert
#define static_assert(__e, __m)
Definition: platform.h:153
cutlass::TensorRef::offset
CUTLASS_HOST_DEVICE LongIndex offset(TensorCoord const &coord) const
Computes the offset of an index from the origin of the tensor.
Definition: tensor_ref.h:331
cutlass::IdentityTensorMapFunc::operator()
CUTLASS_HOST_DEVICE Coord< Rank > operator()(Coord< Rank > const &coord) const
Definition: tensor_ref.h:44
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::LongIndex
LongIndex_ LongIndex
Typically, strides in memory can be very large.
Definition: tensor_ref.h:430
cutlass::Coord
Statically-sized array specifying Coords within a tensor.
Definition: coord.h:49
vector.h
Defines a 1D vector of elements held in the registers of each thread.
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::offset
CUTLASS_HOST_DEVICE LongIndex offset(TensorCoord const &coord) const
Computes the offset of an index from the origin of the tensor.
Definition: tensor_ref.h:575
cutlass::TensorRef::at
CUTLASS_HOST_DEVICE Storage & at(TensorCoord const &coord) const
Returns a reference to the element at a given Coord.
Definition: tensor_ref.h:337
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::const_ref
CUTLASS_HOST_DEVICE ConstTensorRef const_ref() const
Returns a reference to constant-valued tensor.
Definition: tensor_ref.h:515
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::Index
Index_ Index
Index type.
Definition: tensor_ref.h:427
cutlass::IdentityTensorMapFunc::kStorageRank
static int const kStorageRank
Definition: tensor_ref.h:42
cutlass::TensorRef::const_ref
CUTLASS_HOST_DEVICE ConstTensorRef const_ref() const
Returns a reference to constant-valued tensor.
Definition: tensor_ref.h:262
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::map
CUTLASS_HOST_DEVICE StorageCoord map(TensorCoord const &coord) const
Maps a logical coordinate to an n-D array in memory.
Definition: tensor_ref.h:569
cutlass::TensorRef::TensorRef
CUTLASS_HOST_DEVICE TensorRef(Storage *ptr, StorageCoord const &stride)
Definition: tensor_ref.h:229
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::MapFunc
MapFunc_ MapFunc
Mapping function from logical coordinate to internal n-D array.
Definition: tensor_ref.h:421
cutlass::TensorRef::TensorRef
CUTLASS_HOST_DEVICE TensorRef(Storage *ptr, StrideVector const &stride)
Constructs from a single pointer and stride vector.
Definition: tensor_ref.h:222
cutlass::TensorRef::StrideVector
Coord< kStorageRank - 1 > StrideVector
Definition: tensor_ref.h:159
cutlass::TensorRef< Storage_, Rank_, MapFunc_, 1, Index_, LongIndex_ >::TensorCoord
Coord< kRank > TensorCoord
Coordinate in logical tensor space.
Definition: tensor_ref.h:433
cutlass.h
Basic include for CUTLASS macros.
cutlass::TensorRef::operator-
CUTLASS_HOST_DEVICE TensorRef operator-(TensorCoord const &b) const
Returns a TensorRef offset by a given amount.
Definition: tensor_ref.h:379
cutlass::TensorRef::good
CUTLASS_HOST_DEVICE bool good() const
Returns true if the TensorRef may be safely accessed.
Definition: tensor_ref.h:290
cutlass::IdentityTensorMapFunc
Definition: tensor_ref.h:41
cutlass::TensorRef::leading_dim
CUTLASS_HOST_DEVICE Index leading_dim(int idx=0) const
Returns the maximum stride element as the &#39;leading dimension&#39;.
Definition: tensor_ref.h:321
cutlass::TensorRef::operator-=
CUTLASS_HOST_DEVICE TensorRef & operator-=(TensorCoord const &b)
Returns a TensorRef offset by a given amount.
Definition: tensor_ref.h:387