[mlir][Linalg] Bugfix for folder of linalg.transpose (#102888)

Folder of linalg transpose only support tensor type. Fix #102576.

Author: CoTinker

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

@@ -1908,6 +1908,10 @@ void TransposeOp::getEffects(
 
 LogicalResult TransposeOp::fold(FoldAdaptor adaptor,
                                 SmallVectorImpl<OpFoldResult> &result) {
+  // Only the tensor type is supported.
+  if (!isa<TensorType>(getInput().getType()))
+    return failure();
+
   // Single dimension transpose.
   if (getPermutation().size() == 0) {
     result.push_back(getInput());

mlir/test/Dialect/Linalg/canonicalize.mlir

@@ -1216,3 +1216,19 @@ func.func @concats_of_fill(
 //       CHECK:   %[[CONCAT:.+]] = tensor.concat dim(1) %[[EMPTY0]], %[[EMPTY1]]
 //       CHECK:   %[[FILL:.+]] = linalg.fill ins(%[[CST]] : f32) outs(%[[CONCAT]] :
 //       CHECK:   return %[[FILL]]
+
+// -----
+
+func.func @transpose_buffer(%input: memref<?xf32>,
+                            %init: memref<?xf32>) {
+  linalg.transpose ins(%input:memref<?xf32>)
+                   outs(%init:memref<?xf32>)
+                   permutation = [0]
+  func.return
+}
+
+// CHECK-LABEL:   func.func @transpose_buffer(
+//  CHECK-SAME:            %[[VAL_0:.*]]: memref<?xf32>,
+//  CHECK-SAME:            %[[VAL_1:.*]]: memref<?xf32>) {
+//       CHECK:     linalg.transpose ins(%[[VAL_0]] : memref<?xf32>)
+//  CHECK-SAME:       outs(%[[VAL_1]] : memref<?xf32>) permutation = [0]

mlir/test/Dialect/Linalg/loops.mlir

@@ -873,3 +873,39 @@ func.func @lower_to_loops_with_rank_reducing_subviews(
 //       CHECKPARALLEL:     %[[VAL:.+]] = memref.load %{{.+}}[%[[IV]]]
 //       CHECKPARALLEL:     memref.store %[[VAL]], %{{.+}}[%[[IV]]]
 //       CHECKPARALLEL:   }
+
+// -----
+
+func.func @transpose(%input: memref<?xf32>,
+                     %init: memref<?xf32>) {
+  linalg.transpose ins(%input:memref<?xf32>)
+                   outs(%init:memref<?xf32>)
+                   permutation = [0]
+  return
+}
+// CHECK-LABEL:   func.func @transpose(
+// CHECK-SAME:                         %[[VAL_0:.*]]: memref<?xf32>,
+// CHECK-SAME:                         %[[VAL_1:.*]]: memref<?xf32>) {
+//      CHECK:      %[[VAL_2:.*]] = arith.constant 1 : index
+//      CHECK:      %[[VAL_3:.*]] = arith.constant 0 : index
+//      CHECK:      %[[VAL_4:.*]] = memref.dim %[[VAL_0]], %[[VAL_3]] : memref<?xf32>
+//      CHECK:      scf.for %[[VAL_5:.*]] = %[[VAL_3]] to %[[VAL_4]] step %[[VAL_2]] {
+//      CHECK:        %[[VAL_6:.*]] = memref.load %[[VAL_0]]{{\[}}%[[VAL_5]]] : memref<?xf32>
+//      CHECK:        memref.store %[[VAL_6]], %[[VAL_1]]{{\[}}%[[VAL_5]]] : memref<?xf32>
+//      CHECK:      }
+//      CHECK:      return
+//      CHECK:    }
+
+// CHECKPARALLEL-LABEL:   func.func @transpose(
+// CHECKPARALLEL-SAME:                         %[[VAL_0:.*]]: memref<?xf32>,
+// CHECKPARALLEL-SAME:                         %[[VAL_1:.*]]: memref<?xf32>) {
+//      CHECKPARALLEL:      %[[VAL_2:.*]] = arith.constant 1 : index
+//      CHECKPARALLEL:      %[[VAL_3:.*]] = arith.constant 0 : index
+//      CHECKPARALLEL:      %[[VAL_4:.*]] = memref.dim %[[VAL_0]], %[[VAL_3]] : memref<?xf32>
+//      CHECKPARALLEL:      scf.parallel (%[[VAL_5:.*]]) = (%[[VAL_3]]) to (%[[VAL_4]]) step (%[[VAL_2]]) {
+//      CHECKPARALLEL:        %[[VAL_6:.*]] = memref.load %[[VAL_0]]{{\[}}%[[VAL_5]]] : memref<?xf32>
+//      CHECKPARALLEL:        memref.store %[[VAL_6]], %[[VAL_1]]{{\[}}%[[VAL_5]]] : memref<?xf32>
+//      CHECKPARALLEL:        scf.reduce
+//      CHECKPARALLEL:      }
+//      CHECKPARALLEL:      return
+//      CHECKPARALLEL:    }