leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 26 Packages Projects Releases Wiki Activity

CUTLASS 2.1 (#83)

Browse Source 
CUTLASS 2.1 contributes:
- BLAS-style host-side API added to CUTLASS Library
- Planar Complex GEMM kernels targeting Volta and Turing Tensor Cores
- Minor enhancements and bug fixes
This commit is contained in:
Andrew Kerr
2020-04-07 13:51:25 -07:00
committed by GitHub
parent 7c0cd26d13
commit 96dab34ad9
196 changed files with 20653 additions and 1995 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cutlass/epilogue/threadblock/default_epilogue_complex_tensor_op.h
+15 -5
View File
@@ -58,16 +58,25 @@ namespace cutlass {
namespace epilogue {
namespace threadblock {
////////////////////////////////////////////////////////////////////////////////
/// Defines sensible defaults for epilogues for TensorOps.
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Specialization and defines sensible defaults for epilogues for complex*complex case
// 4 real-valued mma operations (Complex)
// A = (ar + j ai), B (br +j bi), D = AB
// D = dr + j di = (ar*br - ai*bi) + j (ar*bi + ai*br)
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
/// Epilouge Shape
typename Shape_,
/// Warp-level mma operator
typename WarpMmaTensorOp_,
/// Number of k partitions
int PartitionsK,
/// Epilogue output operator
typename OutputOp_,
int ElementsPerAccess
>
/// Elements accessed by inner-most loop of AccumulatorFragmentIterator::load()
int ElementsPerAccess,
/// Multiply-add operator
typename Operator_ = arch::OpMultiplyAddComplex>
struct DefaultEpilogueComplexTensorOp {
using Shape = Shape_;
@@ -75,6 +84,7 @@ struct DefaultEpilogueComplexTensorOp {
static int const kPartitionsK = PartitionsK;
using OutputOp = OutputOp_;
static int const kElementsPerAccess = ElementsPerAccess;
using Operator = Operator_;
using ElementOutput = typename OutputOp::ElementOutput;
using LayoutC = typename WarpMmaTensorOp::LayoutC;
include/cutlass/epilogue/threadblock/default_epilogue_planar_complex.h
+197
View File
@@ -0,0 +1,197 @@
/***************************************************************************************************
* 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 Constructs a default epilogue for planar complex outputs.
This template reuses components for real-valued epilogues and applies them to planar complex
output matrices.
*/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/numeric_types.h"
#include "cutlass/array.h"
#include "cutlass/array_planar_complex.h"
#include "cutlass/arch/arch.h"
#include "cutlass/epilogue/thread/linear_combination_planar_complex.h"
#include "cutlass/epilogue/threadblock/default_epilogue_simt.h"
#include "cutlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h"
#include "cutlass/epilogue/threadblock/default_epilogue_tensor_op.h"
#include "cutlass/epilogue/threadblock/epilogue_planar_complex.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
namespace epilogue {
namespace threadblock {
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Defines sensible defaults for epilogues.
template <
typename ThreadblockShape_,
typename WarpMma_,
typename OpcodeClass_,
typename ArchTag_,
int PartitionsK,
typename OutputOp_,
int ElementsPerAccess
>
struct DefaultEpiloguePlanarComplex;
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Defines sensible defaults for epilogues.
template <
typename ThreadblockShape_,
typename WarpMmaOperator_,
int PartitionsK,
typename OutputOp_,
int ElementsPerAccess
>
struct DefaultEpiloguePlanarComplex<
ThreadblockShape_,
WarpMmaOperator_,
arch::OpClassTensorOp,
arch::Sm70,
PartitionsK,
OutputOp_,
ElementsPerAccess> {
using RealEpilogue = DefaultEpilogueVoltaTensorOp<
ThreadblockShape_,
WarpMmaOperator_,
PartitionsK,
OutputOp_,
ElementsPerAccess
>;
using Epilogue = EpiloguePlanarComplex<
ThreadblockShape_,
WarpMmaOperator_,
PartitionsK,
typename RealEpilogue::OutputTileIterator,
typename RealEpilogue::AccumulatorFragmentIterator,
typename RealEpilogue::WarpTileIterator,
typename RealEpilogue::SharedLoadIterator,
OutputOp_,
typename RealEpilogue::Padding
>;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Defines sensible defaults for epilogues.
template <
typename ThreadblockShape_,
typename WarpMmaOperator_,
int PartitionsK,
typename OutputOp_,
int ElementsPerAccess
>
struct DefaultEpiloguePlanarComplex<
ThreadblockShape_,
WarpMmaOperator_,
arch::OpClassTensorOp,
arch::Sm75,
PartitionsK,
OutputOp_,
ElementsPerAccess> {
using RealEpilogue = DefaultEpilogueTensorOp<
ThreadblockShape_,
WarpMmaOperator_,
PartitionsK,
OutputOp_,
ElementsPerAccess
>;
using Epilogue = EpiloguePlanarComplex<
ThreadblockShape_,
WarpMmaOperator_,
PartitionsK,
typename RealEpilogue::OutputTileIterator,
typename RealEpilogue::AccumulatorFragmentIterator,
typename RealEpilogue::WarpTileIterator,
typename RealEpilogue::SharedLoadIterator,
OutputOp_,
typename RealEpilogue::Padding
>;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Defines sensible defaults for epilogues.
template <
typename ThreadblockShape_,
typename WarpMmaOperator_,
typename ArchTag_,
int PartitionsK,
typename OutputOp_,
int ElementsPerAccess
>
struct DefaultEpiloguePlanarComplex<
ThreadblockShape_,
WarpMmaOperator_,
arch::OpClassSimt,
ArchTag_,
PartitionsK,
OutputOp_,
ElementsPerAccess> {
using RealEpilogue = DefaultEpilogueSimt<
ThreadblockShape_,
WarpMmaOperator_,
OutputOp_,
ElementsPerAccess
>;
using Epilogue = EpiloguePlanarComplex<
ThreadblockShape_,
WarpMmaOperator_,
PartitionsK,
typename RealEpilogue::OutputTileIterator,
typename RealEpilogue::AccumulatorFragmentIterator,
typename RealEpilogue::WarpTileIterator,
typename RealEpilogue::SharedLoadIterator,
OutputOp_,
typename RealEpilogue::Padding
>;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace threadblock
} // namespace epilogue
} // namespace cutlass
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/epilogue/threadblock/default_epilogue_wmma_tensor_op.h
+1 -1
View File
@@ -23,7 +23,7 @@
*
**************************************************************************************************/
/*! \file
\brief Epilogue for threadblock scoped GEMMs using Tensor Ops.
\brief Epilogue for threadblock scoped GEMMs using WMMA.
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.
include/cutlass/epilogue/threadblock/default_thread_map_tensor_op.h
+48
View File
@@ -146,6 +146,54 @@ struct DefaultInterleavedThreadMapTensorOp {
////////////////////////////////////////////////////////////////////////////////
/// Defines the optimal thread map for TensorOp accumulator layouts
template <typename ThreadblockShape_, typename WarpShape_, int PartitionsK,
typename Element_, int ElementsPerAccess, int InterleavedK>
struct DefaultInterleavedConvThreadMapTensorOp {
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::MaskedTileIterator satisfying concept
/// InterleavedOutputTileThreadMap
using Type = InterleavedConvOutputTileThreadMap<
MatrixShape<Detail::WarpCount::kM, Detail::WarpCount::kN>,
MatrixShape<WarpShape::kM / Detail::kTensorOpRows,
WarpShape::kN / InterleavedK>,
Detail::kThreads, kElementsPerAccess, sizeof_bits<Element>::value>;
};
////////////////////////////////////////////////////////////////////////////////
} // namespace threadblock
} // namespace epilogue
} // namespace cutlass
include/cutlass/epilogue/threadblock/epilogue.h
+9 -4
View File
@@ -32,7 +32,11 @@
#pragma once
#if defined(__CUDACC_RTC__)
#include <cuda/std/cassert>
#else
#include <assert.h>
#endif
#include "cutlass/cutlass.h"
#include "cutlass/numeric_types.h"
@@ -74,7 +78,7 @@ template <
class Epilogue :
public EpilogueBase<
Shape_,
WarpMmaOperator_,
typename WarpMmaOperator_::Shape,
PartitionsK,
AccumulatorFragmentIterator_,
WarpTileIterator_,
@@ -84,7 +88,7 @@ public:
using Base = EpilogueBase<
Shape_,
WarpMmaOperator_,
typename WarpMmaOperator_::Shape,
PartitionsK,
AccumulatorFragmentIterator_,
WarpTileIterator_,
@@ -171,8 +175,9 @@ public:
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)
OutputTileIterator source_iterator, ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles)
int64_t imag_stride_dest = 0, ///< Arguments required for planar complex case - not used in real-valued case
int64_t imag_stride_src = 0) { ///<
typename OutputTileIterator::Fragment source_fragment;
include/cutlass/epilogue/threadblock/epilogue_base.h
+8 -23
View File
@@ -32,7 +32,11 @@
#pragma once
#if defined(__CUDACC_RTC__)
#include <cuda/std/cassert>
#else
#include <assert.h>
#endif
#include "cutlass/cutlass.h"
#include "cutlass/matrix_shape.h"
@@ -58,7 +62,7 @@ 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)
typename WarpShape_, ///< 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
@@ -68,7 +72,7 @@ class EpilogueBase {
public:
using Shape = Shape_;
using WarpMmaOperator = WarpMmaOperator_;
using WarpShape = WarpShape_;
static int const kPartitionsK = PartitionsK;
using AccumulatorFragmentIterator = AccumulatorFragmentIterator_;
using WarpTileIterator = WarpTileIterator_;
@@ -83,11 +87,10 @@ public:
/// Accumulator element
using ElementAccumulator = typename AccumulatorTile::Element;
/// Number of warps
using WarpCount = gemm::GemmShape<
Shape::kM / WarpMmaOperator::Shape::kM,
Shape::kN / WarpMmaOperator::Shape::kN,
Shape::kM / WarpShape::kM,
Shape::kN / WarpShape::kN,
kPartitionsK
>;
@@ -144,24 +147,6 @@ public:
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:
include/cutlass/epilogue/threadblock/epilogue_planar_complex.h
+397
View File
@@ -0,0 +1,397 @@
/***************************************************************************************************
* 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/array_planar_complex.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 for planar-complex output representations.
///
/// Note, as with most CUTLASS components for planar complex, the template arguments describe
/// the underlying real data type.
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 EpiloguePlanarComplex {
public:
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 = ArrayPlanarComplex<
typename WarpMmaOperator::FragmentC::Element,
WarpMmaOperator::FragmentC::kElements
>;
/// 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>;
/// Shape of each warp-level operation
using WarpShape = typename WarpMmaOperator::Shape;
/// Number of warps
using WarpCount = gemm::GemmShape<
Shape::kM / WarpShape::kM,
Shape::kN / WarpShape::kN,
kPartitionsK
>;
/// Shared memory allocation
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
>;
static int const kImaginaryStride = StorageShape::kCount;
//
// Data members
//
AlignedBuffer<Element, kImaginaryStride * 2> 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}));
}
};
private:
//
// Data members
//
SharedStorage &shared_storage_;
/// Loads fragment from shared memory aligned with output tensor
SharedLoadIterator shared_load_iterator_;
/// Stores a warp's fragment of accumulators to SMEM
WarpTileIterator warp_tile_iterator_;
public:
/// Constructor
CUTLASS_DEVICE
EpiloguePlanarComplex(
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),
shared_load_iterator_(shared_storage.reference(), thread_idx),
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);
}
/// Streams the result to global memory
CUTLASS_DEVICE
void operator()(
OutputOp const &output_op, ///< Output operator
OutputTileIterator destination_iterator_real, ///< Tile iterator for destination
OutputTileIterator destination_iterator_imag, ///< Tile iterator for destination
AccumulatorTile const &accumulators, ///< Complete warp-level accumulator tile
OutputTileIterator source_iterator_real, ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles)
OutputTileIterator source_iterator_imag) { ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles)
typename OutputTileIterator::Fragment source_fragment_real;
typename OutputTileIterator::Fragment source_fragment_imag;
if (!output_op.is_source_needed()) {
source_iterator_real.clear_mask();
source_iterator_imag.clear_mask();
}
source_fragment_real.clear();
source_fragment_imag.clear();
//
// Iterator over warp-level accumulator fragment
//
AccumulatorFragmentIterator accum_fragment_iterator_real(accumulators.real);
AccumulatorFragmentIterator accum_fragment_iterator_imag(accumulators.imag);
//
// Iterate over accumulator tile
//
CUTLASS_PRAGMA_UNROLL
for (int iter = 0; iter < OutputTileIterator::kIterations; ++iter) {
//
// Load the source
//
source_iterator_real.load(source_fragment_real);
source_iterator_imag.load(source_fragment_imag);
++source_iterator_real;
++source_iterator_imag;
//
// Convert and store fragment
//
__syncthreads();
typename AccumulatorFragmentIterator::Fragment accum_fragment_real;
typename AccumulatorFragmentIterator::Fragment accum_fragment_imag;
accum_fragment_iterator_real.load(accum_fragment_real);
accum_fragment_iterator_imag.load(accum_fragment_imag);
++accum_fragment_iterator_real;
++accum_fragment_iterator_imag;
this->warp_tile_iterator_.store(accum_fragment_real);
this->warp_tile_iterator_.store_with_pointer_offset(accum_fragment_imag, SharedStorage::kImaginaryStride);
__syncthreads();
//
// Load fragments from shared memory
//
typename SharedLoadIterator::Fragment aligned_accum_fragment_real[kPartitionsK];
typename SharedLoadIterator::Fragment aligned_accum_fragment_imag[kPartitionsK];
shared_load_iterator_.load(aligned_accum_fragment_real[0]);
shared_load_iterator_.load_with_pointer_offset(aligned_accum_fragment_imag[0], SharedStorage::kImaginaryStride);
// If the number of k-slices is > 1 - perform a reduction amongst the k-slices
static_assert(kPartitionsK == 1, "Sliced-K not supported for planar complex at this time");
//
// Compute the output result
//
typename OutputTileIterator::Fragment output_fragment_real;
typename OutputTileIterator::Fragment output_fragment_imag;
apply_output_operator_(
output_fragment_real,
output_fragment_imag,
output_op,
aligned_accum_fragment_real[0],
aligned_accum_fragment_imag[0],
source_fragment_real,
source_fragment_imag);
//
// Store the final result
//
destination_iterator_real.store(output_fragment_real);
destination_iterator_imag.store(output_fragment_imag);
++destination_iterator_real;
++destination_iterator_imag;
}
}
private:
/// Helper to invoke the output functor over each vector of output
CUTLASS_DEVICE
void apply_output_operator_(
typename OutputTileIterator::Fragment &output_fragment_real,
typename OutputTileIterator::Fragment &output_fragment_imag,
OutputOp const &output_op, ///< Output operator
typename SharedLoadIterator::Fragment const &aligned_accum_fragment_real,
typename SharedLoadIterator::Fragment const &aligned_accum_fragment_imag,
typename OutputTileIterator::Fragment const &source_fragment_real,
typename OutputTileIterator::Fragment const &source_fragment_imag) {
OutputAccessType *output_frag_real_ptr =
reinterpret_cast<OutputAccessType *>(&output_fragment_real);
OutputAccessType *output_frag_imag_ptr =
reinterpret_cast<OutputAccessType *>(&output_fragment_imag);
AccumulatorAccessType const *compute_frag_real_ptr =
reinterpret_cast<AccumulatorAccessType const *>(&aligned_accum_fragment_real);
AccumulatorAccessType const *compute_frag_imag_ptr =
reinterpret_cast<AccumulatorAccessType const *>(&aligned_accum_fragment_imag);
OutputAccessType const *source_frag_real_ptr =
reinterpret_cast<OutputAccessType const *>(&source_fragment_real);
OutputAccessType const *source_frag_imag_ptr =
reinterpret_cast<OutputAccessType const *>(&source_fragment_imag);
int const kOutputOpIterations =
OutputTileIterator::Fragment::kElements / OutputTileIterator::kElementsPerAccess;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < kOutputOpIterations; ++i) {
// Call the output operator
auto result_fragment = output_op(
make_ArrayPlanarComplex(compute_frag_real_ptr[i], compute_frag_imag_ptr[i]),
make_ArrayPlanarComplex(source_frag_real_ptr[i], source_frag_imag_ptr[i])
);
output_frag_real_ptr[i] = result_fragment.real;
output_frag_imag_ptr[i] = result_fragment.imag;
}
}
};
////////////////////////////////////////////////////////////////////////////////
} // namespace threadblock
} // namespace epilogue
} // namespace cutlass
////////////////////////////////////////////////////////////////////////////////
include/cutlass/epilogue/threadblock/interleaved_epilogue.h
-3
View File
@@ -85,9 +85,6 @@ class InterleavedEpilogue {
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;
include/cutlass/epilogue/threadblock/output_tile_thread_map.h
+63 -3
View File
@@ -437,11 +437,10 @@ struct OutputTileOptimalThreadMap {
/// - minimal address arithmetic
/// - minimal predicate calculations
///
template <typename WarpCount_, typename MmaCount_, int Threads,
template <typename WarpCount_, typename Iterations_, int Threads,
int ElementsPerAccess, int ElementSize>
struct InterleavedOutputTileThreadMap {
using WarpCount = WarpCount_;
using MmaCount = MmaCount_;
static int const kWarpSize = 32;
static int const kThreads = Threads;
@@ -460,7 +459,7 @@ struct InterleavedOutputTileThreadMap {
// Output
//
using Iterations = MmaCount;
using Iterations = Iterations_;
using Delta = layout::PitchLinearShape<kWarpSize * kElementsPerAccess, 1>;
@@ -491,6 +490,67 @@ struct InterleavedOutputTileThreadMap {
////////////////////////////////////////////////////////////////////////////////
/// Template metaprogram for partitioning a 4D 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 Iterations_, int Threads,
int ElementsPerAccess, int ElementSize>
struct InterleavedConvOutputTileThreadMap {
using WarpCount = WarpCount_;
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 = Iterations_;
using Delta = MatrixShape<kWarpSize / 4, 4 * kElementsPerAccess>;
/// 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
MatrixCoord warp_footprint{
Delta::kRow * Iterations::kRow,
Delta::kColumn * Iterations::kColumn,
};
MatrixCoord warp_offset{warp_idx % WarpCount::kRow,
warp_idx / WarpCount::kRow};
// Compute per-lane offset
MatrixCoord thread_offset_in_warp{lane_idx / 4,
(lane_idx % 4) * kElementsPerAccess};
MatrixCoord 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
include/cutlass/epilogue/threadblock/predicated_tile_iterator.h
+47 -31
View File
@@ -36,9 +36,9 @@
#include "cutlass/numeric_types.h"
#include "cutlass/array.h"
#include "cutlass/layout/matrix.h"
#include "cutlass/layout/tensor.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"
@@ -107,16 +107,16 @@ public:
// Data members
//
Index stride; ///< stride in bytes between rows
LongIndex 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
LongIndex increment_row; ///< increment quantity (in bytes) to advance when moving between rows
LongIndex increment_group; ///< increment quantity (in bytes) to advance when moving to the next group
LongIndex 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'
LongIndex advance_row; ///< amount to add to move to the next 'row' position
LongIndex advance_group; ///< amount to add to move to the next 'group' position
LongIndex advance_cluster; ///< amount to add to move to the next 'cluster' position
LongIndex advance_tile; ///< amount to add to move to the next 'tile'
//
// Methods
@@ -125,7 +125,7 @@ public:
CUTLASS_HOST_DEVICE
Status initialize(Index stride_) {
stride = stride_;
stride = LongIndex(stride_);
increment_row = stride * ThreadMap::Delta::kRow;
@@ -261,8 +261,8 @@ public:
// Initialize pointer
byte_pointer_ = reinterpret_cast<uint8_t *>(pointer) +
thread_offset.row() * params_.stride +
thread_offset.column() * sizeof(AccessType) / kElementsPerAccess;
LongIndex(thread_offset.row()) * LongIndex(params_.stride) +
LongIndex(thread_offset.column()) * sizeof(AccessType) / kElementsPerAccess;
// Initialize internal state counter
state_[0] = state_[1] = state_[2] = 0;
@@ -276,7 +276,7 @@ public:
/// Loads a fragment from memory
CUTLASS_DEVICE
void load(Fragment &frag) {
void load_with_byte_offset(Fragment &frag, int64_t byte_offset) {
uint8_t *byte_pointer = byte_pointer_;
AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
@@ -299,7 +299,7 @@ public:
bool row_guard = ((row_offset + thread_start_row_) < extent_row_);
AccessType *memory_pointer = reinterpret_cast<AccessType *>(byte_pointer);
AccessType *memory_pointer = reinterpret_cast<AccessType *>(byte_pointer + byte_offset);
CUTLASS_PRAGMA_UNROLL
for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
@@ -328,9 +328,15 @@ public:
}
}
/// Loads a fragment from memory
CUTLASS_DEVICE
void load(Fragment &frag) {
load_with_byte_offset(frag, 0);
}
/// Stores a fragment to memory
CUTLASS_DEVICE
void store(Fragment const &frag) {
void store_with_byte_offset(Fragment const &frag, int64_t byte_offset) {
uint8_t *byte_pointer = byte_pointer_;
AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
@@ -352,7 +358,7 @@ public:
bool row_guard = ((row_offset + thread_start_row_) < extent_row_);
AccessType *memory_pointer = reinterpret_cast<AccessType *>(byte_pointer);
AccessType *memory_pointer = reinterpret_cast<AccessType *>(byte_pointer + byte_offset);
CUTLASS_PRAGMA_UNROLL
for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
@@ -382,6 +388,12 @@ public:
}
}
/// Stores a fragment to memory
CUTLASS_DEVICE
void store(Fragment const &frag) {
store_with_byte_offset(frag, 0);
}
/// Advances to the next position to load or store
CUTLASS_HOST_DEVICE
PredicatedTileIterator &operator++() {
@@ -440,6 +452,7 @@ public:
};
////////////////////////////////////////////////////////////////////////////////
/// Tile iterator used to load output tile from shared memory in epilogue.
///
/// Satisfies: ReadableTileIterator | InterleavedPredicatedTileIterator | ForwardTileIterator
@@ -447,7 +460,7 @@ public:
template <
typename ThreadMap_, ///< Thread map (conept: OutputTileThreadMap)
typename Element_, ///< Element data type
int InterleavedK ///< Number of Interleaved K
int InterleavedN ///< Number of Interleaved N
>
class InterleavedPredicatedTileIterator {
public:
@@ -455,7 +468,7 @@ public:
using Element = Element_;
using Layout = layout::ColumnMajorInterleaved<InterleavedK>;
using Layout = layout::ColumnMajorInterleaved<InterleavedN>;
using TensorRef = TensorRef<Element, Layout>;
using ConstTensorRef = typename TensorRef::ConstTensorRef;
@@ -483,10 +496,10 @@ public:
// Data members
//
Index stride; ///< stride in bytes between columns
LongIndex 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
LongIndex advance_row; ///< amount to add to move to the next 'row' position
LongIndex advance_column; ///< amount to add to move to the next 'column' position
//
// Methods
@@ -494,14 +507,16 @@ public:
CUTLASS_HOST_DEVICE
Status initialize(Index stride_) {
stride = stride_;
stride = LongIndex(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;
advance_column = LongIndex(stride_) - ThreadMap::Iterations::kContiguous *
kElementsPerAccess *
sizeof_bits<Element>::value *
ThreadMap::kWarpSize / 8;
return Status::kSuccess;
}
@@ -602,10 +617,10 @@ public:
):
params_(params) {
TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx) +
TensorCoord(threadblock_offset.contiguous() * InterleavedK,
threadblock_offset.strided() / InterleavedK);
TensorCoord(threadblock_offset.contiguous() * InterleavedN,
threadblock_offset.strided() / InterleavedN);
extent_col_ = extent.strided() / InterleavedK;
extent_col_ = extent.strided() / InterleavedN;
thread_start_col_ = thread_offset.strided();
// Initialize predicates
@@ -613,13 +628,13 @@ public:
for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) {
mask_.predicates[c] =
((thread_offset.contiguous() + ThreadMap::Delta::kContiguous * c) <
(extent.contiguous() * InterleavedK));
(extent.contiguous() * InterleavedN));
}
// Initialize pointer
byte_pointer_ = reinterpret_cast<uint8_t *>(pointer) +
thread_offset.strided() * params_.stride +
thread_offset.contiguous() * sizeof(AccessType) / kElementsPerAccess;
LongIndex(thread_offset.strided()) * LongIndex(params_.stride) +
LongIndex(thread_offset.contiguous()) * sizeof(AccessType) / kElementsPerAccess;
// Initialize internal state counter
iteration_contiguous_ = iteration_strided_ = 0;
@@ -634,6 +649,7 @@ public:
/// 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);