pitch_linear.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"
#include "cutlass/coord.h"
namespace cutlass {
namespace layout {
template <
int Contiguous,
int Strided
>
struct PitchLinearShape {
static int const kContiguous = Contiguous;
static int const kStrided = Strided;
static int const kCount = Contiguous * Strided;
};
struct PitchLinearCoord : public Coord<2, int> {
public:
using Index = int;
using Base = Coord<2, Index>;
private:
static int const kContiguous = 0;
static int const kStrided = 1;
public:
//
// Methods
//
CUTLASS_HOST_DEVICE
PitchLinearCoord() { }
CUTLASS_HOST_DEVICE
PitchLinearCoord(Coord<2, Index> const &coord): Base(coord) { }
CUTLASS_HOST_DEVICE
PitchLinearCoord(Index contiguous_, Index strided_): Base(make_Coord(contiguous_, strided_)) { }
CUTLASS_HOST_DEVICE
Index const & contiguous() const { return this->at(kContiguous); }
CUTLASS_HOST_DEVICE
Index & contiguous() { return this->at(kContiguous); }
CUTLASS_HOST_DEVICE
Index const & strided() const { return this->at(kStrided); }
CUTLASS_HOST_DEVICE
Index & strided() { return this->at(kStrided); }
//
// Coord operators
//
CUTLASS_HOST_DEVICE
PitchLinearCoord operator+(Base const& b) const {
return PitchLinearCoord(Base::operator+(b));
}
CUTLASS_HOST_DEVICE
PitchLinearCoord operator-(Base const& b) const {
return PitchLinearCoord(Base::operator-(b));
}
CUTLASS_HOST_DEVICE
PitchLinearCoord operator*(Base const& b) const {
return PitchLinearCoord(Base::operator*(b));
}
CUTLASS_HOST_DEVICE
PitchLinearCoord operator/(Base const& b) const {
return PitchLinearCoord(Base::operator/(b));
}
CUTLASS_HOST_DEVICE
PitchLinearCoord& operator+=(Base const& b) {
Base::operator+=(b);
return *this;
}
CUTLASS_HOST_DEVICE
PitchLinearCoord& operator-=(Base const& b) {
Base::operator-=(b);
return *this;
}
CUTLASS_HOST_DEVICE
PitchLinearCoord& operator*=(Base const& b) {
Base::operator*=(b);
return *this;
}
CUTLASS_HOST_DEVICE
PitchLinearCoord& operator/=(Base const& b) {
Base::operator/=(b);
return *this;
}
};
class PitchLinear {
public:
static int const kRank = 2;
static int const kStrideRank = 1;
using Index = int32_t;
using LongIndex = int64_t;
using TensorCoord = PitchLinearCoord;
using Stride = Coord<kStrideRank, Index>;
private:
//
// Data members
//
Stride stride_;
public:
//
// Methods
//
CUTLASS_HOST_DEVICE
PitchLinear(Index ldm = 0): stride_(ldm) { }
CUTLASS_HOST_DEVICE
PitchLinear(Stride _stride): stride_(_stride) { }
CUTLASS_HOST_DEVICE
static PitchLinear packed(TensorCoord const &extent) {
return PitchLinear(extent.contiguous());
}
CUTLASS_HOST_DEVICE
LongIndex operator()(TensorCoord const &coord) const {
return coord.contiguous() + coord.strided() * stride_[0];
}
CUTLASS_HOST_DEVICE
TensorCoord inverse(LongIndex index) const {
return make_Coord(
Index(index % stride_[0]),
Index(index / stride_[0])
);
}
CUTLASS_HOST_DEVICE
Stride stride() const {
return stride_;
}
CUTLASS_HOST_DEVICE
Stride & stride() {
return stride_;
}
CUTLASS_HOST_DEVICE
Index stride(int rank) const {
return stride_[rank];
}
CUTLASS_HOST_DEVICE
Index & stride(int rank) {
return stride_[rank];
}
CUTLASS_HOST_DEVICE
LongIndex capacity(TensorCoord const &extent) const {
return extent.strided() * stride_[0];
}
};
} // namespace layout
} // namespace cutlass
cutlass::layout::PitchLinearCoord::Index
int Index
Integer-valued index.
Definition: pitch_linear.h:56
cutlass::layout::PitchLinearShape::kCount
static int const kCount
Definition: pitch_linear.h:46
cutlass
Definition: aligned_buffer.h:35
cutlass::layout::PitchLinearCoord
Coordinate in pitch-linear space.
Definition: pitch_linear.h:52
cutlass::layout::PitchLinearCoord::operator-
CUTLASS_HOST_DEVICE PitchLinearCoord operator-(Base const &b) const
Element-wise subtraction.
Definition: pitch_linear.h:115
cutlass::layout::PitchLinear
Mapping function for pitch-linear memory.
Definition: pitch_linear.h:163
coord.h
A Coord is a coordinate of arbitrary rank into a tensor or matrix.
cutlass::make_Coord
CUTLASS_HOST_DEVICE Coord< 1 > make_Coord(int _0)
Helper to make a 2-element coordinate.
Definition: coord.h:387
cutlass::layout::PitchLinear::stride
CUTLASS_HOST_DEVICE Stride stride() const
Returns the stride of the layout.
Definition: pitch_linear.h:228
cutlass::layout::PitchLinear::packed
static CUTLASS_HOST_DEVICE PitchLinear packed(TensorCoord const &extent)
Helper returns a layout to a tightly packed tensor.
Definition: pitch_linear.h:206
cutlass::operator/=
CUTLASS_HOST_DEVICE half_t & operator/=(half_t &lhs, half_t const &rhs)
Definition: half.h:684
cutlass::operator+=
CUTLASS_HOST_DEVICE half_t & operator+=(half_t &lhs, half_t const &rhs)
Definition: half.h:654
cutlass::layout::PitchLinearShape
Template defining a shape used by pitch-linear operators.
Definition: pitch_linear.h:43
cutlass::operator-=
CUTLASS_HOST_DEVICE half_t & operator-=(half_t &lhs, half_t const &rhs)
Definition: half.h:664
cutlass::layout::PitchLinear::LongIndex
int64_t LongIndex
Long index type used for offsets.
Definition: pitch_linear.h:175
cutlass::layout::PitchLinearCoord::operator+=
CUTLASS_HOST_DEVICE PitchLinearCoord & operator+=(Base const &b)
In-place addition.
Definition: pitch_linear.h:133
cutlass::layout::PitchLinear::PitchLinear
CUTLASS_HOST_DEVICE PitchLinear(Stride _stride)
Constructor.
Definition: pitch_linear.h:202
cutlass::layout::PitchLinearShape::kStrided
static int const kStrided
Definition: pitch_linear.h:45
cutlass::layout::PitchLinearCoord::operator*=
CUTLASS_HOST_DEVICE PitchLinearCoord & operator*=(Base const &b)
In-place multiplication.
Definition: pitch_linear.h:147
cutlass::layout::PitchLinearShape::kContiguous
static int const kContiguous
Definition: pitch_linear.h:44
cutlass::layout::PitchLinear::inverse
CUTLASS_HOST_DEVICE TensorCoord inverse(LongIndex index) const
Returns the logical coordinate given an offset.
Definition: pitch_linear.h:219
cutlass::layout::PitchLinear::stride
CUTLASS_HOST_DEVICE Index & stride(int rank)
Returns the stride of the layout.
Definition: pitch_linear.h:246
cutlass::layout::PitchLinear::stride
CUTLASS_HOST_DEVICE Stride & stride()
Returns the stride of the layout.
Definition: pitch_linear.h:234
CUTLASS_HOST_DEVICE
#define CUTLASS_HOST_DEVICE
Definition: cutlass.h:89
cutlass::layout::PitchLinearCoord::contiguous
CUTLASS_HOST_DEVICE Index const & contiguous() const
Returns the contiguous dimension.
Definition: pitch_linear.h:89
cutlass::layout::PitchLinearCoord::PitchLinearCoord
CUTLASS_HOST_DEVICE PitchLinearCoord(Index contiguous_, Index strided_)
Helper to construct from a row and column.
Definition: pitch_linear.h:85
cutlass::layout::PitchLinear::Index
int32_t Index
Index type used for coordinates.
Definition: pitch_linear.h:172
cutlass::operator*=
CUTLASS_HOST_DEVICE half_t & operator*=(half_t &lhs, half_t const &rhs)
Definition: half.h:674
cutlass::Coord
Statically-sized array specifying Coords within a tensor.
Definition: coord.h:43
cutlass::layout::PitchLinear::operator()
CUTLASS_HOST_DEVICE LongIndex operator()(TensorCoord const &coord) const
Definition: pitch_linear.h:213
cutlass::layout::PitchLinearCoord::operator/=
CUTLASS_HOST_DEVICE PitchLinearCoord & operator/=(Base const &b)
In-place division.
Definition: pitch_linear.h:154
cutlass::layout::PitchLinear::stride
CUTLASS_HOST_DEVICE Index stride(int rank) const
Returns the stride of the layout.
Definition: pitch_linear.h:240
cutlass::layout::PitchLinear::PitchLinear
CUTLASS_HOST_DEVICE PitchLinear(Index ldm=0)
Constructor.
Definition: pitch_linear.h:198
cutlass::layout::PitchLinearCoord::strided
CUTLASS_HOST_DEVICE Index & strided()
Returns the column of the coordinate.
Definition: pitch_linear.h:101
cutlass::layout::PitchLinearCoord::operator/
CUTLASS_HOST_DEVICE PitchLinearCoord operator/(Base const &b) const
Element-wise division.
Definition: pitch_linear.h:127
cutlass::layout::PitchLinearCoord::operator*
CUTLASS_HOST_DEVICE PitchLinearCoord operator*(Base const &b) const
Element-wise multiplication.
Definition: pitch_linear.h:121
cutlass::layout::PitchLinearCoord::PitchLinearCoord
CUTLASS_HOST_DEVICE PitchLinearCoord(Coord< 2, Index > const &coord)
Constructs from Coord<2>
Definition: pitch_linear.h:81
cutlass::layout::PitchLinearCoord::operator+
CUTLASS_HOST_DEVICE PitchLinearCoord operator+(Base const &b) const
Element-wise addition.
Definition: pitch_linear.h:109
cutlass::layout::PitchLinearCoord::contiguous
CUTLASS_HOST_DEVICE Index & contiguous()
Returns the contiguous dimension.
Definition: pitch_linear.h:93
cutlass::layout::PitchLinearCoord::PitchLinearCoord
CUTLASS_HOST_DEVICE PitchLinearCoord()
Default ctor.
Definition: pitch_linear.h:77
cutlass.h
Basic include for CUTLASS.
cutlass::layout::PitchLinearCoord::operator-=
CUTLASS_HOST_DEVICE PitchLinearCoord & operator-=(Base const &b)
In-place subtraction.
Definition: pitch_linear.h:140
cutlass::layout::PitchLinearCoord::strided
CUTLASS_HOST_DEVICE Index const & strided() const
Returns the column of the coordinate.
Definition: pitch_linear.h:97
cutlass::layout::PitchLinear::capacity
CUTLASS_HOST_DEVICE LongIndex capacity(TensorCoord const &extent) const
Compute the number of contiguous elements needed to store a tensor with the given size...
Definition: pitch_linear.h:252