TransposeNode::propagate() improvements

`TransposeNode::convert_predecessor_index()` recomputed the same
information (the number of indices of the predecessor node traversed by
taking a stride along an axis) for each update. Since this node is
already extremely lightweight (its buffer is simply its predecessor's
buffer), we can easily cache this data.

Author: fastbodin

dwave/optimization/include/dwave-optimization/nodes/manipulation.hpp

@@ -537,6 +537,8 @@ class TransposeNode : public ArrayNode {
     const ssize_t ndim_;
     const std::unique_ptr<ssize_t[]> shape_;
     const std::unique_ptr<ssize_t[]> strides_;
+    /// The number of indices per stride in the predecessor array.
+    const std::vector<ssize_t> array_indices_per_stride_;
     const bool contiguous_;
     const ValuesInfo values_info_;
 

dwave/optimization/src/nodes/manipulation.cpp

@@ -250,7 +250,8 @@ void BroadcastToNode::propagate(State& state) const {
             assert(([&]() {
                        std::vector<ssize_t> multi_index =
                                unravel_index(update.index, array_ptr_->shape());
-                       multi_index.insert(multi_index.begin(), this->ndim() - array_ptr_->ndim(), 0);
+                       multi_index.insert(multi_index.begin(), this->ndim() - array_ptr_->ndim(),
+                                          0);
                        const ssize_t assert_index = ravel_multi_index(multi_index, this->shape());
                        return assert_index == index;
                    })() &&
@@ -1464,42 +1465,57 @@ ArrayNode* TransposeNode::predeccesor_check_(ArrayNode* array_ptr) const {
     return array_ptr;
 }
 
-// a TransposeNodes shape and strides are the reverse of its predecessor
+// A TransposeNodes shape and strides are the reverse of its predecessor.
 std::unique_ptr<ssize_t[]> reverse_span_helper(const std::span<const ssize_t> span,
                                                const ssize_t size) {
     std::unique_ptr<ssize_t[]> reverse_span = std::make_unique<ssize_t[]>(size);
     std::reverse_copy(span.begin(), span.end(), reverse_span.get());
     return reverse_span;
 }
 
+std::vector<ssize_t> array_indices_per_stride_helper(const std::span<const ssize_t> array_shape,
+                                                     const ssize_t ndim) {
+    std::vector<ssize_t> axis_index_strides;
+    axis_index_strides.reserve(ndim);
+
+    ssize_t indices_per_stride = 1;
+    // Traverse the array axes in forward order.
+    for (ssize_t i = 0; i < ndim; ++i) {
+        // Record the number of indices traversed when moving along the ith axis.
+        axis_index_strides.push_back(indices_per_stride);
+        // Account for indices in ith axis.
+        indices_per_stride *= array_shape[i];
+    }
+    return axis_index_strides;
+}
+
 TransposeNode::TransposeNode(ArrayNode* array_ptr)
         : array_ptr_(predeccesor_check_(array_ptr)),
           ndim_(array_ptr->ndim()),
           shape_(reverse_span_helper(array_ptr->shape(), ndim_)),
           strides_(reverse_span_helper(array_ptr->strides(), ndim_)),
+          array_indices_per_stride_(array_indices_per_stride_helper(array_ptr->shape(), ndim_)),
           contiguous_(is_contiguous(ndim_, shape_.get(), strides_.get())),
           values_info_(array_ptr) {
     add_predecessor(array_ptr);
 }
 
-// this node simply points to the predecessor buff
+// This node simply points to the predecessor buff.
 double const* TransposeNode::buff(const State& state) const { return array_ptr_->buff(state); }
 
 ssize_t TransposeNode::ndim() const { return ndim_; }
 
 std::span<const ssize_t> TransposeNode::shape(const State& state) const {
-    if (ndim_ <= 1) {  // predecessor is vector and may be dynamic
-        return array_ptr_->shape(state);
-    }
-    // predecessor is (>=2)-D array and shape is static
+    // Predecessor is vector and may be dynamic.
+    if (ndim_ <= 1) return array_ptr_->shape(state);
+    // Predecessor is (>=2)-D array and shape is static.
     return std::span<const ssize_t>(shape_.get(), ndim_);
 }
 
 std::span<const ssize_t> TransposeNode::shape() const {
-    if (ndim_ <= 1) {  // predecessor is vector and may be dynamic
-        return array_ptr_->shape();
-    }
-    // predecessor is (>=2)-D array and shape is fixed
+    // Predecessor is vector and may be dynamic.
+    if (ndim_ <= 1) return array_ptr_->shape();
+    // Predecessor is (>=2)-D array and shape is fixed.
     return std::span<const ssize_t>(shape_.get(), ndim_);
 }
 
@@ -1533,73 +1549,53 @@ class TransposeNodeDiffData : public NodeStateData {
 std::span<const Update> TransposeNode::diff(const State& state) const {
     // If the predecessor is a vector, the transpose does nothing and the diff
     // of this node is simply the diff of the predecessor node.
-    if (ndim_ <= 1) {  // predecessor is vector
-        return array_ptr_->diff(state);
-    }
+    if (ndim_ <= 1) return array_ptr_->diff(state);
     // Otherwise, we use the stored diff data.
     return data_ptr<TransposeNodeDiffData>(state)->diff;
 }
 
 ssize_t TransposeNode::size_diff(const State& state) const { return array_ptr_->size_diff(state); }
 
 void TransposeNode::initialize_state(State& state) const {
-    if (ndim_ <= 1) {
-        return Node::initialize_state(state);  // stateless
-    }
+    if (ndim_ <= 1) return Node::initialize_state(state);  // stateless
     // Construct diff data if predecessor is (>=2)-D array
     emplace_data_ptr<TransposeNodeDiffData>(state);
 }
 
 Update TransposeNode::convert_predecessor_update_(Update update) const {
-    if (ndim_ <= 1) {  // predecessor is vector
-        return update;
-    }
-
     const std::span<const ssize_t> array_shape = array_ptr_->shape();
     ssize_t transpose_flat_index = 0;
-    // when constructing a flat index of the transpose, it is helpful to know
-    // the # of indices contributed when you move along a fixed axes.
-    // `transpose_axis_index_stride` is initialized by the # of indices
-    // contributed when moving along the 0th axis of the transpose.
-    ssize_t transpose_axis_index_stride = std::accumulate(
-            array_shape.begin(), array_shape.end() - 1, 1, std::multiplies<ssize_t>());
-
-    // traverse the predecessor axes in backward (reverse) order and the
-    // transpose axes in forward order
+    assert(ndim_ > 1);
+    assert(array_indices_per_stride_.size() == array_shape.size());
+
+    // Traverse the predecessor axes in backward (reverse) order and the
+    // transpose axes in forward order.
     for (ssize_t i = ndim_ - 1; i >= 0; --i) {
-        // grab predecessor shape along the ith axis
+        // Grab predecessor shape along the ith axis.
         const ssize_t axis_shape = array_shape[i];
         assert(0 <= axis_shape &&
                "all dimensions of (>=2)-D array must be non-negative for transpose operation");
-        // determine the multidimensional index of `flat_index` along the ith
-        // axis of predecessor. Note: this is the multidimensional index along
-        // the (ndim_ - 1 - i)th axis of the transpose
+        // Determine the multidimensional index along the ith axis of
+        // predecessor. Note: this is the multidimensional index along the
+        // (ndim_ - 1 - i)th axis of the transpose.
         const ssize_t multidimensional_index = update.index % axis_shape;
-        // reassign flat_index to the correct index along the (i - 1)th axes of predecessor
+        // Weight the multidimensional index along the (ndim_ - 1 - i)th axis
+        // of the transpose by # of indices contributed by moving along the ith
+        // axis of the predecessor.
+        transpose_flat_index += multidimensional_index * array_indices_per_stride_[i];
+        // Reassign the index to the correct index along the (i - 1)th axes of
+        // predecessor. Note we are using integer division here.
         update.index /= axis_shape;
-
-        // weight the multidimensional index along the (ndim_ - 1 - i)th axis
-        // of the transpose by # of indices contributed by moving along axis
-        transpose_flat_index += multidimensional_index * transpose_axis_index_stride;
-
-        // recall we are traversing the tranpose axes in forward order.
-        // the # of indices contributed by moving along the (ndim - 2 - i)th
-        // axis is the same as (the # of indices contributed by moving along the
-        // (ndim_ - 1 - i)th axis) / shape(ndim_ - i - 1)
-        transpose_axis_index_stride /= array_shape[ndim_ - i - 1];
     }
 
     update.index = transpose_flat_index;
-
     return update;
 }
 
 void TransposeNode::propagate(State& state) const {
     const std::span<const Update> array_diff = array_ptr_->diff(state);
-
-    if (array_diff.empty() || ndim_ <= 1) {
-        return;  // Nothing to do or predecessor is vector (transpose of vector is vector)
-    }
+    // Nothing to do or predecessor is vector (transpose of vector is vector).
+    if (array_diff.empty() || ndim_ <= 1) return;
 
     // Predecessor is a non-dynamic (>=2)-D array.
     std::vector<Update>& transpose_diff = data_ptr<TransposeNodeDiffData>(state)->diff;
@@ -1608,34 +1604,26 @@ void TransposeNode::propagate(State& state) const {
 
     for (const Update& u : array_diff) {
         assert(([&]() {
-            // make a copy of the update
-            Update u_copy = u;
-            // convert flat index of predecessor update to multidimensional indices
-            std::vector<ssize_t> multi_index = unravel_index(u_copy.index, array_ptr_->shape());
-            // reverse multidimensional indices to obtain the multidimensional
-            // transpose indices
+            // Convert flat index of predecessor update to multidimensional indices.
+            std::vector<ssize_t> multi_index = unravel_index(u.index, array_ptr_->shape());
+            // Reverse indices to obtain the transpose indices.
             std::reverse(multi_index.begin(), multi_index.end());
-            // convert multidimensional transpose indices to transpose flat index
-            // and check conversion
-            return ravel_multi_index(multi_index, this->shape()) ==
-                   convert_predecessor_update_(u_copy).index;
+            // Convert to transpose flat index and check conversion.
+            return ravel_multi_index(multi_index, shape()) == convert_predecessor_update_(u).index;
         })());
-        // Make a copy of the update and convert the index to the respective
-        // transpose index
+        // Copy update and convert predecessor index to transpose index.
         transpose_diff.emplace_back(convert_predecessor_update_(u));
     }
 }
 
 void TransposeNode::commit(State& state) const {
-    if (ndim_ > 1) {
-        data_ptr<TransposeNodeDiffData>(state)->commit();
-    }  // otherwise, stateless
+    if (ndim_ > 1) data_ptr<TransposeNodeDiffData>(state)->commit();
+    // otherwise, stateless
 };
 
 void TransposeNode::revert(State& state) const {
-    if (ndim_ > 1) {
-        data_ptr<TransposeNodeDiffData>(state)->revert();
-    }  // otherwise, stateless
+    if (ndim_ > 1) data_ptr<TransposeNodeDiffData>(state)->revert();
+    // otherwise, stateless
 }
 
 }  // namespace dwave::optimization