Padding with border reflect 101

[guillaume-michel]

Hello,
I am trying to do the following with CuTe:

• create a HxW layout row-major (OK)
• create a new layout with the first layout at its center + a halo around it with 101 reflection. this new layout has shape: (H+2halo)x(W+2halo).

What I have so far:

#include <vector>

#include <cute/tensor.hpp>
#include <cute/util/print_tensor.hpp>

using namespace cute;

CUTE_HOST_DEVICE int reflect101(int p, int len) {
    if (len <= 1) {
        return 0;
    }

    while (p < 0 || p >= len) {
        p = (p < 0) ? -p : (2 * len - 2 - p);
    }

    return p;
}

int main() {
    constexpr auto halo = _1;

    constexpr auto h = _4;
    constexpr auto w = _4;

    constexpr auto n = h * w;
    std::vector<int> data(n);

    std::iota(data.begin(), data.end(), 0);

    auto t1 = make_tensor(data.data(), make_shape(h, w), LayoutRight{});

    print(t1);

    print_tensor(t1);

    // t2 = ...?

    return 0;
}

Could you please help me?

[guillaume-michel]

I managed to go a bit further:

#include <cute/tensor.hpp>
#include <cute/layout_composed.hpp>
#include <cute/util/print_tensor.hpp>

#include <numeric>
#include <vector>

#include <spdlog/spdlog.h>

using namespace cute;

namespace {
// 1-D mirror-101 mapper acting on LINEAR indices
template <int Dim, int Halo>
struct Reflect101Index1D {
    static_assert(Dim >= 2, "Reflect101 requires Dim >= 2");

    // Optional: size of codomain (fold target)
    CUTE_HOST_DEVICE
    constexpr auto cosize() const { return Int<Dim>{}; }

    // Map extended index [0..Dim+2*Halo) -> [0..Dim)
    CUTE_HOST_DEVICE
    constexpr int operator()(int idx_ext) const {
        int j = idx_ext - Halo;      // shift: [-Halo, Dim+Halo-1]
        int period = 2 * (Dim - 1);  // triangular-wave period

        int m = j % period;

        if (m < 0) {
            m += period;
        }

        return (m <= Dim - 1)
                   ? m
                   : (period - m);  // 101 reflection (no duplicated endpoints)
    }
};

template <int Dim, int Halo>
CUTE_HOST_DEVICE auto makeReflect101Layout1d() {
    auto ext_id =
        make_layout(Int<Dim + 2 * Halo>{});  // 1-D identity (domain provider)
    return make_composed_layout(Reflect101Index1D<Dim, Halo>{}, Int<0>{},
                                ext_id);
}

// for ADL
template <int Dim, int Halo>
void print(const Reflect101Index1D<Dim, Halo>& l) {
    cute::print("Reflect101Index1D<%d, %d>", Dim, Halo);
}
}  // namespace

int main() {
    constexpr auto halo = Int<2>{};
    constexpr auto h = Int<4>{};
    constexpr auto w = Int<4>{};

    // Base data: 4x4 row-major
    constexpr auto n = h * w;
    std::vector<int> data(n);
    std::iota(data.begin(), data.end(), 0);

    auto t1 = cute::make_tensor(data.data(), cute::make_shape(h, w),
                                cute::LayoutRight{});

    print("t1:\n");
    print_tensor(t1);
    print("\n");

    auto ry = makeReflect101Layout1d<h, halo>();

    spdlog::info("rank ry: {}", rank(ry));
    spdlog::info("shape: {}", shape(ry));
    spdlog::info("size: {}", size(ry));

    auto t = make_tensor(data.data(), ry);

    for (int i = 0; i < size<0>(t); ++i) {
        print(t(i));
        print(" ");
    }
    print("\n");

    print("t:\n");
    print_tensor(t);
    print("\n");

    return 0;
}

It prints the following:

t1:
ptr[32b](0x5757ea13f7e0) o (_4,_4):(_4,_1):
    0    1    2    3
    4    5    6    7
    8    9   10   11
   12   13   14   15

[2025-10-29 12:08:24.521] [info] rank ry: 1
[2025-10-29 12:08:24.522] [info] shape: 8
[2025-10-29 12:08:24.522] [info] size: 8
2 1 0 1 2 3 2 1 
t:
ptr[32b](0x5757ea13f7e0) o Reflect101Index1D<4, 2> o _0 o _8:_1:
    2
    1
    0
    1
    2
    3
    2
    1

which is the expected result.

Now I want to go 2D. How to compose my 1D ComposedLayout for get a 2D layout?
Thanks in advance.