[mlir][vector] Add support for linearizing Extract, ExtractStridedSlice, Shuffle VectorOps in VectorLinearize #88204

charithaintc · 2024-04-09T22:09:34Z

This PR adds support for converting vector.extract_strided_slice and vector.extract operations to equivalent vector.shuffle operations that operates on linearized (1-D) vectors. vector.shuffle operations operating on n-D (n > 1) are also converted to equivalent shuffle operations working on linearized vectors.

…orOps

github-actions · 2024-04-09T22:09:52Z

Thank you for submitting a Pull Request (PR) to the LLVM Project!

This PR will be automatically labeled and the relevant teams will be
notified.

If you wish to, you can add reviewers by using the "Reviewers" section on this page.

If this is not working for you, it is probably because you do not have write
permissions for the repository. In which case you can instead tag reviewers by
name in a comment by using @ followed by their GitHub username.

If you have received no comments on your PR for a week, you can request a review
by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate
is once a week. Please remember that you are asking for valuable time from other developers.

If you have further questions, they may be answered by the LLVM GitHub User Guide.

You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums.

llvmbot · 2024-04-09T22:40:07Z

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-vector

Author: Charitha Saumya (charithaintc)

Changes

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

2 Files Affected:

(modified) mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp (+241-1)
(modified) mlir/test/Dialect/Vector/linearize.mlir (+92)

diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
index b59e9062e5a08e..e5157abd245b5d 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
@@ -15,7 +15,9 @@
 #include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
+#include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/DialectConversion.h"
+#include <numeric>
 
 using namespace mlir;
 
@@ -103,6 +105,234 @@ struct LinearizeVectorizable final
     return success();
   }
 
