From d2f55a3689101c51c1a00bfd4a9875c09cda206e Mon Sep 17 00:00:00 2001
From: Adam Cheney <acheney@nvidia.com>
Date: Mon, 13 Jan 2025 15:56:52 -0500
Subject: [PATCH 1/2] Fix simplify if-else

The if-else optimization observes that at if at least one
true/false block is merely an unconditional jump to the
after block, that the whole if-else can be replaced with
a jump to the after block. But it's important to copy
the phi arguments from the aforementioned unconditional
jump, rather than what is present in the 'true' block,
since the 'true' block might actually just be the after
block itself.

Below, the ifElse() would be replaced with an unconditional
jump to block %39, but with the `phi` arguments copied from
the branch to %29, which is an unrelated block.

ifElse(%38, %39, %40, %39)

block %40:
    unconditionalBranch(%39)

block %39:
    unconditionalBranch(%29, 0 : Float)

block %29(
        [nameHint("ret")]
        param %ret  : Float):

Fixes issue #5972
---
 source/slang/slang-ir-simplify-cfg.cpp | 11 +++++++++--
 tests/bugs/simplify-if-else.slang      | 26 ++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 2 deletions(-)
 create mode 100644 tests/bugs/simplify-if-else.slang

diff --git a/source/slang/slang-ir-simplify-cfg.cpp b/source/slang/slang-ir-simplify-cfg.cpp
index 90d30dcc77..6b3b4e4641 100644
--- a/source/slang/slang-ir-simplify-cfg.cpp
+++ b/source/slang/slang-ir-simplify-cfg.cpp
@@ -490,11 +490,18 @@ static bool trySimplifyIfElse(IRBuilder& builder, IRIfElse* ifElseInst)
     bool isFalseBranchTrivial = false;
     if (isTrivialIfElse(ifElseInst, isTrueBranchTrivial, isFalseBranchTrivial))
     {
-        // If both branches of `if-else` are trivial jumps into after block,
+        // If either branch of `if-else` is a trivial jump into after block,
         // we can get rid of the entire conditional branch and replace it
         // with a jump into the after block.
-        if (auto termInst = as<IRUnconditionalBranch>(ifElseInst->getTrueBlock()->getTerminator()))
+        IRUnconditionalBranch* termInst = as<IRUnconditionalBranch>(ifElseInst->getTrueBlock()->getTerminator());
+        if (!termInst || (termInst->getTargetBlock() != ifElseInst->getAfterBlock()))
         {
+            termInst = as<IRUnconditionalBranch>(ifElseInst->getFalseBlock()->getTerminator());
+        }
+
+        if (termInst)
+        {
+            SLANG_ASSERT(termInst->getTargetBlock() == ifElseInst->getAfterBlock());
             List<IRInst*> args;
             for (UInt i = 0; i < termInst->getArgCount(); i++)
                 args.add(termInst->getArg(i));
diff --git a/tests/bugs/simplify-if-else.slang b/tests/bugs/simplify-if-else.slang
new file mode 100644
index 0000000000..8719a15995
--- /dev/null
+++ b/tests/bugs/simplify-if-else.slang
@@ -0,0 +1,26 @@
+//TEST:SIMPLE(filecheck=CHECK): -stage compute -entry computeMain -target hlsl
+//CHECK: computeMain
+
+//TEST_INPUT:ubuffer(data=[9 9 9 9], stride=4):out,name=outputBuffer
+RWStructuredBuffer<int> outputBuffer;
+
+[numthreads(4, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    vector<float32_t, 4> vvv = vector<float32_t, 4>(0);
+    float32_t ret = 0.0f;
+    if (vvv.y < 1.0f)
+    {
+        ret = 1.0f;
+    }
+    else
+    {
+        if (vvv.y > 1.0f && outputBuffer[3] == 3)
+        {
+            ret = 0.0f;
+        } else {
+            if (true) {}
+        }
+    }
+    outputBuffer[int(dispatchThreadID.x)] = int(ret);
+}

From b42f6b314b945f004c275b900e0e67d99e449727 Mon Sep 17 00:00:00 2001
From: slangbot <186143334+slangbot@users.noreply.github.com>
Date: Tue, 14 Jan 2025 15:05:57 +0000
Subject: [PATCH 2/2] format code

---
 source/slang/slang-ir-simplify-cfg.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/source/slang/slang-ir-simplify-cfg.cpp b/source/slang/slang-ir-simplify-cfg.cpp
index 6b3b4e4641..68d79617a8 100644
--- a/source/slang/slang-ir-simplify-cfg.cpp
+++ b/source/slang/slang-ir-simplify-cfg.cpp
@@ -493,7 +493,8 @@ static bool trySimplifyIfElse(IRBuilder& builder, IRIfElse* ifElseInst)
         // If either branch of `if-else` is a trivial jump into after block,
         // we can get rid of the entire conditional branch and replace it
         // with a jump into the after block.
-        IRUnconditionalBranch* termInst = as<IRUnconditionalBranch>(ifElseInst->getTrueBlock()->getTerminator());
+        IRUnconditionalBranch* termInst =
+            as<IRUnconditionalBranch>(ifElseInst->getTrueBlock()->getTerminator());
         if (!termInst || (termInst->getTargetBlock() != ifElseInst->getAfterBlock()))
         {
             termInst = as<IRUnconditionalBranch>(ifElseInst->getFalseBlock()->getTerminator());