host/tensor_foreach.h
Go to the documentation of this file.
/***************************************************************************************************
* Copyright (c) 2017-2019, 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.
*
**************************************************************************************************/
#pragma once
#include <stdexcept>
#include "cutlass/cutlass.h"
namespace cutlass {
namespace reference {
namespace host {
namespace detail {
template <typename Func, int Rank, int RankRemaining>
struct TensorForEachHelper {
static int const kActiveRank = Rank - RankRemaining - 1;
TensorForEachHelper(
Func &func,
Coord<Rank> const &extent,
Coord<Rank> &coord) {
for (int i = 0; i < extent.at(kActiveRank); ++i) {
coord[kActiveRank] = i;
TensorForEachHelper<Func, Rank, RankRemaining - 1>(func, extent, coord);
}
}
};
template <typename Func, int Rank>
struct TensorForEachHelper<Func, Rank, 0> {
static int const kActiveRank = Rank - 1;
TensorForEachHelper(
Func &func,
Coord<Rank> const &extent,
Coord<Rank> &coord) {
for (int i = 0; i < extent.at(kActiveRank); ++i) {
coord[kActiveRank] = i;
func(coord);
}
}
};
} // namespace detail
template <
typename Func,
int Rank>
void TensorForEach(Coord<Rank> extent, Func & func) {
Coord<Rank> coord;
detail::TensorForEachHelper<Func, Rank, Rank - 1>(func, extent, coord);
}
template <
typename Func,
int Rank>
void TensorForEachLambda(Coord<Rank> extent, Func func) {
Coord<Rank> coord;
detail::TensorForEachHelper<Func, Rank, Rank - 1>(func, extent, coord);
}
template <typename Element, typename Func>
struct BlockForEach {
BlockForEach(
Element *ptr,
size_t capacity,
typename Func::Params params = typename Func::Params()) {
Func func(params);
for (size_t index = 0; index < capacity; ++index) {
ptr[index] = func();
}
}
};
} // namespace host
} // namespace reference
} // namespace cutlass
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 chaning 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.