[mlir][vector] Add pattern to drop unit dims from vector.transpose #102017

MacDue · 2024-08-05T16:55:27Z

Example:

BEFORE:

%transpose = vector.transpose %vector, [3, 0, 1, 2]
  : vector<1x1x4x[4]xf32> to vector<[4]x1x1x4xf32>

AFTER:

%dropDims = vector.shape_cast %vector
  : vector<1x1x4x[4]xf32> to vector<4x[4]xf32>
%transpose = vector.transpose %0, [1, 0]
  : vector<4x[4]xf32> to vector<[4]x4xf32>
%restoreDims = vector.shape_cast %transpose
  : vector<[4]x4xf32> to vector<[4]x1x1x4xf32>

llvmbot · 2024-08-05T16:55:58Z

@llvm/pr-subscribers-mlir-vector

@llvm/pr-subscribers-mlir

Author: Benjamin Maxwell (MacDue)

Changes

Example:

BEFORE:

%transpose = vector.transpose %vector, [3, 0, 1, 2]
  : vector&lt;1x1x4x[4]xf32&gt; to vector&lt;[4]x1x1x4xf32&gt;

AFTER:

%dropDims = vector.shape_cast %vector
  : vector&lt;1x1x4x[4]xf32&gt; to vector&lt;4x[4]xf32&gt;
%transpose = vector.transpose %0, [1, 0]
  : vector&lt;4x[4]xf32&gt; to vector&lt;[4]x4xf32&gt;
%restoreDims = vector.shape_cast %transpose
  : vector&lt;[4]x4xf32&gt; to vector&lt;[4]x1x1x4xf32&gt;

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

3 Files Affected:

(modified) mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h (+5)
(modified) mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp (+68-2)
(modified) mlir/test/Dialect/Vector/vector-transfer-flatten.mlir (+33)

diff --git a/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h b/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
index 40e04b76593a0..67c36bfa06ded 100644
--- a/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
+++ b/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
@@ -120,6 +120,11 @@ inline auto makeVscaleConstantBuilder(PatternRewriter &rewriter, Location loc) {
   };
 }
 
+/// Returns an iterator over the dims (inc scalability) of a VectorType.
+inline auto getDims(VectorType vType) {
+  return llvm::zip_equal(vType.getShape(), vType.getScalableDims());
+}
+
 /// A wrapper for getMixedSizes for vector.transfer_read and
 /// vector.transfer_write Ops (for source and destination, respectively).
 ///
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index 6777e589795c8..6b39cef7899d9 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -1720,6 +1720,72 @@ struct DropUnitDimFromElementwiseOps final
   }
 };
 
+/// A pattern to drop unit dims from vector.transpose.
+///
+/// Example:
+///
+///  BEFORE:
+///  ```mlir
+///  %transpose = vector.transpose %vector, [3, 0, 1, 2]
+///    : vector<1x1x4x[4]xf32> to vector<[4]x1x1x4xf32>
+///  ```
+///
+///  AFTER:
+///  ```mlir
+///  %dropDims = vector.shape_cast %vector
+///    : vector<1x1x4x[4]xf32> to vector<4x[4]xf32>
+///  %transpose = vector.transpose %0, [1, 0]
+///    : vector<4x[4]xf32> to vector<[4]x4xf32>
+///  %restoreDims = vector.shape_cast %transpose
+///    : vector<[4]x4xf32> to vector<[4]x1x1x4xf32>
+///  ```
+struct DropUnitDimsFromTransposeOp final
+    : OpRewritePattern<vector::TransposeOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::TransposeOp op,
+                                PatternRewriter &rewriter) const override {
+    VectorType sourceType = op.getSourceVectorType();
+    VectorType sourceTypeWithoutUnitDims =
+        dropNonScalableUnitDimFromType(sourceType);
+
+    if (sourceType == sourceTypeWithoutUnitDims)
+      return failure();
+
+    // Construct a map from dimIdx -> number of dims dropped before dimIdx.
+    auto sourceDims = llvm::to_vector(vector::getDims(sourceType));
+    SmallVector<int64_t> droppedDimsBefore(sourceType.getRank());
+    int64_t droppedDims = 0;
+    for (auto [i, dim] : llvm::enumerate(sourceDims)) {
+      droppedDimsBefore[i] = droppedDims;
+      if (dim == std::make_tuple(1, false))
+        ++droppedDims;
+    }
+
+    // Drop unit dims from transpose permutation.
+    ArrayRef<int64_t> perm = op.getPermutation();
+    SmallVector<int64_t> newPerm;
+    for (int64_t idx : perm) {
+      if (sourceDims[idx] == std::make_tuple(1, false))
+        continue;
+      newPerm.push_back(idx - droppedDimsBefore[idx]);
+    }
+
+    auto loc = op.getLoc();
+    // Drop the unit dims via shape_cast.
+    auto dropDimsShapeCast = rewriter.create<vector::ShapeCastOp>(
+        loc, sourceTypeWithoutUnitDims, op.getVector());
+    // Create the new transpose.
+    auto tranposeWithoutUnitDims =
+        rewriter.create<vector::TransposeOp>(loc, dropDimsShapeCast, newPerm);
+    // Restore the unit dims via shape cast.
+    rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(
+        op, op.getResultVectorType(), tranposeWithoutUnitDims);
+
+    return failure();
+  }
+};
+
 /// Pattern to eliminate redundant zero-constants added to reduction operands.
 /// It's enough for there to be one initial zero value, so we can eliminate the
 /// extra ones that feed into `vector.reduction <add>`. These get created by the
