Add API for whole program compilation. (#1562)

* Add API for whole program compilation.

This change exposes a new target flag: `SLANG_TARGET_FLAG_GENERATE_WHOLE_PROGRAM` that can be set on a target with `spSetTargetFlags`. When this flag is set, `spCompile` function generates target code for the entire input module instead of just the specified entrypoints. The resulting code will include all the entrypoints defined in the input source.

The resulting whole program code can be retrieved with two new functions: `spGetTargetCodeBlob` and `spGetTargetHostCallable`.

This change also cleans up the unnecessary `entryPointIndices` parameter of `TargetProgram::getOrCreateWholeProgramResult`, and modifies the `cpu-hello-world` example to make use of the new whole-program compilation API to simplify its logic.

* Update comments.

7 files changed, 122 insertions, 37 deletions

diff --git a/examples/cpu-hello-world/main.cpp b/examples/cpu-hello-world/main.cpp
index ee919a713..cf8c57285 100644
--- a/examples/cpu-hello-world/main.cpp
+++ b/examples/cpu-hello-world/main.cpp
@@ -85,6 +85,10 @@ static SlangResult _innerMain(int argc, char** argv)
     // If we wanted a just a shared library/dll, we could have used SLANG_SHARED_LIBRARY.
     int targetIndex = spAddCodeGenTarget(slangRequest, SLANG_HOST_CALLABLE);
 
+    // Set the target flag to indicate that we want to compile all the entrypoints in the
+    // slang shader file into a library.
+    spSetTargetFlags(slangRequest, targetIndex, SLANG_TARGET_FLAG_GENERATE_WHOLE_PROGRAM);
+
     // A compile request can include one or more "translation units," which more or
     // less amount to individual source files (think `.c` files, not the `.h` files they
     // might include).
@@ -99,14 +103,6 @@ static SlangResult _innerMain(int argc, char** argv)
     // There are also variations of this API for adding source code from application-provided buffers.
     //
     spAddTranslationUnitSourceFile(slangRequest, translationUnitIndex, "shader.slang");
-    
-    // Next we will specify the entry points we'd like to compile.
-    // It is often convenient to put more than one entry point in the same file,
-    // and the Slang API makes it convenient to use a single run of the compiler
-    // to compile all entry points. 
-    // 
-    const char entryPointName[] = "computeMain";
-    int computeIndex = spAddEntryPoint(slangRequest, translationUnitIndex, entryPointName,  SLANG_STAGE_COMPUTE);
 
     // Once all of the input options for the compiler have been specified,
     // we can invoke `spCompile` to run the compiler and see if any errors
@@ -133,7 +129,7 @@ static SlangResult _innerMain(int argc, char** argv)
     // Get the 'shared library' (note that this doesn't necessarily have to be implemented as a shared library
     // it's just an interface to executable code).
     ComPtr<ISlangSharedLibrary> sharedLibrary;
-    SLANG_RETURN_ON_FAIL(spGetEntryPointHostCallable(slangRequest, 0, 0, sharedLibrary.writeRef()));
+    SLANG_RETURN_ON_FAIL(spGetTargetHostCallable(slangRequest, 0, sharedLibrary.writeRef()));
 
     // Once we have the sharedLibrary, we no longer need the request
     // unless we want to use reflection, to for example workout how 'UniformState' and 'UniformEntryPointParams' are laid out
@@ -141,10 +137,10 @@ static SlangResult _innerMain(int argc, char** argv)
     spDestroyCompileRequest(slangRequest);
 
     // Get the function we are going to execute 
+    const char entryPointName[] = "computeMain";
     CPPPrelude::ComputeFunc func = (CPPPrelude::ComputeFunc)sharedLibrary->findFuncByName(entryPointName);
     if (!func)
     {
-        spDestroyCompileRequest(slangRequest);
         return SLANG_FAIL;
     }
 
diff --git a/examples/cpu-hello-world/shader.slang b/examples/cpu-hello-world/shader.slang
index 6611962d7..f650c3481 100644
--- a/examples/cpu-hello-world/shader.slang
+++ b/examples/cpu-hello-world/shader.slang
@@ -3,6 +3,7 @@
 //TEST_INPUT:ubuffer(random(float, 4096, -1.0, 1.0), stride=4):name=ioBuffer
 RWStructuredBuffer<float> ioBuffer;
 
+[shader("compute")]
 [numthreads(4, 1, 1)]
 void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
 {
diff --git a/slang.h b/slang.h
index 60e7c43a7..acbee3982 100644
--- a/slang.h
+++ b/slang.h
@@ -584,6 +584,12 @@ extern "C"
            @deprecated This behavior is now enabled unconditionally.
         */
         SLANG_TARGET_FLAG_PARAMETER_BLOCKS_USE_REGISTER_SPACES = 1 << 4,
+
+        /* When set, will generate target code that contains all entrypoints defined
+           in the input source or specified via the `spAddEntryPoint` function in a
+           single output module (library/source file).
+        */
+        SLANG_TARGET_FLAG_GENERATE_WHOLE_PROGRAM = 1 << 8
     };
 
     /*!
@@ -1632,8 +1638,6 @@ extern "C"
     @param targetIndex The index of the target to get code for (default: zero).
     @param outBlob A pointer that will receive the blob of code
     @returns A `SlangResult` to indicate success or failure.
-
-    The lifetime of the output pointer is the same as `request`.
     */
     SLANG_API SlangResult spGetEntryPointCodeBlob(
         SlangCompileRequest*    request,
@@ -1659,6 +1663,34 @@ extern "C"
         int                     targetIndex,
         ISlangSharedLibrary**   outSharedLibrary);
 
+    /** Get the output code associated with a specific target.
+
+    @param request   The request 
+    @param targetIndex The index of the target to get code for (default: zero).
+    @param outBlob A pointer that will receive the blob of code
+    @returns A `SlangResult` to indicate success or failure.
+    */
+    SLANG_API SlangResult spGetTargetCodeBlob(
+        SlangCompileRequest*    request,
+        int                     targetIndex,
+        ISlangBlob**            outBlob);
+
+    /** Get 'callable' functions for a target accessible through the ISlangSharedLibrary interface.
+
+    That the functions remain in scope as long as the ISlangSharedLibrary interface is in scope.
+
+    NOTE! Requires a compilation target of SLANG_HOST_CALLABLE.
+    
+    @param request          The request 
+    @param targetIndex      The index of the target to get code for (default: zero).
+    @param outSharedLibrary A pointer to a ISharedLibrary interface which functions can be queried on.
+    @returns                A `SlangResult` to indicate success or failure.
+    */
+    SLANG_API SlangResult spGetTargetHostCallable(
+        SlangCompileRequest*    request,
+        int                     targetIndex,
+        ISlangSharedLibrary**   outSharedLibrary);
+
     /** Get the output bytecode associated with an entire compile request.
 
     The lifetime of the output pointer is the same as `request` and the last spCompile.
diff --git a/source/slang/slang-compiler.cpp b/source/slang/slang-compiler.cpp
index 3855eda68..ea8540800 100755
--- a/source/slang/slang-compiler.cpp
+++ b/source/slang/slang-compiler.cpp
@@ -2197,24 +2197,19 @@ SlangResult dissassembleDXILUsingDXC(
     }
 
     CompileResult& TargetProgram::_createWholeProgramResult(
-        const List<Int>&        entryPointIndices,
         BackEndCompileRequest*  backEndRequest,
         EndToEndCompileRequest* endToEndRequest)
     {
-        for (auto entryPointIndex = entryPointIndices.begin(); entryPointIndex != entryPointIndices.end(); entryPointIndex++) {
-            if (*entryPointIndex >= m_entryPointResults.getCount())
-                m_entryPointResults.setCount(*entryPointIndex + 1);
+        // We want to call `emitEntryPoints` function to generate code that contains
+        // all the entrypoints defined in `m_program`.
+        // The current logic of `emitEntryPoints` takes a list of entry-point indices to
+        // emit code for, so we construct such a list first.
+        List<Int> entryPointIndices;
+        m_entryPointResults.setCount(m_program->getEntryPointCount());
+        entryPointIndices.setCount(m_program->getEntryPointCount());
+        for (Index i = 0; i < entryPointIndices.getCount(); i++)
+            entryPointIndices[i] = i;
 
-            // It is possible that entry points goot added to the `Program`
-            // *after* we created this `TargetProgram`, so there might be
-            // a request for an entry point that we didn't allocate space for.
-            //
-            // TODO: Change the construction logic so that a `Program` is
-            // constructed all at once rather than incrementally, to avoid
-            // this problem.
-            //
-            //auto entryPoint = m_program->getEntryPoint(*entryPointIndex);
-        }
         auto& result = m_wholeProgramResult;
         result = emitEntryPoints(
             m_program,
@@ -2224,7 +2219,6 @@ SlangResult dissassembleDXILUsingDXC(
             endToEndRequest);
 
         return result;
-
     }
 
     CompileResult& TargetProgram::_createEntryPointResult(
@@ -2256,7 +2250,6 @@ SlangResult dissassembleDXILUsingDXC(
     }
 
     CompileResult& TargetProgram::getOrCreateWholeProgramResult(
-        const List<Int>& entryPointIndices,
         DiagnosticSink* sink)
     {
         auto& result = m_wholeProgramResult;
@@ -2278,7 +2271,6 @@ SlangResult dissassembleDXILUsingDXC(
             m_program);
 
         return _createWholeProgramResult(
-            entryPointIndices,
             backEndRequest,
             nullptr);
     }
@@ -2325,10 +2317,9 @@ SlangResult dissassembleDXILUsingDXC(
         // Generate target code any entry points that
         // have been requested for compilation.
         auto entryPointCount = program->getEntryPointCount();
-        if (targetReq->isWholeProgramRequest)
+        if (targetReq->isWholeProgramRequest())
         {
             targetProgram->_createWholeProgramResult(
-                List<Int>(),
                 compileReq,
                 endToEndReq);
         }
@@ -2497,7 +2488,7 @@ SlangResult dissassembleDXILUsingDXC(
             for (auto targetReq : linkage->targets)
             {
                 Index entryPointCount = program->getEntryPointCount();
-                if (targetReq->isWholeProgramRequest) {
+                if (targetReq->isWholeProgramRequest()) {
                     writeWholeProgramResult(
                         compileRequest,
                         targetReq);
diff --git a/source/slang/slang-compiler.h b/source/slang/slang-compiler.h
index f7a5d98b1..e120d5849 100755
--- a/source/slang/slang-compiler.h
+++ b/source/slang/slang-compiler.h
@@ -1132,7 +1132,10 @@ namespace Slang
         SlangTargetFlags    targetFlags = 0;
         Slang::Profile      targetProfile = Slang::Profile();
         FloatingPointMode   floatingPointMode = FloatingPointMode::Default;
-        bool                isWholeProgramRequest = false;
+        bool isWholeProgramRequest()
+        {
+            return (targetFlags & SLANG_TARGET_FLAG_GENERATE_WHOLE_PROGRAM) != 0;
+        }
 
         Linkage* getLinkage() { return linkage; }
         CodeGenTarget getTarget() { return target; }
@@ -1673,7 +1676,7 @@ namespace Slang
             /// code generation to the given `sink`.
             ///
         CompileResult& getOrCreateEntryPointResult(Int entryPointIndex, DiagnosticSink* sink);
-        CompileResult& getOrCreateWholeProgramResult(const List<Int>& entryPointIndices, DiagnosticSink* sink);
+        CompileResult& getOrCreateWholeProgramResult(DiagnosticSink* sink);
 
 
         CompileResult& getExistingWholeProgramResult()
@@ -1691,7 +1694,6 @@ namespace Slang
         }
 
         CompileResult& _createWholeProgramResult(
-            const List<Int>&        entryPointIndices,
             BackEndCompileRequest*  backEndRequest,
             EndToEndCompileRequest* endToEndRequest);
             /// Internal helper for `getOrCreateEntryPointResult`.
diff --git a/source/slang/slang-options.cpp b/source/slang/slang-options.cpp
index 42003fb62..db56e8115 100644
--- a/source/slang/slang-options.cpp
+++ b/source/slang/slang-options.cpp
@@ -160,7 +160,6 @@ struct OptionsParser
         SlangTargetFlags    targetFlags = 0;
         int                 targetID = -1;
         FloatingPointMode   floatingPointMode = FloatingPointMode::Default;
-        bool                isWholeProgramRequest = false;
 
         // State for tracking command-line errors
         bool conflictingProfilesSet = false;
@@ -1511,7 +1510,7 @@ struct OptionsParser
                 }
                 else
                 {
-                    target->isWholeProgramRequest = true;
+                    target->targetFlags |= SLANG_TARGET_FLAG_GENERATE_WHOLE_PROGRAM;
                     targetInfo->wholeTargetOutputPath = rawOutput.path;
                 }
             }
diff --git a/source/slang/slang.cpp b/source/slang/slang.cpp
index 572c9c3a3..376376f24 100644
--- a/source/slang/slang.cpp
+++ b/source/slang/slang.cpp
@@ -3613,6 +3613,33 @@ static SlangResult _getEntryPointResult(
     return SLANG_OK;
 }
 
+static SlangResult _getWholeProgramResult(
+    SlangCompileRequest* request,
+    int targetIndex,
+    Slang::CompileResult** outCompileResult)
+{
+    using namespace Slang;
+    if (!request)
+        return SLANG_ERROR_INVALID_PARAMETER;
+
+    auto req = Slang::asInternal(request);
+    auto linkage = req->getLinkage();
+    auto program = req->getSpecializedGlobalAndEntryPointsComponentType();
+
+    Index targetCount = linkage->targets.getCount();
+    if ((targetIndex < 0) || (targetIndex >= targetCount))
+    {
+        return SLANG_ERROR_INVALID_PARAMETER;
+    }
+    auto targetReq = linkage->targets[targetIndex];
+
+    auto targetProgram = program->getTargetProgram(targetReq);
+    if (!targetProgram)
+        return SLANG_FAIL;
+    *outCompileResult = &targetProgram->getExistingWholeProgramResult();
+    return SLANG_OK;
+}
+
 SLANG_API SlangResult spGetEntryPointCodeBlob(
         SlangCompileRequest*    request,
         int                     entryPointIndex,
@@ -3648,6 +3675,43 @@ SLANG_API SlangResult spGetEntryPointHostCallable(
     return SLANG_OK;
 }
 
+SLANG_API SlangResult spGetTargetCodeBlob(
+    SlangCompileRequest* request,
+    int targetIndex,
+    ISlangBlob** outBlob)
+{
+    using namespace Slang;
+    if (!outBlob)
+        return SLANG_ERROR_INVALID_PARAMETER;
+    Slang::CompileResult* compileResult = nullptr;
+    SLANG_RETURN_ON_FAIL(
+        _getWholeProgramResult(request, targetIndex, &compileResult));
+
+    ComPtr<ISlangBlob> blob;
+    SLANG_RETURN_ON_FAIL(compileResult->getBlob(blob));
+    *outBlob = blob.detach();
+    return SLANG_OK;
+}
+
+SLANG_API SlangResult spGetTargetHostCallable(
+    SlangCompileRequest* request,
+    int targetIndex,
+    ISlangSharedLibrary** outSharedLibrary)
+{
+    using namespace Slang;
+    if (!outSharedLibrary)
+        return SLANG_ERROR_INVALID_PARAMETER;
+
+    Slang::CompileResult* compileResult = nullptr;
+    SLANG_RETURN_ON_FAIL(
+        _getWholeProgramResult(request, targetIndex, &compileResult));
+
+    ComPtr<ISlangSharedLibrary> sharedLibrary;
+    SLANG_RETURN_ON_FAIL(compileResult->getSharedLibrary(sharedLibrary));
+    *outSharedLibrary = sharedLibrary.detach();
+    return SLANG_OK;
+}
+
 SLANG_API char const* spGetEntryPointSource(
     SlangCompileRequest*    request,
     int                     entryPointIndex)