1#include <cute/tensor.hpp>

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

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

4

5#include <cstdio>

6#include <cstdlib>

7

8using namespace cute;

9

10CUTE_HOST_DEVICE auto make_example_copy()

11{

12 return make_tiled_copy(

13 Copy_Atom<UniversalCopy<uint128_t>, float>{},

14 make_layout(make_shape(_4{}, _8{}), make_stride(_8{}, _1{})),

15 make_layout(make_shape(_1{}, _4{}), make_stride(_4{}, _1{})));

16}

17

18template <class T>

19void show(char const* name, T const& value)

20{

21 std::printf("%s", name);

22 cute::print(value);

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

24}

25

26void require(bool condition, char const* message)

27{

28 if (!condition) {

29 std::fprintf(stderr, "Verification failed: %s\n", message);

30 std::exit(1);

31 }

32}

33

34// Inspect ldmatrix ownership only. Do not execute a device instruction.

35void verify_reference_side()

36{

37 using Atom = Copy_Atom<SM75_U32x4_LDSM_N, uint32_t>;

38 auto tc = make_tiled_copy_impl(

39 Atom{}, Layout<Shape<_32,_4>,Stride<_1,_32>>{}, Shape<_128>{});

40 using TC = decltype(tc);

41 auto source = tc.get_layoutS_TV();

42 auto destination = tc.get_layoutD_TV();

43 show("ldmatrix/AtomLayoutSrc: ", typename TC::AtomLayoutSrc{});

44 show("ldmatrix/AtomLayoutDst: ", typename TC::AtomLayoutDst{});

45 show("ldmatrix/AtomLayoutRef: ", typename TC::AtomLayoutRef{});

46 show("ldmatrix/TiledLayout_TV: ", typename TC::TiledLayout_TV{});

47 show("ldmatrix/get_layoutS_TV: ", source);

48 show("ldmatrix/get_layoutD_TV: ", destination);

49 for (int t = 0; t < 32; ++t) {

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

51 require(source(make_coord(t,v)) == 4*t+v, "ldmatrix source ownership");

52 require(destination(make_coord(t,v)) == t+32*v, "ldmatrix destination ownership");

53 }

54 }

55 require(source(make_coord(8,0)) == destination(make_coord(0,1)),

56 "different lanes must refer to the same ldmatrix element");

57 auto wrong_reference = make_tiled_copy_impl(

58 Atom{}, Layout<Shape<_32,_4>,Stride<_4,_1>>{}, Shape<_128>{});

59 auto wrong_source = wrong_reference.get_layoutS_TV();

60 require(wrong_source(make_coord(8,0)) == 1,

61 "source TV cannot be passed as destination reference TV unchanged");

62 std::printf("Verified ldmatrix reference side: 128 source/destination mappings.\n");

63}

64

65// Host-only verification: enumerate coordinates, check complete coverage, and

66// actually copy through each thread's fragment. No CUDA device is needed.

67template <int M, int N, class TiledCopy>

68void verify_case(TiledCopy const& tc, int values_per_thread, char const* name)

69{

70 alignas(16) float source[M * N];

71 alignas(16) float destination[M * N];

72 int visits[M * N]{};

73 for (int i = 0; i < M * N; ++i) {

74 source[i] = float(i + 1);

75 destination[i] = -1.0f;

76 }

77

78 auto physical = make_layout(make_shape(Int<M>{}, Int<N>{}),

79 make_stride(Int<N>{}, _1{}));

80 auto src = make_tensor(source + 0, physical);

81 auto dst = make_tensor(destination + 0, physical);

82 auto coordinates = make_identity_tensor(shape(physical));

83

84 for (int t = 0; t < 32; ++t) {

85 auto thread = tc.get_slice(t);

86 auto source_coordinates = thread.partition_S(coordinates);

87 auto destination_coordinates = thread.partition_D(coordinates);

88 auto s = thread.partition_S(src);

89 auto d = thread.partition_D(dst);

90 auto fragment = make_fragment_like(d);

91 require(int(size(s)) == values_per_thread * (M / 4) * (N / (8 * values_per_thread)),

92 "per-thread element count");

93

94 for (int rm = 0; rm < M / 4; ++rm) {

95 for (int rn = 0; rn < N / (8 * values_per_thread); ++rn) {

96 for (int v = 0; v < values_per_thread; ++v) {

97 auto coordinate = source_coordinates(v, rm, rn);

98 int m = t / 8 + 4 * rm;

99 int n = values_per_thread * (t % 8) + v + 8 * values_per_thread * rn;

100 require(coordinate == make_coord(m, n), "source coordinate formula");

101 require(destination_coordinates(v, rm, rn) == coordinate, "source/destination agreement");

102 require(int(s.layout()(v, rm, rn)) ==

103 physical(m, n) - physical(t / 8, values_per_thread * (t % 8)),

104 "physical offset relative to sliced pointer");

105 require(++visits[m * N + n] == 1, "non-overlapping thread ownership");

106 }

107 }

108 }

109

110 copy(tc, s, fragment);

111 copy(tc, fragment, d);

112 }

113

114 for (int i = 0; i < M * N; ++i) {

115 require(visits[i] == 1, "complete tile coverage");

116 require(destination[i] == source[i], "copied values");

117 }

118 std::printf("Verified %s: %d coordinates and %d copied floats.\n", name, M * N, M * N);

119}

