1544 lines
48 KiB
C++
1544 lines
48 KiB
C++
/***************************************************************************************************
|
|
* Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
* SPDX-License-Identifier: BSD-3-Clause
|
|
*
|
|
* Redistribution and use in source and binary forms, with or without
|
|
* modification, are permitted provided that the following conditions are met:
|
|
*
|
|
* 1. Redistributions of source code must retain the above copyright notice, this
|
|
* list of conditions and the following disclaimer.
|
|
*
|
|
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
|
* this list of conditions and the following disclaimer in the documentation
|
|
* and/or other materials provided with the distribution.
|
|
*
|
|
* 3. Neither the name of the copyright holder nor the names of its
|
|
* contributors may be used to endorse or promote products derived from
|
|
* this software without specific prior written permission.
|
|
*
|
|
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
|
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
|
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
|
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
|
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
|
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
|
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
|
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
|
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
|
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
|
*
|
|
**************************************************************************************************/
|
|
/*! \file
|
|
\brief Unit tests for thread-level GEMM
|
|
*/
|
|
|
|
#pragma once
|
|
|
|
#include "cutlass/cutlass.h"
|
|
#include "cutlass/aligned_buffer.h"
|
|
#include "cutlass/numeric_types.h"
|
|
#include "cutlass/subbyte_reference.h"
|
|
#include "cutlass/platform/platform.h"
|
|
#include "cutlass/arch/arch.h"
|
|
|
|
#include "cutlass/util/host_tensor.h"
|
|
#include "cutlass/util/tensor_view_io.h"
|
|
|
|
#include "cutlass/util/distribution.h"
|
|
#include "cutlass/util/reference/host/gemm.h"
|
|
#include "cutlass/util/reference/host/gemm_complex.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/host_reorder.h"
|
|
#include "cutlass/util/host_uncompress.h"
|
|
|
|
namespace test {
|
|
namespace gemm {
|
|
namespace warp {
|
|
|
|
/////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
/// Test kernel
|
|
template <typename Mma, typename ThreadblockShape>
|
|
__global__ void kernel(
|
|
typename Mma::ElementC *output_C,
|
|
typename Mma::ElementA const *input_A,
|
|
typename Mma::ElementB const *input_B,
|
|
typename Mma::ElementC const *input_C,
|
|
int iterations = 1) {
|
|
|
|
// Use AlignedBuffer to store trivially copyable objects in unions and __shared__ buffers.
|
|
__shared__ cutlass::AlignedBuffer<
|
|
typename Mma::ElementA, ThreadblockShape::kM * ThreadblockShape::kK> smem_buffer_A;
|
|
|
|
__shared__ cutlass::AlignedBuffer<
|
|
typename Mma::ElementB, ThreadblockShape::kN * ThreadblockShape::kK> smem_buffer_B;
|
|
|
|
if (threadIdx.x == 0) {
|
|
typename Mma::ElementA *smem_ptr_A = smem_buffer_A.data();
|
|
#pragma unroll 1
|
|
for (size_t i = 0; i < smem_buffer_A.size(); ++i) {
|
|
cutlass::ReferenceFactory<typename Mma::ElementA>::get(smem_ptr_A, i) =
|
|
cutlass::ReferenceFactory<typename cutlass::platform::remove_const<
|
|
typename Mma::ElementA>::type>::get(input_A, i);
|
|
}
|
|
|
|
typename Mma::ElementB *smem_ptr_B = smem_buffer_B.data();
|
|
#pragma unroll 1
|
|
for (size_t i = 0; i < smem_buffer_B.size(); ++i) {
|
|
cutlass::ReferenceFactory<typename Mma::ElementB>::get(smem_ptr_B, i) =
|
|
cutlass::ReferenceFactory<typename cutlass::platform::remove_const<
|
|
typename Mma::ElementB>::type>::get(input_B, i);
|
|
}
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
//
|
|
// Construct warp-level matrix product
|
|
//
|
|
|
|
using FragmentA = typename Mma::FragmentA;
|
|
using FragmentB = typename Mma::FragmentB;
|
|
using FragmentC = typename Mma::FragmentC;
|
|
|
|
typename Mma::LayoutA layout_A = Mma::LayoutA::packed({ThreadblockShape::kM, ThreadblockShape::kK});
|
|
typename Mma::LayoutB layout_B = Mma::LayoutB::packed({ThreadblockShape::kK, ThreadblockShape::kN});
|
|
typename Mma::LayoutC layout_C = Mma::LayoutC::packed({Mma::Shape::kM, Mma::Shape::kN});
|
|
|
|
typename Mma::IteratorA iter_A({smem_buffer_A.data(), layout_A}, cutlass::arch::LaneId());
|
|
|
|
typename Mma::IteratorB iter_B({smem_buffer_B.data(), layout_B}, cutlass::arch::LaneId());
|
|
|
|
FragmentA frag_A;
|
|
FragmentB frag_B;
|
|
|
|
FragmentC accum;
|
|
|
|
Mma mma;
|
|
|
|
accum.clear();
|
|
|
|
CUTLASS_PRAGMA_NO_UNROLL
|
|
for (int iter = 0; iter < iterations; ++iter) { // place in loop that is not unrolled
|
|
|
|
CUTLASS_PRAGMA_UNROLL
|
|
for (int k = 0; k < ThreadblockShape::kK;
|
|
k += Mma::Policy::MmaShape::kK) {
|
|
iter_A.load(frag_A);
|
|
iter_B.load(frag_B);
|
|
|
|
++iter_A;
|
|
++iter_B;
|
|
|
|
mma(accum, frag_A, frag_B, accum);
|
|
}
|
|
}
|
|
|
|
typename Mma::IteratorC iter_C({output_C, layout_C}, cutlass::arch::LaneId());
|
|
|
|
iter_C.store(accum);
|
|
}
|
|
|
|
/////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
/// Structure to compute the matrix product
|
|
template <
|
|
/// Warp-level matrix multiply-accumulate
|
|
typename Mma_,
|
|
/// Size of threadblock-scoped shape used to store SMEM
|
|
typename ThreadblockShape_,
|
|
/// The inner product operation performed by GEMM
|
|
typename Operator_ = cutlass::arch::OpMultiplyAdd
|
|
>
|
|
struct Testbed {
|
|
|
|
/// Thread-level matrix multiply-accumulate operator
|
|
using Mma = Mma_;
|
|
using ThreadblockShape = ThreadblockShape_;
|
|
using Operator = Operator_;
|
|
|
|
using Shape = typename Mma::Shape;
|
|
using ElementA = typename Mma::ElementA;
|
|
using LayoutA = typename Mma::LayoutA;
|
|
using ElementB = typename Mma::ElementB;
|
|
using LayoutB = typename Mma::LayoutB;
|
|
using ElementC = typename Mma::ElementC;
|
|
using LayoutC = typename Mma::LayoutC;
|
|
|
|
//
|
|
// Data members
|
|
//
|
|
|
|
cutlass::HostTensor<ElementA, LayoutA> tensor_A;
|
|
cutlass::HostTensor<ElementB, LayoutB> tensor_B;
|
|
cutlass::HostTensor<ElementC, LayoutC> tensor_C;
|
|
cutlass::HostTensor<ElementC, LayoutC> tensor_D_computed;
|
|
cutlass::HostTensor<ElementC, LayoutC> tensor_D_reference;
|
|
|
|
//
|
|
// Methods
|
|
//
|
|
|
|
/// Allocates workspace in device memory
|
|
Testbed() {
|
|
|
|
tensor_A.reset(cutlass::make_Coord(ThreadblockShape::kM, ThreadblockShape::kK));
|
|
tensor_B.reset(cutlass::make_Coord(ThreadblockShape::kK, ThreadblockShape::kN));
|
|
tensor_C.reset(cutlass::make_Coord(Shape::kM, Shape::kN));
|
|
tensor_D_computed.reset(cutlass::make_Coord(Shape::kM, Shape::kN));
|
|
tensor_D_reference.reset(cutlass::make_Coord(Shape::kM, Shape::kN), false);
|
|
}
|
|
|
|
/// Returns true if the CUDA device is sufficient to execute the kernel.
|
|
bool sufficient() const {
|
|
|
|
cudaDeviceProp properties;
|
|
int device_idx;
|
|
cudaError_t result = cudaGetDevice(&device_idx);
|
|
|
|
if (result != cudaSuccess) {
|
|
throw std::runtime_error("cudaGetDevice() API call failed.");
|
|
}
|
|
|
|
result = cudaGetDeviceProperties(&properties, device_idx);
|
|
|
|
if (result != cudaSuccess) {
|
|
throw std::runtime_error("cudaGetDeviceProperties() failed");
|
|
}
|
|
|
|
if (properties.major == 9) {
|
|
// NVIDIA Hopper drops support for several data types
|
|
if (
|
|
cutlass::sizeof_bits<ElementA>::value < 8 ||
|
|
cutlass::sizeof_bits<ElementB>::value < 8 ||
|
|
cutlass::sizeof_bits<ElementC>::value < 8) {
|
|
|
|
return false;
|
|
}
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
|
|
/// Runs the test
|
|
bool run(
|
|
cutlass::Distribution::Kind init_A = cutlass::Distribution::Uniform,
|
|
cutlass::Distribution::Kind init_B = cutlass::Distribution::Uniform) {
|
|
|
|
if (!sufficient()) {
|
|
return true;
|
|
}
|
|
|
|
//
|
|
// initialize device memory
|
|
//
|
|
|
|
if (init_A == cutlass::Distribution::Uniform) {
|
|
int scope_max = 8;
|
|
int scope_min = -8;
|
|
|
|
if (cutlass::sizeof_bits<ElementA>::value == 4) {
|
|
scope_max = 2;
|
|
scope_min = -2;
|
|
} else if (cutlass::sizeof_bits<ElementA>::value == 1) {
|
|
scope_max = 2;
|
|
scope_min = 0;
|
|
}
|
|
|
|
uint64_t seed = 7;
|
|
|
|
cutlass::reference::host::BlockFillRandomUniform(tensor_A.host_data(),
|
|
tensor_A.capacity(), seed, scope_max, scope_min, 0);
|
|
|
|
} else if (init_A == cutlass::Distribution::Sequential) {
|
|
cutlass::reference::host::BlockFillSequential(tensor_A.host_data(),
|
|
tensor_A.capacity());
|
|
} else if (init_A == cutlass::Distribution::Identity) {
|
|
cutlass::reference::host::TensorFillIdentity(tensor_A.host_view());
|
|
} else {
|
|
return false;
|
|
}
|
|
|
|
if (init_B == cutlass::Distribution::Uniform) {
|
|
int scope_max = 8;
|
|
int scope_min = -8;
|
|
|
|
if (cutlass::sizeof_bits<ElementB>::value == 4) {
|
|
scope_max = 2;
|
|
scope_min = -2;
|
|
} else if (cutlass::sizeof_bits<ElementB>::value == 1) {
|
|
scope_max = 2;
|
|
scope_min = 0;
|
|
}
|
|
|
|
uint64_t seed = 7;
|
|
|
|
cutlass::reference::host::BlockFillRandomUniform(tensor_B.host_data(),
|
|
tensor_B.capacity(), seed, scope_max, scope_min, 0);
|
|
|
|
} else if (init_B == cutlass::Distribution::Sequential) {
|
|
cutlass::reference::host::BlockFillSequential(tensor_B.host_data(),
|
|
tensor_B.capacity());
|
|
} else if (init_B == cutlass::Distribution::Identity) {
|
|
cutlass::reference::host::TensorFillIdentity(tensor_B.host_view());
|
|
} else {
|
|
return false;
|
|
}
|
|
|
|
cutlass::reference::host::TensorFill(
|
|
tensor_C.host_view(),
|
|
ElementC(0)
|
|
);
|
|
|
|
cutlass::reference::host::TensorFill(
|
|
tensor_D_computed.host_view(),
|
|
ElementC(0)
|
|
);
|
|
|
|
cutlass::reference::host::TensorFill(
|
|
tensor_D_reference.host_view(),
|
|
ElementC(0)
|
|
);
|
|
|
|
tensor_A.sync_device();
|
|
tensor_B.sync_device();
|
|
tensor_C.sync_device();
|
|
tensor_D_computed.sync_device();
|
|
|
|
// launch kernel
|
|
kernel<Mma, ThreadblockShape><<< dim3(1, 1), dim3(32, 1, 1) >>>(
|
|
tensor_D_computed.device_data(),
|
|
tensor_A.device_data(),
|
|
tensor_B.device_data(),
|
|
tensor_C.device_data());
|
|
|
|
// verify no errors
|
|
cudaError_t result = cudaDeviceSynchronize();
|
|
|
|
EXPECT_EQ(result, cudaSuccess) << "CUDA ERROR: " << cudaGetErrorString(result);
|
|
if (result != cudaSuccess) {
|
|
return false;
|
|
}
|
|
|
|
tensor_D_computed.sync_host();
|
|
|
|
//
|
|
// Reference implementation
|
|
//
|
|
|
|
cutlass::reference::host::Gemm<ElementA, LayoutA, ElementB, LayoutB,
|
|
ElementC, LayoutC, ElementC, ElementC,
|
|
Operator>
|
|
reference_gemm;
|
|
|
|
reference_gemm(
|
|
{Shape::kM, Shape::kN, ThreadblockShape::kK},
|
|
ElementC(1),
|
|
tensor_A.host_ref(),
|
|
tensor_B.host_ref(),
|
|
ElementC(0),
|
|
tensor_D_reference.host_ref()
|
|
);
|
|
|
|
//
|
|
// Verify equivalence
|
|
//
|
|
|
|
// compare
|
|
bool passed = cutlass::reference::host::TensorEquals(
|
|
tensor_D_computed.host_view(),
|
|
tensor_D_reference.host_view()
|
|
);
|
|
|
|
EXPECT_TRUE(passed);
|
|
|
|
if (!passed) {
|
|
|
|
cutlass::TensorView<ElementA, cutlass::layout::ColumnMajor> tensor_A_physical(
|
|
tensor_A.host_data(),
|
|
tensor_A.stride()[0],
|
|
tensor_A.extent());
|
|
|
|
cutlass::TensorView<ElementB, cutlass::layout::RowMajor> tensor_B_physical(
|
|
tensor_B.host_data(),
|
|
tensor_B.stride()[0],
|
|
tensor_B.extent());
|
|
|
|
std::cout <<"cutlass::sizeof_bits<ElementA>::value = "<<cutlass::sizeof_bits<ElementA>::value<<"\n";
|
|
std::cout
|
|
<< "A:\n" << tensor_A.host_view() << "\n\n"
|
|
<< "A(physical - stride: " << tensor_A.stride()[0]
|
|
<< ", extent: " << tensor_A.extent() << "):\n" << tensor_A_physical << "\n\n";
|
|
|
|
std::cout <<"cutlass::sizeof_bits<ElementB>::value = "<<cutlass::sizeof_bits<ElementB>::value<<"\n";
|
|
std::cout
|
|
<< "B:\n" << tensor_B.host_view() << "\n\n"
|
|
<< "B(physical - stride: " << tensor_B.stride()[0]
|
|
<< ", extent: " << tensor_B.extent() << "):\n" << tensor_B_physical << "\n\n";
|
|
|
|
std::cout
|
|
<< "C:\n" << tensor_C.host_view() << "\n\n"
|
|
<< "Reference:\n" << tensor_D_reference.host_view() << "\n\n"
|
|
<< "Computed:\n" << tensor_D_computed.host_view() << std::endl;
|
|
}
|
|
|
|
return passed;
|
|
}
|
|
};
|
|
|
|
/////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
/// Structure to compute the matrix product
|
|
template <
|
|
/// Warp-level matrix multiply-accumulate
|
|
typename Mma_,
|
|
/// Size of threadblock-scoped shape used to store SMEM
|
|
typename ThreadblockShape_
|
|
>
|
|
struct TestbedComplex {
|
|
|
|
/// Thread-level matrix multiply-accumulate operator
|
|
using Mma = Mma_;
|
|
using ThreadblockShape = ThreadblockShape_;
|
|
|
|
using Shape = typename Mma::Shape;
|
|
using ElementA = typename Mma::ElementA;
|
|
using LayoutA = typename Mma::LayoutA;
|
|
using ElementB = typename Mma::ElementB;
|
|
using LayoutB = typename Mma::LayoutB;
|
|
using ElementC = typename Mma::ElementC;
|
|
using LayoutC = typename Mma::LayoutC;
|
|
|
|
//
|
|
// Data members
|
|
//
|
|
|
|
cutlass::HostTensor<ElementA, LayoutA> tensor_A;
|
|
cutlass::HostTensor<ElementB, LayoutB> tensor_B;
|
|
cutlass::HostTensor<ElementC, LayoutC> tensor_C;
|
|
cutlass::HostTensor<ElementC, LayoutC> tensor_D_computed;
|
|
cutlass::HostTensor<ElementC, LayoutC> tensor_D_reference;
|
|
|
|
//
|
|
// Methods
|
|
//
|
|
|
|
/// Allocates workspace in device memory
|
|
TestbedComplex() {
|
|
|
|
tensor_A.reset(cutlass::make_Coord(ThreadblockShape::kM, ThreadblockShape::kK));
|
|
tensor_B.reset(cutlass::make_Coord(ThreadblockShape::kK, ThreadblockShape::kN));
|
|
tensor_C.reset(cutlass::make_Coord(Shape::kM, Shape::kN));
|
|
tensor_D_computed.reset(cutlass::make_Coord(Shape::kM, Shape::kN));
|
|
tensor_D_reference.reset(cutlass::make_Coord(Shape::kM, Shape::kN), false);
|
|
}
|
|
|
|
/// Returns true if the CUDA device is sufficient to execute the kernel.
|
|
bool sufficient() const {
|
|
|
|
cudaDeviceProp properties;
|
|
int device_idx;
|
|
cudaError_t result = cudaGetDevice(&device_idx);
|
|
|
|
if (result != cudaSuccess) {
|
|
throw std::runtime_error("cudaGetDevice() API call failed.");
|
|
}
|
|
|
|
result = cudaGetDeviceProperties(&properties, device_idx);
|
|
|
|
if (result != cudaSuccess) {
|
|
throw std::runtime_error("cudaGetDeviceProperties() failed");
|
|
}
|
|
|
|
if (properties.major == 9) {
|
|
// NVIDIA Hopper drops support for several data types
|
|
if (
|
|
cutlass::sizeof_bits<ElementA>::value < 8 ||
|
|
cutlass::sizeof_bits<ElementB>::value < 8 ||
|
|
cutlass::sizeof_bits<ElementC>::value < 8) {
|
|
|
|
return false;
|
|
}
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
/// Runs the test
|
|
bool run(
|
|
cutlass::Distribution::Kind init_A = cutlass::Distribution::Uniform,
|
|
cutlass::Distribution::Kind init_B = cutlass::Distribution::Uniform) {
|
|
|
|
if (!sufficient()) {
|
|
return true;
|
|
}
|
|
|
|
//
|
|
// initialize device memory
|
|
//
|
|
|
|
if (init_A == cutlass::Distribution::Uniform) {
|
|
uint64_t seed = 7;
|
|
cutlass::reference::host::TensorFillRandomUniform(tensor_A.host_view(),
|
|
seed, 8, -8, 0);
|
|
} else if (init_A == cutlass::Distribution::Sequential) {
|
|
cutlass::reference::host::BlockFillSequential(tensor_A.host_data(),
|
|
tensor_A.capacity());
|
|
} else if (init_A == cutlass::Distribution::Identity) {
|
|
cutlass::reference::host::TensorFillIdentity(tensor_A.host_view());
|
|
} else {
|
|
return false;
|
|
}
|
|
|
|
if (init_B == cutlass::Distribution::Uniform) {
|
|
uint64_t seed = 7;
|
|
cutlass::reference::host::TensorFillRandomUniform(tensor_B.host_view(),
|
|
seed + 16, 8, -8, 0);
|
|
} else if (init_B == cutlass::Distribution::Sequential) {
|
|
cutlass::reference::host::BlockFillSequential(tensor_B.host_data(),
|
|
tensor_B.capacity());
|
|
} else if (init_B == cutlass::Distribution::Identity) {
|
|
cutlass::reference::host::TensorFillIdentity(tensor_B.host_view());
|
|
} else {
|
|
return false;
|
|
}
|
|
|
|
cutlass::reference::host::TensorFill(
|
|
tensor_C.host_view(),
|
|
ElementC(0)
|
|
);
|
|
|
|
cutlass::reference::host::TensorFill(
|
|
tensor_D_computed.host_view(),
|
|
ElementC(0)
|
|
);
|
|
|
|
cutlass::reference::host::TensorFill(
|
|
tensor_D_reference.host_view(),
|
|
ElementC(0)
|
|
);
|
|
|
|
tensor_A.sync_device();
|
|
tensor_B.sync_device();
|
|
tensor_C.sync_device();
|
|
tensor_D_computed.sync_device();
|
|
|
|
// launch kernel
|
|
kernel<Mma, ThreadblockShape><<< dim3(1, 1), dim3(32, 1, 1) >>>(
|
|
tensor_D_computed.device_data(),
|
|
tensor_A.device_data(),
|
|
tensor_B.device_data(),
|
|
tensor_C.device_data());
|
|
|
|
// verify no errors
|
|
cudaError_t result = cudaDeviceSynchronize();
|
|
|
|
EXPECT_EQ(result, cudaSuccess) << "CUDA ERROR: " << cudaGetErrorString(result);
|
|
if (result != cudaSuccess) {
|
|
return false;
|
|
}
|
|
|
|
tensor_D_computed.sync_host();
|
|
|
|
//
|
|
// Reference implementation
|
|
//
|
|
|
|
cutlass::reference::host::GemmComplex(
|
|
{Shape::kM, Shape::kN, ThreadblockShape::kK},
|
|
ElementC(1),
|
|
tensor_A.host_ref(),
|
|
Mma::kTransformA,
|
|
tensor_B.host_ref(),
|
|
Mma::kTransformB,
|
|
ElementC(0),
|
|
tensor_C.host_ref(),
|
|
tensor_D_reference.host_ref()
|
|
);
|
|
|
|
//
|
|
// Verify equivalence
|
|
//
|
|
|
|
// compare
|
|
bool passed = cutlass::reference::host::TensorEquals(
|
|
tensor_D_computed.host_view(),
|
|
tensor_D_reference.host_view()
|
|
);
|
|
|
|
EXPECT_TRUE(passed);
|
|
|
|
if (!passed) {
|
|
|
|
cutlass::TensorView<ElementA, cutlass::layout::ColumnMajor> tensor_A_physical(
|
|
tensor_A.host_data(),
|
|
tensor_A.stride()[0],
|
|
tensor_A.extent());
|
|
|
|
cutlass::TensorView<ElementB, cutlass::layout::RowMajor> tensor_B_physical(
|
|
tensor_B.host_data(),
|
|
tensor_B.stride()[0],
|
|
tensor_B.extent());
|
|
|
|
std::cout <<"cutlass::sizeof_bits<ElementA>::value = "<<cutlass::sizeof_bits<ElementA>::value<<"\n";
|
|
std::cout
|
|
<< "A:\n" << tensor_A.host_view() << "\n\n"
|
|
<< "A(physical - stride: " << tensor_A.stride()[0] << ", extent: " << tensor_A.extent() << "):\n" << tensor_A_physical << "\n\n";
|
|
|
|
std::cout <<"cutlass::sizeof_bits<ElementB>::value = "<<cutlass::sizeof_bits<ElementB>::value<<"\n";
|
|
std::cout
|
|
<< "B:\n" << tensor_B.host_view() << "\n\n"
|
|
<< "B(physical - stride: " << tensor_B.stride()[0] << ", extent: " << tensor_B.extent() <<"):\n" << tensor_B_physical << "\n\n";
|
|
|
|
std::cout
|
|
<< "C:\n" << tensor_C.host_view() << "\n\n"
|
|
<< "Reference:\n" << tensor_D_reference.host_view() << "\n\n"
|
|
<< "Computed:\n" << tensor_D_computed.host_view() << std::endl;
|
|
}
|
|
|
|
return passed;
|
|
}
|
|
};
|
|
|
|
/////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
/// Test kernel
|
|
template <typename Mma, typename ThreadblockShape>
|
|
__global__ void kernel_transform(
|
|
typename Mma::ElementC *output_C,
|
|
typename Mma::ElementA const *input_A,
|
|
typename Mma::ElementB const *input_B,
|
|
typename Mma::ElementC const *input_C,
|
|
int iterations = 1) {
|
|
|
|
// Use AlignedBuffer to store trivially copyable objects in unions and __shared__ buffers.
|
|
__shared__ cutlass::AlignedBuffer<
|
|
typename Mma::ElementA, ThreadblockShape::kM * ThreadblockShape::kK> smem_buffer_A;
|
|
|
|
__shared__ cutlass::AlignedBuffer<
|
|
typename Mma::ElementB, ThreadblockShape::kN * ThreadblockShape::kK> smem_buffer_B;
|
|
|
|
if (threadIdx.x == 0) {
|
|
typename Mma::ElementA *smem_ptr_A = smem_buffer_A.data();
|
|
#pragma unroll 1
|
|
for (size_t i = 0; i < smem_buffer_A.size(); ++i) {
|
|
cutlass::ReferenceFactory<typename Mma::ElementA>::get(smem_ptr_A, i) =
|
|
cutlass::ReferenceFactory<typename cutlass::platform::remove_const<
|
|
typename Mma::ElementA>::type>::get(input_A, i);
|
|
}
|
|
|
|
typename Mma::ElementB *smem_ptr_B = smem_buffer_B.data();
|
|
#pragma unroll 1
|
|
for (size_t i = 0; i < smem_buffer_B.size(); ++i) {
|
|
cutlass::ReferenceFactory<typename Mma::ElementB>::get(smem_ptr_B, i) =
|
|
cutlass::ReferenceFactory<typename cutlass::platform::remove_const<
|
|
typename Mma::ElementB>::type>::get(input_B, i);
|
|
}
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
//
|
|
// Construct warp-level matrix product
|
|
//
|
|
|
|
using FragmentA = typename Mma::FragmentA;
|
|
using FragmentB = typename Mma::FragmentB;
|
|
using FragmentC = typename Mma::FragmentC;
|
|
|
|
using TransformedFragmentA = typename Mma::TransformedFragmentA;
|
|
using TransformedFragmentB = typename Mma::TransformedFragmentB;
|
|
|
|
typename Mma::LayoutA layout_A = Mma::LayoutA::packed({ThreadblockShape::kM, ThreadblockShape::kK});
|
|
typename Mma::LayoutB layout_B = Mma::LayoutB::packed({ThreadblockShape::kK, ThreadblockShape::kN});
|
|
typename Mma::LayoutC layout_C = Mma::LayoutC::packed({Mma::Shape::kM, Mma::Shape::kN});
|
|
|
|
typename Mma::IteratorA iter_A({smem_buffer_A.data(), layout_A}, cutlass::arch::LaneId());
|
|
|
|
typename Mma::IteratorB iter_B({smem_buffer_B.data(), layout_B}, cutlass::arch::LaneId());
|
|
|
|
FragmentA loaded_frag_A;
|
|
FragmentB loaded_frag_B;
|
|
TransformedFragmentA transformed_frag_A;
|
|
TransformedFragmentB transformed_frag_B;
|
|
|
|
FragmentC accum;
|
|
|
|
Mma mma;
|
|
|
|
accum.clear();
|
|
|
|
CUTLASS_PRAGMA_NO_UNROLL
|
|
for (int iter = 0; iter < iterations; ++iter) { // place in loop that is not unrolled
|
|
|
|
CUTLASS_PRAGMA_UNROLL
|
|
for (int k = 0; k < ThreadblockShape::kK;
|
|
k += Mma::Policy::MmaShape::kK) {
|
|
iter_A.load(loaded_frag_A);
|
|
iter_B.load(loaded_frag_B);
|
|
|
|
++iter_A;
|
|
++iter_B;
|
|
|
|
mma.transform(transformed_frag_A, transformed_frag_B, loaded_frag_A,
|
|
loaded_frag_B);
|
|
|
|
mma(accum, transformed_frag_A, transformed_frag_B, accum);
|
|
}
|
|
}
|
|
|
|
typename Mma::IteratorC iter_C({output_C, layout_C}, cutlass::arch::LaneId());
|
|
|
|
iter_C.store(accum);
|
|
}
|
|
|
|
/////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
/// Structure to compute the matrix product
|
|
template <
|
|
/// Warp-level matrix multiply-accumulate
|
|
typename Mma_,
|
|
/// Size of threadblock-scoped shape used to store SMEM
|
|
typename ThreadblockShape_,
|
|
/// The innter product operation performed by GEMM
|
|
typename Operator_ = cutlass::arch::OpMultiplyAdd
|
|
>
|
|
struct TransformTestbed {
|
|
|
|
/// Thread-level matrix multiply-accumulate operator
|
|
using Mma = Mma_;
|
|
using ThreadblockShape = ThreadblockShape_;
|
|
using Operator = Operator_;
|
|
|
|
using Shape = typename Mma::Shape;
|
|
using ElementA = typename Mma::ElementA;
|
|
using LayoutA = typename Mma::LayoutA;
|
|
using ElementB = typename Mma::ElementB;
|
|
using LayoutB = typename Mma::LayoutB;
|
|
using ElementC = typename Mma::ElementC;
|
|
using LayoutC = typename Mma::LayoutC;
|
|
|
|
//
|
|
// Data members
|
|
//
|
|
|
|
cutlass::HostTensor<ElementA, LayoutA> tensor_A;
|
|
cutlass::HostTensor<ElementB, LayoutB> tensor_B;
|
|
cutlass::HostTensor<ElementC, LayoutC> tensor_C;
|
|
cutlass::HostTensor<ElementC, LayoutC> tensor_D_computed;
|
|
cutlass::HostTensor<ElementC, LayoutC> tensor_D_reference;
|
|
|
|
//
|
|
// Methods
|
|
//
|
|
|
|
/// Allocates workspace in device memory
|
|
TransformTestbed() {
|
|
|
|
tensor_A.reset(cutlass::make_Coord(ThreadblockShape::kM, ThreadblockShape::kK));
|
|
tensor_B.reset(cutlass::make_Coord(ThreadblockShape::kK, ThreadblockShape::kN));
|
|
tensor_C.reset(cutlass::make_Coord(Shape::kM, Shape::kN));
|
|
tensor_D_computed.reset(cutlass::make_Coord(Shape::kM, Shape::kN));
|
|
tensor_D_reference.reset(cutlass::make_Coord(Shape::kM, Shape::kN), false);
|
|
}
|
|
|
|
/// Returns true if the CUDA device is sufficient to execute the kernel.
|
|
bool sufficient() const {
|
|
|
|
cudaDeviceProp properties;
|
|
int device_idx;
|
|
cudaError_t result = cudaGetDevice(&device_idx);
|
|
|
|
if (result != cudaSuccess) {
|
|
throw std::runtime_error("cudaGetDevice() API call failed.");
|
|
}
|
|
|
|
result = cudaGetDeviceProperties(&properties, device_idx);
|
|
|
|
if (result != cudaSuccess) {
|
|
throw std::runtime_error("cudaGetDeviceProperties() failed");
|
|
}
|
|
|
|
if (properties.major == 9) {
|
|
// NVIDIA Hopper drops support for several data types
|
|
if (
|
|
cutlass::sizeof_bits<ElementA>::value < 8 ||
|
|
cutlass::sizeof_bits<ElementB>::value < 8 ||
|
|
cutlass::sizeof_bits<ElementC>::value < 8) {
|
|
|
|
return false;
|
|
}
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
/// Runs the test
|
|
bool run(
|
|
cutlass::Distribution::Kind init_A = cutlass::Distribution::Uniform,
|
|
cutlass::Distribution::Kind init_B = cutlass::Distribution::Uniform) {
|
|
|
|
if (!sufficient()) {
|
|
return true;
|
|
}
|
|
|
|
//
|
|
// initialize device memory
|
|
//
|
|
|
|
if (init_A == cutlass::Distribution::Uniform) {
|
|
int scope_max = 8;
|
|
int scope_min = -8;
|
|
|
|
if (cutlass::sizeof_bits<ElementA>::value == 4) {
|
|
scope_max = 2;
|
|
scope_min = -2;
|
|
} else if (cutlass::sizeof_bits<ElementA>::value == 1) {
|
|
scope_max = 2;
|
|
scope_min = 0;
|
|
}
|
|
|
|
uint64_t seed = 7;
|
|
cutlass::reference::host::TensorFillRandomUniform(
|
|
tensor_A.host_view(), seed, scope_max, scope_min, 0);
|
|
} else if (init_A == cutlass::Distribution::Sequential) {
|
|
cutlass::reference::host::BlockFillSequential(tensor_A.host_data(),
|
|
tensor_A.capacity());
|
|
} else if (init_A == cutlass::Distribution::Identity) {
|
|
cutlass::reference::host::TensorFillIdentity(tensor_A.host_view());
|
|
} else {
|
|
return false;
|
|
}
|
|
|
|
if (init_B == cutlass::Distribution::Uniform) {
|
|
int scope_max = 8;
|
|
int scope_min = -8;
|
|
|
|
if (cutlass::sizeof_bits<ElementB>::value == 4) {
|
|
scope_max = 2;
|
|
scope_min = -2;
|
|
} else if (cutlass::sizeof_bits<ElementB>::value == 1) {
|
|
scope_max = 2;
|
|
scope_min = 0;
|
|
}
|
|
|
|
uint64_t seed = 7;
|
|
cutlass::reference::host::TensorFillRandomUniform(
|
|
tensor_B.host_view(), seed + 16, scope_max, scope_min, 0);
|
|
} else if (init_B == cutlass::Distribution::Sequential) {
|
|
cutlass::reference::host::BlockFillSequential(tensor_B.host_data(),
|
|
tensor_B.capacity());
|
|
} else if (init_B == cutlass::Distribution::Identity) {
|
|
cutlass::reference::host::TensorFillIdentity(tensor_B.host_view());
|
|
} else {
|
|
return false;
|
|
}
|
|
|
|
cutlass::reference::host::TensorFill(
|
|
tensor_C.host_view(),
|
|
ElementC(0)
|
|
);
|
|
|
|
cutlass::reference::host::TensorFill(
|
|
tensor_D_computed.host_view(),
|
|
ElementC(0)
|
|
);
|
|
|
|
cutlass::reference::host::TensorFill(
|
|
tensor_D_reference.host_view(),
|
|
ElementC(0)
|
|
);
|
|
|
|
tensor_A.sync_device();
|
|
tensor_B.sync_device();
|
|
tensor_C.sync_device();
|
|
tensor_D_computed.sync_device();
|
|
|
|
// launch kernel
|
|
kernel_transform<Mma, ThreadblockShape><<<dim3(1, 1), dim3(32, 1, 1)>>>(
|
|
tensor_D_computed.device_data(), tensor_A.device_data(),
|
|
tensor_B.device_data(), tensor_C.device_data());
|
|
|
|
// verify no errors
|
|
cudaError_t result = cudaDeviceSynchronize();
|
|
|
|
EXPECT_EQ(result, cudaSuccess) << "CUDA ERROR: " << cudaGetErrorString(result);
|
|
if (result != cudaSuccess) {
|
|
return false;
|
|
}
|
|
|
|
tensor_D_computed.sync_host();
|
|
|
|
//
|
|
// Reference implementation
|
|
//
|
|
|
|
cutlass::reference::host::Gemm<ElementA, LayoutA, ElementB, LayoutB,
|
|
ElementC, LayoutC, ElementC, ElementC,
|
|
Operator>
|
|
reference_gemm;
|
|
|
|
reference_gemm(
|
|
{Shape::kM, Shape::kN, ThreadblockShape::kK},
|
|
ElementC(1),
|
|
tensor_A.host_ref(),
|
|
tensor_B.host_ref(),
|
|
ElementC(0),
|
|
tensor_D_reference.host_ref()
|
|
);
|
|
|
|
//
|
|
// Verify equivalence
|
|
//
|
|
|
|
// compare
|
|
bool passed = cutlass::reference::host::TensorEquals(
|
|
tensor_D_computed.host_view(),
|
|
tensor_D_reference.host_view()
|
|
);
|
|
|
|
EXPECT_TRUE(passed);
|
|
|
|
if (!passed) {
|
|
|
|
cutlass::TensorView<ElementA, cutlass::layout::ColumnMajor> tensor_A_physical(
|
|
tensor_A.host_data(),
|
|
tensor_A.stride()[0],
|
|
tensor_A.extent());
|
|
|
|
cutlass::TensorView<ElementB, cutlass::layout::RowMajor> tensor_B_physical(
|
|
tensor_B.host_data(),
|
|
tensor_B.stride()[0],
|
|
tensor_B.extent());
|
|
|
|
std::cout <<"cutlass::sizeof_bits<ElementA>::value = "<<cutlass::sizeof_bits<ElementA>::value<<"\n";
|
|
std::cout
|
|
<< "A:\n" << tensor_A.host_view() << "\n\n"
|
|
<< "A(physical - stride: " << tensor_A.stride()[0] << ", extent: " << tensor_A.extent() << "):\n" << tensor_A_physical << "\n\n";
|
|
|
|
std::cout <<"cutlass::sizeof_bits<ElementB>::value = "<<cutlass::sizeof_bits<ElementB>::value<<"\n";
|
|
std::cout
|
|
<< "B:\n" << tensor_B.host_view() << "\n\n"
|
|
<< "B(physical - stride: " << tensor_B.stride()[0] << ", extent: " << tensor_B.extent() << "):\n" << tensor_B_physical << "\n\n";
|
|
|
|
std::cout
|
|
<< "C:\n" << tensor_C.host_view() << "\n\n"
|
|
<< "Reference:\n" << tensor_D_reference.host_view() << "\n\n"
|
|
<< "Computed:\n" << tensor_D_computed.host_view() << std::endl;
|
|
}
|
|
|
|
return passed;
|
|
}
|
|
};
|
|
|
|
/////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
/// Structure to compute the matrix product
|
|
template <
|
|
/// Warp-level matrix multiply-accumulate
|
|
typename Mma_,
|
|
/// Size of threadblock-scoped shape used to store SMEM
|
|
typename ThreadblockShape_
|
|
>
|
|
struct TransformedTestbedComplex {
|
|
|
|
/// Thread-level matrix multiply-accumulate operator
|
|
using Mma = Mma_;
|
|
using ThreadblockShape = ThreadblockShape_;
|
|
|
|
using Shape = typename Mma::Shape;
|
|
using ElementA = typename Mma::ElementA;
|
|
using LayoutA = typename Mma::LayoutA;
|
|
using ElementB = typename Mma::ElementB;
|
|
using LayoutB = typename Mma::LayoutB;
|
|
using ElementC = typename Mma::ElementC;
|
|
using LayoutC = typename Mma::LayoutC;
|
|
|
|
//
|
|
// Data members
|
|
//
|
|
|
|
cutlass::HostTensor<ElementA, LayoutA> tensor_A;
|
|
cutlass::HostTensor<ElementB, LayoutB> tensor_B;
|
|
cutlass::HostTensor<ElementC, LayoutC> tensor_C;
|
|
cutlass::HostTensor<ElementC, LayoutC> tensor_D_computed;
|
|
cutlass::HostTensor<ElementC, LayoutC> tensor_D_reference;
|
|
|
|
//
|
|
// Methods
|
|
//
|
|
|
|
/// Allocates workspace in device memory
|
|
TransformedTestbedComplex() {
|
|
|
|
tensor_A.reset(cutlass::make_Coord(ThreadblockShape::kM, ThreadblockShape::kK));
|
|
tensor_B.reset(cutlass::make_Coord(ThreadblockShape::kK, ThreadblockShape::kN));
|
|
tensor_C.reset(cutlass::make_Coord(Shape::kM, Shape::kN));
|
|
tensor_D_computed.reset(cutlass::make_Coord(Shape::kM, Shape::kN));
|
|
tensor_D_reference.reset(cutlass::make_Coord(Shape::kM, Shape::kN), false);
|
|
}
|
|
|
|
/// Returns true if the CUDA device is sufficient to execute the kernel.
|
|
bool sufficient() const {
|
|
|
|
cudaDeviceProp properties;
|
|
int device_idx;
|
|
cudaError_t result = cudaGetDevice(&device_idx);
|
|
|
|
if (result != cudaSuccess) {
|
|
throw std::runtime_error("cudaGetDevice() API call failed.");
|
|
}
|
|
|
|
result = cudaGetDeviceProperties(&properties, device_idx);
|
|
|
|
if (result != cudaSuccess) {
|
|
throw std::runtime_error("cudaGetDeviceProperties() failed");
|
|
}
|
|
|
|
if (properties.major == 9) {
|
|
// NVIDIA Hopper drops support for several data types
|
|
if (
|
|
cutlass::sizeof_bits<ElementA>::value < 8 ||
|
|
cutlass::sizeof_bits<ElementB>::value < 8 ||
|
|
cutlass::sizeof_bits<ElementC>::value < 8) {
|
|
|
|
return false;
|
|
}
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
/// Runs the test
|
|
bool run(
|
|
cutlass::Distribution::Kind init_A = cutlass::Distribution::Uniform,
|
|
cutlass::Distribution::Kind init_B = cutlass::Distribution::Uniform) {
|
|
|
|
if (!sufficient()) {
|
|
return true;
|
|
}
|
|
|
|
//
|
|
// initialize device memory
|
|
//
|
|
|
|
if (init_A == cutlass::Distribution::Uniform) {
|
|
uint64_t seed = 7;
|
|
cutlass::reference::host::TensorFillRandomUniform(tensor_A.host_view(),
|
|
seed, 8, -8, 0);
|
|
} else if (init_A == cutlass::Distribution::Sequential) {
|
|
cutlass::reference::host::BlockFillSequential(tensor_A.host_data(),
|
|
tensor_A.capacity());
|
|
} else if (init_A == cutlass::Distribution::Identity) {
|
|
cutlass::reference::host::TensorFillIdentity(tensor_A.host_view());
|
|
} else {
|
|
return false;
|
|
}
|
|
|
|
if (init_B == cutlass::Distribution::Uniform) {
|
|
uint64_t seed = 7;
|
|
cutlass::reference::host::TensorFillRandomUniform(tensor_B.host_view(),
|
|
seed + 16, 8, -8, 0);
|
|
} else if (init_B == cutlass::Distribution::Sequential) {
|
|
cutlass::reference::host::BlockFillSequential(tensor_B.host_data(),
|
|
tensor_B.capacity());
|
|
} else if (init_B == cutlass::Distribution::Identity) {
|
|
cutlass::reference::host::TensorFillIdentity(tensor_B.host_view());
|
|
} else {
|
|
return false;
|
|
}
|
|
|
|
cutlass::reference::host::TensorFill(
|
|
tensor_C.host_view(),
|
|
ElementC(0)
|
|
);
|
|
|
|
cutlass::reference::host::TensorFill(
|
|
tensor_D_computed.host_view(),
|
|
ElementC(0)
|
|
);
|
|
|
|
cutlass::reference::host::TensorFill(
|
|
tensor_D_reference.host_view(),
|
|
ElementC(0)
|
|
);
|
|
|
|
tensor_A.sync_device();
|
|
tensor_B.sync_device();
|
|
tensor_C.sync_device();
|
|
tensor_D_computed.sync_device();
|
|
|
|
// launch kernel
|
|
kernel_transform<Mma, ThreadblockShape><<< dim3(1, 1), dim3(32, 1, 1) >>>(
|
|
tensor_D_computed.device_data(),
|
|
tensor_A.device_data(),
|
|
tensor_B.device_data(),
|
|
tensor_C.device_data());
|
|
|
|
// verify no errors
|
|
cudaError_t result = cudaDeviceSynchronize();
|
|
|
|
EXPECT_EQ(result, cudaSuccess) << "CUDA ERROR: " << cudaGetErrorString(result);
|
|
if (result != cudaSuccess) {
|
|
return false;
|
|
}
|
|
|
|
tensor_D_computed.sync_host();
|
|
|
|
//
|
|
// Reference implementation
|
|
//
|
|
|
|
cutlass::reference::host::GemmComplex(
|
|
{Shape::kM, Shape::kN, ThreadblockShape::kK},
|
|
ElementC(1),
|
|
tensor_A.host_ref(),
|
|
Mma::kTransformA,
|
|
tensor_B.host_ref(),
|
|
Mma::kTransformB,
|
|
ElementC(0),
|
|
tensor_C.host_ref(),
|
|
tensor_D_reference.host_ref()
|
|
);
|
|
|
|
//
|
|
// Verify equivalence
|
|
//
|
|
|
|
// compare
|
|
bool passed = cutlass::reference::host::TensorEquals(
|
|
tensor_D_computed.host_view(),
|
|
tensor_D_reference.host_view()
|
|
);
|
|
|
|
EXPECT_TRUE(passed);
|
|
|
|
if (!passed) {
|
|
|
|
cutlass::TensorView<ElementA, cutlass::layout::ColumnMajor> tensor_A_physical(
|
|
tensor_A.host_data(),
|
|
tensor_A.stride()[0],
|
|
tensor_A.extent());
|
|
|
|
cutlass::TensorView<ElementB, cutlass::layout::RowMajor> tensor_B_physical(
|
|
tensor_B.host_data(),
|
|
tensor_B.stride()[0],
|
|
tensor_B.extent());
|
|
|
|
std::cout <<"cutlass::sizeof_bits<ElementA>::value = "<<cutlass::sizeof_bits<ElementA>::value<<"\n";
|
|
std::cout
|
|
<< "A:\n" << tensor_A.host_view() << "\n\n"
|
|
<< "A(physical - stride: " << tensor_A.stride()[0] << ", extent: " << tensor_A.extent() << "):\n" << tensor_A_physical << "\n\n";
|
|
|
|
std::cout <<"cutlass::sizeof_bits<ElementB>::value = "<<cutlass::sizeof_bits<ElementB>::value<<"\n";
|
|
std::cout
|
|
<< "B:\n" << tensor_B.host_view() << "\n\n"
|
|
<< "B(physical - stride: " << tensor_B.stride()[0] << ", extent: " << tensor_B.extent() <<"):\n" << tensor_B_physical << "\n\n";
|
|
|
|
std::cout
|
|
<< "C:\n" << tensor_C.host_view() << "\n\n"
|
|
<< "Reference:\n" << tensor_D_reference.host_view() << "\n\n"
|
|
<< "Computed:\n" << tensor_D_computed.host_view() << std::endl;
|
|
}
|
|
|
|
return passed;
|
|
}
|
|
};
|
|
|
|
/////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
/// Test kernel
|
|
template <typename Mma, typename ThreadblockShape>
|
|
__global__ void sparse_kernel(
|
|
typename Mma::ElementC *output_C,
|
|
typename Mma::ElementA const *input_A,
|
|
typename Mma::ElementB const *input_B,
|
|
typename Mma::ElementC const *input_C,
|
|
typename Mma::ElementE const *input_E,
|
|
int iterations = 1) {
|
|
|
|
// Use AlignedBuffer to store trivially copyable objects in unions and __shared__ buffers.
|
|
__shared__ cutlass::AlignedBuffer<typename Mma::ElementA,
|
|
ThreadblockShape::kM *
|
|
ThreadblockShape::kK / Mma::kSparse>
|
|
smem_buffer_A;
|
|
|
|
__shared__ cutlass::AlignedBuffer<
|
|
typename Mma::ElementB, ThreadblockShape::kN * ThreadblockShape::kK> smem_buffer_B;
|
|
|
|
__shared__ cutlass::AlignedBuffer<
|
|
typename Mma::ElementE, Mma::Shape::kM * Mma::Shape::kK /
|
|
Mma::kSparse / Mma::kElementsPerElementE>
|
|
smem_buffer_E;
|
|
|
|
__syncthreads();
|
|
|
|
if (threadIdx.x == 0) {
|
|
typename Mma::ElementA *smem_ptr_A = smem_buffer_A.data();
|
|
#pragma unroll 1
|
|
for (size_t i = 0; i < smem_buffer_A.size(); ++i) {
|
|
cutlass::ReferenceFactory<typename Mma::ElementA>::get(smem_ptr_A, i) =
|
|
cutlass::ReferenceFactory<typename cutlass::platform::remove_const<
|
|
typename Mma::ElementA>::type>::get(input_A, i);
|
|
}
|
|
|
|
typename Mma::ElementB *smem_ptr_B = smem_buffer_B.data();
|
|
#pragma unroll 1
|
|
for (size_t i = 0; i < smem_buffer_B.size(); ++i) {
|
|
cutlass::ReferenceFactory<typename Mma::ElementB>::get(smem_ptr_B, i) =
|
|
cutlass::ReferenceFactory<typename cutlass::platform::remove_const<
|
|
typename Mma::ElementB>::type>::get(input_B, i);
|
|
}
|
|
|
|
typename Mma::ElementE *smem_ptr_E = smem_buffer_E.data();
|
|
#pragma unroll 1
|
|
for (size_t i = 0; i < smem_buffer_E.size(); ++i) {
|
|
cutlass::ReferenceFactory<typename Mma::ElementE>::get(smem_ptr_E, i) =
|
|
cutlass::ReferenceFactory<typename cutlass::platform::remove_const<
|
|
typename Mma::ElementE>::type>::get(input_E, i);
|
|
}
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
//
|
|
// Construct warp-level matrix product
|
|
//
|
|
|
|
using FragmentA = typename Mma::FragmentA;
|
|
using FragmentB = typename Mma::FragmentB;
|
|
using FragmentC = typename Mma::FragmentC;
|
|
using FragmentE = typename Mma::FragmentE;
|
|
|
|
typename Mma::LayoutA layout_A = Mma::LayoutA::packed(
|
|
{ThreadblockShape::kM, ThreadblockShape::kK / Mma::kSparse});
|
|
typename Mma::LayoutB layout_B =
|
|
Mma::LayoutB::packed({ThreadblockShape::kK, ThreadblockShape::kN});
|
|
typename Mma::LayoutC layout_C = Mma::LayoutC::packed({Mma::Shape::kM, Mma::Shape::kN});
|
|
typename Mma::LayoutE layout_E =
|
|
Mma::LayoutE::packed({Mma::Shape::kM * Mma::kInterleaved,
|
|
Mma::Shape::kK / Mma::kSparse /
|
|
Mma::kElementsPerElementE / Mma::kInterleaved});
|
|
|
|
typename Mma::IteratorA iter_A({smem_buffer_A.data(), layout_A}, cutlass::arch::LaneId());
|
|
|
|
typename Mma::IteratorB iter_B({smem_buffer_B.data(), layout_B}, cutlass::arch::LaneId());
|
|
|
|
typename Mma::IteratorE iter_E({smem_buffer_E.data(), layout_E}, cutlass::arch::LaneId());
|
|
|
|
FragmentA frag_A;
|
|
FragmentB frag_B;
|
|
|
|
FragmentC accum;
|
|
|
|
FragmentE frag_E;
|
|
|
|
Mma mma;
|
|
|
|
accum.clear();
|
|
|
|
CUTLASS_PRAGMA_NO_UNROLL
|
|
for (int iter = 0; iter < iterations; ++iter) { // place in loop that is not unrolled
|
|
|
|
CUTLASS_PRAGMA_UNROLL
|
|
for (int k = 0; k < ThreadblockShape::kK;
|
|
k += Mma::Policy::MmaShape::kK) {
|
|
iter_A.load(frag_A);
|
|
iter_B.load(frag_B);
|
|
iter_E.load(frag_E);
|
|
|
|
++iter_A;
|
|
++iter_B;
|
|
++iter_E;
|
|
|
|
mma(accum, frag_A, frag_B, accum, frag_E);
|
|
}
|
|
}
|
|
|
|
typename Mma::IteratorC iter_C({output_C, layout_C}, cutlass::arch::LaneId());
|
|
|
|
iter_C.store(accum);
|
|
}
|
|
|
|
/////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
/// Structure to compute the matrix product
|
|
template <
|
|
/// Warp-level matrix multiply-accumulate
|
|
typename Mma_,
|
|
/// Size of threadblock-scoped shape used to store SMEM
|
|
typename ThreadblockShape_,
|
|
/// The innter product operation performed by GEMM
|
|
typename Operator_ = cutlass::arch::OpMultiplyAdd
|
|
>
|
|
struct SparseTestbed {
|
|
|
|
/// Thread-level matrix multiply-accumulate operator
|
|
using Mma = Mma_;
|
|
using ThreadblockShape = ThreadblockShape_;
|
|
using Operator = Operator_;
|
|
|
|
using Shape = typename Mma::Shape;
|
|
using ElementA = typename Mma::ElementA;
|
|
using LayoutA = typename Mma::LayoutA;
|
|
using ElementB = typename Mma::ElementB;
|
|
using LayoutB = typename Mma::LayoutB;
|
|
using ElementC = typename Mma::ElementC;
|
|
using LayoutC = typename Mma::LayoutC;
|
|
|
|
static int const Sparse = Mma::kSparse;
|
|
static int const MetaSizeInBits = Mma::kMetaSizeInBits;
|
|
static int const MaxID2 = Mma::kMaxID2;
|
|
static int const Interleaved = Mma::kInterleaved;
|
|
|
|
using ElementE = typename Mma::ElementE;
|
|
|
|
static int const ElementsPerElementE = Mma::kElementsPerElementE;
|
|
|
|
using LayoutE = cutlass::layout::RowMajor;
|
|
using ReorderedLayoutE =
|
|
cutlass::layout::ColumnMajorInterleaved<Interleaved>;
|
|
|
|
//
|
|
// Data members
|
|
//
|
|
|
|
cutlass::HostTensor<ElementA, LayoutA> tensor_A;
|
|
cutlass::HostTensor<ElementA, LayoutA> tensor_A_uncompressed;
|
|
cutlass::HostTensor<ElementB, LayoutB> tensor_B;
|
|
cutlass::HostTensor<ElementC, LayoutC> tensor_C;
|
|
cutlass::HostTensor<ElementC, LayoutC> tensor_D_computed;
|
|
cutlass::HostTensor<ElementC, LayoutC> tensor_D_reference;
|
|
cutlass::HostTensor<ElementE, LayoutE> tensor_E;
|
|
cutlass::HostTensor<ElementE, ReorderedLayoutE> tensor_E_reordered;
|
|
|
|
//
|
|
// Methods
|
|
//
|
|
|
|
/// Allocates workspace in device memory
|
|
SparseTestbed() {
|
|
|
|
tensor_A.reset(cutlass::make_Coord(ThreadblockShape::kM,
|
|
ThreadblockShape::kK / Sparse));
|
|
tensor_A_uncompressed.reset(
|
|
cutlass::make_Coord(ThreadblockShape::kM, ThreadblockShape::kK));
|
|
tensor_B.reset(cutlass::make_Coord(ThreadblockShape::kK, ThreadblockShape::kN));
|
|
tensor_C.reset(cutlass::make_Coord(Shape::kM, Shape::kN));
|
|
tensor_D_computed.reset(cutlass::make_Coord(Shape::kM, Shape::kN));
|
|
tensor_D_reference.reset(cutlass::make_Coord(Shape::kM, Shape::kN), false);
|
|
tensor_E.reset(cutlass::make_Coord(
|
|
Shape::kM, Shape::kK / Sparse / ElementsPerElementE));
|
|
tensor_E_reordered.reset(cutlass::make_Coord(
|
|
Shape::kM, Shape::kK / Sparse / ElementsPerElementE));
|
|
}
|
|
|
|
/// Returns true if the CUDA device is sufficient to execute the kernel.
|
|
bool sufficient() const {
|
|
|
|
cudaDeviceProp properties;
|
|
int device_idx;
|
|
cudaError_t result = cudaGetDevice(&device_idx);
|
|
|
|
if (result != cudaSuccess) {
|
|
throw std::runtime_error("cudaGetDevice() API call failed.");
|
|
}
|
|
|
|
result = cudaGetDeviceProperties(&properties, device_idx);
|
|
|
|
if (result != cudaSuccess) {
|
|
throw std::runtime_error("cudaGetDeviceProperties() failed");
|
|
}
|
|
|
|
if (properties.major == 9) {
|
|
// NVIDIA Hopper drops support for several data types
|
|
if (
|
|
cutlass::sizeof_bits<ElementA>::value < 8 ||
|
|
cutlass::sizeof_bits<ElementB>::value < 8 ||
|
|
cutlass::sizeof_bits<ElementC>::value < 8) {
|
|
|
|
return false;
|
|
}
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
/// Runs the test
|
|
bool run(
|
|
cutlass::Distribution::Kind init_A = cutlass::Distribution::Uniform,
|
|
cutlass::Distribution::Kind init_B = cutlass::Distribution::Uniform,
|
|
cutlass::Distribution::Kind init_E = cutlass::Distribution::Uniform) {
|
|
|
|
if (!sufficient()) {
|
|
return true;
|
|
}
|
|
|
|
//
|
|
// initialize device memory
|
|
//
|
|
|
|
if (init_A == cutlass::Distribution::Uniform) {
|
|
int scope_max = 8;
|
|
int scope_min = -8;
|
|
|
|
if (cutlass::sizeof_bits<ElementA>::value == 4) {
|
|
scope_max = 2;
|
|
scope_min = -2;
|
|
} else if (cutlass::sizeof_bits<ElementA>::value == 1) {
|
|
scope_max = 2;
|
|
scope_min = 0;
|
|
}
|
|
|
|
uint64_t seed = 7;
|
|
cutlass::reference::host::TensorFillRandomUniform(
|
|
tensor_A.host_view(), seed, scope_max, scope_min, 0);
|
|
} else if (init_A == cutlass::Distribution::Sequential) {
|
|
cutlass::reference::host::BlockFillSequential(tensor_A.host_data(),
|
|
tensor_A.capacity());
|
|
} else if (init_A == cutlass::Distribution::Identity) {
|
|
cutlass::reference::host::TensorFillIdentity(tensor_A.host_view());
|
|
} else {
|
|
return false;
|
|
}
|
|
|
|
if (init_B == cutlass::Distribution::Uniform) {
|
|
int scope_max = 8;
|
|
int scope_min = -8;
|
|
|
|
if (cutlass::sizeof_bits<ElementB>::value == 4) {
|
|
scope_max = 2;
|
|
scope_min = -2;
|
|
} else if (cutlass::sizeof_bits<ElementB>::value == 1) {
|
|
scope_max = 2;
|
|
scope_min = 0;
|
|
}
|
|
|
|
uint64_t seed = 7;
|
|
cutlass::reference::host::TensorFillRandomUniform(
|
|
tensor_B.host_view(), seed + 16, scope_max, scope_min, 0);
|
|
} else if (init_B == cutlass::Distribution::Sequential) {
|
|
cutlass::reference::host::BlockFillSequential(tensor_B.host_data(),
|
|
tensor_B.capacity());
|
|
} else if (init_B == cutlass::Distribution::Identity) {
|
|
cutlass::reference::host::TensorFillIdentity(tensor_B.host_view());
|
|
} else {
|
|
return false;
|
|
}
|
|
|
|
cutlass::reference::host::TensorFill(
|
|
tensor_C.host_view(),
|
|
ElementC(0)
|
|
);
|
|
|
|
cutlass::reference::host::TensorFill(
|
|
tensor_D_computed.host_view(),
|
|
ElementC(0)
|
|
);
|
|
|
|
cutlass::reference::host::TensorFill(
|
|
tensor_D_reference.host_view(),
|
|
ElementC(0)
|
|
);
|
|
|
|
if (init_E == cutlass::Distribution::Uniform) {
|
|
uint64_t seed = 7;
|
|
cutlass::reference::host::TensorFillRandomSparseMeta(
|
|
tensor_E.host_view(), seed, MetaSizeInBits);
|
|
} else if (init_E == cutlass::Distribution::Identity) {
|
|
uint32_t content = (MaxID2 == 1) ? 0x44444444 : 0x4444;
|
|
cutlass::reference::host::TensorFill(tensor_E.host_view(),
|
|
(ElementE)(content));
|
|
} else {
|
|
return false;
|
|
}
|
|
|
|
cutlass::reorder_meta(
|
|
tensor_E_reordered.host_ref(), tensor_E.host_ref(),
|
|
{Shape::kM, Shape::kN, Shape::kK / Sparse / ElementsPerElementE});
|
|
|
|
tensor_A.sync_device();
|
|
tensor_B.sync_device();
|
|
tensor_C.sync_device();
|
|
tensor_D_computed.sync_device();
|
|
tensor_E_reordered.sync_device();
|
|
|
|
// launch kernel
|
|
sparse_kernel<Mma, ThreadblockShape><<< dim3(1, 1), dim3(32, 1, 1) >>>(
|
|
tensor_D_computed.device_data(),
|
|
tensor_A.device_data(),
|
|
tensor_B.device_data(),
|
|
tensor_C.device_data(),
|
|
tensor_E_reordered.device_data());
|
|
|
|
// verify no errors
|
|
cudaError_t result = cudaDeviceSynchronize();
|
|
|
|
EXPECT_EQ(result, cudaSuccess) << "CUDA ERROR: " << cudaGetErrorString(result);
|
|
if (result != cudaSuccess) {
|
|
return false;
|
|
}
|
|
|
|
tensor_D_computed.sync_host();
|
|
|
|
//
|
|
// Reference implementation
|
|
//
|
|
cutlass::uncompress(tensor_A_uncompressed.host_ref(), tensor_A.host_ref(),
|
|
tensor_E.host_ref(), Shape::kM, Shape::kK);
|
|
|
|
cutlass::reference::host::Gemm<ElementA, LayoutA, ElementB, LayoutB,
|
|
ElementC, LayoutC, ElementC, ElementC,
|
|
Operator>
|
|
reference_gemm;
|
|
|
|
reference_gemm(
|
|
{Shape::kM, Shape::kN, ThreadblockShape::kK},
|
|
ElementC(1),
|
|
tensor_A_uncompressed.host_ref(),
|
|
tensor_B.host_ref(),
|
|
ElementC(0),
|
|
tensor_D_reference.host_ref()
|
|
);
|
|
|
|
//
|
|
// Verify equivalence
|
|
//
|
|
|
|
// compare
|
|
bool passed = cutlass::reference::host::TensorEquals(
|
|
tensor_D_computed.host_view(),
|
|
tensor_D_reference.host_view()
|
|
);
|
|
|
|
EXPECT_TRUE(passed);
|
|
|
|
if (!passed) {
|
|
std::cout <<"cutlass::sizeof_bits<ElementA>::value = "<<cutlass::sizeof_bits<ElementA>::value<<"\n";
|
|
std::cout << "A:\n" << tensor_A.host_view() << "\n\n";
|
|
|
|
std::cout <<"cutlass::sizeof_bits<ElementB>::value = "<<cutlass::sizeof_bits<ElementB>::value<<"\n";
|
|
std::cout << "B:\n" << tensor_B.host_view() << "\n\n";
|
|
|
|
std::cout <<"cutlass::sizeof_bits<ElementB>::value = "<<cutlass::sizeof_bits<ElementE>::value<<"\n";
|
|
std::cout << "E:\n" << tensor_E.host_view() << "\n\n";
|
|
|
|
std::cout
|
|
<< "C:\n" << tensor_C.host_view() << "\n\n"
|
|
<< "Reference:\n" << tensor_D_reference.host_view() << "\n\n"
|
|
<< "Computed:\n" << tensor_D_computed.host_view() << "\n";
|
|
}
|
|
|
|
return passed;
|
|
}
|
|
};
|
|
|
|
/////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
} // namespace warp
|
|
} // namespace gemm
|
|
} // namespace test
|