@@ -1924,8 +1990,8 @@ void mlir::vector::populateShapeCastFoldingPatterns(RewritePatternSet &patterns,
 
 void mlir::vector::populateDropUnitDimWithShapeCastPatterns(
     RewritePatternSet &patterns, PatternBenefit benefit) {
-  patterns.add<DropUnitDimFromElementwiseOps, ShapeCastOpFolder>(
-      patterns.getContext(), benefit);
+  patterns.add<DropUnitDimFromElementwiseOps, DropUnitDimsFromTransposeOp,
+               ShapeCastOpFolder>(patterns.getContext(), benefit);
 }
 
 void mlir::vector::populateBubbleVectorBitCastOpPatterns(
diff --git a/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir b/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir
index 9d16aa46a9f2a..222a05ff70d02 100644
--- a/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir
+++ b/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir
@@ -700,3 +700,36 @@ func.func @negative_out_of_bound_transfer_write(
 }
 // CHECK:     func.func @negative_out_of_bound_transfer_write
 // CHECK-NOT:   memref.collapse_shape
+
+// -----
+
+///----------------------------------------------------------------------------------------
+/// [Pattern: DropUnitDimsFromTransposeOp]
+/// TODO: Move to a dedicated file - there's no "flattening" in the following tests
+///----------------------------------------------------------------------------------------
+
+func.func @transpose_with_internal_unit_dims(%vector: vector<1x1x4x[4]xf32>) -> vector<[4]x1x1x4xf32> {
+  %0 = vector.transpose %vector, [3, 0, 1, 2] : vector<1x1x4x[4]xf32> to vector<[4]x1x1x4xf32>
+  return %0 : vector<[4]x1x1x4xf32>
+}
+
+// CHECK-LABEL: func.func @transpose_with_internal_unit_dims(
+// CHECK-SAME:                                               %[[VEC:.*]]: vector<1x1x4x[4]xf32>)
+// CHECK-NEXT:    %[[DROP_DIMS:.*]] = vector.shape_cast %arg0 : vector<1x1x4x[4]xf32> to vector<4x[4]xf32>
+// CHECK-NEXT:    %[[TRANSPOSE:.*]] = vector.transpose %0, [1, 0] : vector<4x[4]xf32> to vector<[4]x4xf32>
+// CHECK-NEXT:    %[[RESTORE_DIMS:.*]] = vector.shape_cast %1 : vector<[4]x4xf32> to vector<[4]x1x1x4xf32>
+// CHECK-NEXT:    return %[[RESTORE_DIMS]] : vector<[4]x1x1x4xf32>
+
+// -----
+
+func.func @transpose_with_units_dims_before_and_after(%vector: vector<1x1x1x4x[4]x1xf32>) -> vector<[4]x1x1x1x4x1xf32> {
+  %0 = vector.transpose %vector, [4, 1, 0, 2, 3, 5] : vector<1x1x1x4x[4]x1xf32> to vector<[4]x1x1x1x4x1xf32>
+  return %0 : vector<[4]x1x1x1x4x1xf32>
+}
+
+// CHECK-LABEL: func.func @transpose_with_units_dims_before_and_after(
+// CHECK-SAME:                                                        %[[VEC:.*]]: vector<1x1x1x4x[4]x1xf32>)
+// CHECK-NEXT:    %[[DROP_DIMS:.*]] = vector.shape_cast %arg0 : vector<1x1x1x4x[4]x1xf32> to vector<4x[4]xf32>
+// CHECK-NEXT:    %[[TRANSPOSE:.*]] = vector.transpose %0, [1, 0] : vector<4x[4]xf32> to vector<[4]x4xf32>
+// CHECK-NEXT:    %[[RESTORE_DIMS:.*]] = vector.shape_cast %1 : vector<[4]x4xf32> to vector<[4]x1x1x1x4x1xf32>
+// CHECK-NEXT:    return %[[RESTORE_DIMS]] : vector<[4]x1x1x1x4x1xf32>

Example: BEFORE: ```mlir %transpose = vector.transpose %vector, [3, 0, 1, 2] : vector<1x1x4x[4]xf32> to vector<[4]x1x1x4xf32> ``` AFTER: ```mlir %dropDims = vector.shape_cast %vector : vector<1x1x4x[4]xf32> to vector<4x[4]xf32> %transpose = vector.transpose %0, [1, 0] : vector<4x[4]xf32> to vector<[4]x4xf32> %restoreDims = vector.shape_cast %transpose : vector<[4]x4xf32> to vector<[4]x1x1x4xf32> ```

banach-space · 2024-08-06T12:22:50Z

Adding DropUnitDimsFromTransposeOp to populateDropUnitDimWithShapeCastPatterns makes a lot of sense to me in terms of grouping these patterns.

How is this related to #100933? Also, let's make sure that somebody with SPIR-V lenses on takes a look :) CC @kuhar

MacDue · 2024-08-07T09:28:49Z

This can supersede #100933

banach-space

Argh, sorry, forgot to send yesterday.

mlir/test/Dialect/Vector/vector-transfer-flatten.mlir

mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h

kuhar

I don't foresee this causing any issues related to the SPIR-V lowering.

mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h

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

…spose(shape_cast) (#100731)" (#102457) This reverts commit 88accd9. This change can be dropped in favor of just #102017.

dcaballe

Great! Thanks for bearing with me. Appreciate it

MacDue requested review from banach-space, dcaballe, hanhanW and nicolasvasilache as code owners August 5, 2024 16:55

llvmbot added mlir:vectorops mlir mlir:vector labels Aug 5, 2024

MacDue mentioned this pull request Aug 5, 2024

[mlir][vector] Add SwapShapeCastOfTranspose canonicalizer pattern #100933

Closed

MacDue requested a review from c-rhodes August 6, 2024 09:51

MacDue force-pushed the drop_tr_dims branch from 0f41cb1 to d41ea3c Compare August 6, 2024 09:53

banach-space requested a review from kuhar August 6, 2024 12:22

banach-space reviewed Aug 7, 2024

View reviewed changes

mlir/test/Dialect/Vector/vector-transfer-flatten.mlir Outdated Show resolved Hide resolved

mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h Outdated Show resolved Hide resolved

Update tests

3680f20

kuhar approved these changes Aug 7, 2024

View reviewed changes

mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h Outdated Show resolved Hide resolved

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

Fixups

ecc9375

MacDue merged commit da8778e into llvm:main Aug 8, 2024
7 checks passed

MacDue deleted the drop_tr_dims branch August 8, 2024 10:42

MacDue mentioned this pull request Aug 8, 2024

Revert "[mlir][ArmSME] Pattern to swap shape_cast(tranpose) with transpose(shape_cast) (#100731)" #102457

Merged

MacDue added a commit that referenced this pull request Aug 9, 2024

Revert "[mlir][ArmSME] Pattern to swap shape_cast(tranpose) with tran…

fc4485b

…spose(shape_cast) (#100731)" (#102457) This reverts commit 88accd9. This change can be dropped in favor of just #102017.

dcaballe reviewed Aug 20, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[mlir][vector] Add pattern to drop unit dims from vector.transpose #102017

[mlir][vector] Add pattern to drop unit dims from vector.transpose #102017

Uh oh!

MacDue commented Aug 5, 2024

Uh oh!

llvmbot commented Aug 5, 2024 •

edited

Loading

Uh oh!

banach-space commented Aug 6, 2024

Uh oh!

MacDue commented Aug 7, 2024

Uh oh!

banach-space left a comment

Uh oh!

Uh oh!

Uh oh!

kuhar left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dcaballe left a comment

Uh oh!

Uh oh!

[mlir][vector] Add pattern to drop unit dims from vector.transpose #102017

[mlir][vector] Add pattern to drop unit dims from vector.transpose #102017

Uh oh!

Conversation

MacDue commented Aug 5, 2024

Uh oh!

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

Uh oh!

Uh oh!

banach-space commented Aug 6, 2024

Uh oh!

MacDue commented Aug 7, 2024

Uh oh!

banach-space left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

kuhar left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dcaballe left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

llvmbot commented Aug 5, 2024 •

edited

Loading