"src/include/blockwise_4d_tensor_op.hpp" did not exist on "df228b3cf514ec23dcc1decacfc1973e7f9016d9"
ConstantTensorDescriptor.hip.hpp 26.7 KB
Newer Older
Chao Liu's avatar
Chao Liu committed
1
#pragma once
2
#include "common.hip.hpp"
Chao Liu's avatar
Chao Liu committed
3

4
template <class Lengths>
Chao Liu's avatar
Chao Liu committed
5
__host__ __device__ constexpr auto calculate_tensor_strides_default_rank_packed(Lengths)
6
{
7
8
    return reverse_inclusive_scan_sequence(Lengths{}.PopFront(), std::multiplies<index_t>{})
        .PushBack(Number<1>{});
9
10
}

11
template <class Lengths, index_t Align>
Chao Liu's avatar
Chao Liu committed
12
13
__host__ __device__ constexpr auto calculate_tensor_strides_default_rank_aligned(Lengths,
                                                                                 Number<Align>)
Chao Liu's avatar
Chao Liu committed
14
{
15
16
    constexpr index_t L_back_align =
        Align * mod_conv::integer_divide_ceiler<index_t>{}(Lengths{}.Back(), Align);
Chao Liu's avatar
Chao Liu committed
17

Chao Liu's avatar
Chao Liu committed
18
    return calculate_tensor_strides_default_rank_packed(
19
        Lengths{}.Modify(Number<Lengths{}.GetSize() - 1>{}, Number<L_back_align>{}));
20
21
}

22
23
// MemoryRanks of dimensions is for conversion from offset to multi-index
template <class Lengths, class Strides, class MemoryRanks>
Chao Liu's avatar
Chao Liu committed
24
25
struct ConstantTensorDescriptor
{
Chao Liu's avatar
Chao Liu committed
26
27
    using Type = ConstantTensorDescriptor;

28
    static constexpr index_t nDim = Lengths::GetSize();
Chao Liu's avatar
Chao Liu committed
29
30
31

    __host__ __device__ constexpr ConstantTensorDescriptor()
    {
32
33
34
35
36
37
38
39
40
        static_assert(Lengths::GetSize() == Strides::GetSize() &&
                          Lengths::GetSize() == MemoryRanks::GetSize(),
                      "nDim not consistent");

#if 0 // require sequence_sort, but it's not implemented yet
        static_assert(is_same<typename sequence_sort<MemoryRanks>::SortedSeqType,
                              typename arithmetic_sequence_gen<0, nDim, 1>::SeqType>::value,
                      "wrong! invalid MemoryRanks");
#endif
Chao Liu's avatar
Chao Liu committed
41
42
    }

Chao Liu's avatar
Chao Liu committed
43
    __host__ __device__ static constexpr index_t GetNumOfDimension() { return nDim; }
Chao Liu's avatar
Chao Liu committed
44

45
    __host__ __device__ static constexpr auto GetLengths() { return Lengths{}; }
Chao Liu's avatar
Chao Liu committed
46

47
48
49
    __host__ __device__ static constexpr auto GetStrides() { return Strides{}; }

