DistArray: init_{tiles,elements} and fill* are parametrized by fence …

…template parameter that controls whether operation uses local, global, or no fence (default, same as before)

src/TiledArray/array_impl.h

@@ -198,6 +198,17 @@ std::ostream& operator<<(std::ostream& os, const TileConstReference<Impl>& a) {
   return os;
 }
 
+/// Callaback used to update counter (typically, task counter)
+template <typename AtomicInt>
+struct IncrementCounter : public madness::CallbackInterface {
+  AtomicInt& counter;
+  IncrementCounter(AtomicInt& counter) : counter(counter) {}
+  void notify() override {
+    ++counter;
+    delete this;
+  }
+};
+
 }  // namespace detail
 }  // namespace TiledArray
 
@@ -770,20 +781,24 @@ class ArrayImpl : public TensorImpl<Policy>,
   /// \tparam Op The type of the functor/function
   /// \param[in] op The operation used to generate tiles
   /// \param[in] skip_set If false, will throw if any tiles are already set
+  /// \return the total number of tiles that have been (or will be) initialized
   /// \throw TiledArray::Exception if the PIMPL is not set. Strong throw
   ///                              guarantee.
   /// \throw TiledArray::Exception if a tile is already set and skip_set is
   ///                              false. Weak throw guarantee.
