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

CUTLASS 2.5

Browse Source
This commit is contained in:
Andrew Kerr
2021-02-26 09:58:26 -05:00
parent ccb697bac7
commit 0e13748649
771 changed files with 15474 additions and 1715 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cutlass/epilogue/threadblock/default_epilogue_complex_tensor_op.h
+8 -1
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2021, 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:
@@ -39,6 +39,11 @@
#include "cutlass/gemm/gemm.h"
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/epilogue/thread/linear_combination_relu.h"
#include "cutlass/epilogue/thread/linear_combination_gelu.h"
#include "cutlass/epilogue/thread/linear_combination_sigmoid.h"
#include "cutlass/epilogue/thread/linear_combination_planar_complex.h"
#include "cutlass/epilogue/thread/conversion_op.h"
#include "cutlass/epilogue/thread/reduction_op.h"
@@ -88,6 +93,7 @@ struct DefaultEpilogueComplexTensorOp {
using OutputOp = OutputOp_;
static int const kElementsPerAccess = ElementsPerAccess;
using Operator = Operator_;
using ElementOutput = typename OutputOp::ElementOutput;
using LayoutC = typename WarpMmaTensorOp::LayoutC;
using ElementAccumulator = typename WarpMmaTensorOp::ElementC;
@@ -173,6 +179,7 @@ struct DefaultEpilogueComplexTensorOp <Shape_, WarpMmaTensorOp_, PartitionsK,
using OutputOp = OutputOp_;
static int const kElementsPerAccess = ElementsPerAccess;
using Operator = arch::OpMultiplyAddGaussianComplex;
using ElementOutput = typename OutputOp::ElementOutput;
using LayoutC = typename WarpMmaTensorOp::LayoutC;
using ElementAccumulator = typename WarpMmaTensorOp::ElementC;
include/cutlass/epilogue/threadblock/default_epilogue_planar_complex.h
+1 -1
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2021, 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:
include/cutlass/epilogue/threadblock/default_epilogue_simt.h
+6 -1
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2021, 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:
@@ -40,6 +40,11 @@
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/epilogue/thread/linear_combination_clamp.h"
#include "cutlass/epilogue/thread/linear_combination_relu.h"
#include "cutlass/epilogue/thread/linear_combination_gelu.h"
#include "cutlass/epilogue/thread/linear_combination_sigmoid.h"
#include "cutlass/epilogue/thread/linear_combination_planar_complex.h"
#include "cutlass/epilogue/thread/conversion_op.h"
#include "cutlass/epilogue/thread/reduction_op.h"
include/cutlass/epilogue/threadblock/default_epilogue_tensor_op.h
+49 -6
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2021, 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:
@@ -40,6 +40,11 @@
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/epilogue/thread/linear_combination_clamp.h"
#include "cutlass/epilogue/thread/linear_combination_relu.h"
#include "cutlass/epilogue/thread/linear_combination_gelu.h"
#include "cutlass/epilogue/thread/linear_combination_sigmoid.h"
#include "cutlass/epilogue/thread/linear_combination_planar_complex.h"
#include "cutlass/epilogue/thread/conversion_op.h"
#include "cutlass/epilogue/thread/reduction_op.h"
@@ -89,6 +94,32 @@ struct DefaultIteratorsTensorOp {
ThreadMap,
ElementAccumulator
>;
static int const kFragmentsPerIteration = 1;
};
/// Partial specialization for float <= float x 4
template <
typename ThreadblockShape,
typename WarpShape,
typename InstructionShape,
typename ThreadMap
>
struct DefaultIteratorsTensorOp<float, float, 4, ThreadblockShape, WarpShape, InstructionShape, ThreadMap> {
using WarpTileIterator = cutlass::epilogue::warp::TileIteratorTensorOp<
WarpShape,
InstructionShape,
float,
layout::RowMajor
>;
using SharedLoadIterator = cutlass::epilogue::threadblock::SharedLoadIterator<
ThreadMap,
float
>;
static int const kFragmentsPerIteration = 2;
};
/// Partial specialization for half <= float x 8 epilogues avoids shared memory bank conflicts.
@@ -125,6 +156,8 @@ struct DefaultIteratorsTensorOp<
8,
8
>;
static int const kFragmentsPerIteration = 2;
};
/// Partial specialization for int8_t x 16 <= int32_t x 16 epilogues avoids shared memory bank conflicts.
@@ -160,6 +193,8 @@ struct DefaultIteratorsTensorOp<
16,
8
>;
static int const kFragmentsPerIteration = 1;
};
/// Partial specialization for int8_t x 8 <= int32_t x 8 epilogues avoids shared memory bank conflicts.
@@ -195,6 +230,8 @@ struct DefaultIteratorsTensorOp<
8,
8
>;
static int const kFragmentsPerIteration = 1;
};
/// Partial specialization for int8_t x 8 <= int32_t x 8 epilogues avoids shared memory bank conflicts.
@@ -230,6 +267,8 @@ struct DefaultIteratorsTensorOp<
8,
8
>;
static int const kFragmentsPerIteration = 1;
};
} // namespace detail
@@ -251,6 +290,7 @@ struct DefaultEpilogueTensorOp {
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;
@@ -303,6 +343,8 @@ struct DefaultEpilogueTensorOp {
/// Hard-coded padding elements added
using Padding = cutlass::MatrixShape<0, 64 / sizeof_bits<ElementAccumulator>::value * 4>;
static int const kFragmentsPerIteration = (kPartitionsK == 1 ? DefaultIterators::kFragmentsPerIteration : 1);
//
// Define the epilogue
//
@@ -315,7 +357,8 @@ struct DefaultEpilogueTensorOp {
WarpTileIterator,
SharedLoadIterator,
OutputOp,
Padding
Padding,
kFragmentsPerIteration
>;
};
@@ -325,7 +368,7 @@ struct DefaultEpilogueTensorOp {
/// 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>
bool isSplitK = false>
struct DefaultInterleavedEpilogueTensorOp {
using Shape = Shape_;
using WarpMmaTensorOp = WarpMmaTensorOp_;
@@ -362,7 +405,7 @@ struct DefaultInterleavedEpilogueTensorOp {
//
using Epilogue = cutlass::epilogue::threadblock::InterleavedEpilogue<
Shape, WarpMmaTensorOp, kPartitionsK, OutputTileIterator,
AccumulatorFragmentIterator, OutputOp, InterleavedK, IsBetaZero>;
AccumulatorFragmentIterator, OutputOp, InterleavedK>;
};
////////////////////////////////////////////////////////////////////////////////
@@ -371,7 +414,7 @@ struct DefaultInterleavedEpilogueTensorOp {
/// 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>
bool isSplitK = false>
struct DefaultInterleavedConvEpilogue {
using Shape = Shape_;
using WarpMmaTensorOp = WarpMmaTensorOp_;
@@ -408,7 +451,7 @@ struct DefaultInterleavedConvEpilogue {
//
using Epilogue = cutlass::epilogue::threadblock::InterleavedEpilogue<
Shape, WarpMmaTensorOp, kPartitionsK, OutputTileIterator,
AccumulatorFragmentIterator, OutputOp, InterleavedK, IsBetaZero>;
AccumulatorFragmentIterator, OutputOp, InterleavedK>;
};
////////////////////////////////////////////////////////////////////////////////
include/cutlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h
+6 -1
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2021, 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:
@@ -40,6 +40,11 @@
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/epilogue/thread/linear_combination_clamp.h"
#include "cutlass/epilogue/thread/linear_combination_relu.h"
#include "cutlass/epilogue/thread/linear_combination_gelu.h"
#include "cutlass/epilogue/thread/linear_combination_sigmoid.h"
#include "cutlass/epilogue/thread/linear_combination_planar_complex.h"
#include "cutlass/epilogue/thread/conversion_op.h"
#include "cutlass/epilogue/thread/reduction_op.h"
include/cutlass/epilogue/threadblock/default_epilogue_wmma_tensor_op.h
+6 -1
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2021, 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:
@@ -40,6 +40,11 @@
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/epilogue/thread/linear_combination_clamp.h"
#include "cutlass/epilogue/thread/linear_combination_relu.h"
#include "cutlass/epilogue/thread/linear_combination_gelu.h"
#include "cutlass/epilogue/thread/linear_combination_sigmoid.h"
#include "cutlass/epilogue/thread/linear_combination_planar_complex.h"
#include "cutlass/epilogue/thread/conversion_op.h"
#include "cutlass/epilogue/thread/reduction_op.h"
include/cutlass/epilogue/threadblock/default_thread_map_simt.h
+1 -1
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2021, 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:
include/cutlass/epilogue/threadblock/default_thread_map_tensor_op.h
+1 -1
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2021, 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:
include/cutlass/epilogue/threadblock/default_thread_map_volta_tensor_op.h
+1 -1
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2021, 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:
include/cutlass/epilogue/threadblock/default_thread_map_wmma_tensor_op.h
+1 -1
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2021, 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:
include/cutlass/epilogue/threadblock/direct_epilogue_tensor_op.h
+1 -1
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2021, 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:
include/cutlass/epilogue/threadblock/epilogue.h
+77 -48
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2021, 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:
@@ -63,7 +63,7 @@ namespace threadblock {
////////////////////////////////////////////////////////////////////////////////
/// Epilogue operator without splitk
/// Epilogue operator
template <
typename Shape_, ///< Shape of threadblock tile (concept: GemmShape)
typename WarpMmaOperator_, ///< Warp-level MMA operator (concept: gemm::warp::MmaTensorOp)
@@ -73,7 +73,8 @@ template <
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)
typename Padding_, ///< Padding added to SMEM allocation to avoid bank conflicts (concept: MatrixShape)
int FragmentsPerPartition = 1 ///< Used to coarsten the epilogue granularity
>
class Epilogue :
public EpilogueBase<
@@ -82,7 +83,8 @@ class Epilogue :
PartitionsK,
AccumulatorFragmentIterator_,
WarpTileIterator_,
Padding_> {
Padding_,
FragmentsPerPartition> {
public:
@@ -92,7 +94,8 @@ public:
PartitionsK,
AccumulatorFragmentIterator_,
WarpTileIterator_,
Padding_>;
Padding_,
FragmentsPerPartition>;
using Shape = Shape_;
using WarpMmaOperator = WarpMmaOperator_;
@@ -113,7 +116,6 @@ public:
/// Accumulator element
using ElementAccumulator = typename WarpTileIterator::Element;
/// Output element
using ElementOutput = typename OutputTileIterator::Element;
@@ -139,6 +141,9 @@ public:
/// Number of warps
using WarpCount = typename Base::WarpCount;
int const kSmemTiles = Base::kFragmentsPerIteration > 1 ? Base::kFragmentsPerIteration : kPartitionsK;
int const kSmemPointerOffset = Base::SharedStorage::StorageShape::kCount / kSmemTiles;
public:
@@ -166,7 +171,10 @@ public:
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) { }
shared_load_iterator_(shared_storage.reference(), thread_idx)
{
}
/// Streams the result to global memory
CUTLASS_DEVICE
@@ -177,7 +185,7 @@ public:
OutputTileIterator source_iterator) { ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles)
if (!output_op.is_source_needed()) {
compute_source_not_needed_(output_op, destination_iterator, accumulators);
compute_source_not_needed_(output_op, destination_iterator, accumulators);
}
else {
compute_source_needed_(output_op, destination_iterator, accumulators, source_iterator);
@@ -185,6 +193,8 @@ public:
}
private:
static_assert(kPartitionsK == 1 || Base::kFragmentsPerIteration == 1, "One of these must be exactly 1.");
/// Streams the result to global memory
CUTLASS_DEVICE
@@ -205,7 +215,7 @@ private:
//
CUTLASS_PRAGMA_UNROLL
for (int iter = 0; iter < OutputTileIterator::kIterations; ++iter) {
for (int iter = 0; iter < OutputTileIterator::kIterations; iter += Base::kFragmentsPerIteration) {
//
// Convert and store fragment
@@ -213,12 +223,24 @@ private:
__syncthreads();
typename AccumulatorFragmentIterator::Fragment accum_fragment;
accum_fragment_iterator.load(accum_fragment);
++accum_fragment_iterator;
CUTLASS_PRAGMA_UNROLL
for (int p = 0; p < Base::kFragmentsPerIteration; ++p) {
typename AccumulatorFragmentIterator::Fragment accum_fragment;
this->warp_tile_iterator_.store(accum_fragment);
accum_fragment_iterator.load(accum_fragment);
++accum_fragment_iterator;
this->warp_tile_iterator_.store(accum_fragment);
if (p < Base::kFragmentsPerIteration - 1) {
this->warp_tile_iterator_.add_pointer_offset(kSmemPointerOffset);
}
}
if (Base::kFragmentsPerIteration > 1) {
this->warp_tile_iterator_.add_pointer_offset(kSmemPointerOffset * (1 - Base::kFragmentsPerIteration));
}
__syncthreads();
@@ -226,45 +248,53 @@ private:
// Load fragments from shared memory
//
typename SharedLoadIterator::Fragment aligned_accum_fragment[kPartitionsK];
CUTLASS_PRAGMA_UNROLL
for (int p = 0; p < Base::kFragmentsPerIteration; ++p) {
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;
typename SharedLoadIterator::Fragment aligned_accum_fragment[kPartitionsK];
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_.load(aligned_accum_fragment[0]);
if (p < Base::kFragmentsPerIteration - 1) {
shared_load_iterator_.add_pointer_offset(kSmemPointerOffset);
}
else if (kPartitionsK > 1) {
plus <typename SharedLoadIterator::Fragment> add_fragments;
CUTLASS_PRAGMA_UNROLL
for ( int i = 1; i < kPartitionsK; ++i) {
shared_load_iterator_.add_pointer_offset(kSmemPointerOffset);
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_pointer_offset((1 - kPartitionsK) * kSmemPointerOffset);
}
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_source_not_needed_(output_fragment, output_op, aligned_accum_fragment[0]);
//
// Store the final result
//
destination_iterator.store(output_fragment);
++destination_iterator;
}
//
// Compute the output result
//
typename OutputTileIterator::Fragment output_fragment;
apply_output_operator_source_not_needed_(output_fragment, output_op, aligned_accum_fragment[0]);
//
// Store the final result
//
destination_iterator.store(output_fragment);
++destination_iterator;
if (Base::kFragmentsPerIteration > 1) {
shared_load_iterator_.add_pointer_offset(kSmemPointerOffset * (1 - Base::kFragmentsPerIteration));
}
}
}
/// Streams the result to global memory
CUTLASS_DEVICE
@@ -323,19 +353,18 @@ private:
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)
{
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_.add_pointer_offset(kSmemPointerOffset);
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});
shared_load_iterator_.add_pointer_offset((1 - kPartitionsK) * kSmemPointerOffset);
}
//
include/cutlass/epilogue/threadblock/epilogue_base.h
+7 -3
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2021, 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:
@@ -66,7 +66,8 @@ template <
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)
typename Padding_, ///< Padding added to SMEM allocation to avoid bank conflicts (concept: MatrixShape)
int FragmentsPerIteration = 1
>
class EpilogueBase {
public:
@@ -94,6 +95,9 @@ public:
kPartitionsK
>;
/// Use this to control the granularity of one epilogue 'iteration'
static int const kFragmentsPerIteration = FragmentsPerIteration;
public:
/// Shared storage allocation needed by the epilogue
@@ -120,7 +124,7 @@ public:
/// Shape of the shared memory allocation for the epilogue
using StorageShape = MatrixShape<
Shape::kRow + Padding::kRow,
(Shape::kRow + Padding::kRow) * kFragmentsPerIteration,
Shape::kColumn + Padding::kColumn
>;
include/cutlass/epilogue/threadblock/epilogue_planar_complex.h
+1 -1
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2021, 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:
include/cutlass/epilogue/threadblock/epilogue_workspace.h
+1 -1
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2021, 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:
include/cutlass/epilogue/threadblock/interleaved_epilogue.h
+92 -18
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2021, 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:
@@ -73,9 +73,7 @@ template <
/// Output operator
typename OutputOp_,
/// Number of interleaved k
int InterleavedK,
/// Whether Beta is zero
bool IsBetaZero = false>
int InterleavedK>
class InterleavedEpilogue {
public:
using Shape = Shape_;
@@ -149,22 +147,76 @@ class InterleavedEpilogue {
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)
if (!output_op.is_source_needed()) {
compute_source_not_needed_(output_op, destination_iterator, accumulators);
}
else {
compute_source_needed_(output_op, destination_iterator, accumulators, source_iterator);
}
}
/// Streams the result to global memory
CUTLASS_DEVICE
void compute_source_not_needed_(
OutputOp const &output_op, ///< Output operator
OutputTileIterator destination_iterator, ///< Tile iterator for destination
AccumulatorTile const &accumulators ///< Complete warp-level accumulator tile
) {
//
// 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) {
//
// 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_source_not_needed_(output_op, output_fragment, accum_fragment);
//
// Store the final result
//
destination_iterator.set_iteration_index(iter);
destination_iterator.store(output_fragment);
++destination_iterator;
}
}
/// Streams the result to global memory
CUTLASS_DEVICE
void compute_source_needed_(
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();
//
@@ -183,11 +235,9 @@ class InterleavedEpilogue {
// Load the source
//
if (!IsBetaZero) {
source_iterator.set_iteration_index(iter);
source_iterator.load(source_fragment);
++source_iterator;
}
source_iterator.set_iteration_index(iter);
source_iterator.load(source_fragment);
++source_iterator;
//
// Convert fragment
@@ -243,6 +293,30 @@ class InterleavedEpilogue {
output_frag_ptr[i] = output_op(compute_frag_ptr[i], source_frag_ptr[i]);
}
}
/// Helper to invoke the output functor over each vector of output
CUTLASS_DEVICE
void apply_output_operator_source_not_needed_(
OutputOp const &output_op, ///< Output operator
typename OutputTileIterator::Fragment &output_fragment,
typename AccumulatorFragmentIterator::Fragment const
&aligned_accum_fragment) {
OutputAccessType *output_frag_ptr =
reinterpret_cast<OutputAccessType *>(&output_fragment);
AccumulatorAccessType const *compute_frag_ptr =
reinterpret_cast<AccumulatorAccessType const *>(
&aligned_accum_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]);
}
}
};
////////////////////////////////////////////////////////////////////////////////
include/cutlass/epilogue/threadblock/output_tile_thread_map.h
+6 -1
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2021, 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:
@@ -371,6 +371,11 @@ struct OutputTileOptimalThreadMap {
using Shape = Shape_;
using TileShape = MatrixShape<
Shape::kTile * Shape::kCluster * Shape::kGroup * Shape::kRow,
Shape::kColumn
>;
using Iterations = OutputTileShape<
Detail::RowArrangement::kIterationsColumn,
Detail::RowArrangement::kIterationsRow,
include/cutlass/epilogue/threadblock/predicated_tile_iterator.h
+22 -18
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2021, 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:
@@ -176,7 +176,7 @@ private:
/// Internal state counter
int state_[3];
private:
//
@@ -214,6 +214,11 @@ public:
+ ThreadMap::Delta::kColumn * c) < extent.column());
}
// Null pointer performs no accesses
if (!pointer) {
mask_.clear();
}
// Initialize pointer
byte_pointer_ = reinterpret_cast<uint8_t *>(pointer) +
LongIndex(thread_offset.row()) * LongIndex(params_.stride) +
@@ -288,11 +293,12 @@ public:
}
}
/// Loads a fragment from memory
CUTLASS_DEVICE
void load(Fragment &frag) {
load_with_byte_offset(frag, 0);
load_with_byte_offset(frag, 0);
}
/// Stores a fragment to memory
@@ -326,11 +332,10 @@ public:
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];
}
cutlass::arch::global_store<AccessType, sizeof(AccessType)>(
frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column],
(void *)&memory_pointer[column * ThreadMap::Delta::kColumn / kElementsPerAccess],
guard);
}
if (row + 1 < ThreadMap::Iterations::kRow) {
@@ -349,11 +354,12 @@ public:
}
}
/// Stores a fragment to memory
CUTLASS_DEVICE
void store(Fragment const &frag) {
store_with_byte_offset(frag, 0);
store_with_byte_offset(frag, 0);
}
/// Advances to the next position to load or store
@@ -404,7 +410,7 @@ public:
///< Sets the mask
CUTLASS_DEVICE void get_mask(Mask &mask) {
return mask_;
mask = mask_;
}
///< Sets the mask
@@ -644,9 +650,8 @@ public:
bool guard = col_guard && mask_.predicates[iteration_contiguous_];
if (guard) {
*memory_pointer = *frag_ptr;
}
cutlass::arch::global_store<AccessType, sizeof(AccessType)>(
*frag_ptr, (void *)memory_pointer, guard);
}
/// Overrides the internal iteration index
@@ -689,7 +694,7 @@ public:
///< Sets the mask
CUTLASS_DEVICE void get_mask(Mask &mask) {
return mask_;
mask = mask_;
}
///< Sets the mask
@@ -949,9 +954,8 @@ public:
AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
AccessType *memory_pointer = reinterpret_cast<AccessType *>(byte_pointer);
if (guard) {
*memory_pointer = *frag_ptr;
}
cutlass::arch::global_store<AccessType, sizeof(AccessType)>(
*frag_ptr, (void *)memory_pointer, guard);
}
/// Overrides the internal iteration index
@@ -993,7 +997,7 @@ public:
///< Sets the mask
CUTLASS_DEVICE void get_mask(Mask &mask) {
return mask_;
mask = mask_;
}
///< Sets the mask
include/cutlass/epilogue/threadblock/predicated_tile_iterator_params.h
+1 -1
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2021, 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:
include/cutlass/epilogue/threadblock/shared_load_iterator.h
+5 -3
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2021, 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:
@@ -61,7 +61,7 @@ template <
class SharedLoadIterator {
public:
using ThreadMap = ThreadMap_;
using Shape = typename ThreadMap::Shape;
using Shape = typename ThreadMap::TileShape;
using Element = Element_;
@@ -151,7 +151,9 @@ public:
CUTLASS_DEVICE
void add_tile_offset(TensorCoord const &offset) {
add_pointer_offset(offset.row() * stride_ / (sizeof_bits<Element>::value / 8) + offset.column() * Shape::kColumn);
byte_pointer_ +=
offset.row() * Shape::kRow * stride_ +
offset.column() * Shape::kColumn * sizeof_bits<Element>::value / 8;
}
/// Loads a fragment from memory
include/cutlass/epilogue/threadblock/shared_load_iterator_mixed.h
+12 -6
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2021, 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:
@@ -171,7 +171,7 @@ public:
void add_pointer_offset(LongIndex pointer_offset) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < kLoadsPerAccess; ++i) {
pointers_ += pointer_offset / LoadType::kElements;
pointers_[i] += pointer_offset / LoadType::kElements;
}
}
@@ -179,7 +179,9 @@ public:
void add_tile_offset(TensorCoord const &offset) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < kLoadsPerAccess; ++i) {
pointers_[i] += offset.row() * stride_ + offset.column() / LoadType::kElements;
pointers_[i] +=
offset.row() * Shape::kRow * stride_ +
offset.column() * Shape::kColumn / LoadType::kElements;
}
}
@@ -236,7 +238,7 @@ public:
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Partial specialization for int32_t x 16 => int8_t x 16
/// Partial specialization for int32_t x 8 => int8_t x 8
template <
typename ThreadMap_ ///< Thread map (conept: OutputTileThreadMap)
>
@@ -339,7 +341,9 @@ public:
void add_tile_offset(TensorCoord const &offset) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < kLoadsPerAccess; ++i) {
pointers_[i] += offset.row() * stride_ + offset.column() / LoadType::kElements;
pointers_[i] +=
offset.row() * Shape::kRow * stride_ +
offset.column() * Shape::kColumn / LoadType::kElements;
}
}
@@ -497,7 +501,9 @@ public:
void add_tile_offset(TensorCoord const &offset) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < kLoadsPerAccess; ++i) {
pointers_[i] += offset.row() * stride_ + offset.column() / LoadType::kElements;
pointers_[i] +=
offset.row() * Shape::kRow * stride_ +
offset.column() * Shape::kColumn / LoadType::kElements;
}
}