CUTLASS 3.6.0 (#1850)
* v3.6 * update changelog * update readme * fix typo * fixing typos * hopper gemm with weight prefetch --------- Co-authored-by: yuzhai <yuzhai@nvidia.com> Co-authored-by: Haicheng Wu <haichengw@nvidia.com>
This commit is contained in:
Yujia Zhai
co-authored by
yuzhai
Haicheng Wu
yuzhai
Haicheng Wu
parent
0837a2a00a
commit
cc3c29a81a
@@ -28,7 +28,6 @@
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
/*! \file
|
||||
\brief Transform Kernel Universal adapter
|
||||
*/
|
||||
@@ -37,12 +36,25 @@
|
||||
|
||||
// common
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/cluster_launch.hpp"
|
||||
#include "cutlass/device_kernel.h"
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
#include "cutlass/detail/layout.hpp"
|
||||
#include "cutlass/detail/mma.hpp"
|
||||
#include "cutlass/cuda_host_adapter.hpp"
|
||||
|
||||
#include "cutlass/kernel_launch.h"
|
||||
#if !defined(__CUDACC_RTC__)
|
||||
#include "cutlass/cluster_launch.hpp"
|
||||
#include "cutlass/trace.h"
|
||||
#endif // !defined(__CUDACC_RTC__)
|
||||
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass::transform::device {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <class TransformKernel_>
|
||||
class TransformUniversalAdapter
|
||||
{
|
||||
@@ -50,58 +62,73 @@ public:
|
||||
using TransformKernel = TransformKernel_;
|
||||
using Arguments = typename TransformKernel::Arguments;
|
||||
using Params = typename TransformKernel::Params;
|
||||
static bool const kEnableCudaHostAdapter = CUTLASS_ENABLE_CUDA_HOST_ADAPTER;
|
||||
|
||||
|
||||
private:
|
||||
|
||||
/// Kernel API parameters object
|
||||
Params params_;
|
||||
static constexpr bool const EnableCudaHostAdapter = CUTLASS_ENABLE_CUDA_HOST_ADAPTER;
|
||||
|
||||
public:
|
||||
|
||||
/// Access the Params structure
|
||||
Params const& params() const {
|
||||
return this->params_;
|
||||
return params_;
|
||||
}
|
||||
|
||||
/// Determines whether the GEMM can execute the given problem.
|
||||
static Status
|
||||
can_implement(Arguments const& args) {
|
||||
return TransformKernel::can_implement(args);
|
||||
}
|
||||
|
||||
/// Gets the workspace size
|
||||
static size_t
|
||||
get_workspace_size(Arguments const& args) {
|
||||
return TransformKernel::get_workspace_size(args);
|
||||
size_t workspace_bytes = 0;
|
||||
workspace_bytes += TransformKernel::get_workspace_size(args);
|
||||
|
||||
CUTLASS_TRACE_HOST(" workspace_bytes: " << workspace_bytes);
|
||||
|
||||
return workspace_bytes;
|
||||
}
|
||||
|
||||
/// Computes the grid shape
|
||||
static dim3
|
||||
get_grid_shape(Arguments const& args, void* workspace = nullptr) {
|
||||
auto tmp_params = TransformKernel::to_underlying_arguments(args, workspace);
|
||||
return TransformKernel::get_grid_shape(tmp_params);
|
||||
}
|
||||
|
||||
/// Computes the grid shape
|
||||
static dim3
|
||||
get_grid_shape(Params const& params) {
|
||||
return TransformKernel::get_grid_shape(params);
|
||||
}
|
||||
|
||||
|
||||
/// Initializes GEMM state from arguments.
|
||||
Status
|
||||
initialize(
|
||||
Arguments & args,
|
||||
Arguments const& args,
|
||||
void* workspace = nullptr,
|
||||
cudaStream_t stream = nullptr,
|
||||
CudaHostAdapter *cuda_adapter = nullptr) {
|
||||
CudaHostAdapter* cuda_adapter = nullptr) {
|
||||
|
||||
CUTLASS_TRACE_HOST("TransformUniversalAdapter::initialize() - workspace "
|
||||
<< workspace << ", stream: " << (stream ? "non-null" : "null"));
|
||||
<< workspace << ", stream: " << (stream ? "non-null" : "null")
|
||||
<< ", EnableCudaHostAdapter: " << (kEnableCudaHostAdapter ? "True" : "false"));
|
||||
|
||||
// Initialize the workspace
|
||||
Status status = TransformKernel::initialize_workspace(args, workspace, stream, cuda_adapter);
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
|
||||
// Initialize the Params structure
|
||||
this->params_ = TransformKernel::to_underlying_arguments(args, workspace);
|
||||
|
||||
params_ = TransformKernel::to_underlying_arguments(args, workspace);
|
||||
// Don't set the function attributes - require the CudaHostAdapter to set it.
|
||||
if constexpr (EnableCudaHostAdapter) {
|
||||
if constexpr (kEnableCudaHostAdapter) {
|
||||
CUTLASS_ASSERT(cuda_adapter);
|
||||
return Status::kSuccess;
|
||||
}
|
||||
@@ -116,50 +143,59 @@ public:
|
||||
if (smem_size >= (48 << 10)) {
|
||||
CUTLASS_TRACE_HOST(" Setting smem size to " << smem_size);
|
||||
cudaError_t result = cudaFuncSetAttribute(
|
||||
device_kernel<TransformKernel>,
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize,
|
||||
smem_size);
|
||||
device_kernel<TransformKernel>,
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize,
|
||||
smem_size);
|
||||
if (cudaSuccess != result) {
|
||||
result = cudaGetLastError();
|
||||
result = cudaGetLastError(); // to clear the error bit
|
||||
CUTLASS_TRACE_HOST(" cudaFuncSetAttribute() returned error: " << cudaGetErrorString(result));
|
||||
return Status::kErrorInternal;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return Status::kSuccess;
|
||||
}
|
||||
|
||||
static Status
|
||||
run(
|
||||
Params& params,
|
||||
cudaStream_t stream = nullptr,
|
||||
CudaHostAdapter *cuda_adapter = nullptr,
|
||||
int32_t kernel_index = 0) {
|
||||
|
||||
run(Params& params,
|
||||
cudaStream_t stream = nullptr,
|
||||
CudaHostAdapter *cuda_adapter = nullptr,
|
||||
int32_t kernel_index = 0,
|
||||
bool launch_with_pdl = false) {
|
||||
CUTLASS_TRACE_HOST("TransformUniversalAdapter::run()");
|
||||
dim3 const block = TransformKernel::get_block_shape();
|
||||
dim3 const grid = get_grid_shape(params);
|
||||
// Currently only support 1x1x1 for transform kernel.
|
||||
dim3 const cluster = {1,1,1};
|
||||
|
||||
// configure smem size and carveout
|
||||
int smem_size = TransformKernel::SharedStorageSize;
|
||||
|
||||
Status launch_result;
|
||||
|
||||
Status launch_result{ Status::kSuccess };
|
||||
// Use extended launch API only for mainloops that use it
|
||||
if constexpr(TransformKernel::ArchTag::kMinComputeCapability >= 90) {
|
||||
if constexpr (TransformKernel::ArchTag::kMinComputeCapability >= 90) {
|
||||
// Currently only support 1x1x1 for transform kernel.
|
||||
dim3 const cluster = {1,1,1};
|
||||
void* kernel_params[] = {¶ms};
|
||||
|
||||
if constexpr (EnableCudaHostAdapter) {
|
||||
if constexpr (kEnableCudaHostAdapter) {
|
||||
//
|
||||
// Use the cuda host adapter
|
||||
//
|
||||
CUTLASS_ASSERT(cuda_adapter);
|
||||
if (cuda_adapter) {
|
||||
launch_result = cuda_adapter->launch(
|
||||
grid, cluster, block, smem_size, stream, kernel_params, kernel_index);
|
||||
|
||||
if (launch_with_pdl) {
|
||||
CUTLASS_TRACE_HOST(
|
||||
"TransformUniversalAdapter::run() does not support launching with PDL and a custom cuda adapter.");
|
||||
return Status::kErrorInternal;
|
||||
}
|
||||
launch_result = cuda_adapter->launch(grid,
|
||||
cluster,
|
||||
block,
|
||||
smem_size,
|
||||
stream,
|
||||
kernel_params,
|
||||
kernel_index);
|
||||
CUTLASS_TRACE_HOST("Kernel Launch Result" << cutlassGetStatusString(launch_result));
|
||||
}
|
||||
else {
|
||||
return Status::kErrorInternal;
|
||||
@@ -168,18 +204,25 @@ public:
|
||||
else {
|
||||
CUTLASS_ASSERT(cuda_adapter == nullptr);
|
||||
void const* kernel = (void const*) device_kernel<TransformKernel>;
|
||||
launch_result = ClusterLauncher::launch(
|
||||
grid, cluster, block, smem_size, stream, kernel, kernel_params);
|
||||
if constexpr (TransformKernel::ArchTag::kMinComputeCapability == 90) {
|
||||
launch_result = ClusterLauncher::launch(
|
||||
grid, cluster, block, smem_size, stream, kernel, kernel_params, launch_with_pdl);
|
||||
}
|
||||
}
|
||||
}
|
||||
else {
|
||||
launch_result = Status::kSuccess;
|
||||
if constexpr (EnableCudaHostAdapter) {
|
||||
cutlass::arch::synclog_setup();
|
||||
|
||||
if constexpr (kEnableCudaHostAdapter) {
|
||||
CUTLASS_ASSERT(cuda_adapter);
|
||||
if (cuda_adapter) {
|
||||
void* kernel_params[] = {¶ms};
|
||||
|
||||
launch_result = cuda_adapter->launch(
|
||||
grid, block, smem_size, stream, kernel_params, 0);
|
||||
grid, block, smem_size, stream, kernel_params, 0
|
||||
);
|
||||
|
||||
}
|
||||
else {
|
||||
return Status::kErrorInternal;
|
||||
@@ -187,56 +230,74 @@ public:
|
||||
}
|
||||
else {
|
||||
CUTLASS_ASSERT(cuda_adapter == nullptr);
|
||||
device_kernel<TransformKernel><<<grid, block, smem_size, stream>>>(params);
|
||||
cutlass::kernel_launch<TransformKernel>(grid, block, smem_size, stream, params, launch_with_pdl);
|
||||
}
|
||||
}
|
||||
|
||||
cudaError_t result = cudaGetLastError();
|
||||
if (cudaSuccess == result && Status::kSuccess == launch_result) {
|
||||
return Status::kSuccess;
|
||||
} else {
|
||||
CUTLASS_TRACE_HOST(" Kernel launch failed. Reason: " << result);
|
||||
return Status::kErrorInternal;
|
||||
}
|
||||
else if (cudaSuccess != result) {
|
||||
CUTLASS_TRACE_HOST(" Kernel launch failed. Reason: " << cudaGetErrorString(result));
|
||||
}
|
||||
else if (Status::kSuccess != launch_result) {
|
||||
CUTLASS_TRACE_HOST(" Kernel launch failed. Reason: " << cutlassGetStatusString(launch_result));
|
||||
}
|
||||
return Status::kErrorInternal;
|
||||
}
|
||||
|
||||
//
|
||||
// Non-static launch overloads that first create and set the internal params struct of this kernel handle.
|
||||
//
|
||||
|
||||
/// Launches the kernel after first constructing Params internal state from supplied arguments.
|
||||
Status
|
||||
run(
|
||||
Arguments & args,
|
||||
Arguments const& args,
|
||||
void* workspace = nullptr,
|
||||
cudaStream_t stream = nullptr,
|
||||
CudaHostAdapter *cuda_adapter = nullptr,
|
||||
int32_t kernel_index = 0) {
|
||||
|
||||
int32_t kernel_index = 0,
|
||||
bool launch_with_pdl = false
|
||||
) {
|
||||
Status status = initialize(args, workspace, stream, cuda_adapter);
|
||||
|
||||
if (Status::kSuccess == status) {
|
||||
status = run(this->params_, stream, cuda_adapter, kernel_index);
|
||||
status = run(params_, stream, cuda_adapter, kernel_index, launch_with_pdl);
|
||||
}
|
||||
return status;
|
||||
}
|
||||
|
||||
/// Launches the kernel after first constructing Params internal state from supplied arguments.
|
||||
Status
|
||||
operator()(
|
||||
Arguments & args,
|
||||
Arguments const& args,
|
||||
void* workspace = nullptr,
|
||||
cudaStream_t stream = nullptr,
|
||||
CudaHostAdapter *cuda_adapter = nullptr) {
|
||||
return run(args, workspace, stream, cuda_adapter);
|
||||
CudaHostAdapter *cuda_adapter = nullptr,
|
||||
bool launch_with_pdl = false) {
|
||||
return run(args, workspace, stream, cuda_adapter, 0 /*kernel_index*/, launch_with_pdl);
|
||||
}
|
||||
|
||||
/// Overload that allows a user to re-launch the same kernel without updating internal params struct.
|
||||
Status
|
||||
run(
|
||||
cudaStream_t stream = nullptr,
|
||||
CudaHostAdapter *cuda_adapter = nullptr) {
|
||||
return run(this->params_, stream, cuda_adapter);
|
||||
cudaStream_t stream = nullptr,
|
||||
CudaHostAdapter *cuda_adapter = nullptr,
|
||||
bool launch_with_pdl = false) {
|
||||
return run(params_, stream, cuda_adapter, 0 /*kernel_index*/, launch_with_pdl);
|
||||
}
|
||||
|
||||
/// Overload that allows a user to re-launch the same kernel without updating internal params struct.
|
||||
Status
|
||||
operator()(
|
||||
cudaStream_t stream = nullptr,
|
||||
CudaHostAdapter *cuda_adapter = nullptr) {
|
||||
return run(this->params_, stream, cuda_adapter);
|
||||
operator()(cudaStream_t stream = nullptr, CudaHostAdapter *cuda_adapter = nullptr, bool launch_with_pdl = false) {
|
||||
return run(params_, stream, cuda_adapter, 0 /*kernel_index*/, launch_with_pdl);
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace cutlass::transform::device
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -61,13 +61,15 @@ template <
|
||||
FilterFormat SrcFormat,
|
||||
FilterFormat DstFormat,
|
||||
int NumDimensions,
|
||||
class Element,
|
||||
class Element_,
|
||||
int AlignmentBytes = 16
|
||||
>
|
||||
struct ConvFilterFormatTransformer {
|
||||
|
||||
using Element = Element_;
|
||||
static_assert(SrcFormat == FilterFormat::CKTRS, "Currently only source format of CKTRS is supported");
|
||||
static_assert(DstFormat == FilterFormat::CTRSK || DstFormat == FilterFormat::KTRSC, "Currently only destination format of CTRSK/KTRSC is supported");
|
||||
static_assert(AlignmentBytes % static_cast<int>(sizeof(Element)) == 0, "Invalid alignment setting");
|
||||
static_assert(AlignmentBytes > 0 && AlignmentBytes % static_cast<int>(sizeof(Element)) == 0, "Invalid alignment setting");
|
||||
|
||||
// In ktrsc order.
|
||||
using FilterExtent = array<int, NumDimensions>;
|
||||
@@ -108,6 +110,20 @@ struct ConvFilterFormatTransformer {
|
||||
|
||||
static Status
|
||||
can_implement(Arguments const& args) {
|
||||
bool implementable = true;
|
||||
// alignment rule
|
||||
{
|
||||
int contiguous_dim = DstFormat == FilterFormat::CTRSK ? args.filter_extent[0] : args.filter_extent[NumDimensions - 1];
|
||||
int align_element = AlignmentBytes / static_cast<int>(sizeof(Element));
|
||||
|
||||
implementable &= (contiguous_dim % align_element == 0);
|
||||
|
||||
if (!implementable) {
|
||||
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Alignment setting is invalid.\n");
|
||||
return Status::kInvalid;
|
||||
}
|
||||
}
|
||||
|
||||
return Status::kSuccess;
|
||||
}
|
||||
|
||||
@@ -136,7 +152,7 @@ struct ConvFilterFormatTransformer {
|
||||
}
|
||||
|
||||
static Params
|
||||
to_underlying_arguments(Arguments & args, void* workspace) {
|
||||
to_underlying_arguments(Arguments const& args, void* workspace) {
|
||||
auto k = args.filter_extent[0];
|
||||
auto c = args.filter_extent[NumDimensions - 1];
|
||||
auto srt = reverse(take<1,NumDimensions - 1>(args.filter_extent));
|
||||
@@ -192,9 +208,11 @@ struct ConvFilterFormatTransformer {
|
||||
auto kc_coord = DstFormat == FilterFormat::CTRSK ?
|
||||
make_coord(n_idx+i, get<NumDimensions - 2>(srtc_coord)) :
|
||||
make_coord(get<NumDimensions - 2>(srtc_coord), n_idx+i);
|
||||
auto coord = flatten(make_coord(srt_coord, kc_coord));
|
||||
frag(i) = params.src(coord);
|
||||
auto coord = flatten(make_coord(srt_coord, kc_coord));
|
||||
thr_tile_P(i) = elem_less(coord, shape(params.src));
|
||||
if (thr_tile_P(i)) {
|
||||
frag(i) = params.src(coord);
|
||||
}
|
||||
}
|
||||
|
||||
// Copy from RMEM to GMEM
|
||||
|
||||
@@ -0,0 +1,578 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2023 - 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 Compress utils specific for SM90 structure sparse kernels
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cute/container/bit_field.hpp" // cute::bit_field
|
||||
#include "cute/numeric/numeric_types.hpp" // cute::sizeof_bits_v, cute::uint_bit_t
|
||||
#include "cute/tensor.hpp" // cute::Tensor, cute::make_tensor
|
||||
#include "cute/algorithm/cooperative_copy.hpp" // cute::cooperative_copy
|
||||
#include "cutlass/arch/arch.h" // cutlass::arch::Sm90
|
||||
#include "cutlass/cuda_host_adapter.hpp" // cutlass::CudaHostAdapter
|
||||
#include "cutlass/cutlass.h" // cutlass::Status
|
||||
#include "cutlass/gemm/gemm.h" // cutlass::TagToStrideA_t
|
||||
#include "cutlass/fast_math.h" // cutlass::ceil_div, cutlass::round_up
|
||||
#include "cutlass/kernel_hardware_info.h" // cutlass::KernelHardwareInfo
|
||||
#include "cutlass/numeric_size.h" // cutlass::bits_to_bytes
|
||||
#include "cutlass/cuda_host_adapter.hpp" // cutlass::CudaHostAdapter
|
||||
|
||||
namespace cutlass::transform::kernel {
|
||||
|
||||
using namespace cute;
|
||||
|
||||
template<
|
||||
class ProblemShape_,
|
||||
class ElementA_,
|
||||
class LayoutATag_,
|
||||
class SparseConfig_
|
||||
>
|
||||
class SM90StructuredSparseCompressor {
|
||||
public:
|
||||
using SparseConfig = SparseConfig_;
|
||||
using ProblemShape = ProblemShape_;
|
||||
|
||||
// * EltA
|
||||
using ElementA = ElementA_;
|
||||
using ElementAUint = cute::uint_bit_t<cute::sizeof_bits_v<ElementA>>;
|
||||
using ElementAMma = typename SparseConfig::ElementAMma;
|
||||
using ElementAMmaRaw = typename SparseConfig::ElementAMmaRaw;
|
||||
using ElementAMmaRawUnit = cute::uint_bit_t<cute::sizeof_bits_v<ElementAMmaRaw>>;
|
||||
using ElementASparsity = typename SparseConfig::ElementASparsity;
|
||||
using ElementAMmaSparsity = typename SparseConfig::ElementAMmaSparsity;
|
||||
using ElementAUintCompressed = cute::sparse_elem<ElementASparsity{}, ElementAUint>;
|
||||
using LayoutATag = LayoutATag_;
|
||||
using LayoutA = LayoutATag;
|
||||
using StrideA = cutlass::gemm::TagToStrideA_t<LayoutATag>;
|
||||
|
||||
// * EltE
|
||||
using ElementEMma = typename SparseConfig::ElementEMma;
|
||||
using ElementEMmaRaw = typename SparseConfig::ElementEMmaRaw;
|
||||
using ElementEMmaSparsity = typename SparseConfig::ElementEMmaSparsity;
|
||||
// Data Type for storing one chunk's metadata
|
||||
static constexpr int ElementEBitsPerChunk = typename SparseConfig::ElementEBitsPerChunk{};
|
||||
CUTE_STATIC_ASSERT(ElementEBitsPerChunk == 4, "ElementEBitsPerChunk is 4 for SM90");
|
||||
using ElementEChunk = cute::uint_bit_t<ElementEBitsPerChunk>;
|
||||
CUTE_STATIC_ASSERT(cute::is_same_v<ElementEChunk, cute::uint4_t>, "ElementEChunk is uint4_t for SM90");
|
||||
using ElementESparsityPerChunk = Int<ElementEMmaSparsity{} / (cute::sizeof_bits_v<ElementEMmaRaw> / ElementEBitsPerChunk)>;
|
||||
|
||||
// AtomE
|
||||
using TensorEAtom = typename SparseConfig::TensorEAtom;
|
||||
using TensorEAtomK = typename SparseConfig::TensorEAtomK;
|
||||
using TensorEAtomM = typename SparseConfig::TensorEAtomM;
|
||||
|
||||
static constexpr int ElemsARawPerElementAMmaRaw = typename SparseConfig::ElemsARawPerElementAMmaRaw{};
|
||||
static constexpr int LogicalElemsAPerChunk = typename SparseConfig::LogicalElemsAPerChunk{};
|
||||
static constexpr int PhysicalElemsAPerChunk = typename SparseConfig::PhysicalElemsAPerChunk{};
|
||||
static constexpr int LogicalElemsAMmaRawPerChunk = cutlass::ceil_div(LogicalElemsAPerChunk, ElemsARawPerElementAMmaRaw);
|
||||
static constexpr int PhysicalElemsAMmaRawPerChunk = cutlass::ceil_div(PhysicalElemsAPerChunk, ElemsARawPerElementAMmaRaw);
|
||||
|
||||
// * Alignment
|
||||
static constexpr int TensorEAlignmentM = typename SparseConfig::TensorEAlignmentM{};
|
||||
static constexpr int TensorEAlignmentK = typename SparseConfig::TensorEAlignmentK{};
|
||||
static constexpr int TensorAAlignmentK = typename SparseConfig::TensorAAlignmentK{};
|
||||
static constexpr int TensorAAlignmentM = typename SparseConfig::TensorAAlignmentM{};
|
||||
|
||||
// Required by `device_kernel`
|
||||
static constexpr int MaxThreadsPerBlock = TensorEAtomM{};
|
||||
static constexpr int MinBlocksPerMultiprocessor = 1;
|
||||
using ArchTag = arch::Sm90;
|
||||
|
||||
struct SharedStorage {
|
||||
ElementEMma cEsE[cute::size(TensorEAtom{})];
|
||||
ElementAUintCompressed cACsAC[cute::size(TensorEAtom{})];
|
||||
ElementAUint cAsA[cute::size(TensorEAtom{})];
|
||||
};
|
||||
|
||||
static constexpr int SharedStorageSize = sizeof(SharedStorage);
|
||||
|
||||
struct TransformArguments {
|
||||
void const* ptr_A{nullptr};
|
||||
StrideA dA{};
|
||||
void* ptr_ACompress{nullptr};
|
||||
void* ptr_E{nullptr};
|
||||
};
|
||||
|
||||
using TransformParams = TransformArguments;
|
||||
|
||||
struct Arguments {
|
||||
ProblemShape problem_shape{};
|
||||
TransformArguments transform{};
|
||||
KernelHardwareInfo hw_info{};
|
||||
};
|
||||
|
||||
struct Params {
|
||||
ProblemShape problem_shape{};
|
||||
TransformParams transform{};
|
||||
KernelHardwareInfo hw_info{};
|
||||
void* workspace = nullptr;
|
||||
};
|
||||
|
||||
public:
|
||||
static Params
|
||||
to_underlying_arguments(Arguments const& args, void* workspace = nullptr) {
|
||||
CUTLASS_TRACE_HOST("SM90StructuredSparseCompressor::to_underlying_arguments()");
|
||||
return Params{{args.problem_shape},
|
||||
{args.transform.ptr_A, args.transform.dA, args.transform.ptr_ACompress, args.transform.ptr_E},
|
||||
{args.hw_info},
|
||||
workspace};
|
||||
}
|
||||
|
||||
static Status
|
||||
can_implement(Arguments const& args) {
|
||||
auto [M, N, K, L] = args.problem_shape;
|
||||
if (K % LogicalElemsAPerChunk != 0) {
|
||||
CUTLASS_TRACE_HOST("SM90 Sparse Compressor CAN NOT IMPLEMENT: GemmK not multiplier of logical chunk size");
|
||||
return Status::kErrorInvalidProblem;
|
||||
}
|
||||
CUTLASS_TRACE_HOST("SM90StructuredSparseCompressor::can_implement() (True)");
|
||||
return Status::kSuccess;
|
||||
}
|
||||
|
||||
static size_t
|
||||
get_workspace_size(Arguments const& args) {
|
||||
CUTLASS_UNUSED(args);
|
||||
// Backward compatible with host compressor
|
||||
CUTLASS_TRACE_HOST("SM90StructuredSparseCompressor::get_workspace_size() (" << SharedStorageSize << ")");
|
||||
return SharedStorageSize;
|
||||
}
|
||||
|
||||
static Status
|
||||
initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr,
|
||||
CudaHostAdapter *cuda_adapter = nullptr) {
|
||||
CUTLASS_UNUSED(args);
|
||||
CUTLASS_UNUSED(workspace);
|
||||
CUTLASS_UNUSED(stream);
|
||||
CUTLASS_UNUSED(cuda_adapter);
|
||||
CUTLASS_TRACE_HOST("SM90StructuredSparseCompressor::initialize_workspace()");
|
||||
return Status::kSuccess;
|
||||
}
|
||||
|
||||
static dim3
|
||||
get_grid_shape(Params const& params) {
|
||||
constexpr int MaxAlignmentM = cutlass::const_max(TensorEAlignmentM, TensorAAlignmentM);
|
||||
constexpr int MaxAlignmentK = cutlass::const_max(TensorEAlignmentK, TensorAAlignmentK);
|
||||
const auto [GemmM, GemmN, GemmK, GemmL] = params.problem_shape;
|
||||
|
||||
const int GemmMAlignedMax = cutlass::round_up(GemmM, MaxAlignmentM);
|
||||
const int GemmKAlignedMax = cutlass::round_up(GemmK, MaxAlignmentK);
|
||||
|
||||
const int gridDim_X = cutlass::ceil_div(GemmMAlignedMax, TensorEAtomM{});
|
||||
const int gridDim_Y = cutlass::ceil_div(GemmKAlignedMax, TensorEAtomK{});
|
||||
const int gridDim_Z = GemmL;
|
||||
|
||||
CUTLASS_TRACE_HOST("SM90StructuredSparseCompressor::get_grid_shape() ("
|
||||
<< gridDim_X << ", "
|
||||
<< gridDim_Y << ", "
|
||||
<< gridDim_Z << ")");
|
||||
return dim3(gridDim_X, gridDim_Y, gridDim_Z);
|
||||
}
|
||||
|
||||
static dim3
|
||||
get_block_shape() {
|
||||
CUTLASS_TRACE_HOST("SM90StructuredSparseCompressor::get_block_shape() ("
|
||||
<< MaxThreadsPerBlock << ", "
|
||||
<< 1 << ", "
|
||||
<< 1 << ")");
|
||||
return dim3(MaxThreadsPerBlock, 1, 1);
|
||||
}
|
||||
|
||||
CUTE_DEVICE
|
||||
void
|
||||
operator()(Params params, void* smem_buf = nullptr) {
|
||||
run(params, smem_buf);
|
||||
}
|
||||
|
||||
CUTE_DEVICE
|
||||
static void
|
||||
run(Params params, void* smem_buf = nullptr) {
|
||||
structure_sparse_compress(params, smem_buf);
|
||||
}
|
||||
|
||||
private:
|
||||
|
||||
struct MetadataOneChunk1to2 {
|
||||
|
||||
CUTE_DEVICE
|
||||
void set_metadata_bits(int elt_log_idx, int elt_phy_idx) {
|
||||
auto metadata_bits = [&]() -> uint8_t {
|
||||
CUTLASS_ASSERT(elt_log_idx >= 0 && elt_log_idx < 2);
|
||||
switch (elt_log_idx) {
|
||||
case 0:
|
||||
return 0b0100;
|
||||
case 1:
|
||||
return 0b1110;
|
||||
default:
|
||||
CUTE_GCC_UNREACHABLE;
|
||||
}
|
||||
};
|
||||
|
||||
storage_ |= (metadata_bits() << (4 * elt_phy_idx));
|
||||
}
|
||||
|
||||
|
||||
CUTE_DEVICE
|
||||
ElementEChunk storage() const {
|
||||
return ElementEChunk{storage_};
|
||||
}
|
||||
|
||||
private:
|
||||
uint8_t storage_ = 0b0000;
|
||||
};
|
||||
|
||||
struct MetadataOneChunk2to4{
|
||||
|
||||
CUTE_DEVICE
|
||||
void set_metadata_bits(int elt_log_idx, int elt_phy_idx) {
|
||||
auto metadata_bits = [&]() -> uint8_t {
|
||||
CUTLASS_ASSERT(elt_log_idx >= 0 && elt_log_idx < 4);
|
||||
switch (elt_log_idx) {
|
||||
case 0:
|
||||
return 0b00;
|
||||
case 1:
|
||||
return 0b01;
|
||||
case 2:
|
||||
return 0b10;
|
||||
case 3:
|
||||
return 0b11;
|
||||
default:
|
||||
CUTE_GCC_UNREACHABLE;
|
||||
}
|
||||
};
|
||||
|
||||
storage_ |= (metadata_bits() << (2 * elt_phy_idx));
|
||||
}
|
||||
|
||||
CUTE_DEVICE
|
||||
ElementEChunk storage() const {
|
||||
return ElementEChunk{storage_};
|
||||
}
|
||||
|
||||
private:
|
||||
uint8_t storage_ = 0b0000;
|
||||
};
|
||||
|
||||
using MetadataOneChunk = cute::conditional_t<SparseConfig::IsTfmma,
|
||||
MetadataOneChunk1to2,
|
||||
MetadataOneChunk2to4>;
|
||||
|
||||
private:
|
||||
|
||||
CUTE_DEVICE
|
||||
static void
|
||||
structure_sparse_compress(Params params, void* smem_buf) {
|
||||
// * Input Params
|
||||
auto [GemmM, GemmN, GemmK, GemmL] = params.problem_shape;
|
||||
auto [ptr_A, dA, ptr_ACompress, ptr_E] = params.transform;
|
||||
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_buf);
|
||||
|
||||
[[maybe_unused]] const int gridDim_X = gridDim.x;
|
||||
[[maybe_unused]] const int gridDim_Y = gridDim.y;
|
||||
[[maybe_unused]] const int gridDim_Z = gridDim.z;
|
||||
[[maybe_unused]] const int blockDim_X = blockDim.x;
|
||||
|
||||
// * Global Tensor Layout
|
||||
const cute::Layout layout_gA = make_layout(make_shape(GemmM, GemmK, GemmL), dA);
|
||||
const cute::Layout layout_gAC = SparseConfig::fill_layoutA(params.problem_shape);
|
||||
const cute::Layout layout_gE = SparseConfig::fill_layoutE(params.problem_shape);
|
||||
|
||||
// * Construct Global Tensor
|
||||
const cute::Tensor gA = make_tensor(make_gmem_ptr(cute::recast_ptr<ElementAUint>(ptr_A)), layout_gA);
|
||||
cute::Tensor gAC_sparse = make_tensor(make_gmem_ptr(cute::recast_ptr<ElementAUintCompressed>(ptr_ACompress)), layout_gAC );
|
||||
cute::Tensor gAC = cute::recast<ElementAUint>(gAC_sparse);
|
||||
cute::Tensor gE_sparse = make_tensor(make_gmem_ptr(cute::recast_ptr<ElementEMma>(ptr_E)), layout_gE);
|
||||
cute::Tensor gE = cute::recast<ElementEMmaRaw>(gE_sparse);
|
||||
|
||||
// * CTA Tensor Layout
|
||||
using cAsA_layout_row = decltype(make_layout(make_shape(TensorEAtomM{}, TensorEAtomK{}), LayoutRight{}));
|
||||
using cAsA_layout_col = decltype(make_layout(make_shape(TensorEAtomM{}, TensorEAtomK{}), LayoutLeft{}));
|
||||
using cAsA_layout = cute::conditional_t<cute::is_same_v<LayoutATag, layout::RowMajor>, cAsA_layout_row, cAsA_layout_col>;
|
||||
using cACsAC_layout = decltype(make_layout(make_shape(TensorEAtomM{}, TensorEAtomK{} / ElementASparsity{}), LayoutRight{}));
|
||||
using cEsE_layout = decltype(make_layout(make_shape(TensorEAtomM{}, TensorEAtomK{} / ElementEMmaSparsity{}), LayoutRight{}));
|
||||
|
||||
CUTE_STATIC_ASSERT(cute::is_static_v<TensorEAtom>, "TensorEAtom needs to be static");
|
||||
CUTE_STATIC_ASSERT(cute::is_static_v<cAsA_layout>, "cAsA_layout needs to be static");
|
||||
CUTE_STATIC_ASSERT(cute::is_static_v<cACsAC_layout>, "cACsAC_layout needs to be static");
|
||||
CUTE_STATIC_ASSERT(cute::is_static_v<cEsE_layout>, "cEsE_layout needs to be static");
|
||||
|
||||
const int blockIdx_X = blockIdx.x;
|
||||
const int blockIdx_Y = blockIdx.y;
|
||||
const int blockIdx_Z = blockIdx.z;
|
||||
const int threadIdx_X = threadIdx.x;
|
||||
|
||||
// * Construct CTA Tensor
|
||||
const auto cta_coord = make_coord(blockIdx_X, blockIdx_Y, blockIdx_Z);
|
||||
cute::Tensor cAgA = cute::recast<ElementAMmaRawUnit>(local_tile(gA, shape(cAsA_layout{}), cta_coord));
|
||||
cute::Tensor cACgAC = cute::recast<ElementAMmaRawUnit>(local_tile(gAC, shape(cACsAC_layout{}), cta_coord));
|
||||
cute::Tensor cEgE = local_tile(gE, shape(cEsE_layout{}), cta_coord);
|
||||
|
||||
cute::Tensor cAsA = cute::recast<ElementAMmaRawUnit>(make_tensor(make_smem_ptr(cute::recast_ptr<ElementAUint>(shared_storage.cAsA)), cAsA_layout{}));
|
||||
cute::Tensor cACsAC = cute::recast<ElementAMmaRawUnit>(make_tensor(make_smem_ptr(cute::recast_ptr<ElementAUint>(shared_storage.cACsAC)), cACsAC_layout{}));
|
||||
cute::Tensor cEsE = make_tensor(make_smem_ptr(cute::recast_ptr<ElementEMmaRaw>(shared_storage.cEsE)), cEsE_layout{});
|
||||
cute::Tensor cEsE_chunk = cute::recast<ElementEChunk>(cEsE);
|
||||
|
||||
// * Handle in unit of Chunk when compress
|
||||
using OneChunkSizeA = Int<LogicalElemsAMmaRawPerChunk>;
|
||||
using OneChunkSizeAC = Int<PhysicalElemsAMmaRawPerChunk>;
|
||||
using OneChunkSizeE = Int<LogicalElemsAPerChunk / ElementESparsityPerChunk{}>;
|
||||
using NumOneChunkK = Int<cutlass::ceil_div(TensorEAtomK{}, LogicalElemsAPerChunk)>;
|
||||
|
||||
cute::Tensor cAsA_log_chunk = logical_divide(cAsA, make_shape(_, OneChunkSizeA{}));
|
||||
cute::Tensor cACsAC_log_chunk = logical_divide(cACsAC, make_shape(_, OneChunkSizeAC{}));
|
||||
cute::Tensor cEsE_log_chunk = logical_divide(cEsE_chunk, make_shape(_, OneChunkSizeE{}));
|
||||
|
||||
// * Corner Case Handle
|
||||
const auto GemmM_within_Cta = (GemmM - blockIdx_X * TensorEAtomM{} > TensorEAtomM{}) ? TensorEAtomM{} : GemmM - blockIdx_X * TensorEAtomM{};
|
||||
const auto GemmK_within_Cta = ( (GemmK - blockIdx_Y * TensorEAtomK{} > TensorEAtomK{}) ? TensorEAtomK{} : GemmK - blockIdx_Y * TensorEAtomK{} ) / ElemsARawPerElementAMmaRaw;
|
||||
const auto GemmK_NumOneChunk_within_Cta = GemmK_within_Cta / LogicalElemsAMmaRawPerChunk;
|
||||
|
||||
const auto GemmMAlignedAC = cutlass::round_up(GemmM, TensorAAlignmentM);
|
||||
const auto GemmKAlignedAC = cutlass::round_up(GemmK, TensorAAlignmentK);
|
||||
const auto GemmMAlignedAC_within_Cta = (GemmMAlignedAC - blockIdx_X * TensorEAtomM{} > TensorEAtomM{}) ? TensorEAtomM{} : GemmMAlignedAC - blockIdx_X * TensorEAtomM{};
|
||||
const auto GemmKAlignedAC_within_Cta = ( (GemmKAlignedAC - blockIdx_Y * TensorEAtomK{} > TensorEAtomK{}) ? TensorEAtomK{} : GemmKAlignedAC - blockIdx_Y * TensorEAtomK{} ) / ElemsARawPerElementAMmaRaw;
|
||||
|
||||
// * Clear CTA Smem Tensor
|
||||
cooperative_clear<MaxThreadsPerBlock>(threadIdx_X, cACsAC);
|
||||
cooperative_clear<MaxThreadsPerBlock>(threadIdx_X, cEsE);
|
||||
|
||||
// * Input CTA Tensor G to S
|
||||
if (GemmM_within_Cta == TensorEAtomM{} && GemmK_within_Cta == TensorEAtomK{}) {
|
||||
copy_vec_pred<false, LayoutATag>(cAgA, cAsA, threadIdx_X, GemmM_within_Cta, GemmK_within_Cta);
|
||||
}
|
||||
else {
|
||||
copy_vec_pred<true, LayoutATag>(cAgA, cAsA, threadIdx_X, GemmM_within_Cta, GemmK_within_Cta);
|
||||
}
|
||||
|
||||
// * Compress
|
||||
// cACsAC is always row major order
|
||||
// TensorEAtomM threads perform the compression, each thread compress one row
|
||||
const int row_i = threadIdx_X;
|
||||
if (row_i < GemmM_within_Cta) {
|
||||
|
||||
CUTE_UNROLL
|
||||
for (int col_chunk_i = 0; col_chunk_i < NumOneChunkK{}; ++col_chunk_i) {
|
||||
if (col_chunk_i < GemmK_NumOneChunk_within_Cta) {
|
||||
// Compress is handled in unit of ElementAMmaRawUnit
|
||||
cute::Tensor tAsA = cAsA_log_chunk(row_i, make_coord(_, col_chunk_i));
|
||||
cute::Tensor tACsAC = cACsAC_log_chunk(row_i, make_coord(_, col_chunk_i));
|
||||
cute::Tensor tEsE = cEsE_log_chunk(row_i, make_coord(_, col_chunk_i));
|
||||
|
||||
int non_zero_cnt = 0;
|
||||
// None zero element indx
|
||||
// e.g.
|
||||
// 2:4 sparsity [x 0 0 x]
|
||||
// non_zero_elt_log_idx = [0, 3]
|
||||
int non_zero_elt_log_idx[OneChunkSizeAC{}] = { 0 };
|
||||
|
||||
// * Find None Zero Element Idx within Chunk
|
||||
CUTE_UNROLL
|
||||
for (int elt_log_idx = 0; elt_log_idx < OneChunkSizeA{}; ++elt_log_idx) {
|
||||
ElementAMmaRawUnit elem_A = tAsA[elt_log_idx];
|
||||
if ( elem_A != ElementAMmaRawUnit{0} ) {
|
||||
non_zero_elt_log_idx[non_zero_cnt] = elt_log_idx;
|
||||
tACsAC[non_zero_cnt] = elem_A;
|
||||
non_zero_cnt++;
|
||||
}
|
||||
}
|
||||
|
||||
// * Corner Case for 2:4 sparsity
|
||||
if constexpr (cute::sizeof_bits_v<ElementAMmaRawUnit> < 32) {
|
||||
// i.e. [0 0 0 x] -> [(0) 0 0 x]
|
||||
if (non_zero_cnt == 1 && non_zero_elt_log_idx[0] == 3) {
|
||||
tACsAC[1] = tACsAC[0];
|
||||
tACsAC[0] = ElementAMmaRawUnit{0};
|
||||
non_zero_elt_log_idx[0] = 0;
|
||||
non_zero_elt_log_idx[1] = 3;
|
||||
}
|
||||
// i.e. [0 0 x 0] -> [0 0 x (0)]
|
||||
// i.e. [0 x 0 0] -> [0 x 0 (0)]
|
||||
// i.e. [x 0 0 0] -> [x 0 0 (0)]
|
||||
else if (non_zero_cnt == 1) {
|
||||
tACsAC[1] = ElementAMmaRawUnit{0};
|
||||
non_zero_elt_log_idx[1] = 3;
|
||||
}
|
||||
}
|
||||
|
||||
// * Set Metadata Bits
|
||||
MetadataOneChunk metadata_one_chunk;
|
||||
CUTE_UNROLL
|
||||
for (int elt_phy_idx = 0; elt_phy_idx < OneChunkSizeAC{}; elt_phy_idx++) {
|
||||
metadata_one_chunk.set_metadata_bits(non_zero_elt_log_idx[elt_phy_idx], elt_phy_idx);
|
||||
}
|
||||
tEsE[0] = metadata_one_chunk.storage();
|
||||
|
||||
}
|
||||
else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// * Sync after Compress
|
||||
__syncthreads();
|
||||
|
||||
// * Output Cta Tensor S to G
|
||||
if (GemmM_within_Cta > 0 && GemmK_within_Cta > 0) {
|
||||
constexpr int MaxVecBits = 128; // STG.128
|
||||
cute::cooperative_copy<MaxThreadsPerBlock, MaxVecBits>(threadIdx_X, cEsE, cEgE);
|
||||
}
|
||||
|
||||
if (GemmMAlignedAC_within_Cta == TensorEAtomM{} && GemmKAlignedAC_within_Cta == TensorEAtomK{}) {
|
||||
copy_vec_pred<false, LayoutATag>(cACsAC, cACgAC, threadIdx_X, GemmMAlignedAC_within_Cta, (GemmKAlignedAC_within_Cta / ElementASparsity::value));
|
||||
}
|
||||
else {
|
||||
copy_vec_pred<true, LayoutATag>(cACsAC, cACgAC, threadIdx_X, GemmMAlignedAC_within_Cta, (GemmKAlignedAC_within_Cta / ElementASparsity::value));
|
||||
}
|
||||
|
||||
} // end of structure_sparse_compress()
|
||||
|
||||
template<uint32_t NumThreads,
|
||||
typename TensorSrc>
|
||||
CUTE_DEVICE
|
||||
static void
|
||||
cooperative_clear(
|
||||
uint32_t const& tid,
|
||||
TensorSrc dSrc) {
|
||||
|
||||
auto dSrctSrc = local_partition(dSrc, make_layout(make_shape(NumThreads, _1{})), tid);
|
||||
cute::clear(dSrctSrc);
|
||||
|
||||
// Sync all thread data access
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
template <bool pred,
|
||||
typename LayoutTag,
|
||||
typename TensorSrc,
|
||||
typename TensorDst>
|
||||
CUTE_DEVICE
|
||||
static void
|
||||
copy_vec_pred(
|
||||
TensorSrc dSrc,
|
||||
TensorDst dDst,
|
||||
int threadIdx_X,
|
||||
int valid_rows,
|
||||
int valid_cols) {
|
||||
|
||||
constexpr bool IsRowMajor = cute::is_same_v<LayoutTag, cutlass::layout::RowMajor>;
|
||||
using Element = typename TensorSrc::element_type;
|
||||
CUTE_STATIC_ASSERT(cute::is_static_v<decltype(shape(dSrc))>, "shape(dSrc) needs to be static");
|
||||
CUTE_STATIC_ASSERT(cute::is_static_v<decltype(shape(dDst))>, "shape(dDst) needs to be static");
|
||||
CUTE_STATIC_ASSERT(cute::sizeof_bits_v<typename TensorSrc::element_type> == cute::sizeof_bits_v<typename TensorDst::element_type>,
|
||||
"dSrc and dDst need to have same element bit width");
|
||||
CUTE_STATIC_ASSERT(cute::size(dSrc) == cute::size(dDst), "dSrc and dDst need to have same size");
|
||||
|
||||
// ValueShape
|
||||
using ValueShape =
|
||||
cute::conditional_t<IsRowMajor,
|
||||
Shape<Int<1>, Int<128 / sizeof_bits_v<Element>>>,
|
||||
Shape<Int<128 / sizeof_bits_v<Element>>, Int<1>>>
|
||||
;
|
||||
|
||||
constexpr int ValueShapeRows = shape<0>(ValueShape{});
|
||||
constexpr int ValueShapeCols = shape<1>(ValueShape{});
|
||||
|
||||
// ThreadShape
|
||||
using ThreadShape =
|
||||
cute::conditional_t<IsRowMajor,
|
||||
Shape<Int<MaxThreadsPerBlock / (shape<1>(dSrc) / ValueShapeCols)>, Int< (shape<1>(dSrc) / ValueShapeCols)>>,
|
||||
Shape<Int< (shape<0>(dSrc) / ValueShapeRows)>, Int<MaxThreadsPerBlock / (shape<0>(dSrc) / ValueShapeRows)>>>
|
||||
;
|
||||
|
||||
constexpr int ThreadShapeRows = shape<0>(ThreadShape{});
|
||||
constexpr int ThreadShapeCols = shape<1>(ThreadShape{});
|
||||
|
||||
const int threadIdx_X_row = threadIdx_X / ThreadShapeCols;
|
||||
const int threadIdx_X_col = threadIdx_X % ThreadShapeCols;
|
||||
|
||||
// Row Major
|
||||
if constexpr (IsRowMajor) {
|
||||
CUTE_UNROLL
|
||||
for (int iter_row_blk = 0; iter_row_blk < cutlass::ceil_div(shape<0>(dSrc), ThreadShapeRows * ValueShapeRows); ++iter_row_blk) {
|
||||
CUTE_UNROLL
|
||||
for (int col_chunk_i = 0; col_chunk_i < cutlass::ceil_div(shape<1>(dSrc) , ThreadShapeCols * ValueShapeCols); ++col_chunk_i) {
|
||||
CUTE_UNROLL
|
||||
for (int iter_row_thr = 0; iter_row_thr < ValueShapeRows; ++iter_row_thr) {
|
||||
CUTE_UNROLL
|
||||
for (int iter_col_thr = 0; iter_col_thr < ValueShapeCols; ++iter_col_thr) {
|
||||
const int row_i = (iter_row_blk * ThreadShapeRows + threadIdx_X_row) * ValueShapeRows + iter_row_thr;
|
||||
const int col_i = (col_chunk_i * ThreadShapeCols + threadIdx_X_col) * ValueShapeCols + iter_col_thr;
|
||||
if constexpr ( (not pred)
|
||||
) {
|
||||
dDst(row_i, col_i) = dSrc(row_i, col_i);
|
||||
}
|
||||
else {
|
||||
if (row_i < valid_rows && col_i < valid_cols) {
|
||||
dDst(row_i, col_i) = dSrc(row_i, col_i);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// Col Major
|
||||
else {
|
||||
CUTE_UNROLL
|
||||
for (int col_chunk_i = 0; col_chunk_i < cutlass::ceil_div(shape<1>(dSrc) , ThreadShapeCols * ValueShapeCols); ++col_chunk_i) {
|
||||
CUTE_UNROLL
|
||||
for (int iter_row_blk = 0; iter_row_blk < cutlass::ceil_div(shape<0>(dSrc), ThreadShapeRows * ValueShapeRows); ++iter_row_blk) {
|
||||
CUTE_UNROLL
|
||||
for (int iter_col_thr = 0; iter_col_thr < ValueShapeCols; ++iter_col_thr) {
|
||||
CUTE_UNROLL
|
||||
for (int iter_row_thr = 0; iter_row_thr < ValueShapeRows; ++iter_row_thr) {
|
||||
const int row_i = (iter_row_blk * ThreadShapeRows + threadIdx_X_row) * ValueShapeRows + iter_row_thr;
|
||||
const int col_i = (col_chunk_i * ThreadShapeCols + threadIdx_X_col) * ValueShapeCols + iter_col_thr;
|
||||
if constexpr ( (not pred)
|
||||
) {
|
||||
dDst(row_i, col_i) = dSrc(row_i, col_i);
|
||||
}
|
||||
else {
|
||||
if (row_i < valid_rows && col_i < valid_cols) {
|
||||
dDst(row_i, col_i) = dSrc(row_i, col_i);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Sync all thread data access
|
||||
__syncthreads();
|
||||
} // end of copy_vec_pred()
|
||||
|
||||
};
|
||||
|
||||
} // namespace cutlass::transform::kernel
|
||||
@@ -0,0 +1,284 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2023 - 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 Compress utils for structured sparse kernels
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <algorithm> // std::fill
|
||||
#include <array> // std::array
|
||||
#include <random> // std::mt19937
|
||||
|
||||
#include "cute/numeric/numeric_types.hpp" // cute::sizeof_bits_v
|
||||
#include "cute/tensor.hpp" // cute::Tensor, cute::make_tensor
|
||||
#include "cutlass/arch/arch.h" // cutlass::arch::SmXY
|
||||
#include "cutlass/gemm/gemm.h" // cutlass::TagToStrideA_t
|
||||
#include "cutlass/fast_math.h" // cutlass::ceil_div, cutlass::round_up
|
||||
#include "cutlass/numeric_size.h" // cutlass::bits_to_bytes
|
||||
|
||||
#include "cutlass/transform/kernel/sm90_sparse_gemm_compressor.hpp"
|
||||
|
||||
namespace cutlass::transform::kernel {
|
||||
|
||||
template<
|
||||
class ProblemShape_,
|
||||
class ElementA_,
|
||||
class LayoutATag_,
|
||||
class SparseConfig_
|
||||
>
|
||||
class StructuredSparseCompressorUtility {
|
||||
public:
|
||||
using SparseConfig = SparseConfig_;
|
||||
using ProblemShape = ProblemShape_;
|
||||
|
||||
//* EltA
|
||||
using ElementA = ElementA_;
|
||||
using LayoutATag = LayoutATag_;
|
||||
using StrideA = cutlass::gemm::TagToStrideA_t<LayoutATag>;
|
||||
using ElementAMmaRaw = typename SparseConfig::ElementAMmaRaw;
|
||||
using ElementASparsity = typename SparseConfig::ElementASparsity;
|
||||
using ElementAMmaSparsity = typename SparseConfig::ElementAMmaSparsity;
|
||||
|
||||
//* EltE
|
||||
using ElementEMmaRaw = typename SparseConfig::ElementEMmaRaw;
|
||||
using ElementEMmaSparsity = typename SparseConfig::ElementEMmaSparsity;
|
||||
|
||||
//* AtomE
|
||||
using TensorEAtom = typename SparseConfig::TensorEAtom;
|
||||
using TensorEAtomK = typename SparseConfig::TensorEAtomK;
|
||||
using TensorEAtomM = typename SparseConfig::TensorEAtomM;
|
||||
|
||||
static constexpr int ElemsARawPerElementAMmaRaw = typename SparseConfig::ElemsARawPerElementAMmaRaw{};
|
||||
static constexpr int LogicalElemsAPerChunk = typename SparseConfig::LogicalElemsAPerChunk{};
|
||||
static constexpr int PhysicalElemsAPerChunk = typename SparseConfig::PhysicalElemsAPerChunk{};
|
||||
static constexpr int LogicalElemsAMmaRawPerChunk = cutlass::ceil_div(LogicalElemsAPerChunk, ElemsARawPerElementAMmaRaw);
|
||||
static constexpr int PhysicalElemsAMmaRawPerChunk = cutlass::ceil_div(PhysicalElemsAPerChunk, ElemsARawPerElementAMmaRaw);
|
||||
|
||||
//* Alignment
|
||||
static constexpr int TensorEAlignmentM = typename SparseConfig::TensorEAlignmentM{};
|
||||
static constexpr int TensorEAlignmentK = typename SparseConfig::TensorEAlignmentK{};
|
||||
static constexpr int TensorAAlignmentK = typename SparseConfig::TensorAAlignmentK{};
|
||||
static constexpr int TensorAAlignmentM = typename SparseConfig::TensorAAlignmentM{};
|
||||
|
||||
StructuredSparseCompressorUtility() = default;
|
||||
|
||||
StructuredSparseCompressorUtility(ProblemShape problem, StrideA dA) {
|
||||
set_problem_size(problem, dA);
|
||||
}
|
||||
|
||||
void set_problem_size(ProblemShape problem, StrideA dA_) {
|
||||
M = cute::size<0>(problem);
|
||||
K = cute::size<2>(problem);
|
||||
L = cute::size<3>(problem);
|
||||
|
||||
// The following three vars are logical elem count!
|
||||
K_alignedA = round_up(K, TensorAAlignmentK);
|
||||
M_alignedA = round_up(M, TensorAAlignmentM);
|
||||
K_alignedE = round_up(K, TensorEAlignmentK);
|
||||
M_alignedE = round_up(M, TensorEAlignmentM);
|
||||
|
||||
dA = dA_;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Get the TensorE number of ElementE along K after alignment requirement
|
||||
*
|
||||
* @return int : number of ElementE (uint8_t) along K-dim
|
||||
*/
|
||||
int get_metadata_m_physical() const {
|
||||
return M_alignedE;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Get the TensorE number of ElementE along M after alignment requirement
|
||||
*
|
||||
* @return int : number of ElementE (uint8_t) along M-dim
|
||||
*/
|
||||
int get_metadata_k_physical() const {
|
||||
return K_alignedE / ElementEMmaSparsity{};
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Get the TensorACompressed number of ElementA along K after alignment requirement
|
||||
*
|
||||
* @return int : number of ElementA along K-dim
|
||||
*/
|
||||
int get_tensorA_k_physical() const {
|
||||
return K_alignedA / ElementASparsity{};
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Get the TensorACompressed number of ElementA along M after alignment requirement
|
||||
*
|
||||
* @return int : number of ElementA along M-dim
|
||||
*/
|
||||
int get_tensorA_m_physical() const {
|
||||
return M_alignedA;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Get the TensorACompressed Bytes
|
||||
*
|
||||
* @return uint64_t bytes
|
||||
*/
|
||||
uint64_t get_compressed_tensor_A_bytes() const {
|
||||
const auto tensor_a_comp_num_elt_a = get_tensorA_m_physical() * get_tensorA_k_physical() * L;
|
||||
const auto tensor_a_comp_bytes = cutlass::bits_to_bytes<uint64_t>(tensor_a_comp_num_elt_a * cute::sizeof_bits_v<ElementA>);
|
||||
return tensor_a_comp_bytes;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Get the TensorA Bytes
|
||||
*
|
||||
* @return uint64_t bytes
|
||||
*/
|
||||
uint64_t get_raw_tensor_A_bytes() const {
|
||||
const auto tensor_a_num_elt_a = uint64_t(M) * uint64_t(K) * uint64_t(L);
|
||||
const auto tensor_a_bytes = cutlass::bits_to_bytes<uint64_t>(tensor_a_num_elt_a * cute::sizeof_bits_v<ElementA>);
|
||||
return tensor_a_bytes;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Get the TensorE Bytes
|
||||
*
|
||||
* @return uint64_t bytes
|
||||
*/
|
||||
uint64_t get_tensor_E_bytes() const {
|
||||
const auto tensor_e_num_elt_a = uint64_t(get_metadata_m_physical()) * uint64_t(get_metadata_k_physical()) * uint64_t(L);
|
||||
const auto tensor_e_bytes = cutlass::bits_to_bytes<uint64_t>(tensor_e_num_elt_a * cute::sizeof_bits_v<ElementEMmaRaw>);
|
||||
return tensor_e_bytes;
|
||||
}
|
||||
|
||||
constexpr auto fill_layoutA_from_compressor() const {
|
||||
return SparseConfig::fill_layoutA(cute::make_tuple(M,_1{},K,L));
|
||||
}
|
||||
|
||||
constexpr auto fill_layoutE_from_compressor() const {
|
||||
return SparseConfig::fill_layoutE(cute::make_tuple(M,_1{},K,L));
|
||||
}
|
||||
|
||||
void structure_sparse_zero_mask_fill(void* host_a_ptr, uint64_t seed) {
|
||||
|
||||
constexpr int ChunkSize = LogicalElemsAMmaRawPerChunk;
|
||||
using ChunkElement = cute::uint_bit_t<cute::sizeof_bits_v<ElementAMmaRaw>>;
|
||||
|
||||
cute::Tensor gA_eltA = cute::make_tensor(
|
||||
cute::recast_ptr<ElementA>(host_a_ptr),
|
||||
cute::make_layout(make_shape(M, K, L), dA));
|
||||
|
||||
// Input TensorA is handled in unit of ElementAMmaRaw instead of ElementA
|
||||
cute::Tensor gA = cute::recast<ChunkElement>(gA_eltA);
|
||||
|
||||
// Extract out the Chunk from K-mode
|
||||
Tensor gA_chunk = cute::zipped_divide(gA, cute::Shape<_1,cute::Int<ChunkSize>>{}); // (Chunk, Rest)
|
||||
|
||||
// Half of the data is zero to indicate sparsityA = 2
|
||||
std::array<int, ChunkSize> nnzb_indicator{};
|
||||
for (size_t i = 1; i < nnzb_indicator.size(); i += 2) {
|
||||
nnzb_indicator.at(i) = 1;
|
||||
}
|
||||
|
||||
std::mt19937 rng(seed);
|
||||
auto rest_shape = cute::shape<1>(gA_chunk);
|
||||
for (auto iter = cute::make_coord_iterator(rest_shape); iter != cute::ForwardCoordIteratorSentinel{}; ++iter) {
|
||||
std::shuffle(nnzb_indicator.begin(), nnzb_indicator.end(), rng);
|
||||
for (int c = 0; c < size<0>(gA_chunk); ++c) { // for each elem within chunk
|
||||
if (nnzb_indicator[c] == 0) {
|
||||
gA_chunk(c, *iter) = ChunkElement{0};
|
||||
}
|
||||
} // end of within chunk
|
||||
} // end of chunk_idx
|
||||
}
|
||||
|
||||
int M{-1};
|
||||
int K{-1};
|
||||
int L{-1};
|
||||
StrideA dA{};
|
||||
|
||||
private:
|
||||
int K_alignedA{-1};
|
||||
int M_alignedA{-1};
|
||||
int K_alignedE{-1};
|
||||
int M_alignedE{-1};
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template<
|
||||
class ProblemShape,
|
||||
class ElementA,
|
||||
class LayoutATag,
|
||||
class SparseConfig,
|
||||
class ArchTag
|
||||
>
|
||||
struct StructuredSparseCompressorSelector {
|
||||
static_assert(cutlass::detail::dependent_false<ArchTag>,
|
||||
"Could not select a structured sparse compressor for given parameters.");
|
||||
};
|
||||
|
||||
template<
|
||||
class ProblemShape,
|
||||
class ElementA,
|
||||
class LayoutATag,
|
||||
class SparseConfig
|
||||
>
|
||||
struct StructuredSparseCompressorSelector<
|
||||
ProblemShape,
|
||||
ElementA,
|
||||
LayoutATag,
|
||||
SparseConfig,
|
||||
arch::Sm90> {
|
||||
using Compressor = SM90StructuredSparseCompressor<
|
||||
ProblemShape,
|
||||
ElementA,
|
||||
LayoutATag,
|
||||
SparseConfig
|
||||
>;
|
||||
};
|
||||
|
||||
template<
|
||||
class ProblemShape,
|
||||
class ElementA,
|
||||
class LayoutATag,
|
||||
class SparseConfig,
|
||||
class ArchTag
|
||||
>
|
||||
using StructuredSparseCompressor = typename StructuredSparseCompressorSelector<
|
||||
ProblemShape,
|
||||
ElementA,
|
||||
LayoutATag,
|
||||
SparseConfig,
|
||||
ArchTag
|
||||
>::Compressor;
|
||||
|
||||
} // End namespace cutlass::transform::kernel
|
||||
Reference in New Issue
Block a user