1#include <cute/tensor.hpp>

2

3#include <cstdio>

4#include <cstdlib>

5

6using namespace cute;

7

8namespace {

9

10int checks = 0;

11

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

13{

14 if (!condition) {

15 std::fprintf(stderr, "FAILED: %s\n", message);

16 std::exit(1);

17 }

18 ++checks;

19}

20

21template <class T>

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

23{

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

25 cute::print(value);

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

27}

28

29template <class L>

30void table(char const* label, L const& layout)

31{

32 show(label, layout);

33 for (int m = 0; m < int(size<0>(layout)); ++m) {

34 for (int n = 0; n < int(size<1>(layout)); ++n) {

35 std::printf("%4d", int(layout(m,n)));

36 }

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

38 }

39}

40

41} // namespace

42

43int main()

44{

45 auto row = make_layout(Shape<_3,_4>{}, Stride<_4,_1>{});

46 auto column = make_layout(Shape<_3,_4>{}, Stride<_1,_3>{});

47 auto padded = make_layout(Shape<_3,_4>{}, Stride<_6,_1>{});

48 auto broadcast = make_layout(Shape<_3,_4>{}, Stride<_0,_1>{});

49

50 table("row: ", row);

51 table("column: ", column);

52 table("padded: ", padded);

53 table("broadcast: ", broadcast);

54

55 for (int m = 0; m < 3; ++m) {

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

57 require(int(row(m,n)) == 4*m+n, "row offset");

58 require(int(column(m,n)) == m+3*n, "column offset");

59 require(int(padded(m,n)) == 6*m+n, "padded offset");

60 require(int(broadcast(m,n)) == n, "broadcast alias");

61 }

62 }

63

64 require(int(size(padded)) == 12, "padded size");

65 require(int(cosize(padded)) == 16, "padded cosize");

66 require(int(size(broadcast)) == 12, "broadcast size");

67 require(int(cosize(broadcast)) == 4, "broadcast cosize");

68

69 std::printf("row(i):");

70 for (int i = 0; i < 12; ++i) {

71 int m = i % 3;

72 int n = i / 3;

73 require(int(row(i)) == 4*m+n, "integer coordinate uses shape");

74 require(int(row(i)) == int(row(m,n)), "integer and tuple coordinate agree");

75 std::printf(" %d", int(row(i)));

76 }

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

78 require(int(row(5)) == 9, "row(5) is an input coordinate");

79 require(row.get_hier_coord(5) == make_coord(1,1), "index-to-coordinate for compact row layout");

80 require(int(row.get_1d_coord(5)) == 4, "index-to-1D-coordinate for compact row layout");

81

82 auto h = make_layout(Shape<_2,Shape<_2,_2>>{},

83 Stride<_4,Stride<_2,_1>>{});

84 table("hierarchical: ", h);

85 require(int(rank(h)) == 2, "hierarchical rank");

86 require(int(depth(h)) == 2, "hierarchical depth");

87 require(int(size(h)) == 8, "hierarchical size");

88 require(int(size<1>(h)) == 4, "nested mode size");

89 require(int(size<1,0>(h)) == 2, "nested submode size");

90 show("shape<1>(hierarchical): ", shape<1>(h));

91 show("layout<1>(hierarchical): ", layout<1>(h));

92

93 for (int m = 0; m < 2; ++m) {

94 for (int j = 0; j < 4; ++j) {

95 int a = j % 2;

96 int b = j / 2;

97 require(int(h(m,j)) == 4*m+2*a+b, "hierarchical 2D formula");

98 require(int(h(make_coord(m,make_coord(a,b)))) == int(h(m,j)),

99 "natural coordinate and flat subcoordinate agree");

100 require(int(h(m+2*j)) == int(h(m,j)), "hierarchical integer coordinate");

101 }

102 }

103

104 auto flat = flatten(h);

105 auto regrouped = group<1,3>(flat);

106 show("flatten(h): ", flat);

107 show("regrouped: ", regrouped);

108 require(regrouped == h, "regrouping recovers shape and stride");

109 for (int i = 0; i < 8; ++i) {

110 require(int(flat(i)) == int(h(i)), "flatten preserves integer-coordinate mapping");

111 }

112

113 auto left = make_layout(Shape<_2,Shape<_2,_2>>{});

114 auto right = make_layout(Shape<_2,Shape<_2,_2>>{}, LayoutRight{});

115 show("hierarchical LayoutLeft: ", left);

116 show("hierarchical LayoutRight: ", right);

117 require(left == make_layout(Shape<_2,Shape<_2,_2>>{},

118 Stride<_1,Stride<_2,_4>>{}), "default left strides");

119 require(right == h, "right strides across all leaves");

120

121 int columns = 4;

122 auto dynamic = make_layout(make_shape(3,columns), make_stride(columns,1));

