fix reshapeStrides bug (#2553)

Signed-off-by: Soren Lassen <[email protected]>

src/Dialect/ONNX/ElementsAttr/Strides.cpp

@@ -112,18 +112,22 @@ std::optional<SmallVector<int64_t, 4>> reshapeStrides(ArrayRef<int64_t> shape,
         ++a1;
       }
     } while (a1 < rank1 && last == shape[a1] * strides[a1]);
-    assert(total == n * strides[a1 - 1]);
+    assert(total == n * last);
     // Add contiguous strides for axes in reshapedShape with dimSizes product n.
     int64_t n2 = 1;
     while (a2 < rank2 && n2 * reshapedShape[a2] <= n) {
-      n2 *= reshapedShape[a2];
-      total /= reshapedShape[a2];
-      reshapedStrides.push_back(total);
+      if (reshapedShape[a2] == 1) {
+        reshapedStrides.push_back(0);
+      } else {
+        n2 *= reshapedShape[a2];
+        total /= reshapedShape[a2];
+        reshapedStrides.push_back(total);
+      }
       ++a2;
     }
     if (n2 < n)
       return std::nullopt;
-    assert(strides[a1 - 1] == reshapedStrides[a2 - 1]);
+    assert(last == total);
   } while (a1 < rank1);
   assert(a2 == rank2);
   assert(a2 == reshapedStrides.size());

test/unit/Strides/TestStrides.cpp

@@ -35,6 +35,21 @@ class Test {
     return 0;
   }
 
+  // This example triggered a bug.
+  int test_reshapeStrides_unsqueeze_last() {
+    std::cout << "test_reshapeStrides_unsqueeze_last:" << std::endl;
+
+    SmallVector<int64_t, 4> shape{1, 2, 124, 1};
+    SmallVector<int64_t, 4> strides{0, 0, 1, 0};
+
+    SmallVector<int64_t, 4> reshapedShape{1, 2, 124, 1, 1};
+    SmallVector<int64_t, 4> expectedReshapedStrides{0, 0, 1, 0, 0};
+    auto reshapedStrides = reshapeStrides(shape, strides, reshapedShape);
+    assert(reshapedStrides == expectedReshapedStrides);
+
+    return 0;
+  }
+
   int test_reshapeStrides_failure() {
     std::cout << "test_reshapeStrides_failure:" << std::endl;
 
@@ -60,6 +75,7 @@ int main(int argc, char *argv[]) {
   Test test;
   int failures = 0;
   failures += test.test_reshapeStrides_success();
+  failures += test.test_reshapeStrides_unsqueeze_last();
   failures += test.test_reshapeStrides_failure();
   if (failures != 0) {
     std::cerr << failures << " test failures\n";