CUTLASS: tensor_foreach.h Source File

| CUTLASS CUDA Templates for Linear Algebra Subroutines and Solvers

- Main Page
- Modules
- Namespaces
- Classes
- Files

- File List
- File Members

host/tensor_foreach.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 **************************************************************************************************/
25 #pragma once
26
27 #include <stdexcept>
28 #include "cutlass/cutlass.h"
29
30 namespace cutlass {
31 namespace reference {
32 namespace host {
33
35
37 namespace detail {
38
40 template <typename Func, int Rank, int RankRemaining>
41 struct TensorForEachHelper {
42
44 static int const kActiveRank = Rank - RankRemaining - 1;
45
47 TensorForEachHelper(
48 Func &func,
49 Coord<Rank> const &extent,
50 Coord<Rank> &coord) {
51
52 for (int i = 0; i < extent.at(kActiveRank); ++i) {
53 coord[kActiveRank] = i;
54 TensorForEachHelper<Func, Rank, RankRemaining - 1>(func, extent, coord);
55 }
56 }
57 };
58
60 template <typename Func, int Rank>
61 struct TensorForEachHelper<Func, Rank, 0> {
62
64 static int const kActiveRank = Rank - 1;
65
67 TensorForEachHelper(
68 Func &func,
69 Coord<Rank> const &extent,
70 Coord<Rank> &coord) {
71
72 for (int i = 0; i < extent.at(kActiveRank); ++i) {
73 coord[kActiveRank] = i;
74 func(coord);
75 }
76 }
77 };
78
79 } // namespace detail
80
82
84 template <
85 typename Func,
86 int Rank>
87 void TensorForEach(Coord<Rank> extent, Func & func) {
88 Coord<Rank> coord;
89 detail::TensorForEachHelper<Func, Rank, Rank - 1>(func, extent, coord);
90 }
91
93
95 template <
96 typename Func,
97 int Rank>
98 void TensorForEachLambda(Coord<Rank> extent, Func func) {
99 Coord<Rank> coord;
100 detail::TensorForEachHelper<Func, Rank, Rank - 1>(func, extent, coord);
101 }
102
104
105 template <typename Element, typename Func>
106 struct BlockForEach {
107
109 BlockForEach(
110 Element *ptr,
111 size_t capacity,
112 typename Func::Params params = typename Func::Params()) {
113
114 Func func(params);
115
116 for (size_t index = 0; index < capacity; ++index) {
117 ptr[index] = func();
118 }
119 }
120 };
121
123
124 } // namespace host
125 } // namespace reference
126 } // namespace cutlass
127
cutlass::reference::host::BlockForEach
Definition: host/tensor_foreach.h:106
cutlass
Definition: aligned_buffer.h:35
cutlass::reference::host::BlockForEach::BlockForEach
BlockForEach(Element *ptr, size_t capacity, typename Func::Params params=typename Func::Params())
Constructor performs the operation.
Definition: host/tensor_foreach.h:109
cutlass::reference::host::detail::TensorForEachHelper::kActiveRank
static int const kActiveRank
Index of the active rank.
Definition: host/tensor_foreach.h:44
cutlass::reference::host::detail::TensorForEachHelper::TensorForEachHelper
TensorForEachHelper(Func &func, Coord< Rank > const &extent, Coord< Rank > &coord)
Constructor for general rank.
Definition: host/tensor_foreach.h:47
cutlass::reference::host::detail::TensorForEachHelper
Helper to perform for-each operation.
Definition: host/tensor_foreach.h:41
cutlass::reference::host::detail::TensorForEachHelper< Func, Rank, 0 >::TensorForEachHelper
TensorForEachHelper(Func &func, Coord< Rank > const &extent, Coord< Rank > &coord)
Constructor for fastest changing rank.
Definition: host/tensor_foreach.h:67
cutlass::Coord
Statically-sized array specifying Coords within a tensor.
Definition: coord.h:43
cutlass::reference::host::TensorForEachLambda
void TensorForEachLambda(Coord< Rank > extent, Func func)
Iterates over the index space of a tensor and calls a C++ lambda.
Definition: host/tensor_foreach.h:98
cutlass::reference::host::TensorForEach
void TensorForEach(Coord< Rank > extent, Func &func)
Iterates over the index space of a tensor.
Definition: host/tensor_foreach.h:87
cutlass::Coord::at
CUTLASS_HOST_DEVICE Index & at()
Gets the index of a given Coord element.
Definition: coord.h:255
cutlass.h
Basic include for CUTLASS.