Add `requirePrelude()` intrinsic function. (#3250)

* Add `requirePrelude()` intrinsic function.

* Fix.

---------

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

10 files changed, 78 insertions, 17 deletions

diff --git a/source/slang/core.meta.slang b/source/slang/core.meta.slang
index 87697076e..3e51d7028 100644
--- a/source/slang/core.meta.slang
+++ b/source/slang/core.meta.slang
@@ -136,6 +136,8 @@ attribute_syntax [BackwardDifferentiable(order:int = 0)] : BackwardDifferentiabl
 __attributeTarget(FunctionDeclBase)
 attribute_syntax [Differentiable(order:int = 0)] : BackwardDifferentiableAttribute;
 
+__intrinsic_op($(kIROp_RequirePrelude))
+void __requirePrelude(constexpr String preludeText);
 
 /// Interface to denote types as differentiable.
 /// Allows for user-specified differential types as
diff --git a/source/slang/slang-emit-c-like.cpp b/source/slang/slang-emit-c-like.cpp
index fb216ae01..a2c3abb5e 100644
--- a/source/slang/slang-emit-c-like.cpp
+++ b/source/slang/slang-emit-c-like.cpp
@@ -112,6 +112,15 @@ void CLikeSourceEmitter::emitFrontMatterImpl(TargetRequest* targetReq)
     SLANG_UNUSED(targetReq);
 }
 
+void CLikeSourceEmitter::emitPreModuleImpl()
+{
+    for (auto prelude : m_requiredPreludes)
+    {
+        m_writer->emit(prelude->getStringSlice());
+        m_writer->emit("\n");
+    }
+}
+
 //
 // Types
 //
@@ -1856,6 +1865,22 @@ void CLikeSourceEmitter::emitCallExpr(IRCall* inst, EmitOpInfo outerPrec)
     UnownedStringSlice intrinsicDefinition;
     if (findTargetIntrinsicDefinition(funcValue, intrinsicDefinition))
     {
+        // Make sure we register all required preludes for emit.
+        if (auto func = as<IRFunc>(getResolvedInstForDecorations(funcValue)))
+        {
+            for (auto block : func->getBlocks())
+            {
+                for (auto ii : block->getChildren())
+                {
+                    if (auto requirePrelude = as<IRRequirePrelude>(ii))
+                    {
+                        auto preludeTextInst = as<IRStringLit>(requirePrelude->getOperand(0));
+                        if (preludeTextInst)
+                            m_requiredPreludes.add(preludeTextInst);
+                    }
+                }
+            }
+        }
         emitIntrinsicCallExpr(inst, intrinsicDefinition, outerPrec);
     }
     else
diff --git a/source/slang/slang-emit-c-like.h b/source/slang/slang-emit-c-like.h
index ffb4c0244..01b5681b7 100644
--- a/source/slang/slang-emit-c-like.h
+++ b/source/slang/slang-emit-c-like.h
@@ -465,7 +465,7 @@ public:
         /// Emit any declarations, and other material that is needed before the modules contents
         /// For example on targets that don't have built in vector/matrix support, this is where
         /// the appropriate generated declarations occur.
-    virtual void emitPreModuleImpl() {}
+    virtual void emitPreModuleImpl();
 
     virtual void emitRateQualifiersAndAddressSpaceImpl(IRRate* rate, IRIntegerValue addressSpace) { SLANG_UNUSED(rate); SLANG_UNUSED(addressSpace); }
     virtual void emitSemanticsImpl(IRInst* inst, bool allowOffsetLayout) { SLANG_UNUSED(inst); SLANG_UNUSED(allowOffsetLayout); }
@@ -570,6 +570,8 @@ public:
     // Map an IR instruction to the name that we've decided
     // to use for it when emitting code.
     Dictionary<IRInst*, String> m_mapInstToName;
+
+    OrderedHashSet<IRStringLit*> m_requiredPreludes;
 };
 
 }
diff --git a/source/slang/slang-emit-cpp.cpp b/source/slang/slang-emit-cpp.cpp
index ebe5965f4..c48fa369d 100644
--- a/source/slang/slang-emit-cpp.cpp
+++ b/source/slang/slang-emit-cpp.cpp
@@ -1709,6 +1709,7 @@ void CPPSourceEmitter::emitPreModuleImpl()
         m_writer->emit("using namespace SLANG_PRELUDE_NAMESPACE;\n");
         m_writer->emit("#endif\n\n");
     }
+    Super::emitPreModuleImpl();
 }
 
 
diff --git a/source/slang/slang-emit-cuda.cpp b/source/slang/slang-emit-cuda.cpp
index 9364f4441..bd891928c 100644
--- a/source/slang/slang-emit-cuda.cpp
+++ b/source/slang/slang-emit-cuda.cpp
@@ -819,16 +819,6 @@ void CUDASourceEmitter::emitMatrixLayoutModifiersImpl(IRVarLayout* layout)
     Super::emitMatrixLayoutModifiersImpl(layout);
 }
 