    __host__ __device__ static constexpr auto GetMemoryRanks() { return MemoryRanks{}; }
Chao Liu's avatar
Chao Liu committed
50

Chao Liu's avatar
Chao Liu committed
51
    template <index_t I>
52
    __host__ __device__ static constexpr index_t GetLength(Number<I>)
Chao Liu's avatar
Chao Liu committed
53
    {
Chao Liu's avatar
Chao Liu committed
54
        return Lengths{}.Get(Number<I>{});
Chao Liu's avatar
Chao Liu committed
55
56
    }

Chao Liu's avatar
Chao Liu committed
57
    template <index_t I>
58
    __host__ __device__ static constexpr index_t GetStride(Number<I>)
Chao Liu's avatar
Chao Liu committed
59
    {
Chao Liu's avatar
Chao Liu committed
60
        return Strides{}.Get(Number<I>{});
Chao Liu's avatar
Chao Liu committed
61
62
    }

63
64
65
66
67
68
    template <index_t I>
    __host__ __device__ static constexpr index_t GetMemoryRank(Number<I>)
    {
        return MemoryRanks{}.Get(Number<I>{});
    }

Chao Liu's avatar
Chao Liu committed
69
70
71
72
73
74
    template <class T>
    __host__ __device__ static constexpr bool ContainMultipleOriginalDimensions(T)
    {
        return false;
    }

75
    __host__ __device__ static constexpr index_t GetElementSize()
Chao Liu's avatar
Chao Liu committed
76
    {
Chao Liu's avatar
Chao Liu committed
77
        return accumulate_on_sequence(Lengths{}, std::multiplies<index_t>{}, Number<1>{});
78
    }
79

80
    // WRONG! ReorderGivenOld2New is broken
Chao Liu's avatar
Chao Liu committed
81
    template <class Align = Number<1>>
82
    __host__ __device__ static constexpr index_t GetElementSpace(Align align = Align{})
Chao Liu's avatar
Chao Liu committed
83
    {
84
85
86
87
88
89
90
#if 0
        constexpr auto lengths_in_rank = GetLengths().ReorderGivenOld2New(MemoryRank{});
        constexpr auto strides_in_rank = GetStrides().ReorderGivenOld2new(MemoryRank{});

        constexpr index_t element_space_unaligned = accumulate_on_sequence(
            (lengths_in_rank - Number<1>{}) * strides_in_rank, std::plus<index_t>{}, Number<1>{});
#else // WRONG! align shouldbe applied to the last memory rank, not the last tensor dimension
Chao Liu's avatar
Chao Liu committed
91
92
        constexpr index_t element_space_unaligned = accumulate_on_sequence(
            (GetLengths() - Number<1>{}) * GetStrides(), std::plus<index_t>{}, Number<1>{});
93
#endif
Chao Liu's avatar
Chao Liu committed
94
95

        return align.Get() * ((element_space_unaligned + align.Get() - 1) / align.Get());
Chao Liu's avatar
Chao Liu committed
96
    }
Chao Liu's avatar
Chao Liu committed
97

98
    template <index_t NSize>
99
    __host__ __device__ static index_t GetOffsetFromMultiIndex(Array<index_t, NSize> multi_id)
Chao Liu's avatar
Chao Liu committed
100
    {
101
        static_assert(NSize == nDim, "wrong! Dimension not consistent");
Chao Liu's avatar
Chao Liu committed
102

103
        index_t offset = 0;
Chao Liu's avatar
Chao Liu committed
104

105
        static_for<0, nDim, 1>{}([&](auto IDim) {
Chao Liu's avatar
Chao Liu committed
106
            constexpr index_t idim = IDim.Get();
107
            offset += multi_id[idim] * GetStride(IDim);
108
        });
Chao Liu's avatar
Chao Liu committed
109

110
        return offset;
111
112
    }

113
    template <class... Is>
114
    __host__ __device__ static index_t GetOffsetFromMultiIndex(Is... is)
115
    {
116
        return GetOffsetFromMultiIndex(Array<index_t, sizeof...(Is)>{is...});
117
118
    }

119
    template <index_t... Is>
120
    __host__ __device__ static constexpr index_t GetOffsetFromMultiIndex(Sequence<Is...>)
121
122
123
    {
        static_assert(sizeof...(Is) == nDim, "wrong! Dimension not consistent");

Chao Liu's avatar
Chao Liu committed
124
125
        constexpr auto multi_id = Sequence<Is...>{};

Chao Liu's avatar
Chao Liu committed
126
        return accumulate_on_sequence(multi_id * GetStrides(), std::plus<index_t>{}, Number<0>{});
127
128
    }

129
130
#if 0 // ReorderGivenOld2new is broken
    __host__ __device__ static Array<index_t, nDim> GetMultiIndexFromOffset(index_t offset)
Chao Liu's avatar
Chao Liu committed
131
    {
132
133
134
135
        Array<index_t, nDim> ranked_multi_id;

        constexpr auto ranked_strides =
            GetStrides().ReorderGivenOld2new(MemoryRanks{}); // check this
136

137
        // calculate index in each of the dimensions in the order of their rank (not dimension)
138
        static_for<0, nDim - 1, 1>{}([&](auto IDim) {
139
140
141
142
            constexpr index_t idim   = IDim.Get();
            constexpr index_t stride = ranked_strides.Get(Number<idim>{});
            ranked_multi_id[idim]    = offset / stride;
            offset -= ranked_multi_id[idim] * stride;
143
144
        });

145
        ranked_multi_id[nDim - 1] = offset / ranked_strides.Get(Number<nDim - 1>{});
146

147
        return reorder_array_given_new2old(ranked_multi_id, MemoryRanks{}); // check this
Chao Liu's avatar
Chao Liu committed
148
    }
149
#endif
150

151
    __host__ __device__ static Array<index_t, nDim> GetMultiIndexFrom1dIndex(index_t id)
152
    {
153
154
        Array<index_t, nDim> multi_id;

Chao Liu's avatar
Chao Liu committed
155
        constexpr auto dummy_strides = calculate_tensor_strides_default_rank_packed(GetLengths());
156
157
158
159
160
161
162
163
164
165
166
167

        // calculate index in each of the dimensions in the order of their dimension (not rank)
        static_for<0, nDim - 1, 1>{}([&](auto IDim) {
            constexpr index_t idim   = IDim.Get();
            constexpr index_t stride = dummy_strides.Get(Number<idim>{});
            multi_id[idim]           = id / stride;
            id -= multi_id[idim] * stride;
        });

        multi_id[nDim - 1] = id / dummy_strides.Get(Number<nDim - 1>{});

        return multi_id;
168
    }
Chao Liu's avatar
Chao Liu committed
169

170
    // WRONG! Ranks is broken
Chao Liu's avatar
Chao Liu committed
171
    template <index_t... IDims>
Chao Liu's avatar
Chao Liu committed
172
    __host__ __device__ static constexpr auto Extract(Number<IDims>... extract_dims)
Chao Liu's avatar
Chao Liu committed
173
    {
Chao Liu's avatar
Chao Liu committed
174
175
        static_assert(sizeof...(IDims) <= GetNumOfDimension(),
                      "wrong! too many number of dimensions to be extracted");
Chao Liu's avatar
Chao Liu committed
176

177
178
179
180
181
182
183
184
185
186
187
        using extract_lengths = decltype(Lengths{}.Extract(extract_dims...));
        using extract_strides = decltype(Strides{}.Extract(extract_dims...));
        using extract_ranks   = decltype(MemoryRanks{}.Extract(extract_dims...));

#if 0
        using new_ranks = typename sequence_sort<extract_ranks>::Original2SortedType;
#else // WRONG! TODO:: implement sequence_sort
        using new_ranks = typename arithmetic_sequence_gen<0, sizeof...(IDims), 1>::SeqType;
#endif

        return ConstantTensorDescriptor<extract_lengths, extract_strides, new_ranks>{};
Chao Liu's avatar
Chao Liu committed
188
189
    }

Chao Liu's avatar
Chao Liu committed
190
191
192
193
194
195
    template <index_t... IDims>
    __host__ __device__ static constexpr auto Extract(Sequence<IDims...>)
    {
        return Extract(Number<IDims>{}...);
    }

Chao Liu's avatar
Chao Liu committed
196
197
198
    template <index_t IDim, index_t SliceLen>
    __host__ __device__ static constexpr auto Slice(Number<IDim>, Number<SliceLen>)
    {
199
200
201
        using slice_lengths = decltype(Lengths{}.Modify(Number<IDim>{}, Number<SliceLen>{}));

        return ConstantTensorDescriptor<slice_lengths, Strides, MemoryRanks>{};
Chao Liu's avatar
Chao Liu committed
202
203
    }

204
205
206
207
208
209
210
211
212
    template <index_t Threashold, index_t Delta>
    struct f_fold_impl
    {
        __host__ __device__ constexpr index_t operator()(index_t x) const
        {
            return x > Threashold ? x + Delta : x;
        }
    };

Chao Liu's avatar
Chao Liu committed
213
    template <index_t IDim, index_t... FoldIntervals>
Chao Liu's avatar
Chao Liu committed
214
    __host__ __device__ static constexpr auto Fold(Number<IDim>, Number<FoldIntervals>...)
Chao Liu's avatar
Chao Liu committed
215
    {
Chao Liu's avatar
Chao Liu committed
216
217
        constexpr auto fold_intervals = Sequence<FoldIntervals...>{};

Chao Liu's avatar
Chao Liu committed
218
        constexpr index_t fold_intervals_product =
Chao Liu's avatar
Chao Liu committed
219
220
221
222
            accumulate_on_sequence(fold_intervals, std::multiplies<index_t>{}, Number<1>{});

        constexpr auto unfold_length = GetLength(Number<IDim>{});
        constexpr auto unfold_stride = GetStride(Number<IDim>{});
223
        constexpr auto unfold_rank   = GetMemoryRank(Number<IDim>{});
Chao Liu's avatar
Chao Liu committed
224
225
226

        // length of the dimension to be folded needs to be dividable by fold_interval_product,
        // otherwise, folding is invalid
Chao Liu's avatar
Chao Liu committed
227
        static_assert(unfold_length % fold_intervals_product == 0,
Chao Liu's avatar
Chao Liu committed
228
229
230
231
                      "wrong! length on the dimension to be folded cannot be evenly divided!");

        // folded lengths
        constexpr auto fold_lengths =
Chao Liu's avatar
Chao Liu committed
232
            Sequence<unfold_length / fold_intervals_product>{}.Append(fold_intervals);
Chao Liu's avatar
Chao Liu committed
233
234

        // folded strides
Chao Liu's avatar
Chao Liu committed
235
236
        constexpr auto fold_strides =
            Number<unfold_stride>{} *
Chao Liu's avatar
Chao Liu committed
237
238
            reverse_inclusive_scan_sequence(fold_intervals.PushBack(Number<1>{}),
                                            std::multiplies<index_t>{});
Chao Liu's avatar
Chao Liu committed
239

240
241
242
243
244
        // folded_ranks
        constexpr auto fold_ranks =
            typename arithmetic_sequence_gen<unfold_rank,
                                             unfold_rank + fold_intervals.GetSize() + 1,
                                             1>::SeqType{};
Chao Liu's avatar
Chao Liu committed
245

246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
        // increase the ranks that are larger than unfold_rank
        constexpr auto tmp_ranks = transform_sequences(
            f_fold_impl<unfold_rank, fold_intervals.GetSize()>{}, GetMemoryRanks());

        // left and right
        constexpr auto left = typename arithmetic_sequence_gen<0, IDim, 1>::SeqType{};
        constexpr auto right =
            typename arithmetic_sequence_gen<IDim + 1, GetNumOfDimension(), 1>::SeqType{};

        constexpr auto new_lengths =
            GetLengths().Extract(left).Append(fold_lengths).Append(GetLengths().Extract(right));
        constexpr auto new_strides =
            GetStrides().Extract(left).Append(fold_strides).Append(GetStrides().Extract(right));
        constexpr auto new_ranks =
            tmp_ranks.Extract(left).Append(fold_ranks).Append(tmp_ranks.Extract(right));

        static_assert(new_ranks.GetSize() == new_lengths.GetSize(), "wrong!");
        static_assert(fold_ranks.GetSize() == fold_lengths.GetSize(), "wrong!");

        return ConstantTensorDescriptor<decltype(new_lengths),
                                        decltype(new_strides),
                                        decltype(new_ranks)>{};
Chao Liu's avatar
Chao Liu committed
268
269
    }

270
271
272
273
274
275
276
277
278
    template <index_t Threashold, index_t Delta>
    struct f_unfold_impl
    {
        __host__ __device__ constexpr index_t operator()(index_t x) const
        {
            return x > Threashold ? x - Delta : x;
        }
    };

Chao Liu's avatar
Chao Liu committed
279
280
281
    template <index_t FirstUnfoldDim, index_t LastUnfoldDim>
    __host__ __device__ static constexpr auto Unfold(Number<FirstUnfoldDim>, Number<LastUnfoldDim>)
    {
Chao Liu's avatar
Chao Liu committed
282
283
284
285
        static_assert(FirstUnfoldDim >= 0 && LastUnfoldDim < nDim &&
                          FirstUnfoldDim <= LastUnfoldDim,
                      "wrong! should have FirstUnfoldDim <= LastUnfoldDim!");

Chao Liu's avatar
Chao Liu committed
286
#if 0 // cannot compile: compiler complain about constexpr
Chao Liu's avatar
Chao Liu committed
287
288
        // dimensions to be unfold need to be in descending order (w.r.t. strides), and need to be
        // packed in memory, otherwise, unfolding is invalid
Chao Liu's avatar
Chao Liu committed
289
290
        static_for<FirstUnfoldDim, LastUnfoldDim, 1>{}([&](auto IDim_) {
            constexpr auto IDim    = decltype(IDim_){};
291
292
293
            constexpr auto IDim_p1 = IDim + Number<1>{};

            // check stride
Chao Liu's avatar
Chao Liu committed
294
            static_assert(
295
                GetStride(IDim) >= GetStride(IDim_p1),
Chao Liu's avatar
Chao Liu committed
296
297
                "wrong! dimensions to be unfolded need to be in descending order w.r.t strides");

298
299
            // check if packed
            static_assert(GetStride(IDim_p1) * GetLength(IDim_p1) == GetStride(IDim),
Chao Liu's avatar
Chao Liu committed
300
                          "wrong! dimensions to be unfolded need to be packed");
301

Chao Liu's avatar
Chao Liu committed
302
            // check ranks
303
304
305
            static_assert(GetMemoryRank(IDim_p1) == GetMemoryRank(IDim) + 1,
                          "wrong! ranks of dimensions to be unfolded need to be in increasing and "
                          "continuous ranks");
Chao Liu's avatar
Chao Liu committed
306
        });
Chao Liu's avatar
Chao Liu committed
307
#endif
Chao Liu's avatar
Chao Liu committed
308

Chao Liu's avatar
Chao Liu committed
309
        // left and right
310
311
312
313
314
315
316
        constexpr auto left = typename arithmetic_sequence_gen<0, FirstUnfoldDim, 1>::SeqType{};
        constexpr auto middle =
            typename arithmetic_sequence_gen<FirstUnfoldDim, LastUnfoldDim + 1, 1>::SeqType{};
        constexpr auto right =
            typename arithmetic_sequence_gen<LastUnfoldDim + 1, GetNumOfDimension(), 1>::SeqType{};

        // unfolded length, stride and rank
Chao Liu's avatar
Chao Liu committed
317
318
319
320
321
        constexpr index_t unfold_length = accumulate_on_sequence(
            GetLengths().Extract(middle), std::multiplies<index_t>{}, Number<1>{});

        constexpr index_t unfold_stride = GetStride(Number<LastUnfoldDim>{});

322
323
324
325
        constexpr index_t unfold_rank = GetMemoryRank(Number<FirstUnfoldDim>{});

        // decrease the ranks that are larger than the rank of LastUnfoldDim
        constexpr auto tmp_ranks =
Chao Liu's avatar
Chao Liu committed
326
327
328
            transform_sequences(f_unfold_impl<GetMemoryRank(Number<LastUnfoldDim>{}),
                                              LastUnfoldDim - FirstUnfoldDim + 1>{},
                                GetMemoryRanks());
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356

        // new lengths, strides and ranks
        constexpr auto new_lengths = GetLengths()
                                         .Extract(left)
                                         .PushBack(Number<unfold_length>{})
                                         .Append(GetLengths().Extract(right));

        constexpr auto new_strides = GetStrides()
                                         .Extract(left)
                                         .PushBack(Number<unfold_stride>{})
                                         .Append(GetStrides().Extract(right));

        constexpr auto new_ranks = tmp_ranks.Extract(left)
                                       .PushBack(Number<unfold_rank>{})
                                       .Append(tmp_ranks.Extract(right));

        return ConstantTensorDescriptor<decltype(new_lengths),
                                        decltype(new_strides),
                                        decltype(new_ranks)>{};
    }