+private:
+  unsigned targetVectorBitWidth;
+};
+
+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.
+    int64_t extractSliceLen = 1;
+    auto n = extractOp.getSourceVectorType().getRank();
+    auto k = (int64_t)offsets.size();
+    if (n > k) {
+      for (unsigned i = 0; i < n - k; i++) {
+        extractSliceLen *= extractOp.getSourceVectorType().getShape()[i + k];
+      }
+    }
+
+    // get total number of extracted slices
+    int64_t nExtractedSlices = 1;
+    for (auto size : sizes) {
+      nExtractedSlices *= size.cast<IntegerAttr>().getInt();
+    }
+
+    // compute the strides of the source vector considering first k dimensions
+    llvm::SmallVector<int64_t, 4> sourceStrides(k, extractSliceLen);
+    for (int i = k - 2; i >= 0; --i) {
+      sourceStrides[i] = sourceStrides[i + 1] *
+                         extractOp.getSourceVectorType().getShape()[i + 1];
+    }
+    // final shuffle indices has nExtractedElems * extractSliceLen elements
+    llvm::SmallVector<int64_t, 4> indices(nExtractedSlices * extractSliceLen);
+    // compute the strides of the extracted kD vector
+    llvm::SmallVector<int64_t, 4> extractedStrides(k, 1);
+    // compute extractedStrides
+    for (int i = k - 2; i >= 0; --i) {
+      extractedStrides[i] =
+          extractedStrides[i + 1] * sizes[i + 1].cast<IntegerAttr>().getInt();
+    }
+    // iterate over all extracted slices from 0 to nExtractedElems-1
+    // and compute the multi-dimensional index and the corresponding linearized
+    // index within the source vector
+    for (int64_t i = 0; i < nExtractedSlices; ++i) {
+      int64_t index = i;
+      // compute the corresponding multi-dimensional index
+      llvm::SmallVector<int64_t, 4> multiDimIndex(k, 0);
+      for (int64_t j = 0; j < k; ++j) {
+        multiDimIndex[j] = (index / extractedStrides[j]);
+        index -= multiDimIndex[j] * extractedStrides[j];
+      }
+      // compute the corresponding linearized index in the source vector
+      // i.e. shift the multiDimIndex by the offsets
+      int64_t linearizedIndex = 0;
+      for (int64_t j = 0; j < k; ++j) {
+        linearizedIndex +=
+            (offsets[j].cast<IntegerAttr>().getInt() + multiDimIndex[j]) *
+            sourceStrides[j];
+      }
+      // fill the indices array form linearizedIndex to linearizedIndex +
+      // sliceLen
+      for (int64_t j = 0; j < extractSliceLen; ++j) {
+        indices[i * extractSliceLen + j] = linearizedIndex + j;
+      }
+    }
+    // perform a shuffle to extract the kD vector
+    rewriter.replaceOpWithNewOp<vector::ShuffleOp>(
+        extractOp, dstType, srcVector, srcVector,
+        rewriter.getI64ArrayAttr(indices));
+
+    return success();
+  }
+
+private:
+  unsigned targetVectorBitWidth;
+};
+
+struct LinearizeVectorShffle final
+    : public OpConversionPattern<vector::ShuffleOp> {
+  using OpConversionPattern::OpConversionPattern;
+  LinearizeVectorShffle(
+      const TypeConverter &typeConverter, MLIRContext *context,
+      unsigned targetVectBitWidth = std::numeric_limits<unsigned>::max(),
+      PatternBenefit benefit = 1)
+      : OpConversionPattern(typeConverter, context, benefit),
+        targetVectorBitWidth(targetVectBitWidth) {}
+
+  LogicalResult
+  matchAndRewrite(vector::ShuffleOp shuffleOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto dstType = getTypeConverter()->convertType(shuffleOp.getType());
+    auto loc = shuffleOp.getLoc();
+    if (!dstType)
+      return rewriter.notifyMatchFailure(loc, "cannot convert type.");
+
+    if (shuffleOp.getV1VectorType().isScalable() ||
+        shuffleOp.getV2VectorType().isScalable() ||
+        dstType.cast<VectorType>().isScalable())
+      return rewriter.notifyMatchFailure(loc,
+                                         "scalable vectors are not supported.");
+    if (!isLessThanTargetBitWidth(shuffleOp, targetVectorBitWidth))
+      return rewriter.notifyMatchFailure(
+          shuffleOp, "Can't flatten since targetBitWidth <= OpSize");
+
+    auto vec1 = adaptor.getV1();
+    auto vec2 = adaptor.getV2();
+
+    int shuffleSliceLen = 1;
+    int rank = shuffleOp.getV1().getType().getRank();
+
+    // if rank > 1, we need to do the shuffle in the granularity of slices
+    // instead of scalars. Size of the slice is equal to the rank-1 innermost
+    // dims. Mask of the shuffle op specifies which slice to take from the
+    // outermost dim.
+    if (rank > 1) {
+      auto shape = shuffleOp.getV1().getType().getShape();
+      for (unsigned i = 1; i < shape.size(); i++) {
+        shuffleSliceLen *= shape[i];
+      }
+    }
+
+    auto mask = shuffleOp.getMask();
+    auto totalSize = mask.size() * shuffleSliceLen;
+
+    llvm::SmallVector<int64_t, 2> indices(totalSize);
+    for (auto [i, value] :
+         llvm::enumerate(mask.getAsValueRange<IntegerAttr>())) {
+
+      int64_t v = value.getZExtValue();
+      std::iota(indices.begin() + shuffleSliceLen * i,
+                indices.begin() + shuffleSliceLen * (i + 1),
+                shuffleSliceLen * v);
+    }
+
+    rewriter.replaceOpWithNewOp<vector::ShuffleOp>(
+        shuffleOp, dstType, vec1, vec2, rewriter.getI64ArrayAttr(indices));
+
+    return success();
+  }
+
+private:
+  unsigned targetVectorBitWidth;
+};
+
+struct LinearizeVectorExtract final
+    : public OpConversionPattern<vector::ExtractOp> {
+  using OpConversionPattern::OpConversionPattern;
+  LinearizeVectorExtract(
+      const TypeConverter &typeConverter, MLIRContext *context,
+      unsigned targetVectBitWidth = std::numeric_limits<unsigned>::max(),
+      PatternBenefit benefit = 1)
+      : OpConversionPattern(typeConverter, context, benefit),
+        targetVectorBitWidth(targetVectBitWidth) {}
+  LogicalResult
+  matchAndRewrite(vector::ExtractOp extractOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto dstTy = getTypeConverter()->convertType(extractOp.getType());
+    if (!dstTy)
+      return rewriter.notifyMatchFailure(extractOp, "cannot convert type.");
+
+    if (extractOp.getVector().getType().isScalable() ||
+        dstTy.cast<VectorType>().isScalable())
+      return rewriter.notifyMatchFailure(extractOp,
+                                         "scalable vectors are not supported.");
+    if (!isLessThanTargetBitWidth(extractOp, targetVectorBitWidth))
+      return rewriter.notifyMatchFailure(
+          extractOp, "Can't flatten since targetBitWidth <= OpSize");
+
+    // dynamic position is not supported
+    if (extractOp.hasDynamicPosition())
+      return rewriter.notifyMatchFailure(extractOp,
+                                         "dynamic position is not supported.");
+
+    auto shape = extractOp.getVector().getType().getShape();
+    auto size = extractOp.getVector().getType().getNumElements();
+
+    // compute linearized offset
+    int64_t linearizedOffset = 0;
+    auto offsets = extractOp.getStaticPosition();
+    for (auto [i, off] : llvm::enumerate(offsets)) {
+      size /= shape[i];
+      linearizedOffset += offsets[i] * size;
+    }
+
+    llvm::SmallVector<int64_t, 2> indices(size);
+    std::iota(indices.begin(), indices.end(), linearizedOffset);
+    rewriter.replaceOpWithNewOp<vector::ShuffleOp>(
+        extractOp, dstTy, adaptor.getVector(), adaptor.getVector(),
+        rewriter.getI64ArrayAttr(indices));
+
+    return success();
+  }
+
 private:
   unsigned targetVectorBitWidth;
 };
