[mlir][linalg] Fix linalg.transpose region builder.

The region should yield the first argument (input) not the last argument
(output). Also fix a few tests that were affected by this bug.

Differential Revision: https://reviews.llvm.org/D136924

diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 03a959a..896fcf4 100644 (file)
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -1614,7 +1614,7 @@ std::function<void(mlir::ImplicitLocOpBuilder &, mlir::Block &,
 TransposeOp::getRegionBuilder() {
   return [](mlir::ImplicitLocOpBuilder &b, mlir::Block &block,
             mlir::ArrayRef<mlir::NamedAttribute>) {
-    b.create<linalg::YieldOp>(block.getArguments().back());
+    b.create<linalg::YieldOp>(block.getArguments().front());
   };
 }
 

diff --git a/mlir/test/Dialect/Linalg/vectorization.mlir b/mlir/test/Dialect/Linalg/vectorization.mlir
index 39c6d9d..3b351dc 100644 (file)
--- a/mlir/test/Dialect/Linalg/vectorization.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization.mlir
@@ -1556,9 +1556,7 @@ func.func @vectorize_transpose(%arg0: memref<16x32x64xf32>,
   return
 }
 // CHECK-LABEL: func @vectorize_transpose
-// CHECK:         %[[FIRST:.*]] = vector.transpose
-// CHECK-SAME:      [2, 0, 1] : vector<32x64x16xf32> to vector<16x32x64xf32>
-// CHECK-NEXT:    vector.transpose %[[FIRST]]
+// CHECK:         vector.transpose
 // CHECK-SAME:      [1, 2, 0] : vector<16x32x64xf32> to vector<32x64x16xf32>
 
 transform.sequence failures(propagate) {

diff --git a/mlir/test/Interfaces/TilingInterface/lower-to-loops-using-interface.mlir b/mlir/test/Interfaces/TilingInterface/lower-to-loops-using-interface.mlir
index f9dd94e..9addbcc 100644 (file)
--- a/mlir/test/Interfaces/TilingInterface/lower-to-loops-using-interface.mlir
+++ b/mlir/test/Interfaces/TilingInterface/lower-to-loops-using-interface.mlir
@@ -208,7 +208,7 @@ func.func @transpose(%arg0: memref<16x32x64xf32>,
 // CHECK:     scf.for %[[I:.*]] = %[[C0]] to %[[C16]] step %[[C1]] {
 // CHECK:       scf.for %[[J:.*]] = %[[C0]] to %[[C32]] step %[[C1]] {
 // CHECK:         scf.for %[[K:.*]] = %[[C0]] to %[[C64]] step %[[C1]] {
-// CHECK:           %[[ELEM:.*]] = memref.load %[[OUT]][%[[J]], %[[K]], %[[I]]]
+// CHECK:           %[[ELEM:.*]] = memref.load %[[IN]][%[[I]], %[[J]], %[[K]]]
 // CHECK:           memref.store %[[ELEM]], %[[OUT]][%[[J]], %[[K]], %[[I]]]
 
 // -----