tensor_ref.h
Go to the documentation of this file.
/***************************************************************************************************
* Copyright (c) 2017-2018, 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 <typeinfo>
#include <cutlass/coord.h>
#include <cutlass/cutlass.h>
#include <cutlass/vector.h>
namespace cutlass {
template <typename Storage_, int Rank_>
class TensorRef {
public:
typedef Storage_ Storage;
static int const Rank = Rank_;
private:
//
// Data members
//
Storage* ptr_;
Coord<Rank> stride_;
public:
//
// Methods
//
CUTLASS_HOST_DEVICE
TensorRef() : ptr_(nullptr) {}
CUTLASS_HOST_DEVICE
TensorRef(Storage* ptr, Coord<Rank> stride) : ptr_(ptr), stride_(stride) {}
CUTLASS_HOST_DEVICE
void reset(Storage* ptr = nullptr, Coord<Rank> stride = Coord<Rank>(0)) {
ptr_ = ptr;
stride_ = stride;
}
template <typename T>
TensorRef<T, Rank> convert() {
Coord<Rank> converted_stride;
for (int i = 0; i < Rank - 1; ++i) {
converted_stride[i] = stride_[i] * Extent<Storage>::kValue / Extent<T>::kValue;
}
converted_stride[Rank - 1] = stride_[Rank - 1];
return TensorRef<T, Rank>(reinterpret_cast<T*>(ptr_), converted_stride);
}
CUTLASS_HOST_DEVICE
bool good() const { return ptr_ != nullptr; }
CUTLASS_HOST_DEVICE
Storage* data() const { return ptr_; }
CUTLASS_HOST_DEVICE
Coord<Rank> const& stride() const { return stride_; }
CUTLASS_HOST_DEVICE
int const& stride(int dim) const { return stride_.at(dim); }
CUTLASS_HOST_DEVICE
int leading_dim() const { return __NV_STD_MAX(stride_[1], stride_[2]); }
CUTLASS_HOST_DEVICE
long long offset(Coord<Rank> const& coord) const {
return stride_.template dot<long long>(coord);
}
CUTLASS_HOST_DEVICE
Storage& at(Coord<Rank> const& coord) const { return ptr_[offset(coord)]; }
Storage& operator[](Coord<Rank> const& coord) const { return at(coord); }
CUTLASS_HOST_DEVICE
Storage& at(int idx) const { return ptr_[idx]; }
Storage& operator[](int idx) const { return at(idx); }
CUTLASS_HOST_DEVICE
TensorRef& advance(Coord<Rank> const& b) {
ptr_ += offset(b);
return *this;
}
CUTLASS_HOST_DEVICE
TensorRef operator+(Coord<Rank> const& b) const { return TensorRef(ptr_ + offset(b), stride_); }
CUTLASS_HOST_DEVICE
TensorRef operator-(Coord<Rank> const& b) const { return TensorRef(ptr_ - offset(b), stride_); }
};
} // namespace cutlass
cutlass::TensorRef::stride
CUTLASS_HOST_DEVICE int const & stride(int dim) const
Returns the stride of the tensor in the given dimension.
Definition: tensor_ref.h:107
cutlass::TensorRef::operator[]
Storage & operator[](int idx) const
Element-wise accessor.
Definition: tensor_ref.h:131
cutlass
Definition: convert.h:33
cutlass::TensorRef::at
CUTLASS_HOST_DEVICE Storage & at(Coord< Rank > const &coord) const
Returns a reference to the element at a given Coord.
Definition: tensor_ref.h:121
cutlass::TensorRef::advance
CUTLASS_HOST_DEVICE TensorRef & advance(Coord< Rank > const &b)
Adds an offset to the pointer.
Definition: tensor_ref.h:135
cutlass::TensorRef::Rank
static int const Rank
Rank of tensor.
Definition: tensor_ref.h:48
cutlass::TensorRef::operator+
CUTLASS_HOST_DEVICE TensorRef operator+(Coord< Rank > const &b) const
Returns a TensorRef offset by a given amount.
Definition: tensor_ref.h:142
coord.h
A Coord is a coordinate of arbitrary rank into a tensor or matrix.
cutlass::TensorRef::Storage
Storage_ Storage
Data type of individual access.
Definition: tensor_ref.h:45
cutlass::TensorRef::operator-
CUTLASS_HOST_DEVICE TensorRef operator-(Coord< Rank > const &b) const
Returns a TensorRef offset by a given amount.
Definition: tensor_ref.h:146
__NV_STD_MAX
#define __NV_STD_MAX(a, b)
Select maximum(a, b)
Definition: platform.h:155
cutlass::TensorRef::leading_dim
CUTLASS_HOST_DEVICE int leading_dim() const
Returns the maximum stride element as the &#39;leading dimension&#39;.
Definition: tensor_ref.h:111
cutlass::TensorRef::data
CUTLASS_HOST_DEVICE Storage * data() const
Returns the pointer to referenced data.
Definition: tensor_ref.h:99
cutlass::TensorRef::TensorRef
CUTLASS_HOST_DEVICE TensorRef(Storage *ptr, Coord< Rank > stride)
Constructs from a pointer, size, and stride.
Definition: tensor_ref.h:72
cutlass::TensorRef::operator[]
Storage & operator[](Coord< Rank > const &coord) const
Element-wise accessor.
Definition: tensor_ref.h:124
nullptr
#define nullptr
nullptr
Definition: platform.h:136
cutlass::TensorRef::offset
CUTLASS_HOST_DEVICE long long offset(Coord< Rank > const &coord) const
Computes the offset of an index from the origin of the tensor.
Definition: tensor_ref.h:115
cutlass::TensorRef
Structure modeling a pointer and stride into a tensor.
Definition: tensor_ref.h:42
cutlass::TensorRef::convert
TensorRef< T, Rank > convert()
Conversion function.
Definition: tensor_ref.h:83
CUTLASS_HOST_DEVICE
#define CUTLASS_HOST_DEVICE
Definition: cutlass.h:46
cutlass::TensorRef::good
CUTLASS_HOST_DEVICE bool good() const
Returns true if the TensorRef may be safely accessed.
Definition: tensor_ref.h:95
cutlass::Coord< Rank >
vector.h
Defines a 1D vector of elements held in the registers of each thread.
cutlass::TensorRef::reset
CUTLASS_HOST_DEVICE void reset(Storage *ptr=nullptr, Coord< Rank > stride=Coord< Rank >(0))
Updates the pointer, stride, and location within a TensorRef.
Definition: tensor_ref.h:76
cutlass::Coord::at
CUTLASS_HOST_DEVICE int & at()
Gets the index of a given Coord element.
Definition: coord.h:185
cutlass::TensorRef::stride
CUTLASS_HOST_DEVICE Coord< Rank > const & stride() const
Returns the stride of the tensor.
Definition: tensor_ref.h:103
cutlass.h
Basic include for CUTLASS macros.
cutlass::TensorRef::at
CUTLASS_HOST_DEVICE Storage & at(int idx) const
Returns a reference to the element at a given Coord.
Definition: tensor_ref.h:128
cutlass::TensorRef::TensorRef
CUTLASS_HOST_DEVICE TensorRef()
Default ctor.
Definition: tensor_ref.h:68
cutlass::Extent
Returns the extent of a scalar or vector.
Definition: vector.h:161