Fixed a race condition that can occur when there is an unused tile inside an expression. This fix also provides a minor optimization where it avoids communication for unused tiles in certain circumstances.

Author: justusc

src/TiledArray/dist_eval/array_eval.h

@@ -195,6 +195,13 @@ namespace TiledArray {
         }
       }
 
+      /// Discard a tile that is not needed
+
+      /// This function handles the cleanup for tiles that are not needed in
+      /// subsequent computation.
+      /// \param i The index of the tile
+      virtual void discard_tile(size_type i) const { }
+
     private:
 
       value_type make_tile(const typename array_type::value_type& tile, const bool consume) const {

src/TiledArray/dist_eval/binary_eval.h

@@ -105,6 +105,13 @@ namespace TiledArray {
         return TensorImpl_::get_world().gop.template recv<value_type>(source, key);
       }
 
+      /// Discard a tile that is not needed
+
+      /// This function handles the cleanup for tiles that are not needed in
+      /// subsequent computation.
+      /// \param i The index of the tile
+      virtual void discard_tile(size_type i) const { get_tile(i); }
+
     private:
 
       /// Task function for evaluating tiles
@@ -191,9 +198,9 @@ namespace TiledArray {
             } else {
               // Cleanup unused tiles
               if(! left_.is_zero(index))
-                left_.get(index);
+                left_.discard(index);
               if(! right_.is_zero(index))
-                right_.get(index);
+                right_.discard(index);
             }
           }
         }

src/TiledArray/dist_eval/contraction_eval.h

@@ -958,14 +958,22 @@ namespace TiledArray {
             const madness::Group row_group = make_row_group(k);
             const ProcessID group_root = get_row_group_root(k, row_group);
 
-            // Broadcast column k of left_.
-            for(; index < left_end_; index += left_stride_local_) {
-              if(left_.shape().is_zero(index)) continue;
-
-              // Broadcast the tile
-              const madness::DistributedID key(DistEvalImpl_::id(), index);
-              auto tile = get_tile(left_, index);
-              TensorImpl_::get_world().gop.bcast(key, tile, group_root, row_group);
+            if(row_group.size() > 1) {
+              // Broadcast column k of left_.
+              for(; index < left_end_; index += left_stride_local_) {
+                if(left_.shape().is_zero(index)) continue;
+
+                // Broadcast the tile
+                const madness::DistributedID key(DistEvalImpl_::id(), index);
+                auto tile = get_tile(left_, index);
+                TensorImpl_::get_world().gop.bcast(key, tile, group_root, row_group);
+              }
+            } else {
+              // Discard column k of left_.
+              for(; index < left_end_; index += left_stride_local_) {
+                if(left_.shape().is_zero(index)) continue;
+                left_.discard(index);
+              }
             }
 
             break;
@@ -993,14 +1001,22 @@ namespace TiledArray {
             const madness::Group col_group = make_col_group(k);
             const ProcessID group_root = get_col_group_root(k, col_group);
 
-            // Broadcast row k of right_.
-            for(; index < row_end; index += right_stride_local_) {
-              if(right_.shape().is_zero(index)) continue;
+            if(col_group.size() > 1) {
+              // Broadcast row k of right_.
+              for(; index < row_end; index += right_stride_local_) {
+                if(right_.shape().is_zero(index)) continue;
 
-              // Broadcast the tile
-              const madness::DistributedID key(DistEvalImpl_::id(), index + left_.size());
-              auto tile = get_tile(right_, index);
-              TensorImpl_::get_world().gop.bcast(key, tile, group_root, col_group);
+                // Broadcast the tile
+                const madness::DistributedID key(DistEvalImpl_::id(), index + left_.size());
+                auto tile = get_tile(right_, index);
+                TensorImpl_::get_world().gop.bcast(key, tile, group_root, col_group);
+              }
+            } else {
+              // Broadcast row k of right_.
+              for(; index < row_end; index += right_stride_local_) {
+                if(right_.shape().is_zero(index)) continue;
+                right_.discard(index);
+              }
             }
 
             break;
@@ -1761,6 +1777,14 @@ namespace TiledArray {
         return TensorImpl_::get_world().gop.template recv<value_type>(source, key);
       }
 
+
+      /// Discard a tile that is not needed
+
+      /// This function handles the cleanup for tiles that are not needed in
+      /// subsequent computation.
+      /// \param i The index of the tile
+      virtual void discard_tile(size_type i) const { get_tile(i); }
+
     private:
 
       /// Adjust iteration depth based on memory constraints

src/TiledArray/dist_eval/dist_eval.h

@@ -127,9 +127,15 @@ namespace TiledArray {
 
       /// \param i The index of the tile
       /// \return Tile at index i
-      /// \throw TiledArray::Exception When tile \c i is owned by a remote node.
       virtual Future<value_type> get_tile(size_type i) const = 0;
 
+      /// Discard a tile that is not needed
+
+      /// This function handles the cleanup for tiles that are not needed in
+      /// subsequent computation.
+      /// \param i The index of the tile
+      virtual void discard_tile(size_type i) const = 0;
+
       /// Set tensor value
 
       /// This will store \c value at ordinal index \c i . Typically, this
@@ -322,6 +328,13 @@ namespace TiledArray {
       /// \return Tile \c i
       future get(size_type i) const { return pimpl_->get_tile(i); }
 
+      /// Discard a tile that is not needed
+
+      /// This function handles the cleanup for tiles that are not needed in
+      /// subsequent computation.
+      /// \param i The index of the tile
+      virtual void discard(size_type i) const { pimpl_->discard_tile(i); }
+
       /// World object accessor
 
       /// \return A reference to the world object

src/TiledArray/dist_eval/unary_eval.h

@@ -85,6 +85,13 @@ namespace TiledArray {
         return TensorImpl_::get_world().gop.template recv<value_type>(source, key);
       }
 
+      /// Discard a tile that is not needed
+
+      /// This function handles the cleanup for tiles that are not needed in
+      /// subsequent computation.
+      /// \param i The index of the tile
+      virtual void discard_tile(size_type i) const { get_tile(i); }
+
     private:
 
       /// Input tile argument type