data: Cu10 CUTLASS headers from corex-samples
This commit is contained in:
root
345
cat_files/batched_gemm.cu
Normal file
345
cat_files/batched_gemm.cu
Normal file
@@ -0,0 +1,345 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2017-2020, 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.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/layout/matrix.h"
|
||||
#include "cutlass/gemm/device/gemm_batched.h"
|
||||
|
||||
#pragma warning( disable : 4503)
|
||||
|
||||
/*
|
||||
This example demonstrates how to use cutlass to compute a batched strided gemm.
|
||||
In this example, both A and B matrix are non-transpose and column major matrix
|
||||
batched_C = batched_A x batched_B
|
||||
As an example, matrix C can be seen as
|
||||
-----------------------------------------------------------
|
||||
(0,0,0) | (0,0,1) | (0,0,2) | (1,0,0) | (1,0,1) | (1,0,2) |
|
||||
-----------------------------------------------------------
|
||||
(0,1,0) | (0,1,1) | (0,1,2) | (1,1,0) | (1,1,1) | (1,1,2) |
|
||||
-----------------------------------------------------------
|
||||
(0,2,0) | (0,2,1) | (0,2,2) | (1,2,0) | (1,2,1) | (1,2,2) |
|
||||
-----------------------------------------------------------
|
||||
(0,3,0) | (0,3,1) | (0,3,2) | (1,3,0) | (1,3,1) | (1,3,2) |
|
||||
-----------------------------------------------------------
|
||||
(0,4,0) | (0,4,1) | (0,4,2) | (1,4,0) | (1,4,1) | (1,4,2) |
|
||||
-----------------------------------------------------------
|
||||
(0,5,0) | (0,5,1) | (0,5,2) | (1,5,0) | (1,5,1) | (1,5,2) |
|
||||
-----------------------------------------------------------
|
||||
batch 0 | batch 1
|
||||
where we denote each element with (batch_idx, row_idx, column_idx)
|
||||
In this example, batch size is 2, M is 6 and N is 3
|
||||
The stride (batch_stride_C) between the first element of two batches is ldc * n
|
||||
|
||||
matrix A can be seen as
|
||||
---------------------------------------
|
||||
(0,0,0) | (0,0,1) | (1,0,0) | (1,0,1) |
|
||||
---------------------------------------
|
||||
(0,1,0) | (0,1,1) | (1,1,0) | (1,1,1) |
|
||||
---------------------------------------
|
||||
(0,2,0) | (0,2,1) | (1,2,0) | (1,2,1) |
|
||||
---------------------------------------
|
||||
(0,3,0) | (0,3,1) | (1,3,0) | (1,3,1) |
|
||||
---------------------------------------
|
||||
(0,4,0) | (0,4,1) | (1,4,0) | (1,4,1) |
|
||||
---------------------------------------
|
||||
(0,5,0) | (0,5,1) | (1,5,0) | (1,5,1) |
|
||||
---------------------------------------
|
||||
batch 0 | batch 1
|
||||
, where batch size is 2, M is 6 and K is 2
|
||||
The stride (batch_stride_B) between the first element of two batches is lda * k
|
||||
|
||||
matrix B can be seen as
|
||||
-----------------------------
|
||||
(0,0,0) | (0,0,1) | (0,0,2) |
|
||||
----------------------------- batch 0
|
||||
(0,1,0) | (0,1,1) | (0,1,2) |
|
||||
-------------------------------------
|
||||
(1,0,0) | (1,0,1) | (1,0,2) |
|
||||
----------------------------- batch 1
|
||||
(1,1,0) | (1,1,1) | (1,1,2) |
|
||||
-----------------------------
|
||||
, where the batch size is 2, N is 3 and K is 2
|
||||
The stride (batch_stride_C) between the first element of two batches is k
|
||||
|
||||
|
||||
*/
|
||||
|
||||
cudaError_t cutlass_strided_batched_sgemm(
|
||||
int m,
|
||||
int n,
|
||||
int k,
|
||||
float alpha,
|
||||
float const *A,
|
||||
int lda,
|
||||
long long int batch_stride_A,
|
||||
float const *B,
|
||||
int ldb,
|
||||
long long int batch_stride_B,
|
||||
float *C,
|
||||
int ldc,
|
||||
long long int batch_stride_C,
|
||||
float beta,
|
||||
int batch_count) {
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmBatched<
|
||||
float, cutlass::layout::ColumnMajor,
|
||||
float, cutlass::layout::ColumnMajor,
|
||||
float, cutlass::layout::ColumnMajor
|
||||
>;
|
||||
|
||||
Gemm gemm_op;
|
||||
|
||||
cutlass::Status status = gemm_op({
|
||||
{m, n, k},
|
||||
{A, lda},
|
||||
batch_stride_A,
|
||||
{B, ldb},
|
||||
batch_stride_B,
|
||||
{C, ldc},
|
||||
batch_stride_C,
|
||||
{C, ldc},
|
||||
batch_stride_C,
|
||||
{alpha, beta},
|
||||
batch_count
|
||||
});
|
||||
|
||||
if (status != cutlass::Status::kSuccess) {
|
||||
return cudaErrorUnknown;
|
||||
}
|
||||
|
||||
return cudaSuccess;
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
cudaError_t strided_batched_gemm_nn_reference(
|
||||
int m,
|
||||
int n,
|
||||
int k,
|
||||
T alpha,
|
||||
std::vector<T> const &A,
|
||||
int lda,
|
||||
long long int batch_stride_A,
|
||||
std::vector<T> const &B,
|
||||
int ldb,
|
||||
long long int batch_stride_B,
|
||||
std::vector<T> &C,
|
||||
int ldc,
|
||||
long long int batch_stride_C,
|
||||
T beta,
|
||||
int batch_count) {
|
||||
/*
|
||||
strided batched gemm NN
|
||||
*/
|
||||
|
||||
cudaError_t result = cudaSuccess;
|
||||
|
||||
if (A.size() < lda * k * batch_count) {
|
||||
std::cout << "the size of A is too small" << std::endl;
|
||||
return cudaErrorInvalidValue;
|
||||
}
|
||||
if (B.size() < ldb * n) {
|
||||
std::cout << "the size of B is too small" << std::endl;
|
||||
return cudaErrorInvalidValue;
|
||||
}
|
||||
if (C.size() < ldc * n * batch_count) {
|
||||
std::cout << "the size of C is too small" << std::endl;
|
||||
return cudaErrorInvalidValue;
|
||||
}
|
||||
|
||||
for (int batch_idx = 0; batch_idx < batch_count; batch_idx++) {
|
||||
for (int n_idx = 0; n_idx < n; n_idx++) {
|
||||
for (int m_idx = 0; m_idx < m; m_idx++) {
|
||||
T accum = beta * C[batch_idx * batch_stride_C + n_idx * ldc + m_idx];
|
||||
for (int k_idx = 0; k_idx < k; k_idx++) {
|
||||
accum += alpha
|
||||
* A[batch_idx * batch_stride_A + k_idx * lda + m_idx]
|
||||
* B[batch_idx * batch_stride_B + n_idx * ldb + k_idx];
|
||||
}
|
||||
C[batch_idx * batch_stride_C + n_idx * ldc + m_idx] = accum;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
int main() {
|
||||
|
||||
// Arbitrary problem size
|
||||
int const m = 520;
|
||||
int const n = 219;
|
||||
int const k = 129;
|
||||
int const batch_count = 17;
|
||||
|
||||
// A, B are non-transpose, column major
|
||||
int const lda = m;
|
||||
int const ldb = k * batch_count;
|
||||
int const ldc = m;
|
||||
|
||||
int const count_A = batch_count * lda * k;
|
||||
int const count_B = ldb * n;
|
||||
int const count_C = batch_count * ldc * n;
|
||||
|
||||
// the memory is batched along K dimension
|
||||
long long int batch_stride_A = static_cast<long long int>(lda) * static_cast<long long int>(k);
|
||||
long long int batch_stride_B = static_cast<long long int>(k);
|
||||
long long int batch_stride_C = static_cast<long long int>(ldc) * static_cast<long long int>(n);
|
||||
|
||||
// alpha and beta
|
||||
float alpha = 1.0f;
|
||||
float beta = 2.0f;
|
||||
|
||||
cudaError_t result = cudaSuccess;
|
||||
|
||||
// allocate the host memory
|
||||
std::vector<float> host_A(count_A);
|
||||
std::vector<float> host_B(count_B);
|
||||
std::vector<float> host_C(count_C);
|
||||
std::vector<float> result_C(count_C);
|
||||
|
||||
// allocate the device memory
|
||||
float *A;
|
||||
float *B;
|
||||
float *C;
|
||||
|
||||
result = cudaMalloc(&A, count_A * sizeof(float));
|
||||
if (result != cudaSuccess) {
|
||||
std::cerr << "cudaMalloc result = " << result << std::endl;
|
||||
return result;
|
||||
}
|
||||
result = cudaMalloc(&B, count_B * sizeof(float));
|
||||
if (result != cudaSuccess) {
|
||||
std::cerr << "cudaMalloc result = " << result << std::endl;
|
||||
return result;
|
||||
}
|
||||
result = cudaMalloc(&C, count_C * sizeof(float));
|
||||
if (result != cudaSuccess) {
|
||||
std::cerr << "cudaMalloc result = " << result << std::endl;
|
||||
return result;
|
||||
}
|
||||
|
||||
// Limit range to avoid floating-point errors
|
||||
int const kRange = 8;
|
||||
|
||||
// fill A
|
||||
for (int b_idx = 0; b_idx < batch_count; b_idx++) {
|
||||
for (int col_idx = 0; col_idx < k; col_idx++) {
|
||||
for (int row_idx = 0; row_idx < m; row_idx++) {
|
||||
host_A[row_idx + col_idx * lda + b_idx * lda * k] = static_cast<float>((row_idx + col_idx * lda + b_idx * lda * k) % kRange);
|
||||
}
|
||||
}
|
||||
}
|
||||
// fill B
|
||||
for (int b_idx = 0; b_idx < batch_count; b_idx++) {
|
||||
for (int col_idx = 0; col_idx < n; col_idx++) {
|
||||
for (int row_idx = 0; row_idx < k; row_idx++) {
|
||||
host_B[row_idx + col_idx * ldb + b_idx * k] = static_cast<float>(((n + k * ldb + batch_count * k) - (row_idx + col_idx * ldb + b_idx * k)) % kRange);
|
||||
}
|
||||
}
|
||||
}
|
||||
// fill C
|
||||
for (int b_idx = 0; b_idx < batch_count; b_idx++) {
|
||||
for (int col_idx = 0; col_idx < n; col_idx++) {
|
||||
for (int row_idx = 0; row_idx < m; row_idx++) {
|
||||
host_C[row_idx + col_idx * ldc + b_idx * ldc * n] = 1.f;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ref memory
|
||||
std::vector<float> ref_A(host_A);
|
||||
std::vector<float> ref_B(host_B);
|
||||
std::vector<float> ref_C(host_C);
|
||||
// copy host memory to device
|
||||
result = cudaMemcpy(A, host_A.data(), count_A * sizeof(float), cudaMemcpyHostToDevice);
|
||||
if (result != cudaSuccess) {
|
||||
std::cerr << "cudaMemcpy result = " << result << std::endl;
|
||||
return result;
|
||||
}
|
||||
result = cudaMemcpy(B, host_B.data(), count_B * sizeof(float), cudaMemcpyHostToDevice);
|
||||
if (result != cudaSuccess) {
|
||||
std::cerr << "cudaMemcpy result = " << result << std::endl;
|
||||
return result;
|
||||
}
|
||||
result = cudaMemcpy(C, host_C.data(), count_C * sizeof(float), cudaMemcpyHostToDevice);
|
||||
if (result != cudaSuccess) {
|
||||
std::cerr << "cudaMemcpy result = " << result << std::endl;
|
||||
return result;
|
||||
}
|
||||
|
||||
// run cutlass
|
||||
result = cutlass_strided_batched_sgemm(
|
||||
m, n, k, alpha, A, lda, batch_stride_A, B, ldb, batch_stride_B, C, ldc, batch_stride_C,
|
||||
beta, batch_count);
|
||||
if (result != cudaSuccess)
|
||||
return result;
|
||||
|
||||
// copy device memory to host
|
||||
result = cudaMemcpy(result_C.data(), C, count_C * sizeof(float), cudaMemcpyDeviceToHost);
|
||||
if (result != cudaSuccess) {
|
||||
std::cerr << "cudaMemcpy result = " << result << std::endl;
|
||||
return result;
|
||||
}
|
||||
|
||||
//compare with reference code
|
||||
result = strided_batched_gemm_nn_reference(m, n, k, alpha, ref_A, lda, batch_stride_A, ref_B, ldb, batch_stride_B, ref_C, ldc, batch_stride_C,
|
||||
beta, batch_count);
|
||||
if (result != 0)
|
||||
return result;
|
||||
|
||||
// Expect bit-level accuracy for this simple example
|
||||
if (ref_C != result_C) {
|
||||
std::cout << "CUTLASS strided batched gemm does not run correctly" << std::endl;
|
||||
return cudaErrorUnknown;
|
||||
}
|
||||
|
||||
// free memory
|
||||
result = cudaFree(A);
|
||||
if (result != cudaSuccess) {
|
||||
std::cerr << "cudaFree result = " << result << std::endl;
|
||||
return result;
|
||||
}
|
||||
result = cudaFree(B);
|
||||
if (result != cudaSuccess) {
|
||||
std::cerr << "cudaFree result = " << result << std::endl;
|
||||
return result;
|
||||
}
|
||||
result = cudaFree(C);
|
||||
if (result != cudaSuccess) {
|
||||
std::cerr << "cudaFree result = " << result << std::endl;
|
||||
return result;
|
||||
}
|
||||
|
||||
|
||||
if (result == cudaSuccess) {
|
||||
std::cout << "Passed." << std::endl;
|
||||
}
|
||||
|
||||
// Exit.
|
||||
return result == cudaSuccess ? 0 : -1;
|
||||
}
|
||||
4304
cat_files/cutlass_samples_tree.txt
Normal file
4304
cat_files/cutlass_samples_tree.txt
Normal file
File diff suppressed because it is too large
Load Diff
383
cat_files/default_gemm.h
Normal file
383
cat_files/default_gemm.h
Normal file
@@ -0,0 +1,383 @@
|
||||
/***************************************************************************************************
|
||||
* 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:
|
||||
* * 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.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2021 Iluvatar CoreX. All rights reserved.
|
||||
* Copyright Declaration: This software, including all of its code and documentation,
|
||||
* except for the third-party software it contains, is a copyrighted work of Shanghai Iluvatar CoreX
|
||||
* Semiconductor Co., Ltd. and its affiliates ("Iluvatar CoreX") in accordance with the PRC Copyright
|
||||
* Law and relevant international treaties, and all rights contained therein are enjoyed by Iluvatar
|
||||
* CoreX. No user of this software shall have any right, ownership or interest in this software and
|
||||
* any use of this software shall be in compliance with the terms and conditions of the End User
|
||||
* License Agreement.
|
||||
**************************************************************************************************/
|
||||
|
||||
|
||||
/*! \file
|
||||
\brief
|
||||
Default kernel-level GEMM definitions combine threadblock-scoped matrix multiply-add with
|
||||
the appropriate threadblock-scoped epilogue.
|
||||
|
||||
Note, CUTLASS epilogues universally target row-major outputs. Column-major outputs are
|
||||
accommodated by exchanging A and B operands and assuming transposed layouts. Partial
|
||||
specializations here choose 'device::GemmTransposed' to implement this functionality.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
|
||||
#include "cutlass/layout/matrix.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/arch/mma.h"
|
||||
|
||||
#include "cutlass/epilogue/threadblock/epilogue.h"
|
||||
#include "cutlass/epilogue/thread/linear_combination.h"
|
||||
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
#include "cutlass/gemm/kernel/gemm.h"
|
||||
#include "cutlass/gemm/kernel/gemm_pipelined.h"
|
||||
#include "cutlass/gemm/threadblock/default_mma.h"
|
||||
#include "cutlass/gemm/threadblock/default_mma_core_simt.h"
|
||||
#include "cutlass/gemm/threadblock/threadblock_swizzle.h"
|
||||
|
||||
#include "cutlass/epilogue/threadblock/default_epilogue_simt.h"
|
||||
#include "cutlass/epilogue/threadblock/default_epilogue_tensor_op.h"
|
||||
#include "cutlass/transform/threadblock/predicated_tile_iterator.h"
|
||||
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace gemm {
|
||||
namespace kernel {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
/// Element type for A matrix operand
|
||||
typename ElementA_,
|
||||
/// Layout type for A matrix operand
|
||||
typename LayoutA_,
|
||||
/// Access granularity of A matrix in units of elements
|
||||
int kAlignmentA,
|
||||
/// Element type for B matrix operand
|
||||
typename ElementB_,
|
||||
/// Layout type for B matrix operand
|
||||
typename LayoutB_,
|
||||
/// Access granularity of B matrix in units of elements
|
||||
int kAlignmentB,
|
||||
/// Element type for C and D matrix operands
|
||||
typename ElementC_,
|
||||
/// Layout type for C and D matrix operands
|
||||
typename LayoutC_,
|
||||
/// Element type for internal accumulation
|
||||
typename ElementAccumulator,
|
||||
/// Operator class tag
|
||||
typename OperatorClass,
|
||||
/// Tag indicating architecture to tune for
|
||||
typename ArchTag,
|
||||
/// Threadblock-level tile size (concept: GemmShape)
|
||||
typename ThreadblockShape,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename WarpShape,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename InstructionShape,
|
||||
/// Epilogue output operator
|
||||
typename EpilogueOutputOp,
|
||||
/// Threadblock-level swizzling operator
|
||||
typename ThreadblockSwizzle,
|
||||
/// Number of stages used in the pipelined mainloop
|
||||
int Stages,
|
||||
/// If true, kernel is configured to support serial reduction in the
|
||||
/// epilogue
|
||||
bool SplitKSerial,
|
||||
/// Operation performed by GEMM
|
||||
typename Operator>
|
||||
struct DefaultGemm;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Partial specialization for SIMT
|
||||
template <
|
||||
/// Element type for A matrix operand
|
||||
typename ElementA,
|
||||
/// Layout type for A matrix operand
|
||||
typename LayoutA,
|
||||
/// Access granularity of A matrix in units of elements
|
||||
int kAlignmentA,
|
||||
/// Element type for B matrix operand
|
||||
typename ElementB,
|
||||
/// Layout type for B matrix operand
|
||||
typename LayoutB,
|
||||
/// Access granularity of A matrix in units of elements
|
||||
int kAlignmentB,
|
||||
/// Element type for C and D matrix operands
|
||||
typename ElementC,
|
||||
/// Element type for internal accumulation
|
||||
typename ElementAccumulator,
|
||||
/// Tag indicating architecture to tune for
|
||||
typename ArchTag,
|
||||
/// Threadblock-level tile size (concept: GemmShape)
|
||||
typename ThreadblockShape,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename WarpShape,
|
||||
/// Epilogue output operator
|
||||
typename EpilogueOutputOp,
|
||||
/// Threadblock-level swizzling operator
|
||||
typename ThreadblockSwizzle,
|
||||
/// If true, kernel is configured to support serial reduction in the epilogue
|
||||
bool SplitKSerial,
|
||||
/// Operation performed by GEMM
|
||||
typename Operator
|
||||
>
|
||||
struct DefaultGemm<
|
||||
ElementA,
|
||||
LayoutA,
|
||||
kAlignmentA,
|
||||
ElementB,
|
||||
LayoutB,
|
||||
kAlignmentB,
|
||||
ElementC,
|
||||
layout::RowMajor,
|
||||
ElementAccumulator,
|
||||
arch::OpClassSimt,
|
||||
ArchTag,
|
||||
ThreadblockShape,
|
||||
WarpShape,
|
||||
GemmShape<1, 1, 1>,
|
||||
EpilogueOutputOp,
|
||||
ThreadblockSwizzle,
|
||||
2,
|
||||
SplitKSerial,
|
||||
Operator> {
|
||||
/// Define the threadblock-scoped matrix multiply-accumulate
|
||||
using Mma = typename cutlass::gemm::threadblock::DefaultMma<
|
||||
ElementA,
|
||||
LayoutA,
|
||||
kAlignmentA,
|
||||
ElementB,
|
||||
LayoutB,
|
||||
kAlignmentB,
|
||||
ElementAccumulator,
|
||||
layout::RowMajor,
|
||||
arch::OpClassSimt,
|
||||
arch::Sm50,
|
||||
ThreadblockShape,
|
||||
WarpShape,
|
||||
GemmShape<1, 1, 1>,
|
||||
2,
|
||||
Operator>::ThreadblockMma;
|
||||
|
||||
static int const kEpilogueElementsPerAccess = EpilogueOutputOp::kCount;
|
||||
static_assert(kEpilogueElementsPerAccess == 1, "simt epilogue must operate on scalars");
|
||||
|
||||
/// Define the epilogue
|
||||
using Epilogue = typename cutlass::epilogue::threadblock::DefaultEpilogueSimt<
|
||||
ThreadblockShape,
|
||||
typename Mma::Operator,
|
||||
EpilogueOutputOp,
|
||||
kEpilogueElementsPerAccess
|
||||
>::Epilogue;
|
||||
|
||||
/// Define the kernel-level GEMM operator.
|
||||
using GemmKernel = kernel::Gemm<Mma, Epilogue, ThreadblockSwizzle, SplitKSerial>;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
/// Partial specialization for SIMT DP4A
|
||||
|
||||
template <
|
||||
/// Layout type for A matrix operand
|
||||
typename LayoutA,
|
||||
/// Access granularity of A matrix in units of elements
|
||||
int kAlignmentA,
|
||||
/// Layout type for B matrix operand
|
||||
typename LayoutB,
|
||||
/// Access granularity of A matrix in units of elements
|
||||
int kAlignmentB,
|
||||
/// Layout type for C matrix operand
|
||||
typename LayoutC,
|
||||
/// Element type for C and D matrix operands
|
||||
typename ElementC,
|
||||
/// Tag indicating architecture to tune for
|
||||
typename ArchTag,
|
||||
/// Element type for internal accumulation
|
||||
typename ElementAccumulator,
|
||||
/// Threadblock-level tile size (concept: GemmShape)
|
||||
typename ThreadblockShape,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename WarpShape,
|
||||
/// Epilogue output operator
|
||||
typename EpilogueOutputOp,
|
||||
/// Threadblock-level swizzling operator
|
||||
typename ThreadblockSwizzle,
|
||||
/// If true, kernel is configured to support serial reduction in the
|
||||
/// epilogue
|
||||
bool SplitKSerial,
|
||||
/// Operation performed by GEMM
|
||||
typename Operator>
|
||||
struct DefaultGemm<int8_t, LayoutA, kAlignmentA, int8_t, LayoutB, kAlignmentB,
|
||||
ElementC, LayoutC, ElementAccumulator, arch::OpClassSimt,
|
||||
ArchTag, ThreadblockShape, WarpShape, GemmShape<1, 1, 4>,
|
||||
EpilogueOutputOp, ThreadblockSwizzle, 2, SplitKSerial,
|
||||
Operator> {
|
||||
using InstructionShape = GemmShape<1, 1, 4>;
|
||||
using ElementA = int8_t;
|
||||
using ElementB = int8_t;
|
||||
|
||||
using OperatorClass = arch::OpClassSimt;
|
||||
/// Define the threadblock-scoped matrix multiply-accumulate
|
||||
using Mma = typename cutlass::gemm::threadblock::DefaultMma<ElementA,
|
||||
LayoutA,
|
||||
kAlignmentA,
|
||||
ElementB,
|
||||
LayoutB,
|
||||
kAlignmentB,
|
||||
ElementAccumulator,
|
||||
LayoutC,
|
||||
arch::OpClassSimt,
|
||||
arch::Sm50,
|
||||
ThreadblockShape,
|
||||
WarpShape,
|
||||
InstructionShape,
|
||||
2,
|
||||
Operator,
|
||||
false
|
||||
>::ThreadblockMma;
|
||||
|
||||
static int const kEpilogueElementsPerAccess = EpilogueOutputOp::kCount;
|
||||
static_assert(kEpilogueElementsPerAccess == 1, "simt epilogue must operate on scalars");
|
||||
|
||||
/// Define the epilogue
|
||||
using Epilogue = typename cutlass::epilogue::threadblock::DefaultEpilogueSimt<
|
||||
ThreadblockShape,
|
||||
typename Mma::Operator,
|
||||
EpilogueOutputOp,
|
||||
kEpilogueElementsPerAccess
|
||||
>::Epilogue;
|
||||
|
||||
/// Define the kernel-level GEMM operator.
|
||||
using GemmKernel = kernel::Gemm<Mma, Epilogue, ThreadblockSwizzle, SplitKSerial>;
|
||||
};
|
||||
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
/// Partial specialization for BigIsland 1.0 tensor op architecture
|
||||
template <
|
||||
/// Element type for A matrix operand
|
||||
typename ElementA,
|
||||
/// Layout type for A matrix operand
|
||||
typename LayoutA,
|
||||
/// Access granularity of A matrix in units of elements
|
||||
int kAlignmentA,
|
||||
/// Element type for B matrix operand
|
||||
typename ElementB,
|
||||
/// Layout type for B matrix operand
|
||||
typename LayoutB,
|
||||
/// Access granularity of B matrix in units of elements
|
||||
int kAlignmentB,
|
||||
/// Element type for C and D matrix operands
|
||||
typename ElementC,
|
||||
/// Element type for internal accumulation
|
||||
typename ElementAccumulator,
|
||||
/// Tag indicating architecture to tune for
|
||||
typename ArchTag,
|
||||
/// Threadblock-level tile size (concept: GemmShape)
|
||||
typename ThreadblockShape,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename WarpShape,
|
||||
/// Instrcution shape
|
||||
typename InstructionShape,
|
||||
/// Epilogue output operator
|
||||
typename EpilogueOutputOp,
|
||||
/// Threadblock-level swizzling operator
|
||||
typename ThreadblockSwizzle,
|
||||
/// Number of stages used in the pipelined mainloop
|
||||
int Stages,
|
||||
/// If true, kernel is configured to support serial reduction in the epilogue
|
||||
bool SplitKSerial,
|
||||
/// Operation performed by GEMM
|
||||
typename Operator
|
||||
>
|
||||
struct DefaultGemm<
|
||||
ElementA, LayoutA, kAlignmentA,
|
||||
ElementB, LayoutB, kAlignmentB,
|
||||
ElementC, layout::RowMajor,
|
||||
ElementAccumulator,
|
||||
arch::OpClassTensorOp,
|
||||
ArchTag,
|
||||
ThreadblockShape,
|
||||
WarpShape,
|
||||
InstructionShape,
|
||||
EpilogueOutputOp,
|
||||
ThreadblockSwizzle,
|
||||
Stages,
|
||||
SplitKSerial,
|
||||
Operator
|
||||
> {
|
||||
|
||||
/// Define the threadblock-scoped matrix multiply-accumulate
|
||||
using Mma = typename cutlass::gemm::threadblock::DefaultMma<
|
||||
ElementA,
|
||||
LayoutA,
|
||||
kAlignmentA,
|
||||
ElementB,
|
||||
LayoutB,
|
||||
kAlignmentB,
|
||||
ElementAccumulator,
|
||||
layout::RowMajor,
|
||||
arch::OpClassTensorOp,
|
||||
ArchTag,
|
||||
ThreadblockShape,
|
||||
WarpShape,
|
||||
InstructionShape,
|
||||
Stages,
|
||||
Operator
|
||||
>::ThreadblockMma;
|
||||
|
||||
static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK;
|
||||
|
||||
/// FIXME(Peter Han): Probably DefaultEpiloguesTensorOp should be used here, let's see
|
||||
static const int kEpilougeElementsPerAccess = EpilogueOutputOp::kCount;
|
||||
|
||||
/// Define the epilogue
|
||||
using Epilogue = typename cutlass::epilogue::threadblock::DefaultEpilogueTensorOp<
|
||||
ThreadblockShape,
|
||||
typename Mma::Operator,
|
||||
EpilogueOutputOp,
|
||||
kEpilougeElementsPerAccess
|
||||
>::Epilogue;
|
||||
|
||||
/// Define the kernel-level GEMM operator.
|
||||
using GemmKernel = kernel::Gemm<Mma, Epilogue, ThreadblockSwizzle, SplitKSerial>;
|
||||
};
|
||||
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace kernel
|
||||
} // namespace gemm
|
||||
} // namespace cutlass
|
||||
292
cat_files/default_gemm_configuration.h
Normal file
292
cat_files/default_gemm_configuration.h
Normal file
@@ -0,0 +1,292 @@
|
||||
/***************************************************************************************************
|
||||
* 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:
|
||||
* * 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.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2021 Iluvatar CoreX. All rights reserved.
|
||||
* Copyright Declaration: This software, including all of its code and documentation,
|
||||
* except for the third-party software it contains, is a copyrighted work of Shanghai Iluvatar CoreX
|
||||
* Semiconductor Co., Ltd. and its affiliates ("Iluvatar CoreX") in accordance with the PRC Copyright
|
||||
* Law and relevant international treaties, and all rights contained therein are enjoyed by Iluvatar
|
||||
* CoreX. No user of this software shall have any right, ownership or interest in this software and
|
||||
* any use of this software shall be in compliance with the terms and conditions of the End User
|
||||
* License Agreement.
|
||||
**************************************************************************************************/
|
||||
|
||||
/*! \file
|
||||
\brief Definitions for GEMM structures
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/arch/arch.h"
|
||||
#include "cutlass/arch/mma.h"
|
||||
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
#include "cutlass/epilogue/thread/linear_combination.h"
|
||||
#include "cutlass/epilogue/thread/linear_combination_clamp.h"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace gemm {
|
||||
namespace device {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
typename OperatorClass,
|
||||
typename ArchTag,
|
||||
typename ElementA,
|
||||
typename ElementB,
|
||||
typename ElementC,
|
||||
typename ElementAccumulator
|
||||
>
|
||||
struct DefaultGemmConfiguration;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// FIXME(Peter Han): Need to update configuration according to perf results, so
|
||||
/// that could archieve good performance by default.
|
||||
|
||||
template <
|
||||
typename ArchTag,
|
||||
typename ElementA,
|
||||
typename ElementB,
|
||||
typename ElementC,
|
||||
typename ElementAccumulator>
|
||||
struct DefaultGemmConfiguration<
|
||||
arch::OpClassSimt,
|
||||
ArchTag,
|
||||
ElementA,
|
||||
ElementB,
|
||||
ElementC,
|
||||
ElementAccumulator> {
|
||||
|
||||
static int const kAlignmentA = 1;
|
||||
static int const kAlignmentB = 1;
|
||||
using ThreadblockShape = GemmShape<128, 128, 8>;
|
||||
using WarpShape = GemmShape<64, 64, 8>;
|
||||
using InstructionShape = GemmShape<1, 1, 1>;
|
||||
static int const kStages = 2;
|
||||
|
||||
using EpilogueOutputOp = epilogue::thread::LinearCombination<
|
||||
ElementC,
|
||||
1,
|
||||
ElementAccumulator,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
using Operator = arch::OpMultiplyAdd;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
typename ArchTag,
|
||||
typename ElementC>
|
||||
struct DefaultGemmConfiguration<arch::OpClassSimt, ArchTag, int8_t, int8_t, ElementC, int32_t> {
|
||||
|
||||
static int const kAlignmentA = 4;
|
||||
static int const kAlignmentB = 4;
|
||||
using ThreadblockShape = GemmShape<128, 128, 32>;
|
||||
using WarpShape = GemmShape<64, 64, 32>;
|
||||
using InstructionShape = GemmShape<1, 1, 4>;
|
||||
static int const kStages = 2;
|
||||
|
||||
using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp<
|
||||
ElementC,
|
||||
1,
|
||||
int32_t,
|
||||
float
|
||||
>;
|
||||
|
||||
using Operator = arch::OpMultiplyAdd;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
typename ElementC>
|
||||
struct DefaultGemmConfiguration<
|
||||
arch::OpClassTensorOp,
|
||||
arch::Cu10,
|
||||
int8_t,
|
||||
int8_t,
|
||||
ElementC,
|
||||
int32_t> {
|
||||
|
||||
using ElementA = int8_t;
|
||||
using ElementB = int8_t;
|
||||
using ElementAccumulator = int32_t;
|
||||
static int const kAlignmentA = MEMORY_ACCESS_SIZE / sizeof_bits<ElementA>::value;
|
||||
static int const kAlignmentB = MEMORY_ACCESS_SIZE / sizeof_bits<ElementB>::value;
|
||||
|
||||
using ThreadblockShape = GemmShape<256, 256, 32>;
|
||||
using WarpShape = GemmShape<64, 64, 32>;
|
||||
using InstructionShape = GemmShape<16, 16, 16>;
|
||||
static int const kStages = 2;
|
||||
|
||||
using EpilogueOutputOp = epilogue::thread::LinearCombination<
|
||||
ElementC,
|
||||
MEMORY_ACCESS_SIZE / sizeof_bits<ElementC>::value,
|
||||
ElementAccumulator,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
using Operator = arch::OpMultiplyAdd;
|
||||
};
|
||||
|
||||
template <
|
||||
typename ElementC>
|
||||
struct DefaultGemmConfiguration<
|
||||
arch::OpClassTensorOp,
|
||||
arch::Cu10,
|
||||
uint8_t,
|
||||
uint8_t,
|
||||
ElementC,
|
||||
uint32_t> {
|
||||
|
||||
using ElementA = uint8_t;
|
||||
using ElementB = uint8_t;
|
||||
using ElementAccumulator = uint32_t;
|
||||
static int const kAlignmentA = MEMORY_ACCESS_SIZE / sizeof_bits<ElementA>::value;
|
||||
static int const kAlignmentB = MEMORY_ACCESS_SIZE / sizeof_bits<ElementB>::value;
|
||||
|
||||
using ThreadblockShape = GemmShape<256, 256, 32>;
|
||||
using WarpShape = GemmShape<64, 64, 32>;
|
||||
using InstructionShape = GemmShape<16, 16, 16>;
|
||||
static int const kStages = 2;
|
||||
|
||||
using EpilogueOutputOp = epilogue::thread::LinearCombination<
|
||||
ElementC,
|
||||
MEMORY_ACCESS_SIZE / sizeof_bits<ElementC>::value,
|
||||
ElementAccumulator,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
using Operator = arch::OpMultiplyAdd;
|
||||
};
|
||||
|
||||
template <
|
||||
typename ElementC>
|
||||
struct DefaultGemmConfiguration<
|
||||
arch::OpClassTensorOp,
|
||||
arch::Cu10,
|
||||
half_t,
|
||||
half_t,
|
||||
ElementC,
|
||||
float> {
|
||||
|
||||
using ElementA = half_t;
|
||||
using ElementB = half_t;
|
||||
using ElementAccumulator = float;
|
||||
static int const kAlignmentA = MEMORY_ACCESS_SIZE / sizeof_bits<ElementA>::value;
|
||||
static int const kAlignmentB = MEMORY_ACCESS_SIZE / sizeof_bits<ElementB>::value;
|
||||
|
||||
using ThreadblockShape = GemmShape<128, 128, 32>;
|
||||
using WarpShape = GemmShape<32, 32, 32>;
|
||||
using InstructionShape = GemmShape<16, 16, 16>;
|
||||
static int const kStages = 2;
|
||||
|
||||
using EpilogueOutputOp = epilogue::thread::LinearCombination<
|
||||
ElementC,
|
||||
MEMORY_ACCESS_SIZE / sizeof_bits<ElementC>::value,
|
||||
ElementAccumulator,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
using Operator = arch::OpMultiplyAdd;
|
||||
};
|
||||
|
||||
template <
|
||||
typename ElementC>
|
||||
struct DefaultGemmConfiguration<
|
||||
arch::OpClassTensorOp,
|
||||
arch::Cu10,
|
||||
bfloat16_t,
|
||||
bfloat16_t,
|
||||
ElementC,
|
||||
float> {
|
||||
|
||||
using ElementA = bfloat16_t;
|
||||
using ElementB = bfloat16_t;
|
||||
using ElementAccumulator = float;
|
||||
static int const kAlignmentA = 32 / sizeof_bits<ElementA>::value;
|
||||
static int const kAlignmentB = 32 / sizeof_bits<ElementB>::value;
|
||||
|
||||
using ThreadblockShape = GemmShape<128, 128, 32>;
|
||||
using WarpShape = GemmShape<32, 32, 32>;
|
||||
using InstructionShape = GemmShape<16, 16, 16>;
|
||||
static int const kStages = 2;
|
||||
|
||||
using EpilogueOutputOp = epilogue::thread::LinearCombination<
|
||||
ElementC,
|
||||
MEMORY_ACCESS_SIZE / sizeof_bits<ElementC>::value,
|
||||
ElementAccumulator,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
using Operator = arch::OpMultiplyAdd;
|
||||
};
|
||||
|
||||
template <
|
||||
typename ElementC>
|
||||
struct DefaultGemmConfiguration<
|
||||
arch::OpClassTensorOp,
|
||||
arch::Cu10,
|
||||
float,
|
||||
float,
|
||||
ElementC,
|
||||
float> {
|
||||
|
||||
using ElementA = float;
|
||||
using ElementB = float;
|
||||
using ElementAccumulator = float;
|
||||
static int const kAlignmentA = 32 / sizeof_bits<ElementA>::value;
|
||||
static int const kAlignmentB = 32 / sizeof_bits<ElementB>::value;
|
||||
|
||||
using ThreadblockShape = GemmShape<128, 128, 32>;
|
||||
using WarpShape = GemmShape<32, 32, 32>;
|
||||
using InstructionShape = GemmShape<16, 16, 16>;
|
||||
static int const kStages = 2;
|
||||
|
||||
using EpilogueOutputOp = epilogue::thread::LinearCombination<
|
||||
ElementC,
|
||||
MEMORY_ACCESS_SIZE / sizeof_bits<ElementC>::value,
|
||||
ElementAccumulator,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
using Operator = arch::OpMultiplyAdd;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
} // namespace device
|
||||
} // namespace gemm
|
||||
} // namespace cutlass
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
307
cat_files/default_gemm_universal.h
Normal file
307
cat_files/default_gemm_universal.h
Normal file
@@ -0,0 +1,307 @@
|
||||
/***************************************************************************************************
|
||||
* 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:
|
||||
* * 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
|
||||
Default kernel-level GEMM definitions combine threadblock-scoped matrix multiply-add with
|
||||
the appropriate threadblock-scoped epilogue.
|
||||
|
||||
Note, CUTLASS epilogues universally target row-major outputs. Column-major outputs are
|
||||
accommodated by exchanging A and B operands and assuming transposed layouts. Partial
|
||||
specializations here choose 'device::GemmTransposed' to implement this functionality.
|
||||
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
|
||||
#include "cutlass/complex.h"
|
||||
#include "cutlass/layout/matrix.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
|
||||
#include "cutlass/gemm/kernel/gemm_universal.h"
|
||||
#include "cutlass/gemm/kernel/default_gemm.h"
|
||||
#include "cutlass/gemm/kernel/default_gemm_complex.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace gemm {
|
||||
namespace kernel {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
/// Element type for A matrix operand
|
||||
typename ElementA_,
|
||||
/// Layout type for A matrix operand
|
||||
typename LayoutA_,
|
||||
/// Complex elementwise transformation on A operand
|
||||
ComplexTransform TransformA,
|
||||
/// Access granularity of A matrix in units of elements
|
||||
int kAlignmentA,
|
||||
/// Element type for B matrix operand
|
||||
typename ElementB_,
|
||||
/// Layout type for B matrix operand
|
||||
typename LayoutB_,
|
||||
/// Complex elementwise transformation on B operand
|
||||
ComplexTransform TransformB,
|
||||
/// Access granularity of B matrix in units of elements
|
||||
int kAlignmentB,
|
||||
/// Element type for C and D matrix operands
|
||||
typename ElementC_,
|
||||
/// Layout type for C and D matrix operands
|
||||
typename LayoutC_,
|
||||
/// Element type for internal accumulation
|
||||
typename ElementAccumulator,
|
||||
/// Operator class tag
|
||||
typename OperatorClass,
|
||||
/// Tag indicating architecture to tune for
|
||||
typename ArchTag,
|
||||
/// Threadblock-level tile size (concept: GemmShape)
|
||||
typename ThreadblockShape,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename WarpShape,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename InstructionShape,
|
||||
/// Epilogue output operator
|
||||
typename EpilogueOutputOp,
|
||||
/// Threadblock-level swizzling operator
|
||||
typename ThreadblockSwizzle,
|
||||
/// Number of stages used in the pipelined mainloop
|
||||
int Stages,
|
||||
/// Operation performed by GEMM
|
||||
typename Operator,
|
||||
///
|
||||
typename Enable = void
|
||||
>
|
||||
struct DefaultGemmUniversal;
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
//
|
||||
// Real-valued GEMM kernels
|
||||
//
|
||||
|
||||
template <
|
||||
/// Element type for A matrix operand
|
||||
typename ElementA,
|
||||
/// Layout type for A matrix operand
|
||||
typename LayoutA,
|
||||
/// Access granularity of A matrix in units of elements
|
||||
int kAlignmentA,
|
||||
/// Element type for B matrix operand
|
||||
typename ElementB,
|
||||
/// Layout type for B matrix operand
|
||||
typename LayoutB,
|
||||
/// Access granularity of B matrix in units of elements
|
||||
int kAlignmentB,
|
||||
/// Element type for C and D matrix operands
|
||||
typename ElementC,
|
||||
/// Layout type for C and D matrix operands
|
||||
typename LayoutC,
|
||||
/// Element type for internal accumulation
|
||||
typename ElementAccumulator,
|
||||
/// Operator class tag
|
||||
typename OperatorClass,
|
||||
/// Tag indicating architecture to tune for
|
||||
typename ArchTag,
|
||||
/// Threadblock-level tile size (concept: GemmShape)
|
||||
typename ThreadblockShape,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename WarpShape,
|
||||
/// Instruction-level tile size (concept: GemmShape)
|
||||
typename InstructionShape,
|
||||
/// Epilogue output operator
|
||||
typename EpilogueOutputOp,
|
||||
/// Threadblock-level swizzling operator
|
||||
typename ThreadblockSwizzle,
|
||||
/// Number of stages used in the pipelined mainloop
|
||||
int Stages,
|
||||
/// Operation performed by GEMM
|
||||
typename Operator>
|
||||
struct DefaultGemmUniversal<
|
||||
ElementA,
|
||||
LayoutA,
|
||||
ComplexTransform::kNone, // transform A
|
||||
kAlignmentA,
|
||||
ElementB,
|
||||
LayoutB,
|
||||
ComplexTransform::kNone, // transform B
|
||||
kAlignmentB,
|
||||
ElementC,
|
||||
LayoutC,
|
||||
ElementAccumulator,
|
||||
OperatorClass,
|
||||
ArchTag,
|
||||
ThreadblockShape,
|
||||
WarpShape,
|
||||
InstructionShape,
|
||||
EpilogueOutputOp,
|
||||
ThreadblockSwizzle,
|
||||
Stages,
|
||||
Operator,
|
||||
typename std::enable_if< ! cutlass::is_complex<ElementAccumulator>::value>::type
|
||||
> {
|
||||
|
||||
using DefaultGemmKernel = typename kernel::DefaultGemm<
|
||||
ElementA,
|
||||
LayoutA,
|
||||
kAlignmentA,
|
||||
ElementB,
|
||||
LayoutB,
|
||||
kAlignmentB,
|
||||
ElementC,
|
||||
LayoutC,
|
||||
ElementAccumulator,
|
||||
OperatorClass,
|
||||
ArchTag,
|
||||
ThreadblockShape,
|
||||
WarpShape,
|
||||
InstructionShape,
|
||||
EpilogueOutputOp,
|
||||
ThreadblockSwizzle,
|
||||
Stages,
|
||||
true,
|
||||
Operator
|
||||
>::GemmKernel;
|
||||
|
||||
/// Define the kernel in terms of the default kernel
|
||||
using GemmKernel = kernel::GemmUniversal<
|
||||
typename DefaultGemmKernel::Mma,
|
||||
typename DefaultGemmKernel::Epilogue,
|
||||
ThreadblockSwizzle
|
||||
>;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
//
|
||||
// Complex-valued GEMM kernels
|
||||
//
|
||||
|
||||
template <
|
||||
/// Element type for A matrix operand
|
||||
typename ElementA,
|
||||
/// Layout type for A matrix operand
|
||||
typename LayoutA,
|
||||
/// Complex elementwise transformation on A operand
|
||||
ComplexTransform TransformA,
|
||||
/// Access granularity of A matrix in units of elements
|
||||
int kAlignmentA,
|
||||
/// Element type for B matrix operand
|
||||
typename ElementB,
|
||||
/// Layout type for B matrix operand
|
||||
typename LayoutB,
|
||||
/// Complex elementwise transformation on B operand
|
||||
ComplexTransform TransformB,
|
||||
/// Access granularity of B matrix in units of elements
|
||||
int kAlignmentB,
|
||||
/// Element type for C and D matrix operands
|
||||
typename ElementC,
|
||||
/// Layout type for C and D matrix operands
|
||||
typename LayoutC,
|
||||
/// Element type for internal accumulation
|
||||
typename ElementAccumulator,
|
||||
/// Operator class tag
|
||||
typename OperatorClass,
|
||||
/// Tag indicating architecture to tune for
|
||||
typename ArchTag,
|
||||
/// Threadblock-level tile size (concept: GemmShape)
|
||||
typename ThreadblockShape,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename WarpShape,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename InstructionShape,
|
||||
/// Epilogue output operator
|
||||
typename EpilogueOutputOp,
|
||||
/// Threadblock-level swizzling operator
|
||||
typename ThreadblockSwizzle,
|
||||
/// Number of stages used in the pipelined mainloop
|
||||
int Stages,
|
||||
/// Operation performed by GEMM
|
||||
typename Operator
|
||||
>
|
||||
struct DefaultGemmUniversal<
|
||||
ElementA,
|
||||
LayoutA,
|
||||
TransformA,
|
||||
kAlignmentA,
|
||||
ElementB,
|
||||
LayoutB,
|
||||
TransformB,
|
||||
kAlignmentB,
|
||||
ElementC,
|
||||
LayoutC,
|
||||
ElementAccumulator,
|
||||
OperatorClass,
|
||||
ArchTag,
|
||||
ThreadblockShape,
|
||||
WarpShape,
|
||||
InstructionShape,
|
||||
EpilogueOutputOp,
|
||||
ThreadblockSwizzle,
|
||||
Stages,
|
||||
Operator,
|
||||
typename std::enable_if<cutlass::is_complex<ElementAccumulator>::value>::type
|
||||
> {
|
||||
|
||||
using DefaultGemmKernel = typename kernel::DefaultGemmComplex<
|
||||
ElementA,
|
||||
LayoutA,
|
||||
ElementB,
|
||||
LayoutB,
|
||||
ElementC,
|
||||
LayoutC,
|
||||
ElementAccumulator,
|
||||
OperatorClass,
|
||||
ArchTag,
|
||||
ThreadblockShape,
|
||||
WarpShape,
|
||||
InstructionShape,
|
||||
EpilogueOutputOp,
|
||||
ThreadblockSwizzle,
|
||||
Stages,
|
||||
TransformA,
|
||||
TransformB,
|
||||
Operator,
|
||||
false
|
||||
>::GemmKernel;
|
||||
|
||||
/// Define the kernel in terms of the default kernel
|
||||
using GemmKernel = kernel::GemmUniversal<
|
||||
typename DefaultGemmKernel::Mma,
|
||||
typename DefaultGemmKernel::Epilogue,
|
||||
ThreadblockSwizzle
|
||||
>;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace kernel
|
||||
} // namespace gemm
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
114
cat_files/default_mma_core.h
Normal file
114
cat_files/default_mma_core.h
Normal file
@@ -0,0 +1,114 @@
|
||||
/***************************************************************************************************
|
||||
* 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:
|
||||
* * 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.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2021 Iluvatar CoreX. All rights reserved.
|
||||
* Copyright Declaration: This software, including all of its code and documentation,
|
||||
* except for the third-party software it contains, is a copyrighted work of Shanghai Iluvatar CoreX
|
||||
* Semiconductor Co., Ltd. and its affiliates ("Iluvatar CoreX") in accordance with the PRC Copyright
|
||||
* Law and relevant international treaties, and all rights contained therein are enjoyed by Iluvatar
|
||||
* CoreX. No user of this software shall have any right, ownership or interest in this software and
|
||||
* any use of this software shall be in compliance with the terms and conditions of the End User
|
||||
* License Agreement.
|
||||
**************************************************************************************************/
|
||||
|
||||
/*! \file
|
||||
\brief Defines basic properties needed by CTA-level GEMMs assuming expectations about data
|
||||
layout of the global memory fragments, data types, and internal tile sizes.
|
||||
|
||||
Partial specializations for threadblock::Mma operations targeting TensorOp instructions.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/array.h"
|
||||
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/matrix_shape.h"
|
||||
|
||||
#include "cutlass/gemm/warp/mma.h"
|
||||
#include "cutlass/gemm/threadblock/mma_pipelined.h"
|
||||
#include "cutlass/gemm/threadblock/mma_singlestage.h"
|
||||
#include "cutlass/gemm/threadblock/mma_preload.h"
|
||||
#include "cutlass/arch/cache_operation.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace gemm {
|
||||
namespace threadblock {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Template defininng default matrix multiply operators inferred from threadblock tile size,
|
||||
/// global memory data layout, and target math instruction.
|
||||
template <
|
||||
/// Shape of threadblock-scoped matrix multiply operator
|
||||
typename Shape,
|
||||
/// Shape of warp-level matrix multiply operator
|
||||
typename WarpShape,
|
||||
/// Shape of one matrix production operation (concept: GemmShape)
|
||||
typename InstructionShape,
|
||||
/// Element data type of A operand
|
||||
typename ElementA,
|
||||
/// Layout of operand A
|
||||
typename LayoutA,
|
||||
/// Element data type of B operand
|
||||
typename ElementB,
|
||||
/// Layout of operand B
|
||||
typename LayoutB,
|
||||
/// Data type of accumulator
|
||||
typename ElementC,
|
||||
/// Layout of accumulator
|
||||
typename LayoutC,
|
||||
/// Indicates type of math operator (arch::OpClassSimt or arch::OpClassTensorOp)
|
||||
typename OperatorClass,
|
||||
/// Number of stages
|
||||
int Stages = 2,
|
||||
/// Operation performed by MMA
|
||||
typename Operator = cutlass::arch::OpMultiplyAdd,
|
||||
/// Store the accumulators in row major or column major. Row major is used
|
||||
/// when output layout is interleaved.
|
||||
bool AccumulatorsInRowMajor = false,
|
||||
/// Cache operation of operand A
|
||||
cutlass::arch::CacheOperation::Kind CacheOpA =
|
||||
cutlass::arch::CacheOperation::Global,
|
||||
/// Cache operation of operand B
|
||||
cutlass::arch::CacheOperation::Kind CacheOpB =
|
||||
cutlass::arch::CacheOperation::Global,
|
||||
/// per-element transformation for elements of A
|
||||
ComplexTransform TransformA = ComplexTransform::kNone,
|
||||
/// per-element transformation for elements of B
|
||||
ComplexTransform TransformB = ComplexTransform::kNone,
|
||||
bool IsComplex = false // (is_complex<ElementA>::value || is_complex<ElementB>::value)
|
||||
>
|
||||
struct DefaultMmaCore;
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace gemm
|
||||
} // namespace cutlass
|
||||
835
cat_files/default_mma_core_cu10.h
Normal file
835
cat_files/default_mma_core_cu10.h
Normal file
@@ -0,0 +1,835 @@
|
||||
/***************************************************************************************************
|
||||
* 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:
|
||||
* * 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.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2021 Iluvatar CoreX. All rights reserved.
|
||||
* Copyright Declaration: This software, including all of its code and documentation,
|
||||
* except for the third-party software it contains, is a copyrighted work of Shanghai Iluvatar CoreX
|
||||
* Semiconductor Co., Ltd. and its affiliates ("Iluvatar CoreX") in accordance with the PRC Copyright
|
||||
* Law and relevant international treaties, and all rights contained therein are enjoyed by Iluvatar
|
||||
* CoreX. No user of this software shall have any right, ownership or interest in this software and
|
||||
* any use of this software shall be in compliance with the terms and conditions of the End User
|
||||
* License Agreement.
|
||||
**************************************************************************************************/
|
||||
|
||||
/*! \file
|
||||
\brief Defines basic properties needed by CTA-level GEMMs assuming expectations about data
|
||||
layout of the global memory fragments, data types, and internal tile sizes.
|
||||
|
||||
Partial specializations for threadblock::Mma operations targeting TensorOp instructions.
|
||||
|
||||
Aims at TensorOp of the first generation BigIsland.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/array.h"
|
||||
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/matrix_shape.h"
|
||||
|
||||
#include "cutlass/transform/pitch_linear_thread_map.h"
|
||||
#include "cutlass/transform/threadblock/regular_tile_access_iterator_tensor_op.h"
|
||||
#include "cutlass/transform/threadblock/regular_tile_iterator_tensor_op.h"
|
||||
#include "cutlass/layout/tensor_op_multiplicand.h"
|
||||
#include "cutlass/layout/tensor_op_em.h"
|
||||
|
||||
#include "cutlass/gemm/warp/mma_tensor_op_policy.h"
|
||||
#include "cutlass/gemm/warp/mma_tensor_op.h"
|
||||
#include "cutlass/gemm/warp/default_mma_tensor_op.h"
|
||||
#include "cutlass/gemm/threadblock/default_mma_core.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace gemm {
|
||||
namespace threadblock {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
///
|
||||
/// Specialization: A: row-major, B: row-major, TT
|
||||
///
|
||||
/// This uses the default warp-level operator given tile sizes
|
||||
///
|
||||
template <
|
||||
/// Shape of threadblock-scoped matrix multiply operator (concept:
|
||||
/// GemmShape)
|
||||
typename Shape_,
|
||||
/// Shape of warp-level matrix multiply operator (concept: GemmShape)
|
||||
typename WarpShape_,
|
||||
/// Data type of A operand
|
||||
typename ElementA_,
|
||||
/// Data type of B operand
|
||||
typename ElementB_,
|
||||
/// Data type of accumulator
|
||||
typename ElementC_,
|
||||
/// Layout of accumulator
|
||||
typename LayoutC_,
|
||||
/// Stages
|
||||
int Stages,
|
||||
/// Operation performed by GEMM
|
||||
typename Operator_>
|
||||
struct DefaultMmaCore<Shape_,
|
||||
WarpShape_,
|
||||
GemmShape<16, 16, 16>,
|
||||
ElementA_,
|
||||
layout::RowMajor,
|
||||
ElementB_,
|
||||
layout::RowMajor,
|
||||
ElementC_,
|
||||
LayoutC_,
|
||||
arch::OpClassTensorOp,
|
||||
Stages,
|
||||
Operator_> {
|
||||
using Shape = Shape_;
|
||||
using WarpShape = WarpShape_;
|
||||
using InstructionShape = GemmShape<16, 16, 16>;
|
||||
using ElementA = ElementA_;
|
||||
using LayoutA = layout::RowMajor;
|
||||
using ElementB = ElementB_;
|
||||
using LayoutB = layout::RowMajor;
|
||||
using ElementC = ElementC_;
|
||||
using LayoutC = LayoutC_;
|
||||
using OperatorClass = arch::OpClassTensorOp;
|
||||
|
||||
static int const kStages = Stages;
|
||||
|
||||
/// Default Operator
|
||||
using Operator = Operator_;
|
||||
|
||||
/// Warp thread arrangement
|
||||
using WarpThreadArrangement = layout::PitchLinearShape<16, 4>;
|
||||
|
||||
/// Number of warps present
|
||||
using WarpCount = GemmShape<
|
||||
Shape::kM / WarpShape::kM,
|
||||
Shape::kN / WarpShape::kN,
|
||||
Shape::kK / WarpShape::kK
|
||||
>;
|
||||
|
||||
/// Don't support split K within CTA
|
||||
static_assert(Shape::kK == WarpShape::kK,
|
||||
"Threadblock-scoped GEMM shape K should equal warp-scoped GEMM shape K"
|
||||
);
|
||||
|
||||
// Divisibility requirements
|
||||
static_assert(
|
||||
!(Shape::kM % WarpShape::kM) &&
|
||||
!(Shape::kN % WarpShape::kN) &&
|
||||
!(Shape::kK % WarpShape::kK),
|
||||
"Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."
|
||||
);
|
||||
|
||||
// Divisibility requirements
|
||||
static_assert(
|
||||
!(WarpShape::kM % 16) &&
|
||||
!(WarpShape::kN % 16) &&
|
||||
!(WarpShape::kK % 16),
|
||||
"Threadblock-scoped GEMM should be divisible by 16."
|
||||
);
|
||||
|
||||
/// Number of threads per warp
|
||||
static int const kWarpSize = warp::WarpSize<arch::OpClassTensorOp>::value;
|
||||
|
||||
/// Number of threads total
|
||||
static int const kThreads = WarpCount::kCount * kWarpSize;
|
||||
|
||||
/// Size of a threadblock-scoped access
|
||||
static int const kAccessSizeInBits = 32;
|
||||
|
||||
/// Number of A elemnts per access
|
||||
static int const kElementsPerAccessA = kAccessSizeInBits / sizeof_bits<ElementA>::value;
|
||||
|
||||
/// Number of A elemnts per access
|
||||
static int const kElementsPerAccessB = kAccessSizeInBits / sizeof_bits<ElementB>::value;
|
||||
|
||||
//
|
||||
// Shared memory layouts
|
||||
//
|
||||
|
||||
#if BLOCK_LOAD_STORE
|
||||
using SmemLayoutA = layout::TensorOpEm<sizeof_bits<ElementA>::value, LayoutA>;
|
||||
using SmemLayoutB = layout::TensorOpEm<sizeof_bits<ElementB>::value, LayoutB>;
|
||||
#else
|
||||
using SmemLayoutA = layout::TensorOpMultiplicand<sizeof_bits<ElementA>::value, LayoutA>;
|
||||
using SmemLayoutB = layout::TensorOpMultiplicand<sizeof_bits<ElementB>::value, LayoutB>;
|
||||
#endif
|
||||
|
||||
//
|
||||
// Iterators to write to shared memory
|
||||
//
|
||||
|
||||
/// ThreadMap of iterator A
|
||||
///
|
||||
using IteratorThreadMapA = transform::PitchLinear2DThreadTileWarpRakedThreadMap<
|
||||
layout::PitchLinearShape<Shape::kK, Shape::kM>,
|
||||
kThreads,
|
||||
WarpThreadArrangement,
|
||||
layout::PitchLinearShape<kElementsPerAccessA, kElementsPerAccessA>
|
||||
>;
|
||||
|
||||
/// Shared memory iterator to A operand
|
||||
using SmemIteratorA = transform::threadblock::RegularTileIterator<
|
||||
MatrixShape<Shape::kM, Shape::kK>,
|
||||
ElementA,
|
||||
SmemLayoutA,
|
||||
1,
|
||||
IteratorThreadMapA
|
||||
>;
|
||||
|
||||
/// Policy of iterator B
|
||||
using IteratorThreadMapB = transform::PitchLinear2DThreadTileWarpRakedThreadMap<
|
||||
layout::PitchLinearShape<Shape::kN, Shape::kK>,
|
||||
kThreads,
|
||||
WarpThreadArrangement,
|
||||
layout::PitchLinearShape<kElementsPerAccessB, kElementsPerAccessB>
|
||||
>;
|
||||
|
||||
/// Shared memory iterator to B operand
|
||||
using SmemIteratorB = transform::threadblock::RegularTileIterator<
|
||||
MatrixShape<Shape::kK, Shape::kN>,
|
||||
ElementB,
|
||||
SmemLayoutB,
|
||||
0,
|
||||
IteratorThreadMapB
|
||||
>;
|
||||
|
||||
//
|
||||
// Warp-level matrix multiply operator
|
||||
//
|
||||
|
||||
// Define the warp-level tensor op
|
||||
using Policy = gemm::warp::MmaTensorOpPolicy<
|
||||
arch::Mma<
|
||||
gemm::GemmShape<16, 16, 16>,
|
||||
NUM_THREADS_PER_WARP,
|
||||
ElementA,
|
||||
LayoutA,
|
||||
ElementB,
|
||||
LayoutB,
|
||||
ElementC,
|
||||
layout::RowMajor,
|
||||
arch::OpMultiplyAdd
|
||||
>,
|
||||
MatrixShape<1, 1>
|
||||
>;
|
||||
|
||||
using MmaTensorOp = typename gemm::warp::DefaultMmaTensorOp<
|
||||
WarpShape,
|
||||
gemm::GemmShape<16, 16, 16>,
|
||||
ElementA,
|
||||
SmemLayoutA,
|
||||
ElementB,
|
||||
SmemLayoutB,
|
||||
ElementC,
|
||||
LayoutC,
|
||||
arch::OpMultiplyAdd
|
||||
>::Type;
|
||||
|
||||
/// Policy used to define MmaPipelined
|
||||
using MmaPolicy = MmaPolicy<
|
||||
MmaTensorOp,
|
||||
MatrixShape<0, 0>,
|
||||
MatrixShape<0, 0>,
|
||||
WarpCount::kK
|
||||
>;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
///
|
||||
/// Specialization: A: row-major, B: column-major, TN
|
||||
///
|
||||
/// This uses the default warp-level operator given tile sizes
|
||||
///
|
||||
template <
|
||||
/// Shape of threadblock-scoped matrix multiply operator (concept:
|
||||
/// GemmShape)
|
||||
typename Shape_,
|
||||
/// Shape of warp-level matrix multiply operator (concept: GemmShape)
|
||||
typename WarpShape_,
|
||||
/// Data type of A operand
|
||||
typename ElementA_,
|
||||
/// Data type of B operand
|
||||
typename ElementB_,
|
||||
/// Data type of accumulator
|
||||
typename ElementC_,
|
||||
/// Layout of accumulator
|
||||
typename LayoutC_,
|
||||
/// Stages
|
||||
int Stages,
|
||||
/// Operation performed by GEMM
|
||||
typename Operator_>
|
||||
struct DefaultMmaCore<Shape_,
|
||||
WarpShape_,
|
||||
GemmShape<16, 16, 16>,
|
||||
ElementA_,
|
||||
layout::RowMajor,
|
||||
ElementB_,
|
||||
layout::ColumnMajor,
|
||||
ElementC_,
|
||||
LayoutC_,
|
||||
arch::OpClassTensorOp,
|
||||
Stages,
|
||||
Operator_> {
|
||||
using Shape = Shape_;
|
||||
using WarpShape = WarpShape_;
|
||||
using InstructionShape = GemmShape<16, 16, 16>;
|
||||
using ElementA = ElementA_;
|
||||
using LayoutA = layout::RowMajor;
|
||||
using ElementB = ElementB_;
|
||||
using LayoutB = layout::ColumnMajor;
|
||||
using ElementC = ElementC_;
|
||||
using LayoutC = LayoutC_;
|
||||
using OperatorClass = arch::OpClassTensorOp;
|
||||
|
||||
static int const kStages = Stages;
|
||||
|
||||
/// Default Operator
|
||||
using Operator = Operator_;
|
||||
|
||||
/// Warp thread arrangement
|
||||
using WarpThreadArrangement = layout::PitchLinearShape<16, 4>;
|
||||
|
||||
/// Number of warps present
|
||||
using WarpCount = GemmShape<
|
||||
Shape::kM / WarpShape::kM,
|
||||
Shape::kN / WarpShape::kN,
|
||||
Shape::kK / WarpShape::kK
|
||||
>;
|
||||
|
||||
/// Don't support split K within CTA
|
||||
static_assert(Shape::kK == WarpShape::kK,
|
||||
"Threadblock-scoped GEMM shape K should equal warp-scoped GEMM shape K"
|
||||
);
|
||||
|
||||
// Divisibility requirements
|
||||
static_assert(
|
||||
!(Shape::kM % WarpShape::kM) &&
|
||||
!(Shape::kN % WarpShape::kN) &&
|
||||
!(Shape::kK % WarpShape::kK),
|
||||
"Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."
|
||||
);
|
||||
|
||||
// Divisibility requirements
|
||||
static_assert(
|
||||
!(WarpShape::kM % 16) &&
|
||||
!(WarpShape::kN % 16) &&
|
||||
!(WarpShape::kK % 16),
|
||||
"Threadblock-scoped GEMM should be divisible by 16."
|
||||
);
|
||||
|
||||
/// Number of threads per warp
|
||||
static int const kWarpSize = warp::WarpSize<arch::OpClassTensorOp>::value;
|
||||
|
||||
/// Number of threads total
|
||||
static int const kThreads = WarpCount::kCount * kWarpSize;
|
||||
|
||||
/// Size of a threadblock-scoped access
|
||||
static int const kAccessSizeInBits = 32;
|
||||
|
||||
/// Number of A elemnts per access
|
||||
static int const kElementsPerAccessA = kAccessSizeInBits / sizeof_bits<ElementA>::value;
|
||||
|
||||
/// Number of A elemnts per access
|
||||
static int const kElementsPerAccessB = kAccessSizeInBits / sizeof_bits<ElementB>::value;
|
||||
|
||||
//
|
||||
// Shared memory layouts
|
||||
//
|
||||
|
||||
#if BLOCK_LOAD_STORE
|
||||
using SmemLayoutA = layout::TensorOpEm<sizeof_bits<ElementA>::value, LayoutA>;
|
||||
using SmemLayoutB = layout::TensorOpMultiplicand<sizeof_bits<ElementB>::value, LayoutB>;
|
||||
#else
|
||||
using SmemLayoutA = layout::TensorOpMultiplicand<sizeof_bits<ElementA>::value, LayoutA>;
|
||||
using SmemLayoutB = layout::TensorOpMultiplicand<sizeof_bits<ElementB>::value, LayoutB>;
|
||||
#endif
|
||||
|
||||
//
|
||||
|
||||
//
|
||||
// Iterators to write to shared memory
|
||||
//
|
||||
|
||||
/// ThreadMap of iterator A
|
||||
///
|
||||
using IteratorThreadMapA = transform::PitchLinear2DThreadTileWarpRakedThreadMap<
|
||||
layout::PitchLinearShape<Shape::kK, Shape::kM>,
|
||||
kThreads,
|
||||
WarpThreadArrangement,
|
||||
layout::PitchLinearShape<kElementsPerAccessA, kElementsPerAccessA>
|
||||
>;
|
||||
|
||||
/// Shared memory iterator to A operand
|
||||
using SmemIteratorA = transform::threadblock::RegularTileIterator<
|
||||
MatrixShape<Shape::kM, Shape::kK>,
|
||||
ElementA,
|
||||
SmemLayoutA,
|
||||
1,
|
||||
IteratorThreadMapA
|
||||
>;
|
||||
|
||||
/// Policy of iterator B
|
||||
using IteratorThreadMapB = transform::PitchLinear2DThreadTileWarpRakedThreadMap<
|
||||
layout::PitchLinearShape<Shape::kK, Shape::kN>,
|
||||
kThreads,
|
||||
WarpThreadArrangement,
|
||||
layout::PitchLinearShape<kElementsPerAccessB, kElementsPerAccessB>
|
||||
>;
|
||||
|
||||
/// Shared memory iterator to B operand
|
||||
using SmemIteratorB = transform::threadblock::RegularTileIterator<
|
||||
MatrixShape<Shape::kK, Shape::kN>,
|
||||
ElementB,
|
||||
SmemLayoutB,
|
||||
0,
|
||||
IteratorThreadMapB
|
||||
>;
|
||||
|
||||
//
|
||||
// Warp-level matrix multiply operator
|
||||
//
|
||||
|
||||
// Define the warp-level tensor op
|
||||
using Policy = gemm::warp::MmaTensorOpPolicy<
|
||||
arch::Mma<
|
||||
gemm::GemmShape<16, 16, 16>,
|
||||
NUM_THREADS_PER_WARP,
|
||||
ElementA,
|
||||
LayoutA,
|
||||
ElementB,
|
||||
LayoutB,
|
||||
ElementC,
|
||||
layout::RowMajor,
|
||||
arch::OpMultiplyAdd
|
||||
>,
|
||||
MatrixShape<1, 1>
|
||||
>;
|
||||
|
||||
using MmaTensorOp = typename gemm::warp::DefaultMmaTensorOp<
|
||||
WarpShape,
|
||||
gemm::GemmShape<16, 16, 16>,
|
||||
ElementA,
|
||||
SmemLayoutA,
|
||||
ElementB,
|
||||
SmemLayoutB,
|
||||
ElementC,
|
||||
LayoutC,
|
||||
arch::OpMultiplyAdd
|
||||
>::Type;
|
||||
|
||||
/// Policy used to define MmaPipelined
|
||||
using MmaPolicy = MmaPolicy<
|
||||
MmaTensorOp,
|
||||
MatrixShape<0, 0>,
|
||||
MatrixShape<0, 0>,
|
||||
WarpCount::kK
|
||||
>;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
///
|
||||
/// Specialization: A: column-major, B: row-major, NT
|
||||
///
|
||||
/// This uses the default warp-level operator given tile sizes
|
||||
///
|
||||
template <
|
||||
/// Shape of threadblock-scoped matrix multiply operator (concept:
|
||||
/// GemmShape)
|
||||
typename Shape_,
|
||||
/// Shape of warp-level matrix multiply operator (concept: GemmShape)
|
||||
typename WarpShape_,
|
||||
/// Data type of A operand
|
||||
typename ElementA_,
|
||||
/// Data type of B operand
|
||||
typename ElementB_,
|
||||
/// Data type of accumulator
|
||||
typename ElementC_,
|
||||
/// Layout of accumulator
|
||||
typename LayoutC_,
|
||||
/// Stages
|
||||
int Stages,
|
||||
/// Operation performed by GEMM
|
||||
typename Operator_>
|
||||
struct DefaultMmaCore<Shape_,
|
||||
WarpShape_,
|
||||
GemmShape<16, 16, 16>,
|
||||
ElementA_,
|
||||
layout::ColumnMajor,
|
||||
ElementB_,
|
||||
layout::RowMajor,
|
||||
ElementC_,
|
||||
LayoutC_,
|
||||
arch::OpClassTensorOp,
|
||||
Stages,
|
||||
Operator_> {
|
||||
using Shape = Shape_;
|
||||
using WarpShape = WarpShape_;
|
||||
using InstructionShape = GemmShape<16, 16, 16>;
|
||||
using ElementA = ElementA_;
|
||||
using LayoutA = layout::ColumnMajor;
|
||||
using ElementB = ElementB_;
|
||||
using LayoutB = layout::RowMajor;
|
||||
using ElementC = ElementC_;
|
||||
using LayoutC = LayoutC_;
|
||||
using OperatorClass = arch::OpClassTensorOp;
|
||||
|
||||
static int const kStages = Stages;
|
||||
|
||||
/// Default Operator
|
||||
using Operator = Operator_;
|
||||
|
||||
/// Warp thread arrangement
|
||||
using WarpThreadArrangement = layout::PitchLinearShape<16, 4>;
|
||||
|
||||
/// Number of warps present
|
||||
using WarpCount = GemmShape<
|
||||
Shape::kM / WarpShape::kM,
|
||||
Shape::kN / WarpShape::kN,
|
||||
Shape::kK / WarpShape::kK
|
||||
>;
|
||||
|
||||
/// Don't support split K within CTA
|
||||
static_assert(Shape::kK == WarpShape::kK,
|
||||
"Threadblock-scoped GEMM shape K should equal warp-scoped GEMM shape K"
|
||||
);
|
||||
|
||||
// Divisibility requirements
|
||||
static_assert(
|
||||
!(Shape::kM % WarpShape::kM) &&
|
||||
!(Shape::kN % WarpShape::kN) &&
|
||||
!(Shape::kK % WarpShape::kK),
|
||||
"Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."
|
||||
);
|
||||
|
||||
// Divisibility requirements
|
||||
static_assert(
|
||||
!(WarpShape::kM % 16) &&
|
||||
!(WarpShape::kN % 16) &&
|
||||
!(WarpShape::kK % 16),
|
||||
"Threadblock-scoped GEMM should be divisible by 16."
|
||||
);
|
||||
|
||||
/// Number of threads per warp
|
||||
static int const kWarpSize = warp::WarpSize<arch::OpClassTensorOp>::value;
|
||||
|
||||
/// Number of threads total
|
||||
static int const kThreads = WarpCount::kCount * kWarpSize;
|
||||
|
||||
/// Size of a threadblock-scoped access
|
||||
static int const kAccessSizeInBits = 32;
|
||||
|
||||
/// Number of A elemnts per access
|
||||
static int const kElementsPerAccessA = kAccessSizeInBits / sizeof_bits<ElementA>::value;
|
||||
|
||||
/// Number of A elemnts per access
|
||||
static int const kElementsPerAccessB = kAccessSizeInBits / sizeof_bits<ElementB>::value;
|
||||
|
||||
//
|
||||
// Shared memory layouts
|
||||
//
|
||||
|
||||
#if BLOCK_LOAD_STORE
|
||||
using SmemLayoutA = layout::TensorOpMultiplicand<sizeof_bits<ElementA>::value, LayoutA>;
|
||||
using SmemLayoutB = layout::TensorOpEm<sizeof_bits<ElementB>::value, LayoutB>;
|
||||
#else
|
||||
using SmemLayoutA = layout::TensorOpMultiplicand<sizeof_bits<ElementA>::value, LayoutA>;
|
||||
using SmemLayoutB = layout::TensorOpMultiplicand<sizeof_bits<ElementB>::value, LayoutB>;
|
||||
#endif
|
||||
|
||||
//
|
||||
// Iterators to write to shared memory
|
||||
//
|
||||
|
||||
/// ThreadMap of iterator A
|
||||
///
|
||||
using IteratorThreadMapA = transform::PitchLinear2DThreadTileWarpRakedThreadMap<
|
||||
layout::PitchLinearShape<Shape::kM, Shape::kK>,
|
||||
kThreads,
|
||||
WarpThreadArrangement,
|
||||
layout::PitchLinearShape<kElementsPerAccessB, kElementsPerAccessB>
|
||||
>;
|
||||
|
||||
/// Shared memory iterator to A operand
|
||||
using SmemIteratorA = transform::threadblock::RegularTileIterator<
|
||||
MatrixShape<Shape::kM, Shape::kK>,
|
||||
ElementA,
|
||||
SmemLayoutA,
|
||||
1,
|
||||
IteratorThreadMapA
|
||||
>;
|
||||
|
||||
/// Policy of iterator B
|
||||
using IteratorThreadMapB = transform::PitchLinear2DThreadTileWarpRakedThreadMap<
|
||||
layout::PitchLinearShape<Shape::kN, Shape::kK>,
|
||||
kThreads,
|
||||
WarpThreadArrangement,
|
||||
layout::PitchLinearShape<kElementsPerAccessA, kElementsPerAccessA>
|
||||
>;
|
||||
|
||||
/// Shared memory iterator to B operand
|
||||
using SmemIteratorB = transform::threadblock::RegularTileIterator<
|
||||
MatrixShape<Shape::kK, Shape::kN>,
|
||||
ElementB,
|
||||
SmemLayoutB,
|
||||
0,
|
||||
IteratorThreadMapB
|
||||
>;
|
||||
|
||||
//
|
||||
// Warp-level matrix multiply operator
|
||||
//
|
||||
|
||||
// Define the warp-level tensor op
|
||||
using Policy = gemm::warp::MmaTensorOpPolicy<
|
||||
arch::Mma<
|
||||
gemm::GemmShape<16, 16, 16>,
|
||||
NUM_THREADS_PER_WARP,
|
||||
ElementA,
|
||||
LayoutA,
|
||||
ElementB,
|
||||
LayoutB,
|
||||
ElementC,
|
||||
layout::RowMajor,
|
||||
arch::OpMultiplyAdd
|
||||
>,
|
||||
MatrixShape<1, 1>
|
||||
>;
|
||||
|
||||
using MmaTensorOp = typename gemm::warp::DefaultMmaTensorOp<
|
||||
WarpShape,
|
||||
gemm::GemmShape<16, 16, 16>,
|
||||
ElementA,
|
||||
SmemLayoutA,
|
||||
ElementB,
|
||||
SmemLayoutB,
|
||||
ElementC,
|
||||
LayoutC,
|
||||
arch::OpMultiplyAdd
|
||||
>::Type;
|
||||
|
||||
/// Policy used to define MmaPipelined
|
||||
using MmaPolicy = MmaPolicy<
|
||||
MmaTensorOp,
|
||||
MatrixShape<0, 0>,
|
||||
MatrixShape<0, 0>,
|
||||
WarpCount::kK
|
||||
>;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
///
|
||||
/// Specialization: A: column-major, B: column-major, NN
|
||||
///
|
||||
/// This uses the default warp-level operator given tile sizes
|
||||
///
|
||||
template <
|
||||
/// Shape of threadblock-scoped matrix multiply operator (concept:
|
||||
/// GemmShape)
|
||||
typename Shape_,
|
||||
/// Shape of warp-level matrix multiply operator (concept: GemmShape)
|
||||
typename WarpShape_,
|
||||
/// Data type of A operand
|
||||
typename ElementA_,
|
||||
/// Data type of B operand
|
||||
typename ElementB_,
|
||||
/// Data type of accumulator
|
||||
typename ElementC_,
|
||||
/// Layout of accumulator
|
||||
typename LayoutC_,
|
||||
/// Stages
|
||||
int Stages,
|
||||
/// Operation performed by GEMM
|
||||
typename Operator_>
|
||||
struct DefaultMmaCore<Shape_,
|
||||
WarpShape_,
|
||||
GemmShape<16, 16, 16>,
|
||||
ElementA_,
|
||||
layout::ColumnMajor,
|
||||
ElementB_,
|
||||
layout::ColumnMajor,
|
||||
ElementC_,
|
||||
LayoutC_,
|
||||
arch::OpClassTensorOp,
|
||||
Stages,
|
||||
Operator_> {
|
||||
using Shape = Shape_;
|
||||
using WarpShape = WarpShape_;
|
||||
using InstructionShape = GemmShape<16, 16, 16>;
|
||||
using ElementA = ElementA_;
|
||||
using LayoutA = layout::ColumnMajor;
|
||||
using ElementB = ElementB_;
|
||||
using LayoutB = layout::ColumnMajor;
|
||||
using ElementC = ElementC_;
|
||||
using LayoutC = LayoutC_;
|
||||
using OperatorClass = arch::OpClassTensorOp;
|
||||
|
||||
static int const kStages = Stages;
|
||||
|
||||
/// Default Operator
|
||||
using Operator = Operator_;
|
||||
|
||||
/// Warp thread arrangement
|
||||
using WarpThreadArrangement = layout::PitchLinearShape<16, 4>;
|
||||
|
||||
/// Number of warps present
|
||||
using WarpCount = GemmShape<
|
||||
Shape::kM / WarpShape::kM,
|
||||
Shape::kN / WarpShape::kN,
|
||||
Shape::kK / WarpShape::kK
|
||||
>;
|
||||
|
||||
/// Don't support split K within CTA
|
||||
static_assert(Shape::kK == WarpShape::kK,
|
||||
"Threadblock-scoped GEMM shape K should equal warp-scoped GEMM shape K"
|
||||
);
|
||||
|
||||
// Divisibility requirements
|
||||
static_assert(
|
||||
!(Shape::kM % WarpShape::kM) &&
|
||||
!(Shape::kN % WarpShape::kN) &&
|
||||
!(Shape::kK % WarpShape::kK),
|
||||
"Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."
|
||||
);
|
||||
|
||||
// Divisibility requirements
|
||||
static_assert(
|
||||
!(WarpShape::kM % 16) &&
|
||||
!(WarpShape::kN % 16) &&
|
||||
!(WarpShape::kK % 16),
|
||||
"Threadblock-scoped GEMM should be divisible by 16."
|
||||
);
|
||||
|
||||
/// Number of threads per warp
|
||||
static int const kWarpSize = warp::WarpSize<arch::OpClassTensorOp>::value;
|
||||
|
||||
/// Number of threads total
|
||||
static int const kThreads = WarpCount::kCount * kWarpSize;
|
||||
|
||||
/// Size of a threadblock-scoped access
|
||||
static int const kAccessSizeInBits = 32;
|
||||
|
||||
/// Number of A elemnts per access
|
||||
static int const kElementsPerAccessA = kAccessSizeInBits / sizeof_bits<ElementA>::value;
|
||||
|
||||
/// Number of A elemnts per access
|
||||
static int const kElementsPerAccessB = kAccessSizeInBits / sizeof_bits<ElementB>::value;
|
||||
|
||||
//
|
||||
// Shared memory layouts
|
||||
//
|
||||
using SmemLayoutA = layout::TensorOpMultiplicand<sizeof_bits<ElementA>::value, LayoutA>;
|
||||
using SmemLayoutB = layout::TensorOpMultiplicand<sizeof_bits<ElementB>::value, LayoutB>;
|
||||
|
||||
//
|
||||
|
||||
//
|
||||
// Iterators to write to shared memory
|
||||
//
|
||||
|
||||
/// ThreadMap of iterator A
|
||||
///
|
||||
using IteratorThreadMapA = transform::PitchLinear2DThreadTileWarpRakedThreadMap<
|
||||
layout::PitchLinearShape<Shape::kM, Shape::kK>,
|
||||
kThreads,
|
||||
WarpThreadArrangement,
|
||||
layout::PitchLinearShape<kElementsPerAccessA, kElementsPerAccessA>
|
||||
>;
|
||||
|
||||
/// Shared memory iterator to A operand
|
||||
using SmemIteratorA = transform::threadblock::RegularTileIterator<
|
||||
MatrixShape<Shape::kM, Shape::kK>,
|
||||
ElementA,
|
||||
SmemLayoutA,
|
||||
1,
|
||||
IteratorThreadMapA
|
||||
>;
|
||||
|
||||
/// Policy of iterator B
|
||||
using IteratorThreadMapB = transform::PitchLinear2DThreadTileWarpRakedThreadMap<
|
||||
layout::PitchLinearShape<Shape::kK, Shape::kN>,
|
||||
kThreads,
|
||||
WarpThreadArrangement,
|
||||
layout::PitchLinearShape<kElementsPerAccessB, kElementsPerAccessB>
|
||||
>;
|
||||
|
||||
/// Shared memory iterator to B operand
|
||||
using SmemIteratorB = transform::threadblock::RegularTileIterator<
|
||||
MatrixShape<Shape::kK, Shape::kN>,
|
||||
ElementB,
|
||||
SmemLayoutB,
|
||||
0,
|
||||
IteratorThreadMapB
|
||||
>;
|
||||
|
||||
//
|
||||
// Warp-level matrix multiply operator
|
||||
//
|
||||
|
||||
// Define the warp-level tensor op
|
||||
using Policy = gemm::warp::MmaTensorOpPolicy<
|
||||
arch::Mma<
|
||||
gemm::GemmShape<16, 16, 16>,
|
||||
NUM_THREADS_PER_WARP,
|
||||
ElementA,
|
||||
LayoutA,
|
||||
ElementB,
|
||||
LayoutB,
|
||||
ElementC,
|
||||
layout::RowMajor,
|
||||
arch::OpMultiplyAdd
|
||||
>,
|
||||
MatrixShape<1, 1>
|
||||
>;
|
||||
|
||||
using MmaTensorOp = typename gemm::warp::DefaultMmaTensorOp<
|
||||
WarpShape,
|
||||
gemm::GemmShape<16, 16, 16>,
|
||||
ElementA,
|
||||
SmemLayoutA,
|
||||
ElementB,
|
||||
SmemLayoutB,
|
||||
ElementC,
|
||||
LayoutC,
|
||||
arch::OpMultiplyAdd
|
||||
>::Type;
|
||||
|
||||
/// Policy used to define MmaPipelined
|
||||
using MmaPolicy = MmaPolicy<
|
||||
MmaTensorOp,
|
||||
MatrixShape<0, 0>,
|
||||
MatrixShape<0, 0>,
|
||||
WarpCount::kK
|
||||
>;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace gemm
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
148
cat_files/default_mma_tensor_op.h
Normal file
148
cat_files/default_mma_tensor_op.h
Normal file
@@ -0,0 +1,148 @@
|
||||
/***************************************************************************************************
|
||||
* 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:
|
||||
* * 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.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2021 Iluvatar CoreX. All rights reserved.
|
||||
* Copyright Declaration: This software, including all of its code and documentation,
|
||||
* except for the third-party software it contains, is a copyrighted work of Shanghai Iluvatar CoreX
|
||||
* Semiconductor Co., Ltd. and its affiliates ("Iluvatar CoreX") in accordance with the PRC Copyright
|
||||
* Law and relevant international treaties, and all rights contained therein are enjoyed by Iluvatar
|
||||
* CoreX. No user of this software shall have any right, ownership or interest in this software and
|
||||
* any use of this software shall be in compliance with the terms and conditions of the End User
|
||||
* License Agreement.
|
||||
**************************************************************************************************/
|
||||
|
||||
/*! \file
|
||||
\brief Default warp-level GEMM operators selected by data type, size, and layouts of operands.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/gemm/warp/mma_tensor_op.h"
|
||||
|
||||
namespace cutlass {
|
||||
namespace gemm {
|
||||
namespace warp {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
/// Size of the Gemm problem - concept: gemm::GemmShape<>
|
||||
typename WarpShape_,
|
||||
/// Shape of one matrix production operation (concept: GemmShape)
|
||||
typename InstructionShape_,
|
||||
/// Data type of A elements
|
||||
typename ElementA_,
|
||||
/// Layout of A matrix (concept: MatrixLayout)
|
||||
typename LayoutA_,
|
||||
/// Data type of B elements
|
||||
typename ElementB_,
|
||||
/// Layout of B matrix (concept: MatrixLayout)
|
||||
typename LayoutB_,
|
||||
/// Element type of C matrix
|
||||
typename ElementC_,
|
||||
/// Layout of C matrix (concept: MatrixLayout)
|
||||
typename LayoutC_,
|
||||
/// Operator describing the tensor operation
|
||||
typename Operator_ = arch::OpMultiplyAdd,
|
||||
/// Number of partitions along K dimension
|
||||
int PartitionsK = 1,
|
||||
/// Store the accumulators in row major or column major.
|
||||
bool AccumulatorsInRowMajor = true>
|
||||
struct DefaultMmaTensorOp;
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Partial specialization for m-by-n-by-kgroup
|
||||
template <
|
||||
/// Shape of one matrix production operation (concept: GemmShape)
|
||||
typename WarpShape_,
|
||||
/// Data type of A elements
|
||||
typename ElementA,
|
||||
/// Layout of A matrix (concept: MatrixLayout)
|
||||
typename LayoutA,
|
||||
/// Data type of B elements
|
||||
typename ElementB,
|
||||
/// Layout of B matrix (concept: MatrixLayout)
|
||||
typename LayoutB,
|
||||
/// Element type of C matrix
|
||||
typename ElementC,
|
||||
/// Layout of C matrix (concept: MatrixLayout)
|
||||
typename LayoutC,
|
||||
/// Number of partitions along K dimension
|
||||
int PartitionsK,
|
||||
/// Store the accumulators in row major or column major.
|
||||
bool AccumulatorsInRowMajor>
|
||||
struct DefaultMmaTensorOp<
|
||||
WarpShape_,
|
||||
GemmShape<16, 16, 16>,
|
||||
ElementA,
|
||||
LayoutA,
|
||||
ElementB,
|
||||
LayoutB,
|
||||
ElementC,
|
||||
LayoutC,
|
||||
arch::OpMultiplyAdd,
|
||||
PartitionsK,
|
||||
AccumulatorsInRowMajor> {
|
||||
|
||||
/// Warp shape
|
||||
using Shape = WarpShape_;
|
||||
|
||||
using Policy = cutlass::gemm::warp::MmaTensorOpPolicy<
|
||||
cutlass::arch::Mma<GemmShape<16, 16, 16>,
|
||||
64,
|
||||
ElementA,
|
||||
LayoutA,
|
||||
ElementB,
|
||||
LayoutB,
|
||||
ElementC,
|
||||
LayoutC,
|
||||
arch::OpMultiplyAdd>,
|
||||
cutlass::MatrixShape<1, 1> >;
|
||||
|
||||
// Define the warp-level tensor op
|
||||
using Type = cutlass::gemm::warp::MmaTensorOp<
|
||||
WarpShape_,
|
||||
ElementA,
|
||||
LayoutA,
|
||||
ElementB,
|
||||
LayoutB,
|
||||
ElementC,
|
||||
LayoutC,
|
||||
Policy,
|
||||
PartitionsK,
|
||||
AccumulatorsInRowMajor>;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace warp
|
||||
} // namespace gemm
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
726
cat_files/gemm_batched.h
Normal file
726
cat_files/gemm_batched.h
Normal file
@@ -0,0 +1,726 @@
|
||||
/***************************************************************************************************
|
||||
* 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:
|
||||
* * 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.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2021 Iluvatar CoreX. All rights reserved.
|
||||
* Copyright Declaration: This software, including all of its code and documentation,
|
||||
* except for the third-party software it contains, is a copyrighted work of Shanghai Iluvatar CoreX
|
||||
* Semiconductor Co., Ltd. and its affiliates ("Iluvatar CoreX") in accordance with the PRC Copyright
|
||||
* Law and relevant international treaties, and all rights contained therein are enjoyed by Iluvatar
|
||||
* CoreX. No user of this software shall have any right, ownership or interest in this software and
|
||||
* any use of this software shall be in compliance with the terms and conditions of the End User
|
||||
* License Agreement.
|
||||
**************************************************************************************************/
|
||||
|
||||
/*! \file
|
||||
\brief Template for a pipelined GEMM kernel. Does not compute batching or support split-K.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/arch/arch.h"
|
||||
#include "cutlass/device_kernel.h"
|
||||
|
||||
#include "cutlass/gemm/threadblock/threadblock_swizzle.h"
|
||||
#include "cutlass/gemm/kernel/gemm_batched.h"
|
||||
|
||||
#include "cutlass/gemm/kernel/default_gemm.h"
|
||||
#include "cutlass/gemm/device/default_gemm_configuration.h"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace gemm {
|
||||
namespace device {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/*! Gemm device-level operator. This is an interface to efficient CUTLASS GEMM kernels that may
|
||||
be invoked from host code.
|
||||
|
||||
The contributions of this class are:
|
||||
|
||||
1. At compile time, it maps data types and high-level structural parameters onto
|
||||
specific CUTLASS components.
|
||||
|
||||
2. At runtime, it maps logical arguments to GEMM problems to kernel parameters.
|
||||
|
||||
3. At runtime, it launches kernels on the device.
|
||||
|
||||
The intent is to provide a convenient mechanism for interacting with most plausible GEMM
|
||||
configurations for each supported architecture. Consequently, not all parameters are exposed
|
||||
to the top-level interface. Rather, sensible defaults at each level of the CUTLASS hierarchy
|
||||
are selected to tradeoff simplicity of the interface with flexibility. We expect
|
||||
most configurations to be specified at this level. Applications with more exotic requirements
|
||||
may construct their kernels of interest using CUTLASS components at the threadblock, warp,
|
||||
and thread levels of abstraction.
|
||||
|
||||
CUTLASS exposes computations using the functor design pattern in which objects compose some
|
||||
internal state with an overloaded function call operator. This enables decoupling of
|
||||
initialization from execution, possibly reducing overhead during steady state phases of
|
||||
application execution.
|
||||
|
||||
CUTLASS device-level operators expose an Arguments structure encompassing each logical
|
||||
input to the computation. This is distinct from the kernel-level Params structure pattern
|
||||
which contains application-specific precomputed state needed by the device code.
|
||||
|
||||
Example of a CUTLASS GEMM operator implementing the functionality of cuBLAS's SGEMM NN
|
||||
is as follows:
|
||||
|
||||
//
|
||||
// Instantiate the CUTLASS GEMM operator.
|
||||
//
|
||||
|
||||
cutlass::gemm::device::Gemm<
|
||||
float,
|
||||
cutlass::layout::ColumnMajor,
|
||||
float,
|
||||
cutlass::layout::ColumnMajor,
|
||||
float,
|
||||
cutlass::layout::ColumnMajor
|
||||
> gemm_op;
|
||||
|
||||
//
|
||||
// Launch the GEMM operation on the device
|
||||
//
|
||||
|
||||
cutlass::Status status = gemm_op({
|
||||
{m, n, k}, // GemmCoord problem_size,
|
||||
{A, lda}, // TensorRef<float, layout::ColumnMajor> ref_A,
|
||||
{B, ldb}, // TensorRef<float, layout::ColumnMajor> ref_B,
|
||||
{C, ldc}, // TensorRef<float, layout::ColumnMajor> ref_C,
|
||||
{D, ldd}, // TensorRef<float, layout::ColumnMajor> ref_D,
|
||||
{alpha, beta} // EpilogueOutputOp::Params epilogue_op_params
|
||||
});
|
||||
|
||||
|
||||
A simplified view of the template is listed below.
|
||||
|
||||
template <
|
||||
/// Element type for A matrix operand
|
||||
typename ElementA,
|
||||
|
||||
/// Layout type for A matrix operand
|
||||
typename LayoutA,
|
||||
|
||||
/// Element type for B matrix operand
|
||||
typename ElementB,
|
||||
|
||||
/// Layout type for B matrix operand
|
||||
typename LayoutB,
|
||||
|
||||
/// Element type for C and D matrix operands
|
||||
typename ElementC,
|
||||
|
||||
/// Layout type for C and D matrix operands
|
||||
typename LayoutC,
|
||||
|
||||
/// Element type for internal accumulation
|
||||
typename ElementAccumulator,
|
||||
|
||||
/// Operator class tag
|
||||
typename OperatorClass,
|
||||
|
||||
/// Tag indicating architecture to tune for. This is the minimum SM that
|
||||
/// supports the intended feature. The device kernel can be built
|
||||
/// targeting any SM larger than this number.
|
||||
typename ArchTag,
|
||||
|
||||
/// Threadblock-level tile size (concept: GemmShape)
|
||||
typename ThreadblockShape,
|
||||
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename WarpShape,
|
||||
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename InstructionShape,
|
||||
|
||||
/// Epilogue output operator
|
||||
typename EpilogueOutputOp,
|
||||
|
||||
/// Threadblock-level swizzling operator
|
||||
typename ThreadblockSwizzle,
|
||||
|
||||
/// Number of stages used in the pipelined mainloop
|
||||
int Stages
|
||||
>
|
||||
class Gemm;
|
||||
*/
|
||||
template <
|
||||
/// Element type for A matrix operand
|
||||
typename ElementA_,
|
||||
/// Layout type for A matrix operand
|
||||
typename LayoutA_,
|
||||
/// Element type for B matrix operand
|
||||
typename ElementB_,
|
||||
/// Layout type for B matrix operand
|
||||
typename LayoutB_,
|
||||
/// Element type for C and D matrix operands
|
||||
typename ElementC_,
|
||||
/// Layout type for C and D matrix operands
|
||||
typename LayoutC_,
|
||||
/// Element type for internal accumulation
|
||||
typename ElementAccumulator_ = ElementC_,
|
||||
/// Operator class tag
|
||||
typename OperatorClass_ = arch::OpClassSimt,
|
||||
/// Tag indicating architecture to tune for
|
||||
typename ArchTag_ = arch::Sm61,
|
||||
/// Threadblock-level tile size (concept: GemmShape)
|
||||
typename ThreadblockShape_ = typename DefaultGemmConfiguration<
|
||||
OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
|
||||
ElementAccumulator_>::ThreadblockShape,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename WarpShape_ = typename DefaultGemmConfiguration<
|
||||
OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
|
||||
ElementAccumulator_>::WarpShape,
|
||||
/// Instruction-level tile size (concept: GemmShape)
|
||||
typename InstructionShape_ = typename DefaultGemmConfiguration<
|
||||
OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
|
||||
ElementAccumulator_>::InstructionShape,
|
||||
/// Epilogue output operator
|
||||
typename EpilogueOutputOp_ = typename DefaultGemmConfiguration<
|
||||
OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
|
||||
ElementAccumulator_>::EpilogueOutputOp,
|
||||
/// Threadblock-level swizzling operator
|
||||
typename ThreadblockSwizzle_ = threadblock::GemmBatchedIdentityThreadblockSwizzle,
|
||||
/// Number of stages used in the pipelined mainloop
|
||||
int Stages =
|
||||
DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
|
||||
ElementC_, ElementAccumulator_>::kStages,
|
||||
/// Access granularity of A matrix in units of elements
|
||||
int AlignmentA =
|
||||
DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
|
||||
ElementC_, ElementAccumulator_>::kAlignmentA,
|
||||
/// Access granularity of B matrix in units of elements
|
||||
int AlignmentB =
|
||||
DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
|
||||
ElementC_, ElementAccumulator_>::kAlignmentB,
|
||||
/// Operation performed by GEMM
|
||||
typename Operator_ = typename DefaultGemmConfiguration<
|
||||
OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
|
||||
ElementAccumulator_>::Operator
|
||||
>
|
||||
class GemmBatched {
|
||||
public:
|
||||
|
||||
using ElementA = ElementA_;
|
||||
using LayoutA = LayoutA_;
|
||||
using TensorRefA = TensorRef<ElementA const, LayoutA>;
|
||||
using ElementB = ElementB_;
|
||||
using LayoutB = LayoutB_;
|
||||
using TensorRefB = TensorRef<ElementB const, LayoutB>;
|
||||
using ElementC = ElementC_;
|
||||
using LayoutC = LayoutC_;
|
||||
using TensorRefC = TensorRef<ElementC const, LayoutC>;
|
||||
using TensorRefD = TensorRef<ElementC, LayoutC>;
|
||||
using ElementAccumulator = ElementAccumulator_;
|
||||
using OperatorClass = OperatorClass_;
|
||||
using ArchTag = ArchTag_;
|
||||
using ThreadblockShape = ThreadblockShape_;
|
||||
using WarpShape = WarpShape_;
|
||||
using InstructionShape = InstructionShape_;
|
||||
using EpilogueOutputOp = EpilogueOutputOp_;
|
||||
using ThreadblockSwizzle = ThreadblockSwizzle_;
|
||||
static int const kStages = Stages;
|
||||
static int const kAlignmentA = AlignmentA;
|
||||
static int const kAlignmentB = AlignmentB;
|
||||
static int const kAlignmentC = EpilogueOutputOp::kCount;
|
||||
using Operator = Operator_;
|
||||
|
||||
/// Define the kernel
|
||||
using DefaultGemmKernel = typename kernel::DefaultGemm<
|
||||
ElementA,
|
||||
LayoutA,
|
||||
kAlignmentA,
|
||||
ElementB,
|
||||
LayoutB,
|
||||
kAlignmentB,
|
||||
ElementC,
|
||||
LayoutC,
|
||||
ElementAccumulator,
|
||||
OperatorClass,
|
||||
ArchTag,
|
||||
ThreadblockShape,
|
||||
WarpShape,
|
||||
InstructionShape,
|
||||
EpilogueOutputOp,
|
||||
ThreadblockSwizzle,
|
||||
kStages,
|
||||
false,
|
||||
Operator
|
||||
>::GemmKernel;
|
||||
|
||||
using GemmKernel = kernel::GemmBatched<typename DefaultGemmKernel::Mma, typename DefaultGemmKernel::Epilogue, ThreadblockSwizzle>;
|
||||
|
||||
/// Argument structure
|
||||
struct Arguments {
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
GemmCoord problem_size;
|
||||
TensorRef<ElementA const, LayoutA> ref_A;
|
||||
int64_t stride_A;
|
||||
TensorRef<ElementB const, LayoutB> ref_B;
|
||||
int64_t stride_B;
|
||||
TensorRef<ElementC const, LayoutC> ref_C;
|
||||
int64_t stride_C;
|
||||
TensorRef<ElementC, LayoutC> ref_D;
|
||||
int64_t stride_D;
|
||||
typename EpilogueOutputOp::Params epilogue;
|
||||
int batch_count;
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
/// Default ctor
|
||||
CUTLASS_HOST_DEVICE
|
||||
Arguments() { }
|
||||
|
||||
/// Constructs an Arguments structure
|
||||
CUTLASS_HOST_DEVICE
|
||||
Arguments(
|
||||
GemmCoord problem_size_,
|
||||
TensorRef<ElementA const, LayoutA> ref_A_,
|
||||
int64_t stride_A_,
|
||||
TensorRef<ElementB const, LayoutB> ref_B_,
|
||||
int64_t stride_B_,
|
||||
TensorRef<ElementC const, LayoutC> ref_C_,
|
||||
int64_t stride_C_,
|
||||
TensorRef<ElementC, LayoutC> ref_D_,
|
||||
int64_t stride_D_,
|
||||
typename EpilogueOutputOp::Params epilogue_,
|
||||
int batch_count_
|
||||
):
|
||||
problem_size(problem_size_),
|
||||
ref_A(ref_A_),
|
||||
stride_A(stride_A_),
|
||||
ref_B(ref_B_),
|
||||
stride_B(stride_B_),
|
||||
ref_C(ref_C_),
|
||||
stride_C(stride_C_),
|
||||
ref_D(ref_D_),
|
||||
stride_D(stride_D_),
|
||||
epilogue(epilogue_),
|
||||
batch_count(batch_count_) { }
|
||||
};
|
||||
|
||||
private:
|
||||
|
||||
/// Kernel parameters object
|
||||
typename GemmKernel::Params params_;
|
||||
|
||||
public:
|
||||
|
||||
/// Constructs the GEMM.
|
||||
GemmBatched() { }
|
||||
|
||||
/// Determines whether the GEMM can execute the given problem.
|
||||
static Status can_implement(Arguments const &args) {
|
||||
|
||||
if (!TensorRef_aligned(args.ref_A, kAlignmentA) || (args.stride_A % kAlignmentA)) {
|
||||
return Status::kErrorMisalignedOperand;
|
||||
}
|
||||
|
||||
if (!TensorRef_aligned(args.ref_B, kAlignmentB) || (args.stride_B % kAlignmentB)) {
|
||||
return Status::kErrorMisalignedOperand;
|
||||
}
|
||||
|
||||
if (!TensorRef_aligned(args.ref_C, kAlignmentC) || (args.stride_C % kAlignmentC)) {
|
||||
return Status::kErrorMisalignedOperand;
|
||||
}
|
||||
|
||||
if (!TensorRef_aligned(args.ref_D, kAlignmentC) || (args.stride_D % kAlignmentC)) {
|
||||
return Status::kErrorMisalignedOperand;
|
||||
}
|
||||
|
||||
if ((args.problem_size.m() % kAlignmentA) || (args.problem_size.k() % kAlignmentA) ||
|
||||
(args.problem_size.n() % kAlignmentB) || (args.problem_size.k() % kAlignmentB) ||
|
||||
(args.problem_size.m() % kAlignmentC) || (args.problem_size.n() % kAlignmentC)) {
|
||||
|
||||
return Status::kErrorMisalignedOperand;
|
||||
}
|
||||
|
||||
return Status::kSuccess;
|
||||
}
|
||||
|
||||
/// Gets the workspace size
|
||||
static size_t get_workspace_size(Arguments const &args) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
/// Initializes GEMM state from arguments.
|
||||
Status initialize(Arguments const &args, void *workspace = nullptr, cudaStream_t stream = nullptr) {
|
||||
|
||||
// Determine grid shape
|
||||
ThreadblockSwizzle threadblock_swizzle;
|
||||
|
||||
cutlass::gemm::GemmCoord grid_shape = threadblock_swizzle.get_tiled_shape(
|
||||
args.problem_size,
|
||||
{ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK},
|
||||
args.batch_count);
|
||||
|
||||
// Initialize the Params structure
|
||||
params_ = typename GemmKernel::Params{
|
||||
args.problem_size,
|
||||
grid_shape,
|
||||
args.ref_A.non_const_ref(),
|
||||
args.stride_A,
|
||||
args.ref_B.non_const_ref(),
|
||||
args.stride_B,
|
||||
args.ref_C.non_const_ref(),
|
||||
args.stride_C,
|
||||
args.ref_D,
|
||||
args.stride_D,
|
||||
args.epilogue,
|
||||
args.batch_count
|
||||
};
|
||||
|
||||
return Status::kSuccess;
|
||||
}
|
||||
|
||||
/// Lightweight update given a subset of arguments
|
||||
Status update(Arguments const &args, void *workspace = nullptr) {
|
||||
|
||||
params_.ref_A.reset(args.ref_A.non_const_ref().data());
|
||||
params_.ref_B.reset(args.ref_B.non_const_ref().data());
|
||||
params_.ref_C.reset(args.ref_C.non_const_ref().data());
|
||||
params_.ref_D.reset(args.ref_D.data());
|
||||
|
||||
return Status::kSuccess;
|
||||
}
|
||||
|
||||
/// Runs the kernel using initialized state.
|
||||
Status run(cudaStream_t stream = nullptr) {
|
||||
|
||||
ThreadblockSwizzle threadblock_swizzle;
|
||||
|
||||
dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape);
|
||||
// XXX(Peter Han): prealod needs double warps in z direction
|
||||
dim3 block(GemmKernel::kThreadCount, 1, kStages ? 1 : 2);
|
||||
|
||||
cudaError_t result;
|
||||
|
||||
int smem_size = int(sizeof(typename GemmKernel::SharedStorage));
|
||||
/// cudaFuncSetAttribute isn't supported under CUDA-8.0
|
||||
// if (smem_size >= (48 << 10)) {
|
||||
// result = cudaFuncSetAttribute(Kernel<GemmKernel>,
|
||||
// cudaFuncAttributeMaxDynamicSharedMemorySize,
|
||||
// smem_size);
|
||||
|
||||
// if (result != cudaSuccess) {
|
||||
// return Status::kErrorInternal;
|
||||
// }
|
||||
|
||||
// result = cudaFuncSetAttribute(
|
||||
// Kernel<GemmKernel>,
|
||||
// cudaFuncAttributePreferredSharedMemoryCarveout, 100);
|
||||
|
||||
// if (result != cudaSuccess) {
|
||||
// return Status::kErrorInternal;
|
||||
// }
|
||||
// }
|
||||
|
||||
cutlass::Kernel<GemmKernel><<<grid, block, smem_size, stream>>>(params_);
|
||||
|
||||
result = cudaGetLastError();
|
||||
|
||||
return result == cudaSuccess ? Status::kSuccess : Status::kErrorInternal;
|
||||
}
|
||||
|
||||
/// Runs the kernel using initialized state.
|
||||
Status operator()(cudaStream_t stream = nullptr) {
|
||||
return run(stream);
|
||||
}
|
||||
|
||||
/// Runs the kernel using initialized state.
|
||||
Status operator()(
|
||||
Arguments const &args,
|
||||
void *workspace = nullptr,
|
||||
cudaStream_t stream = nullptr) {
|
||||
|
||||
Status status = initialize(args, workspace, stream);
|
||||
|
||||
if (status == Status::kSuccess) {
|
||||
status = run(stream);
|
||||
}
|
||||
|
||||
return status;
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Parital specialization for column-major output exchanges problem size and operand.
|
||||
template <
|
||||
/// Element type for A matrix operand
|
||||
typename ElementA_,
|
||||
/// Layout type for A matrix operand
|
||||
typename LayoutA_,
|
||||
/// Element type for B matrix operand
|
||||
typename ElementB_,
|
||||
/// Layout type for B matrix operand
|
||||
typename LayoutB_,
|
||||
/// Element type for C and D matrix operands
|
||||
typename ElementC_,
|
||||
/// Element type for internal accumulation
|
||||
typename ElementAccumulator_,
|
||||
/// Operator class tag
|
||||
typename OperatorClass_,
|
||||
/// Tag indicating architecture to tune for
|
||||
typename ArchTag_,
|
||||
/// Threadblock-level tile size (concept: GemmShape)
|
||||
typename ThreadblockShape_,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename WarpShape_,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename InstructionShape_,
|
||||
/// Epilogue output operator
|
||||
typename EpilogueOutputOp_,
|
||||
/// Threadblock-level swizzling operator
|
||||
typename ThreadblockSwizzle_,
|
||||
/// Number of stages used in the pipelined mainloop
|
||||
int Stages,
|
||||
/// Access granularity of A matrix in units of elements
|
||||
int AlignmentA,
|
||||
/// Access granularity of B matrix in units of elements
|
||||
int AlignmentB,
|
||||
typename Operator_
|
||||
>
|
||||
class GemmBatched<
|
||||
ElementA_,
|
||||
LayoutA_,
|
||||
ElementB_,
|
||||
LayoutB_,
|
||||
ElementC_,
|
||||
layout::ColumnMajor,
|
||||
ElementAccumulator_,
|
||||
OperatorClass_,
|
||||
ArchTag_,
|
||||
ThreadblockShape_,
|
||||
WarpShape_,
|
||||
InstructionShape_,
|
||||
EpilogueOutputOp_,
|
||||
ThreadblockSwizzle_,
|
||||
Stages,
|
||||
AlignmentA,
|
||||
AlignmentB,
|
||||
Operator_
|
||||
> {
|
||||
public:
|
||||
|
||||
using ElementA = ElementA_;
|
||||
using LayoutA = LayoutA_;
|
||||
using TensorRefA = TensorRef<ElementA const, LayoutA>;
|
||||
using ElementB = ElementB_;
|
||||
using LayoutB = LayoutB_;
|
||||
using TensorRefB = TensorRef<ElementB const, LayoutB>;
|
||||
using ElementC = ElementC_;
|
||||
using LayoutC = layout::ColumnMajor;
|
||||
using TensorRefC = TensorRef<ElementC const, LayoutC>;
|
||||
using TensorRefD = TensorRef<ElementC, LayoutC>;
|
||||
using ElementAccumulator = ElementAccumulator_;
|
||||
using OperatorClass = OperatorClass_;
|
||||
using ArchTag = ArchTag_;
|
||||
using ThreadblockShape = ThreadblockShape_;
|
||||
using WarpShape = WarpShape_;
|
||||
using InstructionShape = InstructionShape_;
|
||||
using EpilogueOutputOp = EpilogueOutputOp_;
|
||||
using ThreadblockSwizzle = ThreadblockSwizzle_;
|
||||
static int const kStages = Stages;
|
||||
|
||||
static int const kAlignmentA = AlignmentA;
|
||||
static int const kAlignmentB = AlignmentB;
|
||||
static int const kAlignmentC = EpilogueOutputOp::kCount;
|
||||
static bool const kSplitKSerial = false;
|
||||
|
||||
//
|
||||
using UnderlyingOperator = GemmBatched<
|
||||
ElementB,
|
||||
typename layout::LayoutTranspose<LayoutB>::type,
|
||||
ElementA,
|
||||
typename layout::LayoutTranspose<LayoutA>::type,
|
||||
ElementC,
|
||||
layout::RowMajor,
|
||||
ElementAccumulator,
|
||||
OperatorClass,
|
||||
ArchTag,
|
||||
ThreadblockShape,
|
||||
WarpShape,
|
||||
InstructionShape,
|
||||
EpilogueOutputOp,
|
||||
ThreadblockSwizzle,
|
||||
Stages,
|
||||
kAlignmentB,
|
||||
kAlignmentA
|
||||
>;
|
||||
|
||||
using UnderlyingArguments = typename UnderlyingOperator::Arguments;
|
||||
using GemmKernel = typename UnderlyingOperator::GemmKernel;
|
||||
|
||||
/// Argument structure
|
||||
struct Arguments {
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
GemmCoord problem_size;
|
||||
TensorRef<ElementA const, LayoutA> ref_A;
|
||||
int64_t stride_A;
|
||||
TensorRef<ElementB const, LayoutB> ref_B;
|
||||
int64_t stride_B;
|
||||
TensorRef<ElementC const, LayoutC> ref_C;
|
||||
int64_t stride_C;
|
||||
TensorRef<ElementC, LayoutC> ref_D;
|
||||
int64_t stride_D;
|
||||
typename EpilogueOutputOp::Params epilogue;
|
||||
int batch_count;
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
/// Default ctor
|
||||
CUTLASS_HOST_DEVICE
|
||||
Arguments() { }
|
||||
|
||||
/// Constructs an Arguments structure
|
||||
CUTLASS_HOST_DEVICE
|
||||
Arguments(
|
||||
GemmCoord problem_size_,
|
||||
TensorRef<ElementA const, LayoutA> ref_A_,
|
||||
int64_t stride_A_,
|
||||
TensorRef<ElementB const, LayoutB> ref_B_,
|
||||
int64_t stride_B_,
|
||||
TensorRef<ElementC const, LayoutC> ref_C_,
|
||||
int64_t stride_C_,
|
||||
TensorRef<ElementC, LayoutC> ref_D_,
|
||||
int64_t stride_D_,
|
||||
typename EpilogueOutputOp::Params epilogue_,
|
||||
int batch_count_
|
||||
):
|
||||
problem_size(problem_size_),
|
||||
ref_A(ref_A_),
|
||||
stride_A(stride_A_),
|
||||
ref_B(ref_B_),
|
||||
stride_B(stride_B_),
|
||||
ref_C(ref_C_),
|
||||
stride_C(stride_C_),
|
||||
ref_D(ref_D_),
|
||||
stride_D(stride_D_),
|
||||
epilogue(epilogue_),
|
||||
batch_count(batch_count_) { }
|
||||
};
|
||||
|
||||
private:
|
||||
|
||||
UnderlyingOperator underlying_operator_;
|
||||
|
||||
public:
|
||||
|
||||
/// Constructs the GEMM.
|
||||
GemmBatched() { }
|
||||
|
||||
/// Helper to construct a transposed equivalent for the underying GEMM operator
|
||||
static UnderlyingArguments to_underlying_arguments(Arguments const &args) {
|
||||
return UnderlyingArguments(
|
||||
{args.problem_size.n(), args.problem_size.m(), args.problem_size.k()},
|
||||
{args.ref_B.data(), args.ref_B.stride(0)},
|
||||
args.stride_B,
|
||||
{args.ref_A.data(), args.ref_A.stride(0)},
|
||||
args.stride_A,
|
||||
{args.ref_C.data(), args.ref_C.stride(0)},
|
||||
args.stride_C,
|
||||
{args.ref_D.data(), args.ref_D.stride(0)},
|
||||
args.stride_D,
|
||||
args.epilogue,
|
||||
args.batch_count
|
||||
);
|
||||
}
|
||||
|
||||
/// Determines whether the GEMM can execute the given problem.
|
||||
static Status can_implement(Arguments const &args) {
|
||||
|
||||
return UnderlyingOperator::can_implement(to_underlying_arguments(args));
|
||||
}
|
||||
|
||||
/// Gets the workspace size
|
||||
static size_t get_workspace_size(Arguments const &args) {
|
||||
|
||||
return UnderlyingOperator::get_workspace_size(to_underlying_arguments(args));
|
||||
}
|
||||
|
||||
/// Initializes GEMM state from arguments.
|
||||
Status initialize(Arguments const &args, void *workspace = nullptr, cudaStream_t stream = nullptr) {
|
||||
|
||||
return underlying_operator_.initialize(to_underlying_arguments(args), workspace);
|
||||
}
|
||||
|
||||
/// Lightweight update given a subset of arguments
|
||||
Status update(Arguments const &args, void *workspace = nullptr) {
|
||||
|
||||
return underlying_operator_.update(to_underlying_arguments(args), workspace);
|
||||
}
|
||||
|
||||
/// Runs the kernel using initialized state.
|
||||
Status run(cudaStream_t stream = nullptr) {
|
||||
|
||||
return underlying_operator_.run(stream);
|
||||
}
|
||||
|
||||
/// Runs the kernel using initialized state.
|
||||
Status operator()(cudaStream_t stream = nullptr) {
|
||||
return run(stream);
|
||||
}
|
||||
|
||||
/// Runs the kernel using initialized state.
|
||||
Status operator()(
|
||||
Arguments const &args,
|
||||
void *workspace = nullptr,
|
||||
cudaStream_t stream = nullptr) {
|
||||
|
||||
Status status = initialize(args, workspace, stream);
|
||||
|
||||
if (status == Status::kSuccess) {
|
||||
status = run(stream);
|
||||
}
|
||||
|
||||
return status;
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace device
|
||||
} // namespace gemm
|
||||
} // namespace cutlass
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
376
cat_files/gemm_universal.h
Normal file
376
cat_files/gemm_universal.h
Normal file
@@ -0,0 +1,376 @@
|
||||
/***************************************************************************************************
|
||||
* 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:
|
||||
* * Redistributions of source code must retain the above copyright notice, this list of
|
||||
* conditions and the following disclaimer.
|
||||
* * Redistributions in binary form must reproduce the above copyright notice, this list of
|
||||
* conditions and the following disclaimer in the documentation and/or other materials
|
||||
* provided with the distribution.
|
||||
* * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
|
||||
* to endorse or promote products derived from this software without specific prior written
|
||||
* permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
|
||||
* IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
|
||||
* FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
|
||||
* BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
|
||||
* OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
|
||||
* STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
/*! \file
|
||||
\brief
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/arch/arch.h"
|
||||
#include "cutlass/device_kernel.h"
|
||||
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
#include "cutlass/gemm/threadblock/threadblock_swizzle.h"
|
||||
#include "cutlass/gemm/kernel/gemm_universal.h"
|
||||
|
||||
#include "cutlass/gemm/kernel/default_gemm_universal.h"
|
||||
#include "cutlass/gemm/device/default_gemm_configuration.h"
|
||||
#include "cutlass/gemm/device/gemm_universal_base.h"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace gemm {
|
||||
namespace device {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/*!
|
||||
The universal GEMM accommodates serial reductions, parallel reductions, batched strided, and
|
||||
batched array variants.
|
||||
*/
|
||||
template <
|
||||
/// Element type for A matrix operand
|
||||
typename ElementA_,
|
||||
/// Layout type for A matrix operand
|
||||
typename LayoutA_,
|
||||
/// Element type for B matrix operand
|
||||
typename ElementB_,
|
||||
/// Layout type for B matrix operand
|
||||
typename LayoutB_,
|
||||
/// Element type for C and D matrix operands
|
||||
typename ElementC_,
|
||||
/// Layout type for C and D matrix operands
|
||||
typename LayoutC_,
|
||||
/// Element type for internal accumulation
|
||||
typename ElementAccumulator_ = ElementC_,
|
||||
/// Operator class tag
|
||||
typename OperatorClass_ = arch::OpClassSimt,
|
||||
/// Tag indicating architecture to tune for. This is the minimum SM that
|
||||
/// supports the intended feature. The device kernel can be built
|
||||
/// targeting any SM larger than this number.
|
||||
typename ArchTag_ = arch::Sm61,
|
||||
/// Threadblock-level tile size (concept: GemmShape)
|
||||
typename ThreadblockShape_ = typename DefaultGemmConfiguration<
|
||||
OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
|
||||
ElementAccumulator_>::ThreadblockShape,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename WarpShape_ = typename DefaultGemmConfiguration<
|
||||
OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
|
||||
ElementAccumulator_>::WarpShape,
|
||||
/// Instruction-level tile size (concept: GemmShape)
|
||||
typename InstructionShape_ = typename DefaultGemmConfiguration<
|
||||
OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
|
||||
ElementAccumulator_>::InstructionShape,
|
||||
/// Epilogue output operator
|
||||
typename EpilogueOutputOp_ = typename DefaultGemmConfiguration<
|
||||
OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
|
||||
ElementAccumulator_>::EpilogueOutputOp,
|
||||
/// Threadblock-level swizzling operator
|
||||
typename ThreadblockSwizzle_ = threadblock::GemmIdentityThreadblockSwizzle<>,
|
||||
/// Number of stages used in the pipelined mainloop
|
||||
int Stages =
|
||||
DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
|
||||
ElementC_, ElementAccumulator_>::kStages,
|
||||
/// Access granularity of A matrix in units of elements
|
||||
int AlignmentA =
|
||||
DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
|
||||
ElementC_, ElementAccumulator_>::kAlignmentA,
|
||||
/// Access granularity of B matrix in units of elements
|
||||
int AlignmentB =
|
||||
DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
|
||||
ElementC_, ElementAccumulator_>::kAlignmentB,
|
||||
/// Operation performed by GEMM
|
||||
typename Operator_ = typename DefaultGemmConfiguration<
|
||||
OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
|
||||
ElementAccumulator_>::Operator,
|
||||
/// Complex elementwise transformation on A operand
|
||||
ComplexTransform TransformA = ComplexTransform::kNone,
|
||||
/// Complex elementwise transformation on B operand
|
||||
ComplexTransform TransformB = ComplexTransform::kNone
|
||||
>
|
||||
class GemmUniversal :
|
||||
GemmUniversalBase<
|
||||
typename kernel::DefaultGemmUniversal<
|
||||
ElementA_,
|
||||
LayoutA_,
|
||||
TransformA,
|
||||
AlignmentA,
|
||||
ElementB_,
|
||||
LayoutB_,
|
||||
TransformB,
|
||||
AlignmentB,
|
||||
ElementC_,
|
||||
LayoutC_,
|
||||
ElementAccumulator_,
|
||||
OperatorClass_,
|
||||
ArchTag_,
|
||||
ThreadblockShape_,
|
||||
WarpShape_,
|
||||
InstructionShape_,
|
||||
EpilogueOutputOp_,
|
||||
ThreadblockSwizzle_,
|
||||
Stages,
|
||||
Operator_
|
||||
>::GemmKernel
|
||||
> {
|
||||
|
||||
public:
|
||||
|
||||
using ElementAccumulator = ElementAccumulator_;
|
||||
using OperatorClass = OperatorClass_;
|
||||
using ArchTag = ArchTag_;
|
||||
using ThreadblockShape = ThreadblockShape_;
|
||||
using WarpShape = WarpShape_;
|
||||
using InstructionShape = InstructionShape_;
|
||||
using EpilogueOutputOp = EpilogueOutputOp_;
|
||||
using ThreadblockSwizzle = ThreadblockSwizzle_;
|
||||
using Operator = Operator_;
|
||||
static int const kStages = Stages;
|
||||
static int const kAlignmentA = AlignmentA;
|
||||
static int const kAlignmentB = AlignmentB;
|
||||
static int const kAlignmentC = EpilogueOutputOp::kCount;
|
||||
static ComplexTransform const kTransformA = TransformA;
|
||||
static ComplexTransform const kTransformB = TransformB;
|
||||
|
||||
using Base = GemmUniversalBase<
|
||||
typename kernel::DefaultGemmUniversal<
|
||||
ElementA_,
|
||||
LayoutA_,
|
||||
TransformA,
|
||||
AlignmentA,
|
||||
ElementB_,
|
||||
LayoutB_,
|
||||
TransformB,
|
||||
AlignmentB,
|
||||
ElementC_,
|
||||
LayoutC_,
|
||||
ElementAccumulator_,
|
||||
OperatorClass_,
|
||||
ArchTag_,
|
||||
ThreadblockShape_,
|
||||
WarpShape_,
|
||||
InstructionShape_,
|
||||
EpilogueOutputOp_,
|
||||
ThreadblockSwizzle_,
|
||||
Stages,
|
||||
Operator_
|
||||
>::GemmKernel
|
||||
>;
|
||||
|
||||
using Arguments = typename Base::Arguments;
|
||||
using GemmKernel = typename Base::GemmKernel;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Parital specialization for column-major output exchanges problem size and operand.
|
||||
template <
|
||||
/// Element type for A matrix operand
|
||||
typename ElementA_,
|
||||
/// Layout type for A matrix operand
|
||||
typename LayoutA_,
|
||||
/// Element type for B matrix operand
|
||||
typename ElementB_,
|
||||
/// Layout type for B matrix operand
|
||||
typename LayoutB_,
|
||||
/// Element type for C and D matrix operands
|
||||
typename ElementC_,
|
||||
/// Element type for internal accumulation
|
||||
typename ElementAccumulator_,
|
||||
/// Operator class tag
|
||||
typename OperatorClass_,
|
||||
/// Tag indicating architecture to tune for. This is the minimum SM that
|
||||
/// supports the intended feature. The device kernel can be built
|
||||
/// targeting any SM larger than this number.
|
||||
typename ArchTag_,
|
||||
/// Threadblock-level tile size (concept: GemmShape)
|
||||
typename ThreadblockShape_,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename WarpShape_,
|
||||
/// Instruction-level tile size (concept: GemmShape)
|
||||
typename InstructionShape_,
|
||||
/// Epilogue output operator
|
||||
typename EpilogueOutputOp_,
|
||||
/// Threadblock-level swizzling operator
|
||||
typename ThreadblockSwizzle_,
|
||||
/// Number of stages used in the pipelined mainloop
|
||||
int Stages,
|
||||
/// Access granularity of A matrix in units of elements
|
||||
int AlignmentA,
|
||||
/// Access granularity of B matrix in units of elements
|
||||
int AlignmentB,
|
||||
/// Operation performed by GEMM
|
||||
typename Operator_,
|
||||
/// Complex elementwise transformation on A operand
|
||||
ComplexTransform TransformA,
|
||||
/// Complex elementwise transformation on B operand
|
||||
ComplexTransform TransformB>
|
||||
class GemmUniversal<ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_,
|
||||
layout::ColumnMajor, // partially specialized on LayoutC
|
||||
ElementAccumulator_, OperatorClass_, ArchTag_, ThreadblockShape_,
|
||||
WarpShape_, InstructionShape_, EpilogueOutputOp_,
|
||||
ThreadblockSwizzle_, Stages, AlignmentA, AlignmentB,
|
||||
Operator_, TransformA, TransformB> {
|
||||
public:
|
||||
|
||||
using ElementA = ElementA_;
|
||||
using LayoutA = LayoutA_;
|
||||
using TensorRefA = TensorRef<ElementA const, LayoutA>;
|
||||
using ElementB = ElementB_;
|
||||
using LayoutB = LayoutB_;
|
||||
using TensorRefB = TensorRef<ElementB const, LayoutB>;
|
||||
using ElementC = ElementC_;
|
||||
using LayoutC = layout::ColumnMajor;
|
||||
using TensorRefC = TensorRef<ElementC const, LayoutC>;
|
||||
using TensorRefD = TensorRef<ElementC, LayoutC>;
|
||||
using ElementAccumulator = ElementAccumulator_;
|
||||
using OperatorClass = OperatorClass_;
|
||||
using ArchTag = ArchTag_;
|
||||
using ThreadblockShape = ThreadblockShape_;
|
||||
using WarpShape = WarpShape_;
|
||||
using InstructionShape = InstructionShape_;
|
||||
using EpilogueOutputOp = EpilogueOutputOp_;
|
||||
using ThreadblockSwizzle = ThreadblockSwizzle_;
|
||||
using Operator = Operator_;
|
||||
static int const kStages = Stages;
|
||||
static int const kAlignmentA = AlignmentA;
|
||||
static int const kAlignmentB = AlignmentB;
|
||||
static ComplexTransform const kTransformA = TransformA;
|
||||
static ComplexTransform const kTransformB = TransformB;
|
||||
|
||||
using UnderlyingOperator = typename GemmUniversal<
|
||||
ElementB,
|
||||
typename layout::LayoutTranspose<LayoutB>::type,
|
||||
ElementA,
|
||||
typename layout::LayoutTranspose<LayoutA>::type,
|
||||
ElementC,
|
||||
layout::RowMajor,
|
||||
ElementAccumulator,
|
||||
OperatorClass,
|
||||
ArchTag,
|
||||
ThreadblockShape,
|
||||
WarpShape,
|
||||
InstructionShape,
|
||||
EpilogueOutputOp,
|
||||
ThreadblockSwizzle,
|
||||
Stages,
|
||||
kAlignmentB,
|
||||
kAlignmentA,
|
||||
Operator,
|
||||
kTransformB,
|
||||
kTransformA
|
||||
>::Base;
|
||||
|
||||
using GemmKernel = typename UnderlyingOperator::GemmKernel;
|
||||
static int const kAlignmentC = EpilogueOutputOp::kCount;
|
||||
|
||||
/// Argument structure
|
||||
using Arguments = typename UnderlyingOperator::Arguments;
|
||||
|
||||
private:
|
||||
|
||||
UnderlyingOperator underlying_operator_;
|
||||
|
||||
public:
|
||||
|
||||
/// Constructs the GEMM.
|
||||
GemmUniversal() { }
|
||||
|
||||
/// Helper to construct a transposed equivalent for the underying GEMM operator
|
||||
static Arguments to_underlying_arguments(Arguments const &args) {
|
||||
return args.transposed_problem();
|
||||
}
|
||||
|
||||
/// Determines whether the GEMM can execute the given problem.
|
||||
static Status can_implement(Arguments const &args) {
|
||||
|
||||
return UnderlyingOperator::can_implement(to_underlying_arguments(args));
|
||||
}
|
||||
|
||||
/// Gets the workspace size
|
||||
static size_t get_workspace_size(Arguments const &args) {
|
||||
|
||||
return UnderlyingOperator::get_workspace_size(to_underlying_arguments(args));
|
||||
}
|
||||
|
||||
/// Computes the grid shape
|
||||
static dim3 get_grid_shape(Arguments const &args) {
|
||||
return UnderlyingOperator::get_grid_shape(to_underlying_arguments(args));
|
||||
}
|
||||
|
||||
/// Computes the maximum number of active blocks per multiprocessor
|
||||
static int maximum_active_blocks(int smem_capacity = -1) {
|
||||
return UnderlyingOperator::maximum_active_blocks(smem_capacity);
|
||||
}
|
||||
|
||||
/// Initializes GEMM state from arguments.
|
||||
Status initialize(Arguments const &args, void *workspace = nullptr, cudaStream_t stream = nullptr) {
|
||||
|
||||
return underlying_operator_.initialize(to_underlying_arguments(args), workspace, stream);
|
||||
}
|
||||
|
||||
/// Lightweight update given a subset of arguments
|
||||
Status update(Arguments const &args, void *workspace = nullptr) {
|
||||
|
||||
return underlying_operator_.update(to_underlying_arguments(args), workspace);
|
||||
}
|
||||
|
||||
/// Runs the kernel using initialized state.
|
||||
Status run(cudaStream_t stream = nullptr) {
|
||||
|
||||
return underlying_operator_.run(stream);
|
||||
}
|
||||
|
||||
/// Runs the kernel using initialized state.
|
||||
Status operator()(cudaStream_t stream = nullptr) {
|
||||
return run(stream);
|
||||
}
|
||||
|
||||
/// Runs the kernel using initialized state.
|
||||
Status operator()(
|
||||
Arguments const &args,
|
||||
void *workspace = nullptr,
|
||||
cudaStream_t stream = nullptr) {
|
||||
|
||||
Status status = initialize(args, workspace, stream);
|
||||
|
||||
if (status == Status::kSuccess) {
|
||||
status = run(stream);
|
||||
}
|
||||
|
||||
return status;
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace device
|
||||
} // namespace gemm
|
||||
} // namespace cutlass
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
1238
cat_files/iluvatar_mma.hpp
Normal file
1238
cat_files/iluvatar_mma.hpp
Normal file
File diff suppressed because it is too large
Load Diff
4058
cat_files/ixinfer.h
Normal file
4058
cat_files/ixinfer.h
Normal file
File diff suppressed because it is too large
Load Diff
394
cat_files/mma_cu10.h
Normal file
394
cat_files/mma_cu10.h
Normal file
@@ -0,0 +1,394 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2021 Iluvatar CoreX. All rights reserved.
|
||||
* Copyright Declaration: This software, including all of its code and documentation,
|
||||
* except for the third-party software it contains, is a copyrighted work of Shanghai Iluvatar CoreX
|
||||
* Semiconductor Co., Ltd. and its affiliates ("Iluvatar CoreX") in accordance with the PRC Copyright
|
||||
* Law and relevant international treaties, and all rights contained therein are enjoyed by Iluvatar
|
||||
* CoreX. No user of this software shall have any right, ownership or interest in this software and
|
||||
* any use of this software shall be in compliance with the terms and conditions of the End User
|
||||
* License Agreement.
|
||||
**************************************************************************************************/
|
||||
|
||||
/*! \file
|
||||
\brief Matrix Multiply for BigIsland 1st generation
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/arch/mma.h"
|
||||
|
||||
#include "cutlass/layout/matrix.h"
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace arch {
|
||||
|
||||
/// BigIsland Tensor Core tile format - EM orinted vector type definitions
|
||||
/// fp32
|
||||
typedef float v4float_t __attribute__((ext_vector_type(4)));
|
||||
/// s32
|
||||
typedef int32_t v4int32_t __attribute__((ext_vector_type(4)));
|
||||
/// u32
|
||||
typedef uint32_t v4uint32_t __attribute__((ext_vector_type(4)));
|
||||
/// fp16
|
||||
typedef uint16_t v4half_t __attribute__((ext_vector_type(4)));
|
||||
/// bf16
|
||||
typedef uint16_t v4bfloat16_t __attribute__((ext_vector_type(4)));
|
||||
/// s8
|
||||
typedef int8_t v4int8_t __attribute__((ext_vector_type(4)));
|
||||
/// u8
|
||||
typedef uint8_t v4uint8_t __attribute__((ext_vector_type(4)));
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
//
|
||||
// Matrix multiply accumulate 161616 - U32 accumulation
|
||||
//
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Matrix multiply-add operation: U32 = U8 * U8 + U32
|
||||
template <typename LayoutA, typename LayoutB, typename LayoutC>
|
||||
struct Mma<
|
||||
gemm::GemmShape<16, 16, 16>,
|
||||
64,
|
||||
uint8_t,
|
||||
LayoutA,
|
||||
uint8_t,
|
||||
LayoutB,
|
||||
uint32_t,
|
||||
LayoutC,
|
||||
OpMultiplyAdd> {
|
||||
|
||||
using Shape = gemm::GemmShape<16, 16, 16>;
|
||||
|
||||
using ElementA = uint8_t;
|
||||
using FragmentA = Array<uint8_t, 4>;
|
||||
|
||||
using ElementB = uint8_t;
|
||||
using FragmentB = Array<uint8_t, 4>;
|
||||
|
||||
using ElementC = uint;
|
||||
using FragmentC = Array<uint, 4>;
|
||||
|
||||
using Operator = OpMultiplyAdd;
|
||||
using ArchTag = arch::Cu10;
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
void operator()(
|
||||
FragmentC &d,
|
||||
FragmentA const &a,
|
||||
FragmentB const &b,
|
||||
FragmentC const &c
|
||||
) const {
|
||||
#if CUTLASS_ARCH_CU10_SUPPORTED
|
||||
v4uint8_t src_A;
|
||||
v4uint8_t src_B;
|
||||
v4uint32_t src_C;
|
||||
v4uint32_t dst_D;
|
||||
|
||||
src_A[0] = a[0];
|
||||
src_A[1] = a[1];
|
||||
src_A[2] = a[2];
|
||||
src_A[3] = a[3];
|
||||
src_B[0] = b[0];
|
||||
src_B[1] = b[1];
|
||||
src_B[2] = b[2];
|
||||
src_B[3] = b[3];
|
||||
src_C[0] = c[0];
|
||||
src_C[1] = c[1];
|
||||
src_C[2] = c[2];
|
||||
src_C[3] = c[3];
|
||||
|
||||
dst_D = __ivcorex_matrix_mad_u32x4_u8x4(src_A, src_B, src_C);
|
||||
|
||||
d[0] = dst_D[0];
|
||||
d[1] = dst_D[1];
|
||||
d[2] = dst_D[2];
|
||||
d[3] = dst_D[3];
|
||||
#else
|
||||
assert(0);
|
||||
#endif
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
//
|
||||
// Matrix multiply accumulate 161616 - S32 accumulation
|
||||
//
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Matrix multiply-add operation: S32 = S8 * S8 + S32
|
||||
template <typename LayoutA, typename LayoutB, typename LayoutC>
|
||||
struct Mma<
|
||||
gemm::GemmShape<16, 16, 16>,
|
||||
64,
|
||||
int8_t,
|
||||
LayoutA,
|
||||
int8_t,
|
||||
LayoutB,
|
||||
int,
|
||||
LayoutC,
|
||||
OpMultiplyAdd> {
|
||||
|
||||
using Shape = gemm::GemmShape<16, 16, 16>;
|
||||
|
||||
using ElementA = int8_t;
|
||||
using FragmentA = Array<int8_t, 4>;
|
||||
|
||||
using ElementB = int8_t;
|
||||
using FragmentB = Array<int8_t, 4>;
|
||||
|
||||
using ElementC = int;
|
||||
using FragmentC = Array<int, 4>;
|
||||
|
||||
using Operator = OpMultiplyAdd;
|
||||
using ArchTag = arch::Cu10;
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
void operator()(
|
||||
FragmentC &d,
|
||||
FragmentA const &a,
|
||||
FragmentB const &b,
|
||||
FragmentC const &c
|
||||
) const {
|
||||
#if CUTLASS_ARCH_CU10_SUPPORTED
|
||||
v4int8_t src_A;
|
||||
v4int8_t src_B;
|
||||
v4int32_t src_C;
|
||||
v4int32_t dst_D;
|
||||
|
||||
src_A[0] = a[0];
|
||||
src_A[1] = a[1];
|
||||
src_A[2] = a[2];
|
||||
src_A[3] = a[3];
|
||||
src_B[0] = b[0];
|
||||
src_B[1] = b[1];
|
||||
src_B[2] = b[2];
|
||||
src_B[3] = b[3];
|
||||
src_C[0] = c[0];
|
||||
src_C[1] = c[1];
|
||||
src_C[2] = c[2];
|
||||
src_C[3] = c[3];
|
||||
|
||||
dst_D = __ivcorex_matrix_mad_i32x4_i8x4(src_A, src_B, src_C);
|
||||
|
||||
d[0] = dst_D[0];
|
||||
d[1] = dst_D[1];
|
||||
d[2] = dst_D[2];
|
||||
d[3] = dst_D[3];
|
||||
#else
|
||||
assert(0);
|
||||
#endif
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
//
|
||||
// Matrix multiply accumulate 161616 - FP32 accumulation
|
||||
//
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Matrix multiply-add operation: FP32 = FP16 * FP16 + FP32
|
||||
template <typename LayoutA, typename LayoutB, typename LayoutC>
|
||||
struct Mma<
|
||||
gemm::GemmShape<16, 16, 16>,
|
||||
64,
|
||||
cutlass::half_t,
|
||||
LayoutA,
|
||||
cutlass::half_t,
|
||||
LayoutB,
|
||||
float,
|
||||
LayoutC,
|
||||
OpMultiplyAdd> {
|
||||
|
||||
using Shape = gemm::GemmShape<16, 16, 16>;
|
||||
|
||||
using ElementA = cutlass::half_t;
|
||||
using FragmentA = Array<half_t, 4>;
|
||||
|
||||
using ElementB = cutlass::half_t;
|
||||
using FragmentB = Array<half_t, 4>;
|
||||
|
||||
using ElementC = float;
|
||||
using FragmentC = Array<float, 4>;
|
||||
|
||||
using Operator = OpMultiplyAdd;
|
||||
using ArchTag = arch::Cu10;
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
void operator()(
|
||||
FragmentC &d,
|
||||
FragmentA const &a,
|
||||
FragmentB const &b,
|
||||
FragmentC const &c
|
||||
) const {
|
||||
v4half_t src_A;
|
||||
v4half_t src_B;
|
||||
v4float_t src_C;
|
||||
v4float_t dst_D;
|
||||
|
||||
src_A[0] = half_t(a[0]).storage;
|
||||
src_A[1] = half_t(a[1]).storage;
|
||||
src_A[2] = half_t(a[2]).storage;
|
||||
src_A[3] = half_t(a[3]).storage;
|
||||
src_B[0] = half_t(b[0]).storage;
|
||||
src_B[1] = half_t(b[1]).storage;
|
||||
src_B[2] = half_t(b[2]).storage;
|
||||
src_B[3] = half_t(b[3]).storage;
|
||||
src_C[0] = c[0];
|
||||
src_C[1] = c[1];
|
||||
src_C[2] = c[2];
|
||||
src_C[3] = c[3];
|
||||
|
||||
dst_D = __ivcorex_matrix_mad_f32x4_f16x4(src_A, src_B, src_C);
|
||||
#if 0
|
||||
if(threadIdx.x == 0)
|
||||
printf(
|
||||
">>> After\n"
|
||||
"A: %f, %f, %f, %f\n"
|
||||
"B: %f, %f, %f, %f\n"
|
||||
"C: %f, %f, %f, %f\n"
|
||||
"D: %f, %f, %f, %f\n\n",
|
||||
float(a[0]), float(a[1]), float(a[2]), float(a[3]),
|
||||
float(b[0]), float(b[1]), float(b[2]), float(b[3]),
|
||||
float(src_C[0]), float(src_C[1]), float(src_C[2]), float(src_C[3]),
|
||||
float(d[0]), float(d[1]), float(d[2]), float(d[3])
|
||||
);
|
||||
#endif
|
||||
|
||||
d[0] = dst_D[0];
|
||||
d[1] = dst_D[1];
|
||||
d[2] = dst_D[2];
|
||||
d[3] = dst_D[3];
|
||||
|
||||
}
|
||||
};
|
||||
|
||||
/// Matrix multiply-add operation: FP32 = BF16 * BF16 + FP32
|
||||
template <typename LayoutA, typename LayoutB, typename LayoutC>
|
||||
struct Mma<
|
||||
gemm::GemmShape<16, 16, 16>,
|
||||
64,
|
||||
bfloat16_t,
|
||||
LayoutA,
|
||||
bfloat16_t,
|
||||
LayoutB,
|
||||
float,
|
||||
LayoutC,
|
||||
OpMultiplyAdd> {
|
||||
|
||||
using Shape = gemm::GemmShape<16, 16, 16>;
|
||||
|
||||
using ElementA = bfloat16_t;
|
||||
using FragmentA = Array<bfloat16_t, 4>;
|
||||
|
||||
using ElementB = bfloat16_t;
|
||||
using FragmentB = Array<bfloat16_t, 4>;
|
||||
|
||||
using ElementC = float;
|
||||
using FragmentC = Array<float, 4>;
|
||||
|
||||
using Operator = OpMultiplyAdd;
|
||||
using ArchTag = arch::Cu10;
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
void operator()(
|
||||
FragmentC &d,
|
||||
FragmentA const &a,
|
||||
FragmentB const &b,
|
||||
FragmentC const &c
|
||||
) const {
|
||||
v4bfloat16_t src_A;
|
||||
v4bfloat16_t src_B;
|
||||
v4float_t src_C;
|
||||
v4float_t dst_D;
|
||||
|
||||
src_A[0] = bfloat16_t(a[0]).storage;
|
||||
src_A[1] = bfloat16_t(a[1]).storage;
|
||||
src_A[2] = bfloat16_t(a[2]).storage;
|
||||
src_A[3] = bfloat16_t(a[3]).storage;
|
||||
src_B[0] = bfloat16_t(b[0]).storage;
|
||||
src_B[1] = bfloat16_t(b[1]).storage;
|
||||
src_B[2] = bfloat16_t(b[2]).storage;
|
||||
src_B[3] = bfloat16_t(b[3]).storage;
|
||||
src_C[0] = c[0];
|
||||
src_C[1] = c[1];
|
||||
src_C[2] = c[2];
|
||||
src_C[3] = c[3];
|
||||
#if __clang_major__ >= 16
|
||||
dst_D = __ivcorex_matrix_mad_f32x4_bf16x4(src_A, src_B, src_C);
|
||||
#else
|
||||
dst_D = __ivcorex_matrix_mad_f32_bf16(src_A, src_B, src_C);
|
||||
#endif
|
||||
d[0] = dst_D[0];
|
||||
d[1] = dst_D[1];
|
||||
d[2] = dst_D[2];
|
||||
d[3] = dst_D[3];
|
||||
}
|
||||
};
|
||||
|
||||
/// Matrix multiply-add operation: FP32 = FP32 * FP32 + FP32
|
||||
template <typename LayoutA, typename LayoutB, typename LayoutC>
|
||||
struct Mma<
|
||||
gemm::GemmShape<16,16,16>,
|
||||
64,
|
||||
float,
|
||||
LayoutA,
|
||||
float,
|
||||
LayoutB,
|
||||
float,
|
||||
LayoutC,
|
||||
OpMultiplyAdd> {
|
||||
|
||||
using Shape = gemm::GemmShape<16,16,16>;
|
||||
|
||||
using ElementA = float;
|
||||
using FragmentA = Array<float, 4>;
|
||||
|
||||
using ElementB = float;
|
||||
using FragmentB = Array<float, 4>;
|
||||
|
||||
using ElementC = float;
|
||||
using FragmentC = Array<float, 4>;
|
||||
|
||||
using Operator = OpMultiplyAdd;
|
||||
using ArchTag = arch::Cu10;
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
void operator()(
|
||||
FragmentC &d,
|
||||
FragmentA const &a,
|
||||
FragmentB const &b,
|
||||
FragmentC const &c
|
||||
) const {
|
||||
v4float_t src_A;
|
||||
v4float_t src_B;
|
||||
v4float_t src_C;
|
||||
v4float_t dst_D;
|
||||
|
||||
src_A[0] = a[0];
|
||||
src_A[1] = a[1];
|
||||
src_A[2] = a[2];
|
||||
src_A[3] = a[3];
|
||||
src_B[0] = b[0];
|
||||
src_B[1] = b[1];
|
||||
src_B[2] = b[2];
|
||||
src_B[3] = b[3];
|
||||
src_C[0] = c[0];
|
||||
src_C[1] = c[1];
|
||||
src_C[2] = c[2];
|
||||
src_C[3] = c[3];
|
||||
|
||||
dst_D = __ivcorex_matrix_mad_f32x4_f32x4(src_A, src_B, src_C);
|
||||
|
||||
d[0] = dst_D[0];
|
||||
d[1] = dst_D[1];
|
||||
d[2] = dst_D[2];
|
||||
d[3] = dst_D[3];
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
}
|
||||
}
|
||||
382
cat_files/mma_tensor_op.h
Normal file
382
cat_files/mma_tensor_op.h
Normal file
@@ -0,0 +1,382 @@
|
||||
/***************************************************************************************************
|
||||
* 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:
|
||||
* * 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.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2021 Iluvatar CoreX. All rights reserved.
|
||||
* Copyright Declaration: This software, including all of its code and documentation,
|
||||
* except for the third-party software it contains, is a copyrighted work of Shanghai Iluvatar CoreX
|
||||
* Semiconductor Co., Ltd. and its affiliates ("Iluvatar CoreX") in accordance with the PRC Copyright
|
||||
* Law and relevant international treaties, and all rights contained therein are enjoyed by Iluvatar
|
||||
* CoreX. No user of this software shall have any right, ownership or interest in this software and
|
||||
* any use of this software shall be in compliance with the terms and conditions of the End User
|
||||
* License Agreement.
|
||||
**************************************************************************************************/
|
||||
|
||||
/*! \file
|
||||
\brief Templates implementing warp-level matrix multiply-accumulate operations targeting
|
||||
Tensor Cores.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/platform/platform.h"
|
||||
|
||||
#include "cutlass/numeric_conversion.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/matrix_shape.h"
|
||||
|
||||
#include "cutlass/arch/mma.h"
|
||||
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
#include "cutlass/gemm/warp/mma.h"
|
||||
|
||||
#include "cutlass/gemm/warp/mma_tensor_op_policy.h"
|
||||
#include "cutlass/gemm/warp/mma_tensor_op_tile_iterator.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace gemm {
|
||||
namespace warp {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace detail {
|
||||
|
||||
template <typename T, typename S, int N, FloatRoundStyle Round>
|
||||
struct ConvertAndPack {
|
||||
|
||||
using Converter = NumericArrayConverter<T, S, N, Round>;
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Array<T, N> operator()(Array<S, N> const &source) {
|
||||
Converter converter;
|
||||
|
||||
return converter(source);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T, int N, FloatRoundStyle Round>
|
||||
struct ConvertAndPack<T, T, N, Round> {
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Array<T, N> operator()(Array<T, N> const &source) {
|
||||
return source;
|
||||
}
|
||||
};
|
||||
|
||||
template <int N, FloatRoundStyle Round>
|
||||
struct ConvertAndPack<bfloat16_t, float, N, Round> {
|
||||
|
||||
using Converter = NumericArrayConverter<bfloat16_t, float, N, Round>;
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Array<bfloat16_t, N> operator()(Array<float, N> const &source) {
|
||||
Converter converter;
|
||||
|
||||
Array<float, N> tmp;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < N; ++i) {
|
||||
int idx = (((i << 1) & 2) | ((i >> 1) & 1) | (i & 0xfffffffc));
|
||||
tmp[i] = source[idx];
|
||||
}
|
||||
|
||||
return converter(tmp);
|
||||
}
|
||||
};
|
||||
|
||||
template <int N, FloatRoundStyle Round>
|
||||
struct ConvertAndPack<half_t, float, N, Round> {
|
||||
|
||||
using Converter = NumericArrayConverter<half_t, float, N, Round>;
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Array<half_t, N> operator()(Array<float, N> const &source) {
|
||||
Converter converter;
|
||||
|
||||
Array<float, N> tmp;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < N; ++i) {
|
||||
int idx = (((i << 1) & 2) | ((i >> 1) & 1) | (i & 0xfffffffc));
|
||||
tmp[i] = source[idx];
|
||||
}
|
||||
|
||||
return converter(tmp);
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace detail
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Structure to compute the matrix product targeting CUDA cores and SIMT math instructions.
|
||||
template <
|
||||
/// Size of the Gemm problem - concept: gemm::GemmShape<>
|
||||
typename Shape_,
|
||||
/// Data type of A elements
|
||||
typename ElementA_,
|
||||
/// Layout of A matrix (concept: MatrixLayout)
|
||||
typename LayoutA_,
|
||||
/// Data type of B elements
|
||||
typename ElementB_,
|
||||
/// Layout of B matrix (concept: MatrixLayout)
|
||||
typename LayoutB_,
|
||||
/// Element type of C matrix
|
||||
typename ElementC_,
|
||||
/// Layout of C matrix (concept: MatrixLayout)
|
||||
typename LayoutC_,
|
||||
/// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy)
|
||||
typename Policy_,
|
||||
/// Number of partitions along K dimension
|
||||
int PartitionsK_ = 1,
|
||||
/// Store the accumulators in row major or column major.
|
||||
/// Iluvatar Tensor Core always stores accumulators in row major
|
||||
bool AccumulatorsInRowMajor = true,
|
||||
/// Used for partial specialization
|
||||
typename Enable = bool
|
||||
>
|
||||
class MmaTensorOp {
|
||||
public:
|
||||
/// Shape of warp-level matrix operation (concept: GemmShape)
|
||||
using Shape = Shape_;
|
||||
|
||||
/// Data type of multiplicand A
|
||||
using ElementA = ElementA_;
|
||||
|
||||
/// Layout of multiplicand A
|
||||
using LayoutA = LayoutA_;
|
||||
|
||||
/// Data type of multiplicand B
|
||||
using ElementB = ElementB_;
|
||||
|
||||
/// Layout of multiplicand B
|
||||
using LayoutB = LayoutB_;
|
||||
|
||||
/// Data type of accumulator matrix C
|
||||
using ElementC = ElementC_;
|
||||
|
||||
/// Layout of accumulator matrix C
|
||||
using LayoutC = LayoutC_;
|
||||
|
||||
/// Shape of the warp in units of thread (concept: MmaLanePolicySimt)
|
||||
using Policy = Policy_;
|
||||
|
||||
/// Underlying matrix multiply operator (concept: arch::Mma)
|
||||
using ArchMmaOperator = typename Policy::Operator;
|
||||
|
||||
/// Architecture tag from underlying instruction
|
||||
using ArchTag = typename ArchMmaOperator::ArchTag;
|
||||
|
||||
/// Indicates class of matrix operator
|
||||
using OperatorClass = arch::OpClassTensorOp;
|
||||
|
||||
/// Shape of underlying instruction
|
||||
using InstructionShape = typename ArchMmaOperator::Shape;
|
||||
|
||||
/// Complex transform on A operand
|
||||
static ComplexTransform const kTransformA = ComplexTransform::kNone;
|
||||
|
||||
/// Complex transform on B operand
|
||||
static ComplexTransform const kTransformB = ComplexTransform::kNone;
|
||||
|
||||
/// Number of threads participating in warp-level matrix product
|
||||
static int const kThreadCount = NUM_THREADS_PER_WARP;
|
||||
|
||||
/// Number of partitions along K dimension
|
||||
static int const kPartitionsK = PartitionsK_;
|
||||
|
||||
public:
|
||||
/// FIXME(Peter Han): workaround to adapt to simt epilogue, need to remove
|
||||
struct ThreadMma {
|
||||
using ElementC = ElementC;
|
||||
};
|
||||
|
||||
/// Iterates over the A operand in memory
|
||||
using IteratorA = MmaTensorOpMultiplicandTileIterator<
|
||||
MatrixShape<Shape::kM, Policy::Operator::Shape::kK>,
|
||||
Operand::kA,
|
||||
ElementA,
|
||||
LayoutA,
|
||||
InstructionShape,
|
||||
kThreadCount,
|
||||
kPartitionsK>;
|
||||
|
||||
/// Storage for A tile
|
||||
using FragmentA = typename IteratorA::Fragment;
|
||||
|
||||
/// Storage for transformed A tile
|
||||
using TransformedFragmentA =
|
||||
Array<typename ArchMmaOperator::ElementA, FragmentA::kElements>;
|
||||
|
||||
/// Iterates over the B operand in memory
|
||||
using IteratorB = MmaTensorOpMultiplicandTileIterator<
|
||||
MatrixShape<Policy::Operator::Shape::kK, Shape::kN>,
|
||||
Operand::kB,
|
||||
ElementB,
|
||||
LayoutB,
|
||||
InstructionShape,
|
||||
kThreadCount,
|
||||
kPartitionsK>;
|
||||
|
||||
/// Storage for B tile
|
||||
using FragmentB = typename IteratorB::Fragment;
|
||||
|
||||
/// Storage for transformed B tile
|
||||
using TransformedFragmentB =
|
||||
Array<typename ArchMmaOperator::ElementB, FragmentB::kElements>;
|
||||
|
||||
/// Iterates over the C operand in memory
|
||||
using IteratorC = MmaTensorOpAccumulatorTileIterator<
|
||||
MatrixShape<Shape::kM, Shape::kN>,
|
||||
ElementC,
|
||||
LayoutC,
|
||||
InstructionShape>;
|
||||
|
||||
/// Storage for C tile
|
||||
using FragmentC = typename IteratorC::Fragment;
|
||||
|
||||
static_assert(
|
||||
!(Shape::kM % Policy::Operator::Shape::kM) &&
|
||||
!(Shape::kN % Policy::Operator::Shape::kN) &&
|
||||
!(Shape::kK % Policy::Operator::Shape::kK),
|
||||
"Shape of warp-level Mma must be divisible by operator shape.");
|
||||
|
||||
using MmaIterations = gemm::GemmShape<
|
||||
(Shape::kM + ArchMmaOperator::Shape::kM - 1) / ArchMmaOperator::Shape::kM,
|
||||
(Shape::kN + ArchMmaOperator::Shape::kN - 1) / ArchMmaOperator::Shape::kN,
|
||||
InstructionShape::kK / Policy::Operator::Shape::kK
|
||||
>;
|
||||
|
||||
public:
|
||||
|
||||
/// Underlying matrix multiply operator (concept: arch::Mma)
|
||||
ArchMmaOperator mma;
|
||||
|
||||
public:
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
/// Ctor
|
||||
CUTLASS_DEVICE
|
||||
MmaTensorOp() {}
|
||||
|
||||
/// Performs a warp-level matrix multiply-accumulate operation
|
||||
CUTLASS_DEVICE
|
||||
void operator()(
|
||||
FragmentC &D,
|
||||
TransformedFragmentA const &A,
|
||||
TransformedFragmentB const &B,
|
||||
FragmentC const &C
|
||||
) const {
|
||||
|
||||
using MmaOperandA = typename ArchMmaOperator::FragmentA;
|
||||
using MmaOperandB = typename ArchMmaOperator::FragmentB;
|
||||
using MmaOperandC = typename ArchMmaOperator::FragmentC;
|
||||
|
||||
D = C;
|
||||
|
||||
MmaOperandA const *ptr_A = reinterpret_cast<MmaOperandA const *>(&A);
|
||||
MmaOperandB const *ptr_B = reinterpret_cast<MmaOperandB const *>(&B);
|
||||
MmaOperandC *ptr_D = reinterpret_cast<MmaOperandC *>(&D);
|
||||
|
||||
// Serpentine visitation order maximizing reuse of Rb
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int k = 0; k < MmaIterations::kK; ++k) {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int m = 0; m < MmaIterations::kM; ++m) {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int n = 0; n < MmaIterations::kN; ++n) {
|
||||
int n_serpentine = ((m % 2) ? (MmaIterations::kN - 1 - n) : n);
|
||||
|
||||
/// assume A is column-major in VRF, B is row-major in VRF
|
||||
if(AccumulatorsInRowMajor) {
|
||||
mma(
|
||||
ptr_D[n_serpentine + m * MmaIterations::kN],
|
||||
ptr_A[m + k * MmaIterations::kM],
|
||||
ptr_B[n_serpentine + k * MmaIterations::kN],
|
||||
ptr_D[n_serpentine + m * MmaIterations::kN]);
|
||||
} else {
|
||||
mma(
|
||||
ptr_D[m + n_serpentine * MmaIterations::kM],
|
||||
ptr_A[m + k * MmaIterations::kM],
|
||||
ptr_B[n_serpentine + k * MmaIterations::kN],
|
||||
ptr_D[m + n_serpentine * MmaIterations::kM]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Transform the mma operands to the required types
|
||||
CUTLASS_DEVICE
|
||||
void transform(TransformedFragmentA &dst_A, TransformedFragmentB &dst_B,
|
||||
FragmentA const &A, FragmentB const &B) const {
|
||||
|
||||
//
|
||||
// Define conversions from source type to instruction type
|
||||
//
|
||||
FloatRoundStyle const kRoundA =
|
||||
PreferredRoundingMode<typename ArchMmaOperator::ElementA,
|
||||
ElementA>::kRound;
|
||||
FloatRoundStyle const kRoundB =
|
||||
PreferredRoundingMode<typename ArchMmaOperator::ElementB,
|
||||
ElementB>::kRound;
|
||||
detail::ConvertAndPack<typename ArchMmaOperator::ElementA, ElementA,
|
||||
FragmentA::kElements / 2, kRoundA>
|
||||
convert_A;
|
||||
NumericArrayConverter<typename ArchMmaOperator::ElementB, ElementB,
|
||||
FragmentB::kElements, kRoundB>
|
||||
convert_B;
|
||||
Array<ElementA, FragmentA::kElements / 2> const *ptr_A =
|
||||
reinterpret_cast<Array<ElementA, FragmentA::kElements / 2> const *>(&A);
|
||||
Array<typename ArchMmaOperator::ElementA, FragmentA::kElements / 2> *
|
||||
ptr_dst_A = reinterpret_cast<Array<typename ArchMmaOperator::ElementA,
|
||||
FragmentA::kElements / 2> *>(&dst_A);
|
||||
|
||||
dst_B = convert_B(B);
|
||||
|
||||
ptr_dst_A[0] = convert_A(ptr_A[0]);
|
||||
ptr_dst_A[1] = convert_A(ptr_A[1]);
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace warp
|
||||
} // namespace gemm
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
71
cat_files/mma_tensor_op_policy.h
Normal file
71
cat_files/mma_tensor_op_policy.h
Normal file
@@ -0,0 +1,71 @@
|
||||
/***************************************************************************************************
|
||||
* 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:
|
||||
* * 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.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2021 Iluvatar CoreX. All rights reserved.
|
||||
* Copyright Declaration: This software, including all of its code and documentation,
|
||||
* except for the third-party software it contains, is a copyrighted work of Shanghai Iluvatar CoreX
|
||||
* Semiconductor Co., Ltd. and its affiliates ("Iluvatar CoreX") in accordance with the PRC Copyright
|
||||
* Law and relevant international treaties, and all rights contained therein are enjoyed by Iluvatar
|
||||
* CoreX. No user of this software shall have any right, ownership or interest in this software and
|
||||
* any use of this software shall be in compliance with the terms and conditions of the End User
|
||||
* License Agreement.
|
||||
**************************************************************************************************/
|
||||
|
||||
/*! \file
|
||||
\brief Policy describing implementation details of warp-level GEMM targeting Tensor Cores.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/matrix_shape.h"
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace gemm {
|
||||
namespace warp {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Policy
|
||||
template <
|
||||
typename Operator_, ///< hardware instruction(s) performing TensorOp (concept: arch::Mma)
|
||||
typename OpDelta_ ///< distance between operations (concept: MatrixShape)
|
||||
>
|
||||
struct MmaTensorOpPolicy {
|
||||
|
||||
using Operator = Operator_; ///< hardware instruction(s) performing TensorOp (concept: arch::Mma)
|
||||
using OpDelta = OpDelta_; ///< distance between operations (concept: MatrixShape)
|
||||
using MmaShape = typename Operator::Shape;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace warp
|
||||
} // namespace gemm
|
||||
} // namespace cutlass
|
||||
5595
cat_files/mma_tensor_op_tile_iterator.h
Normal file
5595
cat_files/mma_tensor_op_tile_iterator.h
Normal file
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user