Using SIDs as neighbour tables in functional API

include/gridtools/fn/sid_neighbor_table.hpp

@@ -0,0 +1,57 @@
+/*
+ * GridTools
+ *
+ * Copyright (c) 2014-2022, ETH Zurich
+ * All rights reserved.
+ *
+ * Please, refer to the LICENSE file in the root directory.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+#pragma once
+
+#include <cstdint>
+#include <gridtools/common/array.hpp>
+#include <gridtools/fn/unstructured.hpp>
+#include <gridtools/sid/concept.hpp>
+#include <type_traits>
+
+namespace gridtools::fn::sid_neighbor_table {
+    namespace sid_neighbor_table_impl_ {
+        template <class IndexDimension, class NeighborDimension, int32_t MaxNumNeighbors, class Sid>
+        struct sid_neighbor_table {
+            Sid sid;
+        };
+
+        template <class IndexDimension, class NeighborDimension, int32_t MaxNumNeighbors, class Sid>
+        auto neighbor_table_neighbors(
+            sid_neighbor_table<IndexDimension, NeighborDimension, MaxNumNeighbors, Sid> const &table, size_t index) {
+            using element_type = sid::element_type<Sid>;
+
+            const auto ptr = sid_get_origin(table.sid);
+            const auto strides = sid_get_strides(table.sid);
+
+            const auto index_stride = at_key<IndexDimension>(strides);
+            const auto neighbour_stride = at_key<NeighborDimension>(strides);
+
+            gridtools::array<element_type, MaxNumNeighbors> neighbors;
+            for (int32_t elementIdx = 0; elementIdx < MaxNumNeighbors; ++elementIdx) {
+                const auto element_ptr = ptr + index * index_stride + elementIdx * neighbour_stride;
+                neighbors[elementIdx] = *element_ptr();
+            }
+            return neighbors;
+        }
+
+        template <class IndexDimension, class NeighborDimension, int32_t MaxNumNeighbors, class Sid>
+        auto as_neighbor_table(Sid sid) {
+            static_assert(gridtools::tuple_util::size<decltype(sid_get_strides(std::declval<Sid>()))>::value == 2,
+                "Neighbor tables must have exactly two dimensions: the index dimension and the neighbor dimension");
+            static_assert(!std::is_same_v<IndexDimension, NeighborDimension>,
+                "The index dimension and the neighbor dimension must be different.");
+
+            return sid_neighbor_table<IndexDimension, NeighborDimension, MaxNumNeighbors, Sid>{std::move(sid)};
+        }
+    } // namespace sid_neighbor_table_impl_
+
+    using sid_neighbor_table_impl_::as_neighbor_table;
+
+} // namespace gridtools::fn::sid_neighbor_table

tests/unit_tests/fn/CMakeLists.txt

@@ -7,6 +7,7 @@ gridtools_add_unit_test(test_fn_run SOURCES test_fn_run.cpp)
 gridtools_add_unit_test(test_fn_column_stage SOURCES test_fn_column_stage.cpp)
 gridtools_add_unit_test(test_fn_stencil_stage SOURCES test_fn_stencil_stage.cpp LABELS fn)
 gridtools_add_unit_test(test_fn_unstructured SOURCES test_fn_unstructured.cpp LABELS fn)
+gridtools_add_unit_test(test_fn_sid_neighbor_table SOURCES test_fn_sid_neighbor_table.cpp LABELS fn)
 
 if(TARGET _gridtools_cuda)
     gridtools_add_unit_test(test_fn_backend_gpu_cuda

tests/unit_tests/fn/test_fn_sid_neighbor_table.cpp

@@ -0,0 +1,49 @@
+/*
+ * GridTools
+ *
+ * Copyright (c) 2014-2022, ETH Zurich
+ * All rights reserved.
+ *
+ * Please, refer to the LICENSE file in the root directory.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+
+#include <gtest/gtest.h>
+
+#include <gridtools/fn/unstructured.hpp>
+#include <gridtools/sid/allocator.hpp>
+#include <gridtools/sid/composite.hpp>
+#include <gridtools/sid/synthetic.hpp>
+
+#include <array>
+#include <cstddef>
+#include <cstdint>
+#include <gridtools/fn/sid_neighbor_table.hpp>
+
+using namespace gridtools;
+using namespace fn;
+using sid_neighbor_table::as_neighbor_table;
+
+using edge_dim_t = unstructured::dim::horizontal;
+using edge_to_cell_dim_t = struct {};
+
+TEST(sid_neighbor_table, correctness) {
+    constexpr size_t numElements = 3;
+    constexpr size_t numNeighbors = 2;
+    std::array<int32_t, numElements *numNeighbors> data = {0, 1, 10, 11, 20, 21};
+    using dim_hymap_t = hymap::keys<edge_dim_t, edge_to_cell_dim_t>;
+    auto contents = sid::synthetic()
+                        .set<sid::property::origin>(sid::host_device::simple_ptr_holder(data.data()))
+                        .set<sid::property::strides>(dim_hymap_t::make_values(numNeighbors, 1));
+    const auto table = as_neighbor_table<edge_dim_t, edge_to_cell_dim_t, numNeighbors>(contents);
+
+    auto [n00, n01] = neighbor_table_neighbors(table, 0);
+    auto [n10, n11] = neighbor_table_neighbors(table, 1);
+    auto [n20, n21] = neighbor_table_neighbors(table, 2);
+    EXPECT_EQ(n00, 0);
+    EXPECT_EQ(n01, 1);
+    EXPECT_EQ(n10, 10);
+    EXPECT_EQ(n11, 11);
+    EXPECT_EQ(n20, 20);
+    EXPECT_EQ(n21, 21);
+}