host/tensor_elementwise.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
// Cutlass includes
#include "cutlass/cutlass.h"
#include "cutlass/functional.h"
#include "tensor_foreach.h"
namespace cutlass {
namespace reference {
namespace host {
namespace detail {
template <
typename ElementA,
typename LayoutA,
typename ElementB,
typename LayoutB,
typename ElementD,
typename LayoutD,
typename BinaryFunc>
struct TensorFuncBinaryOp {
//
// Data members
//
TensorView<ElementD, LayoutD> view_d;
TensorRef<ElementA, LayoutA> ref_a;
TensorRef<ElementB, LayoutB> ref_b;
BinaryFunc func;
//
// Methods
//
TensorFuncBinaryOp() { }
TensorFuncBinaryOp(
TensorView<ElementD, LayoutD> const & view_d_,
TensorRef<ElementA, LayoutA> const & ref_a_,
TensorRef<ElementB, LayoutB> const & ref_b_,
BinaryFunc func = BinaryFunc()
):
view_d(view_d_), view_a(view_a_), view_b(view_b_), func(func) { }
void operator()(Coord<LayoutD::kRank> const &coord) const {
view_d.at(coord) = func(
ElementD(view_a.at(coord)),
ElementD(view_b.at(coord))
);
}
};
} // namespace detail
template <
typename ElementD,
typename LayoutD,
typename ElementA,
typename LayoutA,
typename ElementB,
typename LayoutB
>
void TensorAdd(
TensorView<ElementD, LayoutD> d,
TensorRef<ElementA, LayoutA> a,
TensorRef<ElementB, LayoutB> b
) {
detail::TensorFuncBinaryOp<
ElementD,
LayoutD,
ElementA,
LayoutA,
ElementB,
LayoutB,
cutlass::plus<ElementD>
> func(d, a, b);
TensorForEach(
d.extent(),
func);
}
template <
typename ElementD,
typename LayoutD,
typename ElementA,
typename LayoutA
>
void TensorAdd(
TensorView<ElementD, LayoutD> d,
TensorRef<ElementA, LayoutA> a
) {
TensorAdd(d, d, a);
}
template <
typename ElementD,
typename LayoutD,
typename ElementA,
typename LayoutA,
typename ElementB,
typename LayoutB
>
void TensorSub(
TensorView<ElementD, LayoutD> d,
TensorRef<ElementA, LayoutA> a,
TensorRef<ElementB, LayoutB> b
) {
detail::TensorFuncBinaryOp<
ElementD,
LayoutD,
ElementA,
LayoutA,
ElementB,
LayoutB,
cutlass::minus<ElementD>
> func(d, a, b);
TensorForEach(
d.extent(),
func);
}
template <
typename ElementD,
typename LayoutD,
typename ElementA,
typename LayoutA,
typename ElementB,
typename LayoutB
>
void TensorSub(
TensorView<ElementD, LayoutD> d,
TensorRef<ElementA, LayoutA> a
) {
TensorSub(d, d, a);
}
template <
typename ElementD,
typename LayoutD,
typename ElementA,
typename LayoutA,
typename ElementB,
typename LayoutB
>
void TensorMul(
TensorView<ElementD, LayoutD> d,
TensorRef<ElementA, LayoutA> a,
TensorRef<ElementB, LayoutB> b
) {
detail::TensorFuncBinaryOp<
ElementD,
LayoutD,
ElementA,
LayoutA,
ElementB,
LayoutB,
cutlass::multiplies<ElementD>
> func(d, a, b);
TensorForEach(
d.extent(),
func);
}
template <
typename ElementD,
typename LayoutD,
typename ElementA,
typename LayoutA
>
void TensorMul(
TensorView<ElementD, LayoutD> d,
TensorRef<ElementA, LayoutA> a
) {
TensorMul(d, d, a);
}
template <
typename ElementD,
typename LayoutD,
typename ElementA,
typename LayoutA,
typename ElementB,
typename LayoutB
>
void TensorDiv(
TensorView<ElementD, LayoutD> d,
TensorRef<ElementA, LayoutA> a,
TensorRef<ElementB, LayoutB> b
) {
detail::TensorFuncBinaryOp<
ElementD,
LayoutD,
ElementA,
LayoutA,
ElementB,
LayoutB,
cutlass::divides<ElementD>
> func(d, a, b);
TensorForEach(
d.extent(),
func);
}
template <
typename ElementD,
typename LayoutD,
typename ElementA,
typename LayoutA
>
void TensorDiv(
TensorView<ElementD, LayoutD> d,
TensorRef<ElementA, LayoutA> a
) {
TensorMul(d, d, a);
}
template <
typename ElementD,
typename LayoutD,
typename ElementA,
typename LayoutA,
typename ElementB,
typename LayoutB
>
void TensorModulus(
TensorView<ElementD, LayoutD> d,
TensorRef<ElementA, LayoutA> a,
TensorRef<ElementB, LayoutB> b
) {
detail::TensorFuncBinaryOp<
ElementD,
LayoutD,
ElementA,
LayoutA,
ElementB,
LayoutB,
cutlass::modulus<ElementD>
> func(d, a, b);
TensorForEach(
d.extent(),
func);
}
template <
typename ElementD,
typename LayoutD,
typename ElementA,
typename LayoutA
>
void TensorModulus(
TensorView<ElementD, LayoutD> d,
TensorRef<ElementA, LayoutA> a
) {
TensorMul(d, d, a);
}
} // namespace host
} // namespace reference
} // namespace cutlass
cutlass::reference::host::detail::TensorFuncBinaryOp::operator()
void operator()(Coord< LayoutD::kRank > const &coord) const
Equality check.
Definition: host/tensor_elementwise.h:86
cutlass
Definition: aligned_buffer.h:35
cutlass::reference::host::detail::TensorFuncBinaryOp
Helper to apply a binary operator in place.
Definition: host/tensor_elementwise.h:57
cutlass::reference::host::TensorAdd
void TensorAdd(TensorView< ElementD, LayoutD > d, TensorRef< ElementA, LayoutA > a, TensorRef< ElementB, LayoutB > b)
Adds two tensors and stores in the destination tensor: d = a + b.
Definition: host/tensor_elementwise.h:108
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::TensorDiv
void TensorDiv(TensorView< ElementD, LayoutD > d, TensorRef< ElementA, LayoutA > a, TensorRef< ElementB, LayoutB > b)
Divides two tensors and stores in the destination tensor: d = a ./ b.
Definition: host/tensor_elementwise.h:249
cutlass::reference::host::detail::TensorFuncBinaryOp::func
BinaryFunc func
Definition: host/tensor_elementwise.h:67
cutlass::plus
Definition: functional.h:46
cutlass::reference::host::detail::TensorFuncBinaryOp::TensorFuncBinaryOp
TensorFuncBinaryOp(TensorView< ElementD, LayoutD > const &view_d_, TensorRef< ElementA, LayoutA > const &ref_a_, TensorRef< ElementB, LayoutB > const &ref_b_, BinaryFunc func=BinaryFunc())
Constructor.
Definition: host/tensor_elementwise.h:77
cutlass::TensorView< ElementD, LayoutD >
cutlass::TensorRef< ElementA, LayoutA >
cutlass::reference::host::detail::TensorFuncBinaryOp::ref_b
TensorRef< ElementB, LayoutB > ref_b
Definition: host/tensor_elementwise.h:66
cutlass::multiplies
Definition: functional.h:64
cutlass::reference::host::TensorSub
void TensorSub(TensorView< ElementD, LayoutD > d, TensorRef< ElementA, LayoutA > a, TensorRef< ElementB, LayoutB > b)
Subtracts two tensors and stores in the destination tensor: d = a - b.
Definition: host/tensor_elementwise.h:154
cutlass::reference::host::detail::TensorFuncBinaryOp::view_d
TensorView< ElementD, LayoutD > view_d
View of left-hand-side tensor.
Definition: host/tensor_elementwise.h:64
cutlass::divides
Definition: functional.h:73
cutlass::reference::host::TensorMul
void TensorMul(TensorView< ElementD, LayoutD > d, TensorRef< ElementA, LayoutA > a, TensorRef< ElementB, LayoutB > b)
Multiplies two tensors and stores in the destination tensor: d = a .* b.
Definition: host/tensor_elementwise.h:203
cutlass::Coord
Statically-sized array specifying Coords within a tensor.
Definition: coord.h:43
cutlass::reference::host::TensorModulus
void TensorModulus(TensorView< ElementD, LayoutD > d, TensorRef< ElementA, LayoutA > a, TensorRef< ElementB, LayoutB > b)
Divides two tensors and stores in the destination tensor: d = a ./ b.
Definition: host/tensor_elementwise.h:296
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::minus
Definition: functional.h:55
cutlass::reference::host::detail::TensorFuncBinaryOp::TensorFuncBinaryOp
TensorFuncBinaryOp()
Constructor.
Definition: host/tensor_elementwise.h:74
cutlass.h
Basic include for CUTLASS.
functional.h
Define basic numeric operators with specializations for Array<T, N>. SIMD-ize where possible...
cutlass::reference::host::detail::TensorFuncBinaryOp::ref_a
TensorRef< ElementA, LayoutA > ref_a
Definition: host/tensor_elementwise.h:65
tensor_foreach.h