[Mlir-commits] [mlir] [mlir][vector] fold transpose(poison) -> poison (PR #135675)

[James Newling]

https://github.com/newling updated https://github.com/llvm/llvm-project/pull/135675

>From fa8ebf6c8cc5f4a50aa96056f12811fa3dbdaaf5 Mon Sep 17 00:00:00 2001
From: James Newling <james.newling at gmail.com>
Date: Mon, 14 Apr 2025 13:35:17 -0700
Subject: [PATCH 1/2] Canonicalize transpose(poison) to poison

Signed-off-by: James Newling <james.newling at gmail.com>
---
 mlir/lib/Dialect/Vector/IR/VectorOps.cpp   |  5 +++++
 mlir/test/Dialect/Vector/canonicalize.mlir | 11 +++++++++++
 2 files changed, 16 insertions(+)

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index bee5c1fd6ed58..3a5646267d7ea 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -6009,6 +6009,11 @@ OpFoldResult vector::TransposeOp::fold(FoldAdaptor adaptor) {
     if (attr.isSplat())
       return attr.reshape(getResultVectorType());
 
+
+  // Eliminate poison transpose ops.
+  if (llvm::dyn_cast_if_present<ub::PoisonAttr>(adaptor.getVector()))
+    return ub::PoisonAttr::get(getContext());
+
   // Eliminate identity transpose ops. This happens when the dimensions of the
   // input vector remain in their original order after the transpose operation.
   ArrayRef<int64_t> perm = getPermutation();
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 78b0ea78849e8..bf5d8d5495077 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -2264,6 +2264,17 @@ func.func @transpose_splat2(%arg : f32) -> vector<3x4xf32> {
 
 // -----
 
+// CHECK-LABEL: transpose_poison
+//       CHECK:  %[[POISON:.*]] = ub.poison : vector<4x6xi8>
+//       CHECK:  return %[[POISON]] : vector<4x6xi8>
+func.func @transpose_poison() -> vector<4x6xi8> {
+  %poison = ub.poison : vector<6x4xi8>
+  %transpose = vector.transpose %poison, [1, 0] : vector<6x4xi8> to vector<4x6xi8>
+  return %transpose : vector<4x6xi8>
+}
+
+// -----
+
 // CHECK-LABEL: func.func @insert_1d_constant
 //   CHECK-DAG: %[[ACST:.*]] = arith.constant dense<[9, 1, 2]> : vector<3xi32>
 //   CHECK-DAG: %[[BCST:.*]] = arith.constant dense<[0, 9, 2]> : vector<3xi32>

>From 866cd62163c28f99fad04a017c3e9ba1566fe09b Mon Sep 17 00:00:00 2001
From: James Newling <james.newling at gmail.com>
Date: Mon, 14 Apr 2025 13:50:45 -0700
Subject: [PATCH 2/2] clang-format

---
 mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 1 -
 1 file changed, 1 deletion(-)

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 3a5646267d7ea..e812d3d880ef3 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -6009,7 +6009,6 @@ OpFoldResult vector::TransposeOp::fold(FoldAdaptor adaptor) {
     if (attr.isSplat())
       return attr.reshape(getResultVectorType());
 
-
   // Eliminate poison transpose ops.
   if (llvm::dyn_cast_if_present<ub::PoisonAttr>(adaptor.getVector()))
     return ub::PoisonAttr::get(getContext());