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

[olegshyshkov]

Summary

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

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

 // TransposeOp
 //===----------------------------------------------------------------------===//
 
 std::function<void(mlir::ImplicitLocOpBuilder &, mlir::Block &,
                    mlir::ArrayRef<mlir::NamedAttribute>)>
 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());
   };
 }
 
 void TransposeOp::createRegion(::mlir::OpBuilder &opBuilder,
                                ::mlir::OperationState &odsState) {
   Region *region = odsState.addRegion();
 
   SmallVector<Type> argTypes;

mlir/test/Dialect/Linalg/vectorization.mlir

 
 func.func @vectorize_transpose(%arg0: memref<16x32x64xf32>,
                                %arg1: memref<32x64x16xf32>) {
   linalg.transpose ins(%arg0 : memref<16x32x64xf32>)
                    outs(%arg1 : memref<32x64x16xf32>) permutation = [1, 2, 0]
   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) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.transpose"]} in %arg1
   %1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation
   %2 = transform.structured.vectorize %1
 }

mlir/test/Interfaces/TilingInterface/lower-to-loops-using-interface.mlir

 // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
 // CHECK-DAG: %[[C16:.*]] = arith.constant 16 : index
 // CHECK-DAG: %[[C32:.*]] = arith.constant 32 : index
 // CHECK-DAG: %[[C64:.*]] = arith.constant 64 : index
 
 // 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]]]
 
 // -----
 
 func.func @reduce(%arg0: memref<16x32x64xf32>,
                   %arg1: memref<16x64xf32>) {
   linalg.reduce ins(%arg0 : memref<16x32x64xf32>)
                 outs(%arg1 : memref<16x64xf32>) dimensions = [1]