CUTLASS 2.0 (#62)
CUTLASS 2.0 Substantially refactored for - Better performance, particularly for native Turing Tensor Cores - Robust and durable templates spanning the design space - Encapsulated functionality embodying modern C++11 programming techniques - Optimized containers and data types for efficient, generic, portable device code Updates to: - Quick start guide - Documentation - Utilities - CUTLASS Profiler Native Turing Tensor Cores - Efficient GEMM kernels targeting Turing Tensor Cores - Mixed-precision floating point, 8-bit integer, 4-bit integer, and binarized operands Coverage of existing CUTLASS functionality: - GEMM kernels targeting CUDA and Tensor Cores in NVIDIA GPUs - Volta Tensor Cores through native mma.sync and through WMMA API - Optimizations such as parallel reductions, threadblock rasterization, and intra-threadblock reductions - Batched GEMM operations - Complex-valued GEMMs Note: this commit and all that follow require a host compiler supporting C++11 or greater.
This commit is contained in:
Andrew Kerr
@@ -20,19 +20,8 @@
|
||||
# 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.
|
||||
|
||||
|
||||
set(EXAMPLES_SPLITK_GEMM_SOURCES
|
||||
splitK_gemm.cu
|
||||
)
|
||||
|
||||
if (NOT CUTLASS_NATIVE_CUDA)
|
||||
# cuda_add_executable does not take interface include directories into account
|
||||
# Let's fetch them and pass them to CUDA.
|
||||
get_target_property(CUTLASS_INCLUDES CUTLASS INTERFACE_INCLUDE_DIRECTORIES)
|
||||
include_directories("${CUTLASS_INCLUDES}")
|
||||
endif()
|
||||
|
||||
cutlass_add_executable(
|
||||
cutlass_example_add_executable(
|
||||
06_splitK_gemm
|
||||
${EXAMPLES_SPLITK_GEMM_SOURCES}
|
||||
)
|
||||
splitk_gemm.cu
|
||||
)
|
||||
|
||||
|
||||
@@ -1,298 +0,0 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2017-2019, NVIDIA CORPORATION. All rights reserved.
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without modification, are permitted
|
||||
* provided that the following conditions are met:
|
||||
* * Redistributions of source code must retain the above copyright notice, this list of
|
||||
* conditions and the following disclaimer.
|
||||
* * Redistributions in binary form must reproduce the above copyright notice, this list of
|
||||
* conditions and the following disclaimer in the documentation and/or other materials
|
||||
* provided with the distribution.
|
||||
* * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
|
||||
* to endorse or promote products derived from this software without specific prior written
|
||||
* permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
|
||||
* IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
|
||||
* FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
|
||||
* BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
|
||||
* OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
|
||||
* STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/gemm/device_gemm.h"
|
||||
#include "cutlass/gemm/sgemm_traits.h"
|
||||
#include "cutlass/reduction/batched_reduction_traits.h"
|
||||
#include "cutlass/gemm/device_gemm_traits.h"
|
||||
#pragma warning( disable : 4503)
|
||||
/*
|
||||
This example demonstrates how to use cutlass to compute sgemm with splitK
|
||||
splitK is useful for gemm with small M and N and reasonably large K.
|
||||
Because the sizes of M and N are small, the number of threadblocks we can launch is often limited and
|
||||
results in under utilization of the hardware.
|
||||
splitK allows us to divide a gemm across K dimension by first launching a partitionedK gemm (very similar to batched gemm),
|
||||
storing the intermediate result in workspace and then launching a second reduction kernel.
|
||||
Thus, as demonstrated by function cutlass_splitK_sgemm_nn(), the users need to create two traits, one for the partitionedK gemm,
|
||||
and one for the reduction. The users are also responsible for allocating and releasing the workspace memory. The size of the workspace
|
||||
memory can be queried by calling required_workspace_memory_in_byte().
|
||||
*/
|
||||
|
||||
template<int splits_count>
|
||||
cudaError_t cutlass_splitK_sgemm_nn(float const *A,
|
||||
int lda,
|
||||
float const *B,
|
||||
int ldb,
|
||||
float *C,
|
||||
int ldc,
|
||||
float alpha,
|
||||
float beta,
|
||||
int m,
|
||||
int n,
|
||||
int k) {
|
||||
cudaError_t result = cudaSuccess;
|
||||
|
||||
// create cutlass gemm traits for the first kernel
|
||||
typedef cutlass::gemm::SgemmTraits<cutlass::MatrixLayout::kColumnMajor, /*the layout of A*/
|
||||
cutlass::MatrixLayout::kColumnMajor, /*the layout of B*/
|
||||
cutlass::Shape<8, 128, 128> > /*the tile for each threadblock*/
|
||||
SgemmTraits;
|
||||
|
||||
// create cutlass batched reduction traits for the second kernel
|
||||
// for reduction D = alpha * Reduction(A) + beta * C
|
||||
typedef cutlass::reduction::BatchedReductionTraits<float, /*the scalar type of A in reduction, not to be confused with A in GEMM*/
|
||||
float, /*the scalar type of C in reduction, not to be confused with C in GEMM*/
|
||||
float, /*the scalar type of D in reduction, not to be confused with D in GEMM*/
|
||||
float, /*the scalar type of alpha and beta in reduction*/
|
||||
float, /*the scalar type of accumulation in reduction*/
|
||||
splits_count /*reduction workload*/
|
||||
>
|
||||
BatchedReductionTraits;
|
||||
|
||||
// create a device gemm that packages gemm traits and batched reduction traits
|
||||
typedef cutlass::gemm::SplitkPIGemmTraits<SgemmTraits, BatchedReductionTraits> deviceGemmTraits;
|
||||
|
||||
// kernel class
|
||||
typedef typename deviceGemmTraits::KernelClass deviceGemm;
|
||||
|
||||
// Params ctor requires M, N, K sizes
|
||||
typename deviceGemm::Params deviceGemmParams(m, n, k);
|
||||
|
||||
// query if workspace is needed. the workspace size is sizeof(accumulateType) * M * N * splits_count
|
||||
size_t workspace_size = deviceGemmParams.required_workspace_memory_in_byte();
|
||||
|
||||
// allocate workspace memory
|
||||
float *workspace_ptr;
|
||||
result = cudaMalloc(&workspace_ptr, workspace_size);
|
||||
if (result != cudaSuccess) {
|
||||
std::cerr << "cudaMalloc result = " << result << std::endl;
|
||||
return result;
|
||||
}
|
||||
|
||||
// finish init Params
|
||||
deviceGemmParams.initialize(alpha, /*alpha*/
|
||||
A, /*A*/
|
||||
lda, /*lda*/
|
||||
B, /*B*/
|
||||
ldb, /*ldb*/
|
||||
beta, /*beta*/
|
||||
C, /*C*/
|
||||
ldc, /*ldc*/
|
||||
C, /*D, can point to the same memory with C*/
|
||||
ldc, /*ldc*/
|
||||
workspace_ptr /*ptr to workspace*/
|
||||
);
|
||||
|
||||
// launch the kernel
|
||||
deviceGemm::launch(deviceGemmParams);
|
||||
result = cudaDeviceSynchronize();
|
||||
if (result != cudaSuccess) {
|
||||
std::cerr << "launch result = " << result << std::endl;
|
||||
cudaFree(workspace_ptr);
|
||||
return result;
|
||||
}
|
||||
|
||||
// release the workspace memory
|
||||
result = cudaFree(workspace_ptr);
|
||||
if (result != cudaSuccess) {
|
||||
std::cerr << "cudaFree result = " << result << std::endl;
|
||||
}
|
||||
|
||||
return cudaGetLastError();
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
cudaError_t sgemm_nn_reference(std::vector<T> const &A,
|
||||
int lda,
|
||||
std::vector<T> const &B,
|
||||
int ldb,
|
||||
std::vector<T> &C,
|
||||
int ldc,
|
||||
T alpha,
|
||||
T beta,
|
||||
int m,
|
||||
int n,
|
||||
int k) {
|
||||
/*
|
||||
sgemm
|
||||
*/
|
||||
|
||||
cudaError_t result = cudaSuccess;
|
||||
for (int n_idx = 0; n_idx < n; n_idx++) {
|
||||
for (int m_idx = 0; m_idx < m; m_idx++) {
|
||||
T accum = beta * C[n_idx * ldc + m_idx];
|
||||
for (int k_idx = 0; k_idx < k; k_idx++) {
|
||||
accum += alpha
|
||||
* A[k_idx * lda + m_idx]
|
||||
* B[n_idx * ldb + k_idx];
|
||||
}
|
||||
C[n_idx * ldc + m_idx] = accum;
|
||||
}
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
int main() {
|
||||
int const m = 128;
|
||||
int const n = 128;
|
||||
int const k = 4096;
|
||||
//splits_count should be known at compile time
|
||||
int const splits_count = 80;
|
||||
|
||||
// A, B are non-transpose, column major
|
||||
int const lda = m;
|
||||
int const ldb = k;
|
||||
int const ldc = m;
|
||||
|
||||
int const count_A = lda * k;
|
||||
int const count_B = ldb * n;
|
||||
int const count_C = ldc * 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;
|
||||
}
|
||||
|
||||
// fill A
|
||||
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] = static_cast<float>((row_idx + col_idx) % 10);
|
||||
}
|
||||
}
|
||||
|
||||
// fill B
|
||||
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] = static_cast<float>((row_idx - col_idx) % 5);
|
||||
}
|
||||
}
|
||||
|
||||
// fill C
|
||||
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] = 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_splitK_sgemm_nn<splits_count>(A, lda, B, ldb, C, ldc, alpha, beta, m, n, k);
|
||||
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 = sgemm_nn_reference(ref_A, lda, ref_B, ldb, ref_C, ldc, alpha, beta, m, n, k);
|
||||
if (result != 0)
|
||||
return result;
|
||||
|
||||
if (ref_C != result_C) {
|
||||
std::cout << "CUTLASS splitK 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;
|
||||
}
|
||||
@@ -0,0 +1,305 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2017-2019, NVIDIA CORPORATION. All rights reserved.
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without modification, are permitted
|
||||
* provided that the following conditions are met:
|
||||
* * Redistributions of source code must retain the above copyright notice, this list of
|
||||
* conditions and the following disclaimer.
|
||||
* * Redistributions in binary form must reproduce the above copyright notice, this list of
|
||||
* conditions and the following disclaimer in the documentation and/or other materials
|
||||
* provided with the distribution.
|
||||
* * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
|
||||
* to endorse or promote products derived from this software without specific prior written
|
||||
* permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
|
||||
* IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
|
||||
* FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
|
||||
* BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
|
||||
* OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
|
||||
* STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
/**
|
||||
This example shows how to use split-k version of matrix multiplication using functions and data
|
||||
structures provided by CUTLASS; which we run on a NVIDIA Volta GPU.
|
||||
|
||||
What is split-k?
|
||||
Consider a problem size of M = 128, N = 128, K = 4096. In this case, if my thread-block tile size (a
|
||||
tile can be viewed as a 2d matrix) is 128x128x4096, then we launch a singled a thread-block taking
|
||||
up a single SM of 84 SMs present on V100. Hence the efficiency of computation is really low. So, how
|
||||
to solve it? This is where split-k comes in. It is a way of partitioning K-dimension of matrix
|
||||
multiplication and distribute across multiple SMs and get better efficiency than single SM. In the
|
||||
above example, we can partition K-dimension with split-k factor of 16 i.e., thread-block tile size
|
||||
will be 128x128x256 and will be launching on 16 SMs. Once each thread-block computes their partial
|
||||
inner product (1/16th of output), they accumulate to single output matrix.
|
||||
|
||||
Writing a single high performance matrix multiplication kernel is hard but do-able. Whereas writing
|
||||
high performance kernels at scale which works for multiple problem sizes with good abstractions is
|
||||
really hard. CUTLASS solves this problem by providing simplified abstractions (knobs) to compose
|
||||
multiple sections of gemm kernel. When used properly, the kernels can hit peak performance of GPU
|
||||
easily.
|
||||
|
||||
CUTLASS divides a kernel into hierarchical composable sections. Which means, at each thread, warp
|
||||
and thread-block level, they compute on their own tile-size with higher level of tile sizes being
|
||||
composed from lower level ones. Multiple thread-tiles (tile size each thread computes) can be used
|
||||
to form warp-tiles (tile size each warp computes) and multiple warp tiles can be used to compute
|
||||
threadblock-tile (tile size computed by a threadblock).
|
||||
|
||||
In thie example, we split variable initialization into
|
||||
1. Setting up data properties : describes how matrices are laid out in the memory and how the kernel
|
||||
can view them (logical to physical mapping)
|
||||
2. Setting up computation properties : describes how the above set matrices will be used to compute
|
||||
output of matrix multiplication.
|
||||
|
||||
First, we setup the data types of matrices A, B, C and D along with alpha, beta as the equation for
|
||||
GEMM is D = alpha * A * B + beta * C. In CUTLASS, the kernels first compute A * B and leaves the
|
||||
rest of the computation to end of the kernel as alpha * X + beta * C is a simple element-wise
|
||||
operation on X (A * B) and C. We call this as epilogue of kernel. Hence, we setup data types for
|
||||
alpha and beta to be equal to ElementComputeEpilogue = float. As we want to MMA instructions on
|
||||
Volta and they support only half-precision floating point (fp16 or half), we use data type for
|
||||
elements in input matrix A and B as cutlass::half_t. Volta also supports accumulation of partial dot
|
||||
product to fp32, which can store wider range of numbers, we use it as data type of output matrix
|
||||
elements and accumulation. We convey this to CUTLASS kernel by initializing template variables
|
||||
ElementAccumulator (float), ElementComputeEpilogue (float), ElementInputA (cutlass::half_t),
|
||||
ElementInputB (cutlass::half_t), ElementOutput (float). Communicating just the data type is not
|
||||
enough. As the data is laid out linearly in memory, we have to convey the layout of matrices. We do
|
||||
that by initializing template variable LayoutInputA to column major cutlass variable, LayoutInputB
|
||||
to row major and LayoutOutput to row major. Next, we setup rules to comptue alpha * X + beta * C
|
||||
which is called epilogue of the kernel. We initialize template variable EpilogueOp, which takes the
|
||||
data type of output ElementOutput (int32_t), the number of elements per vector memory access (16),
|
||||
data type of accumulator (int32_t) and data type of computation of linear combination (alpha * X +
|
||||
beta * C).
|
||||
|
||||
Now that we setup the properties of data, we have to setup properties of computation.
|
||||
|
||||
Second, we create template variables of tile sizes for thread-block, warp and mma-op to 128x128x32,
|
||||
64x64x4, 8x8x4 (MxNxK) respectively. When passed to instantiate CUTLASS GEMM kernel, it internally
|
||||
deduce the amount of threads needed per thread-block, amount of shared memory, storing data in
|
||||
bank-conflict free manner, and ton of other variables required to compose, intialize and launch a
|
||||
high performance GEMM kernel. This is the beauty of CUTLASS, it relieves developer from
|
||||
understanding and coding complicated hardware optimizations which can easily go wrong.
|
||||
|
||||
There are few more template variables initialized such as, which threadblock tile of output matrix
|
||||
is done which threadblock launched on an SM, CUDA SM architecture of GPU you want to run on.
|
||||
|
||||
These are all put together to create a template variable which describes CUTLASS GEMM kernel using
|
||||
cutlass::gemm::device::GemmSplitKParallel template.
|
||||
|
||||
The next step is to intialize physical data, instantiate and initialize CUTLASS kernel and run it.
|
||||
We use CUTLASS utilities to initialize, fill, compare matrices as they are simple and doesn't come
|
||||
in the way of learning CUTLASS.
|
||||
|
||||
Once all the matrices are initialized and filled with data, create arguments tuple to launch CUTLASS
|
||||
kernel which takes problem size (M = 5120, N = 4096 and K = 4096), matrices, alpha, beta and the
|
||||
important one, split k-dimension factor. Along with that, we query CUTLASS if any scratch-space
|
||||
memory required by the kernel we instantiated. If yes, we create it and pass it along with other
|
||||
arguments created to intialize CUTLASS kernel then, the kernel is launched.
|
||||
|
||||
In this example, we later on launch a reference gemm kernel (from CUTLASS utilities) to compare if
|
||||
the output from CUTLASS kernel is same as reference GEMM kernel.
|
||||
*/
|
||||
|
||||
#include <iostream>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/gemm/device/gemm_splitk_parallel.h"
|
||||
#include "cutlass/util/host_tensor.h"
|
||||
#include "cutlass/util/reference/device/gemm.h"
|
||||
#include "cutlass/util/reference/host/tensor_compare.h"
|
||||
#include "cutlass/util/reference/host/tensor_copy.h"
|
||||
#include "cutlass/util/reference/host/tensor_fill.h"
|
||||
#include "cutlass/util/tensor_view_io.h"
|
||||
#include "helper.h"
|
||||
|
||||
// The code section below describes datatype for input, output matrices and computation between
|
||||
// elements in input matrices.
|
||||
using ElementAccumulator = float; // <- data type of accumulator
|
||||
using ElementComputeEpilogue = ElementAccumulator; // <- data type of epilogue operations
|
||||
using ElementInputA = cutlass::half_t; // <- data type of elements in input matrix A
|
||||
using ElementInputB = cutlass::half_t; // <- data type of elements in input matrix B
|
||||
using ElementOutput = float; // <- data type of elements in output matrix D
|
||||
|
||||
// The code section below describes matrix layout of input and output matrices. Column Major for
|
||||
// Matrix A, Row Major for Matrix B and Row Major for Matrix C
|
||||
using LayoutInputA = cutlass::layout::ColumnMajor;
|
||||
using LayoutInputB = cutlass::layout::RowMajor;
|
||||
using LayoutOutput = cutlass::layout::RowMajor;
|
||||
|
||||
// This code section describes whether you want to use tensor cores or regular SIMT cores on GPU SM
|
||||
using MMAOp = cutlass::arch::OpClassTensorOp;
|
||||
|
||||
// This code section describes CUDA SM architecture number
|
||||
using SmArch = cutlass::arch::Sm70;
|
||||
|
||||
// This code section describes the tile size a thread block will compute
|
||||
using ShapeMMAThreadBlock =
|
||||
cutlass::gemm::GemmShape<128, 128, 32>; // <- threadblock tile M = 128, N = 128, K = 32
|
||||
// This code section describes tile size a warp will compute
|
||||
using ShapeMMAWarp = cutlass::gemm::GemmShape<64, 64, 32>; // <- warp tile M = 64, N = 64, K = 32
|
||||
// This code section describes the size of MMA op
|
||||
using ShapeMMAOp = cutlass::gemm::GemmShape<8, 8, 4>; // <- MMA Op tile M = 8, N = 8, K = 4
|
||||
|
||||
// This code section describes how threadblocks are scheduled on GPU
|
||||
using SwizzleThreadBlock = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle; // <- ??
|
||||
|
||||
// This code section describes ?
|
||||
using EpilogueOp = cutlass::epilogue::thread::LinearCombination<
|
||||
ElementOutput, // <- data type of output matrix
|
||||
128 / cutlass::sizeof_bits<ElementOutput>::value, // <- This is the number of elements per
|
||||
// vectorized memory access. For half
|
||||
// precision, it's 8 elements. This becomes
|
||||
// the vector width of math instructions in
|
||||
// epilogue too
|
||||
ElementAccumulator, // <- data type of accumulator
|
||||
ElementComputeEpilogue>; // <- data type for alpha/beta in linear combination function
|
||||
|
||||
// Put all the created template variables to create GemmSplitKParallel template variable
|
||||
using Gemm = cutlass::gemm::device::GemmSplitKParallel<ElementInputA,
|
||||
LayoutInputA,
|
||||
ElementInputB,
|
||||
LayoutInputB,
|
||||
ElementOutput,
|
||||
LayoutOutput,
|
||||
ElementAccumulator,
|
||||
MMAOp,
|
||||
SmArch,
|
||||
ShapeMMAThreadBlock,
|
||||
ShapeMMAWarp,
|
||||
ShapeMMAOp,
|
||||
EpilogueOp>;
|
||||
|
||||
int main() {
|
||||
cudaDeviceProp props;
|
||||
CUDA_CHECK(cudaGetDeviceProperties(&props, 0));
|
||||
if (!(props.major >= 7)) {
|
||||
std::cerr << "Volta Tensor Ops must be run on a machine with compute capability at least 70."
|
||||
<< std::endl;
|
||||
return 0;
|
||||
}
|
||||
|
||||
const int length_m = 5120;
|
||||
const int length_n = 4096;
|
||||
const int length_k = 4096;
|
||||
|
||||
// Create a tuple of problem size for matrix multiplication
|
||||
cutlass::gemm::GemmCoord problem_size(length_m, length_n, length_k);
|
||||
|
||||
// Initialize tensors using CUTLASS helper functions
|
||||
cutlass::HostTensor<ElementInputA, LayoutInputA> tensor_a(
|
||||
problem_size.mk()); // <- Create matrix A with dimensions M x K
|
||||
cutlass::HostTensor<ElementInputB, LayoutInputB> tensor_b(
|
||||
problem_size.nk()); // <- Create matrix B with dimensions N x K
|
||||
cutlass::HostTensor<ElementOutput, LayoutOutput> tensor_c(
|
||||
problem_size.mn()); // <- Create matrix C with dimensions M x N
|
||||
cutlass::HostTensor<ElementOutput, LayoutOutput> tensor_d(
|
||||
problem_size.mn()); // <- Create matrix D with dimensions M x N used to store output from
|
||||
// CUTLASS kernel
|
||||
cutlass::HostTensor<ElementOutput, LayoutOutput> tensor_ref_d(
|
||||
problem_size.mn()); // <- Create matrix D with dimensions M x N used to store output from
|
||||
// reference kernel
|
||||
|
||||
// Fill input and output matrices on host using CUTLASS helper functions
|
||||
cutlass::reference::host::TensorFillRandomUniform(
|
||||
tensor_a.host_view(),
|
||||
1,
|
||||
ElementInputA(4),
|
||||
ElementInputA(-4),
|
||||
0); // <- Fill matrix A on host with uniform-distribution random data
|
||||
cutlass::reference::host::TensorFillRandomUniform(
|
||||
tensor_b.host_view(),
|
||||
1,
|
||||
ElementInputB(4),
|
||||
ElementInputB(-4),
|
||||
0); // <- Fill matrix B on host with uniform-distribution random data
|
||||
cutlass::reference::host::TensorFillRandomUniform(
|
||||
tensor_c.host_view(),
|
||||
1,
|
||||
ElementOutput(4),
|
||||
ElementOutput(-4),
|
||||
0); // <- Fill matrix C on host with uniform-distribution random data
|
||||
cutlass::reference::host::TensorFill(
|
||||
tensor_d.host_view()); // <- fill matrix D on host with zeros
|
||||
cutlass::reference::host::TensorFill(
|
||||
tensor_ref_d.host_view()); // <- fill matrix D for reference on host with zeros
|
||||
|
||||
// Copy data from host to GPU
|
||||
tensor_a.sync_device();
|
||||
tensor_b.sync_device();
|
||||
tensor_c.sync_device();
|
||||
tensor_d.sync_device();
|
||||
tensor_ref_d.sync_device();
|
||||
|
||||
// Initialize alpha and beta for dot product computation
|
||||
ElementComputeEpilogue alpha = ElementComputeEpilogue(1);
|
||||
ElementComputeEpilogue beta = ElementComputeEpilogue(0);
|
||||
|
||||
// Split K dimension into 16 partitions
|
||||
int split_k_slices = 16;
|
||||
|
||||
// Create a tuple of gemm kernel arguments. This is later passed as arguments to launch
|
||||
// instantiated CUTLASS kernel
|
||||
typename Gemm::Arguments arguments{problem_size, // <- problem size of matrix multiplication
|
||||
tensor_a.device_ref(), // <- reference to matrix A on device
|
||||
tensor_b.device_ref(), // <- reference to matrix B on device
|
||||
tensor_c.device_ref(), // <- reference to matrix C on device
|
||||
tensor_d.device_ref(), // <- reference to matrix D on device
|
||||
{alpha, beta}, // <- tuple of alpha and beta
|
||||
split_k_slices}; // <- k-dimension split factor
|
||||
|
||||
// Using the arguments, query for extra workspace required for matrix multiplication computation
|
||||
size_t workspace_size = Gemm::get_workspace_size(arguments);
|
||||
|
||||
// Allocate workspace memory
|
||||
cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
|
||||
|
||||
// Instantiate CUTLASS kernel depending on templates
|
||||
Gemm gemm_op;
|
||||
|
||||
// Initialize CUTLASS kernel with arguments and workspace pointer
|
||||
cutlass::Status status = gemm_op.initialize(arguments, workspace.get());
|
||||
CUTLASS_CHECK(status);
|
||||
|
||||
// Launch initialized CUTLASS kernel
|
||||
status = gemm_op();
|
||||
CUTLASS_CHECK(status);
|
||||
|
||||
// Create instantiation for device reference gemm kernel
|
||||
cutlass::reference::device::Gemm<ElementInputA,
|
||||
LayoutInputA,
|
||||
ElementInputB,
|
||||
LayoutInputB,
|
||||
ElementOutput,
|
||||
LayoutOutput,
|
||||
ElementComputeEpilogue,
|
||||
ElementComputeEpilogue>
|
||||
gemm_device;
|
||||
|
||||
// Launch device reference gemm kernel
|
||||
gemm_device(problem_size,
|
||||
alpha,
|
||||
tensor_a.device_ref(),
|
||||
tensor_b.device_ref(),
|
||||
beta,
|
||||
tensor_c.device_ref(),
|
||||
tensor_ref_d.device_ref());
|
||||
|
||||
// Wait for kernels to finish
|
||||
cudaDeviceSynchronize();
|
||||
|
||||
// Copy output data from CUTLASS and reference kernel to host for comparison
|
||||
tensor_d.sync_host();
|
||||
tensor_ref_d.sync_host();
|
||||
|
||||
// Check if output from CUTLASS kernel and reference kernel are equal or not
|
||||
std::cout << (cutlass::reference::host::TensorEquals(tensor_d.host_view(),
|
||||
tensor_ref_d.host_view())
|
||||
? "Passed"
|
||||
: "Failed")
|
||||
<< std::endl;
|
||||
|
||||
CUTLASS_CHECK(status);
|
||||
}
|
||||
Reference in New Issue
Block a user