From 2298848552e3d3ec58667e7befbcc6212cbd7196 Mon Sep 17 00:00:00 2001
From: Xiang Li <python3kgae@outlook.com>
Date: Thu, 6 May 2021 22:39:31 -0700
Subject: [PATCH] Skip structurize when only 1 ret. (#3754)

---
 .../lib/CodeGen/CGHLSLMSFinishCodeGen.cpp     |  2 +-
 .../return/single_ret_in_loop.hlsl            | 22 +++++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)
 create mode 100644 tools/clang/test/HLSLFileCheck/hlsl/control_flow/return/single_ret_in_loop.hlsl

diff --git a/tools/clang/lib/CodeGen/CGHLSLMSFinishCodeGen.cpp b/tools/clang/lib/CodeGen/CGHLSLMSFinishCodeGen.cpp
index b6d096fb2..b0404ea4e 100644
--- a/tools/clang/lib/CodeGen/CGHLSLMSFinishCodeGen.cpp
+++ b/tools/clang/lib/CodeGen/CGHLSLMSFinishCodeGen.cpp
@@ -3149,7 +3149,7 @@ ScopeInfo::ScopeInfo(Function *F) : maxRetLevel(0), bAllReturnsInIf(true) {
 // When all returns is inside if which is not nested, the flow is still
 // structurized even there're more than one return.
 bool ScopeInfo::CanSkipStructurize() {
-  return bAllReturnsInIf && maxRetLevel < 2;
+  return (bAllReturnsInIf && maxRetLevel < 2) || rets.size() < 2;
 }
 
 void ScopeInfo::AddScope(Scope::ScopeKind k, BasicBlock *endScopeBB) {
diff --git a/tools/clang/test/HLSLFileCheck/hlsl/control_flow/return/single_ret_in_loop.hlsl b/tools/clang/test/HLSLFileCheck/hlsl/control_flow/return/single_ret_in_loop.hlsl
new file mode 100644
index 000000000..9f7532eee
--- /dev/null
+++ b/tools/clang/test/HLSLFileCheck/hlsl/control_flow/return/single_ret_in_loop.hlsl
@@ -0,0 +1,22 @@
+// RUN: %dxc -Emain -opt-enable structurize-returns -Tps_6_0 %s | FileCheck %s
+
+// Make sure not crash.
+// CHECK:call %dx.types.ResRet.i32 @dx.op.bufferLoad
+
+struct ST
+{
+ uint a;
+};
+
+
+StructuredBuffer<ST> buf;
+
+uint main(uint i:I) : SV_Target {
+ for(;;)
+ {
+  if(i<buf[i].a)
+   i++;
+  else
+   return buf[i].a;
+ }
+}
\ No newline at end of file