@@ -139,9 +369,19 @@ void mlir::vector::populateVectorLinearizeTypeConversionsAndLegality(
                       ? typeConverter.isLegal(op)
                       : true);
         }
+        if (isa<vector::ShuffleOp>(op)) {
+          return (isLessThanTargetBitWidth(op, targetBitWidth)
+                      ? (typeConverter.isLegal(op) &&
+                         op->getResult(0)
+                                 .getType()
+                                 .cast<mlir::VectorType>()
+                                 .getRank() == 1)
+                      : true);
+        }
         return std::nullopt;
       });
 
-  patterns.add<LinearizeConstant, LinearizeVectorizable>(
+  patterns.add<LinearizeConstant, LinearizeVectorizable, LinearizeVectorShffle,
+               LinearizeVectorExtract, LinearizeVectorExtractStridedSlice>(
       typeConverter, patterns.getContext(), targetBitWidth);
 }
diff --git a/mlir/test/Dialect/Vector/linearize.mlir b/mlir/test/Dialect/Vector/linearize.mlir
index 22be78cd682057..67f0f667a6b205 100644
--- a/mlir/test/Dialect/Vector/linearize.mlir
+++ b/mlir/test/Dialect/Vector/linearize.mlir
@@ -153,3 +153,95 @@ func.func @test_0d_vector() -> vector<f32> {
   // ALL: return %[[CST]]
   return %0 : vector<f32>
 }
