tensor_ref.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"
#include "cutlass/platform/platform.h"
#include "cutlass/subbyte_reference.h"
namespace cutlass {
template <int Rank>
class IdentityTensorLayout {
public:
static int const kRank = Rank;
static int const kStrideRank = Rank;
using Index = int32_t;
using LongIndex = int64_t;
using TensorCoord = Coord<kRank, Index>;
using Stride = Coord<kStrideRank, Index>;
private:
//
// Data members
//
Stride stride_;
public:
//
// Methods
//
CUTLASS_HOST_DEVICE
IdentityTensorLayout(Stride const &stride = Stride()): stride_(stride) { }
CUTLASS_HOST_DEVICE
LongIndex operator()(Coord<Rank> const &coord) const {
return coord.dot(stride_);
}
CUTLASS_HOST_DEVICE
Stride stride() const {
return stride_;
}
CUTLASS_HOST_DEVICE
Stride & stride() {
return stride_;
}
CUTLASS_HOST_DEVICE
LongIndex capacity(TensorCoord const &size) const {
int idx = stride_.max_dim_index();
return stride_[idx] * size[idx];
}
};
/* \brief TensorRef is a template for objects pointing to the start of tensors of arbitrary rank
and layout within memory. A TensorRef combines a pointer and a Layout concept
Examples:
(These examples use helpers for matrix layouts defined in cutlass/layout/matrix.h)
1. Column-major matrix may be represented as a rank=2 tensor:
TensorRef<float, layout::ColumnMajor> A(ptr_A, ldm);
2. Row-major matrix may be represented as a rank=2 tensor:
TensorRef<float, layout::RowMajor> B(ptr_A, ldm);
3. An interleaved matrix may be represented as a rank=2 tensor:
TensorRef<int8_t, layout::ColumnMajorInterleaved<32> > C;
4. A helper exists to define a TensorRef for a contiguous matrix whose layout
is not known at compile time.
int ldm; // leading dimension
layout::Matrix kind; // Could be layout::Matrix::kRowMajor or layout::Matrix::kColumnMajor
TensorRef<int, layout::ContiguousMatrix> E(ptr_E, {ldm, kind});
*/
template <
typename Element_,
typename Layout_
>
class TensorRef {
public:
using Element = Element_;
using Layout = Layout_;
using Reference = typename platform::conditional<
sizeof_bits<Element>::value >= 8,
Element &,
SubbyteReference<Element>
>::type;
static int const kRank = Layout::kRank;
using Index = typename Layout::Index;
using LongIndex = typename Layout::LongIndex;
using TensorCoord = typename Layout::TensorCoord;
using Stride = typename Layout::Stride;
using ConstTensorRef = TensorRef<
typename platform::remove_const<Element>::type const,
Layout>;
using NonConstTensorRef = TensorRef<
typename platform::remove_const<Element>::type,
Layout>;
static_assert(kRank > 0, "Cannot define a zero-rank TensorRef");
private:
Element* ptr_;
Layout layout_;
public:
//
// Methods
//
CUTLASS_HOST_DEVICE
TensorRef(
Element *ptr = nullptr,
Layout const &layout = Layout()
):
ptr_(ptr), layout_(layout) {
}
CUTLASS_HOST_DEVICE
TensorRef(
NonConstTensorRef const &ref
):
ptr_(ref.data()), layout_(ref.layout()) { }
CUTLASS_HOST_DEVICE
ConstTensorRef const_ref() const {
return ConstTensorRef(ptr_, layout_);
}
CUTLASS_HOST_DEVICE
NonConstTensorRef non_const_ref() const {
return NonConstTensorRef(const_cast<typename platform::remove_const<Element>::type *>(ptr_), layout_);
}
CUTLASS_HOST_DEVICE
void reset(Element* ptr = nullptr) {
ptr_ = ptr;
}
CUTLASS_HOST_DEVICE
void reset(Element* ptr, Layout const &layout) {
ptr_ = ptr;
layout_ = layout;
}
CUTLASS_HOST_DEVICE
bool good() const {
return ptr_ != nullptr;
}
CUTLASS_HOST_DEVICE
Element * data() const { return ptr_; }
CUTLASS_HOST_DEVICE
Reference data(LongIndex idx) const {
return ReferenceFactory<typename platform::remove_const<Element>::type,
(sizeof_bits<Element>::value < 8)>::get(ptr_, idx);
}
CUTLASS_HOST_DEVICE
Layout & layout() {
return layout_;
}
CUTLASS_HOST_DEVICE
Layout layout() const {
return layout_;
}
CUTLASS_HOST_DEVICE
Stride stride() const {
return layout_.stride();
}
CUTLASS_HOST_DEVICE
Stride & stride() {
return layout_.stride();
}
CUTLASS_HOST_DEVICE
Index stride(int dim) const {
return layout_.stride().at(dim);
}
CUTLASS_HOST_DEVICE
Index & stride(int dim) {
return layout_.stride().at(dim);
}
CUTLASS_HOST_DEVICE
LongIndex offset(TensorCoord const& coord) const {
return layout_(coord);
}
CUTLASS_HOST_DEVICE
Reference at(TensorCoord const& coord) const {
return data(offset(coord));
}
CUTLASS_HOST_DEVICE
Reference operator[](TensorCoord const& coord) const {
return data(offset(coord));
}
CUTLASS_HOST_DEVICE
TensorRef & add_pointer_offset(LongIndex offset_) {
ptr_ += offset_;
return *this;
}
CUTLASS_HOST_DEVICE
TensorRef & add_coord_offset(TensorCoord const &coord) {
add_pointer_offset(offset(coord));
return *this;
}
CUTLASS_HOST_DEVICE
TensorRef operator+(TensorCoord const& b) const {
TensorRef result(*this);
result.add_coord_offset(b);
return result;
}
CUTLASS_HOST_DEVICE
TensorRef & operator+=(TensorCoord const& b) {
add_coord_offset(b);
return *this;
}
CUTLASS_HOST_DEVICE
TensorRef operator-(TensorCoord const& b) const {
TensorRef result(*this);
result.add_pointer_offset(-offset(b));
return result;
}
CUTLASS_HOST_DEVICE
TensorRef & operator-=(TensorCoord const& b) {
add_pointer_offset(-offset(b));
return *this;
}
};
template <
typename Element,
typename Layout
>
CUTLASS_HOST_DEVICE
TensorRef<Element, Layout> make_TensorRef(Element *ptr, Layout const &layout) {
return TensorRef<Element, Layout>(ptr, layout);
}
//
// Partial specializations to handle degenerate and sub-byte cases.
//
template <
typename Element,
typename Layout
>
bool TensorRef_aligned(TensorRef<Element, Layout> const &ref, int alignment) {
int const kStrideRank = Layout::kStrideRank;
if (reinterpret_cast<uintptr_t>(ref.data()) % alignment) {
return false;
}
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < kStrideRank; ++i) {
if (ref.stride(i) % alignment) {
return false;
}
}
return true;
}
} // namespace cutlass
cutlass::IdentityTensorLayout::stride
CUTLASS_HOST_DEVICE Stride stride() const
Returns the stride of the layout.
Definition: tensor_ref.h:91
cutlass
Definition: aligned_buffer.h:35
cutlass::IdentityTensorLayout::operator()
CUTLASS_HOST_DEVICE LongIndex operator()(Coord< Rank > const &coord) const
Returns the offset of a coordinate in linear memory.
Definition: tensor_ref.h:85
cutlass::TensorRef< Element, Layout >::stride
CUTLASS_HOST_DEVICE Index stride(int dim) const
Returns the layout object&#39;s stride in a given physical dimension.
Definition: tensor_ref.h:289
cutlass::TensorRef< Element, Layout >< Element, Layout >::Stride
typename Layout::Stride Stride
Layout&#39;s stride vector.
Definition: tensor_ref.h:174
cutlass::TensorRef< Element, Layout >::stride
CUTLASS_HOST_DEVICE Index & stride(int dim)
Returns the layout object&#39;s stride in a given physical dimension.
Definition: tensor_ref.h:295
cutlass::TensorRef< Element, Layout >::stride
CUTLASS_HOST_DEVICE Stride & stride()
Returns the layout object&#39;s stride vector.
Definition: tensor_ref.h:283
cutlass::platform::remove_const::type
T type
Definition: platform.h:351
cutlass::IdentityTensorLayout
Definition: tensor_ref.h:45
cutlass::TensorRef< Element, Layout >::data
CUTLASS_HOST_DEVICE Element * data() const
Returns the pointer to referenced data.
Definition: tensor_ref.h:254
cutlass::IdentityTensorLayout::Stride
Coord< kStrideRank, Index > Stride
Stride vector.
Definition: tensor_ref.h:63
cutlass::TensorRef< Element, Layout >::const_ref
CUTLASS_HOST_DEVICE ConstTensorRef const_ref() const
Returns a reference to constant-valued tensor.
Definition: tensor_ref.h:224
cutlass::IdentityTensorLayout::Index
int32_t Index
Index type used for coordinates.
Definition: tensor_ref.h:54
coord.h
A Coord is a coordinate of arbitrary rank into a tensor or matrix.
cutlass::TensorRef< Element, Layout >< Element, Layout >::Layout
Layout Layout
Mapping function from logical coordinate to linear memory.
Definition: tensor_ref.h:152
cutlass::IdentityTensorLayout::stride
CUTLASS_HOST_DEVICE Stride & stride()
Returns the stride of the layout.
Definition: tensor_ref.h:97
cutlass::TensorRef< Element, Layout >::reset
CUTLASS_HOST_DEVICE void reset(Element *ptr, Layout const &layout)
Updates the pointer and layout object.
Definition: tensor_ref.h:241
cutlass::TensorRef< Element, Layout >::operator-=
CUTLASS_HOST_DEVICE TensorRef & operator-=(TensorCoord const &b)
Returns a TensorRef offset by a given amount.
Definition: tensor_ref.h:356
cutlass::TensorRef< Element, Layout >::operator[]
CUTLASS_HOST_DEVICE Reference operator[](TensorCoord const &coord) const
Returns a reference to the element at a given Coord.
Definition: tensor_ref.h:313
platform.h
C++ features that may be otherwise unimplemented for CUDA device functions.
cutlass::TensorRef< Element, Layout >::TensorRef
CUTLASS_HOST_DEVICE TensorRef(Element *ptr=nullptr, Layout const &layout=Layout())
Constructs a TensorRef with a pointer and layout object.
Definition: tensor_ref.h:207
cutlass::TensorRef< Element, Layout >::layout
CUTLASS_HOST_DEVICE Layout layout() const
Returns the layout object.
Definition: tensor_ref.h:271
cutlass::TensorRef< Element, Layout >::add_coord_offset
CUTLASS_HOST_DEVICE TensorRef & add_coord_offset(TensorCoord const &coord)
Adds an offset to each pointer.
Definition: tensor_ref.h:326
cutlass::TensorRef< Element, Layout >< Element, Layout >::Element
Element Element
Data type of individual access.
Definition: tensor_ref.h:149
CUTLASS_PRAGMA_UNROLL
#define CUTLASS_PRAGMA_UNROLL
Definition: cutlass.h:110
cutlass::TensorRef< Element, Layout >::data
CUTLASS_HOST_DEVICE Reference data(LongIndex idx) const
Returns a reference to the element at a given linear index.
Definition: tensor_ref.h:258
cutlass::TensorRef< Element, Layout >::operator-
CUTLASS_HOST_DEVICE TensorRef operator-(TensorCoord const &b) const
Returns a TensorRef offset by a given amount.
Definition: tensor_ref.h:348
cutlass::TensorRef< Element, Layout >::TensorRef
CUTLASS_HOST_DEVICE TensorRef(NonConstTensorRef const &ref)
Converting constructor from TensorRef to non-constant data.
Definition: tensor_ref.h:217
cutlass::TensorRef< Element, Layout >::stride
CUTLASS_HOST_DEVICE Stride stride() const
Returns the layout object&#39;s stride vector.
Definition: tensor_ref.h:277
cutlass::TensorRef< Element, Layout >< Element, Layout >::TensorCoord
typename Layout::TensorCoord TensorCoord
Coordinate in logical tensor space.
Definition: tensor_ref.h:171
cutlass::TensorRef< Element, Layout >::good
CUTLASS_HOST_DEVICE bool good() const
Returns true if the TensorRef is non-null.
Definition: tensor_ref.h:248
cutlass::sizeof_bits
Defines the size of an element in bits.
Definition: numeric_types.h:42
cutlass::TensorRef< Element, Layout >::reset
CUTLASS_HOST_DEVICE void reset(Element *ptr=nullptr)
Updates only the pointer.
Definition: tensor_ref.h:235
cutlass::ReferenceFactory
Definition: subbyte_reference.h:557
cutlass::TensorRef
Definition: tensor_ref.h:146
cutlass::TensorRef< Element, Layout >< Element, Layout >::Reference
typename platform::conditional< sizeof_bits< Element >::value >=8, Element &, SubbyteReference< Element > >::type Reference
Reference type to an element.
Definition: tensor_ref.h:159
CUTLASS_HOST_DEVICE
#define CUTLASS_HOST_DEVICE
Definition: cutlass.h:89
cutlass::TensorRef< Element, Layout >::offset
CUTLASS_HOST_DEVICE LongIndex offset(TensorCoord const &coord) const
Computes the offset of an index from the origin of the tensor.
Definition: tensor_ref.h:301
cutlass::platform::conditional
std::conditional (true specialization)
Definition: platform.h:325
static_assert
#define static_assert(__e, __m)
Definition: platform.h:153
cutlass::IdentityTensorLayout::kRank
static int const kRank
Logical rank of tensor.
Definition: tensor_ref.h:48
cutlass::TensorRef< Element, Layout >::non_const_ref
CUTLASS_HOST_DEVICE NonConstTensorRef non_const_ref() const
Definition: tensor_ref.h:229
cutlass::Coord
Statically-sized array specifying Coords within a tensor.
Definition: coord.h:43
cutlass::IdentityTensorLayout::LongIndex
int64_t LongIndex
Long index type used for offsets.
Definition: tensor_ref.h:57
cutlass::make_TensorRef
CUTLASS_HOST_DEVICE TensorRef< Element, Layout > make_TensorRef(Element *ptr, Layout const &layout)
Constructs a TensorRef, deducing types from arguments.
Definition: tensor_ref.h:368
cutlass::TensorRef< Element, Layout >< Element, Layout >::Index
typename Layout::Index Index
Index type.
Definition: tensor_ref.h:165
cutlass::IdentityTensorLayout::IdentityTensorLayout
CUTLASS_HOST_DEVICE IdentityTensorLayout(Stride const &stride=Stride())
Definition: tensor_ref.h:81
cutlass::TensorRef< Element, Layout >::at
CUTLASS_HOST_DEVICE Reference at(TensorCoord const &coord) const
Returns a reference to the element at a given Coord.
Definition: tensor_ref.h:307
cutlass::SubbyteReference
Definition: subbyte_reference.h:294
cutlass::TensorRef< Element, Layout >::layout
CUTLASS_HOST_DEVICE Layout & layout()
Returns the layout object.
Definition: tensor_ref.h:265
cutlass::TensorRef_aligned
bool TensorRef_aligned(TensorRef< Element, Layout > const &ref, int alignment)
Definition: tensor_ref.h:382
cutlass::TensorRef< Element, Layout >::operator+
CUTLASS_HOST_DEVICE TensorRef operator+(TensorCoord const &b) const
Returns a TensorRef offset by a given amount.
Definition: tensor_ref.h:333
cutlass::TensorRef< Element, Layout >::operator+=
CUTLASS_HOST_DEVICE TensorRef & operator+=(TensorCoord const &b)
Returns a TensorRef offset by a given amount.
Definition: tensor_ref.h:341
cutlass::TensorRef< Element, Layout >::add_pointer_offset
CUTLASS_HOST_DEVICE TensorRef & add_pointer_offset(LongIndex offset_)
Adds an offset to each pointer.
Definition: tensor_ref.h:319
subbyte_reference.h
Provides a mechanism for packing and unpacking elements smaller than one byte.
cutlass::IdentityTensorLayout::kStrideRank
static int const kStrideRank
Rank of stride vector.
Definition: tensor_ref.h:51
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::TensorRef< Element, Layout >< Element, Layout >::LongIndex
typename Layout::LongIndex LongIndex
Long index used for pointer offsets.
Definition: tensor_ref.h:168
cutlass::IdentityTensorLayout::capacity
CUTLASS_HOST_DEVICE LongIndex capacity(TensorCoord const &size) const
Compute the number of contiguous elements needed to store a tensor with the given size...
Definition: tensor_ref.h:103