[mlir][Vector] Use llvm::zip to avoid assertion failed.

[jacquesguan]

Summary

This patch fixes the issue https://github.com/llvm/llvm-project/issues/59455.
We could omit the un-changed dimensions in offsets and sizes, so llvm::zip_equal would fail in this case.

[aartbik]

Thanks for the fix! Can you add the reproducer from the bug as test in the vector dialect test directory?

[mehdi_amini]

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

2974

I don't quite get how do these have a different size?
The bug report has:
%261 = vector.extract_strided_slice %256 {offsets = [3], sizes = [5], strides = [1]} : vector<12x7xi1> to vector<5x7xi1>
where all the array attributes seem to have the same size?

[Mogball]

Please add a repro test as @aartbik suggested.

[jacquesguan]

Add test.

[jacquesguan]

In D139815#3988909, @aartbik wrote:

Thanks for the fix! Can you add the reproducer from the bug as test in the vector dialect test directory?

Done.

In D139815#3990995, @Mogball wrote:

Please add a repro test as @aartbik suggested.

Done.

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

2974

In this case, maskDimSizes contains 2 dimensions but sliceOffsets and sliceSizes only have one. For the last unchanged dimension, we could omit it in offsets and sizes input of extract_strided_slice.

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

populateFromInt64AttrArray(extractStridedSliceOp.getOffsets(),
                                sliceOffsets);
     SmallVector<int64_t, 4> sliceSizes;
     populateFromInt64AttrArray(extractStridedSliceOp.getSizes(), sliceSizes);
 
     // Compute slice of vector mask region.
     SmallVector<int64_t, 4> sliceMaskDimSizes;
     sliceMaskDimSizes.reserve(maskDimSizes.size());
     for (auto [maskDimSize, sliceOffset, sliceSize] :
-         llvm::zip_equal(maskDimSizes, sliceOffsets, sliceSizes)) {
+         llvm::zip(maskDimSizes, sliceOffsets, sliceSizes)) {
       int64_t sliceMaskDimSize = std::max(
           static_cast<int64_t>(0),
           std::min(sliceOffset + sliceSize, maskDimSize) - sliceOffset);
       sliceMaskDimSizes.push_back(sliceMaskDimSize);
     }
+    // Add unchanged dimensions.
+    if (sliceMaskDimSizes.size() < maskDimSizes.size())
+      for (size_t i = sliceMaskDimSizes.size(); i < maskDimSizes.size(); ++i)
+        sliceMaskDimSizes.push_back(maskDimSizes[i]);
     // If any of 'sliceMaskDimSizes' are zero, then set all to zero (masked
     // region is a conjunction of mask dim intervals).
     if (llvm::is_contained(sliceMaskDimSizes, 0))
       sliceMaskDimSizes.assign(maskDimSizes.size(), 0);
 
     // Replace 'extractStridedSliceOp' with ConstantMaskOp with sliced mask
     // region.
     rewriter.replaceOpWithNewOp<ConstantMaskOp>(

mlir/test/Dialect/Vector/canonicalize.mlir

 func.func @extract_from_broadcast(%src: vector<1x1x1xf32>) -> vector<1xf32> {
   %0 = vector.broadcast %src : vector<1x1x1xf32> to vector<1x1x32x1xf32>
 
   //  CHECK-NEXT:   %0 = vector.extract {{.*}}[0, 0] : vector<1x1x1xf32>
   //  CHECK-NEXT:   return %0 : vector<1xf32>
   %1 = vector.extract %0[0, 0, 31] : vector<1x1x32x1xf32>
   return %1: vector<1xf32>
 }
+
+// -----
+// CHECK-LABEL: func.func @extract_strided_slice_of_constant_mask
+func.func @extract_strided_slice_of_constant_mask() -> vector<5x7xi1>{
+  //  CHECK-NEXT:   %[[RES:.*]] = vector.constant_mask [5, 4] : vector<5x7xi1>
+  //  CHECK-NEXT:   return %[[RES]] : vector<5x7xi1>
+  %c4 = arith.constant 4 : index
+  %c10 = arith.constant 10 : index
+  %mask = vector.create_mask %c10, %c4 : vector<12x7xi1>
+  %res = vector.extract_strided_slice %mask {offsets = [3], sizes = [5], strides = [1]} : vector<12x7xi1> to vector<5x7xi1>
+  return %res : vector<5x7xi1>
+}