+
+// -----
+// ALL-LABEL: test_extract_strided_slice_1
+// ALL-SAME: (%[[ORIG_ARG:.*]]: vector<4x8xf32>) -> vector<2x2xf32> {
+func.func @test_extract_strided_slice_1(%arg0 : vector<4x8xf32>) -> vector<2x2xf32> {
+  // DEFAULT: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<4x8xf32> to vector<32xf32>
+  // DEFAULT: %[[SHUFFLE:.*]] = vector.shuffle %[[ARG]], %[[ARG]]
+  // DEFAULT: [4, 5, 12, 13] : vector<32xf32>, vector<32xf32>
+  // DEFAULT: %[[RES:.*]] = vector.shape_cast %[[SHUFFLE]] : vector<4xf32> to vector<2x2xf32>
+  // DEFAULT: return %[[RES]] : vector<2x2xf32
+
+  // BW-128: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<4x8xf32> to vector<32xf32>
+  // BW-128: %[[SHUFFLE:.*]] = vector.shuffle %[[ARG]], %[[ARG]]
+  // BW-128: [4, 5, 12, 13] : vector<32xf32>, vector<32xf32>
+  // BW-128: %[[RES:.*]] = vector.shape_cast %[[SHUFFLE]] : vector<4xf32> to vector<2x2xf32>
+  // BW-128: return %[[RES]] : vector<2x2xf32>
+
+  // BW-0: %[[RES:.*]] = vector.extract_strided_slice %[[ARG:.*]] {offsets = [0, 4], sizes = [2, 2], strides = [1, 1]} : vector<4x8xf32> to vector<2x2xf32>
+  // BW-0: return %[[RES]] : vector<2x2xf32>
+  %0 = vector.extract_strided_slice %arg0 { sizes = [2, 2], strides = [1, 1], offsets = [0, 4]}
+     : vector<4x8xf32> to vector<2x2xf32>
+  return %0 : vector<2x2xf32>
+}
+
+// -----
+// ALL-LABEL: test_extract_strided_slice_2
+// ALL-SAME: (%[[ORIG_ARG:.*]]: vector<2x8x2xf32>) -> vector<1x4x2xf32> {
+func.func @test_extract_strided_slice_2(%arg0 : vector<2x8x2xf32>) -> vector<1x4x2xf32> {
+  // DEFAULT: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<2x8x2xf32> to vector<32xf32>
+  // DEFAULT: %[[SHUFFLE:.*]] = vector.shuffle %[[ARG]], %[[ARG]]
+  // DEFAULT: [20, 21, 22, 23, 24, 25, 26, 27] : vector<32xf32>, vector<32xf32>
+  // DEFAULT: %[[RES:.*]] = vector.shape_cast %[[SHUFFLE]] : vector<8xf32> to vector<1x4x2xf32>
+  // DEFAULT: return %[[RES]] : vector<1x4x2xf32>
+
+  // BW-128: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<2x8x2xf32> to vector<32xf32>
+  // BW-128: %[[SHUFFLE:.*]] = vector.shuffle %[[ARG]], %[[ARG]]
+  // BW-128: [20, 21, 22, 23, 24, 25, 26, 27] : vector<32xf32>, vector<32xf32>
+  // BW-128: %[[RES:.*]] = vector.shape_cast %[[SHUFFLE]] : vector<8xf32> to vector<1x4x2xf32>
+  // BW-128: return %[[RES]] : vector<1x4x2xf32>
+
+  // BW-0: %[[RES:.*]] = vector.extract_strided_slice %[[ORIG_ARG]] {offsets = [1, 2], sizes = [1, 4], strides = [1, 1]} : vector<2x8x2xf32> to vector<1x4x2xf32>
+  // BW-0: return %[[RES]] : vector<1x4x2xf32>
+  %0 = vector.extract_strided_slice %arg0 { offsets = [1, 2], strides = [1, 1], sizes = [1, 4] }
+    : vector<2x8x2xf32> to vector<1x4x2xf32>
+  return %0 : vector<1x4x2xf32>
+}
+
+// -----
+// ALL-LABEL: test_vector_shuffle
+// ALL-SAME: (%[[ORIG_ARG0:.*]]: vector<4x2xf32>, %[[ORIG_ARG1:.*]]: vector<4x2xf32>) -> vector<8x2xf32> {
+func.func @test_vector_shuffle(%arg0: vector<4x2xf32>, %arg1: vector<4x2xf32>) -> vector<8x2xf32> {
+  // DEFAULT: %[[ARG0:.*]] = vector.shape_cast %[[ORIG_ARG0]] : vector<4x2xf32> to vector<8xf32>
+  // DEFAULT: %[[ARG1:.*]] = vector.shape_cast %[[ORIG_ARG1]] : vector<4x2xf32> to vector<8xf32>
+  // DEFAULT: %[[SHUFFLE:.*]] = vector.shuffle %[[ARG0]], %[[ARG1]]
+  // DEFAULT: [0, 1, 8, 9, 2, 3, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15] : vector<8xf32>, vector<8xf32>
+  // DEFAULT: %[[RES:.*]] = vector.shape_cast %[[SHUFFLE]] : vector<16xf32> to vector<8x2xf32>
+  // DEFAULT: return %[[RES]] : vector<8x2xf32>
+
+  // BW-128: %[[ARG0:.*]] = vector.shape_cast %[[ORIG_ARG0]] : vector<4x2xf32> to vector<8xf32>
+  // BW-128: %[[ARG1:.*]] = vector.shape_cast %[[ORIG_ARG1]] : vector<4x2xf32> to vector<8xf32>
+  // BW-128: %[[SHUFFLE:.*]] = vector.shuffle %[[ARG0]], %[[ARG1]]
+  // BW-128: [0, 1, 8, 9, 2, 3, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15] : vector<8xf32>, vector<8xf32>
+  // BW-128: %[[RES:.*]] = vector.shape_cast %[[SHUFFLE]] : vector<16xf32> to vector<8x2xf32>
+  // BW-128: return %[[RES]] : vector<8x2xf32>
+
+  // BW-0: %[[RES:.*]] = vector.shuffle %[[ORIG_ARG0]], %[[ORIG_ARG1]] [0, 4, 1, 5, 2, 6, 3, 7] : vector<4x2xf32>, vector<4x2xf32>
+  // BW-0: return %[[RES]] : vector<8x2xf32>
+  %0 = vector.shuffle %arg0, %arg1 [0, 4, 1, 5, 2, 6, 3, 7] : vector<4x2xf32>, vector<4x2xf32>
+  return %0 : vector<8x2xf32>
+}
+
+// -----
+// ALL-LABEL: test_vector_extract
+// ALL-SAME: (%[[ORIG_ARG:.*]]: vector<2x8x2xf32>) -> vector<8x2xf32> {
+func.func @test_vector_extract(%arg0: vector<2x8x2xf32>) -> vector<8x2xf32> {
+  // DEFAULT: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<2x8x2xf32> to vector<32xf32>
+  // DEFAULT: %[[SHUFFLE:.*]] = vector.shuffle %[[ARG]], %[[ARG]]
+  // DEFAULT: [16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31] : vector<32xf32>, vector<32xf32>
+  // DEFAULT: %[[RES:.*]] = vector.shape_cast %[[SHUFFLE]] : vector<16xf32> to vector<8x2xf32>
+  // DEFAULT: return %[[RES]] : vector<8x2xf32>
+
+  // BW-128: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<2x8x2xf32> to vector<32xf32>
+  // BW-128: %[[SHUFFLE:.*]] = vector.shuffle %[[ARG]], %[[ARG]]
+  // BW-128: [16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31] : vector<32xf32>, vector<32xf32>
+  // BW-128: %[[RES:.*]] = vector.shape_cast %[[SHUFFLE]] : vector<16xf32> to vector<8x2xf32>
+  // BW-128: return %[[RES]] : vector<8x2xf32>
+
+  // BW-0: %[[RES:.*]] = vector.extract %[[ORIG_ARG]][1] : vector<8x2xf32> from vector<2x8x2xf32>
+  // BW-0: return %[[RES]] : vector<8x2xf32>
+  %0 = vector.extract %arg0[1]: vector<8x2xf32> from vector<2x8x2xf32>
+  return %0 : vector<8x2xf32>
+}

