Linearize vector.transpose to vector.shuffle (#714)

Upstream vector dialect already have transform which lowers transpose to linearized shuffle.
Add it to our linearization pipeline.

Author: Hardcode84

lib/Transforms/VectorLinearize.cpp

@@ -13,17 +13,19 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
+#include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/Pass/Pass.h"
 #include "mlir/Support/LogicalResult.h"
+#include "mlir/Transforms/DialectConversion.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/Transforms/Utils/AddDiscriminators.h"
-#include <cstdint>
-#include <imex/Transforms/Passes.h>
 
-#include <mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h>
-#include <mlir/Pass/Pass.h>
-#include <mlir/Transforms/DialectConversion.h>
+#include "imex/Transforms/Passes.h"
+
+#include <cstdint>
 #include <numeric>
 
 namespace imex {
@@ -227,15 +229,23 @@ struct VectorLinearizePass final
     mlir::RewritePatternSet patterns(context);
     mlir::ConversionTarget target(*context);
 
+    typeConverter.addConversion([](mlir::Type type) { return type; });
+
     target.addDynamicallyLegalOp<mlir::vector::ShuffleOp>([&](mlir::Operation
                                                                   *op) {
       return op->getResult(0).getType().cast<mlir::VectorType>().getRank() == 1;
     });
 
+    target.addIllegalOp<mlir::vector::TransposeOp>();
+    target.addLegalOp<mlir::vector::ShapeCastOp>();
+
     patterns.add<VectorExtractStridedSliceConversion, VectorShffleOpConversion,
                  VectorExtractOpConversion>(typeConverter, context);
 
-    typeConverter.addConversion([](mlir::Type type) { return type; });
+    mlir::vector::populateVectorTransposeLoweringPatterns(
+        patterns,
+        mlir::vector::VectorTransformsOptions().setVectorTransposeLowering(
+            mlir::vector::VectorTransposeLowering::Shuffle1D));
     mlir::vector::populateVectorLinearizeTypeConversionsAndLegality(
         typeConverter, patterns, target);
     if (mlir::failed(mlir::applyPartialConversion(getOperation(), target,

test/Transforms/vector-linearize.mlir

@@ -97,4 +97,18 @@ func.func @test_vector_shuffle(%arg0: vector<4x4xf32>, %arg1: vector<4x4xf32>) -
 func.func @test_vector_extract(%arg0: vector<2x8x4xf32>) -> vector<8x4xf32> {
   %0 = vector.extract %arg0[1]: vector<8x4xf32> from vector<2x8x4xf32>
   return %0 : vector<8x4xf32>
-}
+}
+
+// -----
+
+// CHECK-LABEL: test_vector_transpose
+// CHECK-SAME: (%[[ORIG_ARG:.*]]: vector<2x8xf32>) -> vector<8x2xf32>
+// CHECK: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<2x8xf32> to vector<16xf32>
+// CHECK: %[[SHUFFLE:.*]] = vector.shuffle %[[ARG]], %[[ARG]]
+// CHECK: [0, 8, 1, 9, 2, 10, 3, 11, 4, 12, 5, 13, 6, 14, 7, 15] : vector<16xf32>, vector<16xf32>
+// CHECK: %[[RES:.*]] = vector.shape_cast %[[SHUFFLE]] : vector<16xf32> to vector<8x2xf32>
+// CHECK: return %[[RES]] : vector<8x2xf32>
+func.func @test_vector_transpose(%arg: vector<2x8xf32>) -> vector<8x2xf32> {
+  %0 = vector.transpose %arg, [1, 0] : vector<2x8xf32> to vector<8x2xf32>
+  return %0 : vector<8x2xf32>
+}