Remove KernelContext wrapper from CPU/CUDA emit (#1440)

* Remove KernelContext wrapper from CPU/CUDA emit

Currently, the CPU and CUDA C++ targets rely on a `KernelContext` type that is generated during emit, as a way to provide implicit access to things that were global in the input Slang code, but that can't actually be emitted as globals in the target language (because the semantics of global declarations differ).

For example, input like:

```hlsl
ConstantBuffer<Stuff> gStuff; // shader parameter
groupshared int gData[1024];  // thread-group shared variable
static int gCounter = 0;      // "thread-local" global-scope variable

void subroutine() { ... }

[shader("compute")] void computeMain() { ... }
```

would translate to output C++ for CPU a bit like:

```c++
struct KernelContext
{
    ConstantBuffer<Stuff> gStuff;
    int gData[1024];
    int gCounter = 0;

    void subroutine() { ... }

    void computeMain() { ... }
};
```

Note that both `computeMain()` and `subroutine()` are non-`static` members functions on `KernelContext`, so they have an implicit `this` parameter of type `KernelContext`, which allows the bodies of those functions to implicitly reference `gStuff`, etc. by name in their bodies.

Because `KernelContext::computeMain()` is a member function, we end up emitting an additional global-scope function to expose the entry point to the outside world, and that function is responsible for declaring a local `KernelContext` and invoking the generated entry point on it.

This approach has several important drawbacks:

* It complicates the emit logic for CPU and CUDA, with many special cases around when/how things get emitted

* It complicates the implementation of dynamic dispatch, because what seems like a function pointer in Slang IR needs to be a pointer-to-member-function in C++.

* It makes it difficult to have a non-kernel-oriented mode of compilation for CPU where a Slang function with a given signature gets output as a C++ CPU function with the "same" signature (not wrapped up as a member function of `KernelContext`.

This change makes a step toward addressing these issues by making the introducing of the `KernelContext` type be something that is done in an explicit IR pass instead of being handled as part of the last-mile emit logic.

The most important change is the removal of code related to `KernelContext` from the `slang-emit-{cpp,cuda}.{h,cpp}` files, with the equivalent logic instead being handled in a new pass in `slang-ir-explicit-global-context.{h,cpp}`. It should be noted that further cleanups to the emit logic should now be possible; in particular, both the CPU and CUDA emit paths are manually sequencing the `EmitAction`s instead of relying on the default logic, but at this point they should be able to just use the default. The additional cleanups are left for future work.

The explicit IR pass does more or less what one would expect: it identifies global-scope entities (global variables and parameters) that need to be wrapped and turns them into fields of a `KernelContext` type. It then modifies all entry points to initialize a `KernelContext` as part of their startup. Finally, any code that used to refer to the global entities is changed to refer to a field of the context, with the context passed via new function parameters (the new parameter is only added to functions that need it for now).

Transforming global variables into fields of a `KernelContext` type in the IR pass ends up dropping their initial-value expressions (since those were attached as basic blocks on the `IRGlobalVar`). To avoid breaking code that relies on global-scope (but thread-local) variables, this change also adds an explicit pass that takes the initialization logic on all global variables and moves it to explicit logic that runs at the start of every entry point in a linked module (`slang-ir-explicit-global-init.{h,cpp}`). This pass would also be useful when we get back to direct SPIR-V emit, since SPIR-V also requires initialization logic for globals to be emitted into entry points.

One complication that arises when the IR is introducing the types for entry-point parameters, global-scope parameters, and the `KernelContext` type is that it becomes harder for the emit logic to utter the names of those types (they might not even have names, since `IRNameHint`s might get stripped). This created a problem since the wrapper operations that were being generated for CPU were taking `void*` parameters and casting them to the appropriate type. To work around this issue, we have added an explicit IR pass (`slang-ir-entry-point-raw-ptr-params.{h,cpp}`) that transforms the signature of entry points so that any pointer parameters instead become raw pointer (`void*`) parameters, with the casting being handled inside the entry point itself.

One consequence of all the above changes is that for the CUDA target we no longer need a wrapper function to invoke the generated entry point any more, because the IR function for the entry point ends up having the correct/expected signature already. This is also the case for CPU when it comes to the `*_Thread` wrapper function, but this change doesn't try to eliminate the wrapper because of a belief that the `*_Thread`-level interface is going away anyway.

Because the IR is now responsible for ensuring the signature of the IR entry point for CUDA and CPU is what is expected, I needed to modify the `slang-ir-entry-point-uniforms` pass to always create an explicit parameter for the entry point uniforms when compiling for CUDA/CPU, even if there were no `uniform` parameters on the entry point as written. This also ended up requiring some tweaks to the parameter layout logic to ensure that CPU/CUDA targets always treat `ConstantBuffer<T>` as a `T*` even in the case where `T` is an empty `struct` type (which happens when we construct a `struct` type to represent the uniform parameters of an entry point with no uniform parameters...).

There are several future changes that can/should build on this work:

* We should change the generated signatures for CUDA kernels, so that they don't rely on `KernelContext` for global-scope parameters. At that point we can avoid generating a `KernelContext` at all for CUDA, except when a program uses global-scope thread-local variables.

* We should figure out how to make the "ABI" for dynamic-dispatch calls ensure that the kernel context is either always passed, or always *not* passed. Making a hard-and-fast rule as part of the calling convention for dynamic calls would ensure that they access through the context continues to work with dynamic calls (this change might break it in some cases).

* We should figure out how to handle the layout for the `KernelContext` in cases where a program is composed of multiple separately-compiled modules. Right now the layout of the `KernelContext` requires global knowledge (as does the pass that introduces explicit initialization for global-scope thread-locals).

* We should try to further clean up the CPU/CUDA C++ emit logic to fall back on the default emit behavior more, now that the various special-case approaches that were taken are no longer needed

* fixup: restore build files to default configuration

1 files changed, 523 insertions, 0 deletions

diff --git a/source/slang/slang-ir-explicit-global-context.cpp b/source/slang/slang-ir-explicit-global-context.cpp
new file mode 100644
index 000000000..68f23461b
--- /dev/null
+++ b/source/slang/slang-ir-explicit-global-context.cpp
@@ -0,0 +1,523 @@
+// slang-ir-explicit-global-context.cpp
+#include "slang-ir-explicit-global-context.h"
+
+#include "slang-ir-insts.h"
+
+namespace Slang
+{
+
+// The job of this pass is take global-scope declarations
+// that are actually scoped to a single shader thread or
+// thread-group, and wrap them up in an explicit "context"
+// type that gets passed between functions.
+
+struct IntroduceExplicitGlobalContextPass
+{
+    IRModule*       m_module = nullptr;
+    CodeGenTarget   m_target = CodeGenTarget::Unknown;
+
+    SharedIRBuilder*    m_sharedBuilder         = nullptr;
+    IRStructType*       m_contextStructType     = nullptr;
+    IRPtrType*          m_contextStructPtrType  = nullptr;
+
+    IRGlobalParam*      m_globalUniformsParam   = nullptr;
+    List<IRGlobalVar*>  m_globalVars;
+    List<IRFunc*>       m_entryPoints;
+
+    void processModule()
+    {
+        SharedIRBuilder sharedBuilder(m_module);
+        m_sharedBuilder = &sharedBuilder;
+
+        IRBuilder builder(&sharedBuilder);
+
+        // The global context will be represneted by a `struct`
+        // type with a name hint of `KernelContext`.
+        //
+        m_contextStructType = builder.createStructType();
+        builder.addNameHintDecoration(m_contextStructType, UnownedTerminatedStringSlice("KernelContext"));
+
+        // The context will usually be passed around by pointer,
+        // so we get and cache that pointer type up front.
+        //
+        m_contextStructPtrType = builder.getPtrType(m_contextStructType);
+
+        // The transformation we will perform will need to affect
+        // global variables, global shader parameters, and entry-point
+        // function (at the very least), and we start with an explicit
+        // pass to collect these entities into explicit lists to simplify
+        // looping over them later.
+        //
+        for( auto inst : m_module->getGlobalInsts() )
+        {
+            switch( inst->op )
+            {
+            case kIROp_GlobalVar:
+                {
+                    // A "global variable" in HLSL (and thus Slang) is actually
+                    // a weird kind of thread-local variable, and so it cannot
+                    // actually be lowered to a global variable on targets where
+                    // globals behave like, well, globals.
+                    //
+                    auto globalVar = cast<IRGlobalVar>(inst);
+
+                    // One important exception is that CUDA *does* support
+                    // global variables with the `__shared__` qualifer, with
+                    // semantics that exactly match HLSL/Slang `groupshared`.
+                    //
+                    // We thus need to skip processing of global variables
+                    // that were marked `groupshared`. In our current IR,
+                    // this is represented as a variable with the `@GroupShared`
+                    // rate on its type.
+                    //
+                    if( m_target == CodeGenTarget::CUDASource )
+                    {
+                        if( as<IRGroupSharedRate>(globalVar->getRate()) )
+                            continue;
+                    }
+
+                    m_globalVars.add(globalVar);
+                }
+                break;
+
+            case kIROp_GlobalParam:
+                {
+                    // Global parameters are another HLSL/Slang concept
+                    // that doesn't have a parallel in langauges like C/C++.
+                    //
+                    auto globalParam = cast<IRGlobalParam>(inst);
+
+
+                    // One detail we need to be careful about is that as a result
+                    // of legalizing the varying parameters of kernels, we can end
+                    // up with global parameters for varying parameters on CUDA
+                    // (e.g., to represent `threadIdx`. We thus skip any global-scope
+                    // parameters that are varying instead of uniform.
+                    //
+                    auto layoutDecor = globalParam->findDecoration<IRLayoutDecoration>();
+                    SLANG_ASSERT(layoutDecor);
+                    auto layout = as<IRVarLayout>(layoutDecor->getLayout());
+                    SLANG_ASSERT(layout);
+                    if(isVaryingParameter(layout))
+                        continue;
+
+                    // Because of upstream passes, we expect there to be only a
+                    // single global uniform parameter (at most).
+                    //
+                    // Note: If we ever changed out mind about the representation
+                    // and wanted to support multiple global parameters, we could
+                    // easily generalize this code to work with a list.
+                    //
+                    SLANG_ASSERT(!m_globalUniformsParam);
+                    m_globalUniformsParam = globalParam;
+                }
+                break;
+
+            case kIROp_Func:
+                {
+                    // Every entry point function is going to need to be modified,
+                    // so that it can explicit create the context that other
+                    // operations will use.
+
+                    // We need to filter the IR functions to find only those
+                    // that represent entry points.
+                    //
+                    auto func = cast<IRFunc>(inst);
+                    if(!func->findDecoration<IREntryPointDecoration>())
+                        continue;
+
+                    m_entryPoints.add(func);
+                }
+                break;
+            }
+        }
+
+        // Now that we've capture all the relevant global entities from the IR,
+        // we can being to transform them in an appropriate order.
+        //
+        // The first step will be to create fields in the `KernelContext`
+        // type to represent any global parameters or global variables.
+        //
+        // The keys for the fields that are created will be remembered
+        // in a dictionary, so that we can find them later based on
+        // the global parameter/variable.
+        //
+        if( m_globalUniformsParam )
+        {
+            // For the parameter representing all the global uniform shader
+            // parameters, we create a field that exactly matches its type.
+            //
+            createContextStructField(m_globalUniformsParam, m_globalUniformsParam->getFullType());
+        }
+        for( auto globalVar : m_globalVars )
+        {
+            // A `IRGlobalVar` represents a pointer to where the variable is stored,
+            // so we need to create a field of the pointed-to type to represent it.
+            //
+            createContextStructField(globalVar, globalVar->getDataType()->getValueType());
+        }
+
+        // Once all the fields have been created, we can process the entry points.
+        //
+        // Each entry point will create a local `KernelContext` variable and
+        // initialize it based on the parameters passed to the entry point.
+        //
+        // The local variable introduced here will be registered as the representation
+        // of the context to be used in the body of the entry point.
+        //
+        for( auto entryPoint : m_entryPoints )
+        {
+            createContextForEntryPoint(entryPoint);
+        }
+
+        // Now that we've prepared all the entry points, we can make another
+        // pass over the global parameters/variables and start to replace
+        // their use sites with references to the fields of the context.
+        //
+        // Wherever a global parameter/variable is being referenced in a function,
+        // we will need to find or create a context value for that function
+        // to use. The context value for entry points has already been established
+        // above, but other functions will have an explicit context parameter
+        // added on demand.
+        //
+        if( m_globalUniformsParam )
+        {
+            replaceUsesOfGlobalParam(m_globalUniformsParam);
+        }
+        for( auto globalVar : m_globalVars )
+        {
+            replaceUsesOfGlobalVar(globalVar);
+        }
+    }
+
+    // As noted above, we will maintain mappings to record
+    // the key for the context field created for a global
+    // variable parameter, and to record the context pointer
+    // value to use for a function.
+    //
+    Dictionary<IRInst*, IRStructKey*> m_mapInstToContextFieldKey;
+    Dictionary<IRFunc*, IRInst*> m_mapFuncToContextPtr;
+
+    void createContextStructField(IRInst* originalInst, IRType* type)
+    {
+        // Creating a field in the context struct to represent
+        // `originalInst` is straightforward.
+
+        IRBuilder builder(m_sharedBuilder);
+        builder.setInsertBefore(m_contextStructType);
+
+        // We create a "key" for the new field, and then a field
+        // of the appropraite type.
+        //
+        auto key = builder.createStructKey();
+        auto field = builder.createStructField(m_contextStructType, key, type);
+
+        // If the original instruction had a name hint on it,
+        // then we transfer that name hint over to the key,
+        // so that the field will have the name of the former
+        // global variable/parameter.
+        //
+        if( auto nameHint = originalInst->findDecoration<IRNameHintDecoration>() )
+        {
+            nameHint->insertAtStart(key);
+        }
+
+        // Any other decorations on the original instruction
+        // (e.g., pertaining to layout) need to be transferred
+        // over to the field (not the key).
+        //
+        originalInst->transferDecorationsTo(field);
+
+        // We end by making note of the key that was created
+        // for the instruction, so that we can use the key
+        // to access the field later.
+        //
+        m_mapInstToContextFieldKey.Add(originalInst, key);
+    }
+
+    void createContextForEntryPoint(IRFunc* entryPointFunc)
+    {
+        // We can only introduce the explicit context into
+        // entry points that have definitions.
+        //
+        auto firstBlock = entryPointFunc->getFirstBlock();
+        if(!firstBlock)
+            return;
+
+        IRBuilder builder(m_sharedBuilder);
+
+        // The code we introduce will all be added to the start
+        // of the first block of the function.
+        //
+        auto firstOrdinary = firstBlock->getFirstOrdinaryInst();
+        builder.setInsertBefore(firstOrdinary);
+
+        // If there was a global-scope uniform parameter before,
+        // then we need to introduce an explicit parameter onto
+        // each entry-point function to represent it.
+        //
+        IRParam* globalUniformsParam = nullptr;
+        if( m_globalUniformsParam )
+        {
+            globalUniformsParam = builder.createParam(m_globalUniformsParam->getFullType());
+            if( auto nameHint = m_globalUniformsParam->findDecoration<IRNameHintDecoration>() )
+            {
+                builder.addNameHintDecoration(globalUniformsParam, nameHint->getNameOperand());
+            }
+
+            // The new parameter will be the last one in the
+            // parameter list of the entry point.
+            //
+            globalUniformsParam->insertBefore(firstOrdinary);
+        }
+
+        // The `KernelContext` to use inside the entry point
+        // will be a local variable declared in the first block.
+        //
+        auto contextVarPtr = builder.emitVar(m_contextStructType);
+        addKernelContextNameHint(contextVarPtr);
+        m_mapFuncToContextPtr.Add(entryPointFunc, contextVarPtr);
+
+        // If there is a global-scope uniform parameter, then
+        // we need to use our new explicit entry point parameter
+        // to inialize the corresponding field of the `KernelContext`
+        // before moving on with execution of the kernel body.
+        //
+        if(m_globalUniformsParam)
+        {
+            auto fieldKey = m_mapInstToContextFieldKey[m_globalUniformsParam];
+            auto fieldType = globalUniformsParam->getFullType();
+            auto fieldPtrType = builder.getPtrType(fieldType);
+
+            // We compute the addrress of the field and store the
+            // value of the parameter into it.
+            //
+            auto fieldPtr = builder.emitFieldAddress(fieldPtrType, contextVarPtr, fieldKey);
+            builder.emitStore(fieldPtr, globalUniformsParam);
+        }
+
+        // Note: at this point the `KernelContext` has additional
+        // fields for global variables that do not seem to have
+        // been initialized.
+        //
+        // Instead of making this pass take responsibility for initializing
+        // global variables, it is instead expected that clients will
+        // run the pass in `slang-ir-explicit-global-init` first,
+        // in order to move all initialization of globals into the
+        // entry point functions.
+    }
+
+    void replaceUsesOfGlobalParam(IRGlobalParam* globalParam)
+    {
+        IRBuilder builder(m_sharedBuilder);
+
+        // A global shader parameter was mapped to a field
+        // in the context structure, so we find the appropriate key.
+        //
+        auto key = m_mapInstToContextFieldKey[globalParam];
+
+        auto valType = globalParam->getFullType();
+        auto ptrType = builder.getPtrType(valType);
+
+        // We then iterate over the uses of the parameter,
+        // being careful to defend against the use/def information
+        // being changed while we walk it.
+        //
+        IRUse* nextUse = nullptr;
+        for( IRUse* use = globalParam->firstUse; use; use = nextUse )
+        {
+            nextUse = use->nextUse;
+
+            // At each use site, we need to look up the context
+            // pointer that is appropriate for that use.
+            //
+            auto user = use->getUser();
+            auto contextParam = findOrCreateContextPtrForInst(user);
+            builder.setInsertBefore(user);
+
+            // The value of the parameter can be produced by
+            // taking the address of the corresponding field
+            // in the context struct and loading from it.
+            //
+            auto ptr = builder.emitFieldAddress(ptrType, contextParam, key);
+            auto val = builder.emitLoad(valType, ptr);
+            use->set(val);
+        }
+    }
+
+    void replaceUsesOfGlobalVar(IRGlobalVar* globalVar)
+    {
+        IRBuilder builder(m_sharedBuilder);
+
+        // A global variable was mapped to a field
+        // in the context structure, so we find the appropriate key.
+        //
+        auto key = m_mapInstToContextFieldKey[globalVar];
+
+        auto ptrType = globalVar->getDataType();
+
+        // We then iterate over the uses of the variable,
+        // being careful to defend against the use/def information
+        // being changed while we walk it.
+        //
+        IRUse* nextUse = nullptr;
+        for( IRUse* use = globalVar->firstUse; use; use = nextUse )
+        {
+            nextUse = use->nextUse;
+
+            // At each use site, we need to look up the context
+            // pointer that is appropriate for that use.
+            //
+            auto user = use->getUser();
+            auto contextParam = findOrCreateContextPtrForInst(user);
+            builder.setInsertBefore(user);
+
+            // The address of the variable can be produced by
+            // taking the address of the corresponding field
+            // in the context struct.
+            //
+            auto ptr = builder.emitFieldAddress(ptrType, contextParam, key);
+            use->set(ptr);
+        }
+    }
+
+    IRInst* findOrCreateContextPtrForInst(IRInst* inst)
+    {
+        // When looking up the context pointer to use for
+        // an instruction, we need to find the enclosing
+        // function and use whatever context pointer it uses.
+        //
+        for( IRInst* i = inst; i; i = i->getParent() )
+        {
+            if( auto func = as<IRFunc>(i) )
+            {
+                return findOrCreateContextPtrForFunc(func);
+            }
+        }
+
+        // If a non-constant global entity is being referenced by
+        // something that is *not* nested under an IR function, then
+        // we are in trouble.
+        //
+        SLANG_UNEXPECTED("no outer func at use site for global");
+        UNREACHABLE_RETURN(nullptr);
+    }
+
+    IRInst* findOrCreateContextPtrForFunc(IRFunc* func)
+    {
+        // At this point we are being asked to either find or
+        // produce a context pointer for use inside `func`.
+        //
+        // If we already created such a pointer (perhaps because
+        // `func` is an entry point), then we are home free.
+        //
+        if( auto found = m_mapFuncToContextPtr.TryGetValue(func) )
+        {
+            return *found;
+        }
+
+        // Otherwise, we are going to need to introduce an
+        // explicit parameter to `func` to represent the
+        // context.
+        //
+        IRBuilder builder(m_sharedBuilder);
+
+        // We can safely assume that `func` has a body, because
+        // otherwise we wouldn't be getting a request for the
+        // context pointer value to use in its body.
+        //
+        auto firstBlock = func->getFirstBlock();
+        SLANG_ASSERT(firstBlock);
+
+        // We create a new parameter at the end of the parameter
+        // list for `func`, with a type of `KernelContext*`.
+        //
+        IRParam* contextParam = builder.createParam(m_contextStructPtrType);
+        addKernelContextNameHint(contextParam);
+        contextParam->insertBefore(firstBlock->getFirstOrdinaryInst());
+
+        // The new parameter can be registerd as the context value
+        // to be used for `func` right away.
+        //
+        // Note: we register the value *before* modifying locations
+        // that call `func` to protect against a possible infinite-recursion
+        // situation if `func` is recursive along some path.
+        //
+        m_mapFuncToContextPtr.Add(func, contextParam);
+
+        // Any code that calls `func` now needs to be updated to pass
+        // the context parameter.
+        //
+        // TODO: There is an issue here if `func` might be called
+        // dynamically, through something like a witness table.
+        //
+        List<IRUse*> uses;
+        for( auto use = func->firstUse; use; use = use->nextUse )
+        {
+            // We will only fix up calls to `func`, and ignore
+            // other operations that might refer to it.
+            //
+            // TODO: We need to allow things like decorations that might
+            // refer to `func`, but this logic is also going to
+            // ignore things like witness tables that refer to `func`,
+            // or operations that pass `func` as a function pointer
+            // to a higher-order function.
+            //
+            auto call = as<IRCall>(use->getUser());
+            if(!call)
+                continue;
+
+            // We are going to construct a new call to `func`
+            // that has all of the arguments of the original call...
+            //
+            UInt originalArgCount = call->getArgCount();
+            List<IRInst*> args;
+            for( UInt aa = 0; aa < originalArgCount; ++aa )
+            {
+                args.add(call->getArg(aa));
+            }
+
+            // ... plus an additional argument representing
+            // the context pointer at the call site (note that
+            // this step leads to a potential for recursion in this pass;
+            // the maximum depth of the recursion is bounded by the
+            // maximum length of a cycle-free path through the call
+            // graph of the program).
+            //
+            args.add(findOrCreateContextPtrForInst(call));
+
+            // The new call will be emitted right before the old one,
+            // then used to replace it.
+            //
+            builder.setInsertBefore(call);
+            auto newCall = builder.emitCallInst(call->getFullType(), call->getCallee(), args);
+            call->replaceUsesWith(newCall);
+            call->removeAndDeallocate();
+        }
+
+        return contextParam;
+    }
+
+    // Because we have multiple places where instructions representing
+    // the kernel context get introduced, we have factored out a subroutine
+    // for setting up the name hint to be used by those instructions.
+    //
+    void addKernelContextNameHint(IRInst* inst)
+    {
+        IRBuilder builder(m_sharedBuilder);
+        builder.addNameHintDecoration(inst, UnownedTerminatedStringSlice("kernelContext"));
+    }
+};
+
+    /// Collect global-scope variables/paramters to form an explicit context that gets threaded through
+void introduceExplicitGlobalContext(
+    IRModule*       module,
+    CodeGenTarget   target)
+{
+    IntroduceExplicitGlobalContextPass pass;
+    pass.m_module = module;
+    pass.m_target = target;
+    pass.processModule();
+}
+
+}