Add more wave intrinsics. (#3162)

Co-authored-by: Yong He <yhe@nvidia.com>

1 files changed, 31 insertions, 21 deletions

diff --git a/source/slang/slang-ir-inline.cpp b/source/slang/slang-ir-inline.cpp
index f3fa213da..e171e2dd3 100644
--- a/source/slang/slang-ir-inline.cpp
+++ b/source/slang/slang-ir-inline.cpp
@@ -173,6 +173,19 @@ struct InliningPassBase
         return false;
     }
 
+    static bool hasGenericAsmInst(IRInst* func)
+    {
+        auto f = as<IRFunc>(getResolvedInstForDecorations(func));
+        if (!f)
+            return false;
+        for (auto b : f->getBlocks())
+        {
+            if (as<IRGenericAsm>(b->getTerminator()))
+                return true;
+        }
+        return false;
+    }
+
         /// Determine whether `call` can be inlined, and if so write information about it to `outCallSite`
     bool canInline(IRCall* call, CallSiteInfo& outCallSite)
     {
@@ -236,6 +249,24 @@ struct InliningPassBase
         if (callee->findDecoration<IRIntrinsicOpDecoration>())
             return true;
 
+        // We cannot inline a function that is defined by a generic asm inst.
+        if (hasGenericAsmInst(callee))
+            return false;
+
+        for (auto decor : callee->getDecorations())
+        {
+            switch (decor->getOp())
+            {
+            case kIROp_IntrinsicOpDecoration:
+                return true;
+            case kIROp_RequireSPIRVCapabilityDecoration:
+            case kIROp_RequireSPIRVVersionDecoration:
+            case kIROp_RequireGLSLExtensionDecoration:
+            case kIROp_RequireGLSLVersionDecoration:
+                return false;
+            }
+        }
+
         // At this point the `CallSiteInfo` is complete and
         // could be used for inlining, but we have additional
         // checks to make.
@@ -654,19 +685,6 @@ struct InliningPassBase
 
 };
 
-static bool hasGenericAsmInst(IRInst* func)
-{
-    auto f = as<IRFunc>(getResolvedInstForDecorations(func));
-    if (!f)
-        return false;
-    for (auto b : f->getBlocks())
-    {
-        if (as<IRGenericAsm>(b->getTerminator()))
-            return true;
-    }
-    return false;
-}
-
     /// An inlining pass that inlines calls to `[unsafeForceInlineEarly]` functions
 struct MandatoryEarlyInliningPass : InliningPassBase
 {
@@ -681,10 +699,6 @@ struct MandatoryEarlyInliningPass : InliningPassBase
         if (info.callee->findDecoration<IRIntrinsicOpDecoration>())
             return true;
 
-        // Never inline a callee that has genericASM instruction.
-        if (hasGenericAsmInst(info.callee))
-            return false;
-
         if(info.callee->findDecoration<IRUnsafeForceInlineEarlyDecoration>())
             return true;
         return false;
@@ -800,10 +814,6 @@ struct ForceInliningPass : InliningPassBase
 
     bool shouldInline(CallSiteInfo const& info)
     {
-        // Never inline a callee that has genericASM instruction.
-        if (hasGenericAsmInst(info.callee))
-            return false;
-
         if (info.callee->findDecoration<IRForceInlineDecoration>() ||
             info.callee->findDecoration<IRUnsafeForceInlineEarlyDecoration>()||
             info.callee->findDecoration<IRIntrinsicOpDecoration>())