Fix switch to if

Author: wsmoses

src/enzyme_ad/jax/Passes/CanonicalizeLoops.cpp

@@ -546,7 +546,7 @@ class SwitchToIf final : public OpRewritePattern<scf::IndexSwitchOp> {
                                    cmpResult, /*withElseRegion=*/true);
 
     // Move the first case block into the then region
-    Block &firstBlock = switchOp.getCaseBlock(cases.front());
+    Block &firstBlock = switchOp.getCaseBlock(0);
     rewriter.mergeBlocks(&firstBlock, ifOp.thenBlock(),
                          firstBlock.getArguments());
 

test/lit_tests/switch_to_if.mlir

@@ -24,6 +24,30 @@ func.func @switch_to_if(%arg0: index) -> i32 {
   return %0 : i32
 }
 
+// CHECK-LABEL: func @switch_to_if2
+func.func @switch_to_if2(%arg0: index) -> i32 {
+  // CHECK-DAG: %[[CONST:.*]] = arith.constant 20 : index
+  // CHECK-DAG: %[[VAL1:.*]] = arith.constant 42 : i32
+  // CHECK-DAG: %[[VAL2:.*]] = arith.constant 24 : i32
+  // CHECK: %[[CMP:.*]] = arith.cmpi eq, %arg0, %[[CONST]] : index
+  // CHECK: %[[RESULT:.*]] = scf.if %[[CMP]] -> (i32) {
+  // CHECK:   scf.yield %[[VAL1]] : i32
+  // CHECK: } else {
+  // CHECK:   scf.yield %[[VAL2]] : i32
+  // CHECK: }
+  // CHECK: return %[[RESULT]] : i32
+  %0 = scf.index_switch %arg0 -> i32
+    case 20 {
+      %1 = arith.constant 42 : i32
+      scf.yield %1 : i32
+    }
+    default {
+      %1 = arith.constant 24 : i32
+      scf.yield %1 : i32
+    }
+  return %0 : i32
+}
+
 // Should not convert switches with more than 1 case
 // CHECK-LABEL: func @switch_two_cases
 func.func @switch_two_cases(%arg0: index) -> i32 {