[mlir][Linalg] Change linalg.transpose to use the output indexing map as identity.

[MaheshRavishankar]

This makes it consistent with how other linalg operations represent indexing maps.

[llvmbot]

@llvm/pr-subscribers-mlir-linalg

@llvm/pr-subscribers-mlir

Author: None (MaheshRavishankar)

Changes

This makes it consistent with how other linalg operations represent indexing maps.

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp (+3-3)
• (modified) mlir/test/Dialect/Linalg/vectorization.mlir (+8-8)

diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 828a140be75456..ec823d06c98f39 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -1773,9 +1773,9 @@ ArrayAttr TransposeOp::getIndexingMaps() {
   Builder builder(getContext());
   int64_t rank = getInit().getType().getRank();
   return builder.getAffineMapArrayAttr(
-      {builder.getMultiDimIdentityMap(rank),
-       AffineMap::getPermutationMap(
-           llvm::to_vector_of<unsigned>(getPermutation()), getContext())});
+      {inversePermutation(AffineMap::getPermutationMap(
+           llvm::to_vector_of<unsigned>(getPermutation()), getContext())),
+       builder.getMultiDimIdentityMap(rank)});
 }
 
 void TransposeOp::getEffects(
diff --git a/mlir/test/Dialect/Linalg/vectorization.mlir b/mlir/test/Dialect/Linalg/vectorization.mlir
index 610339405d1c2c..d5fb0cbb9c723b 100644
--- a/mlir/test/Dialect/Linalg/vectorization.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization.mlir
@@ -371,16 +371,16 @@ module attributes {transform.with_named_sequence} {
 // CHECK: #[[MAP:.*]] = affine_map<(d0, d1) -> (d1, d0)>
 // CHECK: func @test_masked_vectorize_linalg_transpose
 func.func @test_masked_vectorize_linalg_transpose(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {
-  // CHECK:      %[[C0:.*]] = arith.constant 0 : index
-  // CHECK:      %[[D0:.*]] = tensor.dim %arg0, %[[C0]]
-  // CHECK:      %[[C1:.*]] = arith.constant 1 : index
-  // CHECK:      %[[D1:.*]] = tensor.dim %arg0, %[[C1]]
+  // CHECK-DAG:  %[[C0:.*]] = arith.constant 0 : index
+  // CHECK-DAG:  %[[D0:.*]] = tensor.dim %arg0, %[[C0]]
+  // CHECK-DAG:  %[[C1:.*]] = arith.constant 1 : index
+  // CHECK-DAG:  %[[D1:.*]] = tensor.dim %arg0, %[[C1]]
   // CHECK:      %[[MASK0:.*]] = vector.create_mask %[[D0]], %[[D1]]
-  // CHECK:      %[[LOAD:.*]] = vector.mask %[[MASK0]] { vector.transfer_read %arg0{{.+}} }
-  // CHECK-SAME:   vector<2x4xi1> -> vector<2x4xf32>
+  // CHECK:      %[[LOAD:.*]] = vector.mask %[[MASK0]] { vector.transfer_read %arg0{{.+}} permutation_map = #[[MAP]]{{.+}} }
+  // CHECK-SAME:   vector<4x2xi1> -> vector<2x4xf32>
   // CHECK:      %[[MASK1:.*]] = vector.create_mask %[[D1]], %[[D0]]
-  // CHECK:      %[[WRITE:.*]] = vector.mask %[[MASK1]] { vector.transfer_write %[[LOAD]], %arg1{{.+}} permutation_map = #[[MAP]]{{.+}} }
-  // CHECK-SAME:   vector<4x2xi1> -> tensor<?x?xf32>
+  // CHECK:      %[[WRITE:.*]] = vector.mask %[[MASK1]] { vector.transfer_write %[[LOAD]], %arg1{{.+}} }
+  // CHECK-SAME:   vector<2x4xi1> -> tensor<?x?xf32>
   // CHECK:      return %[[WRITE]]
   %0 = linalg.transpose ins(%arg0 : tensor<?x?xf32>) outs(%arg1 : tensor<?x?xf32>) permutation = [1, 0]
   return %0 : tensor<?x?xf32>

[qedawkins]

LGTM