1036 lines
41 KiB
C++
1036 lines
41 KiB
C++
/***************************************************************************************************
|
|
* Copyright (c) 2023 - 2025 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.
|
|
*
|
|
**************************************************************************************************/
|
|
#pragma once
|
|
|
|
#include "cutlass/cutlass.h"
|
|
#include "cutlass/numeric_conversion.h"
|
|
|
|
#include "cute/util/type_traits.hpp"
|
|
#include "cute/arch/copy_sm90.hpp"
|
|
#include "cute/numeric/arithmetic_tuple.hpp"
|
|
|
|
|
|
/////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
namespace cutlass {
|
|
|
|
// The universal converter
|
|
template <
|
|
class SrcType,
|
|
class DstType,
|
|
class LayoutIn,
|
|
class LayoutOut
|
|
>
|
|
struct LayoutAwareConvertImpl {
|
|
template<class EngineIn, class EngineOut>
|
|
CUTLASS_DEVICE
|
|
static void convert(
|
|
cute::Tensor<EngineIn, LayoutIn> const& src,
|
|
cute::Tensor<EngineOut, LayoutOut> & dst) {
|
|
|
|
static_assert(cute::is_same_v<SrcType, typename EngineIn::value_type> &&
|
|
cute::is_same_v<DstType, typename EngineOut::value_type>);
|
|
static_assert(cute::cosize_v<LayoutIn> == cute::cosize_v<LayoutOut>);
|
|
constexpr int N = decltype(cute::max_common_vector(LayoutIn{}, LayoutOut{})){};
|
|
using SrcArray = cutlass::Array<SrcType, N>;
|
|
using DstArray = cutlass::Array<DstType, N>;
|
|
using Converter = cutlass::NumericArrayConverter<DstType,
|
|
SrcType,
|
|
N,
|
|
cutlass::FloatRoundStyle::round_to_nearest>;
|
|
auto&& src_vm = cute::recast<SrcArray>(src);
|
|
auto&& dst_vm = cute::recast<DstArray>(dst);
|
|
CUTLASS_PRAGMA_UNROLL
|
|
for (int i = 0; i < src_vm.size(); ++i) {
|
|
dst_vm(i) = Converter::convert(src_vm(i));
|
|
}
|
|
}
|
|
};
|
|
|
|
// Specialization for INT4 -> BF16 with [02461357] value order
|
|
template <>
|
|
struct LayoutAwareConvertImpl<
|
|
cutlass::int4b_t,
|
|
cutlass::bfloat16_t,
|
|
cute::Layout<cute::Shape<_2,_4>, cute::Stride<_4,_1>>,
|
|
cute::Layout<_8>
|
|
> {
|
|
template<class EngineIn, class EngineOut>
|
|
CUTLASS_DEVICE
|
|
static void convert(
|
|
cute::Tensor<EngineIn,
|
|
cute::Layout<cute::Shape<_2,_4>, cute::Stride<_4,_1>>
|
|
> const& src,
|
|
cute::Tensor<EngineOut,
|
|
cute::Layout<_8>
|
|
>& dst) {
|
|
|
|
static_assert(cute::is_same_v<cutlass::int4b_t, typename EngineIn::value_type> &&
|
|
cute::is_same_v<cutlass::bfloat16_t, typename EngineOut::value_type>);
|
|
using SrcArray = cutlass::Array<cutlass::int4b_t, 8>;
|
|
using DstArray = cutlass::Array<cutlass::bfloat16_t, 8>;
|
|
using RegArray = cutlass::AlignedArray<uint32_t, 4, sizeof(DstArray)>;
|
|
|
|
auto&& src_reg = cute::recast<uint32_t>(src)(0);
|
|
auto&& r = cute::recast<RegArray>(dst)(0);
|
|
CUTLASS_PRAGMA_UNROLL
|
|
for (size_t ii = 0; ii < RegArray::kElements; ++ii) {
|
|
r[ii] = src_reg >> (4 * (ii));
|
|
static constexpr uint32_t xor_mask = 0x43084308;
|
|
static constexpr uint32_t lo_mask = 0x000F000F;
|
|
static constexpr uint32_t immLut = (0xf0 & 0xcc) ^ 0xaa;
|
|
asm volatile(
|
|
"{\n"
|
|
" lop3.b32 %0, %0, %1, %2, %3;\n"
|
|
"}\n"
|
|
: "+r"(r[ii])
|
|
: "n"(lo_mask), "n"(xor_mask), "n"(immLut));
|
|
static constexpr uint32_t lo_bias = xor_mask; // 0x43084308, {136, 136}
|
|
{
|
|
__nv_bfloat162& bf16x2_val = reinterpret_cast<__nv_bfloat162&>(r[ii]);
|
|
bf16x2_val = __hsub2(bf16x2_val,
|
|
reinterpret_cast<const __nv_bfloat162&>(lo_bias));
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
// Specialization for UINT4 -> BF16 with [02461357] value order
|
|
template <>
|
|
struct LayoutAwareConvertImpl<
|
|
cutlass::uint4b_t,
|
|
cutlass::bfloat16_t,
|
|
cute::Layout<cute::Shape<_2,_4>, cute::Stride<_4,_1>>,
|
|
cute::Layout<_8>
|
|
> {
|
|
template<class EngineIn, class EngineOut>
|
|
CUTLASS_DEVICE
|
|
static void convert(
|
|
cute::Tensor<EngineIn,
|
|
cute::Layout<cute::Shape<_2,_4>, cute::Stride<_4,_1>>
|
|
> const& src,
|
|
cute::Tensor<EngineOut,
|
|
cute::Layout<_8>
|
|
>& dst) {
|
|
|
|
static_assert(cute::is_same_v<cutlass::uint4b_t, typename EngineIn::value_type> &&
|
|
cute::is_same_v<cutlass::bfloat16_t, typename EngineOut::value_type>);
|
|
using SrcArray = cutlass::Array<cutlass::uint4b_t, 8>;
|
|
using DstArray = cutlass::Array<cutlass::bfloat16_t, 8>;
|
|
using RegArray = cutlass::AlignedArray<uint32_t, 4, sizeof(DstArray)>;
|
|
|
|
auto&& src_reg = cute::recast<uint32_t>(src)(0);
|
|
auto&& r = cute::recast<RegArray>(dst)(0);
|
|
CUTLASS_PRAGMA_UNROLL
|
|
for (size_t ii = 0; ii < RegArray::kElements; ++ii) {
|
|
r[ii] = src_reg >> (4 * (ii));
|
|
static constexpr uint32_t or_mask = 0x43004300;
|
|
static constexpr uint32_t lo_mask = 0x000F000F;
|
|
static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa;
|
|
asm volatile(
|
|
"{\n"
|
|
" lop3.b32 %0, %0, %1, %2, %3;\n"
|
|
"}\n"
|
|
: "+r"(r[ii])
|
|
: "n"(lo_mask), "n"(or_mask), "n"(immLut));
|
|
static constexpr uint32_t lo_bias = or_mask; // 0x43004300, {128, 128}
|
|
{
|
|
__nv_bfloat162& bf16x2_val = reinterpret_cast<__nv_bfloat162&>(r[ii]);
|
|
bf16x2_val = __hsub2(bf16x2_val,
|
|
reinterpret_cast<const __nv_bfloat162&>(lo_bias));
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
// Specialization for INT4 -> FP16 with [02461357] value order
|
|
template <>
|
|
struct LayoutAwareConvertImpl<
|
|
cutlass::int4b_t,
|
|
cutlass::half_t,
|
|
cute::Layout<cute::Shape<_2,_4>, cute::Stride<_4,_1>>,
|
|
cute::Layout<_8>
|
|
> {
|
|
template<class EngineIn, class EngineOut>
|
|
CUTLASS_DEVICE
|
|
static void convert(
|
|
cute::Tensor<EngineIn,
|
|
cute::Layout<cute::Shape<_2,_4>, cute::Stride<_4,_1>>
|
|
> const& src,
|
|
cute::Tensor<EngineOut,
|
|
cute::Layout<_8>
|
|
>& dst) {
|
|
|
|
static_assert(cute::is_same_v<cutlass::int4b_t, typename EngineIn::value_type> &&
|
|
cute::is_same_v<cutlass::half_t, typename EngineOut::value_type>);
|
|
using SrcArray = cutlass::Array<cutlass::int4b_t, 8>;
|
|
using DstArray = cutlass::Array<cutlass::half_t, 8>;
|
|
using RegArray = cutlass::AlignedArray<uint32_t, 4, sizeof(DstArray)>;
|
|
|
|
auto&& src_reg = cute::recast<uint32_t>(src)(0);
|
|
auto&& r = cute::recast<RegArray>(dst)(0);
|
|
CUTLASS_PRAGMA_UNROLL
|
|
for (int ii = 0; ii < RegArray::kElements; ii += 2) {
|
|
auto src_ = src_reg >> (4 * (ii));
|
|
r[ii + 0] = src_;
|
|
r[ii + 1] = src_;
|
|
static constexpr uint32_t lo_xor_mask = 0x64086408;
|
|
static constexpr uint32_t hi_xor_mask = 0x64806480;
|
|
static constexpr uint32_t lo_mask = 0x000F000F;
|
|
static constexpr uint32_t hi_mask = 0x00F000F0;
|
|
static constexpr uint32_t immLut = (0xf0 & 0xcc) ^ 0xaa;
|
|
asm volatile(
|
|
"{\n"
|
|
" lop3.b32 %0, %0, %1, %2, %3;\n"
|
|
"}\n"
|
|
: "+r"(r[ii + 0])
|
|
: "n"(lo_mask), "n"(lo_xor_mask), "n"(immLut));
|
|
asm volatile(
|
|
"{\n"
|
|
" lop3.b32 %0, %0, %1, %2, %3;\n"
|
|
"}\n"
|
|
: "+r"(r[ii + 1])
|
|
: "n"(hi_mask), "n"(hi_xor_mask), "n"(immLut));
|
|
static constexpr uint32_t lo_bias = 0x64086408; // {1032, 1032}
|
|
static constexpr uint32_t hi_bias = 0xD480D480; // {-72, -72}
|
|
static constexpr uint32_t hi_scale = 0x2C002C00; // {1/16, 1/16}
|
|
{
|
|
half2& fp16x2_val = reinterpret_cast<__half2&>(r[ii + 0]);
|
|
fp16x2_val = __hsub2(fp16x2_val,
|
|
reinterpret_cast<const half2&>(lo_bias));
|
|
}
|
|
{
|
|
half2& fp16x2_val = reinterpret_cast<__half2&>(r[ii + 1]);
|
|
fp16x2_val = __hfma2(fp16x2_val,
|
|
reinterpret_cast<const half2&>(hi_scale),
|
|
reinterpret_cast<const half2&>(hi_bias));
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
// Specialization for UINT4 -> FP16 with [02461357] value order
|
|
template <>
|
|
struct LayoutAwareConvertImpl<
|
|
cutlass::uint4b_t,
|
|
cutlass::half_t,
|
|
cute::Layout<cute::Shape<_2,_4>, cute::Stride<_4,_1>>,
|
|
cute::Layout<_8>
|
|
> {
|
|
template<class EngineIn, class EngineOut>
|
|
CUTLASS_DEVICE
|
|
static void convert(
|
|
cute::Tensor<EngineIn,
|
|
cute::Layout<cute::Shape<_2,_4>, cute::Stride<_4,_1>>
|
|
> const& src,
|
|
cute::Tensor<EngineOut,
|
|
cute::Layout<_8>
|
|
>& dst) {
|
|
|
|
static_assert(cute::is_same_v<cutlass::uint4b_t, typename EngineIn::value_type> &&
|
|
cute::is_same_v<cutlass::half_t, typename EngineOut::value_type>);
|
|
using SrcArray = cutlass::Array<cutlass::uint4b_t, 8>;
|
|
using DstArray = cutlass::Array<cutlass::half_t, 8>;
|
|
using RegArray = cutlass::AlignedArray<uint32_t, 4, sizeof(DstArray)>;
|
|
|
|
auto&& src_reg = cute::recast<uint32_t>(src)(0);
|
|
auto&& r = cute::recast<RegArray>(dst)(0);
|
|
CUTLASS_PRAGMA_UNROLL
|
|
for (int ii = 0; ii < RegArray::kElements; ii += 2) {
|
|
auto src_ = src_reg >> (4 * (ii));
|
|
r[ii + 0] = src_;
|
|
r[ii + 1] = src_;
|
|
static constexpr uint32_t or_mask = 0x64006400;
|
|
static constexpr uint32_t lo_mask = 0x000F000F;
|
|
static constexpr uint32_t hi_mask = 0x00F000F0;
|
|
static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa;
|
|
asm volatile(
|
|
"{\n"
|
|
" lop3.b32 %0, %0, %1, %2, %3;\n"
|
|
"}\n"
|
|
: "+r"(r[ii])
|
|
: "n"(lo_mask), "n"(or_mask), "n"(immLut));
|
|
asm volatile(
|
|
"{\n"
|
|
" lop3.b32 %0, %0, %1, %2, %3;\n"
|
|
"}\n"
|
|
: "+r"(r[ii + 1])
|
|
: "n"(hi_mask), "n"(or_mask), "n"(immLut));
|
|
static constexpr uint32_t lo_bias = or_mask; // 0x64006400, {1024, 1024}
|
|
static constexpr uint32_t hi_bias = 0xD400D400; // {-64, -64}
|
|
static constexpr uint32_t hi_scale = 0x2C002C00; // {1/16, 1/16}
|
|
{
|
|
half2& fp16x2_val = reinterpret_cast<__half2&>(r[ii + 0]);
|
|
fp16x2_val = __hsub2(fp16x2_val,
|
|
reinterpret_cast<const half2&>(lo_bias));
|
|
}
|
|
{
|
|
half2& fp16x2_val = reinterpret_cast<__half2&>(r[ii + 1]);
|
|
fp16x2_val = __hfma2(fp16x2_val,
|
|
reinterpret_cast<const half2&>(hi_scale),
|
|
reinterpret_cast<const half2&>(hi_bias));
|
|
}
|
|
}
|
|
}
|
|
};
|
|
/*
|
|
// Specialization for E5M2 -> FP16 with [3120] value order
|
|
template <>
|
|
struct LayoutAwareConvertImpl<
|
|
cutlass::float_e5m2_t,
|
|
cutlass::half_t,
|
|
cute::Layout<cute::Shape<_2,_2>, cute::Stride<_2,_1>>,
|
|
cute::Layout<_4>
|
|
> {
|
|
template<class EngineIn, class EngineOut>
|
|
CUTLASS_DEVICE
|
|
static void convert(
|
|
cute::Tensor<EngineIn,
|
|
cute::Layout<cute::Shape<_2,_2>, cute::Stride<_2,_1>>
|
|
> const& src,
|
|
cute::Tensor<EngineOut,
|
|
cute::Layout<_4>
|
|
>& dst) {
|
|
|
|
static_assert(cute::is_same_v<cutlass::float_e5m2_t, typename EngineIn::value_type> &&
|
|
cute::is_same_v<cutlass::half_t, typename EngineOut::value_type>);
|
|
using SrcArray = cutlass::Array<cutlass::float_e5m2_t, 8>;
|
|
using DstArray = cutlass::Array<cutlass::half_t, 8>;
|
|
using RegArray = cutlass::AlignedArray<uint32_t, 4, sizeof(DstArray)>;
|
|
|
|
auto&& src_reg = cute::recast<uint32_t>(src)(0);
|
|
auto&& r = cute::recast<RegArray>(dst)(0);
|
|
CUTLASS_PRAGMA_UNROLL
|
|
for (int ii = 0; ii < RegArray::kElements; ++ii) {
|
|
// in registers: a3, a1, a2, a0
|
|
r[RegArray::kElements - ii - 1] = src_reg << (8 * (ii));
|
|
|
|
static constexpr uint32_t and_mask = 0xFF00FF00;
|
|
asm volatile(
|
|
"{\n"
|
|
" and.b32 %0, %0, %1;\n"
|
|
"}\n"
|
|
: "+r"(r[ii])
|
|
: "n"(and_mask));
|
|
}
|
|
}
|
|
};
|
|
*/
|
|
// Specialization for INT8 -> BF16 with [3120] value order
|
|
template <>
|
|
struct LayoutAwareConvertImpl<
|
|
cutlass::int8_t,
|
|
cutlass::bfloat16_t,
|
|
cute::Layout<cute::Shape<_2,_2>, cute::Stride<_2,_1>>,
|
|
cute::Layout<_4>
|
|
> {
|
|
template<class EngineIn, class EngineOut>
|
|
CUTLASS_DEVICE
|
|
static void convert(
|
|
cute::Tensor<EngineIn,
|
|
cute::Layout<cute::Shape<_2,_2>, cute::Stride<_2,_1>>
|
|
> const& src,
|
|
cute::Tensor<EngineOut,
|
|
cute::Layout<_4>
|
|
>& dst) {
|
|
|
|
static_assert(cute::is_same_v<cutlass::int8_t, typename EngineIn::value_type> &&
|
|
cute::is_same_v<cutlass::bfloat16_t, typename EngineOut::value_type>);
|
|
using SrcArray = cutlass::Array<cutlass::int8_t, 8>;
|
|
using DstArray = cutlass::Array<cutlass::bfloat16_t, 8>;
|
|
using RegArray = cutlass::AlignedArray<uint32_t, 4, sizeof(DstArray)>;
|
|
|
|
auto&& src_reg = cute::recast<uint32_t>(src)(0);
|
|
auto&& r = cute::recast<RegArray>(dst)(0);
|
|
CUTLASS_PRAGMA_UNROLL
|
|
for (int ii = 0; ii < RegArray::kElements; ++ii) {
|
|
uint32_t tmp0, tmp1;
|
|
r[ii] = src_reg >> (8 * (ii));
|
|
static constexpr uint32_t or_mask = 0x43004300;
|
|
static constexpr uint32_t and_mask_0 = 0x007F007F;
|
|
static constexpr uint32_t and_mask_1 = 0x00800080;
|
|
static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa;
|
|
asm volatile(
|
|
"{\n"
|
|
" lop3.b32 %0, %1, %2, %3, %4;\n"
|
|
"}\n"
|
|
: "=r"(tmp0)
|
|
: "r"(r[ii]), "n"(and_mask_0), "n"(or_mask), "n"(immLut));
|
|
asm volatile(
|
|
"{\n"
|
|
" lop3.b32 %0, %1, %2, %3, %4;\n"
|
|
"}\n"
|
|
: "=r"(tmp1)
|
|
: "r"(r[ii]), "n"(and_mask_1), "n"(or_mask), "n"(immLut));
|
|
{
|
|
__nv_bfloat162& bf16x2_val = reinterpret_cast<__nv_bfloat162&>(r[ii]);
|
|
bf16x2_val = __hsub2(reinterpret_cast<__nv_bfloat162 const&>(tmp0),
|
|
reinterpret_cast<__nv_bfloat162 const&>(tmp1));
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
// Specialization for INT8 -> FP16 with [3120] value order
|
|
template <>
|
|
struct LayoutAwareConvertImpl<
|
|
cutlass::int8_t,
|
|
cutlass::half_t,
|
|
cute::Layout<cute::Shape<_2,_2>, cute::Stride<_2,_1>>,
|
|
cute::Layout<_4>
|
|
> {
|
|
template<class EngineIn, class EngineOut>
|
|
CUTLASS_DEVICE
|
|
static void convert(
|
|
cute::Tensor<EngineIn,
|
|
cute::Layout<cute::Shape<_2,_2>, cute::Stride<_2,_1>>
|
|
> const& src,
|
|
cute::Tensor<EngineOut,
|
|
cute::Layout<_4>
|
|
>& dst) {
|
|
|
|
static_assert(cute::is_same_v<cutlass::int8_t, typename EngineIn::value_type> &&
|
|
cute::is_same_v<cutlass::half_t, typename EngineOut::value_type>);
|
|
using SrcArray = cutlass::Array<cutlass::int8_t, 8>;
|
|
using DstArray = cutlass::Array<cutlass::half_t, 8>;
|
|
using RegArray = cutlass::AlignedArray<uint32_t, 4, sizeof(DstArray)>;
|
|
|
|
auto&& src_reg = cute::recast<uint32_t>(src)(0);
|
|
auto&& r = cute::recast<RegArray>(dst)(0);
|
|
CUTLASS_PRAGMA_UNROLL
|
|
for (int ii = 0; ii < RegArray::kElements; ++ii) {
|
|
r[ii] = src_reg >> (8 * (ii));
|
|
static constexpr uint32_t xor_mask = 0x64806480;
|
|
static constexpr uint32_t and_mask = 0x00FF00FF;
|
|
static constexpr uint32_t immLut = (0xf0 & 0xcc) ^ 0xaa;
|
|
asm volatile(
|
|
"{\n"
|
|
" lop3.b32 %0, %0, %1, %2, %3;\n"
|
|
"}\n"
|
|
: "+r"(r[ii])
|
|
: "n"(and_mask), "n"(xor_mask), "n"(immLut));
|
|
{
|
|
static constexpr uint32_t bias = 0x64806480;
|
|
__half2& fp16x2_val = reinterpret_cast<__half2&>(r[ii]);
|
|
fp16x2_val = __hsub2(fp16x2_val,
|
|
reinterpret_cast<__half2 const&>(bias));
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
template <
|
|
class EngineIn,
|
|
class EngineOut,
|
|
class LayoutIn,
|
|
class LayoutOut
|
|
>
|
|
CUTLASS_DEVICE
|
|
void LayoutAwareConvert( // Accept mutable temporaries
|
|
cute::Tensor<EngineIn, LayoutIn> const& src,
|
|
cute::Tensor<EngineOut, LayoutOut> && dst) {
|
|
|
|
LayoutAwareConvert(src, dst);
|
|
}
|
|
template <
|
|
class EngineIn,
|
|
class EngineOut,
|
|
class LayoutIn,
|
|
class LayoutOut
|
|
>
|
|
CUTLASS_DEVICE
|
|
void LayoutAwareConvert(
|
|
cute::Tensor<EngineIn, LayoutIn> const& src,
|
|
cute::Tensor<EngineOut, LayoutOut> & dst) {
|
|
|
|
using SrcType = typename EngineIn::value_type;
|
|
using DstType = typename EngineOut::value_type;
|
|
Tensor src_vm = coalesce(src);
|
|
Tensor dst_vm = coalesce(dst);
|
|
Layout src_layout = src_vm.layout();
|
|
Layout dst_layout = dst_vm.layout();
|
|
LayoutAwareConvertImpl<SrcType,
|
|
DstType,
|
|
decltype(src_layout),
|
|
decltype(dst_layout)>::convert(src_vm, dst_vm);
|
|
}
|
|
|
|
} // namespace cutlass
|
|
|
|
/////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
namespace cutlass {
|
|
namespace detail {
|
|
enum class ConversionMode {
|
|
DirectConvert, // A * B
|
|
ConvertAndScale, // (scale * A) * B
|
|
ConvertAndScaleWithZero // (scale * A + zeros) * B
|
|
};
|
|
} // namespace detail
|
|
} //namespace cutlass
|
|
|
|
/////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
namespace cutlass::gemm::collective::detail {
|
|
|
|
template <class PointerType>
|
|
static constexpr
|
|
CUTLASS_HOST_DEVICE
|
|
auto get_logical_ptr(PointerType const* ptr) {
|
|
return cute::recast_ptr<PointerType const>(ptr);
|
|
}
|
|
template<int Stages, class LayoutAtom, class TileShape, class Stride>
|
|
static constexpr
|
|
CUTLASS_HOST_DEVICE
|
|
auto get_smem_layout(LayoutAtom layout_atom, TileShape const& tile_shape, Stride const& stride) {
|
|
if constexpr (not cute::is_layout<Stride>::value) {
|
|
return tile_to_shape(
|
|
layout_atom,
|
|
append(tile_shape, Int<Stages>{}),
|
|
cute::conditional_t< ::cutlass::gemm::detail::is_major<0,Stride>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{});
|
|
}
|
|
else {
|
|
auto gmem_tile = composition(stride, tile_shape);
|
|
return make_layout_like(append(gmem_tile, make_layout(Int<Stages>{}, 0)));
|
|
}
|
|
}
|
|
template<class Shape, class Stride>
|
|
static constexpr
|
|
CUTLASS_HOST_DEVICE
|
|
auto get_gmem_layout(Shape const& shape, Stride const& stride) {
|
|
if constexpr (not cute::is_layout<Stride>::value) {
|
|
return make_layout(shape, stride);
|
|
}
|
|
else {
|
|
return stride;
|
|
}
|
|
}
|
|
|
|
template<class Collective>
|
|
struct MixedInputUtils {
|
|
private:
|
|
using ConversionMode = cutlass::detail::ConversionMode;
|
|
using KernelSchedule = typename Collective::KernelSchedule;
|
|
using SmemLayoutA = typename Collective::SmemLayoutA;
|
|
using SmemLayoutB = typename Collective::SmemLayoutB;
|
|
using SmemLayoutScale = typename Collective::SmemLayoutScale;
|
|
using SwappedElementA = typename Collective::SwappedElementA;
|
|
using SwappedElementB = typename Collective::SwappedElementB;
|
|
using RealSwappedElementA = typename Collective::RealSwappedElementA;
|
|
using RealSwappedElementB = typename Collective::RealSwappedElementB;
|
|
using ElementScale = typename Collective::ElementScale;
|
|
using ElementZero = typename Collective::ElementZero;
|
|
using SmemCopyAtomScale = typename Collective::SmemCopyAtomScale;
|
|
static constexpr auto KernelConversionMode = Collective::KernelConversionMode;
|
|
static constexpr auto ModeHasScales = Collective::ModeHasScales;
|
|
static constexpr auto UseScaleLookupTable = Collective::UseScaleLookupTable;
|
|
|
|
public:
|
|
static constexpr auto
|
|
elements_per_smem_scale() {
|
|
if constexpr (KernelConversionMode == ConversionMode::DirectConvert) {
|
|
return 0;
|
|
}
|
|
else if constexpr (ModeHasScales) {
|
|
return cute::cosize_v<SmemLayoutScale>;
|
|
}
|
|
else {
|
|
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Type not handled in scale smem allocation.");
|
|
}
|
|
}
|
|
|
|
static constexpr auto
|
|
elements_per_smem_zero() {
|
|
if constexpr (KernelConversionMode == ConversionMode::DirectConvert ||
|
|
KernelConversionMode == ConversionMode::ConvertAndScale ) {
|
|
return 0;
|
|
}
|
|
else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
|
|
return cute::cosize_v<SmemLayoutScale>;
|
|
}
|
|
else {
|
|
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Type not handled in scale smem allocation.");
|
|
}
|
|
}
|
|
|
|
// These methods use some the public members of the class. For that reason, we define them after the public section.
|
|
static constexpr uint32_t
|
|
compute_tma_transaction_bytes_mk() {
|
|
return cutlass::bits_to_bytes(size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * static_cast<uint32_t>(cute::sizeof_bits_v<SwappedElementA>));
|
|
}
|
|
|
|
static constexpr uint32_t
|
|
compute_tma_transaction_bytes_nk() {
|
|
return cutlass::bits_to_bytes(size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast<uint32_t>(cute::sizeof_bits_v<SwappedElementB>));
|
|
}
|
|
|
|
static constexpr uint32_t
|
|
compute_tma_transaction_bytes_extra() {
|
|
if constexpr (KernelConversionMode == ConversionMode::DirectConvert) {
|
|
return 0;
|
|
}
|
|
else if constexpr (ModeHasScales) {
|
|
constexpr uint32_t scale_tx_bytes = cutlass::bits_to_bytes(size<0>(SmemLayoutScale{}) * size<1>(SmemLayoutScale{}) * static_cast<uint32_t>(cute::sizeof_bits_v<ElementScale>));
|
|
static_assert(scale_tx_bytes % 128 == 0, "Each scale stage must be 128B aligned."); // required by TMA
|
|
if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
|
|
return scale_tx_bytes;
|
|
}
|
|
else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
|
|
// Scale and zero share smem layout
|
|
constexpr uint32_t zero_tx_bytes = cutlass::bits_to_bytes(size<0>(SmemLayoutScale{}) * size<1>(SmemLayoutScale{}) * static_cast<uint32_t>(cute::sizeof_bits_v<ElementZero>));
|
|
static_assert(zero_tx_bytes % 128 == 0, "Each zero stage must be 128B aligned."); // required by TMA
|
|
return scale_tx_bytes + zero_tx_bytes;
|
|
}
|
|
else {
|
|
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Type not handled in tma transaction bytes computation.");
|
|
}
|
|
}
|
|
else {
|
|
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Type not handled in tma transaction bytes computation.");
|
|
}
|
|
}
|
|
|
|
/// Utilities to copy A and extra inputs from smem to RF
|
|
template <class SmemTiledCopyA,
|
|
class TensorASmemView,
|
|
class TensorACopyView,
|
|
class... Ts,
|
|
class... Us
|
|
>
|
|
CUTLASS_DEVICE
|
|
static void copy_tensors_MK(
|
|
SmemTiledCopyA const& smem_tiled_copy_A,
|
|
TensorASmemView const& tCsA,
|
|
TensorACopyView& tCrA_copy_view,
|
|
cute::tuple<Ts...> const& partitioned_mma_extra_info,
|
|
cute::tuple<Us...> const& tiled_copy_and_views,
|
|
int k_block,
|
|
int read_stage) {
|
|
|
|
copy(smem_tiled_copy_A, tCsA(_,_,k_block,read_stage), tCrA_copy_view(_,_,k_block));
|
|
|
|
if (k_block == 0) {
|
|
// We are starting a new k-tile so copy the scale
|
|
if constexpr (KernelConversionMode == ConversionMode::DirectConvert) {
|
|
// nothing to do
|
|
}
|
|
else if constexpr (ModeHasScales) {
|
|
auto smem_tiled_copy_S = cute::get<0>(tiled_copy_and_views);
|
|
auto tCrS_copy_view = cute::get<1>(tiled_copy_and_views);
|
|
auto tCsS = cute::get<0>(partitioned_mma_extra_info);
|
|
copy(smem_tiled_copy_S, tCsS(_,_,k_block,read_stage), tCrS_copy_view(_,_,k_block));
|
|
if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
|
|
// Nothing extra to do
|
|
} else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
|
|
auto tCsZ = cute::get<2>(partitioned_mma_extra_info);
|
|
auto tCrZ_copy_view = cute::get<2>(tiled_copy_and_views);
|
|
copy(smem_tiled_copy_S, tCsZ(_,_,k_block,read_stage), tCrZ_copy_view(_,_,k_block));
|
|
} else {
|
|
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in A -> RF path.");
|
|
}
|
|
}
|
|
else {
|
|
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in A -> RF path.");
|
|
}
|
|
}
|
|
}
|
|
|
|
// Helper functions to select packing for conversion
|
|
template <class SrcType,
|
|
class DstType,
|
|
int Cosize>
|
|
struct select_packing { // Naive packing policy
|
|
static constexpr auto value() {
|
|
return Int<cute::gcd(Cosize, 32 / cute::min(sizeof_bits_v<SrcType>, sizeof_bits_v<DstType>))>{};
|
|
}
|
|
};
|
|
|
|
// The core converter uses a lookup table to converts i4 -> 8 bit value.
|
|
template <class EngineIn,
|
|
class LayoutIn,
|
|
class EngineOut,
|
|
class LayoutOut,
|
|
class EngineScale,
|
|
class LayoutScale>
|
|
CUTLASS_DEVICE
|
|
static void lookup_table_convert( // Accept mutable temporaries
|
|
Tensor<EngineIn, LayoutIn> const& src,
|
|
Tensor<EngineOut, LayoutOut> && dst,
|
|
Tensor<EngineScale, LayoutScale> const& scales_neg,
|
|
Tensor<EngineScale, LayoutScale> const& scales_pos) {
|
|
|
|
lookup_table_convert(src, dst, scales_neg, scales_pos);
|
|
}
|
|
template <class EngineIn,
|
|
class LayoutIn,
|
|
class EngineOut,
|
|
class LayoutOut,
|
|
class EngineScale,
|
|
class LayoutScale>
|
|
CUTLASS_DEVICE
|
|
static void lookup_table_convert(
|
|
Tensor<EngineIn, LayoutIn> const& src,
|
|
Tensor<EngineOut, LayoutOut> & dst,
|
|
Tensor<EngineScale, LayoutScale> const& scales_neg,
|
|
Tensor<EngineScale, LayoutScale> const& scales_pos) {
|
|
|
|
constexpr int N = cute::cosize(LayoutIn{});
|
|
static_assert(N == 4 || N == 8);
|
|
static_assert(cosize(LayoutScale{}) <= N / 4,
|
|
"at least 4 consecutive weights must share the same scale.");
|
|
using SrcArray = cutlass::Array<cutlass::int4b_t, 8>;
|
|
using DstArray = cutlass::Array<RealSwappedElementB, 8>;
|
|
using RegArray = cutlass::AlignedArray<uint32_t, N / 4, sizeof(DstArray)>;
|
|
|
|
// View the input as reg
|
|
auto&& src_reg = cute::recast<uint32_t>(src)(0);
|
|
auto&& r = cute::recast<RegArray>(dst)(0);
|
|
|
|
// Determines if to get from the signed or unsigned candidates
|
|
static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa;
|
|
uint32_t sign; // ((reg & 0x88888888) | 0x64206420) >> 1
|
|
asm volatile(
|
|
"{\n"
|
|
" lop3.b32 %0, %1, %2, %3, %4;\n" \
|
|
"}\n"
|
|
: "=r"(sign)
|
|
: "r"(src_reg), "n"(0x88888888), "n"(0x64206420), "n"(immLut)
|
|
);
|
|
sign = sign >> 1;
|
|
|
|
// Ignore sign bit when indexing into LUT
|
|
uint32_t lut_idx = src_reg & 0x77777777;
|
|
Tensor scales_neg_ = cute::filter(scales_neg);
|
|
Tensor scales_pos_ = cute::filter(scales_pos);
|
|
CUTLASS_PRAGMA_UNROLL
|
|
for (int i = 0; i < N / 4; ++i, lut_idx >>=16, sign >>=16) {
|
|
auto&& scale_neg_ = reinterpret_cast<cutlass::Array<uint32_t, 2> const&>(scales_neg_(i));
|
|
auto&& scale_pos_ = reinterpret_cast<cutlass::Array<uint32_t, 2> const&>(scales_pos_(i));
|
|
asm volatile(
|
|
"{\n"
|
|
" .reg .b32 pos, neg ;\n" \
|
|
" prmt .b32 neg, %3, %4, %1 ;\n" \
|
|
" prmt .b32 pos, %5, %6, %1 ;\n" \
|
|
" prmt .b32 %0, pos, neg, %2 ;\n" \
|
|
"}\n"
|
|
: "=r"(r[i])
|
|
: "r"(lut_idx), "r"(sign), "r"(scale_neg_[0]), "r"(scale_neg_[1]), "r"(scale_pos_[0]), "r"(scale_pos_[1])
|
|
);
|
|
}
|
|
}
|
|
|
|
/// Utilities to dequantize A.
|
|
template <class Layout>
|
|
CUTLASS_DEVICE
|
|
static void static_check_scale(Layout const& tensor) {
|
|
static_assert(shape<0>(Layout{}) >= 4 && stride<0>(Layout{}) == 0, "At least 4 adjacent weights in a thread must share the same scale.");
|
|
}
|
|
template <class Engine,
|
|
class Layout>
|
|
CUTLASS_DEVICE
|
|
static void static_check_scale(Tensor<Engine, Layout> const& tensor) {
|
|
static_check_scale(flatten(Layout{}));
|
|
}
|
|
template <class EngineIn,
|
|
class EngineOut,
|
|
class LayoutIn,
|
|
class LayoutOut,
|
|
class... Ts>
|
|
CUTLASS_DEVICE
|
|
static void dequantize_A_kblock(
|
|
Tensor<EngineIn, LayoutIn> const& tCrA_load,
|
|
Tensor<EngineOut, LayoutOut>& tCrA_mma,
|
|
cute::tuple<Ts...>& partitioned_extra_info,
|
|
int const k_block) {
|
|
|
|
static_assert(is_rmem<EngineIn>::value, "Input tensor for A conversion must come from registers");
|
|
static_assert(is_rmem<EngineOut>::value, "Output tensor for A conversion must come from registers");
|
|
static_assert(cosize_v<LayoutIn> == cosize_v<LayoutOut>);
|
|
static_assert(size_v<LayoutIn> == cosize_v<LayoutIn>);
|
|
static_assert(size_v<LayoutOut> == cosize_v<LayoutOut>);
|
|
using SrcType = typename EngineIn::value_type;
|
|
using DstType = typename EngineOut::value_type;
|
|
|
|
Tensor src = tCrA_load(_, _, k_block);
|
|
Tensor dst = tCrA_mma(_, _, k_block);
|
|
|
|
CUTE_STATIC_ASSERT_V(size(src(_, 0)) == cosize(src(_, 0).layout()),
|
|
"The first mode of tensor src must be contiguous in memory");
|
|
// try to make the size of the first mode equal to 32bit
|
|
int constexpr NumValPerSrcReg = cute::min(decltype(size(src(_, 0)))::value,
|
|
ceil_div(32, sizeof_bits_v<SrcType>));
|
|
Tensor src_vm = cute::group_modes<1,-1>(cute::zipped_divide(src, Int<NumValPerSrcReg>{}));
|
|
Tensor dst_vm = cute::group_modes<1,-1>(cute::zipped_divide(dst, Int<NumValPerSrcReg>{}));
|
|
|
|
if constexpr (KernelConversionMode == ConversionMode::DirectConvert) {
|
|
CUTLASS_PRAGMA_UNROLL
|
|
for (int i = 0; i < size<1>(dst_vm); ++i) {
|
|
LayoutAwareConvert(src_vm(_, i), dst_vm(_, i));
|
|
}
|
|
}
|
|
else if constexpr (UseScaleLookupTable) {
|
|
constexpr int num_elements = decltype(size(src))::value;
|
|
static_assert(is_same_v<RealSwappedElementA, cutlass::int4b_t>, "Lookup table only supports int4 being the quant type now.");
|
|
static_assert(sizeof_bits_v<ElementScale> == 64, "Lookup table only supports 8 8bit scale values now.");
|
|
static_assert(num_elements % 4 == 0 && num_elements >= 4, "Lookup table requires a vector size of 4x when converting.");
|
|
|
|
Tensor tCrS_neg = cute::get<1>(partitioned_extra_info);
|
|
auto&& tCrS_pos = cute::get<2>(partitioned_extra_info); // modification to its value is needed
|
|
Tensor scales_neg = tCrS_neg(_, _, k_block);
|
|
Tensor scales_pos = tCrS_pos(_, _, k_block);
|
|
CUTE_STATIC_ASSERT_V(cute::size(src) == cute::size(scales_neg));
|
|
|
|
static_check_scale(scales_neg);
|
|
static_check_scale(scales_pos);
|
|
Tensor scales_neg_vm = cute::group_modes<1,-1>(cute::zipped_divide(scales_neg, Int<NumValPerSrcReg>{}));
|
|
Tensor scales_pos_vm = cute::group_modes<1,-1>(cute::zipped_divide(scales_pos, Int<NumValPerSrcReg>{}));
|
|
|
|
if (k_block == 0) {
|
|
Tensor scales_neg_vm_ = filter(scales_neg_vm);
|
|
Tensor scales_pos_vm_ = filter(scales_pos_vm);
|
|
CUTLASS_PRAGMA_UNROLL
|
|
for (int i = 0; i < size(scales_neg_vm_.layout()); ++i)
|
|
{
|
|
auto&& scale_neg_ = reinterpret_cast<cutlass::Array<uint32_t, 2> const&>(scales_neg_vm_(i));
|
|
auto&& scale_pos_ = reinterpret_cast<cutlass::Array<uint32_t, 2> &>(scales_pos_vm_(i));
|
|
constexpr uint32_t immLut = (0xf0 & 0xcc) ^ 0xaa;
|
|
asm volatile(
|
|
"{\n"
|
|
" lop3 .b32 %0, %2, %4, %5, %6;\n" \
|
|
" xor .b32 %1, %3, %5; \n" \
|
|
"}\n"
|
|
: "=r"(scale_pos_[0]), "=r"(scale_pos_[1])
|
|
: "r"(scale_neg_[0]), "r"(scale_neg_[1]), "n"(0xFFFFFF00), "n"(0x80808080), "n"(immLut)
|
|
);
|
|
}
|
|
}
|
|
CUTLASS_PRAGMA_UNROLL
|
|
for (int i = 0; i < size<1>(dst_vm); ++i) {
|
|
lookup_table_convert(src_vm(_, i), dst_vm(_, i), scales_neg_vm(_, i), scales_pos_vm(_, i));
|
|
}
|
|
}
|
|
else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
|
|
Tensor scales = cute::get<1>(partitioned_extra_info)(_, _, k_block);
|
|
CUTE_STATIC_ASSERT_V(size(src) == size(scales));
|
|
Tensor scales_vm = cute::group_modes<1,-1>(cute::zipped_divide(scales, Int<NumValPerSrcReg>{}));
|
|
|
|
if constexpr (is_same_v<DstType, ElementScale>) {
|
|
CUTLASS_PRAGMA_UNROLL
|
|
for (int i = 0; i < size<1>(dst_vm); ++i) {
|
|
LayoutAwareConvert(src_vm(_, i), dst_vm(_, i));
|
|
CUTLASS_PRAGMA_UNROLL
|
|
for (int j = 0; j < size<0>(dst_vm); ++j) {
|
|
dst_vm(j, i) *= scales_vm(j, i);
|
|
}
|
|
}
|
|
}
|
|
else {
|
|
auto stage = make_tensor_like<ElementScale>(src_vm(_, 0));
|
|
CUTLASS_PRAGMA_UNROLL
|
|
for (int i = 0; i < size<1>(dst_vm); ++i) {
|
|
LayoutAwareConvert(src_vm(_, i), stage);
|
|
CUTLASS_PRAGMA_UNROLL
|
|
for (int j = 0; j < size<0>(dst_vm); ++j) {
|
|
stage(j) *= scales_vm(j, i);
|
|
}
|
|
LayoutAwareConvert(stage, dst_vm(_, i));
|
|
}
|
|
}
|
|
}
|
|
else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
|
|
static_assert(is_same_v<ElementScale, ElementZero>, "ElementScale and ElementZero must be the same.");
|
|
Tensor scales = cute::get<1>(partitioned_extra_info)(_, _, k_block);
|
|
Tensor zeros = cute::get<3>(partitioned_extra_info)(_, _, k_block);
|
|
CUTE_STATIC_ASSERT_V(size(src) == size(scales));
|
|
CUTE_STATIC_ASSERT_V(size(src) == size(zeros));
|
|
Tensor scales_vm = cute::group_modes<1,-1>(cute::zipped_divide(scales, Int<NumValPerSrcReg>{}));
|
|
Tensor zeros_vm = cute::group_modes<1,-1>(cute::zipped_divide(zeros, Int<NumValPerSrcReg>{}));
|
|
|
|
if constexpr (is_same_v<DstType, ElementScale>) {
|
|
CUTLASS_PRAGMA_UNROLL
|
|
for (int i = 0; i < size<1>(dst_vm); ++i) {
|
|
LayoutAwareConvert(src_vm(_, i), dst_vm(_, i));
|
|
CUTLASS_PRAGMA_UNROLL
|
|
for (int j = 0; j < size<0>(dst_vm); ++j) {
|
|
dst_vm(j, i) = dst_vm(j, i) * scales_vm(j, i) + zeros_vm(j, i);
|
|
}
|
|
}
|
|
}
|
|
else {
|
|
auto stage = make_tensor_like<ElementScale>(src_vm(_, 0));
|
|
CUTLASS_PRAGMA_UNROLL
|
|
for (int i = 0; i < size<1>(dst_vm); ++i) {
|
|
LayoutAwareConvert(src_vm(_, i), stage);
|
|
CUTLASS_PRAGMA_UNROLL
|
|
for (int j = 0; j < size<0>(dst_vm); ++j) {
|
|
stage(j) = stage(j) * scales_vm(j, i) + zeros_vm(j, i);
|
|
}
|
|
LayoutAwareConvert(stage, dst_vm(_, i));
|
|
}
|
|
}
|
|
}
|
|
else {
|
|
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "No A data is loaded.");
|
|
}
|
|
}
|
|
|
|
|
|
/// Utilities for any additional inputs inside of the TMA load
|
|
template <
|
|
class Params,
|
|
class TensorStorage,
|
|
class... Ts
|
|
>
|
|
CUTLASS_DEVICE
|
|
static auto partition_extra_tma_inputs(
|
|
Params const& mainloop_params,
|
|
cute::tuple<Ts...> const& load_inputs,
|
|
TensorStorage& shared_tensors,
|
|
uint2 const& cluster_local_block_id,
|
|
int const m_coord,
|
|
int const l_coord) {
|
|
|
|
if constexpr (KernelConversionMode == ConversionMode::DirectConvert) {
|
|
return cute::make_tuple();
|
|
}
|
|
else if constexpr (ModeHasScales) {
|
|
Tensor sS = make_tensor(make_smem_ptr(shared_tensors.smem_scale.begin()), SmemLayoutScale{}); // (BLK_M,BLK_K,PIPE)
|
|
Tensor gS_mkl = get<2>(load_inputs);
|
|
auto block_tma_s = mainloop_params.tma_load_scale.get_slice(cluster_local_block_id.y);
|
|
Tensor gS = gS_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k)
|
|
|
|
Tensor tSgS = block_tma_s.partition_S(gS);
|
|
Tensor tSsS = block_tma_s.partition_D(sS); // (TMA,TMA_M,TMA_K,PIPE)
|
|
if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
|
|
return cute::make_tuple(tSgS, tSsS);
|
|
}
|
|
else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
|
|
Tensor sZ = make_tensor(make_smem_ptr(shared_tensors.smem_zero.begin()), SmemLayoutScale{}); // (BLK_M,BLK_K,PIPE)
|
|
Tensor gZ_mkl = get<3>(load_inputs);
|
|
auto block_tma_z = mainloop_params.tma_load_zero.get_slice(cluster_local_block_id.y);
|
|
Tensor gZ = gZ_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k)
|
|
|
|
Tensor tZgZ = block_tma_z.partition_S(gZ);
|
|
Tensor tZsZ = block_tma_z.partition_D(sZ); // (TMA,TMA_M,TMA_K,PIPE)
|
|
return cute::make_tuple(tSgS, tSsS, tZgZ, tZsZ);
|
|
}
|
|
else {
|
|
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled for input partitioning.");
|
|
}
|
|
}
|
|
else {
|
|
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled for input partitioning.");
|
|
}
|
|
}
|
|
|
|
/// Utilities for partitioning extra inputs for loading from smem in the mainloop.
|
|
template <
|
|
class ThreadMma,
|
|
class TensorStorage
|
|
>
|
|
CUTLASS_DEVICE
|
|
static auto partition_extra_mma_info(
|
|
ThreadMma const& mma_thread_slice,
|
|
TensorStorage& shared_tensors) {
|
|
|
|
if constexpr (KernelConversionMode == ConversionMode::DirectConvert) {
|
|
// nothing to do
|
|
return cute::make_tuple();
|
|
}
|
|
else if constexpr (UseScaleLookupTable) {
|
|
Tensor sS = make_tensor(make_smem_ptr(shared_tensors.smem_scale.begin()), SmemLayoutScale{});// (BLK_M,BLK_SCALE_K,PIPE)
|
|
Tensor tCsS = mma_thread_slice.partition_A(sS);
|
|
Tensor tCrS_neg = make_tensor<ElementScale>(mma_thread_slice.partition_fragment_A(sS(_,_,Int<0>{})).layout());
|
|
Tensor tCrS_pos = make_tensor<ElementScale>(mma_thread_slice.partition_fragment_A(sS(_,_,Int<0>{})).layout());
|
|
|
|
if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
|
|
return cute::make_tuple(tCsS, tCrS_neg, tCrS_pos);
|
|
}
|
|
}
|
|
else if constexpr (ModeHasScales) {
|
|
Tensor sS = make_tensor(make_smem_ptr(shared_tensors.smem_scale.begin()), SmemLayoutScale{});// (BLK_M,BLK_SCALE_K,PIPE)
|
|
Tensor tCsS = mma_thread_slice.partition_A(sS);
|
|
Tensor tCrS = make_tensor<ElementScale>(mma_thread_slice.partition_fragment_A(sS(_,_,Int<0>{})).layout());
|
|
|
|
if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
|
|
return cute::make_tuple(tCsS, tCrS);
|
|
}
|
|
else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
|
|
Tensor sZ = make_tensor(make_smem_ptr(shared_tensors.smem_zero.begin()), SmemLayoutScale{});// (BLK_M,BLK_SCALE_K,PIPE)
|
|
Tensor tCsZ = mma_thread_slice.partition_A(sZ);
|
|
Tensor tCrZ = make_tensor<ElementZero>(mma_thread_slice.partition_fragment_A(sZ(_,_,Int<0>{})).layout());
|
|
return cute::make_tuple(tCsS, tCrS, tCsZ, tCrZ);
|
|
}
|
|
else {
|
|
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in A -> RF path.");
|
|
}
|
|
}
|
|
else {
|
|
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in A -> RF path.");
|
|
}
|
|
}
|
|
|
|
/// Returns the tiled copy and copy views for the extra inputs.
|
|
template <class TiledMma, class... Ts>
|
|
CUTLASS_DEVICE
|
|
static auto retile_extra_mma_info(
|
|
TiledMma const& tiled_mma,
|
|
cute::tuple<Ts...>& partitioned_extra_info,
|
|
int const warp_group_thread_idx) {
|
|
|
|
if constexpr (KernelConversionMode == ConversionMode::DirectConvert) {
|
|
// nothing to do
|
|
return cute::make_tuple();
|
|
}
|
|
else if constexpr (ModeHasScales) {
|
|
auto smem_tiled_copy_S = make_tiled_copy_A(SmemCopyAtomScale{}, tiled_mma);
|
|
auto smem_thr_copy_S = smem_tiled_copy_S.get_thread_slice(warp_group_thread_idx);
|
|
Tensor tCrS_copy_view = smem_thr_copy_S.retile_D(cute::get<1>(partitioned_extra_info)); // (CPY,CPY_M,CPY_K)
|
|
|
|
if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
|
|
return cute::make_tuple(smem_tiled_copy_S, tCrS_copy_view);
|
|
}
|
|
else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
|
|
Tensor tCrZ_copy_view = smem_thr_copy_S.retile_D(cute::get<3>(partitioned_extra_info)); // (CPY,CPY_M,CPY_K)
|
|
return cute::make_tuple(smem_tiled_copy_S, tCrS_copy_view, tCrZ_copy_view);
|
|
}
|
|
else {
|
|
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in A -> RF path.");
|
|
}
|
|
}
|
|
else {
|
|
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in A -> RF path.");
|
|
}
|
|
}
|
|
};
|
|
|
|
} // cutlass::gemm::collective::detail
|