    template <class MapNew2Old>
    __host__ __device__ static constexpr auto ReorderGivenNew2Old(MapNew2Old)
    {
        return ConstantTensorDescriptor<decltype(Lengths{}.ReorderGivenNew2Old(MapNew2Old{})),
                                        decltype(Strides{}.ReorderGivenNew2Old(MapNew2Old{})),
                                        decltype(
                                            MemoryRanks{}.ReorderGivenNew2Old(MapNew2Old{}))>{};
Chao Liu's avatar
Chao Liu committed
357
358
    }

359
360
361
#if 0 // require sequence_sort, which is not implemented yet
    template <class MapOld2New>
    __host__ __device__ static constexpr auto ReorderGivenOld2New(MapOld2New)
Chao Liu's avatar
Chao Liu committed
362
    {
363
364
365
366
        return ConstantTensorDescriptor<decltype(Lengths{}.ReorderGivenOld2New(MapOld2New{})),
                                        decltype(Strides{}.ReorderGivenOld2New(MapOld2New{})),
                                        decltype(
                                            MemoryRanks{}.ReorderGivenOld2New(MapOld2New{}))>{};
Chao Liu's avatar
Chao Liu committed
367
    }
368
#endif
Chao Liu's avatar
Chao Liu committed
369
};
Chao Liu's avatar
Chao Liu committed
370
371