banach-space

Thanks for the contribution!

I've only been able to scan this very quickly - will take another look later, but mostly looks good. In the meantime, please make sure that you follow the guidelines. In particular:

https://mlir.llvm.org/getting_started/Contributing/#commit-messages

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

mlir/test/Dialect/Vector/linearize.mlir

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

charithaintc · 2024-04-12T19:39:58Z

Thanks for the contribution!

I've only been able to scan this very quickly - will take another look later, but mostly looks good. In the meantime, please make sure that you follow the guidelines. In particular:

https://mlir.llvm.org/getting_started/Contributing/#commit-messages

Hi @banach-space,

Thank you very much for the review. I have addressed the comments. Please let me know if there are any additional concerns.

dcaballe

I did a first pass. LG! A few comments

dcaballe · 2024-04-12T21:45:02Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

        return std::nullopt;
      });

-  patterns.add<LinearizeConstant, LinearizeVectorizable>(
+  patterns.add<LinearizeConstant, LinearizeVectorizable, LinearizeVectorShffle,
+               LinearizeVectorExtract, LinearizeVectorExtractStridedSlice>(


Could we move this to a different populate method? Flattening an elementwise op is potentially a no-op but generating shuffles for extract ops could have a noticeable impact in performance

I added a new populate method called populateVectorLinearizeToShuffleRewritePatterns. If you have a better name suggestion for this please let me know.

dcaballe · 2024-04-12T21:45:24Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

+
+/// This pattern converts the vector.shuffle operation that works on nD (n > 1)
+/// vectors to a vector.shuffle operation that works on linearized vectors.
+struct LinearizeVectorShffle final


Shffle -> Shuffle?

dcaballe · 2024-04-12T21:46:57Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

+      unsigned targetVectBitWidth = std::numeric_limits<unsigned>::max(),
+      PatternBenefit benefit = 1)
+      : OpConversionPattern(typeConverter, context, benefit),
+        targetVectorBitWidth(targetVectBitWidth) {}


Thanks for adding support for the target vector bitwidth!

dcaballe · 2024-04-12T21:48:18Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

+    // here, extraction granularity is 8.
+    int64_t extractSliceLen = 1;
+    auto n = extractOp.getSourceVectorType().getRank();
+    auto k = (int64_t)offsets.size();


nit: we should spell out auto when the type is not literally redundant on the RHS.

dcaballe · 2024-04-12T21:50:39Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

+        extractSliceLen *= extractOp.getSourceVectorType().getShape()[i + k];
+      }
+    }
+    // get total number of extracted slices


