[Mlir-commits] [mlir] [mlir][vector] Add support for linearizing Extract, ExtractStridedSlice, Shuffle VectorOps in VectorLinearize (PR #88204)

[Andrzej Warzyński]

================
@@ -103,6 +106,236 @@ struct LinearizeVectorizable final
     return success();
   }
 
+private:
+  unsigned targetVectorBitWidth;
+};
+
+/// This pattern converts the vector.extract_strided_slice operation to a
+/// vector.shuffle operation that works on a linearized vector.
+struct LinearizeVectorExtractStridedSlice final
+    : public mlir::OpConversionPattern<mlir::vector::ExtractStridedSliceOp> {
+  using OpConversionPattern::OpConversionPattern;
+  LinearizeVectorExtractStridedSlice(
+      const TypeConverter &typeConverter, MLIRContext *context,
+      unsigned targetVectBitWidth = std::numeric_limits<unsigned>::max(),
+      PatternBenefit benefit = 1)
+      : OpConversionPattern(typeConverter, context, benefit),
+        targetVectorBitWidth(targetVectBitWidth) {}
+
+  LogicalResult
+  matchAndRewrite(vector::ExtractStridedSliceOp extractOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto dstType = getTypeConverter()->convertType(extractOp.getType());
+    auto loc = extractOp.getLoc();
+    if (!dstType)
+      return rewriter.notifyMatchFailure(loc, "cannot convert type.");
+    if (extractOp.getVector().getType().isScalable() ||
+        dstType.cast<VectorType>().isScalable())
+      return rewriter.notifyMatchFailure(loc,
+                                         "scalable vectors are not supported.");
+    if (!isLessThanTargetBitWidth(extractOp, targetVectorBitWidth))
+      return rewriter.notifyMatchFailure(
+          extractOp, "Can't flatten since targetBitWidth <= OpSize");
+
+    auto offsets = extractOp.getOffsets().getValue();
+    auto sizes = extractOp.getSizes().getValue();
+    auto strides = extractOp.getStrides().getValue();
+    if (!isConstantIntValue(strides[0], 1))
+      return rewriter.notifyMatchFailure(
+          extractOp, "Strided slice with stride != 1 is not supported.");
+    Value srcVector = adaptor.getVector();
+    // If kD offsets are specified for nd source vector (n > k), the granularity
+    // of the extraction is greater than 1. In this case last (n-k) dimensions
+    // form the extraction granularity.
+    // example :
+    //  %0 = vector.extract_strided_slice %src { offsets = [0, 0], sizes = [2,
+    //  2],
+    //     strides = [1, 1]} : vector<4x8x8xf32> to vector<2x2x8xf32>
+    // here, extraction granularity is 8.
----------------
banach-space wrote:

```suggestion
    // If kD offsets are specified for nD source vector (n > k), the granularity
    // of the extraction is greater than 1. In this case last (n-k) dimensions
    // form the extraction granularity.
    // example :
    //  vector.extract_strided_slice %s {
    //        offsets = [0, 0], sizes = [2, 2], strides = [1, 1]} : 
    //        vector<4x8x8xf32> to vector<2x2x8xf32>
    // here, extraction granularity is 8.
```

1. Please use consistent style in your comments (capitalisation).
2. Aavoid "random" line wrapping (e.g. in `sizes = [2, 2]`) -  that makes it hard to read.
3. This comment defines "granularity", but the code that follows computes `extractSliceLen`. Please make it more consistent with the implementation.

https://github.com/llvm/llvm-project/pull/88204