template <class Lengths>
Chao Liu's avatar
Chao Liu committed
372
__host__ __device__ constexpr auto make_ConstantTensorDescriptor_default_rank_packed(Lengths)
Chao Liu's avatar
Chao Liu committed
373
{
Chao Liu's avatar
Chao Liu committed
374
    using Strides     = decltype(calculate_tensor_strides_default_rank_packed(Lengths{}));
375
376
    using MemoryRanks = typename arithmetic_sequence_gen<0, Lengths::GetSize(), 1>::SeqType;
    return ConstantTensorDescriptor<Lengths, Strides, MemoryRanks>{};
Chao Liu's avatar
Chao Liu committed
377
378
379
}

template <class Lengths, class Strides>
Chao Liu's avatar
Chao Liu committed
380
__host__ __device__ constexpr auto make_ConstantTensorDescriptor_default_rank(Lengths, Strides)
Chao Liu's avatar
Chao Liu committed
381
{
382
383
    using MemoryRanks = typename arithmetic_sequence_gen<0, Lengths::GetSize(), 1>::SeqType;
    return ConstantTensorDescriptor<Lengths, Strides, MemoryRanks>{};
Chao Liu's avatar
Chao Liu committed
384
385
}

Chao Liu's avatar
Chao Liu committed
386
template <class Lengths, index_t Align>
Chao Liu's avatar
Chao Liu committed
387
388
__host__ __device__ constexpr auto make_ConstantTensorDescriptor_default_rank_aligned(Lengths,
                                                                                      Number<Align>)
