leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 29 Packages Projects Releases Wiki Activity

fix wrong A/BLayout in MMA_Traits for binary mma and append other MMA_Traits support (#1856)

Browse Source 
* fix wrong A/BLayout in  MMA_Traits<SM80_16x8x256_S32U1U1S32_TN_XORPOPC> and append support for  m8n8k128, m16n8k128  mma.and.popc in MMA_Traits instantiation

* add "print" template for  subbyte_reference<T>
This commit is contained in:
Caleb_Du
2024-10-24 14:38:35 -04:00
committed by GitHub
parent be692b48b0
commit a424ca6cf9
3 changed files with 155 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cute/atom/mma_traits_sm80.hpp
+51 -4
View File
@@ -433,10 +433,57 @@ struct MMA_Traits<SM80_16x8x256_S32U1U1S32_TN_XORPOPC>
using Shape_MNK = Shape<_16,_8,_256>;
using ThrID = Layout<_32>;
using ALayout = Layout<Shape <_32,Shape < _8, _4,_2, _2>>,
Stride<_64,Stride<_64,_16,_8,_2048>>>;
using BLayout = Layout<Shape <_32,Shape <_32, _2>>,
Stride<_32,Stride< _1,_1024>>>;
using ALayout = Layout<Shape<Shape<_4,_8>,Shape<_32,_2,_2>>,
Stride<Stride<_512,_1>,Stride<_16,_8,_2048>>>;
using BLayout = Layout<Shape<Shape <_4,_8>,Shape<_32,_2>>,
Stride<Stride<_256,_1>,Stride< _8,_1024>>>;
using CLayout = SM80_16x8_Row;
};
template <>
struct MMA_Traits<SM80_16x8x256_S32U1U1S32_TN_ANDPOPC>
:MMA_Traits<SM80_16x8x256_S32U1U1S32_TN_XORPOPC> {};
template<>
struct MMA_Traits<SM80_8x8x128_S32U1U1S32_TN_XORPOPC>
{
using ValTypeD = int32_t;
using ValTypeA = cute::uint1b_t;
using ValTypeB = cute::uint1b_t;
using ValTypeC = int32_t;
using Shape_MNK = Shape<_8,_8,_128>;
using ThrID = Layout<_32>;
using ALayout = Layout<Shape<Shape<_4,_8>,_32>,
Stride<Stride<_256,_1>,_8>>;
using BLayout = Layout<Shape<Shape<_4,_8>,_32>,
Stride<Stride<_256,_1>,_8>>;
using CLayout = SM80_8x8_Row;
};
template <>
struct MMA_Traits<SM80_8x8x128_S32U1U1S32_TN_ANDPOPC>
:MMA_Traits<SM80_8x8x128_S32U1U1S32_TN_XORPOPC> {};
template<>
struct MMA_Traits<SM80_16x8x128_S32U1U1S32_TN_XORPOPC>
{
using ValTypeD = int32_t;
using ValTypeA = cute::uint1b_t;
using ValTypeB = cute::uint1b_t;
using ValTypeC = int32_t;
using Shape_MNK = Shape<_16,_8,_128>;
using ThrID = Layout<_32>;
using ALayout = Layout<Shape<Shape<_4,_8>,Shape<_32,_2>>,
Stride<Stride<_512,_1>,Stride<Stride<_16,_8>>>>;
using BLayout = Layout<Shape <Shape<_4,_8>,_32>,
Stride<Stride<_256,_1>,_8>>;
using CLayout = SM80_16x8_Row;
};
template <>
struct MMA_Traits<SM80_16x8x128_S32U1U1S32_TN_ANDPOPC>
:MMA_Traits<SM80_16x8x128_S32U1U1S32_TN_XORPOPC> {};
} // end namespace cute