CUTLASS
CUDA Templates for Linear Algebra Subroutines and Solvers
epilogue_workspace.h
Go to the documentation of this file.
1 /***************************************************************************************************
2  * Copyright (c) 2017-2019, 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  **************************************************************************************************/
46 #pragma once
47 
48 #include "cutlass/cutlass.h"
49 #include "cutlass/numeric_types.h"
50 #include "cutlass/array.h"
51 
53 
54 namespace cutlass {
55 namespace epilogue {
56 
58 
59 template <
60  typename Shape_,
61  int WarpCount,
62  typename FragmentC_
63 >
65 public:
66 
67  using Shape = Shape_;
68  using FragmentC = FragmentC_;
69  using ElementC = typename FragmentC::value_type;
70 
71  static int const kWarpCount = WarpCount;
72 
74  static int const kAccessSizeInBits = 128;
75 
77  static int const kWarpSize = 32;
78 
80  static int const kElementsPerAccess =
81  kAccessSizeInBits / sizeof_bits<ElementC>::value;
82 
84  static int const kIterations = FragmentC::kElements / kElementsPerAccess;
85 
87  !(FragmentC::kElements % kElementsPerAccess),
88  "The number of accumulators must be divisible by the access size.");
89 
91  static int const kWarpAccesses = kIterations * kWarpSize;
92 
94  static int const kThreadblockAccesses = kWarpAccesses * kWarpCount;
95 
97  struct Params {
98 
101 
103  int stride_n;
104 
106  int stride_k;
107 
108  //
109  // Methods
110  //
111 
114  ElementC *ptr_C,
115  int stride_n_,
116  int stride_k_
117  ):
118  ptr_C(ptr_C), stride_n(stride_n_ / kElementsPerAccess), stride_k(stride_k_ / kElementsPerAccess) {
119 
120  }
121  };
122 
124  struct SharedStorage {
125  // Intentionally empty
126  };
127 
128 private:
129 
130  struct alignas((kAccessSizeInBits / 8)) AccessType {
131  Array<ElementC, kElementsPerAccess> storage;
132  };
133 
135  AccessType *pointer_;
136 
138  int stride_n_;
139 
141  int stride_k_;
142 
143 public:
144 
146  CUTLASS_DEVICE
148  Params const &params,
149  SharedStorage &,
150  int warp_idx,
151  int lane_idx
152 
153  ):
154  pointer_(reinterpret_cast<AccessType *>(params.ptr_C)),
155  stride_n_(params.stride_n),
156  stride_k_(params.stride_k) {
157 
158  // Add per-thread offset
159  pointer_ += lane_idx + warp_idx * kWarpAccesses;
160  }
161 
163  CUTLASS_DEVICE
165  cutlass::gemm::GemmCoord problem_size,
166  cutlass::gemm::GemmCoord tb_tile_coord,
167  FragmentC const &accum) {
168 
169  // Compute offset for entire threadblock (note, per-thread offset has been folded in already)
170  AccessType *pointer = pointer_ +
171  tb_tile_coord.m() * kThreadblockAccesses +
172  tb_tile_coord.n() * stride_n_ +
173  tb_tile_coord.k() * stride_k_;
174 
175  // Cast to vectorized view of accumulator fragments
176  AccessType const * src_pointer = reinterpret_cast<AccessType const *>(&accum);
177 
178  // Write out accumulators at full speed
180  for (int i = 0; i < kIterations; ++i) {
181  pointer[i * kWarpSize] = src_pointer[i];
182  }
183  }
184 };
185 
187 
188 } // namespace epilogue
189 } // namespace cutlass
190 
Definition: aligned_buffer.h:35
Shared storage allocation needed by the epilogue.
Definition: epilogue_workspace.h:124
static int const kAccessSizeInBits
Optimize for 128b accesses.
Definition: epilogue_workspace.h:74
Definition: include/cutlass/gemm/gemm.h:94
static int const kWarpAccesses
Total number of vectorized accesses in warp (in units of vector)
Definition: epilogue_workspace.h:91
CUTLASS_HOST_DEVICE Params(ElementC *ptr_C, int stride_n_, int stride_k_)
Definition: epilogue_workspace.h:113
CUTLASS_HOST_DEVICE Index const & n() const
Returns the GEMM N coordinate.
Definition: include/cutlass/gemm/gemm.h:137
static int const kIterations
Number of stores per thread.
Definition: epilogue_workspace.h:84
CUTLASS_DEVICE void operator()(cutlass::gemm::GemmCoord problem_size, cutlass::gemm::GemmCoord tb_tile_coord, FragmentC const &accum)
Streams the result to global memory.
Definition: epilogue_workspace.h:164
Shape_ Shape
Definition: epilogue_workspace.h:67
CUTLASS_HOST_DEVICE Index const & k() const
Returns the GEMM K coordinate.
Definition: include/cutlass/gemm/gemm.h:145
typename FragmentC::value_type ElementC
Definition: epilogue_workspace.h:69
Statically sized array of elements that accommodates all CUTLASS-supported numeric types and is safe ...
#define CUTLASS_PRAGMA_UNROLL
Definition: cutlass.h:110
Defines the size of an element in bits.
Definition: numeric_types.h:42
ElementC * ptr_C
Pointer to C matrix.
Definition: epilogue_workspace.h:100
FragmentC_ FragmentC
Definition: epilogue_workspace.h:68
#define CUTLASS_HOST_DEVICE
Definition: cutlass.h:89
Top-level include for all CUTLASS numeric types.
CUTLASS_DEVICE EpilogueWorkspace(Params const &params, SharedStorage &, int warp_idx, int lane_idx)
Constructor.
Definition: epilogue_workspace.h:147
#define static_assert(__e, __m)
Definition: platform.h:153
static int const kWarpCount
Definition: epilogue_workspace.h:71
int stride_n
Stride between tiles along the GEMM N dimension (in units of vectors)
Definition: epilogue_workspace.h:103
static int const kElementsPerAccess
Vector length of accesses.
Definition: epilogue_workspace.h:80
CUTLASS_HOST_DEVICE Index const & m() const
Returns the GEMM M coordinate.
Definition: include/cutlass/gemm/gemm.h:129
int stride_k
Stride between tiles along the GEMM K dimension (in units of vectors)
Definition: epilogue_workspace.h:106
Parameters structure.
Definition: epilogue_workspace.h:97
Definition: epilogue_workspace.h:64
Basic include for CUTLASS.
static int const kWarpSize
Warp size from the perspective of memory operations.
Definition: epilogue_workspace.h:77
static int const kThreadblockAccesses
Total number of vectorized accesses in threadblock tile (in units of vector)
Definition: epilogue_workspace.h:94