120

121template <int M, int N>

122CUTE_DEVICE void device_copy(float const* source, float* destination)

123{

124 auto tc = make_example_copy();

125 auto physical = make_layout(make_shape(Int<M>{}, Int<N>{}),

126 make_stride(Int<N>{}, _1{}));

127 auto src = make_tensor(make_gmem_ptr(source), physical);

128 auto dst = make_tensor(make_gmem_ptr(destination), physical);

129 auto thread = tc.get_slice(threadIdx.x);

130 auto s = thread.partition_S(src);

131 auto d = thread.partition_D(dst);

132 auto fragment = make_fragment_like(d);

133 copy(tc, s, fragment);

134 copy(tc, fragment, d);

135}

136

137// These kernels are compiled for PTX inspection; main() does not launch them.

138// Launch contract: one block of 32 threads; 16-byte aligned device pointers.

139extern "C" __global__ void cute_copy_4x32(float const* src, float* dst)

140{

141 device_copy<4, 32>(src, dst);

142}

143

144extern "C" __global__ void cute_copy_8x128(float const* src, float* dst)

145{

146 device_copy<8, 128>(src, dst);

147}

148

149int main()

150{

151 verify_reference_side();

152 auto threads = make_layout(make_shape(_4{}, _8{}), make_stride(_8{}, _1{}));

153 auto values = make_layout(make_shape(_1{}, _4{}), make_stride(_4{}, _1{}));

154 auto logical = logical_product(threads, values);

155 auto mn = raked_product(threads, values);

156 auto inv = right_inverse(mn);

157 auto tv = inv.with_shape(make_shape(size(threads), size(values)));

158 auto tiler = product_each(shape(mn));

159 auto tc = make_example_copy();

160 using TC = decltype(tc);

161

162 show("thr_layout: ", threads);

163 show("val_layout: ", values);

164 show("logical_product: ", logical);

165 show("layout_mn: ", mn);

166 show("right_inverse: ", inv);

167 auto mn_coalesced = coalesce(mn);

168 auto inverse_preprod = cute::fold(shape(mn_coalesced), cute::tuple<_1>{},

169 [](auto acc, auto extent) { return append(acc, extent * back(acc)); });

170 auto inverse_strides = wrap(stride(mn_coalesced));

171 using InverseSorted = detail::SortByKey<decltype(inverse_strides), tuple_seq<decltype(inverse_strides)>>;

172 show("right_inverse/coalesced input: ", mn_coalesced);

173 show("right_inverse/prefix products: ", inverse_preprod);

174 show("right_inverse/stride-sorted modes: ", typename InverseSorted::val_type{});

175 show("right_inverse/input strides: ", inverse_strides);

176 for (int i = 0; i < size(inv); ++i) {

177 require(mn(inv(i)) == i, "right inverse must recover every TV code");

178 require(inv(i) == 16 * (i % 8) + i / 8, "right inverse coordinate formula");

179 }

180 show("layout_tv: ", tv);

181 show("tiler: ", tiler);

182 show("blocked_product: ", blocked_product(threads, values));

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

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

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

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

187 show("get_layoutS_TV: ", tc.get_layoutS_TV());

188 show("get_layoutD_TV: ", tc.get_layoutD_TV());

189

190 auto physical = make_layout(make_shape(_4{}, _32{}), make_stride(_32{}, _1{}));

191 alignas(16) float storage[128]{};

192 auto tensor = make_tensor(storage + 0, physical);

193 auto coordinate = make_identity_tensor(shape(physical));

194 auto split = zipped_divide(physical, tiler);

195 auto atom_tv = zipped_divide(tv, make_shape(typename TC::AtomNumThr{}, typename TC::AtomNumVal{}));

196 auto src_to_ref = right_inverse(typename TC::AtomLayoutRef{}).compose(typename TC::AtomLayoutSrc{});

197 auto target_tv = atom_tv.compose(src_to_ref, _);

198 auto thread_value = coalesce(zip(target_tv), Shape<_1, Shape<_1, _1>>{});

199 auto tv_tensor = split.compose(thread_value, _);

200 show("zipped_divide(physical,tiler): ", split);

201 show("atom_layout_TV: ", atom_tv);

202 show("source_to_reference: ", src_to_ref);

203 show("trg_layout_TV: ", target_tv);

204 show("zip(trg_layout_TV): ", zip(target_tv));

205 show("thrval2mn: ", thread_value);

206 show("coalesce without profile: ", coalesce(zip(target_tv)));

207 show("coalesce profile (1,1): ", coalesce(zip(target_tv), Shape<_1,_1>{}));

208 show("coalesce profile (7,(7,7)): ", coalesce(zip(target_tv), Shape<_7,Shape<_7,_7>>{}));

209 for (int t = 0; t < 32; ++t) {

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

211 auto before = target_tv(make_coord(make_coord(_0{}, v), make_coord(t, _0{})));

212 auto after_zip = zip(target_tv)(make_coord(make_coord(_0{}, t), make_coord(v, _0{})));

213 auto after_coalesce = thread_value(make_coord(t, make_coord(v, _0{})));

214 require(before == after_zip && before == after_coalesce,

215 "zip and profile coalesce must preserve corresponding coordinate mappings");

216 }

217 }

