Turn off structurize-returns when cleanup blocks are present. (#4927)

structurize-returns uses scope-end blocks recorded during codegen to transform the control flow. When "cleanup" blocks are generated (for example for lifetime-markers), structurize-returns as is cannot transform the control flow safely. This change disables structurize-returns and emits a warning when cleanup blocks are detected in the affected function.

(cherry picked from commit 93f43ec)

Author: adam-yang

tools/clang/include/clang/Basic/DiagnosticGroups.td

@@ -798,4 +798,5 @@ def HLSLPayloadAccessQualifer: DiagGroup<"payload-access-qualifier", [
      HLSLPayloadAccessQualiferCall
   ]>;
 def HLSLSemanticIdentifierCollision : DiagGroup<"semantic-identifier-collision">;
+def HLSLStructurizeExitsLifetimeMarkersConflict: DiagGroup<"structurize-exits-lifetime-markers-conflict">;
 // HLSL Change Ends

tools/clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -7754,6 +7754,9 @@ def err_hlsl_logical_binop_scalar : Error<
    "operands for short-circuiting logical binary operator must be scalar, for non-scalar types use '%select{and|or}0'">;
 def err_hlsl_ternary_scalar : Error<
    "condition for short-circuiting ternary operator must be scalar, for non-scalar types use 'select'">;
+def warn_hlsl_structurize_exits_lifetime_markers_conflict : Warning <
+   "structurize-returns skipped function '%0' due to incompatibility with lifetime markers. Use -disable-lifetime-markers to enable structurize-exits on this function.">,
+  InGroup< HLSLStructurizeExitsLifetimeMarkersConflict >;
 // HLSL Change Ends
 
 // SPIRV Change Starts

tools/clang/lib/CodeGen/CGCleanup.cpp

@@ -19,6 +19,7 @@
 
 #include "CGCleanup.h"
 #include "CodeGenFunction.h"
+#include "CGHLSLRuntime.h"    // HLSL Change
 
 using namespace clang;
 using namespace CodeGen;
@@ -435,6 +436,7 @@ static llvm::BasicBlock *CreateNormalEntry(CodeGenFunction &CGF,
   if (!Entry) {
     Entry = CGF.createBasicBlock("cleanup");
     Scope.setNormalBlock(Entry);
+    CGF.CGM.getHLSLRuntime().MarkCleanupBlock(CGF, Entry); // HLSL Change
   }
   return Entry;
 }

tools/clang/lib/CodeGen/CGHLSLMS.cpp

@@ -304,6 +304,7 @@ class CGMSHLSLRuntime : public CGHLSLRuntime {
   void MarkSwitchStmt(CodeGenFunction &CGF, SwitchInst *switchInst,
                       BasicBlock *endSwitch) override;
   void MarkReturnStmt(CodeGenFunction &CGF, BasicBlock *bbWithRet) override;
+  void MarkCleanupBlock(CodeGenFunction &CGF, llvm::BasicBlock *cleanupBB) override;
   void MarkLoopStmt(CodeGenFunction &CGF, BasicBlock *loopContinue,
                      BasicBlock *loopExit) override;
   CGHLSLMSHelper::Scope* MarkScopeEnd(CodeGenFunction &CGF) override;
@@ -2393,7 +2394,7 @@ void CGMSHLSLRuntime::AddHLSLFunctionInfo(Function *F, const FunctionDecl *FD) {
     F->addFnAttr(Twine("exp-", Attr->getName()).str(), Attr->getValue());
   }
 
-  m_ScopeMap[F] = ScopeInfo(F);
+  m_ScopeMap[F] = ScopeInfo(F, FD->getLocation());
 }
 
 void CGMSHLSLRuntime::RemapObsoleteSemantic(DxilParameterAnnotation &paramInfo, bool isPatchConstantFunction) {
@@ -6233,6 +6234,10 @@ void CGMSHLSLRuntime::MarkIfStmt(CodeGenFunction &CGF, BasicBlock *endIfBB) {
     Scope->AddIf(endIfBB);
 }
 
+void CGMSHLSLRuntime::MarkCleanupBlock(CodeGenFunction &CGF, llvm::BasicBlock *cleanupBB) {
+  if (ScopeInfo *Scope = GetScopeInfo(CGF.CurFn))
+    Scope->AddCleanupBB(cleanupBB);
+}
 
 void CGMSHLSLRuntime::MarkSwitchStmt(CodeGenFunction &CGF,
                                      SwitchInst *switchInst,

tools/clang/lib/CodeGen/CGHLSLMSFinishCodeGen.cpp

@@ -29,6 +29,7 @@
 #include "clang/Basic/LangOptions.h"
 #include "clang/Frontend/CodeGenOptions.h"
 #include "clang/Parse/ParseHLSL.h" // root sig would be in Parser if part of lang
+#include "clang/Sema/SemaDiagnostic.h"
 #include "CodeGenFunction.h"
 
 #include "dxc/DXIL/DxilConstants.h"
@@ -3309,7 +3310,7 @@ void AddDxBreak(Module &M,
 
 namespace CGHLSLMSHelper {
 
-ScopeInfo::ScopeInfo(Function *F) : maxRetLevel(0), bAllReturnsInIf(true) {
+ScopeInfo::ScopeInfo(Function *F, clang::SourceLocation loc) : maxRetLevel(0), bAllReturnsInIf(true), sourceLoc(loc) {
   Scope FuncScope;
   FuncScope.kind = Scope::ScopeKind::FunctionScope;
   FuncScope.EndScopeBB = nullptr;
@@ -3549,12 +3550,21 @@ static void ChangePredBranch(BasicBlock *BB, BasicBlock *NewBB) {
 //   }
 //   return vRet;
 // }
-void StructurizeMultiRetFunction(Function *F, ScopeInfo &ScopeInfo,
+void StructurizeMultiRetFunction(Function *F, clang::DiagnosticsEngine &Diags, ScopeInfo &ScopeInfo,
                                  bool bWaveEnabledStage,
                                  SmallVector<BranchInst *, 16> &DxBreaks) {
 
   if (ScopeInfo.CanSkipStructurize())
     return;
+
+  // If there are cleanup blocks generated for lifetime markers, do
+  // not structurize returns. The scope info recorded is no longer correct.
+  if (ScopeInfo.HasCleanupBlocks()) {
+    Diags.Report(ScopeInfo.GetSourceLocation(), clang::diag::warn_hlsl_structurize_exits_lifetime_markers_conflict)
+      << F->getName();
+    return;
+  }
+
   // Get bbWithRets.
   auto &rets = ScopeInfo.GetRetScopes();
 
@@ -3722,7 +3732,8 @@ void StructurizeMultiRet(Module &M, clang::CodeGen::CodeGenModule &CGM,
     auto it = ScopeMap.find(&F);
     if (it == ScopeMap.end())
       continue;
-    StructurizeMultiRetFunction(&F, it->second, bWaveEnabledStage, DxBreaks);
+
+    StructurizeMultiRetFunction(&F, CGM.getDiags(), it->second, bWaveEnabledStage, DxBreaks);
   }
 }
 

tools/clang/lib/CodeGen/CGHLSLMSHelper.h

@@ -122,27 +122,32 @@ struct Scope {
 class ScopeInfo {
 public:
   ScopeInfo(){}
-  ScopeInfo(llvm::Function *F);
+  ScopeInfo(llvm::Function *F, clang::SourceLocation loc);
   void AddIf(llvm::BasicBlock *endIfBB);
   void AddSwitch(llvm::BasicBlock *endSwitchBB);
   void AddLoop(llvm::BasicBlock *loopContinue, llvm::BasicBlock *endLoopBB);
   void AddRet(llvm::BasicBlock *bbWithRet);
+  void AddCleanupBB(llvm::BasicBlock *cleanupBB) { hasCleanupBlocks = true; }
   Scope &EndScope(bool bScopeFinishedWithRet);
   Scope &GetScope(unsigned i);
   const llvm::SmallVector<unsigned, 2> &GetRetScopes() { return rets; }
   void LegalizeWholeReturnedScope();
   llvm::SmallVector<Scope, 16> &GetScopes() { return scopes; }
   bool CanSkipStructurize();
   void dump();
+  bool HasCleanupBlocks() const { return hasCleanupBlocks; }
+  clang::SourceLocation GetSourceLocation() const { return sourceLoc; }
 
 private:
   void AddScope(Scope::ScopeKind k, llvm::BasicBlock *endScopeBB);
+  bool hasCleanupBlocks = false;
   llvm::SmallVector<unsigned, 2> rets;
   unsigned maxRetLevel;
   bool bAllReturnsInIf;
   llvm::SmallVector<unsigned, 8> scopeStack;
   // save all scopes.
   llvm::SmallVector<Scope, 16> scopes;
+  clang::SourceLocation sourceLoc;
 };
 
 // Map from value to resource properties.

tools/clang/lib/CodeGen/CGHLSLRuntime.h

@@ -140,6 +140,7 @@ class CGHLSLRuntime {
       clang::QualType FnRetTy,
       const std::function<void(const VarDecl *, llvm::Value *)> &TmpArgMap) = 0;
   virtual void MarkIfStmt(CodeGenFunction &CGF, llvm::BasicBlock *endIfBB) = 0;
+  virtual void MarkCleanupBlock(CodeGenFunction &CGF, llvm::BasicBlock *cleanupBB) = 0;
   virtual void MarkSwitchStmt(CodeGenFunction &CGF,
                               llvm::SwitchInst *switchInst,
                               llvm::BasicBlock *endSwitch) = 0;

tools/clang/test/HLSLFileCheck/hlsl/control_flow/return/lifetime-markers.hlsl

@@ -0,0 +1,68 @@
+// RUN: %dxc -fdisable-loc-tracking -E main -opt-enable structurize-returns -T cs_6_0 -enable-lifetime-markers -fcgl %s | FileCheck %s -check-prefix=FCGL
+// RUN: %dxc -fdisable-loc-tracking -E main -opt-enable structurize-returns -T cs_6_0 -enable-lifetime-markers %s | FileCheck %s
+// RUN: %dxc -fdisable-loc-tracking -E main -opt-enable structurize-returns -T cs_6_0 -enable-lifetime-markers %s | FileCheck %s -check-prefix=WARNING -input-file=stderr
+// RUN: %dxc -fdisable-loc-tracking -E main -opt-enable structurize-returns -T cs_6_0 -disable-lifetime-markers -fcgl %s | FileCheck %s -check-prefix=NO-LIFETIME
+
+// Regression test for a bug where program structure is completely messed up when lifetime-markers are enabled and
+// -opt-enable structurize-returns is on. The scope information recorded during codegen that structurize-returns uses
+// to modify the control flow is incorrect if lifetime-markers are enabled. This test checks that 
+
+//=================================
+// The fcgl test checks the return condition alloca bReturn is not generated and the cleanup code for lifetime-markers
+// is present.
+// FCGL-NOT: bReturned
+// FCGL: %cleanup.dest
+
+//=================================
+// The non-fcgl test checks the shader is compiled correctly (the bug causes irreducible flow)
+// CHECK-DAG: @main
+
+//=================================
+// Check a warning was emitted.
+// WARNING: structurize-returns skipped function 'main' due to incompatibility with lifetime markers. Use -disable-lifetime-markers to enable structurize-exits on this function.
+
+//=================================
+// The last test makes sure structurize-returns runs as expected
+// NO-LIFETIME: @main
+// NO-LIFETIME: %bReturned = alloca
+
+struct D {
+ float3 d_member;
+};
+
+struct A {
+  float4 a_member;
+};
+
+struct B {
+    uint flags;
+} ;
+
+struct C {
+    uint c_member;
+};
+
+StructuredBuffer   <D> srv0 : register(t0) ;
+StructuredBuffer   <B> srv1 : register(t1) ;
+RWStructuredBuffer <C> uav0 : register(u0) ;
+
+[RootSignature("DescriptorTable(SRV(t0,numDescriptors=10)),DescriptorTable(UAV(u0,numDescriptors=10))")]
+[numthreads (64, 1, 1)]
+void main(uint3 dtid : SV_DispatchThreadID) {
+ if (dtid.x < 10) {
+    A decal = (A)0;
+    {
+      const D d = srv0[0];
+      if (!d.d_member.x)
+        return;
+    }
+    B b = srv1[0];
+    if (b.flags & 1) {
+      InterlockedMax(uav0[0].c_member, 10) ;
+      return;
+    }
+
+    InterlockedMax(uav0[0].c_member, 20) ;
+  }
+  InterlockedMax(uav0[0].c_member, 30) ;
+}