502 ippl does consistently consider striding in indexes

src/Expression/IpplExpressions.h

@@ -5,6 +5,8 @@
 #ifndef IPPL_EXPRESSIONS_H
 #define IPPL_EXPRESSIONS_H
 
+#include <Kokkos_Core.hpp>
+
 #include <type_traits>
 
 namespace ippl {
@@ -41,7 +43,7 @@ namespace ippl {
          * of the BareField class.
          */
         template <typename E, size_t N = sizeof(E)>
-        struct CapturedExpression {
+        struct alignas(E) CapturedExpression {
             constexpr static unsigned dim = E::dim;
 
             template <typename... Args>
@@ -50,7 +52,7 @@ namespace ippl {
                 return reinterpret_cast<const E&>(*this)(args...);
             }
 
-            char buffer[N];
+            alignas(E) char buffer[N];
         };
 
         /*!
@@ -94,6 +96,22 @@ namespace ippl {
         template <typename T>
         struct isExpression<Scalar<T>> : std::true_type {};
 
+        template <typename E, typename = void>
+        struct ExecutionSpaceOf {
+            using type = Kokkos::DefaultExecutionSpace;
+        };
+
+        template <typename E>
+        struct ExecutionSpaceOf<E, std::void_t<typename E::execution_space>> {
+            using type = typename E::execution_space;
+        };
+
+        template <typename E1, typename E2>
+        struct BinaryExecutionSpace {
+            using type = std::conditional_t<(E1::dim != 0), typename ExecutionSpaceOf<E1>::type,
+                                            typename ExecutionSpaceOf<E2>::type>;
+        };
+
     }  // namespace detail
 }  // namespace ippl
 

src/Expression/IpplOperations.h

@@ -71,6 +71,7 @@ namespace ippl {
     struct fun : public detail::Expression<fun<E>, sizeof(E)> {                \
         constexpr static unsigned dim = E::dim;                                \
         using value_type              = typename E::value_type;                \
+        using execution_space         = typename detail::ExecutionSpaceOf<E>::type; \
                                                                                \
         KOKKOS_FUNCTION                                                        \
         fun(const E& u)                                                        \
@@ -132,6 +133,7 @@ namespace ippl {
     struct fun : public detail::Expression<fun<E1, E2>, sizeof(E1) + sizeof(E2)> {             \
         constexpr static unsigned dim = std::max(E1::dim, E2::dim);                            \
         using value_type              = typename E1::value_type;                               \
+        using execution_space         = typename detail::BinaryExecutionSpace<E1, E2>::type;   \
                                                                                                \
         KOKKOS_FUNCTION                                                                        \
         fun(const E1& u, const E2& v)                                                          \

src/Field/BareField.h

@@ -20,6 +20,7 @@
 
 #include "Field/HaloCells.h"
 #include "FieldLayout/FieldLayout.h"
+#include "Index/SIndex.h"
 
 namespace ippl {
     class Index;
@@ -166,6 +167,14 @@ namespace ippl {
             return dview_m(args...);
         }
 
+        auto operator[](const Domain_t& domain);
+
+        auto operator[](const Index& index);
+
+        auto operator[](int index);
+
+        auto operator[](const SIndex<Dim>& sindex);
+
         view_type& getView() { return dview_m; }
 
         const view_type& getView() const { return dview_m; }

src/Field/BareField.hpp

@@ -14,6 +14,8 @@
 
 #include "Utility/Inform.h"
 #include "Utility/IpplInfo.h"
+#include "Field/IndexedBareField.h"
+#include "Field/SparseIndexedBareField.h"
 namespace Kokkos {
     template <typename T, unsigned Dim>
     struct reduction_identity<ippl::Vector<T, Dim>> {
@@ -139,6 +141,29 @@ namespace ippl {
         Kokkos::resize(dview_m, args...);
     }
 
+    template <typename T, unsigned Dim, class... ViewArgs>
+    auto BareField<T, Dim, ViewArgs...>::operator[](const Domain_t& domain) {
+        return IndexedBareField<BareField<T, Dim, ViewArgs...>>(*this, domain);
+    }
+
+    template <typename T, unsigned Dim, class... ViewArgs>
+    auto BareField<T, Dim, ViewArgs...>::operator[](const Index& index) {
+        Domain_t domain = getDomain();
+        domain[0] = index;
+        return IndexedBareField<BareField<T, Dim, ViewArgs...>, 1>(*this, domain);
+    }
+
+    template <typename T, unsigned Dim, class... ViewArgs>
+    auto BareField<T, Dim, ViewArgs...>::operator[](int index) {
+        static_assert(Dim == 1, "Integer overload is only available for one-dimensional fields.");
+        return operator[](Index(index, index));
+    }
+
+    template <typename T, unsigned Dim, class... ViewArgs>
+    auto BareField<T, Dim, ViewArgs...>::operator[](const SIndex<Dim>& sindex) {
+        return SparseIndexedBareField<BareField<T, Dim, ViewArgs...>>(*this, sindex);
+    }
+
     template <typename T, unsigned Dim, class... ViewArgs>
     void BareField<T, Dim, ViewArgs...>::fillHalo() {
         if (layout_m->comm.size() > 1) {

src/Field/IndexedBareField.h

@@ -0,0 +1,173 @@
+//
+// Class IndexedBareField
+//   Lightweight indexed view into a BareField.
+//
+#ifndef IPPL_INDEXED_BARE_FIELD_H
+#define IPPL_INDEXED_BARE_FIELD_H
+
+#include <Kokkos_Core.hpp>
+
+#include "Expression/IpplExpressions.h"
+#include "Index/NDIndex.h"
+#include "Utility/ParallelDispatch.h"
+
+namespace ippl {
+
+    template <typename Field, unsigned Brackets = Field::dim>
+    class IndexedBareField
+        : public detail::Expression<IndexedBareField<Field, Brackets>,
+                                    sizeof(Field) + sizeof(NDIndex<Field::dim>)> {
+    public:
+        using field_type      = Field;
+        using value_type      = typename field_type::value_type;
+        using view_type       = typename field_type::view_type;
+        using execution_space = typename field_type::execution_space;
+        using Domain_t        = NDIndex<field_type::dim>;
+
+        constexpr static unsigned dim = field_type::dim;
+        constexpr static size_t expression_size = sizeof(Field) + sizeof(Domain_t);
+
+        IndexedBareField(field_type& field, const Domain_t& domain)
+            : view_m(field.getView())
+            , owned_m(field.getOwned())
+            , domain_m(domain)
+            , nghost_m(field.getNghost()) {}
+
+        IndexedBareField(const IndexedBareField&) = default;
+
+        IndexedBareField(const view_type& view, const Domain_t& owned, const Domain_t& domain,
+                         int nghost)
+            : view_m(view)
+            , owned_m(owned)
+            , domain_m(domain)
+            , nghost_m(nghost) {}
+
+        auto operator[](const Index& index) const {
+            static_assert(Brackets < dim, "Too many Index brackets for field dimension.");
+            Domain_t domain = domain_m;
+            domain[Brackets] = index;
+            return IndexedBareField<Field, Brackets + 1>(view_m, owned_m, domain, nghost_m);
+        }
+
+        auto operator[](int index) const {
+            return operator[](Index(index, index));
+        }
+
+        template <typename... Args>
+        KOKKOS_INLINE_FUNCTION value_type operator()(Args... args) const {
+            static_assert(Brackets == dim, "IndexedBareField expression requires all dimensions.");
+            static_assert(sizeof...(Args) == dim);
+            typename RangePolicy<dim, execution_space>::index_array_type rel{args...};
+            return apply(view_m, relativeToView(rel));
+        }
+
+        IndexedBareField& operator=(value_type value) {
+            static_assert(Brackets == dim, "IndexedBareField assignment requires all dimensions.");
+            assign(ValueAssign{value});
+            return *this;
+        }
+
+        template <typename E, size_t N>
+        IndexedBareField& operator=(const detail::Expression<E, N>& expr) {
+            static_assert(Brackets == dim, "IndexedBareField assignment requires all dimensions.");
+            using capture_type = detail::CapturedExpression<E, N>;
+            capture_type expr_ = reinterpret_cast<const capture_type&>(expr);
+            assign(ExpressionAssign<E, N>{expr_});
+            return *this;
+        }
+
+        IndexedBareField& operator=(const IndexedBareField& rhs) {
+            static_assert(Brackets == dim, "IndexedBareField assignment requires all dimensions.");
+            const detail::Expression<IndexedBareField<Field, Brackets>, expression_size>& expr = rhs;
+            return operator=(expr);
+        }
+
+        const Domain_t& getDomain() const { return domain_m; }
+
+    private:
+        view_type view_m;
+        Domain_t owned_m;
+        Domain_t domain_m;
+        int nghost_m;
+
+    public:
+        struct ValueAssign {
+            value_type value;
+
+            template <typename Coords>
+            KOKKOS_INLINE_FUNCTION value_type operator()(const Coords&) const {
+                return value;
+            }
+        };
+
+        template <typename E, size_t N>
+        struct ExpressionAssign {
+            detail::CapturedExpression<E, N> expr;
+
+            template <typename Coords>
+            KOKKOS_INLINE_FUNCTION auto operator()(const Coords& rel) const {
+                return apply(expr, rel);
+            }
+        };
+
+    private:
+        template <typename Coords>
+        KOKKOS_INLINE_FUNCTION auto relativeToView(const Coords& rel) const {
+            typename RangePolicy<dim, execution_space>::index_array_type viewCoords;
+            for (unsigned d = 0; d < dim; ++d) {
+                const auto global = domain_m[d].first() + rel[d] * domain_m[d].stride();
+                viewCoords[d] = (global - owned_m[d].first()) / owned_m[d].stride() + nghost_m;
+            }
+            return viewCoords;
+        }
+
+        template <typename Coords>
+        KOKKOS_INLINE_FUNCTION auto globalToRelative(const Coords& global) const {
+            typename RangePolicy<dim, execution_space>::index_array_type rel;
+            for (unsigned d = 0; d < dim; ++d) {
+                rel[d] = (global[d] - domain_m[d].first()) / domain_m[d].stride();
+            }
+            return rel;
+        }
+
+        template <typename Coords>
+        KOKKOS_INLINE_FUNCTION auto globalToView(const Coords& global) const {
+            typename RangePolicy<dim, execution_space>::index_array_type viewCoords;
+            for (unsigned d = 0; d < dim; ++d) {
+                viewCoords[d] = (global[d] - owned_m[d].first()) / owned_m[d].stride() + nghost_m;
+            }
+            return viewCoords;
+        }
+
+    public:
+        template <typename Functor>
+        void assign(const Functor& functor) {
+            Domain_t local = domain_m.intersect(owned_m);
+            if (local.empty()) {
+                return;
+            }
+
+            using range_policy_type = RangePolicy<dim, execution_space>;
+            using index_type        = typename range_policy_type::index_type;
+            Kokkos::Array<index_type, dim> begin, end;
+            for (unsigned d = 0; d < dim; ++d) {
+                begin[d] = 0;
+                end[d]   = local[d].length();
+            }
+
+            ippl::parallel_for(
+                "IndexedBareField::operator=", createRangePolicy<dim, execution_space>(begin, end),
+                KOKKOS_CLASS_LAMBDA(const typename range_policy_type::index_array_type& args) {
+                    typename range_policy_type::index_array_type global;
+                    for (unsigned d = 0; d < dim; ++d) {
+                        global[d] = local[d].first() + args[d] * local[d].stride();
+                    }
+                    const auto rel = globalToRelative(global);
+                    apply(view_m, globalToView(global)) = functor(rel);
+                });
+        }
+    };
+
+}  // namespace ippl
+
+#endif