-void CUDASourceEmitter::emitPreModuleImpl()
-{
-    SourceWriter* writer = getSourceWriter();
-
-    // Emit generated types/functions
-
-    writer->emit("\n");
-}
-
-
 bool CUDASourceEmitter::tryEmitGlobalParamImpl(IRGlobalParam* varDecl, IRType* varType)
 {
     // A global shader parameter in the IR for CUDA output will
diff --git a/source/slang/slang-emit-cuda.h b/source/slang/slang-emit-cuda.h
index a63d087bf..097d7b741 100644
--- a/source/slang/slang-emit-cuda.h
+++ b/source/slang/slang-emit-cuda.h
@@ -66,8 +66,6 @@ protected:
     virtual void emitLayoutSemanticsImpl(IRInst* inst, char const* uniformSemanticSpelling) SLANG_OVERRIDE;
     virtual void emitParameterGroupImpl(IRGlobalParam* varDecl, IRUniformParameterGroupType* type) SLANG_OVERRIDE;
     virtual void emitEntryPointAttributesImpl(IRFunc* irFunc, IREntryPointDecoration* entryPointDecor) SLANG_OVERRIDE;
-    
-    virtual void emitPreModuleImpl() SLANG_OVERRIDE;
 
     virtual void emitRateQualifiersAndAddressSpaceImpl(IRRate* rate, IRIntegerValue addressSpace) SLANG_OVERRIDE;
     virtual void emitSemanticsImpl(IRInst* inst, bool allowOffsetLayout) SLANG_OVERRIDE;
diff --git a/source/slang/slang-ir-autodiff-primal-hoist.cpp b/source/slang/slang-ir-autodiff-primal-hoist.cpp
index f8c4cee66..511055713 100644
--- a/source/slang/slang-ir-autodiff-primal-hoist.cpp
+++ b/source/slang/slang-ir-autodiff-primal-hoist.cpp
@@ -1558,11 +1558,8 @@ RefPtr<HoistedPrimalsInfo> ensurePrimalAvailability(
                     defBlockIndices.clear();
                 }
             }
-            if (const auto ptrInst = as<IRPtrTypeBase>(instToStore->getDataType()))
+            if (IRVar* varToStore = as<IRVar>(instToStore))
             {
-                IRVar* varToStore = as<IRVar>(instToStore);
-                SLANG_RELEASE_ASSERT(varToStore);
-
                 auto storeUse = findLatestUniqueWriteUse(varToStore);
 
                 bool isIndexedStore = (storeUse && defBlockIndices.getCount() > 0);
diff --git a/source/slang/slang-ir-inst-defs.h b/source/slang/slang-ir-inst-defs.h
index e64dfdf47..0865a9170 100644
--- a/source/slang/slang-ir-inst-defs.h
+++ b/source/slang/slang-ir-inst-defs.h
@@ -550,6 +550,8 @@ INST(TargetSwitch, targetSwitch, 1, 0)
 // A generic asm inst has an return semantics that terminates the control flow.
 INST(GenericAsm, GenericAsm, 1, 0)
 
+INST(RequirePrelude, RequirePrelude, 1, 0)
+
 INST(discard, discard, 0, 0)
 
 /* IRUnreachable */
diff --git a/source/slang/slang-ir-insts.h b/source/slang/slang-ir-insts.h
index 038681a58..e2bb53a78 100644
--- a/source/slang/slang-ir-insts.h
+++ b/source/slang/slang-ir-insts.h
@@ -3035,6 +3035,12 @@ struct IRGenericAsm : IRInst
     UnownedStringSlice getAsm() { return as<IRStringLit>(getOperand(0))->getStringSlice(); }
 };
 
+struct IRRequirePrelude : IRInst
+{
+    IR_LEAF_ISA(RequirePrelude)
+    UnownedStringSlice getPrelude() { return as<IRStringLit>(getOperand(0))->getStringSlice(); }
+};
+
 struct IRBuilderSourceLocRAII;
 
 struct IRBuilder
diff --git a/tests/cuda/require-prelude.slang b/tests/cuda/require-prelude.slang
new file mode 100644
index 000000000..a5e7fb19b
--- /dev/null
+++ b/tests/cuda/require-prelude.slang
@@ -0,0 +1,38 @@
+
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-cpu -compute -output-using-type
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-cuda -compute -output-using-type 
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-cpu -compute -output-using-type
+
+int myIntrinsicFunc()
+{
+    __target_switch
+    {
+    case cpp:
+    case hlsl:
+    case cuda:
+        __requirePrelude(R"(#define MY_CUDA_INTRINSIC 100)");
+        __intrinsic_asm "(MY_CUDA_INTRINSIC)";
+    }
+}
+
+int myIntrinsicFunc1()
+{
+    __target_switch
+    {
+    case cpp:
+    case hlsl:
+    case cuda:
+        __requirePrelude(R"(#define MY_CUDA_INTRINSIC 100)");
+        __intrinsic_asm "(MY_CUDA_INTRINSIC + 1)";
+    }
+}
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<uint> outputBuffer;
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    // CHECK: 201
+    outputBuffer[0] = myIntrinsicFunc() + myIntrinsicFunc1();
+}