Revert "[mlir][vector] Make TransposeOpLowering configurable (#73915)"

[banach-space]

Reverting a workaround intended specifically for SPRI-V. That workaround
emerged from this discussion:

• [mlir][vector] Add vector.transpose with unit-dim to vector.shape_cast pattern #72105

AFAIK, it hasn't been required in practice. This is based on IREE
(https://github.com/openxla/iree), which has just bumped it's fork of LLVM
without using it (*).

(*) iree-org/iree@cef31e7

This reverts commit bbd2b08.

[llvmbot]

@llvm/pr-subscribers-mlir-vector

@llvm/pr-subscribers-mlir

Author: Andrzej Warzyński (banach-space)

Changes

Reverting a workaround intended specifically for SPRI-V. That workaround
emerged from this discussion:

• [mlir][vector] Add vector.transpose with unit-dim to vector.shape_cast pattern #72105

AFAIK, it hasn't been required in practice. This is based on IREE
(https://github.com/openxla/iree), which has just bumped it's fork of LLVM
without using it (*).

(*) iree-org/iree@cef31e7

This reverts commit bbd2b08.

---

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

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/Vector/Transforms/VectorTransforms.h (-10)
• (modified) mlir/lib/Dialect/Vector/Transforms/LowerVectorTranspose.cpp (+16-18)

diff --git a/mlir/include/mlir/Dialect/Vector/Transforms/VectorTransforms.h b/mlir/include/mlir/Dialect/Vector/Transforms/VectorTransforms.h
index 41ffc929946027..08d3bb157a0e39 100644
--- a/mlir/include/mlir/Dialect/Vector/Transforms/VectorTransforms.h
+++ b/mlir/include/mlir/Dialect/Vector/Transforms/VectorTransforms.h
@@ -59,16 +59,6 @@ struct VectorTransformsOptions {
     vectorTransferSplit = opt;
     return *this;
   }
-
-  /// Option to control if vector.transpose can lower to a vector.shape_cast.
-  /// TODO: ATM it's not possible to lower `vector.shape_cast` to SPIR-V
-  /// and hence the need for this opt-out. Once the missing support has been
-  /// added, this option can be removed.
-  bool useShapeCast = true;
-  VectorTransformsOptions &setUseShapeCast(bool opt = true) {
-    useShapeCast = opt;
-    return *this;
-  }
 };
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorTranspose.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorTranspose.cpp
index 4d43a76c4a4efc..97f6caca1b25cc 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorTranspose.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorTranspose.cpp
@@ -334,24 +334,22 @@ class TransposeOpLowering : public OpRewritePattern<vector::TransposeOp> {
       return rewriter.notifyMatchFailure(
           op, "Options specifies lowering to shuffle");
 
-    if (vectorTransformOptions.useShapeCast) {
-      // Replace:
-      //   vector.transpose %0, [1, 0] : vector<nx1x<eltty>> to
-      //                                 vector<1xnxelty>
-      // with:
-      //   vector.shape_cast %0 : vector<nx1x<eltty>> to vector<1xnxelty>
-      //
-      // Source with leading unit dim (inverse) is also replaced. Unit dim must
-      // be fixed. Non-unit can be scalable.
-      if (resType.getRank() == 2 &&
-          ((resType.getShape().front() == 1 &&
-            !resType.getScalableDims().front()) ||
-           (resType.getShape().back() == 1 &&
-            !resType.getScalableDims().back())) &&
-          transp == ArrayRef<int64_t>({1, 0})) {
-        rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(op, resType, input);
-        return success();
-      }
+    // Replace:
+    //   vector.transpose %0, [1, 0] : vector<nx1x<eltty>> to
+    //                                 vector<1xnxelty>
+    // with:
+    //   vector.shape_cast %0 : vector<nx1x<eltty>> to vector<1xnxelty>
+    //
+    // Source with leading unit dim (inverse) is also replaced. Unit dim must
+    // be fixed. Non-unit can be scalable.
+    if (resType.getRank() == 2 &&
+        ((resType.getShape().front() == 1 &&
+          !resType.getScalableDims().front()) ||
+         (resType.getShape().back() == 1 &&
+          !resType.getScalableDims().back())) &&
+        transp == ArrayRef<int64_t>({1, 0})) {
+      rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(op, resType, input);
+      return success();
     }
 
     if (inputType.isScalable())

[hanhanW]

If I understand correctly, it is still needed by IREE. IREE pings llvm-project at https://github.com/shark-infra/llvm-project/commits/6cdda88cccc2e65803a893e3bca15bec9cf5d3f7 It includes the commit. If we revert it, it will break IREE in the next integrate.

[MacDue]

If I understand correctly, it is still needed by IREE. IREE pings llvm-project at https://github.com/shark-infra/llvm-project/commits/6cdda88cccc2e65803a893e3bca15bec9cf5d3f7 It includes the commit. If we revert it, it will break IREE in the next integrate.

The change reverted here is not used in IREE. It added a config option to disabled using shape_casts in LowerVectorTranspose.
This was done to avoid issues for SPIR-V in IREE, however, instead the shape_cast lowerings were just enabled for SPIR-V (see here). So this option is entirely unused in IREE.

[dcaballe]

Cool! Nicolas was right after all... :)

[hanhanW]

I see, that makes sense to me now!

[antiagainst]

I was trying to do this and you beat me. :) Thanks!

Cool! Nicolas was right after all... :)

Yup that works, thanks for the suggestion! We explicitly control when to perform transpose lowering and thus we can slot the lowering for type_cast in there too. If this is happening as a global canonicalization we would not be able to do that.

Anyway, glad this is resolved to unblock everybody. We will keep discussing some of the underlying issues surfaced from this and hopefully get to a much better long term solution. :)