[mlir][ArmSME] Replace use of isa with isa_and_present (#82798)

`op` can be null here, in which case this should just return a null value back.
llvm · Feb 26, 2024 · 8cfb716 · 8cfb716
1 parent cb2dd02
commit 8cfb716
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git 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
@@ -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>
+}