1// Numerical checks for cute_layout_tuple_complement_walkthrough_zh.md.

2// Runs on the host; no GPU instructions or product/divide operations are used.

3#include <cute/tensor.hpp>

4#include <cstdio>

5#include <cstdlib>

6#include <type_traits>

7

8using namespace cute;

9

10namespace {

11int checks = 0;

12void expect(bool ok, char const* label) {

13 if (!ok) {

14 std::fprintf(stderr,"FAILED: %s\n",label);

15 std::exit(1);

16 }

17 ++checks;

18}

19template<class T> void show(char const* label, T const& x) {

20 std::printf("%-25s",label); cute::print(x); std::printf("\n");

21}

22template<class L>

23void json_values(FILE* f,char const* key,L const& l,bool comma=true) {

24 std::fprintf(f," \"%s\": [",key);

25 for(int i=0;i<int(size(l));++i) std::fprintf(f,"%s%d",i?",":"",int(l(i)));

26 std::fprintf(f,"]%s\n",comma?",":"");

27}

28template<class L>

29void json_matrix(FILE* f,char const* key,L const& l) {

30 std::fprintf(f," \"%s\": [",key);

31 for(int m=0;m<int(size<0>(l));++m) {

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

33 for(int n=0;n<int(size<1>(l));++n)

34 std::fprintf(f,"%s%d",n?",":"",int(l(m,n)));

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

36 }

37 std::fprintf(f,"],\n");

38}

39}

40

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

42 auto tree=make_tuple(_2{},make_tuple(_3{},_4{}));

43 expect(int(rank(tree))==2,"tuple rank");

44 expect(int(depth(tree))==2,"tuple depth");

45 expect(int(size(tree))==24,"tuple product of leaves");

46 expect(int(get<1,0>(tree))==3,"nested tuple access");

47 expect(int(rank(_6{}))==1 && int(depth(_6{}))==0,"scalar IntTuple");

48 expect(int(depth(make_tuple(_6{})))==1,"one-element tuple depth");

49 static_assert(std::is_same_v<decltype(make_shape(_2{},_3{})),

50 decltype(make_stride(_2{},_3{}))>);

51 static_assert(std::is_same_v<decltype(make_shape(_2{},_3{})),

52 decltype(make_coord(_2{},_3{}))>);

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

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

55 auto padded=make_layout(Shape<_2,_3>{},Stride<_5,_1>{});

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

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

58 expect(int(row(m,n))==3*m+n,"row-major formula");

59 expect(int(col(m,n))==m+2*n,"column-major formula");

60 expect(int(padded(m,n))==5*m+n,"padded formula");

61 expect(int(broadcast(m,n))==n,"broadcast formula");

62 int i=m+2*n;

63 expect(int(row(i))==int(row(m,n)),"integer coordinate is shape order");

64 }

65 expect(int(size(padded))==6 && int(cosize(padded))==8,"padded size and cosize");

66 expect(int(size(broadcast))==6 && int(cosize(broadcast))==3,"broadcast size and cosize");

67 auto nested_shape=Shape<_2,Shape<_2,_3>>{};

68 auto nested_stride=Stride<_3,Stride<_6,_1>>{};

69 auto h=make_layout(nested_shape,nested_stride);

70 expect(bool(congruent(nested_shape,nested_stride)),"shape/stride tree agreement");

71 expect(!bool(congruent(nested_shape,make_stride(_3{},_6{},_1{}))),"flat stride has wrong tree");

72 expect(int(rank(h))==2 && int(depth(h))==2,"nested layout rank and depth");

73 expect(int(size(h))==12 && int(size<1>(h))==6,"nested layout sizes");

74 expect(int(get<1,1>(shape(h)))==3,"shape leaf access");

75 for(int m=0;m<2;++m) for(int a=0;a<2;++a) for(int b=0;b<3;++b) {

76 int j=a+2*b;

77 int i=m+2*j;

78 int expected=3*m+6*a+b;

79 expect(int(h(make_coord(m,make_coord(a,b))))==expected,"fully nested coordinate");

80 expect(int(h(m,j))==expected,"partially flattened coordinate");

81 expect(int(h(i))==expected,"integer nested coordinate");

82 expect(int(flatten(h)(i))==expected,"flatten preserves 1D input mapping");

83 }

84 auto left=make_layout(nested_shape);

85 auto right=make_layout(nested_shape,LayoutRight{});

86 expect(bool(left==make_layout(nested_shape,Stride<_1,Stride<_2,_4>>{})),"LayoutLeft leaf strides");

87 expect(bool(right==make_layout(nested_shape,Stride<_6,Stride<_3,_1>>{})),"LayoutRight leaf strides");

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

89 auto compact_c=coalesce(compact);

90 expect(int(rank(compact_c))==1,"coalesce rank");

91 for(int i=0;i<6;++i) expect(int(compact_c(i))==int(compact(i)),"coalesce input mapping");

92 int n=3;

93 auto mixed=make_layout(make_shape(_2{},n),make_stride(n,_1{}));

94 static_assert(is_static_v<decltype(row)>);

95 static_assert(!is_static_v<decltype(mixed)>);

96 for(int i=0;i<6;++i) expect(int(mixed(i))==int(row(i)),"mixed static/dynamic mapping");

97

98 show("tuple",tree);

99 show("row",row); show("column",col); show("padded",padded); show("broadcast",broadcast);

100 show("nested",h); show("nested mode 1",layout<1>(h));

101 show("default nested strides",left); show("LayoutRight nested",right);

102 show("flatten(nested)",flatten(h)); show("coalesced compact",compact_c); show("mixed",mixed);

103 if(argc>1) {

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

105 if(!f) return 2;

106 std::fprintf(f,"{\n");

107 json_matrix(f,"row",row);json_matrix(f,"column",col);

108 json_matrix(f,"padded",padded);json_matrix(f,"broadcast",broadcast);json_matrix(f,"nested",h);

109 json_values(f,"row_integer",row);

110 std::fprintf(f," \"checks\": %d\n}\n",checks);

111 std::fclose(f);

112 }

113 std::printf("Verified %d tuple and layout mapping checks.\n",checks);

114}

115