Feature: Smart STK Field Types

include/ngp_utils/SmartFieldRef.h

@@ -0,0 +1,183 @@
+// Copyright 2017 National Technology & Engineering Solutions of Sandia, LLC
+// (NTESS), National Renewable Energy Laboratory, University of Texas Austin,
+// Northwest Research Associates. Under the terms of Contract DE-NA0003525
+// with NTESS, the U.S. Government retains certain rights in this software.
+//
+// This software is released under the BSD 3-clause license. See LICENSE file
+// for more details.
+//
+
+#ifndef SMARTFIELDREF_H
+#define SMARTFIELDREF_H
+#include <Kokkos_Macros.hpp>
+#include <stk_mesh/base/Ngp.hpp>
+#include <stk_mesh/base/NgpField.hpp>
+
+namespace sierra::nalu {
+
+struct READ
+{
+};
+struct WRITE
+{
+};
+struct READ_WRITE
+{
+};
+
+struct HOST
+{
+};
+struct DEVICE
+{
+};
+
+template <typename MEMSPACE, typename SCOPE, typename T>
+class SmartFieldRef
+{
+};
+
+template <typename SCOPE, typename T>
+class SmartFieldRef<DEVICE, SCOPE, T>
+{
+public:
+  SmartFieldRef(stk::mesh::NgpField<T>& ngpField) : fieldRef_(ngpField) {}
+
+  SmartFieldRef(const SmartFieldRef& src)
+    : fieldRef_(src.fieldRef_), is_copy_constructed_(true)
+  {
+    if (is_read())
+      fieldRef_.sync_to_device();
+    else
+      fieldRef_.clear_sync_state();
+  }
+
+  // device implementations should only ever execute inside a
+  // kokkos::paralle_for and hence be captured by a lambda. Therefore we only
+  // ever need to sync copies that will have been snatched up through lambda
+  // capture.
+  ~SmartFieldRef()
+  {
+    if (is_copy_constructed_ && is_write()) {
+      fieldRef_.modify_on_device();
+    }
+  }
+
+  KOKKOS_FUNCTION
+  unsigned get_ordinal() const { return fieldRef_.get_ordinal(); }
+
+  // TODO make it so these accessors are read only for read type i.e. const
+  // correct and give clear compile or runtime error for programming mistakes
+  KOKKOS_FUNCTION
+  T& get(stk::mesh::FastMeshIndex index, int component)
+  {
+    return fieldRef_.get(index, component);
+  }
+
+  template <typename MeshIndex>
+  KOKKOS_FUNCTION T& get(MeshIndex index, int component)
+  {
+    return fieldRef_.get(index, component);
+  }
+
+  KOKKOS_FUNCTION
+  T& operator()(stk::mesh::FastMeshIndex index, int component)
+  {
+    return fieldRef_.get(index, component);
+  }
+
+  template <typename MeshIndex>
+  KOKKOS_FUNCTION T& operator()(MeshIndex index, int component)
+  {
+    return fieldRef_.operator()(index, component);
+  }
+
+private:
+  bool is_read()
+  {
+    return std::is_same<SCOPE, READ>::value ||
+           std::is_same<SCOPE, READ_WRITE>::value;
+  }
+
+  bool is_write()
+  {
+    return std::is_same<SCOPE, WRITE>::value ||
+           std::is_same<SCOPE, READ_WRITE>::value;
+  }
+
+  stk::mesh::NgpField<T>& fieldRef_;
+  const bool is_copy_constructed_{false};
+};
+
+template <typename SCOPE, typename T>
+class SmartFieldRef<HOST, SCOPE, T>
+{
+public:
+  SmartFieldRef(stk::mesh::HostField<T>& ngpField) : fieldRef_(ngpField) {}
+
+  SmartFieldRef(const SmartFieldRef& src)
+    : fieldRef_(src.fieldRef_), is_copy_constructed_(true)
+  {
+    if (is_read())
+      fieldRef_.sync_to_host();
+    else
+      fieldRef_.clear_sync_state();
+  }
+
+  // TODO is a copy construction appropriate for host instances?
+  // ideally we will still be using parallel_for's but there are lots of
+  // places where we don't yet.  It would be nice to make this work with
+  // the old model as well
+  //
+  // maybe something like SPACE=OldHost, change the ctor to take a FieldBase ptr
+  // AND then remove the copy construction check?
+  ~SmartFieldRef()
+  {
+    if (is_copy_constructed_ && is_write()) {
+      fieldRef_.modify_on_host();
+    }
+  }
+
+  unsigned get_ordinal() const { return fieldRef_.get_ordinal(); }
+
+  T& get(stk::mesh::FastMeshIndex index, int component)
+  {
+    return fieldRef_.get(index, component);
+  }
+
+  template <typename MeshIndex>
+  T& get(MeshIndex index, int component)
+  {
+    return fieldRef_.get(index, component);
+  }
+
+  T& operator()(stk::mesh::FastMeshIndex index, int component)
+  {
+    return fieldRef_.get(index, component);
+  }
+
+  template <typename MeshIndex>
+  T& operator()(MeshIndex index, int component)
+  {
+    return fieldRef_.operator()(index, component);
+  }
+
+private:
+  bool is_read()
+  {
+    return std::is_same<SCOPE, READ>::value ||
+           std::is_same<SCOPE, READ_WRITE>::value;
+  }
+
+  bool is_write()
+  {
+    return std::is_same<SCOPE, WRITE>::value ||
+           std::is_same<SCOPE, READ_WRITE>::value;
+  }
+
+  stk::mesh::HostField<T>& fieldRef_;
+  const bool is_copy_constructed_{false};
+};
+} // namespace sierra::nalu
+
+#endif

