1#include <cute/tensor.hpp>

2#include <cute/atom/copy_atom.hpp>

3#include <cute/atom/copy_traits_sm75.hpp>

4#include <cstdio>

5

6using namespace cute;

7

8template <class T>

9void show(char const* label, T const& value)

10{

11 std::printf("%s", label);

12 cute::print(value);

13 cute::print("\n");

14}

15

16int main()

17{

18 half_t storage[1024]{};

19 auto sB = make_tensor(

20 make_smem_ptr(storage),

21 make_layout(make_shape(_32{},_32{}), make_stride(_32{},_1{})));

22 auto coordB = make_identity_tensor(make_shape(_32{},_32{}));

23 auto mma_k = make_tiled_mma(

24 SM80_16x8x16_F32F16F16F32_TN{},

25 Layout<Shape<_1,_1,_1>>{}, Tile<_16,_8,_32>{});

26 auto copy_k = make_tiled_copy_B(

27 Copy_Atom<SM75_U32x4_LDSM_N,half_t>{}, mma_k);

28 using TC = decltype(copy_k);

29

30 int checked = 0;

31 int different_groups = 0;

32 for (int lane = 0; lane < 32; ++lane) {

33 auto thr_mma = mma_k.get_slice(lane);

34 auto thr_copy = copy_k.get_slice(lane);

35 auto tCrB = thr_mma.partition_fragment_B(sB);

36 auto tXrB = thr_copy.retile_D(tCrB);

37 auto direct = zipped_divide(tCrB, make_layout(_8{}));

38 auto mma_coords = thr_mma.partition_B(coordB);

39 auto copy_coords = thr_copy.partition_D(coordB);

40 auto retiled_coords = thr_copy.retile_D(mma_coords);

41

42 for (int n = 0; n < 4; ++n) {

43 for (int a = 0; a < 8; ++a) {

44 auto expected = mma_coords(a%4,n,a/4);

45 auto actual = copy_coords(make_coord(a,0),n,0);

46 auto retiled = retiled_coords(make_coord(a,0),n,0);

47 auto slot = tXrB.layout()(make_coord(a,0),n,0);

48 if (!(expected == actual && actual == retiled) ||

49 int(slot) != a+8*n) {

50 std::printf("K gather mismatch: lane=%d n=%d a=%d\n",

51 lane,n,a);

52 return 1;

53 }

54 ++checked;

55 }

56 }

57

58 bool different = false;

59 for (int a = 0; a < 8; ++a) {

60 int expected_slot = int(tXrB.layout()(make_coord(a,0),0,0));

61 int direct_slot = int(direct.layout()(a,0));

62 different |= expected_slot != direct_slot;

63 }

64 different_groups += different;

65

66 if (lane == 0) {

67 constexpr int R = decltype(tCrB)::rank;

68 auto V = size<0>(tCrB);

69 auto mn = right_inverse(typename TC::TiledLayout_TV{})

70 .with_shape(shape(typename TC::Tiler_MN{}));

71 auto fmn = upcast<TC::TiledNumThr{} * V>(mn);

72 auto gather = right_inverse(fmn);

73 auto values = logical_product(make_layout(V),gather);

74 auto fv = zipped_divide(values,make_layout(typename TC::AtomNumVal{}));

75 auto tile = prepend(product_each(shape(fmn)),V);

76 auto tt = zipped_divide(tCrB.layout(),tile);

77 auto vt = tt.compose(fv,_);

78 auto selection = append<R>(Int<0>{},_);

79 show("B input: ",tCrB.layout());

80 show("B retile: ",tXrB.layout());

81 show("B direct: ",direct.layout());

82 show("B TiledNumVal: ",typename TC::TiledNumVal{});

83 show("B AtomNumVal: ",typename TC::AtomNumVal{});

84 show("B inverse with NK shape: ",mn);

85 show("B frg_layout_nk: ",fmn);

86 show("B gather shape: ",product_each(shape(fmn)));

87 show("B gather: ",gather);

88 show("B product: ",values);

89 show("B frg_layout_v: ",fv);

90 show("B tile: ",tile);

91 show("B t_tensor: ",tt);

92 show("B v_tensor: ",vt);

93 show("B selection: ",selection);

94 show("B result: ",vt(_,selection));

95 for (int a = 0; a < 8; ++a) {

96 std::printf("r[%d] -> B",a);

97 cute::print(copy_coords(make_coord(a,0),0,0));

98 cute::print("\n");

99 }

100 std::printf("K group slots:");

101 for (int a = 0; a < 8; ++a) {

102 std::printf(" %d",int(tXrB.layout()(make_coord(a,0),0,0)));

103 }

104 std::printf("\nDirect group slots:");

105 for (int a = 0; a < 8; ++a) {

106 std::printf(" %d",int(direct.layout()(a,0)));

107 }

108 std::printf("\n");

109 }

110 }

111 if (different_groups != 32) {

112 return 1;

113 }

114 std::printf("Verified %d K-gather matches; %d direct groups differ.\n",

115 checked,different_groups);

116}

117