[FoldSCFIf] fixed issues when no else

Author: kumasento

tools/polymer/lib/Transforms/FoldSCFIf.cc

@@ -53,13 +53,17 @@ static void foldSCFIf(scf::IfOp ifOp, FuncOp f, OpBuilder &b) {
   };
 
   cloneAfter(ifOp.thenBlock(), thenResults);
-  cloneAfter(ifOp.elseBlock(), elseResults);
 
-  for (auto ifResult : enumerate(ifOp.getResults())) {
-    Value newResult =
-        b.create<SelectOp>(loc, ifOp.condition(), thenResults[ifResult.index()],
-                           elseResults[ifResult.index()]);
-    ifResult.value().replaceAllUsesWith(newResult);
+  // Only an if op can have results when an else block is present.
+  if (ifOp.elseBlock()) {
+    cloneAfter(ifOp.elseBlock(), elseResults);
+
+    for (auto ifResult : enumerate(ifOp.getResults())) {
+      Value newResult = b.create<SelectOp>(loc, ifOp.condition(),
+                                           thenResults[ifResult.index()],
+                                           elseResults[ifResult.index()]);
+      ifResult.value().replaceAllUsesWith(newResult);
+    }
   }
 
   ifOp.erase();

tools/polymer/test/polymer-opt/FoldSCFIf/no-result-no-else.mlir

@@ -0,0 +1,12 @@
+// RUN: polymer-opt %s -fold-scf-if | FileCheck %s
+
+func @foo(%a: f32, %b: f32, %c: i1) {
+  scf.if %c {
+    %0 = arith.addf %a, %b : f32
+  } 
+  return
+}
+
+// CHECK: func @foo
+// CHECK-NEXT: arith.addf
+// CHECK-NEXT: return

tools/polymer/test/polymer-opt/FoldSCFIf/no-result.mlir

@@ -0,0 +1,15 @@
+// RUN: polymer-opt %s -fold-scf-if | FileCheck %s
+func @foo(%a: f32, %b: f32, %c: i1) {
+  scf.if %c {
+    %0 = arith.addf %a, %b : f32
+  } else {
+    %0 = arith.mulf %a, %b : f32
+  }
+
+  return
+}
+
+// CHECK: func @foo
+// CHECK-NEXT: arith.addf
+// CHECK-NEXT: arith.mulf
+// CHECK-NEXT: return