ANIKET SHIVAM
co-authored by
Aniket Shivam
Aniket Shivam
parent
9b8166e3f0
commit
d572cc1aab
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,510 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. 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.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
/*! \file
|
||||
\brief Defines additional layout functions used in Permute GEMM example to simplify
|
||||
computing reference permutations of 4/5D tensors when source data is column-major.
|
||||
*/
|
||||
#pragma once
|
||||
#if defined(__CUDACC_RTC__)
|
||||
#include <cuda/std/cassert>
|
||||
#else
|
||||
#include "assert.h"
|
||||
#endif
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/layout/pitch_linear.h"
|
||||
#include "cutlass/layout/matrix.h"
|
||||
#include "cutlass/coord.h"
|
||||
#include "cutlass/tensor_coord.h"
|
||||
|
||||
namespace cutlass {
|
||||
namespace layout {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Mapping function for 4-D CWHN tensors.
|
||||
class TensorCWHN {
|
||||
public:
|
||||
/// Logical rank of tensor
|
||||
static int const kRank = 4;
|
||||
|
||||
/// Rank of stride vector
|
||||
static int const kStrideRank = 3;
|
||||
|
||||
/// Index type used for coordinates
|
||||
using Index = int32_t;
|
||||
|
||||
/// Long index type used for offsets
|
||||
using LongIndex = int64_t;
|
||||
|
||||
/// Logical coordinate (n, h, w, c)
|
||||
using TensorCoord = Tensor4DCoord;
|
||||
|
||||
/// Stride vector
|
||||
using Stride = Coord<kStrideRank>;
|
||||
|
||||
private:
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Stride data member - [n, hn, whn]
|
||||
Stride stride_;
|
||||
|
||||
public:
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_HOST_DEVICE
|
||||
TensorCWHN(Stride const &stride = Stride(0)): stride_(stride) { }
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_HOST_DEVICE
|
||||
TensorCWHN(
|
||||
typename Stride::Index stride_h, ///< number of elements between adjacent N coordinates
|
||||
typename Stride::Index stride_w, ///< number of elements between adjacent C coordinates
|
||||
typename Stride::Index stride_c ///< number of elements between adjacent W coordinates
|
||||
):
|
||||
stride_(make_Coord(stride_h, stride_w, stride_c)) { }
|
||||
|
||||
/// Constructor
|
||||
// Once convolutions implement 64b stride this ctor can be deleted
|
||||
CUTLASS_HOST_DEVICE
|
||||
TensorCWHN(Coord<kStrideRank, LongIndex> const &stride):
|
||||
stride_(make_Coord(
|
||||
static_cast<typename Stride::Index>(stride[0]),
|
||||
static_cast<typename Stride::Index>(stride[1]),
|
||||
static_cast<typename Stride::Index>(stride[2]))
|
||||
) { }
|
||||
|
||||
/// Helper returns a layout to a tightly packed WCNH tensor.
|
||||
CUTLASS_HOST_DEVICE
|
||||
static TensorCWHN packed(TensorCoord const &extent) {
|
||||
return TensorCWHN(
|
||||
make_Coord(
|
||||
extent.n(),
|
||||
extent.h() * extent.n(),
|
||||
extent.w() * extent.h() * extent.n()
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
/// Returns the offset of a coordinate (n, h, w, c) in linear memory.
|
||||
CUTLASS_HOST_DEVICE
|
||||
LongIndex operator()(TensorCoord const &coord) const {
|
||||
return coord.n() +
|
||||
LongIndex(stride_[0] * coord.h()) +
|
||||
LongIndex(stride_[1] * coord.w()) +
|
||||
LongIndex(stride_[2] * coord.c());
|
||||
}
|
||||
|
||||
/// Returns the offset of a pitchlinear coordinate in linear memory.
|
||||
CUTLASS_HOST_DEVICE
|
||||
LongIndex operator()(PitchLinearCoord coord) const {
|
||||
return coord.contiguous() + LongIndex(coord.strided() * stride_[2]);
|
||||
}
|
||||
|
||||
/// Returns the stride of the layout
|
||||
CUTLASS_HOST_DEVICE
|
||||
Stride stride() const {
|
||||
return stride_;
|
||||
}
|
||||
|
||||
/// Returns the stride of the layout
|
||||
CUTLASS_HOST_DEVICE
|
||||
Stride & stride() {
|
||||
return stride_;
|
||||
}
|
||||
|
||||
/// Compute the number of contiguous elements needed to store a tensor with the given size
|
||||
CUTLASS_HOST_DEVICE
|
||||
LongIndex capacity(TensorCoord const &extent) const {
|
||||
// it does not make sense if the extent is larger than stride
|
||||
// and we could not rely on the capacity calculation in such cases
|
||||
// we could move this checkers to debug code only
|
||||
if ((extent.n() > stride_[0])
|
||||
|| (extent.h() * stride_[0] > stride_[1])
|
||||
|| (extent.w() * stride_[1] > stride_[2])) {
|
||||
assert(0);
|
||||
}
|
||||
return extent.c() * stride_[2];
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Mapping function for 4-D NHCW tensors.
|
||||
class TensorNHCW {
|
||||
public:
|
||||
/// Logical rank of tensor
|
||||
static int const kRank = 4;
|
||||
|
||||
/// Rank of stride vector
|
||||
static int const kStrideRank = 3;
|
||||
|
||||
/// Index type used for coordinates
|
||||
using Index = int32_t;
|
||||
|
||||
/// Long index type used for offsets
|
||||
using LongIndex = int64_t;
|
||||
|
||||
/// Logical coordinate (n, h, w, c)
|
||||
using TensorCoord = Tensor4DCoord;
|
||||
|
||||
/// Stride vector
|
||||
using Stride = Coord<kStrideRank>;
|
||||
|
||||
private:
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Stride data member - [w, cw, hcw]
|
||||
Stride stride_;
|
||||
|
||||
public:
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_HOST_DEVICE
|
||||
TensorNHCW(Stride const &stride = Stride(0)): stride_(stride) { }
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_HOST_DEVICE
|
||||
TensorNHCW(
|
||||
typename Stride::Index stride_c, ///< number of elements between adjacent C coordinates
|
||||
typename Stride::Index stride_h, ///< number of elements between adjacent H coordinates
|
||||
typename Stride::Index stride_n ///< number of elements between adjacent N coordinates
|
||||
):
|
||||
stride_(make_Coord(stride_c, stride_h, stride_n)) { }
|
||||
|
||||
/// Constructor
|
||||
// Once convolutions implement 64b stride this ctor can be deleted
|
||||
CUTLASS_HOST_DEVICE
|
||||
TensorNHCW(Coord<kStrideRank, LongIndex> const &stride):
|
||||
stride_(make_Coord(
|
||||
static_cast<typename Stride::Index>(stride[0]),
|
||||
static_cast<typename Stride::Index>(stride[1]),
|
||||
static_cast<typename Stride::Index>(stride[2]))
|
||||
) { }
|
||||
|
||||
/// Helper returns a layout to a tightly packed WCNH tensor.
|
||||
CUTLASS_HOST_DEVICE
|
||||
static TensorNHCW packed(TensorCoord const &extent) {
|
||||
return TensorNHCW(
|
||||
make_Coord(
|
||||
extent.w(),
|
||||
extent.c() * extent.w(),
|
||||
extent.h() * extent.c() * extent.w()
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
/// Returns the offset of a coordinate (n, h, w, c) in linear memory.
|
||||
CUTLASS_HOST_DEVICE
|
||||
LongIndex operator()(TensorCoord const &coord) const {
|
||||
return coord.w() +
|
||||
LongIndex(stride_[0] * coord.c()) +
|
||||
LongIndex(stride_[1] * coord.h()) +
|
||||
LongIndex(stride_[2] * coord.n());
|
||||
}
|
||||
|
||||
/// Returns the offset of a pitchlinear coordinate in linear memory.
|
||||
CUTLASS_HOST_DEVICE
|
||||
LongIndex operator()(PitchLinearCoord coord) const {
|
||||
return coord.contiguous() + LongIndex(coord.strided() * stride_[2]);
|
||||
}
|
||||
|
||||
/// Returns the stride of the layout
|
||||
CUTLASS_HOST_DEVICE
|
||||
Stride stride() const {
|
||||
return stride_;
|
||||
}
|
||||
|
||||
/// Returns the stride of the layout
|
||||
CUTLASS_HOST_DEVICE
|
||||
Stride & stride() {
|
||||
return stride_;
|
||||
}
|
||||
|
||||
/// Compute the number of contiguous elements needed to store a tensor with the given size
|
||||
CUTLASS_HOST_DEVICE
|
||||
LongIndex capacity(TensorCoord const &extent) const {
|
||||
// it does not make sense if the extent is larger than stride
|
||||
// and we could not rely on the capacity calculation in such cases
|
||||
// we could move this checkers to debug code only
|
||||
if ((extent.w() > stride_[0])
|
||||
|| (extent.c() * stride_[0] > stride_[1])
|
||||
|| (extent.h() * stride_[1] > stride_[2])) {
|
||||
assert(0);
|
||||
}
|
||||
return extent.n() * stride_[2];
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Mapping function for 4-D NHCW tensors.
|
||||
class TensorNCWH {
|
||||
public:
|
||||
/// Logical rank of tensor
|
||||
static int const kRank = 4;
|
||||
|
||||
/// Rank of stride vector
|
||||
static int const kStrideRank = 3;
|
||||
|
||||
/// Index type used for coordinates
|
||||
using Index = int32_t;
|
||||
|
||||
/// Long index type used for offsets
|
||||
using LongIndex = int64_t;
|
||||
|
||||
/// Logical coordinate (n, h, w, c)
|
||||
using TensorCoord = Tensor4DCoord;
|
||||
|
||||
/// Stride vector
|
||||
using Stride = Coord<kStrideRank>;
|
||||
|
||||
private:
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Stride data member - [h, wh, cwh]
|
||||
Stride stride_;
|
||||
|
||||
public:
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_HOST_DEVICE
|
||||
TensorNCWH(Stride const &stride = Stride(0)): stride_(stride) { }
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_HOST_DEVICE
|
||||
TensorNCWH(
|
||||
typename Stride::Index stride_w, ///< number of elements between adjacent C coordinates
|
||||
typename Stride::Index stride_c, ///< number of elements between adjacent H coordinates
|
||||
typename Stride::Index stride_n ///< number of elements between adjacent N coordinates
|
||||
):
|
||||
stride_(make_Coord(stride_w, stride_c, stride_n)) { }
|
||||
|
||||
/// Constructor
|
||||
// Once convolutions implement 64b stride this ctor can be deleted
|
||||
CUTLASS_HOST_DEVICE
|
||||
TensorNCWH(Coord<kStrideRank, LongIndex> const &stride):
|
||||
stride_(make_Coord(
|
||||
static_cast<typename Stride::Index>(stride[0]),
|
||||
static_cast<typename Stride::Index>(stride[1]),
|
||||
static_cast<typename Stride::Index>(stride[2]))
|
||||
) { }
|
||||
|
||||
/// Helper returns a layout to a tightly packed WCNH tensor.
|
||||
CUTLASS_HOST_DEVICE
|
||||
static TensorNCWH packed(TensorCoord const &extent) {
|
||||
return TensorNCWH(
|
||||
make_Coord(
|
||||
extent.h(),
|
||||
extent.w() * extent.h(),
|
||||
extent.c() * extent.w() * extent.h()
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
/// Returns the offset of a coordinate (n, h, w, c) in linear memory.
|
||||
CUTLASS_HOST_DEVICE
|
||||
LongIndex operator()(TensorCoord const &coord) const {
|
||||
return coord.h() +
|
||||
LongIndex(stride_[0] * coord.w()) +
|
||||
LongIndex(stride_[1] * coord.c()) +
|
||||
LongIndex(stride_[2] * coord.n());
|
||||
}
|
||||
|
||||
/// Returns the offset of a pitchlinear coordinate in linear memory.
|
||||
CUTLASS_HOST_DEVICE
|
||||
LongIndex operator()(PitchLinearCoord coord) const {
|
||||
return coord.contiguous() + LongIndex(coord.strided() * stride_[2]);
|
||||
}
|
||||
|
||||
/// Returns the stride of the layout
|
||||
CUTLASS_HOST_DEVICE
|
||||
Stride stride() const {
|
||||
return stride_;
|
||||
}
|
||||
|
||||
/// Returns the stride of the layout
|
||||
CUTLASS_HOST_DEVICE
|
||||
Stride & stride() {
|
||||
return stride_;
|
||||
}
|
||||
|
||||
/// Compute the number of contiguous elements needed to store a tensor with the given size
|
||||
CUTLASS_HOST_DEVICE
|
||||
LongIndex capacity(TensorCoord const &extent) const {
|
||||
// it does not make sense if the extent is larger than stride
|
||||
// and we could not rely on the capacity calculation in such cases
|
||||
// we could move this checkers to debug code only
|
||||
if ((extent.h() > stride_[0])
|
||||
|| (extent.w() * stride_[0] > stride_[1])
|
||||
|| (extent.c() * stride_[1] > stride_[2])) {
|
||||
assert(0);
|
||||
}
|
||||
return extent.n() * stride_[2];
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Mapping function for 5-D CWHDN tensors.
|
||||
class TensorCWHDN {
|
||||
public:
|
||||
/// Logical rank of tensor
|
||||
static int const kRank = 5;
|
||||
|
||||
/// Rank of stride vector
|
||||
static int const kStrideRank = 4;
|
||||
|
||||
/// Index type used for coordinates
|
||||
using Index = int32_t;
|
||||
|
||||
/// Long index type used for offsets
|
||||
using LongIndex = int64_t;
|
||||
|
||||
/// Logical coordinate (n, d, h, w, c)
|
||||
using TensorCoord = Tensor5DCoord;
|
||||
|
||||
/// Stride vector
|
||||
using Stride = Coord<kStrideRank>;
|
||||
|
||||
private:
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Stride data member - [n, dn, hdn, whdn]
|
||||
Stride stride_;
|
||||
|
||||
public:
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_HOST_DEVICE
|
||||
TensorCWHDN(Stride const &stride = Stride(0)): stride_(stride) { }
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_HOST_DEVICE
|
||||
TensorCWHDN(
|
||||
typename Stride::Index n,
|
||||
typename Stride::Index dn,
|
||||
typename Stride::Index hdn,
|
||||
typename Stride::Index whdn):
|
||||
stride_(make_Coord(n, dn, hdn, whdn)) { }
|
||||
|
||||
/// Constructor
|
||||
// Once convolutions implement 64b stride this ctor can be deleted
|
||||
CUTLASS_HOST_DEVICE
|
||||
TensorCWHDN(Coord<kStrideRank, LongIndex> const &stride):
|
||||
stride_(make_Coord(
|
||||
static_cast<typename Stride::Index>(stride[0]),
|
||||
static_cast<typename Stride::Index>(stride[1]),
|
||||
static_cast<typename Stride::Index>(stride[2]),
|
||||
static_cast<typename Stride::Index>(stride[3]))
|
||||
) { }
|
||||
|
||||
/// Helper returns a layout to a tightly packed CWHDN tensor.
|
||||
CUTLASS_HOST_DEVICE
|
||||
static TensorCWHDN packed(TensorCoord const &extent) {
|
||||
return TensorCWHDN(
|
||||
make_Coord(
|
||||
extent.n(),
|
||||
extent.d() * extent.n(),
|
||||
extent.h() * extent.d() * extent.n(),
|
||||
extent.w() * extent.h() * extent.d() * extent.n()
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
/// Returns the offset of a coordinate (n, d, h, w, c) in linear memory.
|
||||
CUTLASS_HOST_DEVICE
|
||||
LongIndex operator()(TensorCoord const &coord) const {
|
||||
return coord.n() +
|
||||
LongIndex(stride_[0] * coord.d()) +
|
||||
LongIndex(stride_[1] * coord.h()) +
|
||||
LongIndex(stride_[2] * coord.w()) +
|
||||
LongIndex(stride_[3] * coord.c());
|
||||
}
|
||||
|
||||
/// Returns the offset of a pitchlinear coordinate in linear memory.
|
||||
CUTLASS_HOST_DEVICE
|
||||
LongIndex operator()(PitchLinearCoord coord) const {
|
||||
return coord.contiguous() + LongIndex(coord.strided() * stride_[3]);
|
||||
}
|
||||
|
||||
/// Returns the stride of the layout
|
||||
CUTLASS_HOST_DEVICE
|
||||
Stride stride() const {
|
||||
return stride_;
|
||||
}
|
||||
|
||||
/// Returns the stride of the layout
|
||||
CUTLASS_HOST_DEVICE
|
||||
Stride & stride() {
|
||||
return stride_;
|
||||
}
|
||||
|
||||
/// Compute the number of contiguous elements needed to store a tensor with the given size
|
||||
CUTLASS_HOST_DEVICE
|
||||
LongIndex capacity(TensorCoord const &extent) const {
|
||||
// it does not make sense if the extent is larger than stride
|
||||
// and we could not rely on the capacity calculation in such cases
|
||||
// we could move this checkers to debug code only
|
||||
if ((extent.n() > stride_[0])
|
||||
|| (extent.d() * stride_[0] > stride_[1])
|
||||
|| (extent.h() * stride_[1] > stride_[2])
|
||||
|| (extent.w() * stride_[2] > stride_[3])) {
|
||||
assert(0);
|
||||
}
|
||||
return extent.c() * stride_[3];
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace layout
|
||||
} // namespace cutlass
|
||||
@@ -0,0 +1,344 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. 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.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Contains additional metadata about layout permute functions used in the example.
|
||||
*/
|
||||
|
||||
#include "cutlass/tensor_coord.h"
|
||||
#include "cutlass/layout/permute.h"
|
||||
|
||||
/// Additional permutation metadata to facilitate testing/printing
|
||||
template<typename PermuteLayout>
|
||||
struct PermuteInfo;
|
||||
|
||||
/// Specialization for default case (no permute). Other specializations must follow this template.
|
||||
template<>
|
||||
struct PermuteInfo<cutlass::layout::NoPermute> {
|
||||
|
||||
/// Whether this is a BMM or GEMM permutation (NoPermute can actually be either)
|
||||
static bool constexpr kBatched = false;
|
||||
|
||||
/// Minimal divisor for row extent
|
||||
static int constexpr kRowFactor = 1;
|
||||
|
||||
/// Minimum divisor for column extent
|
||||
static int constexpr kColumnFactor = 1;
|
||||
|
||||
/// Minimum divisor for batch size dimension
|
||||
static int constexpr kBatchFactor = 1;
|
||||
|
||||
/// Tensor layout used in permutation operation
|
||||
using Layout = cutlass::layout::PackedVectorLayout;
|
||||
|
||||
static std::string name() {
|
||||
return "NoPermute";
|
||||
}
|
||||
|
||||
/// User-friendly description of the permute operation
|
||||
static std::string desc() {
|
||||
return "no permutation";
|
||||
}
|
||||
|
||||
/// Infer original higher-rank tensor shape from GEMM/BMM matrix extents.
|
||||
/// For direct (output) permutations, must be a simple reshape of extent.
|
||||
/// For inverse (input) permutations, must return shape *before* permute operation.
|
||||
/// In case of NoPermute, simply use a linear (rank 1) view of the memory
|
||||
static Layout::TensorCoord original_shape(cutlass::MatrixCoord extent, int batch_count) {
|
||||
return Layout::TensorCoord(extent.row() * extent.column() * batch_count);
|
||||
}
|
||||
|
||||
/// Compute the permuted higher-rank tensor shape from the original shape.
|
||||
static Layout::TensorCoord permute(Layout::TensorCoord const &s) {
|
||||
return s;
|
||||
}
|
||||
};
|
||||
|
||||
template<int D1>
|
||||
struct PermuteInfo<cutlass::layout::Tensor4DPermuteBMM0213RowMajor<D1>> {
|
||||
|
||||
static bool constexpr kBatched = true;
|
||||
static int constexpr kRowFactor = 1;
|
||||
static int constexpr kColumnFactor = 1;
|
||||
static int constexpr kBatchFactor = D1;
|
||||
|
||||
using Layout = cutlass::layout::TensorNHWC;
|
||||
|
||||
static std::string name() {
|
||||
return "Tensor4DPermuteBMM0213<" + std::to_string(D1) + ">";
|
||||
}
|
||||
|
||||
static std::string desc() {
|
||||
return "batched GEMM permutation [0, 2, 1, 3]";
|
||||
}
|
||||
|
||||
static Layout::TensorCoord original_shape(cutlass::MatrixCoord extent, int batch_count) {
|
||||
int D0 = batch_count / D1;
|
||||
int D2 = extent.row();
|
||||
int D3 = extent.column();
|
||||
return {D0, D1, D2, D3};
|
||||
}
|
||||
|
||||
static Layout::TensorCoord permute(Layout::TensorCoord const &s) {
|
||||
return {s[0], s[2], s[1], s[3]};
|
||||
}
|
||||
};
|
||||
|
||||
template<int D1>
|
||||
struct PermuteInfo<cutlass::layout::Tensor4DPermuteBMM0213RowMajorInverse<D1>>
|
||||
: public PermuteInfo<cutlass::layout::Tensor4DPermuteBMM0213RowMajor<D1>> {
|
||||
|
||||
static bool constexpr kBatched = true;
|
||||
static int constexpr kRowFactor = 1;
|
||||
static int constexpr kColumnFactor = D1;
|
||||
static int constexpr kBatchFactor = 1;
|
||||
|
||||
using Base = PermuteInfo<cutlass::layout::Tensor4DPermuteBMM0213RowMajor<D1>>;
|
||||
using Layout = typename Base::Layout;
|
||||
|
||||
static typename Layout::TensorCoord original_shape(cutlass::MatrixCoord extent, int batch_count) {
|
||||
int D0 = batch_count;
|
||||
int D2 = extent.row();
|
||||
int D3 = extent.column() / D1;
|
||||
return {D0, D1, D2, D3};
|
||||
}
|
||||
};
|
||||
|
||||
template<int D1>
|
||||
struct PermuteInfo<cutlass::layout::Tensor4DPermuteBMM0321ColumnMajor<D1>> {
|
||||
|
||||
static bool constexpr kBatched = true;
|
||||
static int constexpr kRowFactor = 1;
|
||||
static int constexpr kColumnFactor = 1;
|
||||
static int constexpr kBatchFactor = D1;
|
||||
|
||||
using Layout = cutlass::layout::TensorNHCW;
|
||||
|
||||
static std::string name() {
|
||||
return "Tensor4DPermuteBMM0321<" + std::to_string(D1) + ">";
|
||||
}
|
||||
|
||||
static std::string desc() {
|
||||
return "batched GEMM permutation [0, 3, 2, 1]";
|
||||
}
|
||||
|
||||
static Layout::TensorCoord original_shape(cutlass::MatrixCoord extent, int batch_count) {
|
||||
int D0 = batch_count / D1;
|
||||
int D2 = extent.row();
|
||||
int D3 = extent.column();
|
||||
return {D0, D1, D2, D3};
|
||||
}
|
||||
|
||||
static Layout::TensorCoord permute(Layout::TensorCoord const &s) {
|
||||
return {s[0], s[3], s[2], s[1]};
|
||||
}
|
||||
};
|
||||
|
||||
template<int D1>
|
||||
struct PermuteInfo<cutlass::layout::Tensor4DPermuteBMM0321ColumnMajorInverse<D1>>
|
||||
: public PermuteInfo<cutlass::layout::Tensor4DPermuteBMM0321ColumnMajor<D1>> {
|
||||
|
||||
static bool constexpr kBatched = true;
|
||||
static int constexpr kRowFactor = D1;
|
||||
static int constexpr kColumnFactor = 1;
|
||||
static int constexpr kBatchFactor = 1;
|
||||
|
||||
using Base = PermuteInfo<cutlass::layout::Tensor4DPermuteBMM0321ColumnMajor<D1>>;
|
||||
using Layout = typename Base::Layout;
|
||||
|
||||
static typename Layout::TensorCoord original_shape(cutlass::MatrixCoord extent, int batch_count) {
|
||||
int D0 = batch_count;
|
||||
int D2 = extent.row() / D1;
|
||||
int D3 = extent.column();
|
||||
return {D0, D1, D2, D3};
|
||||
}
|
||||
};
|
||||
|
||||
template<int D1, int D2>
|
||||
struct PermuteInfo<cutlass::layout::Tensor4DPermute0213RowMajor<D1, D2>> {
|
||||
|
||||
static bool constexpr kBatched = false;
|
||||
static int constexpr kRowFactor = D1;
|
||||
static int constexpr kColumnFactor = D2;
|
||||
static int constexpr kBatchFactor = 1;
|
||||
|
||||
using Layout = cutlass::layout::TensorNHWC;
|
||||
|
||||
static std::string name() {
|
||||
return "Tensor4DPermute0213<" + std::to_string(D1) + "," + std::to_string(D2) + ">";
|
||||
}
|
||||
|
||||
static std::string desc() {
|
||||
return "normal GEMM permutation [0, 2, 1, 3]";
|
||||
}
|
||||
|
||||
static Layout::TensorCoord original_shape(cutlass::MatrixCoord extent, int batch_count) {
|
||||
int D0 = extent.row() / D1;
|
||||
int D3 = extent.column() / D2;
|
||||
return {D0, D1, D2, D3};
|
||||
}
|
||||
|
||||
static Layout::TensorCoord permute(Layout::TensorCoord const &s) {
|
||||
return {s[0], s[2], s[1], s[3]};
|
||||
}
|
||||
};
|
||||
|
||||
template<int D1, int D2>
|
||||
struct PermuteInfo<cutlass::layout::Tensor4DPermute0213RowMajorInverse<D1, D2>>
|
||||
: public PermuteInfo<cutlass::layout::Tensor4DPermute0213RowMajor<D1, D2>> {
|
||||
|
||||
static bool constexpr kBatched = false;
|
||||
static int constexpr kRowFactor = D2;
|
||||
static int constexpr kColumnFactor = D1;
|
||||
static int constexpr kBatchFactor = 1;
|
||||
|
||||
using Base = PermuteInfo<cutlass::layout::Tensor4DPermute0213RowMajor<D1, D2>>;
|
||||
using Layout = typename Base::Layout;
|
||||
|
||||
static typename Layout::TensorCoord original_shape(cutlass::MatrixCoord extent, int batch_count) {
|
||||
int D0 = extent.row() / D2;
|
||||
int D3 = extent.column() / D1;
|
||||
return {D0, D1, D2, D3};
|
||||
}
|
||||
};
|
||||
|
||||
template<int D1, int D2>
|
||||
struct PermuteInfo<cutlass::layout::Tensor4DPermute0213ColumnMajor<D1, D2>>
|
||||
: public PermuteInfo<cutlass::layout::Tensor4DPermute0213RowMajor<D1, D2>> {
|
||||
using Layout = cutlass::layout::TensorCWHN;
|
||||
};
|
||||
|
||||
template<int D1, int D2>
|
||||
struct PermuteInfo<cutlass::layout::Tensor4DPermute0213ColumnMajorInverse<D1, D2>>
|
||||
: public PermuteInfo<cutlass::layout::Tensor4DPermute0213RowMajorInverse<D1, D2>> {
|
||||
using Layout = cutlass::layout::TensorCWHN;
|
||||
};
|
||||
|
||||
template<int T1, int T2, int T3>
|
||||
struct PermuteInfo<cutlass::layout::Tensor5DPermute20314RowMajor<T1, T2, T3>> {
|
||||
|
||||
static bool constexpr kBatched = false;
|
||||
static int constexpr kRowFactor = T1;
|
||||
static int constexpr kColumnFactor = T2 * T3;
|
||||
static int constexpr kBatchFactor = 1;
|
||||
|
||||
using Layout = cutlass::layout::TensorNDHWC;
|
||||
|
||||
static std::string name() {
|
||||
return "Tensor5DPermute20314<" + std::to_string(T1) + "," + std::to_string(T2) + "," + std::to_string(T3) + ">";
|
||||
}
|
||||
|
||||
static std::string desc() {
|
||||
return "normal GEMM permutation [2, 0, 3, 1, 4]";
|
||||
}
|
||||
|
||||
static Layout::TensorCoord original_shape(cutlass::MatrixCoord extent, int batch_count)
|
||||
{
|
||||
int const T0 = extent.row() / T1;
|
||||
int const T4 = extent.column() / (T2 * T3);
|
||||
return {T0, T1, T2, T3, T4};
|
||||
}
|
||||
|
||||
static Layout::TensorCoord permute(Layout::TensorCoord const &s)
|
||||
{
|
||||
return {s[2], s[0], s[3], s[1], s[4]};
|
||||
}
|
||||
};
|
||||
|
||||
template<int T1, int T2, int T3>
|
||||
struct PermuteInfo<cutlass::layout::Tensor5DPermute20314RowMajorInverse<T1, T2, T3>>
|
||||
: public PermuteInfo<cutlass::layout::Tensor5DPermute20314RowMajor<T1, T2, T3>> {
|
||||
|
||||
static bool constexpr kBatched = false;
|
||||
static int constexpr kRowFactor = T2;
|
||||
static int constexpr kColumnFactor = T1 * T3;
|
||||
static int constexpr kBatchFactor = 1;
|
||||
|
||||
using Base = PermuteInfo<cutlass::layout::Tensor5DPermute20314RowMajor<T1, T2, T3>>;
|
||||
using Layout = typename Base::Layout;
|
||||
|
||||
static typename Layout::TensorCoord original_shape(cutlass::MatrixCoord extent, int batch_count) {
|
||||
int const T0 = extent.row() / T2;
|
||||
int const T4 = extent.column() / (T1 * T3);
|
||||
return {T0, T1, T2, T3, T4};
|
||||
}
|
||||
};
|
||||
|
||||
template<int T1, int T2, int T3>
|
||||
struct PermuteInfo<cutlass::layout::Tensor5DPermute02413ColumnMajor<T1, T2, T3>> {
|
||||
|
||||
static bool constexpr kBatched = false;
|
||||
static int constexpr kRowFactor = T1;
|
||||
static int constexpr kColumnFactor = T2 * T3;
|
||||
static int constexpr kBatchFactor = 1;
|
||||
|
||||
using Layout = cutlass::layout::TensorCWHDN;
|
||||
|
||||
static std::string name() {
|
||||
return "Tensor5DPermute02413<" + std::to_string(T1) + "," + std::to_string(T2) + "," + std::to_string(T3) + ">";
|
||||
}
|
||||
|
||||
static std::string desc() {
|
||||
return "normal GEMM permutation [0, 2, 4, 1, 3]";
|
||||
}
|
||||
|
||||
using Coord = cutlass::Tensor5DCoord;
|
||||
|
||||
static Layout::TensorCoord original_shape(cutlass::MatrixCoord extent, int batch_count)
|
||||
{
|
||||
int const T0 = extent.row() / T1;
|
||||
int const T4 = extent.column() / (T2 * T3);
|
||||
return {T0, T1, T2, T3, T4};
|
||||
}
|
||||
|
||||
static Layout::TensorCoord permute(Layout::TensorCoord const &s)
|
||||
{
|
||||
return {s[0], s[2], s[4], s[1], s[3]};
|
||||
}
|
||||
};
|
||||
|
||||
template<int T1, int T2, int T3>
|
||||
struct PermuteInfo<cutlass::layout::Tensor5DPermute02413ColumnMajorInverse<T1, T2, T3>>
|
||||
: public PermuteInfo<cutlass::layout::Tensor5DPermute02413ColumnMajor<T1, T2, T3>> {
|
||||
|
||||
static bool constexpr kBatched = false;
|
||||
static int constexpr kRowFactor = T2;
|
||||
static int constexpr kColumnFactor = T1 * T3;
|
||||
static int constexpr kBatchFactor = 1;
|
||||
|
||||
using Base = PermuteInfo<cutlass::layout::Tensor5DPermute02413ColumnMajor<T1, T2, T3>>;
|
||||
using Layout = typename Base::Layout;
|
||||
|
||||
static typename Layout::TensorCoord original_shape(cutlass::MatrixCoord extent, int batch_count) {
|
||||
int const T0 = extent.row() / T2;
|
||||
int const T4 = extent.column() / (T1 * T3);
|
||||
return {T0, T1, T2, T3, T4};
|
||||
}
|
||||
};
|
||||
Reference in New Issue
Block a user