Don't rewrite entrypoint in to borrow during ir lowering. (#9869)

If the user compile a shader via `findAndCheckEntrypoint` API on a
function that does not have a `[shader()]` attribute, we could be
generating invalid spirv code.

There is logic in lower-to-ir that silently changes the parameter
passing convention of entrypoint in to borrow in, which is fundamentally
conflicting with the compilation model. When we load a module, we
compile every function down to IR. If we want to lower entrypoint
differently, we must know some func is entrypoint during lower-to-ir of
a module. But since this info is not provided at `loadModule` time, we
fundamentally cannot do this kind of transform when lowering a module to
IR.

The right thing to do is to implement an IR pass that turns these in
parameters into borrow in once we know a function is entrypoint. This PR
deletes the special logic in lower-to-ir that does the problematic
rewrite, and use a dedicated pass derived from the existing
`transformParamsToConstRef` pass to do the rewrite after linking, where
entrypoints are always specified.

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: csyonghe

source/slang/slang-emit.cpp

@@ -746,6 +746,8 @@ Result linkAndOptimizeIR(
     if (!isKhronosTarget(targetRequest) && requiredLoweringPassSet.glslSSBO)
         SLANG_PASS(lowerGLSLShaderStorageBufferObjectsToStructuredBuffers, sink);
 
+    SLANG_PASS(translateEntryPointInParamToBorrow, sink);
+
     if (requiredLoweringPassSet.globalVaryingVar)
         SLANG_PASS(translateGlobalVaryingVar, codeGenContext);
 

source/slang/slang-ir-transform-params-to-constref.cpp

@@ -20,7 +20,7 @@ struct TransformParamsToConstRefContext
     }
 
     // Check if a type should be transformed (struct, array, or other composite types)
-    bool shouldTransformParam(IRParam* param)
+    virtual bool shouldTransformParam(IRParam* param)
     {
         auto type = param->getDataType();
         if (!type)
@@ -97,6 +97,7 @@ struct TransformParamsToConstRefContext
                                 auto fieldExtract = as<IRFieldExtract>(userInst);
                                 builder.setInsertBefore(fieldExtract);
                                 auto fieldAddr = builder.emitFieldAddress(
+                                    builder.getPtrType(userInst->getDataType()),
                                     loadInst->getPtr(),
                                     fieldExtract->getField());
                                 auto loadFieldAddr = as<IRLoad>(builder.emitLoad(fieldAddr));
@@ -120,6 +121,7 @@ struct TransformParamsToConstRefContext
                                 auto getElement = as<IRGetElement>(userInst);
                                 builder.setInsertBefore(getElement);
                                 auto getElementPtr = builder.emitElementAddress(
+                                    builder.getPtrType(userInst->getDataType()),
                                     loadInst->getPtr(),
                                     getElement->getIndex());
                                 auto loadElementPtr = as<IRLoad>(builder.emitLoad(getElementPtr));
@@ -253,7 +255,7 @@ struct TransformParamsToConstRefContext
     }
 
     // Check if function should be excluded from transformation
-    bool shouldProcessFunction(IRFunc* func)
+    virtual bool shouldProcessFunction(IRFunc* func)
     {
         // Skip functions without definitions
         if (!func->isDefinition())
@@ -421,4 +423,58 @@ SlangResult transformParamsToConstRef(IRModule* module, DiagnosticSink* sink)
     return context.processModule();
 }
 
