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 auto mma = make_tiled_mma(SM80_16x8x16_F32F16F16F32_TN{});

19 auto tiled_copy = make_tiled_copy_B(

20 Copy_Atom<SM75_U32x1_LDSM_N, half_t>{}, mma);

21 auto copy_x2 = make_tiled_copy_B(

22 Copy_Atom<SM75_U32x2_LDSM_N, half_t>{}, mma);

23 using TC = decltype(tiled_copy);

24

25 // 主机占位存储：只检查布局，不执行设备加载或读取 fragment 的值。

26 half_t storage[128]{};

27 auto sB = make_tensor(

28 make_smem_ptr(storage),

29 make_layout(make_shape(_8{}, _16{}), make_stride(_16{}, _1{})));

30 auto coordB = make_identity_tensor(make_shape(_8{}, _16{}));

31

32 show("TiledLayout_TV: ", typename TC::TiledLayout_TV{});

33 show("Tiler_NK: ", typename TC::Tiler_MN{});

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

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

36

37 int checked_dst = 0;

38 int checked_src = 0;

39 int checked_x2 = 0;

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

41 auto thr_mma = mma.get_slice(lane);

42 auto thr_copy = tiled_copy.get_slice(lane);

43 auto tCrB = thr_mma.partition_fragment_B(sB);

44 auto mma_coords = thr_mma.partition_B(coordB);

45 auto tXsB = thr_copy.partition_S(sB);

46 auto tXrB = thr_copy.retile_D(tCrB);

47 auto copy_coords = thr_copy.partition_D(coordB);

48 auto src_coords = thr_copy.partition_S(coordB);

49 auto retiled_coords = thr_copy.retile_D(mma_coords);

50 auto direct_zip = zipped_divide(tCrB, make_layout(_2{}));

51 auto direct_mode = logical_divide(

52 tCrB, make_tile(make_layout(_2{})));

53

54 for (int b = 0; b < 2; ++b) {

55 for (int a = 0; a < 2; ++a) {

56 int v = a + 2*b;

57 auto expected = make_coord(lane/4, 2*(lane%4) + a + 8*b);

58 auto mma_coord = mma_coords(v,0,0);

59 auto copy_coord = copy_coords(make_coord(a,b),0,0);

60 auto retiled_coord = retiled_coords(make_coord(a,b),0,0);

61 int slot = int(tXrB.layout()(make_coord(a,b),0,0));

62 if (!(expected == mma_coord && mma_coord == copy_coord &&

63 copy_coord == retiled_coord) || slot != v ||

64 int(direct_zip.layout()(a,b)) != slot ||

65 int(direct_mode.layout()(make_coord(a,b),0,0)) != slot) {

66 std::printf("B destination mismatch: lane=%d a=%d b=%d\n",

67 lane,a,b);

68 return 1;

69 }

70 ++checked_dst;

71 }

72 for (int u = 0; u < 8; ++u) {

73 auto expected = make_coord(lane%8, u+8*b);

74 int ref_lane = 4*(lane%8)+u/2;

75 auto ref_coords = mma.get_slice(ref_lane).partition_B(coordB);

76 auto src_coord = src_coords(make_coord(u,b),0,0);

77 if (!(src_coord == expected &&

78 src_coord == ref_coords(u%2+2*b,0,0))) {

79 std::printf("B source mismatch: lane=%d u=%d b=%d\n",

80 lane,u,b);

81 return 1;

82 }

83 ++checked_src;

84 }

85 }

86

87 auto thr_x2 = copy_x2.get_slice(lane);

88 auto tXrB_x2 = thr_x2.retile_D(tCrB);

89 auto x2_coords = thr_x2.partition_D(coordB);

90 for (int v = 0; v < 4; ++v) {

91 if (!(x2_coords(make_coord(v,0),0,0) == mma_coords(v,0,0)) ||

92 int(tXrB_x2.layout()(make_coord(v,0),0,0)) != v) {

93 std::printf("B x2 destination mismatch: lane=%d v=%d\n",lane,v);

94 return 1;

95 }

96 ++checked_x2;

97 }

98

99 if (lane >= 2) {

100 continue;

101 }

102 std::printf("lane %d\n",lane);

103 show("tCrB: ",tCrB.layout());

104 show("tXsB: ",tXsB.layout());

105 show("tXrB: ",tXrB.layout());

106 show("direct_zip: ",direct_zip.layout());

107 show("direct_mode: ",direct_mode.layout());

108 show("tXrB_x2: ",tXrB_x2.layout());

109

110 if (lane == 0) {

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

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

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

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

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

116 auto gather = right_inverse(fmn);

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

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

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

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

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

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

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

124 show("frg_layout_nk: ",fmn);

125 show("gather shape: ",product_each(shape(fmn)));

126 show("inverse frg_layout_nk: ",gather);

127 show("logical_product: ",values);

128 show("frg_layout_v: ",fv);

129 show("tile: ",tile);

130 show("t_tensor: ",tt);

131 show("v_tensor: ",vt);

132 show("selection: ",selection);

133 show("result: ",vt(_,selection));

134 }

135 std::printf("Source group 0: ");

136 for (int u = 0; u < 8; ++u) {

137 cute::print(src_coords(make_coord(u,0),0,0));

138 cute::print(" ");

139 }

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

141 for (int v = 0; v < 4; ++v) {

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

143 cute::print(mma_coords(v,0,0));

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

145 }

146 for (int b = 0; b < 2; ++b) {

147 std::printf("Copy group %d:",b);

148 for (int a = 0; a < 2; ++a) {

149 std::printf(" r[%d]",int(tXrB.layout()(make_coord(a,b),0,0)));

150 }

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

152 }

153 }

154 std::printf("Verified %d B destination matches and %d source references; "

155 "%d x2 destination matches.\n",checked_dst,checked_src,checked_x2);

156}

157