-  template <HostExecutor Exec = HostExecutor::Default, typename Op>
-  void init_tiles(Op&& op, bool skip_set = false) {
+  template <HostExecutor Exec = HostExecutor::Default, Fence fence = Fence::No,
+            typename Op>
+  std::int64_t init_tiles(Op&& op, bool skip_set = false) {
     // lifetime management of op depends on whether it is a lvalue ref (i.e. has
     // an external owner) or an rvalue ref
     // - if op is an lvalue ref: pass op to tasks
     // - if op is an rvalue ref pass make_shared_function(op) to tasks
     auto op_shared_handle = make_op_shared_handle(std::forward<Op>(op));
 
+    std::int64_t ntiles_initialized{0};
     auto it = this->pmap()->begin();
     const auto end = this->pmap()->end();
+    std::atomic<std::int64_t> ntask_completed{0};
     for (; it != end; ++it) {
       const auto& index = *it;
       if (!this->is_zero(index)) {
@@ -792,19 +807,39 @@ class ArrayImpl : public TensorImpl<Policy>,
           if (fut.probe()) continue;
         }
         if constexpr (Exec == HostExecutor::MADWorld) {
-          Future<value_type> tile = this->world().taskq.add(
-              [this_sptr = this->shared_from_this(),
-               index = ordinal_type(index), op_shared_handle]() -> value_type {
+          Future<value_type> tile =
+              this->world().taskq.add([this_sptr = this->shared_from_this(),
+                                       index = ordinal_type(index),
+                                       op_shared_handle, this]() -> value_type {
                 return op_shared_handle(
                     this_sptr->trange().make_tile_range(index));
               });
+          ++ntiles_initialized;
+          if constexpr (fence == Fence::Local) {
+            tile.register_callback(
+                new IncrementCounter<decltype(ntask_completed)>(
+                    ntask_completed));
+          }
           set(index, std::move(tile));
         } else {
           static_assert(Exec == HostExecutor::Thread);
           set(index, op_shared_handle(this->trange().make_tile_range(index)));
+          ++ntiles_initialized;
         }
       }
     }
+
+    if constexpr (fence == Fence::Local) {
+      if constexpr (Exec == HostExecutor::MADWorld) {
+        if (ntiles_initialized > 0)
+          this->world().await([&ntask_completed, ntiles_initialized]() {
+            return ntask_completed == ntiles_initialized;
+          });
+      }
+    } else if constexpr (fence == Fence::Global) {
+      this->world().gop.fence();
+    }
+    return ntiles_initialized;
   }
 
 };  // class ArrayImpl

src/TiledArray/dist_array.h

@@ -906,23 +906,29 @@ class DistArray : public madness::archive::ParallelSerializableObject {
   ///                              guarantee.
   /// \throw TiledArray::Exception if skip_set is false and a local tile is
   ///                              already set. Weak throw guarantee.
-  void fill_local(const element_type& value = element_type(),
-                  bool skip_set = false) {
-    init_tiles(
+  template <Fence fence = Fence::No>
+  std::int64_t fill_local(const element_type& value = element_type(),
+                          bool skip_set = false) {
+    return init_tiles<HostExecutor::Default, fence>(
         [value](const range_type& range) { return value_type(range, value); },
         skip_set);
   }
 
   /// Fill all local tiles with the specified value
 
+  /// \tparam fence If Fence::No, the operation will return early,
+  ///         before the tasks have completed
   /// \param[in] value What each local tile should be filled with.
   /// \param[in] skip_set If false, will throw if any tiles are already set
+  /// \return the total number of tiles that have been (or will be) initialized
   /// \throw TiledArray::Exception if the PIMPL is uninitialized. Strong throw
   ///                              guarantee.
   /// \throw TiledArray::Exception if skip_set is false and a local tile is
   ///                              already set. Weak throw guarantee.
-  void fill(const element_type& value = numeric_type(), bool skip_set = false) {
-    fill_local(value, skip_set);
+  template <Fence fence = Fence::No>
+  std::int64_t fill(const element_type& value = numeric_type(),
+                    bool skip_set = false) {
+    return fill_local<fence>(value, skip_set);
   }
 
   /// Fill all local tiles with random values
@@ -934,18 +940,21 @@ class DistArray : public madness::archive::ParallelSerializableObject {
   /// generate random values of type T this function will be disabled via SFINAE
   /// and attempting to use it will lead to a compile-time error.
   ///
+  /// \tparam fence If Fence::No, the operation will return early,
+  ///         before the tasks have completed
   /// \tparam T The type of random value to generate. Defaults to
   ///           element_type.
   /// \param[in] skip_set If false, will throw if any tiles are already set
+  /// \return the total number of tiles that have been (or will be) initialized
   /// \throw TiledArray::Exception if the PIMPL is not initialized. Strong
   ///                              throw guarantee.
   /// \throw TiledArray::Exception if skip_set is false and a local tile is
   ///                              already initialized. Weak throw guarantee.
   template <HostExecutor Exec = HostExecutor::Default,
-            typename T = element_type,
+            typename T = element_type, Fence fence = Fence::No,
             typename = detail::enable_if_can_make_random_t<T>>
-  void fill_random(bool skip_set = false) {
-    init_elements<Exec>(
+  std::int64_t fill_random(bool skip_set = false) {
+    return init_elements<Exec, fence>(
         [](const auto&) { return detail::MakeRandom<T>::generate_value(); });
   }
 
@@ -978,16 +987,20 @@ class DistArray : public madness::archive::ParallelSerializableObject {
   ///        return tile;
   ///     });
   /// \endcode
+  /// \tparam fence If Fence::No, the operation will return early,
+  ///         before the tasks have completed
   /// \tparam Op The type of the functor/function
   /// \param[in] op The operation used to generate tiles
   /// \param[in] skip_set If false, will throw if any tiles are already set
   /// \throw TiledArray::Exception if the PIMPL is not set. Strong throw
   ///                              guarantee.
   /// \throw TiledArray::Exception if a tile is already set and skip_set is
   ///                              false. Weak throw guarantee.
-  template <HostExecutor Exec = HostExecutor::Default, typename Op>
-  void init_tiles(Op&& op, bool skip_set = false) {
-    impl_ref().template init_tiles<Exec>(std::forward<Op>(op), skip_set);
+  template <HostExecutor Exec = HostExecutor::Default, Fence fence = Fence::No,
+            typename Op>
+  std::int64_t init_tiles(Op&& op, bool skip_set = false) {
+    return impl_ref().template init_tiles<Exec, fence>(std::forward<Op>(op),
+                                                       skip_set);
   }
 
   /// Initialize elements of local, non-zero tiles with a user provided functor
@@ -1009,15 +1022,17 @@ class DistArray : public madness::archive::ParallelSerializableObject {
   /// \tparam Op Type of the function/functor which will generate the elements.
   /// \param[in] op The operation used to generate elements
   /// \param[in] skip_set If false, will throw if any tiles are already set
+  /// \return the total number of tiles that have been (or will be) initialized
   /// \throw TiledArray::Exception if the PIMPL is not initialized. Strong
   ///                              throw guarnatee.
   /// \throw TiledArray::Exception if skip_set is false and a local, non-zero
   ///                              tile is already initialized. Weak throw
   ///                              guarantee.
-  template <HostExecutor Exec = HostExecutor::Default, typename Op>
-  void init_elements(Op&& op, bool skip_set = false) {
+  template <HostExecutor Exec = HostExecutor::Default, Fence fence = Fence::No,
+            typename Op>
+  std::int64_t init_elements(Op&& op, bool skip_set = false) {
     auto op_shared_handle = make_op_shared_handle(std::forward<Op>(op));
-    init_tiles<Exec>(
+    return init_tiles<Exec, fence>(
         [op = std::move(op_shared_handle)](
             const TiledArray::Range& range) -> value_type {
           // Initialize the tile with the given range object

src/TiledArray/fwd.h

@@ -203,6 +203,14 @@ using Array
 
 enum class HostExecutor { Thread, MADWorld, Default = MADWorld };
 
+/// fence types
+enum class Fence {
+  Global,  //!< global fence (`world.gop.fence()`)
+  Local,   //!< local fence (all local work done, equivalent to
+           //!< `world.taskq.fence() in absence of active messages)
+  No       //!< no fence
+};
+
 namespace conversions {
 
 /// user defined conversions