Make shape cache invalidation more conservative. (#1670)

Repeated calls to nvte_shape should not invalidate
previous data pointers.

It would be possible to avoid unnecessary comparisons
by duplicating some of the logic from shape() so that
the cache is only relevant when columnwise shapes are
involved. Whether this code duplication is preferable
to the comparisons that arise from by value semantics
of reusing shape is a judgment call.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

transformer_engine/common/common.h

@@ -195,7 +195,18 @@ struct Tensor {
   }
 
   const std::vector<size_t> &rowwise_shape_ref() const {
-    rowwise_shape_cache = shape();
+    auto shape_queried = shape();
+    // This method is primarily designed for nvte_shape.
+    // An unfortunate consequence of unconditionally assigning
+    // values to rowwise_shape_cache without a check is that
+    // repeated calls to rowwise_shape_ref are likely to
+    // invalidate the data pointers from previous calls.
+    // If the shape has changed, then invalidating is necessary
+    // in at least some cases, but we want to keep the data
+    // valid otherwise.
+    if (rowwise_shape_cache != shape_queried) {
+      rowwise_shape_cache = std::move(shape_queried);
+    }
     return rowwise_shape_cache;
   }