Skip to content

Commit a97267b

Browse files
taozha2taozha2
committed
add int8/tf32 transpose A intrinsic
1 parent 9fd9dc3 commit a97267b

File tree

3 files changed

+155
-7
lines changed

3 files changed

+155
-7
lines changed

‎include/cute/arch/xe_copy_1B.hpp

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,15 @@ SYCL_DEVICE_BUILTIN(
106106
intptr_t baseoffset, int width_minus_one, int height_minus_one,
107107
int pitch_minus_one, cute::intel::coord_t coord));
108108

109+
// 8bits NO transform transpose
110+
SYCL_DEVICE_BUILTIN(
111+
cute::intel::ushort8 __builtin_IB_subgroup_block_read_cacheopts_transpose_u8_m32k8(
112+
intptr_t baseoffset, int width_minus_one, int height_minus_one,
113+
int pitch_minus_one, cute::intel::coord_t coord, int cache = 0));
114+
SYCL_DEVICE_BUILTIN(
115+
cute::intel::ushort4 __builtin_IB_subgroup_block_read_cacheopts_transpose_u8_m32k4(
116+
intptr_t baseoffset, int width_minus_one, int height_minus_one,
117+
int pitch_minus_one, cute::intel::coord_t coord, int cache = 0));
109118

110119
// 8bits VNNI transform No transpose
111120
SYCL_DEVICE_BUILTIN(
@@ -443,6 +452,45 @@ struct XE_2D_U8x32x32_LD_N {
443452
}
444453
};
445454

455+
struct XE_2D_U8x4x32_LD_T {
456+
using BlockShape = Shape<_4, _32>;
457+
using inst_dtype = uint8_t;
458+
static constexpr bool is_transpose = true;
459+
460+
template <class T>
461+
CUTE_HOST_DEVICE static void copy(const void *baseoffset, int width,
462+
int height, int pitch, intel::coord_t coord,
463+
T *dst) {
464+
#if defined(SYCL_INTEL_TARGET)
465+
static_assert(sizeof(T) == 1, "Expected T to have size 1");
466+
*reinterpret_cast<intel::ushort4 *>(dst) =
467+
__builtin_IB_subgroup_block_read_cacheopts_transpose_u8_m32k4(
468+
(intptr_t)(baseoffset), width - 1, height - 1, pitch - 1, coord);
469+
#else
470+
CUTE_INVALID_CONTROL_PATH("Trying to use block loads on non-PVC hardware");
471+
#endif
472+
}
473+
};
474+
475+
struct XE_2D_U8x8x32_LD_T {
476+
using BlockShape = Shape<_8, _32>;
477+
using inst_dtype = uint8_t;
478+
static constexpr bool is_transpose = true;
479+
480+
template <class T>
481+
CUTE_HOST_DEVICE static void copy(const void *baseoffset, int width,
482+
int height, int pitch, intel::coord_t coord,
483+
T *dst) {
484+
#if defined(SYCL_INTEL_TARGET)
485+
static_assert(sizeof(T) == 1, "Expected T to have size 1");
486+
*reinterpret_cast<intel::ushort8 *>(dst) =
487+
__builtin_IB_subgroup_block_read_cacheopts_transpose_u8_m32k8(
488+
(intptr_t)(baseoffset), width - 1, height - 1, pitch - 1, coord);
489+
#else
490+
CUTE_INVALID_CONTROL_PATH("Trying to use block loads on non-PVC hardware");
491+
#endif
492+
}
493+
};
446494
struct XE_2D_U4x16x16_LD_T {
447495
using BlockShape = Shape<_16, _16>;
448496
using inst_dtype = uint32_t;

‎include/cute/arch/xe_copy_4B.hpp

Lines changed: 33 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -117,9 +117,10 @@ SYCL_DEVICE_BUILTIN(
117117
int pitch_minus_one, cute::intel::coord_t coord));
118118

