tensor_copy.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.
*
**************************************************************************************************/
/* \file
\brief Defines host-side elementwise operations on TensorView.
*/
#pragma once
// Standard Library includes
#include <utility>
// Cutlass includes
#include "cutlass/cutlass.h"
#include "tensor_foreach.h"
namespace cutlass {
namespace reference {
namespace host {
namespace detail {
template <
typename DstElement,
typename SrcElement
>
struct TrivialConvert {
TrivialConvert() { }
DstElement operator()(SrcElement src) const {
return DstElement(src);
}
};
template <
typename DstElement,
typename DstLayout,
typename SrcElement,
typename SrcLayout,
typename F
>
struct TensorCopyIf {
using DstTensorView = TensorView<DstElement, DstLayout>;
using SrcTensorView = TensorView<SrcElement, SrcLayout>;
//
// Data members
//
DstTensorView dst;
SrcTensorView src;
F convert;
//
// Methods
//
TensorCopyIf() { }
TensorCopyIf(
DstTensorView const &dst_,
SrcTensorView const &src_,
F const &convert_): dst(dst_), src(src_), convert(convert_) {}
void operator()(Coord<DstLayout::kRank> const &coord) {
if (dst.contains(coord) && src.contains(coord)) {
dst.at(coord) = convert(src.at(coord));
}
}
};
} // namespace detail
template <
typename DstElement,
typename DstLayout,
typename SrcElement,
typename SrcLayout,
typename F
>
void TensorCopy(
TensorView<DstElement, DstLayout> dst,
TensorView<SrcElement, SrcLayout> src,
F const &transform) {
using CopyIf = detail::TensorCopyIf<
DstElement,
DstLayout,
SrcElement,
SrcLayout,
F>;
CopyIf copy_if(dst, src, transform);
TensorForEach(dst.extent(), copy_if);
}
template <
typename DstElement,
typename DstLayout,
typename SrcElement,
typename SrcLayout,
typename F
>
void TensorCopy(
TensorView<DstElement, DstLayout> dst,
TensorRef<SrcElement, SrcLayout> src,
F const &transform) {
using CopyIf = detail::TensorCopyIf<
DstElement,
DstLayout,
SrcElement,
SrcLayout,
F>;
TensorView<SrcElement, SrcLayout> src_view(src, dst.extent());
CopyIf copy_if(dst, src_view, transform);
TensorForEach(dst.extent(), copy_if);
}
template <
typename DstElement,
typename DstLayout,
typename SrcElement,
typename SrcLayout,
typename F
>
void TensorCopy(
TensorRef<DstElement, DstLayout> dst,
TensorView<SrcElement, SrcLayout> src,
F const &transform) {
using CopyIf = detail::TensorCopyIf<
DstElement,
DstLayout,
SrcElement,
SrcLayout,
F>;
TensorView<DstElement, DstLayout> dst_view(dst, src.extent());
CopyIf copy_if(dst_view, src, transform);
TensorForEach(src.extent(), copy_if);
}
template <
typename DstElement,
typename DstLayout,
typename SrcElement,
typename SrcLayout
>
void TensorCopy(
TensorView<DstElement, DstLayout> dst,
TensorView<SrcElement, SrcLayout> src) {
detail::TrivialConvert<DstElement, SrcElement> convert;
TensorCopy(dst, src, convert);
}
template <
typename DstElement,
typename DstLayout,
typename SrcElement,
typename SrcLayout,
typename F
>
void TensorCopy(
TensorView<DstElement, DstLayout> dst,
TensorRef<SrcElement, SrcLayout> src) {
detail::TrivialConvert<DstElement, SrcElement> convert;
TensorCopy(dst, src, convert);
}
template <
typename DstElement,
typename DstLayout,
typename SrcElement,
typename SrcLayout
>
void TensorCopy(
TensorRef<DstElement, DstLayout> dst,
TensorView<SrcElement, SrcLayout> src) {
detail::TrivialConvert<DstElement, SrcElement> convert;
TensorCopy(dst, src, convert);
}
} // namespace host
} // namespace reference
} // namespace cutlass
cutlass
Definition: aligned_buffer.h:35
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::reference::host::TensorCopy
void TensorCopy(TensorView< DstElement, DstLayout > dst, TensorView< SrcElement, SrcLayout > src, F const &transform)
Copies elements from one tensor view into another, satisfying bounds of each tensor.
Definition: tensor_copy.h:112
cutlass::reference::host::detail::TensorCopyIf::TensorCopyIf
TensorCopyIf()
Definition: tensor_copy.h:85
cutlass::reference::host::detail::TrivialConvert::operator()
DstElement operator()(SrcElement src) const
Definition: tensor_copy.h:55
cutlass::reference::host::detail::TensorCopyIf::operator()
void operator()(Coord< DstLayout::kRank > const &coord)
Copies based on destination and source bounds.
Definition: tensor_copy.h:93
cutlass::reference::host::detail::TrivialConvert
Helper to convert between types.
Definition: tensor_copy.h:51
cutlass::reference::host::detail::TensorCopyIf::TensorCopyIf
TensorCopyIf(DstTensorView const &dst_, SrcTensorView const &src_, F const &convert_)
Definition: tensor_copy.h:87
cutlass::TensorView< DstElement, DstLayout >
cutlass::TensorRef< SrcElement, SrcLayout >
cutlass::reference::host::detail::TensorCopyIf
Helper to conditionally copy between tensor views.
Definition: tensor_copy.h:68
cutlass::Coord
Statically-sized array specifying Coords within a tensor.
Definition: coord.h:43
cutlass::reference::host::detail::TensorCopyIf::convert
F convert
Definition: tensor_copy.h:79
cutlass::TensorRef::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::reference::host::TensorForEach
void TensorForEach(Coord< Rank > extent, Func &func)
Iterates over the index space of a tensor.
Definition: host/tensor_foreach.h:87
cutlass::TensorView::contains
CUTLASS_HOST_DEVICE bool contains(TensorCoord const &coord) const
Determines whether a location is within a tensor.
Definition: tensor_view.h:175
cutlass::reference::host::detail::TrivialConvert::TrivialConvert
TrivialConvert()
Definition: tensor_copy.h:53
cutlass::reference::host::detail::TensorCopyIf::src
SrcTensorView src
Definition: tensor_copy.h:78
cutlass::reference::host::detail::TensorCopyIf::dst
DstTensorView dst
Definition: tensor_copy.h:77
cutlass.h
Basic include for CUTLASS.
tensor_foreach.h