123 auto mixed = make_layout(make_shape(_3{},columns), make_stride(columns,_1{}));

124 show("static row: ", row);

125 show("dynamic row: ", dynamic);

126 show("mixed row: ", mixed);

127 static_assert(is_static<decltype(row)>::value);

128 static_assert(!is_static<decltype(dynamic)>::value);

129 static_assert(!is_static<decltype(mixed)>::value);

130 for (int i = 0; i < 12; ++i) {

131 require(int(dynamic(i)) == int(row(i)), "dynamic row mapping");

132 require(int(mixed(i)) == int(row(i)), "mixed row mapping");

133 }

134

135 int storage[16];

136 for (int i = 0; i < 16; ++i) {

137 storage[i] = 100+i;

138 }

139 auto tensor = make_tensor(storage+0, padded);

140 auto tensor_row = tensor(1,_);

141 auto tensor_column = tensor(_,2);

142 auto row_layout = padded(1,_);

143 auto [row_slice, row_offset] = slice_and_offset(make_coord(1,_), padded);

144 show("row slice layout: ", row_layout);

145 show("slice_and_offset layout: ", row_slice);

146 show("slice_and_offset offset: ", row_offset);

147 show("tensor row layout: ", tensor_row.layout());

148 show("tensor column layout: ", tensor_column.layout());

149 require(tensor_row.data() == storage+6, "row view pointer offset");

150 require(tensor_column.data() == storage+2, "column view pointer offset");

151 require(int(row_offset) == 6, "row fixed-coordinate contribution");

152 require(int(row_layout(0)) == 0, "Layout slice excludes fixed-coordinate offset");

153 for (int j = 0; j < 4; ++j) {

154 require(int(row_slice(j)) == j, "row slice stride");

155 require(tensor_row(j) == storage[6+j], "row view values");

156 }

157 for (int m = 0; m < 3; ++m) {

158 require(tensor_column(m) == storage[6*m+2], "column view values");

159 }

160 tensor_row(2) = -7;

161 require(storage[8] == -7, "view write aliases original storage");

162 tensor_row(2) = 108;

163

164 // A data layout and a thread/value layout answer different questions.

165 // Here a tiny 2x4 tile is distributed across four logical threads.

166 auto tile_data = make_layout(Shape<_2,_4>{}, Stride<_4,_1>{});

167 auto tv = make_layout(Shape<_4,_2>{}, Stride<_1,_4>{});

168 auto tv_to_address = composition(tile_data,tv);

169 show("data layout: ", tile_data);

170 show("thread/value layout: ", tv);

171 show("thread/value -> data address: ", tv_to_address);

172 int visits[8]{};

173 for (int t = 0; t < 4; ++t) {

174 for (int v = 0; v < 2; ++v) {

175 int q = int(tv(t,v));

176 int m = q%2;

177 int n = q/2;

178 int address = int(tile_data(m,n));

179 require(int(tv_to_address(t,v)) == address, "thread/value composition");

180 require(++visits[address] == 1, "non-overlapping tiny tile ownership");

181 std::printf("t=%d v=%d -> q=%d -> (%d,%d) -> offset=%d\n",

182 t,v,q,m,n,address);

183 }

184 }

185 for (int i = 0; i < 8; ++i) {

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

187 }

188

189 // Read actual SM80 MMA metadata on the host; no GPU instruction is run.

190 auto mma = make_tiled_mma(SM80_16x8x16_F32F16F16F32_TN{});

191 auto mma_tv = mma.get_layoutA_TV();

192 auto mma_physical = make_layout(Shape<_16,_16>{}, Stride<_16,_1>{});

193 show("SM80 A atom layout: ", typename MMA_Traits<SM80_16x8x16_F32F16F16F32_TN>::ALayout{});

194 show("SM80 A thread/value layout: ", mma_tv);

195 int mma_visits[256]{};

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

197 for (int v = 0; v < 8; ++v) {

198 int a = v%2;

199 int b = (v/2)%2;

200 int c = v/4;

201 int q = int(mma_tv(lane,v));

202 int m = lane/4 + 8*b;

203 int k = 2*(lane%4) + a + 8*c;

204 require(q == m+16*k, "MMA A TV logical coordinate");

205 require(++mma_visits[q] == 1, "unique MMA A ownership");

206 if (lane < 2) {

207 std::printf("lane=%d v=%d -> A(%d,%d) q=%d physical_offset=%d\n",

208 lane,v,m,k,q,int(mma_physical(m,k)));

209 }

210 }

211 }

212 for (int i = 0; i < 256; ++i) {

213 require(mma_visits[i] == 1, "complete MMA A coverage");

214 }

215

216 std::printf("Verified %d checks; 8 tiny-tile values and 256 MMA A values covered.\n", checks);

217 std::printf("All examples run on the host; no CUDA kernel was launched.\n");

218}

219