unit_tests/CMakeLists.txt

@@ -27,6 +27,7 @@ target_sources(${utest_ex_name} PRIVATE
    ${CMAKE_CURRENT_SOURCE_DIR}/UnitTestPecletFunction.C
    ${CMAKE_CURRENT_SOURCE_DIR}/UnitTestRadarPattern.C
    ${CMAKE_CURRENT_SOURCE_DIR}/UnitTestRealm.C
+   ${CMAKE_CURRENT_SOURCE_DIR}/UnitTestSmartFieldRef.C
    ${CMAKE_CURRENT_SOURCE_DIR}/UnitTestScanningLidarPattern.C
    ${CMAKE_CURRENT_SOURCE_DIR}/UnitTestScratchViews.C
    ${CMAKE_CURRENT_SOURCE_DIR}/UnitTestShmemAlignment.C

unit_tests/UnitTestSmartFieldRef.C

@@ -0,0 +1,91 @@
+// Copyright 2017 National Technology & Engineering Solutions of Sandia, LLC
+// (NTESS), National Renewable Energy Laboratory, University of Texas Austin,
+// Northwest Research Associates. Under the terms of Contract DE-NA0003525
+// with NTESS, the U.S. Government retains certain rights in this software.
+//
+// This software is released under the BSD 3-clause license. See LICENSE file
+// for more details.
+//
+
+#include <Kokkos_Macros.hpp>
+#include <gtest/gtest.h>
+#include "UnitTestUtils.h"
+#include <stk_mesh/base/GetNgpField.hpp>
+#include "ngp_utils/SmartFieldRef.h"
+
+class TestSmartFieldRef : public Hex8Mesh
+{
+protected:
+  void SetUp()
+  {
+    fill_mesh_and_initialize_test_fields();
+
+    auto* field = fieldManager->get_field_ptr<ScalarFieldType>("scalarQ");
+
+    ngpField_ = &stk::mesh::get_updated_ngp_field<double>(*field);
+
+    initSyncsHost_ = ngpField_->num_syncs_to_host();
+    initSyncsDevice_ = ngpField_->num_syncs_to_device();
+  }
+
+  stk::mesh::NgpField<double>* ngpField_;
+  int initSyncsHost_{0};
+  int initSyncsDevice_{0};
+};
+
+template <typename T>
+void
+lambda_impl(T& ptr)
+{
+  Kokkos::parallel_for(
+    1, KOKKOS_LAMBDA(int) { ptr.get_ordinal(); });
+}
+
+namespace sierra::nalu {
+TEST_F(TestSmartFieldRef, device_read_write_mod_sync_with_lambda)
+{
+  ngpField_->modify_on_host();
+
+  ASSERT_TRUE(ngpField_->need_sync_to_device());
+
+  // TODO can we get rid of the double template param some how?
+  auto sPtr = SmartFieldRef<DEVICE, READ_WRITE, double>(*ngpField_);
+  lambda_impl(sPtr);
+
+  EXPECT_FALSE(ngpField_->need_sync_to_device());
+  EXPECT_TRUE(ngpField_->need_sync_to_host());
+  EXPECT_EQ(initSyncsDevice_ + 1, ngpField_->num_syncs_to_device());
+  EXPECT_EQ(initSyncsHost_ + 0, ngpField_->num_syncs_to_host());
+}
+
+TEST_F(TestSmartFieldRef, device_write_clear_mod_with_lambda)
+{
+  ngpField_->modify_on_host();
+
+  ASSERT_TRUE(ngpField_->need_sync_to_device());
+
+  auto sPtr = SmartFieldRef<DEVICE, WRITE, double>(*ngpField_);
+  lambda_impl(sPtr);
+
+  EXPECT_FALSE(ngpField_->need_sync_to_device());
+  EXPECT_TRUE(ngpField_->need_sync_to_host());
+  EXPECT_EQ(initSyncsDevice_ + 0, ngpField_->num_syncs_to_device());
+  EXPECT_EQ(initSyncsHost_ + 0, ngpField_->num_syncs_to_host());
+}
+
+TEST_F(TestSmartFieldRef, device_read_mod_no_sync_with_lambda)
+{
+  ngpField_->modify_on_host();
+
+  ASSERT_TRUE(ngpField_->need_sync_to_device());
+
+  auto sPtr = SmartFieldRef<DEVICE, READ, double>(*ngpField_);
+  lambda_impl(sPtr);
+
+  EXPECT_FALSE(ngpField_->need_sync_to_device());
+  EXPECT_FALSE(ngpField_->need_sync_to_host());
+  EXPECT_EQ(initSyncsDevice_ + 1, ngpField_->num_syncs_to_device());
+  EXPECT_EQ(initSyncsHost_ + 0, ngpField_->num_syncs_to_host());
+}
+
+} // namespace sierra::nalu