119119
// 32bits No transform No transpose
120-
SYCL_DEVICE_BUILTIN(cute::intel::uint __builtin_IB_subgroup_block_read_flat_u32_m1k16v1(
121-
intptr_t baseoffset, int width_minus_one, int height_minus_one,
122-
int pitch_minus_one, cute::intel::coord_t coord));
120+
SYCL_DEVICE_BUILTIN(
121+
cute::intel::uint __builtin_IB_subgroup_block_read_flat_u32_m1k16v1(
122+
intptr_t baseoffset, int width_minus_one, int height_minus_one,
123+
int pitch_minus_one, cute::intel::coord_t coord));
123124
SYCL_DEVICE_BUILTIN(
124125
cute::intel::uint2 __builtin_IB_subgroup_block_read_flat_u32_m2k16v1(
125126
intptr_t baseoffset, int width_minus_one, int height_minus_one,
@@ -142,9 +143,10 @@ SYCL_DEVICE_BUILTIN(
142143
int pitch_minus_one, cute::intel::coord_t coord));
143144

144145
// 32bits No transform Transpose
145-
SYCL_DEVICE_BUILTIN(cute::intel::uint __builtin_IB_subgroup_block_read_flat_transpose_u32_k1(
146-
intptr_t baseoffset, int width_minus_one, int height_minus_one,
147-
int pitch_minus_one, cute::intel::coord_t coord));
146+
SYCL_DEVICE_BUILTIN(
147+
cute::intel::uint __builtin_IB_subgroup_block_read_flat_transpose_u32_k1(
148+
intptr_t baseoffset, int width_minus_one, int height_minus_one,
149+
int pitch_minus_one, cute::intel::coord_t coord));
148150
SYCL_DEVICE_BUILTIN(
149151
cute::intel::uint2 __builtin_IB_subgroup_block_read_flat_transpose_u32_k2(
150152
intptr_t baseoffset, int width_minus_one, int height_minus_one,
@@ -157,6 +159,10 @@ SYCL_DEVICE_BUILTIN(
157159
cute::intel::uint8 __builtin_IB_subgroup_block_read_flat_transpose_u32_k8(
158160
intptr_t baseoffset, int width_minus_one, int height_minus_one,
159161
int pitch_minus_one, cute::intel::coord_t coord));
162+
SYCL_DEVICE_BUILTIN(
163+
cute::intel::uint4 __builtin_IB_subgroup_block_read_flat_transpose_u32_m8k8(
164+
intptr_t baseoffset, int width_minus_one, int height_minus_one,
165+
int pitch_minus_one, cute::intel::coord_t coord));
160166

161167
// 32bits
162168
SYCL_DEVICE_BUILTIN(void __builtin_IB_subgroup_block_write_flat_u32_m1k16v1(
@@ -710,6 +716,27 @@ struct XE_2D_U32x16x8_LD_T {
710716
};
711717
};
712718

719+
struct XE_2D_TF32x8x8_LD_T {
720+
using BlockShape = Shape<_8, _8>;
721+
using ValueShape = Shape<_4, _16>;
722+
723+
static constexpr bool is_transpose = true;
724+
725+
template <class T>
726+
CUTE_HOST_DEVICE static void copy(const void *baseoffset, int width,
727+
int height, int pitch, intel::coord_t coord,
728+
T *dst) {
729+
#if defined(SYCL_INTEL_TARGET)
730+
static_assert(sizeof(T) == 4, "Expected T to have size 4");
731+
*reinterpret_cast<intel::uint4 *>(dst) =
732+
__builtin_IB_subgroup_block_read_flat_transpose_u32_m8k8(
733+
(intptr_t)(baseoffset), width - 1, height - 1, pitch - 1, coord);
734+
#else
735+
CUTE_INVALID_CONTROL_PATH("Trying to use block loads on non-PVC hardware");
736+
#endif
737+
}
738+
};
739+
713740
struct XE_2D_U32x1x16_ST_N {
714741
using BlockShape = Shape<_1, _16>;
715742

‎include/cute/atom/copy_traits_xe.hpp

Lines changed: 74 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1399,7 +1399,25 @@ struct Copy_Traits_<XE_2D_TF32x32x16_LD_N, args_t...>
13991399
};
14001400

14011401
template <class... args_t>
1402-
struct Copy_Traits_<XE_2D_U32x1x16_LD_N, args_t...>
1402+
struct Copy_Traits<XE_2D_TF32x8x8_LD_T, args_t...>
1403+
: XE_2D_LD_Unpack<XE_2D_TF32x8x8_LD_T, args_t...> {
1404+
using ThrID = Layout<_16>;
1405+
// Map from (src-thr,src-val) to bit
1406+
using SrcLayout = Layout<Shape <_16, Shape <_4, _32>>,
1407+
Stride< _0, Stride<_32, _1>>>;
1408+
// Map from (dst-thr,dst-val) to bit
1409+
using DstLayout = Layout<Shape <_16, Shape <_4, _32>>,
1410+
Stride< _32, Stride<_32, _1>>>;
1411+
// Reference map from (thr,val) to bit
1412+
using RefLayout = DstLayout;
1413+
1414+
template <class... ArgTs>
1415+
Copy_Traits(ArgTs... args)
1416+
: XE_2D_LD_Unpack<XE_2D_TF32x8x8_LD_T, args_t...>(args...) {}
1417+
};
1418+
1419+
template <class... args_t>
1420+
struct Copy_Traits<XE_2D_U32x1x16_LD_N, args_t...>
14031421
: XE_2D_LD_Unpack<XE_2D_U32x1x16_LD_N, args_t...> {
14041422
using ThrID = Layout<_16>;
14051423
// Map from (src-thr,src-val) to bit
@@ -1688,6 +1706,61 @@ struct Copy_Traits_<XE_2D_U16x16x16_LD_T, args_t...>
16881706
: XE_2D_LD_Unpack<XE_2D_U16x16x16_LD_T, args_t...>(args...) {}
16891707
};
16901708

1709+
template <class... args_t>
1710+
struct Copy_Traits_<XE_2D_U8x32x16_LD_T, args_t...>
1711+
: XE_2D_LD_Unpack<XE_2D_U8x32x16_LD_T, args_t...> {
1712+
using ThrID = Layout<_16>;
1713+
// Map from (src-thr,src-val) to bit
1714+
using SrcLayout = Layout<Shape <_16,_16>,
1715+
Stride< _0, _1>>;
1716+
// Map from (dst-thr,dst-val) to bit
1717+
using DstLayout = Layout<Shape < _16,Shape <_16,_16>>,
1718+
Stride<_256,Stride< _1,_16>>>;
1719+
// Reference map from (thr,val) to bit
1720+
using RefLayout = DstLayout;
1721+
1722+
template <class... ArgT>
1723+
Copy_Traits_(ArgT... args)
1724+
: XE_2D_LD_Unpack<XE_2D_U8x32x16_LD_T, args_t...>(args...) {}
1725+
};
1726+
1727+
template <class... args_t>
1728+
struct Copy_Traits_<XE_2D_U8x8x32_LD_T, args_t...>
1729+
: XE_2D_LD_Unpack<XE_2D_U8x8x32_LD_T, args_t...> {
1730+
using ThrID = Layout<_16>;
1731+
// Map from (src-thr,src-val) to bit
1732+
using SrcLayout = Layout<Shape <_16,Shape <_8, _2, _8>>,
1733+
Stride<_0, Stride<_1, _8, _16>>>;
1734+
// Map from (dst-thr,dst-val) to bit
1735+
using DstLayout = Layout<Shape < _16,Shape <_8, _2, _8>>,
1736+
Stride<_256,Stride<_1, _8, _16>>>;
1737+
// Reference map from (thr,val) to bit
1738+
using RefLayout = DstLayout;
1739+
1740+
template <class... ArgT>
1741+
Copy_Traits_(ArgT... args)
1742+
: XE_2D_LD_Unpack<XE_2D_U8x8x32_LD_T, args_t...>(args...) {}
1743+
};
1744+
1745+
template <class... args_t>
1746+
struct Copy_Traits<XE_2D_U8x4x32_LD_T, args_t...>
1747+
: XE_2D_LD_Unpack<XE_2D_U8x4x32_LD_T, args_t...> {
1748+
using ThrID = Layout<_16>;
1749+
// Map from (src-thr,src-val) to bit
1750+
using SrcLayout = Layout<Shape <_16,Shape <_8, _2, _4>>,
1751+
Stride<_0, Stride<_1, _8, _16>>>;
1752+
// Map from (dst-thr,dst-val) to bit
1753+
using DstLayout = Layout<Shape < _16,Shape <_8, _2, _4>>,
1754+
Stride<_256,Stride<_1, _8, _16>>>;
1755+
// Reference map from (thr,val) to bit
1756+
using RefLayout = DstLayout;
1757+
1758+
template <class... ArgT>
1759+
Copy_Traits(ArgT... args)
1760+
: XE_2D_LD_Unpack<XE_2D_U8x4x32_LD_T, args_t...>(args...) {}
1761+
};
1762+
};
1763+
16911764
// template<class... args_t>
16921765
// struct Copy_Traits<XE_2D_U32x16x1_LD_T, args_t...>
16931766
// : XE_2D_LD_Unpack<XE_2D_U32x16x1_LD_T, args_t...> {

0 commit comments

Comments
 (0)