ANIKET SHIVAM
co-authored by
Aniket Shivam
Aniket Shivam
parent
ca23ff7924
commit
b72cbf957d
@@ -0,0 +1,326 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2017 - 2022 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 layout functions used by GEMM+permute path for common tensor or matrix formats.
|
||||
|
||||
Like Layout functions, permute layout functions map logical coordinates to linear memory. They often require additional
|
||||
data to describe strides between elements.
|
||||
|
||||
Permute layout functions must implement all members in the interface of NoPermute<> defined in this file. Address offset
|
||||
computation lies in operator() with private member variables {col_permute_, row_permute_ and stride_permute_} as new addresses after permute op.
|
||||
*/
|
||||
#pragma once
|
||||
#if defined(__CUDACC_RTC__)
|
||||
#include <cuda/std/cassert>
|
||||
#else
|
||||
#include "assert.h"
|
||||
#endif
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/fast_math.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 {
|
||||
|
||||
class NoPermute {
|
||||
public:
|
||||
/// Index type used for coordinates
|
||||
using Index = int32_t;
|
||||
|
||||
/// Long index type used for offsets
|
||||
using LongIndex = int64_t;
|
||||
|
||||
private:
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
MatrixCoord extent_;
|
||||
|
||||
Index stride_unit_; // sizeof(AccessType) / kElementsPerAccess in epilogue's predicated_tile_iterator
|
||||
|
||||
Index stride_permute_;
|
||||
|
||||
Index col_permute_;
|
||||
Index row_permute_;
|
||||
|
||||
public:
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_HOST_DEVICE
|
||||
NoPermute() { }
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_HOST_DEVICE
|
||||
NoPermute(MatrixCoord extent, Index stride_init): extent_(extent) { }
|
||||
|
||||
/// Computes the address offset after Permute Op in Bytes
|
||||
CUTLASS_HOST_DEVICE
|
||||
LongIndex operator()(MatrixCoord offset_init) { return 0; }
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
//
|
||||
// Defines permute layouts of various tensor formats.
|
||||
//
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Permute layout function for 4-D permuted tensors with output matrix (dimension as [M, N]) reshaped
|
||||
/// as [M/D1, D1, D2, N/D2]. Then perform permute([0, 2, 1, 3]) on the corresponding output tensor.
|
||||
template <int D1, int D2>
|
||||
class Tensor4DPermute0213 {
|
||||
public:
|
||||
/// Index type used for coordinates
|
||||
using Index = int32_t;
|
||||
|
||||
/// Long index type used for offsets
|
||||
using LongIndex = int64_t;
|
||||
|
||||
private:
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
MatrixCoord extent_;
|
||||
|
||||
Index stride_permute_;
|
||||
|
||||
Index col_permute_;
|
||||
Index row_permute_;
|
||||
|
||||
public:
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_HOST_DEVICE
|
||||
Tensor4DPermute0213() { }
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_HOST_DEVICE
|
||||
Tensor4DPermute0213(MatrixCoord extent, Index stride_init): extent_(extent) {
|
||||
|
||||
/// Update stride_permute with stride_init
|
||||
stride_permute_ = stride_init / D2 * D1; // stride in Elements
|
||||
|
||||
}
|
||||
|
||||
/// Computes the address offset after Permute Op in Bytes
|
||||
CUTLASS_HOST_DEVICE
|
||||
LongIndex operator()(MatrixCoord offset_init) {
|
||||
// Permute as torch.permute(X1, [0, 2, 1, 3]) -> 4D Tensor indices as [i,j,k,l], the dimension of X
|
||||
// is [D0, D1, D2, D3], after permutation the dim of X1 is [D0, D2, D1, D3].
|
||||
assert(extent_.row() % D1 == 0);
|
||||
assert(extent_.column() % D2 == 0);
|
||||
|
||||
int D3 = extent_.column() / D2;
|
||||
|
||||
Index col_init = offset_init.column();
|
||||
Index row_init = offset_init.row();
|
||||
|
||||
int l = col_init % D3;
|
||||
int k = col_init / D3;
|
||||
int j = row_init % D1;
|
||||
int i = row_init / D1;
|
||||
|
||||
// After the Permute Op
|
||||
col_permute_ = l + j * D3;
|
||||
row_permute_ = k + i * D2;
|
||||
|
||||
return LongIndex(row_permute_) * LongIndex(stride_permute_) + LongIndex(col_permute_);
|
||||
}
|
||||
|
||||
/// Return D1
|
||||
CUTLASS_HOST_DEVICE
|
||||
Index d1() const {
|
||||
return D1;
|
||||
}
|
||||
|
||||
/// Return D2
|
||||
CUTLASS_HOST_DEVICE
|
||||
Index d2() const {
|
||||
return D2;
|
||||
}
|
||||
};
|
||||
|
||||
/// Permute layout function for 4-D permuted tensors for BMM with BMM output tensor (dimension as [B, M, N]) reshaped
|
||||
/// as [B/D1, D1, M, N]. Then perform permute([0, 2, 1, 3]) on the corresponding whole BMM output tensor.
|
||||
template <int D1>
|
||||
class Tensor4DPermuteBMM0213 {
|
||||
public:
|
||||
/// Index type used for coordinates
|
||||
using Index = int32_t;
|
||||
|
||||
/// Long index type used for offsets
|
||||
using LongIndex = int64_t;
|
||||
|
||||
private:
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
MatrixCoord extent_;
|
||||
|
||||
Index stride_permute_;
|
||||
|
||||
Index col_permute_;
|
||||
Index row_permute_;
|
||||
|
||||
public:
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_HOST_DEVICE
|
||||
Tensor4DPermuteBMM0213() { }
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_HOST_DEVICE
|
||||
Tensor4DPermuteBMM0213(MatrixCoord extent, Index stride_init): extent_(extent) {
|
||||
|
||||
/// Update stride_permute with stride_init
|
||||
stride_permute_ = stride_init * D1; // stride in Elements
|
||||
|
||||
}
|
||||
|
||||
/// Computes the address offset after Permute Op in Bytes
|
||||
CUTLASS_HOST_DEVICE
|
||||
LongIndex operator()(MatrixCoord offset_init) {
|
||||
|
||||
// The batch index for BMM
|
||||
Index BMM_batch_idx = blockIdx.z;
|
||||
|
||||
// Permute as torch.permute(X1, [0, 2, 1, 3]) -> 4D Tensor indices as [i,j,k,l], the dimension of X
|
||||
// is [D0, D1, D2, D3], after permutation the dim of X1 is [D0, D2, D1, D3].
|
||||
int D2 = extent_.row();
|
||||
int D3 = extent_.column();
|
||||
|
||||
Index col_init = offset_init.column();
|
||||
Index row_init = offset_init.row();
|
||||
|
||||
int l = col_init;
|
||||
int k = row_init;
|
||||
int j = BMM_batch_idx % D1;
|
||||
int i = BMM_batch_idx / D1;
|
||||
|
||||
// After the Permute Op
|
||||
col_permute_ = l + j * D3;
|
||||
row_permute_ = k + i * D2;
|
||||
|
||||
return LongIndex(row_permute_) * LongIndex(stride_permute_) + LongIndex(col_permute_);
|
||||
}
|
||||
|
||||
/// Return D1
|
||||
CUTLASS_HOST_DEVICE
|
||||
Index d1() const {
|
||||
return D1;
|
||||
}
|
||||
};
|
||||
|
||||
/// Permute layout function for 5-D permuted tensors with output matrix (dimension as [M, N]) reshaped
|
||||
/// as [M/T1, T1, T2, T3, N/T3]. Then perform permute([2, 0, 3, 1, 4]) on the corresponding output tensor.
|
||||
template <int T1, int T2, int T3>
|
||||
class Tensor5DPermute20314 {
|
||||
public:
|
||||
/// Index type used for coordinates
|
||||
using Index = int32_t;
|
||||
|
||||
/// Long index type used for offsets
|
||||
using LongIndex = int64_t;
|
||||
|
||||
private:
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
MatrixCoord extent_;
|
||||
|
||||
Index stride_permute_;
|
||||
|
||||
Index col_permute_;
|
||||
Index row_permute_;
|
||||
|
||||
public:
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_HOST_DEVICE
|
||||
Tensor5DPermute20314() { }
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_HOST_DEVICE
|
||||
Tensor5DPermute20314(MatrixCoord extent, Index stride_init): extent_(extent) {
|
||||
|
||||
/// Update stride_permute with stride_init
|
||||
stride_permute_ = stride_init / T2 * T1; // stride in Elements
|
||||
|
||||
}
|
||||
|
||||
/// Computes the address offset after Permute Op in Bytes
|
||||
CUTLASS_HOST_DEVICE
|
||||
LongIndex operator()(MatrixCoord offset_init) {
|
||||
|
||||
// Permute as torch.permute(X1, [2, 0, 3, 1, 4]) -> 5D Tensor indices as [i,j,k,l,m], the dimension of X
|
||||
// is [T0, T1, T2, T3, T4], after permutation the dim of X1 is [T2, T0, T3, T1, T4].
|
||||
int T0 = extent_.row() / T1;
|
||||
int T4 = extent_.column() / T2 / T3;
|
||||
|
||||
Index col_init = offset_init.column();
|
||||
Index row_init = offset_init.row();
|
||||
|
||||
int m = col_init % T4;
|
||||
int l = int(col_init / T4) % T3;
|
||||
int k = int(col_init / T4) / T3;
|
||||
int j = row_init % T1;
|
||||
int i = row_init / T1;
|
||||
|
||||
// After the Permute Op
|
||||
col_permute_ = m + j * T4 + l * T1 * T4;
|
||||
row_permute_ = i + k * T0;
|
||||
|
||||
return LongIndex(row_permute_) * LongIndex(stride_permute_) + LongIndex(col_permute_);
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace layout
|
||||
} // namespace cutlass
|
||||
@@ -48,6 +48,8 @@ namespace layout {
|
||||
|
||||
/// Template based on element size (in bits) - defined in terms of pitch-linear
|
||||
/// memory and Crosswise size (in elements).
|
||||
/// This one is the base class of all Ampere/Turing fp16/bf16/int8/int4/int1
|
||||
/// tensor core kernels. tf32 TN uses this too.
|
||||
template <int ElementSize, int Crosswise>
|
||||
struct TensorOpMultiplicand {
|
||||
/// Logical rank of tensor
|
||||
@@ -321,6 +323,7 @@ struct TensorOpMultiplicandCongruous {
|
||||
|
||||
/// Template based on element size (in bits) - defined in terms of pitch-linear
|
||||
/// memory and Crosswise size (in elements).
|
||||
/// This one is just for TF32 NT kernel.
|
||||
template <int Crosswise>
|
||||
struct TensorOpMultiplicandCongruous<32, Crosswise> {
|
||||
/// Logical rank of tensor
|
||||
|
||||
@@ -29,7 +29,7 @@
|
||||
*
|
||||
**************************************************************************************************/
|
||||
/*! \file
|
||||
\brief
|
||||
\brief layouts needed by Ampere fp64 tensor core kernels.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
Reference in New Issue
Block a user