coord.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 "cutlass/cutlass.h"
namespace cutlass {
template <
int Rank_,
typename Index_ = int,
typename LongIndex_ = int64_t
>
struct Coord {
public:
//
// Type and constant definitions
//
static int const kRank = Rank_;
using Index = Index_;
using LongIndex = LongIndex_;
private:
//
// Data members
//
Index idx[kRank];
public:
//
// Methods
//
CUTLASS_HOST_DEVICE
explicit Coord(Index value = Index(0)) {
for (int i = 0; i < kRank; ++i) {
idx[i] = value;
}
}
CUTLASS_HOST_DEVICE
Coord(Index const (&_idx)[kRank]) {
for (int i = 0; i < kRank; ++i) {
idx[i] = _idx[i];
}
}
CUTLASS_HOST_DEVICE
Coord(Coord<kRank, Index, LongIndex> const &coord) {
for (int i = 0; i < kRank; ++i) {
idx[i] = coord[i];
}
}
template <int Slice>
CUTLASS_HOST_DEVICE
Coord<Slice> slice(int start = 0, Index identity = 0) const {
Coord<Slice> result;
for (int i = 0; i < Slice; ++i) {
if (i + start < kRank) {
result[i] = idx[i + start];
}
else {
result[i] = identity;
}
}
return result;
}
CUTLASS_HOST_DEVICE
int min_dim_index() const {
int i = 0;
for (int j = 1; j < kRank; ++j) {
if (idx[j] < idx[i]) {
i = j;
}
}
return i;
}
CUTLASS_HOST_DEVICE
int max_dim_index() const {
int i = 0;
for (int j = 1; j < kRank; ++j) {
if (idx[j] > idx[i]) {
i = j;
}
}
return i;
}
CUTLASS_HOST_DEVICE
explicit operator bool() const {
for (int i = 0; i < kRank; ++i) {
if (idx[i]) {
return true;
}
}
return false;
}
CUTLASS_HOST_DEVICE
bool operator!() const {
for (int i = 0; i < kRank; ++i) {
if (idx[i]) {
return false;
}
}
return true;
}
CUTLASS_HOST_DEVICE
Coord operator+(Coord const& b) const {
Coord c;
for (int i = 0; i < kRank; ++i) {
c.idx[i] = idx[i] + b.idx[i];
}
return c;
}
CUTLASS_HOST_DEVICE
Coord operator-(Coord const& b) const {
Coord c;
for (int i = 0; i < kRank; ++i) {
c.idx[i] = idx[i] - b.idx[i];
}
return c;
}
CUTLASS_HOST_DEVICE
Coord operator*(Coord const& b) const {
Coord c;
for (int i = 0; i < kRank; ++i) {
c.idx[i] = idx[i] * b.idx[i];
}
return c;
}
CUTLASS_HOST_DEVICE
Coord operator/(Coord const& b) const {
Coord c;
for (int i = 0; i < kRank; ++i) {
c.idx[i] = idx[i] / b.idx[i];
}
return c;
}
CUTLASS_HOST_DEVICE
Coord& operator+=(Coord const& b) {
for (int i = 0; i < kRank; ++i) {
idx[i] += b.idx[i];
}
return *this;
}
CUTLASS_HOST_DEVICE
Coord& operator-=(Coord const& b) {
for (int i = 0; i < kRank; ++i) {
idx[i] -= b.idx[i];
}
return *this;
}
CUTLASS_HOST_DEVICE
Coord& operator*=(Coord const& b) {
for (int i = 0; i < kRank; ++i) {
idx[i] *= b.idx[i];
}
return *this;
}
CUTLASS_HOST_DEVICE
Coord& operator/=(Coord const& b) {
for (int i = 0; i < kRank; ++i) {
idx[i] /= b.idx[i];
}
return *this;
}
CUTLASS_HOST_DEVICE Index& operator[](int dim) { return idx[dim]; }
CUTLASS_HOST_DEVICE Index const& operator[](int dim) const { return idx[dim]; }
CUTLASS_HOST_DEVICE
LongIndex dot(Coord const& b, LongIndex sum = LongIndex(0)) const {
for (int i = 0; i < kRank; ++i) {
sum += idx[i] * b.idx[i];
}
return sum;
}
template <int Dim>
CUTLASS_HOST_DEVICE Index& at() {
return idx[Dim];
}
CUTLASS_HOST_DEVICE
Index& at(int dim) { return idx[dim]; }
template <int Dim>
CUTLASS_HOST_DEVICE Index const& at() const {
return idx[Dim];
}
CUTLASS_HOST_DEVICE
Index const& at(int dim) const { return idx[dim]; }
CUTLASS_HOST_DEVICE
bool operator==(Coord const& b) const {
bool equal = true;
for (int i = 0; equal && i < kRank; ++i) {
equal = (idx[i] == b.idx[i]);
}
return equal;
}
CUTLASS_HOST_DEVICE
bool operator!=(Coord const& b) const { return !(*this == b); }
CUTLASS_HOST_DEVICE
Coord& clamp(Coord const& max, Coord const& min = Coord()) {
for (int i = 0; i < kRank; ++i) {
idx[i] = __NV_STD_MAX(__NV_STD_MIN(idx[i], max.idx[i]), min.idx[i]);
}
return *this;
}
CUTLASS_HOST_DEVICE
Index sum() const {
Index sum_(idx[0]);
for (int i = 1; i < kRank; ++i) {
sum_ += idx[i];
}
return sum_;
}
CUTLASS_HOST_DEVICE
LongIndex product() const {
LongIndex product_(idx[0]);
for (int i = 1; i < kRank; ++i) {
product_ *= idx[i];
}
return product_;
}
CUTLASS_HOST_DEVICE
bool operator<(Coord const &b) const {
for (int i = 0; i < kRank; ++i) {
if (!(idx[i] < b[i])) {
return false;
}
}
return true;
}
CUTLASS_HOST_DEVICE
bool operator<=(Coord const &b) const {
for (int i = 0; i < kRank; ++i) {
if (!(idx[i] <= b[i])) {
return false;
}
}
return true;
}
CUTLASS_HOST_DEVICE
bool operator>(Coord const &b) const {
return !(*this <= b);
}
CUTLASS_HOST_DEVICE
bool operator>=(Coord const &b) const {
return !(*this < b);
}
};
} // namespace cutlass
namespace cutlass {
template <int Rank, typename Index>
CUTLASS_HOST_DEVICE
Coord<Rank, Index> operator/(Index s, Coord<Rank, Index> coord) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < Rank; ++i) {
coord[i] = s / coord[i];
}
return coord;
}
template <int Rank, typename Index>
CUTLASS_HOST_DEVICE
Coord<Rank, Index> operator/(Coord<Rank, Index> coord, Index s) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < Rank; ++i) {
coord[i] /= s;
}
return coord;
}
//
// Integer-valued make_Coord
//
CUTLASS_HOST_DEVICE
Coord<1> make_Coord(int _0) {
int values[1] = {_0};
return Coord<1>(values);
}
CUTLASS_HOST_DEVICE
Coord<2> make_Coord(int _0, int _1) {
int values[2] = {_0, _1};
return Coord<2>(values);
}
CUTLASS_HOST_DEVICE
Coord<3> make_Coord(int _0, int _1, int _2) {
int values[3] = {_0, _1, _2};
return Coord<3>(values);
}
CUTLASS_HOST_DEVICE
Coord<4> make_Coord(int _0, int _1, int _2, int _3) {
int values[4] = {_0, _1, _2, _3};
return Coord<4>(values);
}
} // namespace cutlass
cutlass::Coord::operator-
CUTLASS_HOST_DEVICE Coord operator-(Coord const &b) const
Element-wise subtraction.
Definition: coord.h:174
cutlass::Coord::min_dim_index
CUTLASS_HOST_DEVICE int min_dim_index() const
Returns the index of the dimension with least value.
Definition: coord.h:118
cutlass::platform::max
CUTLASS_HOST_DEVICE constexpr const T & max(const T &a, const T &b)
std::max
Definition: platform.h:189
cutlass::Coord::at
CUTLASS_HOST_DEVICE Index const & at() const
Gets the index of a given Coord element.
Definition: coord.h:265
cutlass::Coord::kRank
static int const kRank
Number of elements in Coord.
Definition: coord.h:52
cutlass::Coord::operator==
CUTLASS_HOST_DEVICE bool operator==(Coord const &b) const
Determines if two Coord<> objects are equal.
Definition: coord.h:275
cutlass
Definition: aligned_buffer.h:35
cutlass::Coord::operator+
CUTLASS_HOST_DEVICE Coord operator+(Coord const &b) const
Element-wise addition.
Definition: coord.h:164
cutlass::Coord::operator[]
CUTLASS_HOST_DEVICE Index const & operator[](int dim) const
Member access operator.
Definition: coord.h:242
cutlass::Coord::Coord
CUTLASS_HOST_DEVICE Coord(Index value=Index(0))
Default ctor initializes uniformly.
Definition: coord.h:77
cutlass::Coord::operator!
CUTLASS_HOST_DEVICE bool operator!() const
Returns true if Coord is uniformly zero.
Definition: coord.h:153
cutlass::Coord::operator*=
CUTLASS_HOST_DEVICE Coord & operator*=(Coord const &b)
In-place multiplication.
Definition: coord.h:222
cutlass::make_Coord
CUTLASS_HOST_DEVICE Coord< 1 > make_Coord(int _0)
Helper to make a 2-element coordinate.
Definition: coord.h:387
cutlass::Coord::clamp
CUTLASS_HOST_DEVICE Coord & clamp(Coord const &max, Coord const &min=Coord())
Clamps a coordinate to a range specified by maximum and minimum values.
Definition: coord.h:289
cutlass::Coord::operator<
CUTLASS_HOST_DEVICE bool operator<(Coord const &b) const
Less than operator.
Definition: coord.h:318
cutlass::Coord::operator!=
CUTLASS_HOST_DEVICE bool operator!=(Coord const &b) const
Not equal.
Definition: coord.h:285
cutlass::Coord::operator-=
CUTLASS_HOST_DEVICE Coord & operator-=(Coord const &b)
In-place subtraction.
Definition: coord.h:213
cutlass::Coord::at
CUTLASS_HOST_DEVICE Index & at(int dim)
Access via index; may limit unrolling potential.
Definition: coord.h:261
cutlass::Coord::Coord
CUTLASS_HOST_DEVICE Coord(Coord< kRank, Index, LongIndex > const &coord)
Copy constructor.
Definition: coord.h:93
cutlass::Coord< 4 >::Index
int Index
Index type used to store elements.
Definition: coord.h:55
CUTLASS_PRAGMA_UNROLL
#define CUTLASS_PRAGMA_UNROLL
Definition: cutlass.h:110
__NV_STD_MAX
#define __NV_STD_MAX(a, b)
Select maximum(a, b)
Definition: platform.h:163
cutlass::Coord::product
CUTLASS_HOST_DEVICE LongIndex product() const
Returns the product of all elements.
Definition: coord.h:308
cutlass::Coord::operator[]
CUTLASS_HOST_DEVICE Index & operator[](int dim)
Member access operator.
Definition: coord.h:239
__NV_STD_MIN
#define __NV_STD_MIN(a, b)
Select minimum(a, b)
Definition: platform.h:168
cutlass::Coord::operator<=
CUTLASS_HOST_DEVICE bool operator<=(Coord const &b) const
Less than or equals operator.
Definition: coord.h:329
cutlass::Coord::operator>
CUTLASS_HOST_DEVICE bool operator>(Coord const &b) const
Greater than operator.
Definition: coord.h:340
cutlass::Coord::slice
CUTLASS_HOST_DEVICE Coord< Slice > slice(int start=0, Index identity=0) const
Definition: coord.h:103
cutlass::Coord::operator*
CUTLASS_HOST_DEVICE Coord operator*(Coord const &b) const
Element-wise multiplication.
Definition: coord.h:184
cutlass::Coord::Coord
CUTLASS_HOST_DEVICE Coord(Index const (&_idx)[kRank])
Constructs from an array of integers.
Definition: coord.h:85
CUTLASS_HOST_DEVICE
#define CUTLASS_HOST_DEVICE
Definition: cutlass.h:89
cutlass::Coord::operator/=
CUTLASS_HOST_DEVICE Coord & operator/=(Coord const &b)
In-place division.
Definition: coord.h:231
cutlass::platform::min
CUTLASS_HOST_DEVICE constexpr const T & min(const T &a, const T &b)
std::min
Definition: platform.h:183
cutlass::Coord::operator/
CUTLASS_HOST_DEVICE Coord operator/(Coord const &b) const
Element-wise division.
Definition: coord.h:194
cutlass::Coord::sum
CUTLASS_HOST_DEVICE Index sum() const
Returns the sum of all elements.
Definition: coord.h:298
cutlass::Coord
Statically-sized array specifying Coords within a tensor.
Definition: coord.h:43
cutlass::Coord::at
CUTLASS_HOST_DEVICE Index const & at(int dim) const
Access via index; may limit unrolling potential.
Definition: coord.h:271
cutlass::Coord::operator+=
CUTLASS_HOST_DEVICE Coord & operator+=(Coord const &b)
In-place addition.
Definition: coord.h:204
cutlass::Coord< 4 >::LongIndex
int64_t LongIndex
Type used to represent linear offsets.
Definition: coord.h:58
cutlass::Coord::at
CUTLASS_HOST_DEVICE Index & at()
Gets the index of a given Coord element.
Definition: coord.h:255
cutlass::Coord::dot
CUTLASS_HOST_DEVICE LongIndex dot(Coord const &b, LongIndex sum=LongIndex(0)) const
Computes the dot product with anotherCoord object.
Definition: coord.h:246
cutlass::Coord::max_dim_index
CUTLASS_HOST_DEVICE int max_dim_index() const
Returns the index of the dimension with greatest value.
Definition: coord.h:130
cutlass.h
Basic include for CUTLASS.
cutlass::Coord::operator>=
CUTLASS_HOST_DEVICE bool operator>=(Coord const &b) const
Greater than or equals operator.
Definition: coord.h:346