CUTLASS 2.0 (#62)
CUTLASS 2.0 Substantially refactored for - Better performance, particularly for native Turing Tensor Cores - Robust and durable templates spanning the design space - Encapsulated functionality embodying modern C++11 programming techniques - Optimized containers and data types for efficient, generic, portable device code Updates to: - Quick start guide - Documentation - Utilities - CUTLASS Profiler Native Turing Tensor Cores - Efficient GEMM kernels targeting Turing Tensor Cores - Mixed-precision floating point, 8-bit integer, 4-bit integer, and binarized operands Coverage of existing CUTLASS functionality: - GEMM kernels targeting CUDA and Tensor Cores in NVIDIA GPUs - Volta Tensor Cores through native mma.sync and through WMMA API - Optimizations such as parallel reductions, threadblock rasterization, and intra-threadblock reductions - Batched GEMM operations - Complex-valued GEMMs Note: this commit and all that follow require a host compiler supporting C++11 or greater.
This commit is contained in:
Andrew Kerr
@@ -0,0 +1,145 @@
|
||||
/***************************************************************************************************
|
||||
* 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 Epilogue for threadblock scoped complex GEMMs using Tensor Ops.
|
||||
|
||||
The epilogue rearranges the result of a matrix product through shared memory to match canonical
|
||||
tensor layouts in global memory. Epilogues support conversion and reduction operations.
|
||||
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/array.h"
|
||||
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
|
||||
#include "cutlass/epilogue/thread/linear_combination.h"
|
||||
#include "cutlass/epilogue/thread/conversion_op.h"
|
||||
#include "cutlass/epilogue/thread/reduction_op.h"
|
||||
|
||||
#include "cutlass/transform/threadblock/regular_tile_iterator_pitch_linear.h"
|
||||
|
||||
#include "cutlass/epilogue/warp/fragment_iterator_complex_tensor_op.h"
|
||||
#include "cutlass/epilogue/warp/tile_iterator_tensor_op.h"
|
||||
#include "cutlass/epilogue/threadblock/default_thread_map_tensor_op.h"
|
||||
#include "cutlass/epilogue/threadblock/predicated_tile_iterator.h"
|
||||
#include "cutlass/epilogue/threadblock/shared_load_iterator.h"
|
||||
|
||||
#include "cutlass/epilogue/threadblock/epilogue.h"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace threadblock {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Defines sensible defaults for epilogues for TensorOps.
|
||||
template <
|
||||
typename Shape_,
|
||||
typename WarpMmaTensorOp_,
|
||||
int PartitionsK,
|
||||
typename OutputOp_,
|
||||
int ElementsPerAccess
|
||||
>
|
||||
struct DefaultEpilogueComplexTensorOp {
|
||||
|
||||
using Shape = Shape_;
|
||||
using WarpMmaTensorOp = WarpMmaTensorOp_;
|
||||
static int const kPartitionsK = PartitionsK;
|
||||
using OutputOp = OutputOp_;
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
|
||||
using ElementOutput = typename OutputOp::ElementOutput;
|
||||
using LayoutC = typename WarpMmaTensorOp::LayoutC;
|
||||
using ElementAccumulator = typename WarpMmaTensorOp::ElementC;
|
||||
|
||||
//
|
||||
// Thread map
|
||||
//
|
||||
|
||||
using OutputTileThreadMap = typename cutlass::epilogue::threadblock::DefaultThreadMapTensorOp<
|
||||
Shape,
|
||||
typename WarpMmaTensorOp::Shape,
|
||||
kPartitionsK,
|
||||
ElementOutput,
|
||||
kElementsPerAccess
|
||||
>::Type;
|
||||
|
||||
using OutputTileIterator = cutlass::epilogue::threadblock::PredicatedTileIterator<
|
||||
OutputTileThreadMap,
|
||||
ElementOutput
|
||||
>;
|
||||
|
||||
using AccumulatorFragmentIterator = cutlass::epilogue::warp::FragmentIteratorComplexTensorOp<
|
||||
typename WarpMmaTensorOp::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::ElementC,
|
||||
typename WarpMmaTensorOp::Policy::Operator::FragmentC,
|
||||
LayoutC
|
||||
>;
|
||||
|
||||
using WarpTileIterator = cutlass::epilogue::warp::TileIteratorTensorOp<
|
||||
typename WarpMmaTensorOp::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::Shape,
|
||||
ElementAccumulator,
|
||||
LayoutC
|
||||
>;
|
||||
|
||||
using SharedLoadIterator = cutlass::epilogue::threadblock::SharedLoadIterator<
|
||||
typename OutputTileThreadMap::CompactedThreadMap,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
/// Hard-coded padding elements added
|
||||
using Padding = cutlass::MatrixShape<0, 0>;
|
||||
|
||||
//
|
||||
// Define the epilogue
|
||||
//
|
||||
using Epilogue = cutlass::epilogue::threadblock::Epilogue<
|
||||
Shape,
|
||||
WarpMmaTensorOp,
|
||||
kPartitionsK,
|
||||
OutputTileIterator,
|
||||
AccumulatorFragmentIterator,
|
||||
WarpTileIterator,
|
||||
SharedLoadIterator,
|
||||
OutputOp,
|
||||
Padding
|
||||
>;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,145 @@
|
||||
/***************************************************************************************************
|
||||
* 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 Epilogue for threadblock scoped GEMMs using SIMT.
|
||||
|
||||
The epilogue rearranges the result of a matrix product through shared memory to match canonical
|
||||
tensor layouts in global memory. Epilogues support conversion and reduction operations.
|
||||
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/array.h"
|
||||
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
|
||||
#include "cutlass/epilogue/thread/linear_combination.h"
|
||||
#include "cutlass/epilogue/thread/conversion_op.h"
|
||||
#include "cutlass/epilogue/thread/reduction_op.h"
|
||||
|
||||
#include "cutlass/transform/threadblock/regular_tile_iterator_pitch_linear.h"
|
||||
|
||||
#include "cutlass/epilogue/warp/fragment_iterator_simt.h"
|
||||
#include "cutlass/epilogue/warp/tile_iterator_simt.h"
|
||||
#include "cutlass/epilogue/threadblock/default_thread_map_simt.h"
|
||||
|
||||
#include "cutlass/epilogue/threadblock/predicated_tile_iterator.h"
|
||||
#include "cutlass/epilogue/threadblock/shared_load_iterator.h"
|
||||
#include "cutlass/epilogue/threadblock/epilogue.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace threadblock {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Defines sensible defaults for epilogues for SimtOps.
|
||||
template <
|
||||
typename Shape_,
|
||||
typename WarpMmaSimt_,
|
||||
typename OutputOp_,
|
||||
int ElementsPerAccess
|
||||
>
|
||||
struct DefaultEpilogueSimt {
|
||||
|
||||
using Shape = Shape_;
|
||||
using WarpMmaSimt = WarpMmaSimt_;
|
||||
using OutputOp = OutputOp_;
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
static const int kPartitionsK = Shape::kK / WarpMmaSimt::Shape::kK;
|
||||
|
||||
using ElementOutput = typename OutputOp::ElementOutput;
|
||||
using LayoutC = typename WarpMmaSimt::LayoutC;
|
||||
using ElementAccumulator = typename WarpMmaSimt::ElementC;
|
||||
|
||||
//
|
||||
// Thread map
|
||||
//
|
||||
|
||||
using OutputTileThreadMap = typename cutlass::epilogue::threadblock::DefaultThreadMapSimt<
|
||||
Shape,
|
||||
typename WarpMmaSimt::Shape,
|
||||
typename WarpMmaSimt::Policy,
|
||||
kPartitionsK,
|
||||
ElementOutput,
|
||||
kElementsPerAccess
|
||||
>::Type;
|
||||
|
||||
using OutputTileIterator = cutlass::epilogue::threadblock::PredicatedTileIterator<
|
||||
OutputTileThreadMap,
|
||||
ElementOutput
|
||||
>;
|
||||
|
||||
using AccumulatorFragmentIterator = cutlass::epilogue::warp::FragmentIteratorSimt<
|
||||
typename WarpMmaSimt::Shape,
|
||||
typename WarpMmaSimt::ThreadMma,
|
||||
layout::RowMajor,
|
||||
typename WarpMmaSimt::Policy
|
||||
>;
|
||||
|
||||
using WarpTileIterator = cutlass::epilogue::warp::TileIteratorSimt<
|
||||
typename WarpMmaSimt::Shape,
|
||||
typename WarpMmaSimt::ThreadMma,
|
||||
ElementAccumulator,
|
||||
layout::RowMajor,
|
||||
typename WarpMmaSimt::Policy
|
||||
>;
|
||||
|
||||
using SharedLoadIterator = cutlass::epilogue::threadblock::SharedLoadIterator<
|
||||
typename OutputTileThreadMap::CompactedThreadMap,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
/// Hard-coded padding elements added
|
||||
using Padding = typename WarpTileIterator::Padding;
|
||||
|
||||
//
|
||||
// Define the epilogue
|
||||
//
|
||||
using Epilogue = cutlass::epilogue::threadblock::Epilogue<
|
||||
Shape,
|
||||
WarpMmaSimt,
|
||||
kPartitionsK,
|
||||
OutputTileIterator,
|
||||
AccumulatorFragmentIterator,
|
||||
WarpTileIterator,
|
||||
SharedLoadIterator,
|
||||
OutputOp,
|
||||
Padding
|
||||
>;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,191 @@
|
||||
/***************************************************************************************************
|
||||
* 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 Epilogue for threadblock scoped GEMMs using Tensor Ops.
|
||||
|
||||
The epilogue rearranges the result of a matrix product through shared memory to match canonical
|
||||
tensor layouts in global memory. Epilogues support conversion and reduction operations.
|
||||
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/array.h"
|
||||
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
|
||||
#include "cutlass/epilogue/thread/linear_combination.h"
|
||||
#include "cutlass/epilogue/thread/conversion_op.h"
|
||||
#include "cutlass/epilogue/thread/reduction_op.h"
|
||||
|
||||
#include "cutlass/transform/threadblock/regular_tile_iterator_pitch_linear.h"
|
||||
|
||||
#include "cutlass/epilogue/warp/fragment_iterator_tensor_op.h"
|
||||
#include "cutlass/epilogue/warp/tile_iterator_tensor_op.h"
|
||||
#include "cutlass/epilogue/threadblock/default_thread_map_tensor_op.h"
|
||||
#include "cutlass/epilogue/threadblock/predicated_tile_iterator.h"
|
||||
#include "cutlass/epilogue/threadblock/shared_load_iterator.h"
|
||||
|
||||
#include "cutlass/epilogue/threadblock/epilogue.h"
|
||||
#include "cutlass/epilogue/threadblock/interleaved_epilogue.h"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace threadblock {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Defines sensible defaults for epilogues for TensorOps.
|
||||
template <
|
||||
typename Shape_,
|
||||
typename WarpMmaTensorOp_,
|
||||
int PartitionsK,
|
||||
typename OutputOp_,
|
||||
int ElementsPerAccess
|
||||
>
|
||||
struct DefaultEpilogueTensorOp {
|
||||
|
||||
using Shape = Shape_;
|
||||
using WarpMmaTensorOp = WarpMmaTensorOp_;
|
||||
static int const kPartitionsK = PartitionsK;
|
||||
using OutputOp = OutputOp_;
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
|
||||
using ElementOutput = typename OutputOp::ElementOutput;
|
||||
using LayoutC = typename WarpMmaTensorOp::LayoutC;
|
||||
using ElementAccumulator = typename WarpMmaTensorOp::ElementC;
|
||||
|
||||
//
|
||||
// Thread map
|
||||
//
|
||||
|
||||
using OutputTileThreadMap = typename cutlass::epilogue::threadblock::DefaultThreadMapTensorOp<
|
||||
Shape,
|
||||
typename WarpMmaTensorOp::Shape,
|
||||
kPartitionsK,
|
||||
ElementOutput,
|
||||
kElementsPerAccess
|
||||
>::Type;
|
||||
|
||||
using OutputTileIterator = cutlass::epilogue::threadblock::PredicatedTileIterator<
|
||||
OutputTileThreadMap,
|
||||
ElementOutput
|
||||
>;
|
||||
|
||||
using AccumulatorFragmentIterator = cutlass::epilogue::warp::FragmentIteratorTensorOp<
|
||||
typename WarpMmaTensorOp::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::ElementC,
|
||||
typename WarpMmaTensorOp::Policy::Operator::FragmentC,
|
||||
LayoutC
|
||||
>;
|
||||
|
||||
using WarpTileIterator = cutlass::epilogue::warp::TileIteratorTensorOp<
|
||||
typename WarpMmaTensorOp::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::Shape,
|
||||
ElementAccumulator,
|
||||
LayoutC
|
||||
>;
|
||||
|
||||
using SharedLoadIterator = cutlass::epilogue::threadblock::SharedLoadIterator<
|
||||
typename OutputTileThreadMap::CompactedThreadMap,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
/// Hard-coded padding elements added
|
||||
using Padding = cutlass::MatrixShape<0, 64 / sizeof_bits<ElementAccumulator>::value * 4>;
|
||||
|
||||
//
|
||||
// Define the epilogue
|
||||
//
|
||||
using Epilogue = cutlass::epilogue::threadblock::Epilogue<
|
||||
Shape,
|
||||
WarpMmaTensorOp,
|
||||
kPartitionsK,
|
||||
OutputTileIterator,
|
||||
AccumulatorFragmentIterator,
|
||||
WarpTileIterator,
|
||||
SharedLoadIterator,
|
||||
OutputOp,
|
||||
Padding
|
||||
>;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Defines sensible defaults for epilogues for TensorOps which uses
|
||||
/// intereleaved output layout. For this case, shared memory is not needed.
|
||||
template <typename Shape_, typename WarpMmaTensorOp_, int PartitionsK,
|
||||
typename OutputOp_, int ElementsPerAccess, int InterleavedK,
|
||||
bool IsBetaZero = false, bool isSplitK = false>
|
||||
struct DefaultInterleavedEpilogueTensorOp {
|
||||
using Shape = Shape_;
|
||||
using WarpMmaTensorOp = WarpMmaTensorOp_;
|
||||
static int const kPartitionsK = PartitionsK;
|
||||
using OutputOp = OutputOp_;
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
|
||||
using ElementOutput = typename OutputOp::ElementOutput;
|
||||
using LayoutC = typename WarpMmaTensorOp::LayoutC;
|
||||
using ElementAccumulator = typename WarpMmaTensorOp::ElementC;
|
||||
|
||||
//
|
||||
// Thread map
|
||||
//
|
||||
using OutputTileThreadMap = typename cutlass::epilogue::threadblock::
|
||||
DefaultInterleavedThreadMapTensorOp<
|
||||
Shape, typename WarpMmaTensorOp::Shape, kPartitionsK, ElementOutput,
|
||||
kElementsPerAccess, InterleavedK>::Type;
|
||||
|
||||
using OutputTileIterator =
|
||||
cutlass::epilogue::threadblock::InterleavedPredicatedTileIterator<
|
||||
OutputTileThreadMap, ElementOutput, InterleavedK>;
|
||||
|
||||
using AccumulatorFragmentIterator =
|
||||
cutlass::epilogue::warp::FragmentIteratorTensorOp<
|
||||
typename WarpMmaTensorOp::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::ElementC,
|
||||
typename WarpMmaTensorOp::Policy::Operator::FragmentC,
|
||||
LayoutC>;
|
||||
|
||||
//
|
||||
// Define the epilogue
|
||||
//
|
||||
using Epilogue = cutlass::epilogue::threadblock::InterleavedEpilogue<
|
||||
Shape, WarpMmaTensorOp, kPartitionsK, OutputTileIterator,
|
||||
AccumulatorFragmentIterator, OutputOp, InterleavedK, IsBetaZero>;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
} // namespace threadblock
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,150 @@
|
||||
/***************************************************************************************************
|
||||
* 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 Epilogue for threadblock scoped GEMMs using Tensor Ops on Volta.
|
||||
|
||||
The epilogue rearranges the result of a matrix product through shared memory to match canonical
|
||||
tensor layouts in global memory. Epilogues support conversion and reduction operations.
|
||||
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/array.h"
|
||||
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
|
||||
#include "cutlass/epilogue/thread/linear_combination.h"
|
||||
#include "cutlass/epilogue/thread/conversion_op.h"
|
||||
#include "cutlass/epilogue/thread/reduction_op.h"
|
||||
|
||||
#include "cutlass/transform/threadblock/regular_tile_iterator_pitch_linear.h"
|
||||
#include "cutlass/epilogue/threadblock/predicated_tile_iterator.h"
|
||||
#include "cutlass/epilogue/threadblock/shared_load_iterator.h"
|
||||
|
||||
#include "cutlass/epilogue/warp/fragment_iterator_volta_tensor_op.h"
|
||||
#include "cutlass/epilogue/warp/tile_iterator_volta_tensor_op.h"
|
||||
#include "cutlass/epilogue/threadblock/default_thread_map_volta_tensor_op.h"
|
||||
|
||||
#include "cutlass/epilogue/threadblock/epilogue.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace threadblock {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Defines sensible defaults for epilogues for TensorOps.
|
||||
template <
|
||||
typename Shape_,
|
||||
typename WarpMmaTensorOp_,
|
||||
int PartitionsK,
|
||||
typename OutputOp_,
|
||||
int ElementsPerAccess
|
||||
>
|
||||
struct DefaultEpilogueVoltaTensorOp {
|
||||
|
||||
using Shape = Shape_;
|
||||
using WarpMmaTensorOp = WarpMmaTensorOp_;
|
||||
static int const kPartitionsK = PartitionsK;
|
||||
using OutputOp = OutputOp_;
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
|
||||
using ElementOutput = typename OutputOp::ElementOutput;
|
||||
using LayoutC = typename WarpMmaTensorOp::LayoutC;
|
||||
using ElementAccumulator = typename WarpMmaTensorOp::ElementC;
|
||||
|
||||
//
|
||||
// Thread map
|
||||
//
|
||||
|
||||
using OutputTileThreadMap = typename cutlass::epilogue::threadblock::DefaultThreadMapVoltaTensorOp<
|
||||
Shape,
|
||||
typename WarpMmaTensorOp::Shape,
|
||||
kPartitionsK,
|
||||
ElementOutput,
|
||||
kElementsPerAccess,
|
||||
ElementAccumulator
|
||||
>::Type;
|
||||
|
||||
using OutputTileIterator = cutlass::epilogue::threadblock::PredicatedTileIterator<
|
||||
OutputTileThreadMap,
|
||||
ElementOutput
|
||||
>;
|
||||
|
||||
using AccumulatorFragmentIterator = cutlass::epilogue::warp::FragmentIteratorVoltaTensorOp<
|
||||
typename WarpMmaTensorOp::Shape,
|
||||
gemm::GemmShape<32, 32, 4>,
|
||||
ElementAccumulator,
|
||||
LayoutC
|
||||
>;
|
||||
|
||||
using WarpTileIterator = cutlass::epilogue::warp::TileIteratorVoltaTensorOp<
|
||||
typename WarpMmaTensorOp::Shape,
|
||||
gemm::GemmShape<32, 32, 4>,
|
||||
ElementAccumulator,
|
||||
LayoutC
|
||||
>;
|
||||
|
||||
static int const kSharedMemAlignment = sizeof_bits<ElementAccumulator>::value * WarpTileIterator::kElementsPerAccess / 8;
|
||||
|
||||
static_assert(kSharedMemAlignment == 8, "Shared memory alignment must be 8B");
|
||||
|
||||
using SharedLoadIterator = cutlass::epilogue::threadblock::SharedLoadIterator<
|
||||
typename OutputTileThreadMap::CompactedThreadMap,
|
||||
ElementAccumulator,
|
||||
kSharedMemAlignment
|
||||
>;
|
||||
|
||||
/// Hard-coded padding elements added
|
||||
using Padding = typename WarpTileIterator::Padding;
|
||||
|
||||
//
|
||||
// Define the epilogue
|
||||
//
|
||||
using Epilogue = cutlass::epilogue::threadblock::Epilogue<
|
||||
Shape,
|
||||
WarpMmaTensorOp,
|
||||
kPartitionsK,
|
||||
OutputTileIterator,
|
||||
AccumulatorFragmentIterator,
|
||||
WarpTileIterator,
|
||||
SharedLoadIterator,
|
||||
OutputOp,
|
||||
Padding
|
||||
>;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,147 @@
|
||||
/***************************************************************************************************
|
||||
* 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 Epilogue for threadblock scoped GEMMs using Tensor Ops.
|
||||
|
||||
The epilogue rearranges the result of a matrix product through shared memory to match canonical
|
||||
tensor layouts in global memory. Epilogues support conversion and reduction operations.
|
||||
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/array.h"
|
||||
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
|
||||
#include "cutlass/epilogue/thread/linear_combination.h"
|
||||
#include "cutlass/epilogue/thread/conversion_op.h"
|
||||
#include "cutlass/epilogue/thread/reduction_op.h"
|
||||
|
||||
#include "cutlass/transform/threadblock/regular_tile_iterator_pitch_linear.h"
|
||||
|
||||
#include "cutlass/epilogue/warp/fragment_iterator_wmma_tensor_op.h"
|
||||
#include "cutlass/epilogue/warp/tile_iterator_wmma_tensor_op.h"
|
||||
#include "cutlass/epilogue/threadblock/default_thread_map_wmma_tensor_op.h"
|
||||
#include "cutlass/epilogue/threadblock/predicated_tile_iterator.h"
|
||||
#include "cutlass/epilogue/threadblock/shared_load_iterator.h"
|
||||
|
||||
#include "cutlass/epilogue/threadblock/epilogue.h"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace threadblock {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Defines sensible defaults for epilogues for WMMA TensorOps.
|
||||
template <
|
||||
typename Shape_,
|
||||
typename WarpMmaTensorOp_,
|
||||
int PartitionsK,
|
||||
typename OutputOp_,
|
||||
int ElementsPerAccess
|
||||
>
|
||||
struct DefaultEpilogueWmmaTensorOp {
|
||||
|
||||
using Shape = Shape_;
|
||||
using WarpMmaTensorOp = WarpMmaTensorOp_;
|
||||
static int const kPartitionsK = PartitionsK;
|
||||
using OutputOp = OutputOp_;
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
|
||||
using ElementOutput = typename OutputOp::ElementOutput;
|
||||
using LayoutC = typename WarpMmaTensorOp::LayoutC;
|
||||
using ElementAccumulator = typename WarpMmaTensorOp::ElementC;
|
||||
|
||||
//
|
||||
// Thread map
|
||||
//
|
||||
|
||||
using OutputTileThreadMap = typename cutlass::epilogue::threadblock::DefaultThreadMapWmmaTensorOp<
|
||||
Shape,
|
||||
typename WarpMmaTensorOp::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::Shape,
|
||||
kPartitionsK,
|
||||
ElementOutput,
|
||||
kElementsPerAccess
|
||||
>::Type;
|
||||
|
||||
using OutputTileIterator = cutlass::epilogue::threadblock::PredicatedTileIterator<
|
||||
OutputTileThreadMap,
|
||||
ElementOutput
|
||||
>;
|
||||
|
||||
using AccumulatorFragmentIterator = cutlass::epilogue::warp::FragmentIteratorWmmaTensorOp<
|
||||
typename WarpMmaTensorOp::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::ElementC,
|
||||
typename WarpMmaTensorOp::Policy::Operator::FragmentC,
|
||||
LayoutC
|
||||
>;
|
||||
|
||||
using WarpTileIterator = cutlass::epilogue::warp::TileIteratorWmmaTensorOp<
|
||||
typename WarpMmaTensorOp::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::FragmentC,
|
||||
LayoutC
|
||||
>;
|
||||
|
||||
using SharedLoadIterator = cutlass::epilogue::threadblock::SharedLoadIterator<
|
||||
typename OutputTileThreadMap::CompactedThreadMap,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
/// Hard-coded padding elements added
|
||||
using Padding = typename WarpTileIterator::Padding;
|
||||
|
||||
//
|
||||
// Define the epilogue
|
||||
//
|
||||
using Epilogue = cutlass::epilogue::threadblock::Epilogue<
|
||||
Shape,
|
||||
WarpMmaTensorOp,
|
||||
kPartitionsK,
|
||||
OutputTileIterator,
|
||||
AccumulatorFragmentIterator,
|
||||
WarpTileIterator,
|
||||
SharedLoadIterator,
|
||||
OutputOp,
|
||||
Padding
|
||||
>;
|
||||
};
|
||||
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,121 @@
|
||||
/***************************************************************************************************
|
||||
* 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
|
||||
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "predicated_tile_iterator.h"
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace threadblock {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Defines the optimal thread map for SIMT accumulator layouts
|
||||
template <
|
||||
typename ThreadblockShape_,
|
||||
typename WarpShape_,
|
||||
typename MmaSimtPolicy_,
|
||||
int PartitionsK,
|
||||
typename Element_,
|
||||
int ElementsPerAccess
|
||||
>
|
||||
struct DefaultThreadMapSimt {
|
||||
|
||||
using ThreadblockShape = ThreadblockShape_;
|
||||
using WarpShape = WarpShape_;
|
||||
using MmaSimtPolicy = MmaSimtPolicy_;
|
||||
static int const kPartitionsK = PartitionsK;
|
||||
using Element = Element_;
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
|
||||
//
|
||||
// Definitions
|
||||
//
|
||||
|
||||
struct Detail {
|
||||
|
||||
static int const kWarpSize = 32;
|
||||
|
||||
static_assert(
|
||||
!(ThreadblockShape::kM % WarpShape::kM) &&
|
||||
!(ThreadblockShape::kM % WarpShape::kM), "Divisibility");
|
||||
|
||||
/// Number of warps
|
||||
using WarpCount = gemm::GemmShape<
|
||||
ThreadblockShape::kM / WarpShape::kM,
|
||||
ThreadblockShape::kN / WarpShape::kN,
|
||||
kPartitionsK
|
||||
>;
|
||||
|
||||
/// Computes number of thread-level matrix multiplies are needed to span a warp
|
||||
static int const kGroupCount =
|
||||
WarpShape::kM / (MmaSimtPolicy::WarpShape::kRow * MmaSimtPolicy::LaneMmaShape::kM);
|
||||
|
||||
/// Number of participating threads
|
||||
static int const kThreads = WarpCount::kCount * kWarpSize;
|
||||
|
||||
/// Number of iterations
|
||||
static int const kIterations = MmaSimtPolicy::LaneMmaShape::kM * kGroupCount;
|
||||
};
|
||||
|
||||
//
|
||||
// ThreadMap
|
||||
//
|
||||
|
||||
/// ThreadMap to be used by epilogue::PredicatedTileIterator satisfying concept OutputTileThreadMap
|
||||
using Type = OutputTileOptimalThreadMap<
|
||||
OutputTileShape< // Shape
|
||||
ThreadblockShape::kN,
|
||||
1,
|
||||
MmaSimtPolicy::WarpShape::kRow,
|
||||
Detail::WarpCount::kM,
|
||||
1>,
|
||||
OutputTileShape< // Count
|
||||
1,
|
||||
MmaSimtPolicy::LaneMmaShape::kM,
|
||||
Detail::kGroupCount,
|
||||
1,
|
||||
Detail::kIterations>,
|
||||
Detail::kThreads,
|
||||
kElementsPerAccess,
|
||||
sizeof_bits<Element>::value
|
||||
>;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,153 @@
|
||||
/***************************************************************************************************
|
||||
* 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
|
||||
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "predicated_tile_iterator.h"
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
#include "cutlass/layout/pitch_linear.h"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace threadblock {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Defines the optimal thread map for TensorOp accumulator layouts
|
||||
template <
|
||||
typename ThreadblockShape_,
|
||||
typename WarpShape_,
|
||||
int PartitionsK,
|
||||
typename Element_,
|
||||
int ElementsPerAccess
|
||||
>
|
||||
struct DefaultThreadMapTensorOp {
|
||||
|
||||
using ThreadblockShape = ThreadblockShape_;
|
||||
using WarpShape = WarpShape_;
|
||||
static int const kPartitionsK = PartitionsK;
|
||||
using Element = Element_;
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
|
||||
//
|
||||
// Definitions
|
||||
//
|
||||
|
||||
struct Detail {
|
||||
|
||||
/// Tensor Operations fundamentally perform operations on 8 rows
|
||||
static int const kTensorOpRows = 8;
|
||||
static int const kWarpSize = 32;
|
||||
|
||||
static_assert(
|
||||
!(ThreadblockShape::kM % WarpShape::kM) &&
|
||||
!(ThreadblockShape::kM % WarpShape::kM), "Divisibility");
|
||||
|
||||
/// Number of warps
|
||||
using WarpCount = gemm::GemmShape<
|
||||
ThreadblockShape::kM / WarpShape::kM,
|
||||
ThreadblockShape::kN / WarpShape::kN,
|
||||
kPartitionsK
|
||||
>;
|
||||
|
||||
/// Number of participating threads
|
||||
static int const kThreads = WarpCount::kCount * kWarpSize;
|
||||
};
|
||||
|
||||
//
|
||||
// ThreadMap
|
||||
//
|
||||
|
||||
/// ThreadMap to be used by epilogue::PredicatedTileIterator satisfying concept OutputTileThreadMap
|
||||
using Type = OutputTileOptimalThreadMap <
|
||||
OutputTileShape<ThreadblockShape::kN, Detail::kTensorOpRows, Detail::WarpCount::kM, 1, 1>,
|
||||
OutputTileShape<1, WarpShape::kM / Detail::kTensorOpRows, 1, 1, WarpShape::kM / Detail::kTensorOpRows>,
|
||||
Detail::kThreads,
|
||||
kElementsPerAccess,
|
||||
sizeof_bits<Element>::value
|
||||
>;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Defines the optimal thread map for TensorOp accumulator layouts
|
||||
template <typename ThreadblockShape_, typename WarpShape_, int PartitionsK,
|
||||
typename Element_, int ElementsPerAccess, int InterleavedK>
|
||||
struct DefaultInterleavedThreadMapTensorOp {
|
||||
using ThreadblockShape = ThreadblockShape_;
|
||||
using WarpShape = WarpShape_;
|
||||
static int const kPartitionsK = PartitionsK;
|
||||
using Element = Element_;
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
static int const kInterleavedK = InterleavedK;
|
||||
|
||||
//
|
||||
// Definitions
|
||||
//
|
||||
|
||||
struct Detail {
|
||||
/// Tensor Operations fundamentally perform operations on 8 rows
|
||||
static int const kTensorOpRows = 8;
|
||||
static int const kWarpSize = 32;
|
||||
|
||||
static_assert(!(ThreadblockShape::kM % WarpShape::kM) &&
|
||||
!(ThreadblockShape::kM % WarpShape::kM),
|
||||
"Divisibility");
|
||||
|
||||
/// Number of warps
|
||||
using WarpCount =
|
||||
gemm::GemmShape<ThreadblockShape::kM / WarpShape::kM,
|
||||
ThreadblockShape::kN / WarpShape::kN, kPartitionsK>;
|
||||
|
||||
/// Number of participating threads
|
||||
static int const kThreads = WarpCount::kCount * kWarpSize;
|
||||
};
|
||||
|
||||
//
|
||||
// ThreadMap
|
||||
//
|
||||
|
||||
/// ThreadMap to be used by epilogue::PredicatedTileIterator satisfying concept
|
||||
/// InterleavedOutputTileThreadMap
|
||||
using Type = InterleavedOutputTileThreadMap<
|
||||
layout::PitchLinearShape<Detail::WarpCount::kM, Detail::WarpCount::kN>,
|
||||
layout::PitchLinearShape<WarpShape::kM / Detail::kTensorOpRows,
|
||||
WarpShape::kN / InterleavedK>,
|
||||
Detail::kThreads, kElementsPerAccess, sizeof_bits<Element>::value>;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,222 @@
|
||||
/***************************************************************************************************
|
||||
* 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
|
||||
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "predicated_tile_iterator.h"
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace threadblock {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Defines the optimal thread map for TensorOp accumulator layouts
|
||||
template <
|
||||
typename ThreadblockShape,
|
||||
typename WarpShape,
|
||||
int PartitionsK,
|
||||
typename ElementOutput,
|
||||
int ElementsPerAccess,
|
||||
typename ElementAccumulator
|
||||
>
|
||||
struct DefaultThreadMapVoltaTensorOp;
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Defines the optimal thread map for TensorOp accumulator layouts
|
||||
template <
|
||||
typename ThreadblockShape_,
|
||||
typename WarpShape_,
|
||||
int PartitionsK,
|
||||
typename ElementOutput_,
|
||||
int ElementsPerAccess
|
||||
>
|
||||
struct DefaultThreadMapVoltaTensorOp<
|
||||
ThreadblockShape_,
|
||||
WarpShape_,
|
||||
PartitionsK,
|
||||
ElementOutput_,
|
||||
ElementsPerAccess,
|
||||
half_t> {
|
||||
|
||||
using ThreadblockShape = ThreadblockShape_;
|
||||
using WarpShape = WarpShape_;
|
||||
static int const kPartitionsK = PartitionsK;
|
||||
using ElementOutput = ElementOutput_;
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
using ElementAccumulator = half_t;
|
||||
|
||||
//
|
||||
// Definitions
|
||||
//
|
||||
|
||||
struct Detail {
|
||||
|
||||
static int const kTensorOpRows = 16;
|
||||
static int const kWarpSize = 32;
|
||||
static int const kInterleavedTilesM = WarpShape::kM / 32;
|
||||
|
||||
static_assert(
|
||||
!(ThreadblockShape::kM % WarpShape::kM) &&
|
||||
!(ThreadblockShape::kM % WarpShape::kM), "Divisibility");
|
||||
|
||||
/// Number of warps
|
||||
using WarpCount = gemm::GemmShape<
|
||||
ThreadblockShape::kM / WarpShape::kM,
|
||||
ThreadblockShape::kN / WarpShape::kN,
|
||||
kPartitionsK
|
||||
>;
|
||||
|
||||
/// Number of participating threads
|
||||
static int const kThreads = WarpCount::kCount * kWarpSize;
|
||||
|
||||
using Shape = cutlass::epilogue::threadblock::OutputTileShape<
|
||||
ThreadblockShape::kN, // column
|
||||
4, // row
|
||||
4, // group
|
||||
WarpCount::kM, // cluster
|
||||
1 // tile
|
||||
>;
|
||||
|
||||
/// Number of iterations per subspace
|
||||
using Count = cutlass::epilogue::threadblock::OutputTileShape<
|
||||
1, // column
|
||||
2, // row
|
||||
kInterleavedTilesM, // group
|
||||
1, // cluster
|
||||
WarpShape::kM / kTensorOpRows // iterations
|
||||
>;
|
||||
};
|
||||
|
||||
//
|
||||
// ThreadMap
|
||||
//
|
||||
|
||||
/// ThreadMap to be used by epilogue::PredicatedTileIterator satisfying concept OutputTileThreadMap
|
||||
using Type = OutputTileOptimalThreadMap <
|
||||
typename Detail::Shape,
|
||||
typename Detail::Count,
|
||||
Detail::kThreads,
|
||||
kElementsPerAccess,
|
||||
sizeof_bits<ElementOutput>::value
|
||||
>;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Defines the optimal thread map for TensorOp accumulator layouts
|
||||
template <
|
||||
typename ThreadblockShape_,
|
||||
typename WarpShape_,
|
||||
int PartitionsK,
|
||||
typename ElementOutput_,
|
||||
int ElementsPerAccess
|
||||
>
|
||||
struct DefaultThreadMapVoltaTensorOp<
|
||||
ThreadblockShape_,
|
||||
WarpShape_,
|
||||
PartitionsK,
|
||||
ElementOutput_,
|
||||
ElementsPerAccess,
|
||||
float> {
|
||||
|
||||
using ThreadblockShape = ThreadblockShape_;
|
||||
using WarpShape = WarpShape_;
|
||||
static int const kPartitionsK = PartitionsK;
|
||||
using ElementOutput = ElementOutput_;
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
using ElementAccumulator = float;
|
||||
|
||||
//
|
||||
// Definitions
|
||||
//
|
||||
|
||||
struct Detail {
|
||||
|
||||
static int const kTensorOpRows = 16;
|
||||
static int const kWarpSize = 32;
|
||||
static int const kInterleavedTilesM = WarpShape::kM / 32;
|
||||
|
||||
static_assert(
|
||||
!(ThreadblockShape::kM % WarpShape::kM) &&
|
||||
!(ThreadblockShape::kM % WarpShape::kM), "Divisibility");
|
||||
|
||||
/// Number of warps
|
||||
using WarpCount = gemm::GemmShape<
|
||||
ThreadblockShape::kM / WarpShape::kM,
|
||||
ThreadblockShape::kN / WarpShape::kN,
|
||||
kPartitionsK
|
||||
>;
|
||||
|
||||
/// Number of participating threads
|
||||
static int const kThreads = WarpCount::kCount * kWarpSize;
|
||||
|
||||
using Shape = cutlass::epilogue::threadblock::OutputTileShape<
|
||||
ThreadblockShape::kN, // column
|
||||
4, // row
|
||||
4, // group
|
||||
WarpCount::kM, // cluster
|
||||
1 // tile
|
||||
>;
|
||||
|
||||
/// Number of iterations per subspace
|
||||
using Count = cutlass::epilogue::threadblock::OutputTileShape<
|
||||
1, // column
|
||||
2, // row
|
||||
kInterleavedTilesM, // group
|
||||
1, // cluster
|
||||
WarpShape::kM / kTensorOpRows // iterations
|
||||
>;
|
||||
};
|
||||
|
||||
//
|
||||
// ThreadMap
|
||||
//
|
||||
|
||||
/// ThreadMap to be used by epilogue::PredicatedTileIterator satisfying concept OutputTileThreadMap
|
||||
using Type = OutputTileOptimalThreadMap <
|
||||
typename Detail::Shape,
|
||||
typename Detail::Count,
|
||||
Detail::kThreads,
|
||||
kElementsPerAccess,
|
||||
sizeof_bits<ElementOutput>::value
|
||||
>;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,107 @@
|
||||
/***************************************************************************************************
|
||||
* 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
|
||||
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "predicated_tile_iterator.h"
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
#include "cutlass/layout/pitch_linear.h"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace threadblock {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Defines the optimal thread map for Wmma TensorOp accumulator layouts
|
||||
template <
|
||||
typename ThreadblockShape_,
|
||||
typename WarpShape_,
|
||||
typename InstructionShape_,
|
||||
int PartitionsK,
|
||||
typename Element_,
|
||||
int ElementsPerAccess
|
||||
>
|
||||
struct DefaultThreadMapWmmaTensorOp {
|
||||
|
||||
using ThreadblockShape = ThreadblockShape_;
|
||||
using WarpShape = WarpShape_;
|
||||
using InstructionShape = InstructionShape_;
|
||||
static int const kPartitionsK = PartitionsK;
|
||||
using Element = Element_;
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
|
||||
//
|
||||
// Definitions
|
||||
//
|
||||
|
||||
struct Detail {
|
||||
|
||||
/// Wmma Tensor Operations fundamentally perform operations on InstructionShape::kM rows
|
||||
static int const kTensorOpRows = InstructionShape::kM;
|
||||
static int const kWarpSize = 32;
|
||||
|
||||
static_assert(
|
||||
!(ThreadblockShape::kM % WarpShape::kM) &&
|
||||
!(ThreadblockShape::kM % WarpShape::kM), "Divisibility");
|
||||
|
||||
/// Number of warps
|
||||
using WarpCount = gemm::GemmShape<
|
||||
ThreadblockShape::kM / WarpShape::kM,
|
||||
ThreadblockShape::kN / WarpShape::kN,
|
||||
kPartitionsK
|
||||
>;
|
||||
|
||||
/// Number of participating threads
|
||||
static int const kThreads = WarpCount::kCount * kWarpSize;
|
||||
};
|
||||
|
||||
//
|
||||
// ThreadMap
|
||||
//
|
||||
|
||||
/// ThreadMap to be used by epilogue::PredicatedTileIterator satisfying concept OutputTileThreadMap
|
||||
using Type = OutputTileOptimalThreadMap <
|
||||
OutputTileShape<ThreadblockShape::kN, Detail::kTensorOpRows, Detail::WarpCount::kM, 1, 1>,
|
||||
OutputTileShape<1, WarpShape::kM / Detail::kTensorOpRows, 1, 1, WarpShape::kM / Detail::kTensorOpRows>,
|
||||
Detail::kThreads,
|
||||
kElementsPerAccess,
|
||||
sizeof_bits<Element>::value
|
||||
>;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,252 @@
|
||||
/***************************************************************************************************
|
||||
* 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 Epilogue for tensor operations
|
||||
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/array.h"
|
||||
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace threadblock {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Epilogue operator
|
||||
template <
|
||||
typename Shape_, /// Threadblock-scoped tile size (concept: GemmShape)
|
||||
typename Operator_, /// Warp-scoped epilogue components (concept: gemm::warp::Mma)
|
||||
int PartitionsK, /// Number of partitions of the threadblock's K dimension
|
||||
typename Element_, /// Data type of source and destination tensor
|
||||
typename OutputOp_, /// Function object computing final output
|
||||
typename ConvertOp_ /// Conversion operator between accumulators and shared memory
|
||||
>
|
||||
class DirectEpilogueTensorOp {
|
||||
public:
|
||||
|
||||
using Shape = Shape_;
|
||||
using Operator = Operator_;
|
||||
|
||||
/// Number of warps spanning threadblock-scoped tile
|
||||
using WarpCount = gemm::GemmShape<
|
||||
Shape::kM / Operator::Shape::kM,
|
||||
Shape::kN / Operator::Shape::kN,
|
||||
PartitionsK,
|
||||
>;
|
||||
|
||||
static_assert(PartitionsK == 1,
|
||||
"Direct epilogue cannot be used with when the threadblock tile is partitioned along the K dimension.");
|
||||
|
||||
/// Accumulator tile is really the warp-scoped tile
|
||||
using FragmentC = typename Operator::FragmentC;
|
||||
|
||||
/// Data type of output tensor
|
||||
using Element = Element_;
|
||||
|
||||
/// Output layout is always row-major
|
||||
using Layout = layout::RowMajor;
|
||||
|
||||
/// Function operator computing final output
|
||||
using OutputOp = OutputOp_;
|
||||
|
||||
/// Conversion operator to shared memory
|
||||
using ConvertOp = ConvertOp_;
|
||||
|
||||
/// Reference to source and destination tensors
|
||||
using TensorRef = TensorRef<Element, Layout::kRank, Layout>;
|
||||
|
||||
public:
|
||||
|
||||
/// Parameters structure for host-constructible state
|
||||
struct Params {
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
TensorRef destination_ref;
|
||||
TensorRef source_ref;
|
||||
|
||||
typename OutputOp::Params output_op;
|
||||
typename ConvertOp::Params convert_op;
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
/// Constructs a Params object
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(
|
||||
TensorRef destination_ref_,
|
||||
TensorRef source_ref_,
|
||||
typename OutputOp::Params output_op_,
|
||||
typename ConvertOp::Params convert_op_
|
||||
):
|
||||
destination_ref(destination_ref_),
|
||||
source_ref(source_ref_),
|
||||
output_op(output_op_),
|
||||
convert_op(convert_op_) {
|
||||
|
||||
}
|
||||
|
||||
/// Constructs a Params object
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(
|
||||
TensorRef destination_ref_,
|
||||
TensorRef source_ref_,
|
||||
typename OutputOp::Params output_op_
|
||||
):
|
||||
Params(
|
||||
destination_ref,
|
||||
source_ref,
|
||||
output_op,
|
||||
ConvertOp::Params()
|
||||
) { }
|
||||
};
|
||||
|
||||
/// Shared storage allocation needed by the epilogue
|
||||
struct SharedStorage { };
|
||||
|
||||
private:
|
||||
|
||||
OutputOp output_op; ///< output operator
|
||||
ConvertOp convert_op; ///< conversion operator
|
||||
|
||||
TensorRef destination_ref_; ///< destination tensor reference
|
||||
TensorRef source_ref_; ///< source tensor reference
|
||||
|
||||
MatrixCoord warp_origin_; ///< position of the warp within the threadblock
|
||||
|
||||
public:
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_DEVICE
|
||||
DirectEpilogueTensorOp(
|
||||
Params const ¶ms, ///< Host-constructable params object
|
||||
SharedStorage &shared_storage, ///< Shared storage object
|
||||
int thread_idx, ///< ID of a thread within the threadblock
|
||||
int warp_idx, ///< ID of warp within threadblock
|
||||
int lane_idx ///< Id of thread within warp
|
||||
):
|
||||
output_op(params.output_op),
|
||||
convert_op(params.convert_op),
|
||||
destination_ref_(params.destination_ref),
|
||||
source_ref_(params.source_ref) {
|
||||
|
||||
|
||||
// Compute warp location within threadblock tile by mapping the warp_id to three coordinates:
|
||||
//
|
||||
// _m: the warp's position within the threadblock along the M dimension
|
||||
// _n: the warp's position within the threadblock along the N dimension
|
||||
// _k: the warp's position within the threadblock along the K dimension
|
||||
|
||||
int warp_mn = warp_idx % (WarpCount::kM * WarpCount::kN);
|
||||
int warp_m = warp_mn % WarpCount::kM;
|
||||
int warp_n = warp_mn / WarpCount::kM;
|
||||
|
||||
warp_origin_ = MatrixCoord{
|
||||
warp_m * Operator::Shape::kM,
|
||||
warp_n * Operator::Shape::kN
|
||||
};
|
||||
|
||||
destination_ref_.add_coord_offset(warp_origin_);
|
||||
source_ref_.add_coord_offset(warp_origin_);
|
||||
}
|
||||
|
||||
/// Streams the result to global memory
|
||||
CUTLASS_DEVICE
|
||||
void operator()(
|
||||
gemm::GemmCoord problem_size, ///< Problem size of GEMM (units of ElementC)
|
||||
gemm::GemmCoord tb_tile_coord, ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles)
|
||||
FragmentC const &accumulators) { ///< Accumulator tile
|
||||
|
||||
MatrixCoord thread_origin =
|
||||
MatrixCoord{tb_tile_coord.m() * Shape::kM, tb_tile_coord.n() * Shape::kN} + warp_origin_;
|
||||
|
||||
/// Number of mma operations performed
|
||||
using MmaIterations = MatrixShape<
|
||||
Operator::Shape::kM / Operator::Policy::Operator::Shape::kM,
|
||||
Operator::Shape::kN / Operator::Policy::Operator::Shape::kN
|
||||
>;
|
||||
|
||||
// Assume accumulator tile is an arrangement of 8-by-8 tiles replicated over the entire
|
||||
// shape, with each quad mapped to one row and each thread mapped to 1/4 of the elements
|
||||
// of that row. The accumulators within one row are assumed to be consecutive.
|
||||
int const kElementsPerAccess = Operator::Policy::Operator::Shape::kN / 4;
|
||||
int const kRowsPerTile = 8;
|
||||
int const kAccumulatorRows = Operator::Policy::Operator::Shape::kM / kRowsPerTile;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int mma_n = 0; mma_n < MmaIterations::kN; ++mma_n) {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int mma_m = 0; mma_m < MmaIterations::kM; ++mma_m) {
|
||||
|
||||
int mma_accum_start = kAccumulatorRows * kElementsPerAccess *
|
||||
(mma_m * MmaIterations::kN + mma_n);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int row = 0; row < kAccumulatorRows; ++row) {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int col = 0; col < kElementsPerAccess; ++col) {
|
||||
|
||||
int accum_m = mma_m * Operator::Policy::Operator::Shape::kM + row * kRowsPerTile;
|
||||
int accum_n = mma_n * Operator::Policy::Operator::Shape::kN + col;
|
||||
int idx = mma_accum_start + row * kElementsPerAccess + col;
|
||||
|
||||
MatrixCoord accum_coord = MatrixCoord{accum_m, accum_n};
|
||||
|
||||
MatrixCoord thread_coord = thread_origin + accum_coord;
|
||||
|
||||
if (thread_coord < MatrixCoord{problem_size.m(), problem_size.n()}) {
|
||||
|
||||
typename ConvertOp::result_type converted_accum = convert_op(accumulators[idx]);
|
||||
|
||||
typename OutputOp::result_type output = output_op(converted_accum, source_ref_.at(accum_coord));
|
||||
|
||||
destination_ref_.at(accum_coord) = output;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,300 @@
|
||||
/***************************************************************************************************
|
||||
* 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 Epilogue for threadblock scoped GEMMs using Tensor Ops.
|
||||
|
||||
The epilogue rearranges the result of a matrix product through shared memory to match canonical
|
||||
tensor layouts in global memory. Epilogues support conversion and reduction operations.
|
||||
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <assert.h>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/layout/vector.h"
|
||||
#include "cutlass/layout/tensor.h"
|
||||
#include "cutlass/tensor_coord.h"
|
||||
#include "cutlass/aligned_buffer.h"
|
||||
#include "cutlass/functional.h"
|
||||
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
|
||||
#include "cutlass/transform/pitch_linear_thread_map.h"
|
||||
#include "cutlass/transform/threadblock/regular_tile_iterator.h"
|
||||
|
||||
#include "cutlass/epilogue/threadblock/epilogue_base.h"
|
||||
#include "cutlass/epilogue/threadblock/predicated_tile_iterator.h"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace threadblock {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Epilogue operator without splitk
|
||||
template <
|
||||
typename Shape_, ///< Shape of threadblock tile (concept: GemmShape)
|
||||
typename WarpMmaOperator_, ///< Warp-level MMA operator (concept: gemm::warp::MmaTensorOp)
|
||||
int PartitionsK, ///< Number of partitions of the K dimension
|
||||
typename OutputTileIterator_, ///< Tile iterator reading and writing output tensors
|
||||
typename AccumulatorFragmentIterator_, ///< Fragment iterator selecting accumulators
|
||||
typename WarpTileIterator_, ///< Warp-scoped tile iterator writing accumulators to SMEM
|
||||
typename SharedLoadIterator_, ///< Threadblock-scoped tile iterator loading from SMEM
|
||||
typename OutputOp_, ///< Output operator
|
||||
typename Padding_ ///< Padding added to SMEM allocation to avoid bank conflicts (concept: MatrixShape)
|
||||
>
|
||||
class Epilogue :
|
||||
public EpilogueBase<
|
||||
Shape_,
|
||||
WarpMmaOperator_,
|
||||
PartitionsK,
|
||||
AccumulatorFragmentIterator_,
|
||||
WarpTileIterator_,
|
||||
Padding_> {
|
||||
|
||||
public:
|
||||
|
||||
using Base = EpilogueBase<
|
||||
Shape_,
|
||||
WarpMmaOperator_,
|
||||
PartitionsK,
|
||||
AccumulatorFragmentIterator_,
|
||||
WarpTileIterator_,
|
||||
Padding_>;
|
||||
|
||||
using Shape = Shape_;
|
||||
using WarpMmaOperator = WarpMmaOperator_;
|
||||
static int const kPartitionsK = PartitionsK;
|
||||
using OutputTileIterator = OutputTileIterator_;
|
||||
using AccumulatorFragmentIterator = AccumulatorFragmentIterator_;
|
||||
using WarpTileIterator = WarpTileIterator_;
|
||||
using SharedLoadIterator = SharedLoadIterator_;
|
||||
using OutputOp = OutputOp_;
|
||||
using Padding = Padding_;
|
||||
|
||||
/// Output layout is always row-major
|
||||
using Layout = layout::RowMajor;
|
||||
using LongIndex = typename Layout::LongIndex;
|
||||
|
||||
/// The complete warp-level accumulator tile
|
||||
using AccumulatorTile = typename Base::AccumulatorTile;
|
||||
|
||||
/// Accumulator element
|
||||
using ElementAccumulator = typename WarpTileIterator::Element;
|
||||
|
||||
|
||||
/// Output element
|
||||
using ElementOutput = typename OutputTileIterator::Element;
|
||||
|
||||
/// Output access size
|
||||
static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess;
|
||||
|
||||
/// Tensor reference to destination tensor
|
||||
using TensorRef = typename OutputTileIterator::TensorRef;
|
||||
|
||||
/// Tensor reference to sync tensor
|
||||
using SyncTensorRef = typename cutlass::TensorRef<int, cutlass::layout::PackedVectorLayout>;
|
||||
|
||||
/// Const tensor reference to source tensor
|
||||
using ConstTensorRef = typename OutputTileIterator::ConstTensorRef;
|
||||
|
||||
/// Array type used to output
|
||||
using OutputAccessType = Array<
|
||||
typename OutputTileIterator::Element, OutputTileIterator::kElementsPerAccess>;
|
||||
|
||||
/// Array type used by output functor
|
||||
using AccumulatorAccessType = Array<typename WarpTileIterator::Element, OutputTileIterator::kElementsPerAccess>;
|
||||
|
||||
/// Number of warps
|
||||
using WarpCount = typename Base::WarpCount;
|
||||
|
||||
public:
|
||||
|
||||
|
||||
static_assert(SharedLoadIterator::Fragment::kElements == OutputTileIterator::Fragment::kElements,
|
||||
"Mismatch between shared load iterator and output tile iterator.");
|
||||
|
||||
static_assert(OutputTileIterator::kElementsPerAccess, "OutputTileIterator::kElementsPerAccess must not be zero.");
|
||||
|
||||
static_assert(!(OutputTileIterator::Fragment::kElements % OutputTileIterator::kElementsPerAccess),
|
||||
"Divisibility");
|
||||
|
||||
private:
|
||||
|
||||
/// Loads fragment from shared memory aligned with output tensor
|
||||
SharedLoadIterator shared_load_iterator_;
|
||||
|
||||
public:
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_DEVICE
|
||||
Epilogue(
|
||||
typename Base::SharedStorage &shared_storage, ///< Shared storage object
|
||||
int thread_idx, ///< ID of a thread within the threadblock
|
||||
int warp_idx, ///< ID of warp within threadblock
|
||||
int lane_idx ///< Id of thread within warp
|
||||
):
|
||||
Base(shared_storage, thread_idx, warp_idx, lane_idx),
|
||||
shared_load_iterator_(shared_storage.reference(), thread_idx) { }
|
||||
|
||||
/// Streams the result to global memory
|
||||
CUTLASS_DEVICE
|
||||
void operator()(
|
||||
OutputOp const &output_op, ///< Output operator
|
||||
OutputTileIterator destination_iterator, ///< Tile iterator for destination
|
||||
AccumulatorTile const &accumulators, ///< Complete warp-level accumulator tile
|
||||
OutputTileIterator source_iterator) { ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles)
|
||||
|
||||
|
||||
typename OutputTileIterator::Fragment source_fragment;
|
||||
|
||||
if (!output_op.is_source_needed()) {
|
||||
source_iterator.clear_mask();
|
||||
}
|
||||
|
||||
source_fragment.clear();
|
||||
|
||||
//
|
||||
// Iterator over warp-level accumulator fragment
|
||||
//
|
||||
|
||||
AccumulatorFragmentIterator accum_fragment_iterator(accumulators);
|
||||
|
||||
//
|
||||
// Iterate over accumulator tile
|
||||
//
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int iter = 0; iter < OutputTileIterator::kIterations; ++iter) {
|
||||
|
||||
//
|
||||
// Load the source
|
||||
//
|
||||
|
||||
source_iterator.load(source_fragment);
|
||||
++source_iterator;
|
||||
|
||||
//
|
||||
// Convert and store fragment
|
||||
//
|
||||
|
||||
__syncthreads();
|
||||
|
||||
typename AccumulatorFragmentIterator::Fragment accum_fragment;
|
||||
|
||||
accum_fragment_iterator.load(accum_fragment);
|
||||
++accum_fragment_iterator;
|
||||
|
||||
this->warp_tile_iterator_.store(accum_fragment);
|
||||
|
||||
__syncthreads();
|
||||
|
||||
//
|
||||
// Load fragments from shared memory
|
||||
//
|
||||
|
||||
typename SharedLoadIterator::Fragment aligned_accum_fragment[kPartitionsK];
|
||||
|
||||
shared_load_iterator_.load(aligned_accum_fragment[0]);
|
||||
|
||||
// If the number of k-slices is > 1 - perform a reduction amongst the k-slices
|
||||
if (kPartitionsK > 1)
|
||||
{
|
||||
plus <typename SharedLoadIterator::Fragment> add_fragments;
|
||||
const int tile_row_offset = Base::SharedStorage::StorageShape::kRow / PartitionsK;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for ( int i = 1; i < kPartitionsK; ++i) {
|
||||
shared_load_iterator_.add_tile_offset({tile_row_offset , 0});
|
||||
shared_load_iterator_.load(aligned_accum_fragment[i]);
|
||||
aligned_accum_fragment[0] = add_fragments(aligned_accum_fragment[0], aligned_accum_fragment[i]);
|
||||
}
|
||||
|
||||
shared_load_iterator_.add_tile_offset({-1 * (kPartitionsK-1) * tile_row_offset, 0});
|
||||
}
|
||||
|
||||
//
|
||||
// Compute the output result
|
||||
//
|
||||
|
||||
typename OutputTileIterator::Fragment output_fragment;
|
||||
|
||||
apply_output_operator_(output_fragment, output_op, aligned_accum_fragment[0], source_fragment);
|
||||
|
||||
|
||||
//
|
||||
// Store the final result
|
||||
//
|
||||
|
||||
destination_iterator.store(output_fragment);
|
||||
++destination_iterator;
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
|
||||
/// Helper to invoke the output functor over each vector of output
|
||||
CUTLASS_DEVICE
|
||||
void apply_output_operator_(
|
||||
typename OutputTileIterator::Fragment &output_fragment,
|
||||
OutputOp const &output_op, ///< Output operator
|
||||
typename SharedLoadIterator::Fragment const &aligned_accum_fragment,
|
||||
typename OutputTileIterator::Fragment const &source_fragment) {
|
||||
|
||||
OutputAccessType *output_frag_ptr =
|
||||
reinterpret_cast<OutputAccessType *>(&output_fragment);
|
||||
|
||||
AccumulatorAccessType const *compute_frag_ptr =
|
||||
reinterpret_cast<AccumulatorAccessType const *>(&aligned_accum_fragment);
|
||||
|
||||
OutputAccessType const *source_frag_ptr =
|
||||
reinterpret_cast<OutputAccessType const *>(&source_fragment);
|
||||
|
||||
int const kOutputOpIterations =
|
||||
OutputTileIterator::Fragment::kElements / OutputTileIterator::kElementsPerAccess;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kOutputOpIterations; ++i) {
|
||||
|
||||
// Call the output operator
|
||||
output_frag_ptr[i] = output_op(compute_frag_ptr[i], source_frag_ptr[i]);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,214 @@
|
||||
/***************************************************************************************************
|
||||
* 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 Epilogue for threadblock scoped GEMMs using Tensor Ops.
|
||||
|
||||
The epilogue rearranges the result of a matrix product through shared memory to match canonical
|
||||
tensor layouts in global memory. Epilogues support conversion and reduction operations.
|
||||
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <assert.h>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/matrix_shape.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/layout/vector.h"
|
||||
#include "cutlass/layout/tensor.h"
|
||||
#include "cutlass/tensor_coord.h"
|
||||
#include "cutlass/aligned_buffer.h"
|
||||
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
|
||||
#include "cutlass/transform/pitch_linear_thread_map.h"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace threadblock {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Base class for epilogues defining warp-level
|
||||
template <
|
||||
typename Shape_, ///< Shape of threadblock tile (concept: GemmShape)
|
||||
typename WarpMmaOperator_, ///< Warp-level MMA operator (concept: gemm::warp::MmaTensorOp)
|
||||
int PartitionsK, ///< Number of partitions of the K dimension
|
||||
typename AccumulatorFragmentIterator_, ///< Fragment iterator selecting accumulators
|
||||
typename WarpTileIterator_, ///< Warp-scoped tile iterator writing accumulators to SMEM
|
||||
typename Padding_ ///< Padding added to SMEM allocation to avoid bank conflicts (concept: MatrixShape)
|
||||
>
|
||||
class EpilogueBase {
|
||||
public:
|
||||
|
||||
using Shape = Shape_;
|
||||
using WarpMmaOperator = WarpMmaOperator_;
|
||||
static int const kPartitionsK = PartitionsK;
|
||||
using AccumulatorFragmentIterator = AccumulatorFragmentIterator_;
|
||||
using WarpTileIterator = WarpTileIterator_;
|
||||
using Padding = Padding_;
|
||||
|
||||
/// Output layout is always row-major
|
||||
using Layout = layout::RowMajor;
|
||||
|
||||
/// The complete warp-level accumulator tile
|
||||
using AccumulatorTile = typename AccumulatorFragmentIterator::AccumulatorTile;
|
||||
|
||||
/// Accumulator element
|
||||
using ElementAccumulator = typename AccumulatorTile::Element;
|
||||
|
||||
|
||||
/// Number of warps
|
||||
using WarpCount = gemm::GemmShape<
|
||||
Shape::kM / WarpMmaOperator::Shape::kM,
|
||||
Shape::kN / WarpMmaOperator::Shape::kN,
|
||||
kPartitionsK
|
||||
>;
|
||||
|
||||
public:
|
||||
|
||||
/// Shared storage allocation needed by the epilogue
|
||||
struct SharedStorage {
|
||||
|
||||
//
|
||||
// Type definitions
|
||||
//
|
||||
|
||||
/// Element type of shared memory
|
||||
using Element = typename WarpTileIterator::Element;
|
||||
|
||||
/// Tensor reference to shared memory allocation
|
||||
using TensorRef = typename WarpTileIterator::TensorRef;
|
||||
|
||||
/// Layout of shared memory allocation
|
||||
using Layout = typename WarpTileIterator::Layout;
|
||||
|
||||
/// Logical shape of the shared memory tile written to by all warps.
|
||||
using Shape = MatrixShape<
|
||||
WarpCount::kM * WarpTileIterator::Shape::kRow * WarpCount::kK,
|
||||
WarpCount::kN * WarpTileIterator::Shape::kColumn
|
||||
>;
|
||||
|
||||
/// Shape of the shared memory allocation for the epilogue
|
||||
using StorageShape = MatrixShape<
|
||||
Shape::kRow + Padding::kRow,
|
||||
Shape::kColumn + Padding::kColumn
|
||||
>;
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
AlignedBuffer<Element, StorageShape::kCount> storage;
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
/// Returns a pointer to the shared memory buffer
|
||||
CUTLASS_DEVICE
|
||||
Element *data() {
|
||||
return storage.data();
|
||||
}
|
||||
|
||||
/// Returns a tensor reference to the shared memory buffer
|
||||
CUTLASS_DEVICE
|
||||
TensorRef reference() {
|
||||
return TensorRef(
|
||||
storage.data(),
|
||||
Layout::packed({StorageShape::kRow, StorageShape::kColumn}));
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void debug_print() {
|
||||
if (threadIdx.x == 0) {
|
||||
|
||||
#pragma unroll 1
|
||||
for (int r = 0; r < Shape::kRow; ++r) {
|
||||
|
||||
#pragma unroll 1
|
||||
for (int c = 0; c < Shape::kColumn; ++c) {
|
||||
|
||||
printf("%d ", int(storage.data()[r * StorageShape::kColumn + c]));
|
||||
}
|
||||
printf("\n");
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
};
|
||||
|
||||
protected:
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
SharedStorage &shared_storage_;
|
||||
|
||||
/// Stores a warp's fragment of accumulators to SMEM
|
||||
WarpTileIterator warp_tile_iterator_;
|
||||
|
||||
public:
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_DEVICE
|
||||
EpilogueBase(
|
||||
SharedStorage &shared_storage, ///< Shared storage object
|
||||
int thread_idx, ///< ID of a thread within the threadblock
|
||||
int warp_idx, ///< ID of warp within threadblock
|
||||
int lane_idx ///< Id of thread within warp
|
||||
):
|
||||
shared_storage_(shared_storage),
|
||||
warp_tile_iterator_(shared_storage.reference(), lane_idx) {
|
||||
|
||||
// Compute warp location within threadblock tile by mapping the warp_id to three coordinates:
|
||||
//
|
||||
// _m: the warp's position within the threadblock along the M dimension
|
||||
// _n: the warp's position within the threadblock along the N dimension
|
||||
// _k: the warp's position within the threadblock along the K dimension
|
||||
|
||||
int warp_k = warp_idx / (WarpCount::kM * WarpCount::kN);
|
||||
int warp_mn = warp_idx % (WarpCount::kM * WarpCount::kN);
|
||||
int warp_m = warp_mn % WarpCount::kM;
|
||||
int warp_n = warp_mn / WarpCount::kM;
|
||||
|
||||
MatrixCoord warp_offset{warp_k * WarpCount::kM + warp_m, warp_n};
|
||||
|
||||
warp_tile_iterator_.add_tile_offset(warp_offset);
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,191 @@
|
||||
/***************************************************************************************************
|
||||
* 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 Epilogue for threadblock scoped GEMMs.
|
||||
|
||||
This does not attempt to target any particular output layout. Instead, each threadblock
|
||||
streams out its accumulator elements using 128b store operations. This assumes all threadblocks
|
||||
have unique output tiles.
|
||||
|
||||
The target data layout is:
|
||||
- threadblock indices mapped to linear offsets as (m, n, k), where m is fastest-changing
|
||||
- threadblock output space partitioned into warps; each warp's region is contiguous
|
||||
- per-thread accumulators partitioned into 128b accesses
|
||||
- output memory striped across the threads of a warp
|
||||
|
||||
This enables very fast streaming of data, completely limited by the memory system. No predication
|
||||
or data exchange is performed, and each threadblock is assumed to have a full region of memory
|
||||
to write to.
|
||||
|
||||
This epilogue establishes an upper bound for epilogue performance and is suitable for
|
||||
reductions across the GEMM K dimension which require a separate workspace.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/array.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
typename Shape_, ///< shape of accumulator tile (concept: MatrixShape)
|
||||
int WarpCount, ///< number of warps
|
||||
typename FragmentC_ ///< warp-level GEMM operator (concept: gemm::warp::Mma)
|
||||
>
|
||||
class EpilogueWorkspace {
|
||||
public:
|
||||
|
||||
using Shape = Shape_;
|
||||
using FragmentC = FragmentC_;
|
||||
using ElementC = typename FragmentC::value_type;
|
||||
|
||||
static int const kWarpCount = WarpCount;
|
||||
|
||||
/// Optimize for 128b accesses
|
||||
static int const kAccessSizeInBits = 128;
|
||||
|
||||
/// Warp size from the perspective of memory operations
|
||||
static int const kWarpSize = 32;
|
||||
|
||||
/// Vector length of accesses
|
||||
static int const kElementsPerAccess =
|
||||
kAccessSizeInBits / sizeof_bits<ElementC>::value;
|
||||
|
||||
/// Number of stores per thread
|
||||
static int const kIterations = FragmentC::kElements / kElementsPerAccess;
|
||||
|
||||
static_assert(
|
||||
!(FragmentC::kElements % kElementsPerAccess),
|
||||
"The number of accumulators must be divisible by the access size.");
|
||||
|
||||
/// Total number of vectorized accesses in warp (in units of vector)
|
||||
static int const kWarpAccesses = kIterations * kWarpSize;
|
||||
|
||||
/// Total number of vectorized accesses in threadblock tile (in units of vector)
|
||||
static int const kThreadblockAccesses = kWarpAccesses * kWarpCount;
|
||||
|
||||
/// Parameters structure
|
||||
struct Params {
|
||||
|
||||
/// Pointer to C matrix
|
||||
ElementC *ptr_C;
|
||||
|
||||
/// Stride between tiles along the GEMM N dimension (in units of vectors)
|
||||
int stride_n;
|
||||
|
||||
/// Stride between tiles along the GEMM K dimension (in units of vectors)
|
||||
int stride_k;
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(
|
||||
ElementC *ptr_C, ///< Pointer to C matrix
|
||||
int stride_n_, ///< Stride between tiles along the GEMM N dimension (in units of ElementC)
|
||||
int stride_k_ ///< Stride between tiles along the GEMM K dimension (in units of ElementC)
|
||||
):
|
||||
ptr_C(ptr_C), stride_n(stride_n_ / kElementsPerAccess), stride_k(stride_k_ / kElementsPerAccess) {
|
||||
|
||||
}
|
||||
};
|
||||
|
||||
/// Shared storage allocation needed by the epilogue
|
||||
struct SharedStorage {
|
||||
// Intentionally empty
|
||||
};
|
||||
|
||||
private:
|
||||
|
||||
struct alignas((kAccessSizeInBits / 8)) AccessType {
|
||||
Array<ElementC, kElementsPerAccess> storage;
|
||||
};
|
||||
|
||||
/// Constant reference to parameters object
|
||||
AccessType *pointer_;
|
||||
|
||||
/// Stride between tiles along the n dimension (in vectors)
|
||||
int stride_n_;
|
||||
|
||||
/// Stride between tiles along the k dimension (in vectors)
|
||||
int stride_k_;
|
||||
|
||||
public:
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_DEVICE
|
||||
EpilogueWorkspace(
|
||||
Params const ¶ms, ///< Host-constructable params object
|
||||
SharedStorage &, ///< Shared storage object
|
||||
int warp_idx, ///< ID of warp within threadblock
|
||||
int lane_idx ///< Id of thread within warp
|
||||
|
||||
):
|
||||
pointer_(reinterpret_cast<AccessType *>(params.ptr_C)),
|
||||
stride_n_(params.stride_n),
|
||||
stride_k_(params.stride_k) {
|
||||
|
||||
// Add per-thread offset
|
||||
pointer_ += lane_idx + warp_idx * kWarpAccesses;
|
||||
}
|
||||
|
||||
/// Streams the result to global memory
|
||||
CUTLASS_DEVICE
|
||||
void operator()(
|
||||
cutlass::gemm::GemmCoord problem_size, ///< Problem size of GEMM (units of ElementC)
|
||||
cutlass::gemm::GemmCoord tb_tile_coord, ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles)
|
||||
FragmentC const &accum) { ///< Accumulator tile
|
||||
|
||||
// Compute offset for entire threadblock (note, per-thread offset has been folded in already)
|
||||
AccessType *pointer = pointer_ +
|
||||
tb_tile_coord.m() * kThreadblockAccesses +
|
||||
tb_tile_coord.n() * stride_n_ +
|
||||
tb_tile_coord.k() * stride_k_;
|
||||
|
||||
// Cast to vectorized view of accumulator fragments
|
||||
AccessType const * src_pointer = reinterpret_cast<AccessType const *>(&accum);
|
||||
|
||||
// Write out accumulators at full speed
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kIterations; ++i) {
|
||||
pointer[i * kWarpSize] = src_pointer[i];
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,257 @@
|
||||
/***************************************************************************************************
|
||||
* 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 Epilogue for threadblock scoped GEMMs using Tensor Ops.
|
||||
|
||||
The epilogue rearranges the result of a matrix product through shared memory to match canonical
|
||||
tensor layouts in global memory. Epilogues support conversion and reduction operations.
|
||||
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <assert.h>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/layout/vector.h"
|
||||
#include "cutlass/layout/tensor.h"
|
||||
#include "cutlass/tensor_coord.h"
|
||||
#include "cutlass/aligned_buffer.h"
|
||||
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
|
||||
#include "cutlass/transform/pitch_linear_thread_map.h"
|
||||
#include "cutlass/transform/threadblock/regular_tile_iterator.h"
|
||||
|
||||
#include "cutlass/epilogue/threadblock/epilogue_base.h"
|
||||
#include "cutlass/epilogue/threadblock/predicated_tile_iterator.h"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace threadblock {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Epilogue operator without splitk
|
||||
template <
|
||||
/// Shape of threadblock tile (concept: GemmShape)
|
||||
typename Shape_,
|
||||
/// Warp-level MMA operator (concept: gemm::warp::MmaTensorOp)
|
||||
typename WarpMmaOperator_,
|
||||
/// Number of partitions of the K dimension
|
||||
int PartitionsK,
|
||||
/// Tile iterator reading and writing output tensors
|
||||
typename OutputTileIterator_,
|
||||
/// Fragment iterator selecting accumulators
|
||||
typename AccumulatorFragmentIterator_,
|
||||
/// Output operator
|
||||
typename OutputOp_,
|
||||
/// Number of interleaved k
|
||||
int InterleavedK,
|
||||
/// Whether Beta is zero
|
||||
bool IsBetaZero = false>
|
||||
class InterleavedEpilogue {
|
||||
public:
|
||||
using Shape = Shape_;
|
||||
using WarpMmaOperator = WarpMmaOperator_;
|
||||
static int const kPartitionsK = PartitionsK;
|
||||
using AccumulatorFragmentIterator = AccumulatorFragmentIterator_;
|
||||
using OutputTileIterator = OutputTileIterator_;
|
||||
using OutputOp = OutputOp_;
|
||||
|
||||
/// Output layout is always row-major
|
||||
using Layout = layout::ColumnMajorInterleaved<InterleavedK>;
|
||||
|
||||
/// The complete warp-level accumulator tile
|
||||
using AccumulatorTile = typename AccumulatorFragmentIterator::AccumulatorTile;
|
||||
|
||||
/// Accumulator element
|
||||
using ElementAccumulator = typename AccumulatorTile::Element;
|
||||
|
||||
/// Output element
|
||||
using ElementOutput = typename OutputTileIterator::Element;
|
||||
|
||||
/// Output access size
|
||||
static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess;
|
||||
|
||||
/// Tensor reference to destination tensor
|
||||
using TensorRef = typename OutputTileIterator::TensorRef;
|
||||
|
||||
/// Tensor reference to sync tensor
|
||||
using SyncTensorRef =
|
||||
typename cutlass::TensorRef<int, cutlass::layout::PackedVectorLayout>;
|
||||
|
||||
/// Const tensor reference to source tensor
|
||||
using ConstTensorRef = typename OutputTileIterator::ConstTensorRef;
|
||||
|
||||
/// Array type used to output
|
||||
using OutputAccessType = Array<typename OutputTileIterator::Element,
|
||||
OutputTileIterator::kElementsPerAccess>;
|
||||
|
||||
/// Array type used by output functor
|
||||
using AccumulatorAccessType =
|
||||
Array<ElementAccumulator, OutputTileIterator::kElementsPerAccess>;
|
||||
|
||||
/// Number of warps
|
||||
using WarpCount =
|
||||
gemm::GemmShape<Shape::kM / WarpMmaOperator::Shape::kM,
|
||||
Shape::kN / WarpMmaOperator::Shape::kN, kPartitionsK>;
|
||||
|
||||
public:
|
||||
static_assert(OutputTileIterator::kElementsPerAccess,
|
||||
"This must not be zero.");
|
||||
|
||||
static_assert(!(OutputTileIterator::Fragment::kElements %
|
||||
OutputTileIterator::kElementsPerAccess),
|
||||
"Divisibility");
|
||||
|
||||
/// Shared storage allocation needed by the epilogue
|
||||
struct SharedStorage {};
|
||||
|
||||
|
||||
public:
|
||||
/// Constructor
|
||||
CUTLASS_DEVICE
|
||||
InterleavedEpilogue(
|
||||
SharedStorage &shared_storage, ///< Shared storage object
|
||||
int thread_idx, ///< ID of a thread within the threadblock
|
||||
int warp_idx, ///< ID of warp within threadblock
|
||||
int lane_idx ///< Id of thread within warp
|
||||
) {}
|
||||
|
||||
/// Streams the result to global memory
|
||||
CUTLASS_DEVICE
|
||||
void operator()(
|
||||
OutputOp const &output_op, ///< Output operator
|
||||
OutputTileIterator destination_iterator, ///< Tile iterator for destination
|
||||
AccumulatorTile const &accumulators, ///< Complete warp-level accumulator tile
|
||||
OutputTileIterator source_iterator) { ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles)
|
||||
|
||||
//
|
||||
// Predicated tile iterators constructed from members
|
||||
//
|
||||
|
||||
if (IsBetaZero && output_op.is_source_needed())
|
||||
assert(0);
|
||||
|
||||
typename OutputTileIterator::Fragment source_fragment;
|
||||
|
||||
if (!IsBetaZero) {
|
||||
if (!output_op.is_source_needed()) {
|
||||
source_iterator.clear_mask();
|
||||
}
|
||||
}
|
||||
|
||||
source_fragment.clear();
|
||||
|
||||
//
|
||||
// Iterator over warp-level accumulator fragment
|
||||
//
|
||||
|
||||
AccumulatorFragmentIterator accum_fragment_iterator(accumulators);
|
||||
|
||||
//
|
||||
// Iterate over accumulator tile
|
||||
//
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int iter = 0; iter < OutputTileIterator::kIterations; ++iter) {
|
||||
//
|
||||
// Load the source
|
||||
//
|
||||
|
||||
if (!IsBetaZero) {
|
||||
source_iterator.set_iteration_index(iter);
|
||||
source_iterator.load(source_fragment);
|
||||
++source_iterator;
|
||||
}
|
||||
|
||||
//
|
||||
// Convert fragment
|
||||
//
|
||||
|
||||
typename AccumulatorFragmentIterator::Fragment accum_fragment;
|
||||
|
||||
accum_fragment_iterator.load(accum_fragment);
|
||||
++accum_fragment_iterator;
|
||||
|
||||
//
|
||||
// Compute the output result
|
||||
//
|
||||
|
||||
typename OutputTileIterator::Fragment output_fragment;
|
||||
apply_output_operator_(output_op, output_fragment, accum_fragment, source_fragment);
|
||||
|
||||
//
|
||||
// Store the final result
|
||||
//
|
||||
|
||||
destination_iterator.set_iteration_index(iter);
|
||||
destination_iterator.store(output_fragment);
|
||||
++destination_iterator;
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
/// Helper to invoke the output functor over each vector of output
|
||||
CUTLASS_DEVICE
|
||||
void apply_output_operator_(
|
||||
OutputOp const &output_op, ///< Output operator
|
||||
typename OutputTileIterator::Fragment &output_fragment,
|
||||
typename AccumulatorFragmentIterator::Fragment const
|
||||
&aligned_accum_fragment,
|
||||
typename OutputTileIterator::Fragment const &source_fragment) {
|
||||
OutputAccessType *output_frag_ptr =
|
||||
reinterpret_cast<OutputAccessType *>(&output_fragment);
|
||||
|
||||
AccumulatorAccessType const *compute_frag_ptr =
|
||||
reinterpret_cast<AccumulatorAccessType const *>(
|
||||
&aligned_accum_fragment);
|
||||
|
||||
OutputAccessType const *source_frag_ptr =
|
||||
reinterpret_cast<OutputAccessType const *>(&source_fragment);
|
||||
|
||||
int const kOutputOpIterations = OutputTileIterator::Fragment::kElements /
|
||||
OutputTileIterator::kElementsPerAccess;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kOutputOpIterations; ++i) {
|
||||
// Call the output operator
|
||||
output_frag_ptr[i] = output_op(compute_frag_ptr[i], source_frag_ptr[i]);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,496 @@
|
||||
/***************************************************************************************************
|
||||
* 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 Metaprogram for determining the mapping of output elements to threads for epilogue tiles.
|
||||
|
||||
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/layout/matrix.h"
|
||||
#include "cutlass/matrix_shape.h"
|
||||
#include "cutlass/tensor_ref.h"
|
||||
#include "cutlass/fast_math.h"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace threadblock {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Tuple defining point in output tile
|
||||
template <
|
||||
int Column,
|
||||
int Row,
|
||||
int Group,
|
||||
int Cluster,
|
||||
int Tile
|
||||
>
|
||||
struct OutputTileShape {
|
||||
static int const kColumn = Column;
|
||||
static int const kRow = Row;
|
||||
static int const kGroup = Group;
|
||||
static int const kCluster = Cluster;
|
||||
static int const kTile = Tile;
|
||||
|
||||
static int const kCount = kColumn * kRow * kGroup * kCluster * kTile;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
typename ThreadMap_,
|
||||
typename Shape_,
|
||||
typename Iterations_,
|
||||
typename Delta_,
|
||||
typename Count_
|
||||
>
|
||||
struct OutputTileThreadMap {
|
||||
|
||||
/// Conventional thread map (concept: ThreadMap)
|
||||
using ThreadMap = ThreadMap_;
|
||||
|
||||
/// Number of threads participating in the operation
|
||||
static int const kThreads = ThreadMap::kThreads;
|
||||
|
||||
/// Number of scalar elements per access
|
||||
static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
|
||||
|
||||
/// Shape of the tile
|
||||
using Shape = Shape_;
|
||||
|
||||
/// Iterations performed by each thread
|
||||
using Iterations = Iterations_;
|
||||
|
||||
/// Delta between accesses
|
||||
using Delta = Delta_;
|
||||
|
||||
/// Number of iterator iterations
|
||||
using Count = Count_;
|
||||
|
||||
/// Initial offset function
|
||||
CUTLASS_HOST_DEVICE
|
||||
static MatrixCoord initial_offset(int thread_idx) {
|
||||
|
||||
using Index = typename layout::PitchLinearCoord::Index;
|
||||
|
||||
layout::PitchLinearCoord coord = ThreadMap::initial_offset(thread_idx);
|
||||
|
||||
Index cluster = coord.strided() / (Shape::kGroup * Shape::kRow);
|
||||
Index cluster_residual = coord.strided() % (Shape::kGroup * Shape::kRow);
|
||||
|
||||
Index group = cluster_residual / (Shape::kRow);
|
||||
Index row = cluster_residual % (Shape::kRow);
|
||||
|
||||
return MatrixCoord{
|
||||
row + group * Shape::kRow * Count::kRow
|
||||
+ cluster * Shape::kGroup * Count::kGroup * Shape::kRow * Count::kRow,
|
||||
coord.contiguous()
|
||||
};
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace detail {
|
||||
|
||||
/// RowArrangement determines how one or more warps cover a region of consecutive rows.
|
||||
template <
|
||||
typename Shape,
|
||||
int WarpsRemaining,
|
||||
int ElementsPerAccess,
|
||||
int ElementSize,
|
||||
bool Is2dTile
|
||||
>
|
||||
struct RowArrangement;
|
||||
|
||||
/// RowArrangement in which each warp's access is a 1D tiled arrangement.
|
||||
template <
|
||||
typename Shape,
|
||||
int WarpsRemaining,
|
||||
int ElementsPerAccess,
|
||||
int ElementSize
|
||||
>
|
||||
struct RowArrangement<Shape, WarpsRemaining, ElementsPerAccess, ElementSize, false> {
|
||||
static int const kWarpSize = 32;
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
static int const kElementSize = ElementSize;
|
||||
|
||||
static int const kIterationsRow = 1;
|
||||
static int const kDeltaRow = 1;
|
||||
static int const kIterationsColumn = Shape::kColumn / kElementsPerAccess / kWarpSize;
|
||||
static int const kDeltaColumn = kWarpSize * kElementsPerAccess;
|
||||
|
||||
static int const kAccessWidth = kWarpSize;
|
||||
static int const kAccessRows = 1;
|
||||
static int const kWarpPartitionsRow = 1;
|
||||
static int const kWarpPartitionsColumn = WarpsRemaining;
|
||||
};
|
||||
|
||||
/// RowArrangement in which each warp's access is a 2D tiled arrangement.
|
||||
template <
|
||||
typename Shape,
|
||||
int WarpsRemaining,
|
||||
int ElementsPerAccess,
|
||||
int ElementSize
|
||||
>
|
||||
struct RowArrangement<Shape, WarpsRemaining, ElementsPerAccess, ElementSize, true> {
|
||||
|
||||
static int const kMemoryAccessSize = 128;
|
||||
static int const kWarpSize = 32;
|
||||
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
static int const kElementSize = ElementSize;
|
||||
|
||||
struct Detail {
|
||||
static int const kShapeRow = Shape::kRow / WarpsRemaining;
|
||||
static int const kShapeWidth = Shape::kColumn / kElementsPerAccess;
|
||||
|
||||
static int const kTargetMemoryAccessWidth =
|
||||
kMemoryAccessSize / (kElementsPerAccess * kElementSize / 8);
|
||||
|
||||
static int const kTargetAccessRows = kWarpSize / kTargetMemoryAccessWidth;
|
||||
};
|
||||
|
||||
static int const kAccessWidth =
|
||||
(Detail::kTargetAccessRows > Detail::kShapeRow ?
|
||||
kWarpSize / Detail::kShapeRow
|
||||
: const_min(
|
||||
Detail::kShapeWidth,
|
||||
const_min(kWarpSize, kMemoryAccessSize / (kElementsPerAccess * kElementSize / 8))
|
||||
));
|
||||
|
||||
static int const kAccessRows =
|
||||
(Detail::kTargetAccessRows > Detail::kShapeRow ?
|
||||
Detail::kShapeRow
|
||||
: const_min(Shape::kRow, kWarpSize / kAccessWidth));
|
||||
|
||||
static int const kIterationsRow = Detail::kShapeRow / kAccessRows;
|
||||
static int const kDeltaRow = kAccessRows;
|
||||
|
||||
static int const kIterationsColumn = Detail::kShapeWidth / kAccessWidth;
|
||||
static int const kDeltaColumn = kAccessWidth * kElementsPerAccess;
|
||||
|
||||
static_assert( kAccessWidth * kElementsPerAccess <= Shape::kColumn, "Accessing too many elements per access");
|
||||
static_assert( kIterationsColumn > 0, "Iteration Count Column must be > 0" );
|
||||
static_assert( kIterationsRow > 0, "Iteration Count Row must be > 0" );
|
||||
|
||||
static int const kWarpPartitionsRow = 1;
|
||||
static int const kWarpPartitionsColumn = 1;
|
||||
};
|
||||
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Template metaprogram for partitioning a 4D space across warps to achieve several performance
|
||||
/// objectives:
|
||||
///
|
||||
/// - coalesced memory accesses in units of 128 Byte lines
|
||||
/// - minimal address arithmetic
|
||||
/// - minimal predicate calculations
|
||||
///
|
||||
template <
|
||||
typename Shape_,
|
||||
typename Count_,
|
||||
int Threads,
|
||||
int ElementsPerAccess,
|
||||
int ElementSize
|
||||
>
|
||||
struct OutputTileOptimalThreadMap {
|
||||
|
||||
using Shape = Shape_;
|
||||
using Count = Count_;
|
||||
|
||||
static int const kWarpSize = 32;
|
||||
static int const kThreads = Threads;
|
||||
static int const kWarpCount = kThreads / kWarpSize;
|
||||
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
static int const kElementSize = ElementSize;
|
||||
|
||||
//
|
||||
// Metaprogram computation
|
||||
//
|
||||
|
||||
struct Detail {
|
||||
|
||||
// Clusters
|
||||
static int const kIterationsCluster =
|
||||
((Shape::kCluster > kWarpCount) ?
|
||||
Shape::kCluster / kWarpCount
|
||||
: 1);
|
||||
|
||||
static int const kDeltaCluster =
|
||||
((Shape::kCluster > kWarpCount) ?
|
||||
Shape::kRow * Count::kRow * Shape::kGroup * Count::kGroup * Shape::kCluster / kIterationsCluster
|
||||
: 1);
|
||||
|
||||
static int const kCompactedDeltaCluster =
|
||||
((Shape::kCluster > kWarpCount) ?
|
||||
Shape::kRow * Shape::kGroup * Shape::kCluster / kIterationsCluster
|
||||
: 1);
|
||||
|
||||
static int const kWarpPartitionsCluster =
|
||||
((Shape::kCluster > kWarpCount) ?
|
||||
kWarpCount
|
||||
: kWarpCount / Shape::kCluster);
|
||||
|
||||
static int const kWarpsRemainingForGroups =
|
||||
((Shape::kCluster > kWarpCount) ? 1 : kWarpCount / Shape::kCluster);
|
||||
|
||||
// Groups
|
||||
static int const kIterationsGroup =
|
||||
((Shape::kGroup > kWarpsRemainingForGroups) ?
|
||||
Shape::kGroup / kWarpsRemainingForGroups
|
||||
: 1);
|
||||
|
||||
static int const kDeltaGroup =
|
||||
((Shape::kGroup > kWarpsRemainingForGroups) ?
|
||||
Shape::kRow * Count::kRow * Shape::kGroup / kIterationsGroup
|
||||
: 1);
|
||||
|
||||
static int const kCompactedDeltaGroup =
|
||||
((Shape::kGroup > kWarpsRemainingForGroups) ?
|
||||
Shape::kRow * Shape::kGroup / kIterationsGroup
|
||||
: 1);
|
||||
|
||||
static int const kWarpPartitionsGroup =
|
||||
((Shape::kGroup > kWarpsRemainingForGroups) ?
|
||||
1
|
||||
: kWarpsRemainingForGroups / Shape::kGroup);
|
||||
|
||||
static int const kWarpsRemainingForRows =
|
||||
((Shape::kGroup > kWarpsRemainingForGroups) ?
|
||||
1
|
||||
: kWarpsRemainingForGroups / Shape::kGroup);
|
||||
|
||||
// Rows
|
||||
using RowArrangement = detail::RowArrangement<
|
||||
Shape,
|
||||
kWarpsRemainingForRows,
|
||||
kElementsPerAccess,
|
||||
kElementSize,
|
||||
(Shape::kRow > kWarpsRemainingForRows)
|
||||
>;
|
||||
|
||||
// Warp partitions
|
||||
using WarpPartitions = OutputTileShape<
|
||||
RowArrangement::kWarpPartitionsColumn,
|
||||
RowArrangement::kWarpPartitionsRow,
|
||||
kWarpPartitionsGroup,
|
||||
kWarpPartitionsCluster,
|
||||
1>;
|
||||
|
||||
static int const kAccessWidth = RowArrangement::kAccessWidth;
|
||||
static int const kAccessRows = RowArrangement::kAccessRows;
|
||||
};
|
||||
|
||||
//
|
||||
// Output
|
||||
//
|
||||
|
||||
using Iterations = OutputTileShape<
|
||||
Detail::RowArrangement::kIterationsColumn,
|
||||
Detail::RowArrangement::kIterationsRow,
|
||||
Detail::kIterationsGroup,
|
||||
Detail::kIterationsCluster,
|
||||
1>;
|
||||
|
||||
using Delta = OutputTileShape<
|
||||
Detail::RowArrangement::kDeltaColumn,
|
||||
Detail::RowArrangement::kDeltaRow,
|
||||
Detail::kDeltaGroup,
|
||||
Detail::kDeltaCluster,
|
||||
1>;
|
||||
|
||||
/// Initial offset function
|
||||
CUTLASS_HOST_DEVICE
|
||||
static MatrixCoord initial_offset(int thread_idx) {
|
||||
|
||||
int warp_idx = thread_idx / kWarpSize;
|
||||
int lane_idx = thread_idx % kWarpSize;
|
||||
|
||||
// Compute warp location
|
||||
int cluster_idx = warp_idx / Detail::WarpPartitions::kCluster;
|
||||
int residual_cluster = warp_idx % Detail::WarpPartitions::kCluster;
|
||||
|
||||
int group_idx = residual_cluster / Detail::WarpPartitions::kGroup;
|
||||
int residual_group = residual_cluster % Detail::WarpPartitions::kGroup;
|
||||
|
||||
int row_idx = residual_group / Detail::WarpPartitions::kRow;
|
||||
int col_idx = residual_group % Detail::WarpPartitions::kRow;
|
||||
|
||||
// Compute per-lane offset
|
||||
int lane_row_offset = lane_idx / Detail::kAccessWidth;
|
||||
int lane_col_offset = lane_idx % Detail::kAccessWidth;
|
||||
|
||||
// Compute coordinate in output space
|
||||
int cluster_offset = cluster_idx * Shape::kRow * Count::kRow * Shape::kGroup * Count::kGroup;
|
||||
int group_offset = group_idx * Shape::kRow * Count::kRow;
|
||||
int row_offset = row_idx * Iterations::kRow * Detail::kAccessRows;
|
||||
int column_offset = col_idx * Iterations::kColumn * Detail::kAccessWidth * kElementsPerAccess;
|
||||
|
||||
return MatrixCoord(
|
||||
cluster_offset + group_offset + row_offset + lane_row_offset,
|
||||
(column_offset + lane_col_offset) * kElementsPerAccess
|
||||
);
|
||||
}
|
||||
|
||||
/// Compacted thread map in which the 4D region is contiguous
|
||||
struct CompactedThreadMap {
|
||||
|
||||
|
||||
using Shape = Shape_;
|
||||
|
||||
using Iterations = OutputTileShape<
|
||||
Detail::RowArrangement::kIterationsColumn,
|
||||
Detail::RowArrangement::kIterationsRow,
|
||||
Detail::kIterationsGroup,
|
||||
Detail::kIterationsCluster,
|
||||
1>;
|
||||
|
||||
using Delta = OutputTileShape<
|
||||
Detail::RowArrangement::kDeltaColumn,
|
||||
Detail::RowArrangement::kDeltaRow,
|
||||
Detail::kCompactedDeltaGroup,
|
||||
Detail::kCompactedDeltaCluster,
|
||||
1>;
|
||||
|
||||
/// Number of elements within each vector access
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
|
||||
/// Number of threads
|
||||
static int const kThreads = Threads;
|
||||
|
||||
/// Function to compute each thread's initial offset
|
||||
CUTLASS_HOST_DEVICE
|
||||
static MatrixCoord initial_offset(int thread_idx) {
|
||||
|
||||
int warp_idx = thread_idx / kWarpSize;
|
||||
int lane_idx = thread_idx % kWarpSize;
|
||||
|
||||
// Compute warp location
|
||||
int cluster_idx = warp_idx / Detail::WarpPartitions::kCluster;
|
||||
int residual_cluster = warp_idx % Detail::WarpPartitions::kCluster;
|
||||
|
||||
int group_idx = residual_cluster / Detail::WarpPartitions::kGroup;
|
||||
int residual_group = residual_cluster % Detail::WarpPartitions::kGroup;
|
||||
|
||||
int row_idx = residual_group / Detail::WarpPartitions::kRow;
|
||||
int col_idx = residual_group % Detail::WarpPartitions::kRow;
|
||||
|
||||
// Compute per-lane offset
|
||||
int lane_row_offset = lane_idx / Detail::kAccessWidth;
|
||||
int lane_col_offset = lane_idx % Detail::kAccessWidth;
|
||||
|
||||
// Compute coordinate in output space
|
||||
int cluster_offset = cluster_idx * Shape::kRow * Shape::kGroup;
|
||||
int group_offset = group_idx * Shape::kRow;
|
||||
int row_offset = row_idx * Iterations::kRow * Detail::kAccessRows;
|
||||
int column_offset = col_idx * Iterations::kColumn * Detail::kAccessWidth * kElementsPerAccess;
|
||||
|
||||
MatrixCoord coord(
|
||||
cluster_offset + group_offset + row_offset + lane_row_offset,
|
||||
(column_offset + lane_col_offset) * kElementsPerAccess
|
||||
);
|
||||
|
||||
return coord;
|
||||
}
|
||||
};
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Template metaprogram for partitioning a 3D interleaved layout across warps
|
||||
/// to achieve several performance objectives:
|
||||
///
|
||||
/// - coalesced memory accesses in units of 64 Byte lines
|
||||
/// - minimal address arithmetic
|
||||
/// - minimal predicate calculations
|
||||
///
|
||||
template <typename WarpCount_, typename MmaCount_, int Threads,
|
||||
int ElementsPerAccess, int ElementSize>
|
||||
struct InterleavedOutputTileThreadMap {
|
||||
using WarpCount = WarpCount_;
|
||||
using MmaCount = MmaCount_;
|
||||
|
||||
static int const kWarpSize = 32;
|
||||
static int const kThreads = Threads;
|
||||
static int const kWarpCount = kThreads / kWarpSize;
|
||||
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
static int const kElementSize = ElementSize;
|
||||
|
||||
//
|
||||
// Metaprogram computation
|
||||
//
|
||||
|
||||
struct Detail {};
|
||||
|
||||
//
|
||||
// Output
|
||||
//
|
||||
|
||||
using Iterations = MmaCount;
|
||||
|
||||
using Delta = layout::PitchLinearShape<kWarpSize * kElementsPerAccess, 1>;
|
||||
|
||||
/// Initial offset function
|
||||
CUTLASS_HOST_DEVICE
|
||||
static layout::PitchLinearCoord initial_offset(int thread_idx) {
|
||||
int warp_idx = thread_idx / kWarpSize;
|
||||
int lane_idx = thread_idx % kWarpSize;
|
||||
|
||||
// Compute warp location
|
||||
layout::PitchLinearCoord warp_footprint{
|
||||
Delta::kContiguous * Iterations::kContiguous,
|
||||
Delta::kStrided * Iterations::kStrided};
|
||||
|
||||
layout::PitchLinearCoord warp_offset{warp_idx % WarpCount::kContiguous,
|
||||
warp_idx / WarpCount::kContiguous};
|
||||
|
||||
// Compute per-lane offset
|
||||
layout::PitchLinearCoord thread_offset_in_warp{
|
||||
lane_idx * kElementsPerAccess, 0};
|
||||
|
||||
layout::PitchLinearCoord thread_offset_in_threadblock_tile =
|
||||
warp_footprint * warp_offset + thread_offset_in_warp;
|
||||
|
||||
return thread_offset_in_threadblock_tile;
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
@@ -0,0 +1,725 @@
|
||||
/***************************************************************************************************
|
||||
* 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 Epilogue for threadblock scoped GEMMs using Tensor Ops.
|
||||
|
||||
The epilogue rearranges the result of a matrix product through shared memory to match canonical
|
||||
tensor layouts in global memory. Epilogues support conversion and reduction operations.
|
||||
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/layout/matrix.h"
|
||||
#include "cutlass/matrix_shape.h"
|
||||
#include "cutlass/tensor_ref.h"
|
||||
|
||||
#include "cutlass/transform/pitch_linear_thread_map.h"
|
||||
#include "cutlass/epilogue/threadblock/output_tile_thread_map.h"
|
||||
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace epilogue {
|
||||
namespace threadblock {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Tile iterator used to load output tile from shared memory in epilogue.
|
||||
///
|
||||
/// Satisfies: ReadableTileIterator | PredicatedTileIterator | ForwardTileIterator
|
||||
///
|
||||
template <
|
||||
typename ThreadMap_, ///< Thread map (conept: OutputTileThreadMap)
|
||||
typename Element_ ///< Element data type
|
||||
>
|
||||
class PredicatedTileIterator {
|
||||
public:
|
||||
using ThreadMap = ThreadMap_;
|
||||
using Shape = typename ThreadMap::Shape;
|
||||
|
||||
using Element = Element_;
|
||||
|
||||
using Layout = layout::RowMajor;
|
||||
using TensorRef = TensorRef<Element, Layout>;
|
||||
using ConstTensorRef = typename TensorRef::ConstTensorRef;
|
||||
|
||||
using Index = typename Layout::Index;
|
||||
using LongIndex = typename Layout::LongIndex;
|
||||
using TensorCoord = MatrixCoord;
|
||||
|
||||
static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
|
||||
static int const kThreads = ThreadMap::kThreads;
|
||||
static int const kIterations = ThreadMap::Count::kTile;
|
||||
|
||||
static_assert( ThreadMap::Iterations::kRow > 0,"ThreadMap::Iterations::kRow must be > 0");
|
||||
static_assert( ThreadMap::Iterations::kGroup > 0,"ThreadMap::Iterations::kGroup must be > 0");
|
||||
static_assert( ThreadMap::Iterations::kCluster > 0,"ThreadMap::Iterations::kCluster must be > 0");
|
||||
static_assert( ThreadMap::Iterations::kColumn > 0,"ThreadMap::Iterations::kColumn must be > 0");
|
||||
|
||||
/// Fragment object
|
||||
using Fragment = Array<
|
||||
Element,
|
||||
ThreadMap::Iterations::kColumn *
|
||||
ThreadMap::Iterations::kRow *
|
||||
ThreadMap::Iterations::kGroup *
|
||||
ThreadMap::Iterations::kCluster * ThreadMap::kElementsPerAccess>;
|
||||
|
||||
/// Memory access size
|
||||
using AccessType = AlignedArray<Element, ThreadMap::kElementsPerAccess>;
|
||||
|
||||
//
|
||||
// Parameters struct
|
||||
//
|
||||
|
||||
struct Params {
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
Index stride; ///< stride in bytes between rows
|
||||
|
||||
Index increment_row; ///< increment quantity (in bytes) to advance when moving between rows
|
||||
Index increment_group; ///< increment quantity (in bytes) to advance when moving to the next group
|
||||
Index increment_cluster; ///< increment quantity (in bytes) to advance when moving to the next cluster
|
||||
|
||||
Index advance_row; ///< amount to add to move to the next 'row' position
|
||||
Index advance_group; ///< amount to add to move to the next 'group' position
|
||||
Index advance_cluster; ///< amount to add to move to the next 'cluster' position
|
||||
Index advance_tile; ///< amount to add to move to the next 'tile'
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Status initialize(Index stride_) {
|
||||
|
||||
stride = stride_;
|
||||
|
||||
increment_row = stride * ThreadMap::Delta::kRow;
|
||||
|
||||
increment_group = stride * ThreadMap::Delta::kGroup
|
||||
- stride * ThreadMap::Delta::kRow * (ThreadMap::Iterations::kRow - 1);
|
||||
|
||||
increment_cluster = stride * ThreadMap::Delta::kCluster
|
||||
- stride * ThreadMap::Delta::kGroup * (ThreadMap::Iterations::kGroup - 1)
|
||||
- stride * ThreadMap::Delta::kRow * (ThreadMap::Iterations::kRow - 1);
|
||||
|
||||
advance_row = stride * ThreadMap::Shape::kRow;
|
||||
|
||||
advance_group = stride * (ThreadMap::Shape::kGroup - 1) * ThreadMap::Shape::kRow * ThreadMap::Count::kRow;
|
||||
|
||||
advance_cluster =
|
||||
stride *
|
||||
ThreadMap::Count::kGroup * ThreadMap::Shape::kGroup * ThreadMap::Count::kRow * ThreadMap::Shape::kRow;;
|
||||
|
||||
advance_tile =
|
||||
stride *
|
||||
ThreadMap::Shape::kGroup *
|
||||
ThreadMap::Shape::kRow *
|
||||
ThreadMap::Shape::kCluster *
|
||||
ThreadMap::Shape::kTile;
|
||||
|
||||
return Status::kSuccess;
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params() {
|
||||
initialize(0);
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(Layout const &layout) {
|
||||
|
||||
initialize(layout.stride(0) * int(sizeof(AccessType)) / kElementsPerAccess);
|
||||
}
|
||||
};
|
||||
|
||||
/// Mask object
|
||||
struct Mask {
|
||||
|
||||
static int const kCount = ThreadMap::Iterations::kColumn;
|
||||
|
||||
/// Predicate state
|
||||
bool predicates[kCount];
|
||||
|
||||
//
|
||||
// Mask
|
||||
//
|
||||
CUTLASS_HOST_DEVICE
|
||||
Mask() {
|
||||
enable();
|
||||
}
|
||||
|
||||
///< Efficiently disables all accesses guarded by mask
|
||||
CUTLASS_HOST_DEVICE void clear() {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kCount; ++i) {
|
||||
predicates[i] = false;
|
||||
}
|
||||
}
|
||||
|
||||
///< CUTLASS_HOST_DEVICE enables all accesses guarded by mask
|
||||
CUTLASS_DEVICE void enable() {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kCount; ++i) {
|
||||
predicates[i] = true;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
private:
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Parameters structure containing reference and precomputed state.
|
||||
Params params_;
|
||||
|
||||
/// Byte-level pointer
|
||||
uint8_t *byte_pointer_;
|
||||
|
||||
/// Array of boolean values to contain steady-state predicates
|
||||
Mask mask_;
|
||||
|
||||
/// Extent of the matrix tile in rows
|
||||
Index extent_row_;
|
||||
|
||||
/// A thread's starting row position (assuming steady-state predicates have been computed)
|
||||
Index thread_start_row_;
|
||||
|
||||
/// Internal state counter
|
||||
int state_[3];
|
||||
|
||||
private:
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
public:
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_DEVICE
|
||||
PredicatedTileIterator(
|
||||
Params const & params,
|
||||
Element *pointer,
|
||||
TensorCoord extent,
|
||||
int thread_idx,
|
||||
TensorCoord threadblock_offset = TensorCoord()
|
||||
):
|
||||
params_(params) {
|
||||
|
||||
TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx) + threadblock_offset;
|
||||
|
||||
extent_row_ = extent.row();
|
||||
thread_start_row_ = thread_offset.row();
|
||||
|
||||
// Initialize predicates
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int c = 0; c < ThreadMap::Iterations::kColumn; ++c) {
|
||||
|
||||
mask_.predicates[c] = ((thread_offset.column()
|
||||
+ ThreadMap::Delta::kColumn * c) < extent.column());
|
||||
}
|
||||
|
||||
// Initialize pointer
|
||||
byte_pointer_ = reinterpret_cast<uint8_t *>(pointer) +
|
||||
thread_offset.row() * params_.stride +
|
||||
thread_offset.column() * sizeof(AccessType) / kElementsPerAccess;
|
||||
|
||||
// Initialize internal state counter
|
||||
state_[0] = state_[1] = state_[2] = 0;
|
||||
}
|
||||
|
||||
/// Adds a pointer offset in units of Element
|
||||
CUTLASS_HOST_DEVICE
|
||||
void add_pointer_offset(LongIndex pointer_offset) {
|
||||
byte_pointer_ += pointer_offset * sizeof_bits<Element>::value / 8;
|
||||
}
|
||||
|
||||
/// Loads a fragment from memory
|
||||
CUTLASS_DEVICE
|
||||
void load(Fragment &frag) {
|
||||
|
||||
uint8_t *byte_pointer = byte_pointer_;
|
||||
AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) {
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
|
||||
|
||||
int frag_row_idx =
|
||||
(row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster));
|
||||
|
||||
int row_offset = row * ThreadMap::Delta::kRow
|
||||
+ group * ThreadMap::Delta::kGroup
|
||||
+ cluster * ThreadMap::Delta::kCluster;
|
||||
|
||||
bool row_guard = ((row_offset + thread_start_row_) < extent_row_);
|
||||
|
||||
AccessType *memory_pointer = reinterpret_cast<AccessType *>(byte_pointer);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
|
||||
|
||||
bool guard = row_guard && mask_.predicates[column];
|
||||
|
||||
if (guard) {
|
||||
frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column] =
|
||||
memory_pointer[column * ThreadMap::Delta::kColumn / kElementsPerAccess];
|
||||
}
|
||||
}
|
||||
|
||||
if (row + 1 < ThreadMap::Iterations::kRow) {
|
||||
byte_pointer += params_.increment_row;
|
||||
}
|
||||
}
|
||||
|
||||
if (group + 1 < ThreadMap::Iterations::kGroup) {
|
||||
byte_pointer += params_.increment_group;
|
||||
}
|
||||
}
|
||||
|
||||
if (cluster + 1 < ThreadMap::Iterations::kCluster) {
|
||||
byte_pointer += params_.increment_cluster;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Stores a fragment to memory
|
||||
CUTLASS_DEVICE
|
||||
void store(Fragment const &frag) {
|
||||
uint8_t *byte_pointer = byte_pointer_;
|
||||
AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) {
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
|
||||
|
||||
int frag_row_idx =
|
||||
(row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster));
|
||||
|
||||
int row_offset = row * ThreadMap::Delta::kRow
|
||||
+ group * ThreadMap::Delta::kGroup
|
||||
+ cluster * ThreadMap::Delta::kCluster;
|
||||
|
||||
bool row_guard = ((row_offset + thread_start_row_) < extent_row_);
|
||||
|
||||
AccessType *memory_pointer = reinterpret_cast<AccessType *>(byte_pointer);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
|
||||
|
||||
bool guard = row_guard && mask_.predicates[column];
|
||||
|
||||
if (guard) {
|
||||
|
||||
memory_pointer[column * ThreadMap::Delta::kColumn / kElementsPerAccess] =
|
||||
frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column];
|
||||
}
|
||||
}
|
||||
|
||||
if (row + 1 < ThreadMap::Iterations::kRow) {
|
||||
byte_pointer += params_.increment_row;
|
||||
}
|
||||
}
|
||||
|
||||
if (group + 1 < ThreadMap::Iterations::kGroup) {
|
||||
byte_pointer += params_.increment_group;
|
||||
}
|
||||
}
|
||||
|
||||
if (cluster + 1 < ThreadMap::Iterations::kCluster) {
|
||||
byte_pointer += params_.increment_cluster;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Advances to the next position to load or store
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedTileIterator &operator++() {
|
||||
|
||||
++state_[0];
|
||||
byte_pointer_ += params_.advance_row;
|
||||
thread_start_row_ += ThreadMap::Shape::kRow;
|
||||
|
||||
if (state_[0] == ThreadMap::Count::kRow) {
|
||||
|
||||
state_[0] = 0;
|
||||
++state_[1];
|
||||
byte_pointer_ += params_.advance_group;
|
||||
|
||||
thread_start_row_ += (ThreadMap::Shape::kGroup - 1) *
|
||||
ThreadMap::Shape::kRow * ThreadMap::Count::kRow;
|
||||
|
||||
if (state_[1] == ThreadMap::Count::kGroup) {
|
||||
|
||||
state_[1] = 0;
|
||||
++state_[2];
|
||||
byte_pointer_ += params_.advance_cluster;
|
||||
|
||||
thread_start_row_ += ThreadMap::Count::kGroup *
|
||||
ThreadMap::Shape::kGroup * ThreadMap::Count::kRow * ThreadMap::Shape::kRow;
|
||||
|
||||
if (state_[2] == ThreadMap::Count::kCluster) {
|
||||
state_[2] = 0;
|
||||
byte_pointer_ += params_.advance_tile;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return *this;
|
||||
}
|
||||
|
||||
///< Efficiently disables all accesses guarded by mask
|
||||
CUTLASS_DEVICE void clear_mask() {
|
||||
mask_.clear();
|
||||
}
|
||||
|
||||
///< Efficiently enables all accesses guarded by mask
|
||||
CUTLASS_DEVICE void enable_mask() {
|
||||
mask_.enable();
|
||||
}
|
||||
|
||||
///< Sets the mask
|
||||
CUTLASS_DEVICE void get_mask(Mask &mask) {
|
||||
return mask_;
|
||||
}
|
||||
|
||||
///< Sets the mask
|
||||
CUTLASS_DEVICE void set_mask(Mask const &mask) {
|
||||
mask_ = mask;
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
/// Tile iterator used to load output tile from shared memory in epilogue.
|
||||
///
|
||||
/// Satisfies: ReadableTileIterator | InterleavedPredicatedTileIterator | ForwardTileIterator
|
||||
///
|
||||
template <
|
||||
typename ThreadMap_, ///< Thread map (conept: OutputTileThreadMap)
|
||||
typename Element_, ///< Element data type
|
||||
int InterleavedK ///< Number of Interleaved K
|
||||
>
|
||||
class InterleavedPredicatedTileIterator {
|
||||
public:
|
||||
using ThreadMap = ThreadMap_;
|
||||
|
||||
using Element = Element_;
|
||||
|
||||
using Layout = layout::ColumnMajorInterleaved<InterleavedK>;
|
||||
using TensorRef = TensorRef<Element, Layout>;
|
||||
using ConstTensorRef = typename TensorRef::ConstTensorRef;
|
||||
|
||||
using Index = typename Layout::Index;
|
||||
using LongIndex = typename Layout::LongIndex;
|
||||
using TensorCoord = layout::PitchLinearCoord;
|
||||
|
||||
static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
|
||||
static int const kThreads = ThreadMap::kThreads;
|
||||
static int const kIterations = ThreadMap::Iterations::kCount;
|
||||
|
||||
/// Fragment object
|
||||
using Fragment = Array<Element, ThreadMap::kElementsPerAccess>;
|
||||
|
||||
/// Memory access size
|
||||
using AccessType = AlignedArray<Element, ThreadMap::kElementsPerAccess>;
|
||||
|
||||
//
|
||||
// Parameters struct
|
||||
//
|
||||
|
||||
struct Params {
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
Index stride; ///< stride in bytes between columns
|
||||
|
||||
Index advance_row; ///< amount to add to move to the next 'row' position
|
||||
Index advance_column; ///< amount to add to move to the next 'column' position
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Status initialize(Index stride_) {
|
||||
stride = stride_;
|
||||
|
||||
advance_row =
|
||||
ThreadMap::Delta::kContiguous * sizeof_bits<Element>::value / 8;
|
||||
|
||||
advance_column =
|
||||
stride_ - ThreadMap::Iterations::kContiguous * kElementsPerAccess *
|
||||
sizeof_bits<Element>::value * ThreadMap::kWarpSize / 8;
|
||||
|
||||
return Status::kSuccess;
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params() {
|
||||
initialize(0);
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(Layout const &layout) {
|
||||
|
||||
initialize(layout.stride(0) * int(sizeof(AccessType)) / kElementsPerAccess);
|
||||
}
|
||||
};
|
||||
|
||||
/// Mask object
|
||||
struct Mask {
|
||||
static int const kCount = (ThreadMap::Iterations::kContiguous < 8)
|
||||
? 8
|
||||
: ThreadMap::Iterations::kContiguous;
|
||||
|
||||
/// Predicate state
|
||||
bool predicates[kCount];
|
||||
|
||||
//
|
||||
// Mask
|
||||
//
|
||||
CUTLASS_HOST_DEVICE
|
||||
Mask() {
|
||||
enable();
|
||||
}
|
||||
|
||||
///< Efficiently disables all accesses guarded by mask
|
||||
CUTLASS_HOST_DEVICE void clear() {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kCount; ++i) {
|
||||
predicates[i] = false;
|
||||
}
|
||||
}
|
||||
|
||||
///< CUTLASS_HOST_DEVICE enables all accesses guarded by mask
|
||||
CUTLASS_DEVICE void enable() {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kCount; ++i) {
|
||||
predicates[i] = true;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
private:
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Parameters structure containing reference and precomputed state.
|
||||
Params params_;
|
||||
|
||||
/// Byte-level pointer
|
||||
uint8_t *byte_pointer_;
|
||||
|
||||
/// Array of boolean values to contain steady-state predicates
|
||||
Mask mask_;
|
||||
|
||||
/// Extent of the matrix tile in columns
|
||||
Index extent_col_;
|
||||
|
||||
/// A thread's starting column position (assuming steady-state predicates have
|
||||
/// been computed)
|
||||
Index thread_start_col_;
|
||||
|
||||
/// Internal iteration counter
|
||||
int iteration_contiguous_;
|
||||
|
||||
int iteration_strided_;
|
||||
|
||||
private:
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
public:
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_DEVICE
|
||||
InterleavedPredicatedTileIterator(
|
||||
Params const & params,
|
||||
Element *pointer,
|
||||
TensorCoord extent,
|
||||
int thread_idx,
|
||||
TensorCoord threadblock_offset
|
||||
):
|
||||
params_(params) {
|
||||
TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx) +
|
||||
TensorCoord(threadblock_offset.contiguous() * InterleavedK,
|
||||
threadblock_offset.strided() / InterleavedK);
|
||||
|
||||
extent_col_ = extent.strided() / InterleavedK;
|
||||
thread_start_col_ = thread_offset.strided();
|
||||
|
||||
// Initialize predicates
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) {
|
||||
mask_.predicates[c] =
|
||||
((thread_offset.contiguous() + ThreadMap::Delta::kContiguous * c) <
|
||||
(extent.contiguous() * InterleavedK));
|
||||
}
|
||||
|
||||
// Initialize pointer
|
||||
byte_pointer_ = reinterpret_cast<uint8_t *>(pointer) +
|
||||
thread_offset.strided() * params_.stride +
|
||||
thread_offset.contiguous() * sizeof(AccessType) / kElementsPerAccess;
|
||||
|
||||
// Initialize internal state counter
|
||||
iteration_contiguous_ = iteration_strided_ = 0;
|
||||
}
|
||||
|
||||
/// Adds a pointer offset in units of Element
|
||||
CUTLASS_HOST_DEVICE
|
||||
void add_pointer_offset(LongIndex pointer_offset) {
|
||||
byte_pointer_ += pointer_offset * sizeof_bits<Element>::value / 8;
|
||||
}
|
||||
|
||||
/// Loads a fragment from memory
|
||||
CUTLASS_DEVICE
|
||||
void load(Fragment &frag) {
|
||||
uint8_t *byte_pointer = byte_pointer_;
|
||||
AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
|
||||
AccessType *memory_pointer = reinterpret_cast<AccessType *>(byte_pointer);
|
||||
|
||||
int col_offset = iteration_strided_ * ThreadMap::Delta::kStrided;
|
||||
|
||||
bool col_guard = ((thread_start_col_ + col_offset) < extent_col_);
|
||||
|
||||
bool guard = col_guard && mask_.predicates[iteration_contiguous_];
|
||||
|
||||
if (guard) {
|
||||
*frag_ptr = *memory_pointer;
|
||||
}
|
||||
}
|
||||
|
||||
/// Stores a fragment to memory
|
||||
CUTLASS_DEVICE
|
||||
void store(Fragment const &frag) {
|
||||
uint8_t *byte_pointer = byte_pointer_;
|
||||
AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
|
||||
AccessType *memory_pointer = reinterpret_cast<AccessType *>(byte_pointer);
|
||||
|
||||
int col_offset = iteration_strided_ * ThreadMap::Delta::kStrided;
|
||||
|
||||
bool col_guard = ((thread_start_col_ + col_offset) < extent_col_);
|
||||
|
||||
bool guard = col_guard && mask_.predicates[iteration_contiguous_];
|
||||
|
||||
if (guard) {
|
||||
*memory_pointer = *frag_ptr;
|
||||
}
|
||||
}
|
||||
|
||||
/// Overrides the internal iteration index
|
||||
CUTLASS_HOST_DEVICE
|
||||
void set_iteration_index(int iteration) {
|
||||
iteration_contiguous_ = iteration % ThreadMap::Iterations::kContiguous;
|
||||
iteration_strided_ = iteration / ThreadMap::Iterations::kContiguous;
|
||||
}
|
||||
|
||||
/// Advances to the next position to load or store
|
||||
CUTLASS_HOST_DEVICE
|
||||
InterleavedPredicatedTileIterator &operator++() {
|
||||
|
||||
++iteration_contiguous_;
|
||||
byte_pointer_ += params_.advance_row;
|
||||
|
||||
if (iteration_contiguous_ == ThreadMap::Iterations::kContiguous) {
|
||||
|
||||
iteration_contiguous_ = 0;
|
||||
++iteration_strided_;
|
||||
byte_pointer_ += params_.advance_column;
|
||||
|
||||
if (iteration_strided_ == ThreadMap::Iterations::kStrided) {
|
||||
iteration_strided_ = 0;
|
||||
}
|
||||
}
|
||||
|
||||
return *this;
|
||||
}
|
||||
|
||||
///< Efficiently disables all accesses guarded by mask
|
||||
CUTLASS_DEVICE void clear_mask() {
|
||||
mask_.clear();
|
||||
}
|
||||
|
||||
///< Efficiently enables all accesses guarded by mask
|
||||
CUTLASS_DEVICE void enable_mask() {
|
||||
mask_.enable();
|
||||
}
|
||||
|
||||
///< Sets the mask
|
||||
CUTLASS_DEVICE void get_mask(Mask &mask) {
|
||||
return mask_;
|
||||
}
|
||||
|
||||
///< Sets the mask
|
||||
CUTLASS_DEVICE void set_mask(Mask const &mask) {
|
||||
mask_ = mask;
|
||||
}
|
||||
};
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,201 @@
|
||||
/***************************************************************************************************
|
||||
* 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 Epilogue for threadblock scoped GEMMs using Tensor Ops.
|
||||
|
||||
The epilogue rearranges the result of a matrix product through shared memory to match canonical
|
||||
tensor layouts in global memory. Epilogues support conversion and reduction operations.
|
||||
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/layout/matrix.h"
|
||||
#include "cutlass/matrix_shape.h"
|
||||
#include "cutlass/tensor_ref.h"
|
||||
|
||||
#include "cutlass/epilogue/threadblock/output_tile_thread_map.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace threadblock {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Tile iterator used to load output tile from shared memory in epilogue.
|
||||
///
|
||||
/// Satisfies: ReadableTileIterator
|
||||
///
|
||||
template <
|
||||
typename ThreadMap_, ///< Thread map (conept: OutputTileThreadMap)
|
||||
typename Element_, ///< Element data type
|
||||
int MaxAlignment = ThreadMap_::kElementsPerAccess * sizeof_bits<Element_>::value / 8
|
||||
>
|
||||
class SharedLoadIterator {
|
||||
public:
|
||||
using ThreadMap = ThreadMap_;
|
||||
using Shape = typename ThreadMap::Shape;
|
||||
|
||||
using Element = Element_;
|
||||
|
||||
using Layout = layout::RowMajor;
|
||||
using TensorRef = TensorRef<Element, Layout>;
|
||||
using ConstTensorRef = typename TensorRef::ConstTensorRef;
|
||||
|
||||
using Index = typename Layout::Index;
|
||||
using LongIndex = typename Layout::LongIndex;
|
||||
using TensorCoord = MatrixCoord;
|
||||
|
||||
static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
|
||||
|
||||
static int const kMinAlignment = ThreadMap_::kElementsPerAccess * sizeof_bits<Element_>::value / 8;
|
||||
|
||||
static int const kAlignment = (MaxAlignment < kMinAlignment ? MaxAlignment : kMinAlignment);
|
||||
|
||||
static int const kThreads = ThreadMap::kThreads;
|
||||
|
||||
/// Fragment object
|
||||
using Fragment = Array<
|
||||
Element,
|
||||
ThreadMap::Iterations::kColumn *
|
||||
ThreadMap::Iterations::kRow *
|
||||
ThreadMap::Iterations::kGroup *
|
||||
ThreadMap::Iterations::kCluster *
|
||||
ThreadMap::kElementsPerAccess>;
|
||||
|
||||
/// Memory access size
|
||||
using AccessType = AlignedArray<
|
||||
Element,
|
||||
ThreadMap::kElementsPerAccess,
|
||||
kAlignment>;
|
||||
|
||||
private:
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Byte-level pointer
|
||||
uint8_t *byte_pointer_;
|
||||
|
||||
/// Stride along adjacent rows
|
||||
int stride_;
|
||||
|
||||
public:
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_DEVICE
|
||||
SharedLoadIterator(
|
||||
TensorRef ref,
|
||||
int thread_idx
|
||||
):
|
||||
byte_pointer_(reinterpret_cast<uint8_t *>(ref.data())),
|
||||
stride_((ref.stride(0) * sizeof_bits<Element>::value) / 8) {
|
||||
|
||||
TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx);
|
||||
|
||||
// Initialize pointer
|
||||
byte_pointer_ +=
|
||||
thread_offset.row() * stride_ +
|
||||
thread_offset.column() * sizeof(AccessType) / kElementsPerAccess;
|
||||
|
||||
int byte_offset = thread_offset.row() * stride_ +
|
||||
thread_offset.column() * sizeof(AccessType) / kElementsPerAccess;
|
||||
}
|
||||
|
||||
/// Adds a pointer offset in units of Element
|
||||
CUTLASS_HOST_DEVICE
|
||||
void add_pointer_offset(LongIndex pointer_offset) {
|
||||
byte_pointer_ += pointer_offset * sizeof_bits<Element>::value / 8;
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void add_tile_offset(TensorCoord const &offset) {
|
||||
add_pointer_offset(offset.row() * stride_ / (sizeof_bits<Element>::value / 8) + offset.column() * Shape::kColumn);
|
||||
}
|
||||
|
||||
/// Loads a fragment from memory
|
||||
CUTLASS_DEVICE
|
||||
void load_with_pointer_offset(Fragment &frag, Index pointer_offset) {
|
||||
|
||||
AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) {
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
|
||||
|
||||
uint8_t const *byte_pointer = byte_pointer_ +
|
||||
row * ThreadMap::Delta::kRow * stride_ +
|
||||
group * ThreadMap::Delta::kGroup* stride_ +
|
||||
cluster * ThreadMap::Delta::kCluster * stride_ +
|
||||
pointer_offset * sizeof_bits<Element>::value / 8;
|
||||
|
||||
int frag_row_idx =
|
||||
(row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster));
|
||||
|
||||
AccessType const *memory_pointer = reinterpret_cast<AccessType const *>(byte_pointer);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
|
||||
|
||||
int frag_idx = frag_row_idx * ThreadMap::Iterations::kColumn + column;
|
||||
|
||||
frag_ptr[frag_idx] =
|
||||
memory_pointer[column * ThreadMap::Delta::kColumn / kElementsPerAccess];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Loads a fragment
|
||||
CUTLASS_DEVICE
|
||||
void load(Fragment &frag) {
|
||||
|
||||
load_with_pointer_offset(frag, 0);
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
Reference in New Issue
Block a user