Chao Liu's avatar
Chao Liu committed
389
{
390
    using Strides =
Chao Liu's avatar
Chao Liu committed
391
        decltype(calculate_tensor_strides_default_rank_aligned(Lengths{}, Number<Align>{}));
392
393
    using MemoryRanks = typename arithmetic_sequence_gen<0, Lengths::GetSize(), 1>::SeqType;
    return ConstantTensorDescriptor<Lengths, Strides, MemoryRanks>{};
Chao Liu's avatar
Chao Liu committed
394
395
}

Chao Liu's avatar
Chao Liu committed
396
template <class TDesc>
Chao Liu's avatar
Chao Liu committed
397
__host__ __device__ void print_ConstantTensorDescriptor(TDesc, const char* s)
Chao Liu's avatar
Chao Liu committed
398
{
Chao Liu's avatar
Chao Liu committed
399
    constexpr index_t ndim = TDesc::GetNumOfDimension();
Chao Liu's avatar
Chao Liu committed
400

Chao Liu's avatar
Chao Liu committed
401
    static_assert(ndim >= 2 && ndim <= 10, "wrong!");
Chao Liu's avatar
Chao Liu committed
402

Chao Liu's avatar
Chao Liu committed
403
    static_if<ndim == 2>{}([&](auto fwd) {
Chao Liu's avatar
Chao Liu committed
404
405
406
        constexpr auto I0 = Number<0>{};
        constexpr auto I1 = Number<1>{};

Chao Liu's avatar
Chao Liu committed
407
408
409
        constexpr auto desc = fwd(TDesc{});

        printf("%s dim %u, lengths {%u %u}, strides {%u %u}, ranks {%u %u}\n",
Chao Liu's avatar
Chao Liu committed
410
               s,
Chao Liu's avatar
Chao Liu committed
411
               desc.GetNumOfDimension(),
Chao Liu's avatar
Chao Liu committed
412
413
414
               desc.GetLength(I0),
               desc.GetLength(I1),
               desc.GetStride(I0),
Chao Liu's avatar
Chao Liu committed
415
416
417
418
419
420
               desc.GetStride(I1),
               desc.GetMemoryRank(I0),
               desc.GetMemoryRank(I1));
    });

    static_if<ndim == 3>{}([&](auto fwd) {
421
422
423
424
        constexpr auto I0 = Number<0>{};
        constexpr auto I1 = Number<1>{};
        constexpr auto I2 = Number<2>{};

Chao Liu's avatar
Chao Liu committed
425
426
427
        constexpr auto desc = fwd(TDesc{});

        printf("%s dim %u, lengths {%u %u %u}, strides {%u %u %u}, ranks {%u %u %u}\n",
428
               s,
Chao Liu's avatar
Chao Liu committed
429
               desc.GetNumOfDimension(),
430
431
432
433
434
               desc.GetLength(I0),
               desc.GetLength(I1),
               desc.GetLength(I2),
               desc.GetStride(I0),
               desc.GetStride(I1),
Chao Liu's avatar
Chao Liu committed
435
436
437
438
439
440
441
               desc.GetStride(I2),
               desc.GetMemoryRank(I0),
               desc.GetMemoryRank(I1),
               desc.GetMemoryRank(I2));
    });

    static_if<ndim == 4>{}([&](auto fwd) {
Chao Liu's avatar
Chao Liu committed
442
443
444
445
446
        constexpr auto I0 = Number<0>{};
        constexpr auto I1 = Number<1>{};
        constexpr auto I2 = Number<2>{};
        constexpr auto I3 = Number<3>{};

Chao Liu's avatar
Chao Liu committed
447
448
449
        constexpr auto desc = fwd(TDesc{});

        printf("%s dim %u, lengths {%u %u %u %u}, strides {%u %u %u %u}, ranks {%u %u %u %u}\n",
Chao Liu's avatar
Chao Liu committed
450
               s,
Chao Liu's avatar
Chao Liu committed
451
               desc.GetNumOfDimension(),
Chao Liu's avatar
Chao Liu committed
452
453
454
455
456
457
458
               desc.GetLength(I0),
               desc.GetLength(I1),
               desc.GetLength(I2),
               desc.GetLength(I3),
               desc.GetStride(I0),
               desc.GetStride(I1),
               desc.GetStride(I2),
Chao Liu's avatar
Chao Liu committed
459
460
461
462
463
464
465
466
               desc.GetStride(I3),
               desc.GetMemoryRank(I0),
               desc.GetMemoryRank(I1),
               desc.GetMemoryRank(I2),
               desc.GetMemoryRank(I3));
    });

    static_if<ndim == 5>{}([&](auto fwd) {
467
468
469
470
471
472
        constexpr auto I0 = Number<0>{};
        constexpr auto I1 = Number<1>{};
        constexpr auto I2 = Number<2>{};
        constexpr auto I3 = Number<3>{};
        constexpr auto I4 = Number<4>{};

Chao Liu's avatar
Chao Liu committed
473
474
475
476
        constexpr auto desc = fwd(TDesc{});

        printf("%s dim %u, lengths {%u %u %u %u %u}, strides {%u %u %u %u %u}, ranks {%u %u %u %u "
               "%u}\n",
477
               s,
Chao Liu's avatar
Chao Liu committed
478
               desc.GetNumOfDimension(),
479
480
481
482
483
484
485
486
487
               desc.GetLength(I0),
               desc.GetLength(I1),
               desc.GetLength(I2),
               desc.GetLength(I3),
               desc.GetLength(I4),
               desc.GetStride(I0),
               desc.GetStride(I1),
               desc.GetStride(I2),
               desc.GetStride(I3),
Chao Liu's avatar
Chao Liu committed
488
489
490
491
492
493
494
495
496
               desc.GetStride(I4),
               desc.GetMemoryRank(I0),
               desc.GetMemoryRank(I1),
               desc.GetMemoryRank(I2),
               desc.GetMemoryRank(I3),
               desc.GetMemoryRank(I4));
    });

    static_if<ndim == 6>{}([&](auto fwd) {
497
498
499
500
501
502
503
        constexpr auto I0 = Number<0>{};
        constexpr auto I1 = Number<1>{};
        constexpr auto I2 = Number<2>{};
        constexpr auto I3 = Number<3>{};
        constexpr auto I4 = Number<4>{};
        constexpr auto I5 = Number<5>{};

Chao Liu's avatar
Chao Liu committed
504
505
506
507
        constexpr auto desc = fwd(TDesc{});

        printf("%s dim %u, lengths {%u %u %u %u %u %u}, strides {%u %u %u %u %u %u}, ranks {%u %u "
               "%u %u %u %u}\n",
508
               s,
Chao Liu's avatar
Chao Liu committed
509
               desc.GetNumOfDimension(),
510
511
512
513
514
515
516
517
518
519
520
               desc.GetLength(I0),
               desc.GetLength(I1),
               desc.GetLength(I2),
               desc.GetLength(I3),
               desc.GetLength(I4),
               desc.GetLength(I5),
               desc.GetStride(I0),
               desc.GetStride(I1),
               desc.GetStride(I2),
               desc.GetStride(I3),
               desc.GetStride(I4),
Chao Liu's avatar
Chao Liu committed
521
522
523
524
525
526
527
528
529
530
               desc.GetStride(I5),
               desc.GetMemoryRank(I0),
               desc.GetMemoryRank(I1),
               desc.GetMemoryRank(I2),
               desc.GetMemoryRank(I3),
               desc.GetMemoryRank(I4),
               desc.GetMemoryRank(I5));
    });

    static_if<ndim == 7>{}([&](auto fwd) {
531
532
533
534
535
536
537
538
        constexpr auto I0 = Number<0>{};
        constexpr auto I1 = Number<1>{};
        constexpr auto I2 = Number<2>{};
        constexpr auto I3 = Number<3>{};
        constexpr auto I4 = Number<4>{};
        constexpr auto I5 = Number<5>{};
        constexpr auto I6 = Number<6>{};

Chao Liu's avatar
Chao Liu committed
539
540
541
542
        constexpr auto desc = fwd(TDesc{});

        printf("%s dim %u, lengths {%u %u %u %u %u %u %u}, strides {%u %u %u %u %u %u %u}, ranks "
               "{%u %u %u %u %u %u %u}\n",
543
               s,
Chao Liu's avatar
Chao Liu committed
544
               desc.GetNumOfDimension(),
545
546
547
548
549
550
551
552
553
554
555
556
557
               desc.GetLength(I0),
               desc.GetLength(I1),
               desc.GetLength(I2),
               desc.GetLength(I3),
               desc.GetLength(I4),
               desc.GetLength(I5),
               desc.GetLength(I6),
               desc.GetStride(I0),
               desc.GetStride(I1),
               desc.GetStride(I2),
               desc.GetStride(I3),
               desc.GetStride(I4),
               desc.GetStride(I5),
Chao Liu's avatar
Chao Liu committed
558
559
560
561
562
563
564
565
566
567
568
               desc.GetStride(I6),
               desc.GetMemoryRank(I0),
               desc.GetMemoryRank(I1),
               desc.GetMemoryRank(I2),
               desc.GetMemoryRank(I3),
               desc.GetMemoryRank(I4),
               desc.GetMemoryRank(I5),
               desc.GetMemoryRank(I6));
    });

    static_if<ndim == 8>{}([&](auto fwd) {
569
570
571
572
573
574
575
576
577
        constexpr auto I0 = Number<0>{};
        constexpr auto I1 = Number<1>{};
        constexpr auto I2 = Number<2>{};
        constexpr auto I3 = Number<3>{};
        constexpr auto I4 = Number<4>{};
        constexpr auto I5 = Number<5>{};
        constexpr auto I6 = Number<6>{};
        constexpr auto I7 = Number<7>{};

Chao Liu's avatar
Chao Liu committed
578
579
580
581
        constexpr auto desc = fwd(TDesc{});

        printf("%s dim %u, lengths {%u %u %u %u %u %u %u %u}, strides {%u %u %u %u %u %u %u %u}, "
               "ranks {%u %u %u %u %u %u %u %u}\n",
582
               s,
Chao Liu's avatar
Chao Liu committed
583
               desc.GetNumOfDimension(),
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
               desc.GetLength(I0),
               desc.GetLength(I1),
               desc.GetLength(I2),
               desc.GetLength(I3),
               desc.GetLength(I4),
               desc.GetLength(I5),
               desc.GetLength(I6),
               desc.GetLength(I7),
               desc.GetStride(I0),
               desc.GetStride(I1),
               desc.GetStride(I2),
               desc.GetStride(I3),
               desc.GetStride(I4),
               desc.GetStride(I5),
               desc.GetStride(I6),
Chao Liu's avatar
Chao Liu committed
599
600
601
602
603
604
605
606
607
608
609
610
               desc.GetStride(I7),
               desc.GetMemoryRank(I0),
               desc.GetMemoryRank(I1),
               desc.GetMemoryRank(I2),
               desc.GetMemoryRank(I3),
               desc.GetMemoryRank(I4),
               desc.GetMemoryRank(I5),
               desc.GetMemoryRank(I6),
               desc.GetMemoryRank(I7));
    });

    static_if<ndim == 9>{}([&](auto fwd) {
Chao Liu's avatar
Chao Liu committed
611
612
613
614
615
616
617
618
619
620
        constexpr auto I0 = Number<0>{};
        constexpr auto I1 = Number<1>{};
        constexpr auto I2 = Number<2>{};
        constexpr auto I3 = Number<3>{};
        constexpr auto I4 = Number<4>{};
        constexpr auto I5 = Number<5>{};
        constexpr auto I6 = Number<6>{};
        constexpr auto I7 = Number<7>{};
        constexpr auto I8 = Number<8>{};

Chao Liu's avatar
Chao Liu committed
621
622
        constexpr auto desc = fwd(TDesc{});

Chao Liu's avatar
tidy yp  
Chao Liu committed
623
        printf("%s dim %u, lengths {%u %u %u %u %u %u %u %u %u}, strides {%u %u %u %u %u %u %u %u "
Chao Liu's avatar
Chao Liu committed
624
               "%u}, ranks {%u %u %u %u %u %u %u %u %u}\n",
Chao Liu's avatar
Chao Liu committed
625
               s,
Chao Liu's avatar
Chao Liu committed
626
               desc.GetNumOfDimension(),
Chao Liu's avatar
Chao Liu committed
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
               desc.GetLength(I0),
               desc.GetLength(I1),
               desc.GetLength(I2),
               desc.GetLength(I3),
               desc.GetLength(I4),
               desc.GetLength(I5),
               desc.GetLength(I6),
               desc.GetLength(I7),
               desc.GetLength(I8),
               desc.GetStride(I0),
               desc.GetStride(I1),
               desc.GetStride(I2),
               desc.GetStride(I3),
               desc.GetStride(I4),
               desc.GetStride(I5),
               desc.GetStride(I6),
               desc.GetStride(I7),
Chao Liu's avatar
Chao Liu committed
644
645
646
647
648
649
650
651
652
653
654
655
656
               desc.GetStride(I8),
               desc.GetMemoryRank(I0),
               desc.GetMemoryRank(I1),
               desc.GetMemoryRank(I2),
               desc.GetMemoryRank(I3),
               desc.GetMemoryRank(I4),
               desc.GetMemoryRank(I5),
               desc.GetMemoryRank(I6),
               desc.GetMemoryRank(I7),
               desc.GetMemoryRank(I8));
    });

    static_if<ndim == 10>{}([&](auto fwd) {
Chao Liu's avatar
Chao Liu committed
657
658
659
660
661
662
663
664
665
666
667
        constexpr auto I0 = Number<0>{};
        constexpr auto I1 = Number<1>{};
        constexpr auto I2 = Number<2>{};
        constexpr auto I3 = Number<3>{};
        constexpr auto I4 = Number<4>{};
        constexpr auto I5 = Number<5>{};
        constexpr auto I6 = Number<6>{};
        constexpr auto I7 = Number<7>{};
        constexpr auto I8 = Number<8>{};
        constexpr auto I9 = Number<9>{};

Chao Liu's avatar
Chao Liu committed
668
669
        constexpr auto desc = fwd(TDesc{});

Chao Liu's avatar
tidy yp  
Chao Liu committed
670
        printf("%s dim %u, lengths {%u %u %u %u %u %u %u %u %u %u}, strides {%u %u %u %u %u %u %u "
Chao Liu's avatar
Chao Liu committed
671
               "%u %u %u}, ranks {%u %u %u %u %u %u %u %u %u %u}\n",
Chao Liu's avatar
Chao Liu committed
672
               s,
Chao Liu's avatar
Chao Liu committed
673
               desc.GetNumOfDimension(),
Chao Liu's avatar
Chao Liu committed
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
               desc.GetLength(I0),
               desc.GetLength(I1),
               desc.GetLength(I2),
               desc.GetLength(I3),
               desc.GetLength(I4),
               desc.GetLength(I5),
               desc.GetLength(I6),
               desc.GetLength(I7),
               desc.GetLength(I8),
               desc.GetLength(I9),
               desc.GetStride(I0),
               desc.GetStride(I1),
               desc.GetStride(I2),
               desc.GetStride(I3),
               desc.GetStride(I4),
               desc.GetStride(I5),
               desc.GetStride(I6),
               desc.GetStride(I7),
               desc.GetStride(I8),
Chao Liu's avatar
Chao Liu committed
693
694
695
696
697
698
699
700
701
702
703
704
               desc.GetStride(I9),
               desc.GetMemoryRank(I0),
               desc.GetMemoryRank(I1),
               desc.GetMemoryRank(I2),
               desc.GetMemoryRank(I3),
               desc.GetMemoryRank(I4),
               desc.GetMemoryRank(I5),
               desc.GetMemoryRank(I6),
               desc.GetMemoryRank(I7),
               desc.GetMemoryRank(I8),
               desc.GetMemoryRank(I9));
    });
Chao Liu's avatar
Chao Liu committed
705
}