+struct EntryPointInParamToBorrowContext : public TransformParamsToConstRefContext
+{
+    EntryPointInParamToBorrowContext(IRModule* module, DiagnosticSink* sink)
+        : TransformParamsToConstRefContext(module, sink)
+    {
+    }
+    virtual bool shouldProcessFunction(IRFunc* func) override
+    {
+        if (!func->isDefinition())
+            return false;
+        if (func->findDecoration<IREntryPointDecoration>() != nullptr)
+            return true;
+        return false;
+    }
+    virtual bool shouldTransformParam(IRParam* param) override
+    {
+        auto type = param->getDataType();
+        if (as<IRPointerLikeType>(type))
+            return false;
+        if (as<IRPtrTypeBase>(type))
+            return false;
+        if (as<IRMeshOutputType>(type))
+            return false;
+        if (as<IRHLSLPatchType>(type))
+            return false;
+
+        // Skip uniform parameters.
+        // We expect all entry-point parameters to have layout information,
+        // but we will be defensive and skip parameters without the required
+        // information when we are in a release build.
+        //
+        auto layoutDecoration = param->findDecoration<IRLayoutDecoration>();
+        SLANG_ASSERT(layoutDecoration);
+        if (!layoutDecoration)
+            return false;
+        auto paramLayout = as<IRVarLayout>(layoutDecoration->getLayout());
+        SLANG_ASSERT(paramLayout);
+        if (!paramLayout)
+            return false;
+        if (!isVaryingParameter(paramLayout))
+            return false;
+
+        // If we reach here, we are dealing with a varying in parameter.
+        // We need to rewrite it to be a `borrow in` parameter.
+        return true;
+    }
+};
+
+SlangResult translateEntryPointInParamToBorrow(IRModule* module, DiagnosticSink* sink)
+{
+    EntryPointInParamToBorrowContext context(module, sink);
+    return context.processModule();
+}
+
 } // namespace Slang

source/slang/slang-ir-transform-params-to-constref.h

@@ -9,4 +9,6 @@ class DiagnosticSink;
 
 SlangResult transformParamsToConstRef(IRModule* module, DiagnosticSink* sink);
 
+SlangResult translateEntryPointInParamToBorrow(IRModule* module, DiagnosticSink* sink);
+
 } // namespace Slang

source/slang/slang-lower-to-ir.cpp

@@ -3552,104 +3552,6 @@ void maybeAddReturnDestinationParam(ParameterLists* ioParameterLists, Type* resu
     }
 }
 