nit: we should use capital letters and periods in comments per coding guidelines.

fixed all comments.

charithaintc · 2024-04-12T23:11:31Z

I did a first pass. LG! A few comments

Hi @dcaballe,

Thanks for the review. I have addressed all the comments.

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

mlir/test/Dialect/Vector/linearize.mlir

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

banach-space · 2024-04-15T15:10:35Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

+  LogicalResult
+  matchAndRewrite(vector::ExtractStridedSliceOp extractOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto dstType = getTypeConverter()->convertType(extractOp.getType());


Why is this needed?

yes. this is not needed. I removed the dstType checks.

nit: spell out auto

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

banach-space · 2024-04-15T15:33:12Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

+    // 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.


Suggested change

// 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.

// 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.

Please use consistent style in your comments (capitalisation).

Aavoid "random" line wrapping (e.g. in sizes = [2, 2]) - that makes it hard to read.

This comment defines "granularity", but the code that follows computes extractSliceLen. Please make it more consistent with the implementation.

renamed to extractGranularitySize. I think the work granularity captures the meaning well. so I like to keep it.

banach-space · 2024-04-15T15:33:51Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

+    //  2],
+    //     strides = [1, 1]} : vector<4x8x8xf32> to vector<2x2x8xf32>
+    // here, extraction granularity is 8.
+    int64_t extractSliceLen = 1;


This represents "flattened slice len", right? Why not flattenedSliceLen? Naming is hard 🤷🏻

addressed above. basically the extraction can happen only at first k dimensions of an n-D vectors. the remaining dimensions are extracted in full. so the last n-k dimensions form the extraction granularity. this is the meaning I wanted to capture in the naming.

banach-space · 2024-04-15T15:34:33Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

+    int64_t extractSliceLen = 1;
+    auto n = extractOp.getSourceVectorType().getRank();
+    int64_t k = (int64_t)offsets.size();
+    if (n > k) {


Is this if is needed here? k is guaranteed to be <= n, right?

addressed above. n > k is true for the following example.

%0 = vector.extract_strided_slice %arg0 { offsets = [1, 2], strides = [1, 1], sizes = [1, 4] } : vector<2x8x2xf32> to vector<1x4x2xf32>

In this case, only first 2 dims are extracted using the offsets, sizes etc. last dim is extracted in full.
https://mlir.llvm.org/docs/Dialects/Vector/#vectorextract_strided_slice-vectorextractstridedsliceop

Right, so there are 2 cases:

k < n

k = n

right? I think that you can get rid of if by doing something like:

int64_t outRank = (int64_t)offsets.size(); int64_t k = outRank; while (k < n) { extractGranularitySize *= extractOp.getSourceVectorType().getShape()[outRank + (n-k)]; ++k; }

This way there's only one case. Also, I'd avoid variable names like k and n for anything other than loop iterators. Perhaps numKDims and numMDims? Or inputRank/outputRank? Or kD/nD? The last suggestion would be consistent with comments and already a nice improvement :)

banach-space · 2024-04-15T15:39:45Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

+    for (auto size : sizes) {
+      nExtractedSlices *= size.cast<IntegerAttr>().getInt();
+    }


Could you try using computeSuffixProduct instead?

tried this. but this require constructing an ArrayRef<int64_t>. So I decided to use llvm::for_each.

charithaintc · 2024-04-16T21:02:46Z

Hi, @banach-space, @dcaballe, @Hardcode84. I think I have addressed all the comments. Can you please review the modifications and/or approve.

Thanks!

github-actions · 2024-04-16T21:13:13Z

✅ With the latest revision this PR passed the C/C++ code formatter.

Hardcode84 · 2024-04-16T21:57:03Z

LGTM, but wait for other reviewers

dcaballe

LGTM (module some nits that you can address before landing)

dcaballe · 2024-04-17T06:23:17Z

mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h

@@ -389,6 +389,12 @@ void populateVectorLinearizeTypeConversionsAndLegality(
    TypeConverter &typeConverter, RewritePatternSet &patterns,
    ConversionTarget &target, unsigned targetBitWidth);

+/// Populates patterns for linearizing ND (N >= 2) vector operations to 1D
+/// vector shuffle operations.
+void populateVectorLinearizeToShuffleRewritePatterns(


Extract ops are "shuffle-like" ops so perhaps we can call this populateVectorLinearizeShuffleLikeOpsPatterns or similar

renamed to populateVectorLinearizeShuffleLikeOpsPatterns

dcaballe · 2024-04-17T06:24:25Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

+  LogicalResult
+  matchAndRewrite(vector::ExtractStridedSliceOp extractOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto dstType = getTypeConverter()->convertType(extractOp.getType());


nit: spell out auto

dcaballe · 2024-04-17T06:25:31Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

+    //      vector<4x8x8xf32> to vector<2x2x8xf32>
+    // Here, extraction granularity is 8.
+    int64_t extractGranularitySize = 1;
+    auto n = extractOp.getSourceVectorType().getRank();


nit: please, spell out the type when it's not literally redundant in the statement.

dcaballe · 2024-04-17T06:26:11Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

+    auto n = extractOp.getSourceVectorType().getRank();
+    int64_t k = (int64_t)offsets.size();
+    if (n > k) {
+      for (unsigned i = 0; i < n - k; i++) {


nit: use pre-increment per coding guidelines

dcaballe · 2024-04-17T06:27:22Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

+    }
+    // Get total number of extracted slices.
+    int64_t nExtractedSlices = 1;
+    llvm::for_each(sizes, [&](Attribute size) {


for (Attribute size : sizes) {...} ?

@banach-space recommended using computeSuffixProduct. But that does not take AttrArray as argument. So I am not sure to use it for this case. Changed back to a for loop for now.

dcaballe · 2024-04-17T06:29:27Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

+          shuffleOp, "Can't flatten since targetBitWidth <= OpSize");
+
+    auto vec1 = adaptor.getV1();
+    auto vec2 = adaptor.getV2();


nit: spell out auto

dcaballe · 2024-04-17T06:29:41Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

+    // outermost dim.
+    if (rank > 1) {
+      auto shape = shuffleOp.getV1().getType().getShape();
+      for (unsigned i = 1; i < shape.size(); i++) {


pre-increment

banach-space

LGTM, left a few final suggestions, thanks!

banach-space · 2024-04-18T13:00:03Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

+      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


Suggested change

// If kD offsets are specified for nd source vector (n > k), the granularity

// If kD offsets are specified for nD source vector (n > k), the granularity

banach-space · 2024-04-18T14:45:23Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

+    int64_t extractSliceLen = 1;
+    auto n = extractOp.getSourceVectorType().getRank();
+    int64_t k = (int64_t)offsets.size();
+    if (n > k) {


Right, so there are 2 cases:

k < n

k = n

right? I think that you can get rid of if by doing something like:

int64_t outRank = (int64_t)offsets.size(); int64_t k = outRank; while (k < n) { extractGranularitySize *= extractOp.getSourceVectorType().getShape()[outRank + (n-k)]; ++k; }

This way there's only one case. Also, I'd avoid variable names like k and n for anything other than loop iterators. Perhaps numKDims and numMDims? Or inputRank/outputRank? Or kD/nD? The last suggestion would be consistent with comments and already a nice improvement :)

charithaintc · 2024-04-18T18:12:31Z

Hi @Hardcode84, @banach-space, @dcaballe,

Thanks for the insightful reviews. This is my first PR in llvm-project and I learned a lot 😃 👍

Can you please help me with the merge process.

Thanks!

github-actions · 2024-04-18T18:14:07Z

@charithaintc Congratulations on having your first Pull Request (PR) merged into the LLVM Project!

Your changes will be combined with recent changes from other authors, then tested
by our build bots. If there is a problem with a build, you may receive a report in an email or a comment on this PR.

Please check whether problems have been caused by your change specifically, as
the builds can include changes from many authors. It is not uncommon for your
change to be included in a build that fails due to someone else's changes, or
infrastructure issues.

How to do this, and the rest of the post-merge process, is covered in detail here.

If your change does cause a problem, it may be reverted, or you can revert it yourself.
This is a normal part of LLVM development. You can fix your changes and open a new PR to merge them again.

If you don't get any reports, no action is required from you. Your changes are working as expected, well done!

…inearize Building on top of llvm#88204, this commit adds support for InsertOp.

…inearize (#92370) Building on top of [#88204](#88204), this PR adds support for converting `vector.insert` into an equivalent `vector.shuffle` operation that operates on linearized (1-D) vectors.

…inearize (llvm#92370) Building on top of [llvm#88204](llvm#88204), this PR adds support for converting `vector.insert` into an equivalent `vector.shuffle` operation that operates on linearized (1-D) vectors.

charithaintc added 4 commits April 9, 2024 14:04

add linearize patterns for Extract, ExtractStridedSlice, Shuffle Vect…

dc63b10

…orOps

Merge branch 'main' of https://github.com/llvm/llvm-project

24a0545

Merge branch 'main' into vector_linearize_mods

1de7dc4

remove comments

de748c0

fix test

962243c

charithaintc marked this pull request as ready for review April 9, 2024 22:39

charithaintc requested review from hanhanW and nicolasvasilache as code owners April 9, 2024 22:39

llvmbot added mlir:vectorops mlir mlir:vector labels Apr 9, 2024

Hardcode84 requested review from Hardcode84, dcaballe and bviyer April 9, 2024 22:40

banach-space reviewed Apr 12, 2024

View reviewed changes

address comments

e20be00

Merge branch 'main' into vector_linearize_mods

ded7ae0

dcaballe reviewed Apr 12, 2024

View reviewed changes

address comments

0d9406d

fix tests

e807c7a

Hardcode84 reviewed Apr 13, 2024

View reviewed changes

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp Outdated Show resolved Hide resolved

banach-space reviewed Apr 15, 2024

View reviewed changes

charithaintc added 4 commits April 16, 2024 11:13

Merge branch 'main' into vector_linearize_mods

186001a

fix

92374a2

add comments

c45c507

fix

1ac90fc

charithaintc added 2 commits April 16, 2024 13:51

fix

0b96134

fix comments

7046090

fix comments

780896e

fix

ec36971

dcaballe approved these changes Apr 17, 2024

View reviewed changes

charithaintc added 2 commits April 17, 2024 09:44

Merge branch 'main' into vector_linearize_mods

f282393

address comments

32f39da

banach-space approved these changes Apr 18, 2024

View reviewed changes

charithaintc added 3 commits April 18, 2024 10:36

Merge branch 'main' into vector_linearize_mods

e486a72

fix

9d908c8

Merge branch 'main' into vector_linearize_mods

5bbf22c

Hardcode84 merged commit c577f91 into llvm:main Apr 18, 2024
3 of 4 checks passed

akroviakov added a commit to akroviakov/llvm-project that referenced this pull request May 16, 2024

[mlir][vector] Add support for linearizing Insert VectorOp in VectorL…

1fa0571

…inearize Building on top of llvm#88204, this commit adds support for InsertOp.

akroviakov mentioned this pull request May 16, 2024

[mlir][vector] Add support for linearizing Insert VectorOp in VectorLinearize #92370

Merged

akroviakov added a commit to akroviakov/llvm-project that referenced this pull request May 16, 2024

[mlir][vector] Add support for linearizing Insert VectorOp in VectorL…

51115c2

…inearize Building on top of llvm#88204, this commit adds support for InsertOp.

akroviakov added a commit to akroviakov/llvm-project that referenced this pull request May 16, 2024

[mlir][vector] Add support for linearizing Insert VectorOp in VectorL…

18c92bb

…inearize Building on top of llvm#88204, this commit adds support for InsertOp.

akroviakov added a commit to akroviakov/llvm-project that referenced this pull request May 16, 2024

[mlir][vector] Add support for linearizing Insert VectorOp in VectorL…

d991dc9

…inearize Building on top of llvm#88204, this commit adds support for InsertOp.

akroviakov added a commit to akroviakov/llvm-project that referenced this pull request May 16, 2024

[mlir][vector] Add support for linearizing Insert VectorOp in VectorL…

41a5598

…inearize Building on top of llvm#88204, this commit adds support for InsertOp.

	// If kD offsets are specified for nd source vector (n > k), the granularity
	// If kD offsets are specified for nD source vector (n > k), the granularity

[mlir][vector] Add support for linearizing Extract, ExtractStridedSlice, Shuffle VectorOps in VectorLinearize #88204

[mlir][vector] Add support for linearizing Extract, ExtractStridedSlice, Shuffle VectorOps in VectorLinearize #88204

Uh oh!

Conversation

charithaintc commented Apr 9, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Apr 9, 2024

Uh oh!

llvmbot commented Apr 9, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

banach-space left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

charithaintc commented Apr 12, 2024

Uh oh!

dcaballe left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

charithaintc commented Apr 12, 2024

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

charithaintc commented Apr 16, 2024

Uh oh!

github-actions bot commented Apr 16, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Hardcode84 commented Apr 16, 2024

charithaintc commented Apr 9, 2024 •

edited

Loading

llvmbot commented Apr 9, 2024 •

edited

Loading

github-actions bot commented Apr 16, 2024 •

edited

Loading