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_A(

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

21 using TC = decltype(tiled_copy);

22

23 // 主机占位存储：仅用于构造 tensor，不执行设备加载。

24 half_t storage[256]{};

25 auto sA = make_tensor(

26 make_smem_ptr(storage),

27 make_layout(make_shape(_16{}, _16{}),

28 make_stride(_16{}, _1{})));

29 auto coordA = make_identity_tensor(make_shape(_16{}, _16{}));

30

31 auto dst_to_ref = right_inverse(typename TC::AtomLayoutRef{})

32 .compose(typename TC::AtomLayoutDst{});

33 auto split_tv = zipped_divide(

34 typename TC::TiledLayout_TV{},

35 make_shape(typename TC::AtomNumThr{}, typename TC::AtomNumVal{}));

36 auto destination_tv = coalesce(

37 zip(split_tv.compose(dst_to_ref, _)), Shape<_1,Shape<_1,_1>>{});

38

39 show("AtomLayoutSrc: ", typename TC::AtomLayoutSrc{});

40 show("AtomLayoutDst: ", typename TC::AtomLayoutDst{});

41 show("AtomLayoutRef: ", typename TC::AtomLayoutRef{});

42 show("DstToRef: ", dst_to_ref);

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

44 show("Destination grouped TV: ", destination_tv);

45

46 int checked_dst = 0;

47 int checked_src = 0;

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

49 auto thr_mma = mma.get_slice(lane);

50 auto thr_copy = tiled_copy.get_slice(lane);

51 auto tCrA = thr_mma.partition_fragment_A(sA);

52 auto coords = thr_mma.partition_A(coordA);

53 auto tXsA = thr_copy.partition_S(sA);

54 auto tXrA = thr_copy.retile_D(tCrA);

55 auto direct_zip = zipped_divide(tCrA, make_layout(_2{}));

56 auto direct_mode = logical_divide(

57 tCrA, make_tile(make_layout(_2{})));

58 auto copy_coords = thr_copy.partition_D(coordA);

59 auto src_coords = thr_copy.partition_S(coordA);

60 auto retiled_coords = thr_copy.retile_D(coords);

61

62 for (int b = 0; b < 4; ++b) {

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

64 auto x = coords(a+2*b, 0, 0);

65 auto y = copy_coords(make_coord(a,b), 0, 0);

66 auto z = retiled_coords(make_coord(a,b), 0, 0);

67 auto slot = tXrA.layout()(make_coord(a,b), 0, 0);

68 auto zip_slot = direct_zip.layout()(a,b);

69 auto mode_slot = direct_mode.layout()(make_coord(a,b),0,0);

70 if (!(x == y && y == z) || int(slot) != a+2*b ||

71 int(zip_slot) != int(slot) || int(mode_slot) != int(slot)) {

72 std::printf("Destination mismatch: lane=%d a=%d b=%d\n",

73 lane, a, b);

74 return 1;

75 }

76 ++checked_dst;

77 }

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

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

80 int ref_value = u%2+2*b;

81 auto ref_coords = mma.get_slice(ref_lane).partition_A(coordA);

82 auto x = src_coords(make_coord(u,b), 0, 0);

83 auto y = ref_coords(ref_value, 0, 0);

84 if (!(x == y)) {

85 std::printf("Source mismatch: lane=%d u=%d b=%d\n",

86 lane, u, b);

87 return 1;

88 }

89 ++checked_src;

90 }

91 }

92

93 // 只打印前两个 lane；全部 32 个 lane 都执行上面的检查。

94 if (lane >= 2) {

95 continue;

96 }

97

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

99 show("tCrA: ", tCrA.layout());

100 show("tXsA: ", tXsA.layout());

101 show("tXrA: ", tXrA.layout());

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

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

104

105 if (lane == 0) {

106 constexpr int R = decltype(tCrA)::rank;

107 auto V = size<0>(tCrA);

108 auto inv = right_inverse(typename TC::TiledLayout_TV{});

109 auto mn = inv.with_shape(shape(typename TC::Tiler_MN{}));

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

111 auto inv_fmn = right_inverse(fmn);

112 auto values = logical_product(make_layout(V), inv_fmn);

113 auto fv = zipped_divide(

114 values, make_layout(typename TC::AtomNumVal{}));

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

116 auto tt = zipped_divide(tCrA.layout(), tile);

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

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

119 auto result = vt(_, selection);

120

121 show("inverse with MN shape: ", mn);

122 show("frg_layout_mn: ", fmn);

123 show("inverse frg_layout_mn: ", inv_fmn);

124 show("logical_product: ", values);

125 show("frg_layout_v: ", fv);

126 show("tile: ", tile);

127 show("t_tensor: ", tt);

128 show("v_tensor: ", vt);

129 show("selection: ", selection);

130 show("result: ", result);

131 }

132

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

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

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

136 cute::print(" ");

137 }

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

139

140 for (int v = 0; v < size<0>(tCrA); ++v) {

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

142 cute::print(coords(v, 0, 0));

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

144 }

145

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

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

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

149 auto slot = tXrA.layout()(make_coord(a, b), 0, 0);

150 std::printf(" r[%d]", int(slot));

151 }

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

153 }

154 }

155 std::printf("Verified %d destination matches and %d source references.\n",

156 checked_dst, checked_src);

157}

158