fixup

ecmwf · Oct 14, 2024 · e3bcb94 · e3bcb94
1 parent c02e4a9
commit e3bcb94
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/src/atlas/array/ArraySpec.cc b/src/atlas/array/ArraySpec.cc
@@ -47,12 +47,18 @@ ArraySpec::ArraySpec(const ArrayShape& shape, const ArrayAlignment& alignment):
     shape_.resize(rank_);
     strides_.resize(rank_);
     layout_.resize(rank_);
+    device_strides_.resize(rank_);
+    device_strides_[rank_ - 1] = 1;
     for (int j = rank_ - 1; j >= 0; --j) {
         shape_[j]   = shape[j];
         strides_[j] = allocated_size_;
         layout_[j]  = j;
         size_ *= size_t(shape_[j]);
         allocated_size_ *= size_t(aligned_shape[j]);
+        if( j < rank_ - 1) {
+            // Assume contiguous device data!
+            device_strides_[j] = strides_[j+1] * shape[j];
+        }
     }
     ATLAS_ASSERT(allocated_size_ == compute_aligned_size(size_t(shape_[0]) * size_t(strides_[0]), size_t(alignment)));
     contiguous_     = (size_ == allocated_size_);