From c787c4b82ba76f87069911f203eb192060b5264f Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Mon, 28 Aug 2023 21:24:49 -0700
Subject: Add `target_switch` and `intrinsic_asm` statement. (#3154)

* Add `target_switch` and `__intrinsic_asm` statement.

* Cleanup.

* WaveGetActiveMask, WaveGetActiveMask, WaveCountBits.

* WaveIsFirstLane.

* More wave intrinsics.

* wave intrinsics.

* merge fix.

* Fix.

* Fix.

* Update test.

* update test.

* Fix.

---------

Co-authored-by: Yong He <yhe@nvidia.com>
---
 source/slang/slang-ir-specialize-target-switch.cpp | 67 ++++++++++++++++++++++
 1 file changed, 67 insertions(+)
 create mode 100644 source/slang/slang-ir-specialize-target-switch.cpp

(limited to 'source/slang/slang-ir-specialize-target-switch.cpp')

diff --git a/source/slang/slang-ir-specialize-target-switch.cpp b/source/slang/slang-ir-specialize-target-switch.cpp
new file mode 100644
index 000000000..2593389b1
--- /dev/null
+++ b/source/slang/slang-ir-specialize-target-switch.cpp
@@ -0,0 +1,67 @@
+#include "slang-ir-specialize-target-switch.h"
+#include "slang-ir.h"
+#include "slang-ir-insts.h"
+#include "slang-compiler.h"
+#include "slang-capability.h"
+#include "slang-ir-dce.h"
+
+namespace Slang
+{
+    void specializeTargetSwitch(TargetRequest* target, IRGlobalValueWithCode* code)
+    {
+        bool changed = false;
+        for (auto block : code->getBlocks())
+        {
+            if (auto targetSwitch = as<IRTargetSwitch>(block->getTerminator()))
+            {
+                CapabilitySet bestCapSet = CapabilitySet::makeInvalid();
+                IRBlock* targetBlock = nullptr;
+                for (UInt i = 0; i < targetSwitch->getCaseCount(); i++)
+                {
+                    auto cap = (CapabilityAtom)getIntVal(targetSwitch->getCaseValue(i));
+                    CapabilitySet capSet;
+                    if (cap == CapabilityAtom::Invalid)
+                        capSet = CapabilitySet::makeEmpty();
+                    else
+                        capSet = CapabilitySet(cap);
+                    if (capSet.isBetterForTarget(bestCapSet, target->getTargetCaps()))
+                    {
+                        targetBlock = targetSwitch->getCaseBlock(i);
+                        bestCapSet = capSet;
+                    }
+                }
+                SLANG_ASSERT(targetBlock);
+                IRBuilder builder(targetSwitch);
+                builder.setInsertBefore(targetSwitch);
+                builder.emitBranch(targetBlock);
+                targetSwitch->removeAndDeallocate();
+                changed = true;
+            }
+        }
+        if (changed)
+        {
+            // Remove unreachable blocks after specialization.
+            eliminateDeadCode(code);
+        }
+    }
+
+    void specializeTargetSwitch(TargetRequest* target, IRModule* module)
+    {
+        for (auto globalInst : module->getGlobalInsts())
+        {
+            if (auto code = as<IRGlobalValueWithCode>(globalInst))
+            {
+                specializeTargetSwitch(target, code);
+                if (auto gen = as<IRGeneric>(code))
+                {
+                    auto retVal = findGenericReturnVal(gen);
+                    if (auto innerCode = as<IRGlobalValueWithCode>(retVal))
+                    {
+                        specializeTargetSwitch(target, innerCode);
+                    }
+                }
+            }
+        }
+    }
+
+}
-- 
cgit v1.2.3