tensor_view.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
#if !defined(__CUDACC_RTC__)
#include <cmath>
#endif
#include "cutlass/cutlass.h"
#include "cutlass/tensor_ref.h"
namespace cutlass {
template <
typename Element_,
typename Layout_
>
class TensorView : public TensorRef<Element_, Layout_> {
public:
using Base = cutlass::TensorRef<Element_, Layout_>;
using Layout = Layout_;
using ConstTensorRef = typename Base::ConstTensorRef;
using TensorRef = Base;
using Element = Element_;
using Reference = Element &;
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 ConstTensorView = TensorView<
typename platform::remove_const<Element>::type const,
Layout>;
using NonConstTensorView = TensorView<
typename platform::remove_const<Element>::type,
Layout>;
static_assert(kRank > 0, "Cannot define a zero-rank TensorRef");
private:
TensorCoord extent_;
public:
//
// Methods
//
CUTLASS_HOST_DEVICE
TensorView(TensorCoord const &extent = TensorCoord()): extent_(extent) {
}
CUTLASS_HOST_DEVICE
TensorView(
Element *ptr,
Layout const &layout,
TensorCoord const &extent
):
Base(ptr, layout), extent_(extent) {
}
CUTLASS_HOST_DEVICE
TensorView(
TensorRef const &ref,
TensorCoord const &extent
):
Base(ref), extent_(extent) {
}
CUTLASS_HOST_DEVICE
TensorView(
NonConstTensorView const &view
):
Base(view), extent_(view.extent_) { }
CUTLASS_HOST_DEVICE
void reset(Element* ptr, Layout const &layout, TensorCoord size) {
Base::reset(ptr, layout);
this->resize(extent_);
}
CUTLASS_HOST_DEVICE
void resize(TensorCoord extent) {
this->extent_ = extent;
}
CUTLASS_HOST_DEVICE
TensorCoord const& extent() const { return extent_; }
CUTLASS_HOST_DEVICE
Index extent(int dim) const { return extent_.at(dim); }
CUTLASS_HOST_DEVICE
bool contains(TensorCoord const& coord) const {
CUTLASS_PRAGMA_UNROLL
for (int dim = 0; dim < kRank; ++dim) {
if (!(coord[dim] >= 0 && coord[dim] < extent(dim))) {
return false;
}
}
return true;
}
CUTLASS_HOST_DEVICE
TensorRef ref() const {
return TensorRef(this->data(), this->layout());
}
CUTLASS_HOST_DEVICE
ConstTensorRef const_ref() const {
return ConstTensorRef(this->data(), this->layout());
}
CUTLASS_HOST_DEVICE
ConstTensorView const_view() const {
return ConstTensorView(const_ref(), extent_);
}
CUTLASS_HOST_DEVICE
TensorView subview(
TensorCoord extent,
TensorCoord const& location = TensorCoord()
) const {
return TensorView(ref(), extent.clamp(extent_ - location)).add_coord_offset(location);
}
CUTLASS_HOST_DEVICE
size_t capacity() const {
return Base::layout().capacity(extent_);
}
CUTLASS_HOST_DEVICE
TensorView operator+(
TensorCoord const& b
) const {
TensorView result(*this);
result.add_pointer_offset(this->offset(b));
return result;
}
CUTLASS_HOST_DEVICE
TensorView& operator+=(
TensorCoord const& b
) {
this->add_pointer_offset(this->offset(b));
return *this;
}
CUTLASS_HOST_DEVICE
TensorView operator-(
TensorCoord const& b
) const {
TensorRef result(*this);
result.add_pointer_offset(-this->offset(b));
return result;
}
CUTLASS_HOST_DEVICE
TensorView& operator-=(
TensorCoord const& b
) {
this->add_pointer_offset(-this->offset(b));
return *this;
}
};
template <
typename Element,
typename Layout
>
CUTLASS_HOST_DEVICE TensorView<Element, Layout> make_TensorView(
Element *ptr,
Layout const &layout,
typename Layout::TensorCoord const &extent) {
return TensorView<Element, Layout>(ptr, layout, extent);
}
} // namespace cutlass
cutlass::TensorView< Element, Layout >::operator+=
CUTLASS_HOST_DEVICE TensorView & operator+=(TensorCoord const &b)
Returns a TensorRef offset by a given amount.
Definition: tensor_view.h:232
cutlass
Definition: aligned_buffer.h:35
cutlass::TensorView< Element, Layout >::capacity
CUTLASS_HOST_DEVICE size_t capacity() const
Returns the number of scalar elements needed to store tensor.
Definition: tensor_view.h:215
tensor_ref.h
Defines a structure containing strides, bounds, and a pointer to tensor data.
cutlass::platform::remove_const::type
T type
Definition: platform.h:351
cutlass::TensorRef::data
CUTLASS_HOST_DEVICE Element * data() const
Returns the pointer to referenced data.
Definition: tensor_ref.h:254
cutlass::TensorView::extent
CUTLASS_HOST_DEVICE TensorCoord const & extent() const
Returns the extent of the view (the size along each logical dimension).
Definition: tensor_view.h:167
cutlass::TensorView::kRank
static int const kRank
Logical rank of tensor index space.
Definition: tensor_view.h:78
cutlass::TensorView::resize
CUTLASS_HOST_DEVICE void resize(TensorCoord extent)
Changes the size of the view without affecting pointer or layout.
Definition: tensor_view.h:161
cutlass::TensorView< Element, Layout >::operator+
CUTLASS_HOST_DEVICE TensorView operator+(TensorCoord const &b) const
Returns a TensorView offset by a given amount.
Definition: tensor_view.h:221
cutlass::TensorView< Element, Layout >::operator-=
CUTLASS_HOST_DEVICE TensorView & operator-=(TensorCoord const &b)
Returns a TensorRef offset by a given amount.
Definition: tensor_view.h:253
cutlass::TensorView< Element, Layout >::operator-
CUTLASS_HOST_DEVICE TensorView operator-(TensorCoord const &b) const
Returns a TensorRef offset by a given amount.
Definition: tensor_view.h:242
cutlass::TensorView::TensorRef
Base TensorRef
Underlying TensorRef type.
Definition: tensor_view.h:69
cutlass::TensorRef::ConstTensorRef
TensorRef< typename platform::remove_const< Element >::type const, Layout > ConstTensorRef
TensorRef to constant data.
Definition: tensor_ref.h:179
cutlass::TensorRef::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::TensorView< Element, Layout >< Element, Layout >::Element
Element Element
Data type of individual access.
Definition: tensor_view.h:72
CUTLASS_PRAGMA_UNROLL
#define CUTLASS_PRAGMA_UNROLL
Definition: cutlass.h:110
cutlass::TensorView::ConstTensorView
TensorView< typename platform::remove_const< Element >::type const, Layout > ConstTensorView
TensorView pointing to constant memory.
Definition: tensor_view.h:95
cutlass::TensorView
Definition: tensor_view.h:56
cutlass::TensorView< Element, Layout >::reset
CUTLASS_HOST_DEVICE void reset(Element *ptr, Layout const &layout, TensorCoord size)
Updates the pointer and layout object.
Definition: tensor_view.h:154
cutlass::TensorView< Element, Layout >< Element, Layout >::TensorCoord
typename Layout::TensorCoord TensorCoord
Coordinate in logical tensor space.
Definition: tensor_view.h:87
cutlass::TensorView< Element, Layout >< Element, Layout >::Reference
Element & Reference
Reference type to an element.
Definition: tensor_view.h:75
cutlass::TensorRef::reset
CUTLASS_HOST_DEVICE void reset(Element *ptr=nullptr)
Updates only the pointer.
Definition: tensor_ref.h:235
cutlass::TensorRef
Definition: tensor_ref.h:146
cutlass::TensorView::ref
CUTLASS_HOST_DEVICE TensorRef ref() const
Returns a TensorRef pointing to the first element of the tensor.
Definition: tensor_view.h:187
cutlass::TensorView< Element, Layout >< Element, Layout >::Stride
typename Layout::Stride Stride
Coordinate in storage n-D array.
Definition: tensor_view.h:90
CUTLASS_HOST_DEVICE
#define CUTLASS_HOST_DEVICE
Definition: cutlass.h:89
cutlass::TensorRef::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
static_assert
#define static_assert(__e, __m)
Definition: platform.h:153
cutlass::TensorView< Element, Layout >::extent
CUTLASS_HOST_DEVICE Index extent(int dim) const
Returns the extent along a particular logical dimension.
Definition: tensor_view.h:171
cutlass::TensorView< Element, Layout >::subview
CUTLASS_HOST_DEVICE TensorView subview(TensorCoord extent, TensorCoord const &location=TensorCoord()) const
Returns a Tensor_view given location and size quantities.
Definition: tensor_view.h:205
cutlass::TensorView< Element, Layout >::TensorView
CUTLASS_HOST_DEVICE TensorView(TensorRef const &ref, TensorCoord const &extent)
Constructs a TensorView object.
Definition: tensor_view.h:137
cutlass::TensorView< Element, Layout >::TensorView
CUTLASS_HOST_DEVICE TensorView(NonConstTensorView const &view)
Converting constructor from TensorRef to non-constant data.
Definition: tensor_view.h:147
cutlass::TensorView< Element, Layout >< Element, Layout >::Index
typename Layout::Index Index
Index type.
Definition: tensor_view.h:81
cutlass::TensorView::Base
cutlass::TensorRef< Element_, Layout_ > Base
Base tensor reference.
Definition: tensor_view.h:60
cutlass::TensorView< Element, Layout >::const_view
CUTLASS_HOST_DEVICE ConstTensorView const_view() const
Returns a TensorView to const data.
Definition: tensor_view.h:199
cutlass::TensorView< Element, Layout >< Element, Layout >::ConstTensorRef
typename Base::ConstTensorRef ConstTensorRef
TensorRef pointing to constant memory.
Definition: tensor_view.h:66
cutlass::TensorRef::layout
CUTLASS_HOST_DEVICE Layout & layout()
Returns the layout object.
Definition: tensor_ref.h:265
cutlass::TensorView< Element, Layout >::TensorView
CUTLASS_HOST_DEVICE TensorView(Element *ptr, Layout const &layout, TensorCoord const &extent)
Constructs a TensorView object.
Definition: tensor_view.h:126
cutlass::TensorView< Element, Layout >::TensorView
CUTLASS_HOST_DEVICE TensorView(TensorCoord const &extent=TensorCoord())
Constructs a TensorView object.
Definition: tensor_view.h:120
cutlass::make_TensorView
CUTLASS_HOST_DEVICE TensorView< Element, Layout > make_TensorView(Element *ptr, Layout const &layout, typename Layout::TensorCoord const &extent)
Constructs a TensorRef, deducing types from arguments.
Definition: tensor_view.h:269
cutlass::TensorView< Element, Layout >::contains
CUTLASS_HOST_DEVICE bool contains(TensorCoord const &coord) const
Determines whether a location is within a tensor.
Definition: tensor_view.h:175
cutlass::TensorRef::add_pointer_offset
CUTLASS_HOST_DEVICE TensorRef & add_pointer_offset(LongIndex offset_)
Adds an offset to each pointer.
Definition: tensor_ref.h:319
cutlass.h
Basic include for CUTLASS.
cutlass::TensorView< Element, Layout >::const_ref
CUTLASS_HOST_DEVICE ConstTensorRef const_ref() const
Returns a TensorRef pointing to the first element of the tensor.
Definition: tensor_view.h:193
cutlass::TensorView< Element, Layout >< Element, Layout >::LongIndex
typename Layout::LongIndex LongIndex
Long index used for pointer offsets.
Definition: tensor_view.h:84
cutlass::TensorView< Element, Layout >< Element, Layout >::Layout
Layout Layout
Mapping function from logical coordinate to internal n-D array.
Definition: tensor_view.h:63