-/// Does the given `declRef` appear to be a declaration of an entry point?
-///
-/// This function does a best-effort job of detecting whether something is
-/// a shader entry point, but it cannot answer the question definitively,
-/// because the compilation model of Slang allows functions to be tagged
-/// as entry points separately from their declaration in the AST.
-///
-bool doesDeclAppearToBeAnEntryPoint(DeclRef<CallableDecl> const& declRef)
-{
-    auto decl = declRef.getDecl();
-    if (decl->hasModifier<EntryPointAttribute>())
-        return true;
-    if (decl->hasModifier<NumThreadsAttribute>())
-        return true;
-    return false;
-}
-
-/// Does a parameter appear to be a declaration of an entry-point varying input.
-///
-/// The `paramInfo` represents the parameter being inspected, and
-/// the `funcDeclRef` should represent the outer declaration that
-/// might or might not be an entry point.
-///
-/// This function is only able to do a best-effort check for what is an
-/// entry point (see `doesDeclAppearToBeAnEntryPoint()`), and cannot be
-/// fully robust in detection.
-///
-bool doesParamAppearToBeAnEntryPointVaryingInput(
-    IRLoweringParameterInfo const& paramInfo,
-    DeclRef<CallableDecl> const& funcDeclRef)
-{
-    // If the outer declaration doesn't appear to be an entry
-    // point, then it seems this isn't an entry point parameter
-    // at all, much less an input.
-    //
-    if (!doesDeclAppearToBeAnEntryPoint(funcDeclRef))
-        return false;
-
-    // We are only intereste in parameters that would otherwise
-    // be lowered to just use `in`.
-    //
-    if (paramInfo.actualParamPassingModeToUse != ParamPassingMode::In)
-        return false;
-
-    // We are only concerned with varying parameters, so `uniform`
-    // parameters aren't relevant.
-    //
-    if (paramInfo.decl->findModifier<HLSLUniformModifier>())
-        return false;
-
-    // Certain types conceptually represent uniform parameters even
-    // if they aren't declared with `uniform`, so we also want to
-    // ignore those.
-    //
-    // TODO: It would be best for logic like this to be factored into
-    // a shared subroutine somewhere, so that we don't have multiple
-    // places in the code doing ad hoc checks for a particular list
-    // of types, that might change over time.
-    //
-    if (as<HLSLPatchType>(paramInfo.type))
-        return false;
-
-    if (as<MeshOutputType>(paramInfo.type))
-        return false;
-
-    return true;
-}
-
-/// Modify the parameter passing mode for the given parameter, if it can
-/// be detected that it seems to be used as a varying input to an entry point.
-///
-/// If a varying `in` parameter is detected, its actual parameter-passing
-/// mode will be changed to `borrow in`.
-///
-/// This function is only able to do a best-effort check for what is an
-/// entry point (see `doesDeclAppearToBeAnEntryPoint()`), and cannot be
-/// fully robust in detection.
-///
-void maybeModifyParamPassingModeForDetectedEntryPointVaryingInput(
-    IRLoweringParameterInfo& ioParamInfo,
-    DeclRef<CallableDecl> const& funcDeclRef)
-{
-    // If we cannot detect that this seems to be a varying `in`
-    // parameter of an entry point, then we'll just skip it.
-    //
-    // There's no way for this code to be sure it isn't skipping
-    // over a function it shouldn't, but there's nothign we can
-    // do about it from here.
-    //
-    if (!doesParamAppearToBeAnEntryPointVaryingInput(ioParamInfo, funcDeclRef))
-        return;
-
-    // We basically just want to change the parameter from `in`
-    // to `borrow in`, so that it is an immutable by-reference parameter.
-    //
-    ioParamInfo.actualParamPassingModeToUse = ParamPassingMode::BorrowIn;
-}
-
 //
 // And here is our function that will do the recursive walk:
 void collectParameterLists(
@@ -3894,20 +3796,6 @@ void collectParameterLists(
     {
         auto paramInfo = getParameterInfo(context, paramDeclRef);
 
-        // One unfortunate wrinkle that arises is that all the downstream
-        // logic in the Slang IR *really* wants the varying input parameters
-        // to a shader entry point to use `ParamPassingMode::BorrowIn`,
-        // so that they don't force copying, but syntactically such parameters
-        // are conventionally declared using `in` (meaning `ParamPassingMode::In`).
-        //
-        // We include logic here to switch up the parameter-passing mode
-        // in the case where we can statically detect that something is
-        // being compiled as an entry point. Note, however, that this logic
-        // is inherently fragile, since not every entry point that a user specifies
-        // is guaranteed to have had a `[shader(...)]` attribute on it.
-        //
-        maybeModifyParamPassingModeForDetectedEntryPointVaryingInput(paramInfo, callableDeclRef);
-
         ioParameterLists->params.add(paramInfo);
     }
 

tests/spirv/debug-return-types.slang

@@ -18,7 +18,7 @@ float4 main(VSOutput input) : SV_TARGET {
 }
 
 // CHECK: %[[VECTOR:[0-9]+]] = OpExtInst %void %{{[a-zA-Z0-9_]+}} DebugTypeVector
-// CHECK: %[[VSOUTPUT:[0-9]+]] = OpExtInst %void %{{[a-zA-Z0-9_]+}} DebugTypePointer
+// CHECK: %[[VSOUTPUT:[0-9]+]] = OpExtInst %void %{{[a-zA-Z0-9_]+}} DebugTypeComposite
 // CHECK: {{.*}} = OpExtInst %void %{{[a-zA-Z0-9_]+}} DebugTypeFunction %{{[a-zA-Z0-9_]+}} %[[VECTOR]] %[[VSOUTPUT]]
 // CHECK: %[[MATRIX:[0-9]+]] = OpExtInst %void %{{[a-zA-Z0-9_]+}} DebugTypeMatrix
 // CHECK: {{.*}} = OpExtInst %void %{{[a-zA-Z0-9_]+}} DebugTypeFunction %{{[a-zA-Z0-9_]+}} %[[MATRIX]] %[[MATRIX]]

tools/slang-unit-test/unit-test-entrypoint-compile.cpp

@@ -0,0 +1,86 @@
+// unit-test-entrypoint-compile.cpp
+
+#include "../../source/core/slang-io.h"
+#include "../../source/core/slang-process.h"
+#include "slang-com-ptr.h"
+#include "slang.h"
+#include "unit-test/slang-unit-test.h"
+
+#include <stdio.h>
+#include <stdlib.h>
+
+using namespace Slang;
+
+SLANG_UNIT_TEST(entryPointCompile)
+{
+    ComPtr<slang::IGlobalSession> globalSession;
+    SLANG_CHECK_ABORT(
+        slang_createGlobalSession(SLANG_API_VERSION, globalSession.writeRef()) == SLANG_OK);
+    slang::TargetDesc targetDesc = {};
+    // Request SPIR-V disassembly so we can check the content.
+    targetDesc.format = SLANG_SPIRV_ASM;
+    targetDesc.profile = globalSession->findProfile("spirv_1_5");
+    slang::SessionDesc sessionDesc = {};
+    sessionDesc.targetCount = 1;
+    sessionDesc.targets = &targetDesc;
+
+    List<slang::CompilerOptionEntry> optionEntries;
+    {
+        slang::CompilerOptionEntry entry;
+        entry.name = slang::CompilerOptionName::EnableEffectAnnotations;
+        entry.value.kind = slang::CompilerOptionValueKind::Int;
+        entry.value.intValue0 = 1;
+        optionEntries.add(entry);
+    }
+
+    sessionDesc.compilerOptionEntries = optionEntries.getBuffer();
+    sessionDesc.compilerOptionEntryCount = (uint32_t)optionEntries.getCount();
+
+    ComPtr<slang::ISession> session;
+    SLANG_CHECK_ABORT(globalSession->createSession(sessionDesc, session.writeRef()) == SLANG_OK);
+
+    ComPtr<slang::IBlob> diagnosticBlob;
+    String userSourceBody = R"(
+        struct VS_INPUT 
+        { 
+            float3 vPositionOs : POSITION; 
+            uint nInstanceIdx : TEXCOORD13; 
+            uint instanceId : SV_InstanceID; 
+        }
+
+        struct PS_INPUT 
+        {
+            float4 vPositionPs : SV_Position ; 
+        }
+
+        StructuredBuffer<float3> gBuffer;
+
+        PS_INPUT vsMain(VS_INPUT i) 
+        { 
+            PS_INPUT o = {}; 
+            float3 v = gBuffer[i.instanceId].xyz ; 
+            o.vPositionPs = float4(v, 1.0f);
+            return o ; 
+        })";
+    auto srcBlob = StringBlob::moveCreate(_Move(userSourceBody));
+    auto module = session->loadModuleFromSource("m", "m.slang", srcBlob, diagnosticBlob.writeRef());
+    SLANG_CHECK_ABORT(module != nullptr);
+
+    ComPtr<slang::IEntryPoint> entryPoint;
+    module->findAndCheckEntryPoint(
+        "vsMain",
+        SLANG_STAGE_VERTEX,
+        entryPoint.writeRef(),
+        diagnosticBlob.writeRef());
+    SLANG_CHECK_ABORT(entryPoint != nullptr);
+
+    ComPtr<slang::IComponentType> linkedProgram;
+    entryPoint->link(linkedProgram.writeRef(), diagnosticBlob.writeRef());
+    SLANG_CHECK_ABORT(linkedProgram != nullptr);
+
+    ComPtr<slang::IBlob> code;
+    linkedProgram->getEntryPointCode(0, 0, code.writeRef(), diagnosticBlob.writeRef());
+    SLANG_CHECK_ABORT(code != nullptr);
+
+    SLANG_CHECK_ABORT(code->getBufferSize() != 0);
+}