Fixed a crash in unroller when exiting blocks use switch instead of branch (microsoft#5779)

The code to structurize loop exits assumed exiting blocks always use
BranchInst and never SwitchInst. This caused a crash. This change makes
the structurizer preemptively check whether this assumption is true
before attempting.

Author: adam-yang

lib/Transforms/Scalar/DxilLoopUnroll.cpp

@@ -137,11 +137,13 @@ class DxilLoopUnroll : public LoopPass {
     GetPassOptionUnsigned(O, "MaxIterationAttempt", &MaxIterationAttempt,
                           false);
     GetPassOptionBool(O, "OnlyWarnOnFail", &OnlyWarnOnFail, false);
+    GetPassOptionBool(O, "StructurizeLoopExits", &StructurizeLoopExits, false);
   }
   void dumpConfig(raw_ostream &OS) override {
     LoopPass::dumpConfig(OS);
     OS << ",MaxIterationAttempt=" << MaxIterationAttempt;
     OS << ",OnlyWarnOnFail=" << OnlyWarnOnFail;
+    OS << ",StructurizeLoopExits=" << StructurizeLoopExits;
   }
   void RecursivelyRemoveLoopOnSuccess(LPPassManager &LPM, Loop *L);
   void RecursivelyRecreateSubLoopForIteration(LPPassManager &LPM, LoopInfo *LI,

lib/Transforms/Scalar/DxilRemoveUnstructuredLoopExits.cpp

@@ -112,6 +112,9 @@
 //     but the code will execute for one iteration even if the exit condition
 //     is met.
 //
+//   - If any exiting block uses a switch statement to conditionally exit the
+//     loop, we currently do not handle that case.
+//
 // These limitations can be fixed in the future as needed.
 //
 //===----------------------------------------------------------------------===//
@@ -561,6 +564,17 @@ bool hlsl::RemoveUnstructuredLoopExits(
   if (!L->isLCSSAForm(*DT))
     return false;
 
+  // Check that the exiting blocks in the loop have BranchInst terminators (as
+  // opposed to SwitchInst). At the moment we only handle BranchInst case.
+  {
+    llvm::SmallVector<BasicBlock *, 4> exiting_blocks;
+    L->getExitingBlocks(exiting_blocks);
+    for (BasicBlock *BB : exiting_blocks) {
+      if (!isa<BranchInst>(BB->getTerminator()))
+        return false;
+    }
+  }
+
   // Give up if loop is not rotated somehow
   if (BasicBlock *latch = L->getLoopLatch()) {
     if (!cast<BranchInst>(latch->getTerminator())->isConditional())

tools/clang/test/HLSLFileCheck/hlsl/control_flow/attributes/unroll/unroll_switch_exit_crash.hlsl

@@ -0,0 +1,39 @@
+// RUN: %dxc %s /T ps_6_0 | FileCheck %s
+
+// Regression test for a crash when unrolling a loop whose exiting block
+// uses a switch and not a branch.
+
+// CHECK: call float @dx.op.unary.f32(i32 13
+// CHECK: call float @dx.op.unary.f32(i32 13
+// CHECK: call float @dx.op.unary.f32(i32 13
+// CHECK: call float @dx.op.unary.f32(i32 13
+// CHECK-NOT: call float @dx.op.unary.f32(i32 13
+
+Texture1D<float> t0 : register(t0);
+
+float foo(int cond) {
+  float ret = 0;
+  [unroll]
+  for (int i = 0; i < 4; i++) {
+    ret += 1;
+    switch (cond) {
+      case 0:
+        ret += 10;
+        break;
+      case 1:
+        ret += 20;
+        break;
+      default:
+        return 42;
+    }
+    ret += sin(t0[cond + i]);
+  }
+  return ret;
+}
+
+[RootSignature("DescriptorTable(SRV(t0))")]
+float main (int cond : COND) : SV_Target {
+  float ret = 0;
+  ret += foo(cond);
+  return ret;
+}

tools/clang/test/HLSLFileCheck/hlsl/control_flow/attributes/unroll/unroll_switch_exit_crash.ll

@@ -0,0 +1,61 @@
+; RUN: %opt %s -dxil-loop-unroll,StructurizeLoopExits=1 -S | FileCheck %s
+
+; CHECK: call float @llvm.sin.f32(
+; CHECK: call float @llvm.sin.f32(
+; CHECK: call float @llvm.sin.f32(
+; CHECK: call float @llvm.sin.f32(
+; CHECK-NOT: call float @llvm.sin.f32(
+
+target datalayout = "e-m:e-p:32:32-i1:32-i8:32-i16:32-i32:32-i64:64-f16:32-f32:32-f64:64-n8:16:32:64"
+target triple = "dxil-ms-dx"
+
+declare float @llvm.sin.f32(float %Val) #0
+
+; Function Attrs: nounwind
+define float @main(i32 %cond) #1 {
+entry:
+  br label %for.body
+
+for.body:
+  %i.0 = phi i32 [ 0, %entry ], [ %inc, %sw.epilog ]
+  %ret.0 = phi float [ 0.000000e+00, %entry ], [ %add7, %sw.epilog ]
+  %add = fadd fast float %ret.0, 1.000000e+00
+  switch i32 %cond, label %end [
+    i32 0, label %sw.bb
+    i32 1, label %sw.bb.3
+  ]
+
+sw.bb:
+  %add2 = fadd fast float %add, 1.000000e+01
+  br label %sw.epilog
+
+sw.bb.3:
+  %add4 = fadd fast float %add, 2.000000e+01
+  br label %sw.epilog
+
+sw.epilog:
+  %ret.1 = phi float [ %add4, %sw.bb.3 ], [ %add2, %sw.bb ]
+  %Sin = call float @llvm.sin.f32(float 0.0)
+  %add7 = fadd fast float %ret.1, %Sin
+  %inc = add nsw i32 %i.0, 1
+  %cmp = icmp slt i32 %inc, 4
+  br i1 %cmp, label %for.body, label %end, !llvm.loop !3
+
+end:                           ; preds = %sw.epilog, %for.body
+  %retval.0 = phi float [ 4.200000e+01, %for.body ], [ %add7, %sw.epilog ]
+  ret float %retval.0
+}
+
+attributes #0 = { nounwind readnone }
+attributes #1 = { nounwind }
+attributes #2 = { nounwind readonly }
+
+!llvm.module.flags = !{!0}
+!pauseresume = !{!1}
+!llvmdent = !{!2}
+
+!0 = !{i32 2, !"Debug Info Version", i32 3}
+!1 = !{!"hlsl-dxilemit", !"hlsl-dxilload"}
+!2 = !{!"dxc(private) 1.7.0.14003 (main, a5b0488bc-dirty)"}
+!3 = distinct !{!3, !4}
+!4 = !{!"llvm.loop.unroll.full"}

utils/hct/hctdb.py

@@ -2250,6 +2250,7 @@ def add_pass(name, type_name, doc, opts):
         add_pass('dxil-loop-unroll', 'DxilLoopUnroll', 'DxilLoopUnroll', [
             {'n':'MaxIterationAttempt', 't':'unsigned', 'c':1, 'd':'Maximum number of iterations to attempt when iteratively unrolling.'},
             {'n':'OnlyWarnOnFail', 't':'bool', 'c':1, 'd':'Whether to just warn when unrolling fails.'},
+            {'n':'StructurizeLoopExits', 't':'bool', 'c':1, 'd':'Whether the unroller should try to structurize loop exits first.'}
         ])
         add_pass('dxil-erase-dead-region', 'DxilEraseDeadRegion', 'DxilEraseDeadRegion', [])
         add_pass('dxil-remove-dead-blocks', 'DxilRemoveDeadBlocks', 'DxilRemoveDeadBlocks', [])