[mlir][ArmSME] Replace use of isa with isa_and_present

[MacDue]

op can be null here, in which case this should just return a null value back.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-sme

Author: Benjamin Maxwell (MacDue)

Changes

op can be null here, in which case this should just return a null value back.

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/ArmSME/Transforms/VectorLegalization.cpp (+1-1)
• (modified) mlir/test/Dialect/ArmSME/vector-legalization.mlir (+11)

diff --git a/mlir/lib/Dialect/ArmSME/Transforms/VectorLegalization.cpp b/mlir/lib/Dialect/ArmSME/Transforms/VectorLegalization.cpp
index 26dfb38263372b..11f8bc04b21844 100644
--- a/mlir/lib/Dialect/ArmSME/Transforms/VectorLegalization.cpp
+++ b/mlir/lib/Dialect/ArmSME/Transforms/VectorLegalization.cpp
@@ -459,7 +459,7 @@ struct LiftIllegalVectorTransposeToMemory
   }
 
   static Value getExtensionSource(Operation *op) {
-    if (isa<arith::ExtSIOp, arith::ExtUIOp, arith::ExtFOp>(op))
+    if (isa_and_present<arith::ExtSIOp, arith::ExtUIOp, arith::ExtFOp>(op))
       return op->getOperand(0);
     return {};
   }
diff --git a/mlir/test/Dialect/ArmSME/vector-legalization.mlir b/mlir/test/Dialect/ArmSME/vector-legalization.mlir
index 11888c675f0b04..bf0b58ff4cf073 100644
--- a/mlir/test/Dialect/ArmSME/vector-legalization.mlir
+++ b/mlir/test/Dialect/ArmSME/vector-legalization.mlir
@@ -377,3 +377,14 @@ func.func @lift_illegal_transpose_to_memory_with_in_bounds_attr(%a: index, %b: i
   %legalType = vector.transpose %illegalRead, [1, 0] : vector<[8]x4xf32> to vector<4x[8]xf32>
   return %legalType : vector<4x[8]xf32>
 }
+
+// -----
+
+// The pass should do nothing (and not crash).
+// CHECK-LABEL: @illegal_transpose_no_defining_source_op
+func.func @illegal_transpose_no_defining_source_op(%vec: vector<[4]x1xf32>) -> vector<1x[4]xf32>
+{
+  // CHECK: vector.transpose
+  %0 = vector.transpose %vec, [1, 0] : vector<[4]x1xf32> to vector<1x[4]xf32>
+  return %0 : vector<1x[4]xf32>
+}