[VectorExt] Fix LayoutIterator iteration for step != 0 (#16133)

diff --git a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/VectorExt/IR/VectorExtOps.h b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/VectorExt/IR/VectorExtOps.h
index ba9bcf0..e1d490b 100644
--- a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/VectorExt/IR/VectorExtOps.h
+++ b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/VectorExt/IR/VectorExtOps.h

@@ -31,8 +31,15 @@
     position += stride;
     return *this;
   }
+
+  bool operator==(const DimensionalIterator &other) const {
+    return position == other.position;
+  }
   bool operator!=(const DimensionalIterator &other) const {
-    return position != other.position;
+    return !(*this == other);
+  }
+  bool operator<(const DimensionalIterator &other) const {
+    return position < other.position;
   }
 
 private:

diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/VectorExt/IR/VectorExtOps.cpp b/llvm-external-projects/iree-dialects/lib/Dialect/VectorExt/IR/VectorExtOps.cpp
index 1c57d06..8894d69 100644
--- a/llvm-external-projects/iree-dialects/lib/Dialect/VectorExt/IR/VectorExtOps.cpp
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/VectorExt/IR/VectorExtOps.cpp

@@ -56,8 +56,8 @@
   for (auto &[dim, it] : state) {
     if (frozenDimensions.contains(dim))
       continue;
-    if (it != ranges[dim].end()) {
-      ++it;
+    ++it;
+    if (it < ranges[dim].end()) {
       break;
     }
     it = ranges[dim].begin();
@@ -89,7 +89,7 @@
     if (isLaneDimension(dim))
       continue;
     int64_t stride = strides.contains(dim) ? strides[dim] : 1;
-    ranges[dim] = DimensionalRange(0, shape - 1, stride);
+    ranges[dim] = DimensionalRange(0, shape, stride);
     state[dim] = ranges[dim].begin();
   }
 }