1// Host-only CuTe layout examples used by cute_layout_product_walkthrough_zh.md.

2// Build: nvcc -std=c++17 -I include cute_layout_product_trace.cu -o /tmp/cute_layout_product_trace

3#include <cute/tensor.hpp>

4#include <cstdio>

5#include <cstdlib>

6

7using namespace cute;

8

9int checked = 0;

10void expect(int actual, int expected, char const* label) {

11 if (actual != expected) {

12 std::fprintf(stderr, "%s: got %d, expected %d\n", label, actual, expected);

13 std::exit(1);

14 }

15 ++checked;

16}

17

18template<class T>

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

20 std::printf("%-23s", name);

21 cute::print(x);

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

23}

24

25int main(int argc, char** argv) {

26 auto dense = logical_product(Layout<_3,_1>{}, Layout<_4,_1>{});

27 auto gaps = logical_product(Layout<_3,_2>{}, Layout<_4,_1>{});

28 auto ordered = logical_product(Layout<_3,_1>{},

29 Layout<Shape<_2,_2>,Stride<_2,_1>>{});

30 auto gather = logical_product(Layout<_4,_1>{}, Layout<_2,_2>{});

31 auto a = make_layout(Shape<_2,_3>{}, Stride<_3,_1>{});

32 auto b = make_layout(Shape<_2,_2>{}, Stride<_1,_2>{});

33 auto lp = logical_product(a,b);

34 auto zp = zipped_product(a,b);

35 auto tp = tiled_product(a,b);

36 auto fp = flat_product(a,b);

37 auto bp = blocked_product(a,b);

38 auto rp = raked_product(a,b);

39 auto target = tile_to_shape(a, Shape<_4,_6>{});

40 auto by_mode = make_tile(Layout<_2,_3>{}, Layout<_2,_4>{});

41 auto ml = logical_product(a, by_mode);

42 auto mz = zipped_product(a, by_mode);

43 auto mt = tiled_product(a, by_mode);

44 auto mf = flat_product(a, by_mode);

45 auto naive = logical_product(a, make_tile(_2{},_2{}));

46 auto threads = Layout<_4,_1>{};

47 auto values = Layout<_2,_1>{};

48 auto copy_blocked = blocked_product(threads, values);

49 auto copy_raked = raked_product(threads, values);

50 auto copy_tv = right_inverse(copy_raked).with_shape(Shape<_4,_2>{});

51 auto mma_threads = tiled_product(Layout<_32,_1>{},

52 make_layout(Shape<_2,_2,_1>{}));

53

54 for (int j=0;j<4;++j) for (int i=0;i<3;++i) {

55 expect(dense(i,j), i+3*j, "dense");

56 expect(gaps(i,j), 2*i+j%2+6*(j/2), "gaps");

57 expect(ordered(i,j), i+3*(2*(j%2)+j/2), "ordered");

58 }

59 for (int j=0;j<2;++j) for (int i=0;i<4;++i)

60 expect(gather(i,j),i+8*j,"retile gather");

61 for (int q=0;q<2;++q) for (int p=0;p<2;++p)

62 for (int n=0;n<3;++n) for (int m=0;m<2;++m) {

63 int expected = 3*m+n+6*p+12*q;

64 auto local = make_coord(m,n);

65 auto tile = make_coord(p,q);

66 expect(lp(local,tile),expected,"logical");

67 expect(zp(local,tile),expected,"zipped");

68 expect(tp(local,p,q),expected,"tiled");

69 expect(fp(m,n,p,q),expected,"flat");

70 expect(bp(make_coord(m,p),make_coord(n,q)),expected,"blocked nested");

71 expect(bp(m+2*p,n+3*q),expected,"blocked scalar");

72 expect(rp(make_coord(p,m),make_coord(q,n)),expected,"raked nested");

73 expect(rp(p+2*m,q+2*n),expected,"raked scalar");

74 expect(target(m+2*p,n+3*q),expected,"tile_to_shape");

75 expect(ml(make_coord(m,p),make_coord(n,q)),expected,"by-mode logical");

76 expect(mz(local,tile),expected,"by-mode zipped");

77 expect(mt(local,p,q),expected,"by-mode tiled");

78 expect(mf(m,n,p,q),expected,"by-mode flat");

79 }

80 expect(naive(make_coord(0,1),make_coord(0,0)),1,"naive alias 1");

81 expect(naive(make_coord(0,0),make_coord(1,0)),1,"naive alias 2");

82 for (int v=0;v<2;++v) for (int t=0;t<4;++t) {

83 expect(copy_blocked(t+4*v),t+4*v,"copy blocked");

84 expect(copy_raked(v+2*t),t+4*v,"copy raked");

85 expect(copy_tv(t,v),v+2*t,"copy TV inverse");

86 }

87 for (int q=0;q<2;++q) for (int p=0;p<2;++p) for (int lane=0;lane<32;++lane)

88 expect(mma_threads(lane,p,q,0),lane+32*p+64*q,"MMA thread ID");

89

90 if (argc>1) {

91 FILE* f=std::fopen(argv[1],"w");

92 if (!f) return 2;

93 std::fprintf(f,"{\n \"dense\": [");

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

95 std::fprintf(f,"%s[",j?",":"");

96 for (int i=0;i<3;++i) std::fprintf(f,"%s%d",i?",":"",int(dense(i,j)));

97 std::fprintf(f,"]");

98 }

99 std::fprintf(f,"],\n \"gaps\": [");

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

101 std::fprintf(f,"%s[",j?",":"");

102 for (int i=0;i<3;++i) std::fprintf(f,"%s%d",i?",":"",int(gaps(i,j)));

103 std::fprintf(f,"]");

104 }

105 std::fprintf(f,"],\n \"blocked\": [");

106 for (int m=0;m<4;++m) {

107 std::fprintf(f,"%s[",m?",":"");

108 for (int n=0;n<6;++n) std::fprintf(f,"%s%d",n?",":"",int(bp(m,n)));

109 std::fprintf(f,"]");

110 }

111 std::fprintf(f,"],\n \"raked\": [");

112 for (int m=0;m<4;++m) {

113 std::fprintf(f,"%s[",m?",":"");

114 for (int n=0;n<6;++n) std::fprintf(f,"%s%d",n?",":"",int(rp(m,n)));

115 std::fprintf(f,"]");

116 }

117 std::fprintf(f,"],\n \"checked\": %d\n}\n",checked);

118 std::fclose(f);

119 }

120

121 show("dense",dense); show("complement(3:1)",complement(Layout<_3,_1>{}));

122 show("complement(3:2)",complement(Layout<_3,_2>{})); show("gaps",gaps);

123 show("ordered repetitions",ordered); show("retile gather",gather);

124 show("A",a); show("B",b); show("logical",lp); show("zipped",zp);

125 show("tiled",tp); show("flat",fp); show("blocked",bp); show("raked",rp);

126 show("tile_to_shape",target); show("by-mode logical",ml); show("by-mode zipped",mz);

127 show("by-mode tiled",mt); show("by-mode flat",mf); show("naive by-mode",naive);

128 show("copy blocked",copy_blocked); show("copy raked",copy_raked);

129 show("copy TV",copy_tv); show("MMA threads",mma_threads);

130 std::printf("Verified %d coordinate mappings. No GPU required.\n", checked);

131}

132