218 show("tv_tensor: ", tv_tensor);

219 show("tidfrg_S: ", TC::tidfrg_S(physical));

220 show("partition_S: ", tc.get_slice(0).partition_S(tensor).layout());

221 show("partition_D: ", tc.get_slice(0).partition_D(tensor).layout());

222

223 for (int t : {0, 1, 7, 8, 31}) {

224 auto part = tc.get_slice(t).partition_S(coordinate);

225 std::printf("thread %d:", t);

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

227 cute::print(" ");

228 cute::print(part(v, 0, 0));

229 int compact_coordinate = tv(t, v);

230 require(compact_coordinate == t / 8 + 4 * (4 * (t % 8) + v), "compact logical coordinate");

231 require(int(mn(compact_coordinate)) == t + 32 * v, "MN-to-TV encoding");

232 }

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

234 }

235

236 auto large_physical = make_layout(make_shape(_8{}, _128{}), make_stride(_128{}, _1{}));

237 alignas(16) float large_storage[1024]{};

238 auto large_tensor = make_tensor(large_storage + 0, large_physical);

239 show("large partition_S: ", tc.get_slice(0).partition_S(large_tensor).layout());

240

241 auto eight_values = make_tiled_copy(

242 Copy_Atom<UniversalCopy<uint128_t>, float>{}, threads,

243 make_layout(make_shape(_1{}, _8{}), make_stride(_8{}, _1{})));

244 using TC8 = decltype(eight_values);

245 show("eight-value TiledLayout_TV: ", typename TC8::TiledLayout_TV{});

246 show("eight-value tiler: ", typename TC8::Tiler_MN{});

247 alignas(16) float eight_storage[256]{};

248 auto eight_tensor = make_tensor(eight_storage + 0,

249 make_layout(make_shape(_4{}, _64{}), make_stride(_64{}, _1{})));

250 show("eight-value partition_S: ", eight_values.get_slice(0).partition_S(eight_tensor).layout());

251

252 // Changing the physical layout preserves logical ownership while changing

253 // offsets. A strided column-major partition cannot use this 128-bit Atom.

254 auto column_major = make_tensor(storage + 0, make_layout(make_shape(_4{}, _32{})));

255 show("column-major partition_S: ", tc.get_slice(0).partition_S(column_major).layout());

256 for (int t = 0; t < 32; ++t) {

257 auto part = tc.get_slice(t).partition_S(column_major);

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

259 require(int(part.layout()(v, 0, 0)) == 4 * v, "column-major value stride");

260 }

261 }

262

263 verify_case<4, 32>(tc, 4, "base tile");

264 verify_case<8, 128>(tc, 4, "repeated tile");

265 verify_case<4, 64>(eight_values, 8, "two Atoms per thread");

266 return 0;

267}

268