Cutlass CUDA Templates for Linear Algebra Subroutines and Solvers

tile_allocation.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/shape.h"
31 #include "cutlass/fragment.h"
32 #include "cutlass/tensor_ref.h"
33 #include "cutlass/tensor_view.h"
34 #include "cutlass/zip_tensor_ref.h"
35
36 namespace cutlass {
37
39
41 template <typename Scalar_, typename Shape_>
42 struct TileAllocation {
43 //
44 // Type definitions
45 //
46
48 typedef Scalar_ Scalar;
49
51 typedef typename StorageType<sizeof(Scalar)>::Type Storage;
52
54 typedef Shape_ Shape;
55
57 typedef typename ShapeStrides<Shape, 1>::Shape Strides;
58
60 typedef TensorRef<Scalar const, 4> ConstTensorRef;
61
63 typedef TensorRef<Scalar, 4> TensorRef;
64
66 typedef TensorView<Scalar const, 4> ConstTensorView;
67
69 typedef TensorView<Scalar, 4> TensorView;
70
71 //
72 // Data members
73 //
74
76 Storage storage[Shape::kD][Shape::kH][Shape::kW][Shape::kC];
77
78 //
79 // Methods
80 //
81
83 CUTLASS_DEVICE
84 Scalar *data() { return reinterpret_cast<Scalar *>(&storage[0][0][0][0]); }
85
87 CUTLASS_DEVICE
88 Scalar const *data() const { return reinterpret_cast<Scalar const *>(&storage[0][0][0][0]); }
89
91 CUTLASS_DEVICE
92 TensorRef reference() {
93 return TensorRef(data(), make_Coord(Strides::kD, Strides::kH, Strides::kW, Strides::kC));
94 }
95
97 CUTLASS_DEVICE
98 ConstTensorRef reference() const {
99 return ConstTensorRef(data(), make_Coord(Strides::kD, Strides::kH, Strides::kW, Strides::kC));
100 }
101
103 CUTLASS_DEVICE
104 TensorView view() {
105 return TensorView(
106 data(),
107 make_Coord(Strides::kD, Strides::kH, Strides::kW, Strides::kC),
108 make_Coord(Shape::kD, Shape::kH, Shape::kW, Shape::kC));
109 }
110
112 CUTLASS_DEVICE
113 ConstTensorView view() const {
114 return TensorView(
115 data(),
116 make_Coord(Strides::kD, Strides::kH, Strides::kW, Strides::kC),
117 make_Coord(Shape::kD, Shape::kH, Shape::kW, Shape::kC));
118 }
119 };
120
122
124 template <typename First_, typename Second_>
125 struct ZipTileAllocation {
126 //
127 // Type definitions
128 //
129
131 typedef First_ First;
132
134 typedef Second_ Second;
135
137 typedef ZipTensorRef<typename First::TensorRef, typename Second::TensorRef> TensorRef;
138
140 typedef ZipTensorRef<typename First::ConstTensorRef, typename Second::ConstTensorRef>
141 ConstTensorRef;
142
143 //
144 // Data members
145 //
146
148 First first;
149
151 Second second;
152
153 //
154 // Methods
155 //
156
158 CUTLASS_DEVICE
159 TensorRef reference() { return TensorRef(first.reference(), second.reference()); }
160
162 CUTLASS_DEVICE
163 ConstTensorRef reference() const { return ConstTensorRef(first.reference(), second.reference()); }
164 };
165
167
168 } // namespace cutlass
cutlass
Definition: convert.h:33
cutlass::ZipTileAllocation::reference
CUTLASS_DEVICE TensorRef reference()
Returns a TensorRef object pointing to the data.
Definition: tile_allocation.h:159
cutlass::ZipTileAllocation::ConstTensorRef
ZipTensorRef< typename First::ConstTensorRef, typename Second::ConstTensorRef > ConstTensorRef
Defines the tensor reference for this allocation.
Definition: tile_allocation.h:141
cutlass::TileAllocation::view
CUTLASS_DEVICE TensorView view()
Returns a TensorView object pointing to the data.
Definition: tile_allocation.h:104
cutlass::ShapeStrides::Shape
Shape< Shape_::kH *Shape_::kW *Shape_::kC, Shape_::kW *Shape_::kC, Shape_::kC, elementsPerAccess > Shape
Definition: shape.h:170
tensor_ref.h
Defines a structure containing strides, bounds, and a pointer to tensor data.
cutlass::ZipTensorRef
Definition: zip_tensor_ref.h:38
cutlass::TileAllocation::TensorView
TensorView< Scalar, 4 > TensorView
View of memory.
Definition: tile_allocation.h:69
cutlass::TileAllocation::TensorRef
TensorRef< Scalar, 4 > TensorRef
Defines the tensor reference for this allocation.
Definition: tile_allocation.h:63
cutlass::TileAllocation::data
CUTLASS_DEVICE Scalar const * data() const
Returns a const pointer to the raw data.
Definition: tile_allocation.h:88
cutlass::make_Coord
CUTLASS_HOST_DEVICE Coord< 1 > make_Coord(int _0)
Helper to make a 2-element coordinate.
Definition: coord.h:368
cutlass::TileAllocation::reference
CUTLASS_DEVICE TensorRef reference()
Returns a TensorRef object pointing to the data.
Definition: tile_allocation.h:92
cutlass::ZipTileAllocation::Second
Second_ Second
Second tensor allocation.
Definition: tile_allocation.h:134
cutlass::TileAllocation::storage
Storage storage[Shape::kD][Shape::kH][Shape::kW][Shape::kC]
Storage.
Definition: tile_allocation.h:76
cutlass::TileAllocation::view
CUTLASS_DEVICE ConstTensorView view() const
Returns a TensorView object pointing to the data.
Definition: tile_allocation.h:113
cutlass::ZipTileAllocation::TensorRef
ZipTensorRef< typename First::TensorRef, typename Second::TensorRef > TensorRef
Defines the tensor reference for this allocation.
Definition: tile_allocation.h:137
cutlass::TileAllocation::Strides
ShapeStrides< Shape, 1 >::Shape Strides
Strides.
Definition: tile_allocation.h:57
tensor_view.h
Defines a structure containing strides and a pointer to tensor data.
cutlass::Shape::kH
static int const kH
The height of the cube.
Definition: shape.h:68
cutlass::ZipTileAllocation::second
Second second
Second tensor allocation.
Definition: tile_allocation.h:151
cutlass::StorageType
Definition: fragment.h:76
cutlass::Shape::kC
static int const kC
The number of scalars per element.
Definition: shape.h:72
cutlass::TileAllocation::ConstTensorRef
TensorRef< Scalar const, 4 > ConstTensorRef
Defines the tensor reference for this allocation.
Definition: tile_allocation.h:60
cutlass::ZipTileAllocation::first
First first
First tensor allocation.
Definition: tile_allocation.h:148
cutlass::TileAllocation::reference
CUTLASS_DEVICE ConstTensorRef reference() const
Returns a TensorRef object pointing to the data.
Definition: tile_allocation.h:98
zip_tensor_ref.h
Defines a structure containing a pair of TensorRef-like objects.
cutlass::TileAllocation::Scalar
Scalar_ Scalar
Scalar element.
Definition: tile_allocation.h:48
cutlass::ZipTileAllocation
Manages a pair of tile allocations as if they are one allocation.
Definition: tile_allocation.h:125
cutlass::TileAllocation::Shape
Shape_ Shape
Size of the allocation in units of scalars.
Definition: tile_allocation.h:54
cutlass::TileAllocation
Class for storing a tile in memory and accessing it through a tensor ref.
Definition: tile_allocation.h:42
cutlass::TileAllocation::data
CUTLASS_DEVICE Scalar * data()
Returns a pointer to the raw data.
Definition: tile_allocation.h:84
cutlass::TileAllocation::ConstTensorView
TensorView< Scalar const, 4 > ConstTensorView
View of memory.
Definition: tile_allocation.h:66
cutlass::ZipTileAllocation::reference
CUTLASS_DEVICE ConstTensorRef reference() const
Returns a TensorRef object pointing to the data.
Definition: tile_allocation.h:163
cutlass::Shape::kW
static int const kW
The width of the cube.
Definition: shape.h:70
cutlass::Shape::kD
static int const kD
The depth of the cube.
Definition: shape.h:66
shape.h
Defines Shape implementing the Layout concept for representing a 4D hypercube of objects.
cutlass::TileAllocation::Storage
StorageType< sizeof(Scalar)>::Type Storage
The actual storage (may differ from the scalar type)
Definition: tile_allocation.h:51
fragment.h
Defines Fragment, a statically-sized array for storing parts of matrices within a thread&#39;s registers...
cutlass::ZipTileAllocation::First
First_ First
First tensor allocation.
Definition: tile_allocation.h:131