[Mlir-commits] [mlir] [mlir][cf] Canonicalize cf.switch cases targeting ub.unreachable (PR #234387)

[Mehdi Amini]

================
@@ -804,6 +808,60 @@ dropSwitchCasesThatMatchDefault(SwitchOp op, PatternRewriter &rewriter) {
   return success();
 }
 
+/// Drop cases whose destination block contains only ub.unreachable. If the
+/// default destination is unreachable, the last remaining case becomes the new
+/// default destination.
+/// switch %flag : i32, [
+///   default: ^bb1,
+///   42: ^bb2,
+///   43: ^bb3
+/// ]
+/// ^bb3:
+///   ub.unreachable
+/// -> switch %flag : i32, [
+///   default: ^bb1,
+///   42: ^bb2
+/// ]
+static LogicalResult dropUnreachableSwitchCases(SwitchOp op,
+                                                PatternRewriter &rewriter) {
+  std::optional<DenseIntElementsAttr> caseValues = op.getCaseValues();
+  if (!caseValues)
+    return failure();
+
+  SmallVector<APInt> newCaseValues;
+  SmallVector<Block *> newCaseDestinations;
+  SmallVector<ValueRange> newCaseOperands;
+  auto caseDests = op.getCaseDestinations();
+  for (const auto &it : llvm::enumerate(caseValues->getValues<APInt>())) {
+    if (isUnreachableBlock(caseDests[it.index()]))
+      continue;
+    newCaseValues.push_back(it.value());
+    newCaseDestinations.push_back(caseDests[it.index()]);
+    newCaseOperands.push_back(op.getCaseOperands(it.index()));
+  }
+
+  Block *defaultDest = op.getDefaultDestination();
+  ValueRange defaultOperands = op.getDefaultOperands();
+  if (isUnreachableBlock(defaultDest) && !newCaseDestinations.empty()) {
+    // Any value not matching a remaining case is undefined behavior, so the
+    // last remaining case can serve as the default destination.
+    defaultDest = newCaseDestinations.pop_back_val();
+    defaultOperands = newCaseOperands.pop_back_val();
+    newCaseValues.pop_back();
+  } else if (newCaseDestinations.size() == caseValues->size()) {
+    return failure();
+  }
----------------
joker-eph wrote:

Seems like there is a missing case here: `if (isUnreachableBlock(defaultDest) && newCaseDestinations.empty()) {` ; then we can just fold entirely to unreachable?

https://github.com/llvm/llvm-project/pull/234387