Enable HLSL/GLSL "rewrite" + IR-based Slang codegen (#300)

The big picture here is that the AST-to-AST pass in `ast-legalize` will now detect when a declaration being referenced comes from an `import`ed module, and (if IR codegen is enabled), it will trigger cloning of the IR for the chosen symbol into an IR module that will sit alongside the legalized AST.

Then, during HLSL/GLSL code emit, we emit all the IR-based code first, and then the AST-based code. Whenever the AST code references a symbol that was lowered via IR (we keep track of these) we emit the mangled name of the IR symbol.

Notes/details:

- A lot of the logic for cloning IR symbols referenced by the AST matches the same logic that would clone them for completely IR-based codegen, so I tried to hoist out the common logic and share it (e.g., so that we apply the same guaranteed transformations in both cases). This required basically rewriting the logic in `emit.cpp` that decomposed the various cases.

- There is a new compute test case added to test this functionality. `tests/compute/rewriter.hlsl` confirms that we can use the `-no-checking` mode for the HLSL code, but still make use of a library of Slang code that employs generics, etc.
  - Adding this test case required adding a new compute test mode that invokes `render-test` with the `-hlsl-rewrite` flag.

- It turns out that the existing `tests/render/cross-compile0.hlsl` test should have been using this functionality already. It was opting into the use of the IR via `-use-ir`, and  the `render-test` application already tries to set `-no-checking` for non-Slang input languages by default. Fixing the code path this test triggers means that it is now a second test of rewriter+IR codegen.

- The `translateDeclRef` logic in `ast-legalize.cpp` seemed sloppy in places, and would potentially clone declarations, when declaration references were desired. I tried to clean a bit of this up, so some call sites are now changed.

- This change tries to clean up some work around cloning of global values
  - All global value kinds (not just functions) now go through the logic of trying to pick a "best" definition, so that they can be used when we are linking multiple modules
  - The logic for registering cloned values has been unified a bit, so that clients always pass in an `IROriginalValuesForClone` that either wraps a single value (maybe just null), or an `IRSpecSymbol*` that gives a list of values to regsiter the new value as a clone for.
  - I made one piece of code that was cloning witness tables as part of generic specializations *not* register a clone. I think this is correct because we may specialize the same generic multiple ways, so registering any values we clone is not the right idea, but I might be missing something...
  - I also reorganized this logic so that it would be easier to clone a global value when we only know its mangled name (which is the case when it is the AST that triggers cloning)

- I made sure that when loading a module via `import`, the translation unit for the new module copies the `-use-ir` flag from the overall compile request, if it is present (otherwise we wouldn't generate IR for loaded modules at all... oops).

- Note that `getSpecializedGlobalValueForDeclRef()`, which is the main routine used by the AST legalization to trigger cloning of an IR value does *not* currently handle declaration references that require specialization.

- This change does *not* deal with trying to unify the type legalization logic between the AST-to-AST rewriter and the IR-based codegen, so if you call an imported function with types that require legalization, Bad Things are expected to happen right now.

1 files changed, 231 insertions, 88 deletions

diff --git a/source/slang/ir.cpp b/source/slang/ir.cpp
index ae7b71172..3a8aabd85 100644
--- a/source/slang/ir.cpp
+++ b/source/slang/ir.cpp
@@ -3041,9 +3041,40 @@ namespace Slang
         IRValue*    clonedValue,
         IRValue*    originalValue)
     {
+        if(!originalValue)
+            return;
         context->getClonedValues().Add(originalValue, clonedValue);
     }
 
+    // Information on values to use when registering a cloned value
+    struct IROriginalValuesForClone
+    {
+        IRValue*        originalVal = nullptr;
+        IRSpecSymbol*   sym = nullptr;
+
+        IROriginalValuesForClone() {}
+
+        IROriginalValuesForClone(IRValue* originalValue)
+            : originalVal(originalValue)
+        {}
+
+        IROriginalValuesForClone(IRSpecSymbol* symbol)
+            : sym(symbol)
+        {}
+    };
+
+    void registerClonedValue(
+        IRSpecContextBase*              context,
+        IRValue*                        clonedValue,
+        IROriginalValuesForClone const& originalValues)
+    {
+        registerClonedValue(context, clonedValue, originalValues.originalVal);
+        for( auto s = originalValues.sym; s; s = s->nextWithSameName )
+        {
+            registerClonedValue(context, clonedValue, s->irGlobalValue);
+        }
+    }
+
     void cloneDecorations(
         IRSpecContextBase*  context,
         IRValue*        clonedValue,
@@ -3100,9 +3131,7 @@ namespace Slang
     };
 
 
-    IRGlobalVar* cloneGlobalVar(IRSpecContext* context, IRGlobalVar* originalVar);
-    IRFunc* cloneFunc(IRSpecContext* context, IRFunc* originalFunc);
-    IRWitnessTable* cloneWitnessTable(IRSpecContext* context, IRWitnessTable* originalVar);
+    IRGlobalValue* cloneGlobalValue(IRSpecContext* context, IRGlobalValue* originalVal);
     RefPtr<Substitutions> cloneSubstitutions(
         IRSpecContext*  context,
         Substitutions*  subst);
@@ -3117,16 +3146,9 @@ namespace Slang
         switch (originalValue->op)
         {
         case kIROp_global_var:
-            return cloneGlobalVar(this, (IRGlobalVar*)originalValue);
-            break;
-
         case kIROp_Func:
-            return cloneFunc(this, (IRFunc*)originalValue);
-            break;
-
         case kIROp_witness_table:
-            return cloneWitnessTable(this, (IRWitnessTable*)originalValue);
-            break;
+            return cloneGlobalValue(this, (IRGlobalValue*) originalValue);
 
         case kIROp_boolConst:
             {
@@ -3334,10 +3356,13 @@ namespace Slang
         IRGlobalValueWithCode*  clonedValue,
         IRGlobalValueWithCode*  originalValue);
 
-    IRGlobalVar* cloneGlobalVar(IRSpecContext* context, IRGlobalVar* originalVar)
+    IRGlobalVar* cloneGlobalVarImpl(
+        IRSpecContext*  context,
+        IRGlobalVar*    originalVar,
+        IROriginalValuesForClone const& originalValues)
     {
         auto clonedVar = context->builder->createGlobalVar(context->maybeCloneType(originalVar->getType()->getValueType()));
-        registerClonedValue(context, clonedVar, originalVar);
+        registerClonedValue(context, clonedVar, originalValues);
 
         auto mangledName = originalVar->mangledName;
         clonedVar->mangledName = mangledName;
@@ -3360,10 +3385,13 @@ namespace Slang
         return clonedVar;
     }
 
-    IRWitnessTable* cloneWitnessTable(IRSpecContext* context, IRWitnessTable* originalTable)
+    IRWitnessTable* cloneWitnessTableImpl(
+        IRSpecContext*  context,
+        IRWitnessTable* originalTable,
+        IROriginalValuesForClone const& originalValues)
     {
         auto clonedTable = context->builder->createWitnessTable();
-        registerClonedValue(context, clonedTable, originalTable);
+        registerClonedValue(context, clonedTable, originalValues);
 
         auto mangledName = originalTable->mangledName;
         clonedTable->mangledName = mangledName;
@@ -3384,6 +3412,13 @@ namespace Slang
         return clonedTable;
     }
 
+    IRWitnessTable* cloneWitnessTableWithoutRegistering(
+        IRSpecContext*  context,
+        IRWitnessTable* originalTable)
+    {
+        return cloneWitnessTableImpl(context, originalTable, IROriginalValuesForClone());
+    }
+
     void cloneGlobalValueWithCodeCommon(
         IRSpecContextBase*      context,
         IRGlobalValueWithCode*  clonedValue,
@@ -3556,15 +3591,6 @@ namespace Slang
         return clonedFunc;
     }
 
-
-    IRFunc* cloneSimpleFunc(IRSpecContextBase* context, IRFunc* originalFunc)
-    {
-        auto clonedFunc = context->builder->createFunc();
-        registerClonedValue(context, clonedFunc, originalFunc);
-        cloneFunctionCommon(context, clonedFunc, originalFunc);
-        return clonedFunc;
-    }
-
     // Get a string form of the target so that we can
     // use it to match against target-specialization modifiers
     //
@@ -3687,57 +3713,111 @@ namespace Slang
         return false;
     }
 
-    IRFunc* cloneFunc(IRSpecContext* context, IRFunc* originalFunc)
+    IRFunc* cloneFuncImpl(
+        IRSpecContext*  context,
+        IRFunc*         originalFunc,
+        IROriginalValuesForClone const& originalValues)
     {
-        // We are being asked to clone a particular function, but in
-        // the IR that comes out of the front-end there could still
-        // be multiple, target-specific, declarations of any given
-        // function, all of which share the same mangled name.
-        auto mangledName = originalFunc->mangledName;
+        auto clonedFunc = context->builder->createFunc();
+        registerClonedValue(context, clonedFunc, originalValues);
+        cloneFunctionCommon(context, clonedFunc, originalFunc);
+        return clonedFunc;
+    }
+
+    // Directly clone a global value, based on a single definition/declaration, `originalVal`.
+    // The symbol `sym` will thread together other declarations of the same value, and
+    // we will register the new value as the cloned version of all of those.
+    IRGlobalValue* cloneGlobalValueImpl(
+        IRSpecContext*  context,
+        IRGlobalValue*  originalVal,
+        IRSpecSymbol*   sym)
+    {
+        if( !originalVal )
+        {
+            SLANG_UNEXPECTED("cloning a null value");
+            UNREACHABLE_RETURN(nullptr);
+        }
+
+        switch( originalVal->op )
+        {
+        case kIROp_Func:
+            return cloneFuncImpl(context, (IRFunc*) originalVal, sym);
+
+        case kIROp_global_var:
+            return cloneGlobalVarImpl(context, (IRGlobalVar*)originalVal, sym);
 
+        case kIROp_witness_table:
+            return cloneWitnessTableImpl(context, (IRWitnessTable*)originalVal, sym);
+
+        default:
+            SLANG_UNEXPECTED("unknown global value kind");
+            UNREACHABLE_RETURN(nullptr);
+        }
+
+    }
+
+    // Clone a global value, which has the given `mangledName`.
+    // The `originalVal` is a known global IR value with that name, if one is available.
+    // (It is okay for this parameter to be null).
+    IRGlobalValue* cloneGlobalValueWithMangledName(
+        IRSpecContext*  context,
+        String const&   mangledName,
+        IRGlobalValue*  originalVal)
+    {
         if(mangledName.Length() == 0)
         {
-            return cloneSimpleFunc(context, originalFunc);
+            // If there is no mangled name, then we assume this is a local symbol,
+            // and it can't possibly have multiple declarations.
+            return cloneGlobalValueImpl(context, originalVal, nullptr);
         }
 
         //
-        // We will scan through all of the available function declarations
-        // with the same mangled name as `originalFunc` and try
+        // We will scan through all of the available declarations
+        // with the same mangled name as `originalVal` and try
         // to pick the "best" one for our target.
 
         RefPtr<IRSpecSymbol> sym;
-        if( !context->getSymbols().TryGetValue(originalFunc->mangledName, sym) )
+        if( !context->getSymbols().TryGetValue(mangledName, sym) )
         {
             // This shouldn't happen!
-            SLANG_UNEXPECTED("no matching function registered");
-            UNREACHABLE_RETURN(cloneSimpleFunc(context, originalFunc));
+            SLANG_UNEXPECTED("no matching values registered");
+            UNREACHABLE_RETURN(cloneGlobalValueImpl(context, originalVal, nullptr));
         }
 
-        // We will try to track the "best" definition we can find.
-        IRFunc* bestFunc = (IRFunc*) sym->irGlobalValue;
-
+        // We will try to track the "best" declaration we can find.
+        //
+        // Generally, one declaration wil lbe better than another if it is
+        // more specialized for the chosen target. Otherwise, we simply favor
+        // definitions over declarations.
+        //
+        IRGlobalValue* bestVal = sym->irGlobalValue;
         for( auto ss = sym->nextWithSameName; ss; ss = ss->nextWithSameName )
         {
-            IRFunc* newFunc = (IRFunc*) ss->irGlobalValue;
-            if(isBetterForTarget(context, newFunc, bestFunc))
-                bestFunc = newFunc;
+            IRGlobalValue* newVal = ss->irGlobalValue;
+            if(isBetterForTarget(context, newVal, bestVal))
+                bestVal = newVal;
         }
 
-        // All right, we are now in a position to clone the "best"
-        // definition that was found.
-        auto clonedFunc = context->builder->createFunc();
-
-        // The resulting function will be used as the cloned version
-        // of every declaration/definition in the original IR.
-        for( auto ss = sym; ss; ss = ss->nextWithSameName )
-        {
-            registerClonedValue(context, clonedFunc, ss->irGlobalValue);
-        }
+        return cloneGlobalValueImpl(context, bestVal, sym);
+    }
 
-        // Clone the "best" definition into our context
-        cloneFunctionCommon(context, clonedFunc, bestFunc);
+    IRGlobalValue* cloneGlobalValueWithMangledName(IRSpecContext* context, String const& mangledName)
+    {
+        return cloneGlobalValueWithMangledName(context, mangledName, nullptr);
+    }
 
-        return clonedFunc;
+    // Clone a global value, where `originalVal` is one declaration/definition, but we might
+    // have to consider others, in order to find the "best" version of the symbol.
+    IRGlobalValue* cloneGlobalValue(IRSpecContext* context, IRGlobalValue* originalVal)
+    {
+        // We are being asked to clone a particular global value, but in
+        // the IR that comes out of the front-end there could still
+        // be multiple, target-specific, declarations of any given
+        // global value, all of which share the same mangled name.
+        return cloneGlobalValueWithMangledName(
+            context,
+            originalVal->mangledName,
+            originalVal);
     }
 
     StructTypeLayout* getGlobalStructLayout(
@@ -3857,7 +3937,7 @@ namespace Slang
                             globalVar = globalVar->getNextValue();
                         }
                         SLANG_ASSERT(table);
-                        table = cloneWitnessTable(context, (IRWitnessTable*)(table));
+                        table = cloneWitnessTableWithoutRegistering(context, (IRWitnessTable*)(table));
                         IRProxyVal * tableVal = new IRProxyVal();
                         tableVal->inst = table;
                         paramSubst->witnessTables.Add(KeyValuePair<RefPtr<Type>, RefPtr<Val>>(subtypeWitness->sup, tableVal));
@@ -3868,55 +3948,57 @@ namespace Slang
         return globalParamSubst;
     }
 
-    IRModule* specializeIRForEntryPoint(
+    struct IRSpecializationState
+    {
+        ProgramLayout*      programLayout;
+        CodeGenTarget       target;
+        TargetRequest*      targetReq;
+
+        IRModule* irModule = nullptr;
+        RefPtr<ProgramLayout> newProgramLayout;
+
+        IRSharedSpecContext sharedContextStorage;
+        IRSpecContext contextStorage;
+
+        IRSharedSpecContext* getSharedContext() { return &sharedContextStorage; }
+        IRSpecContext* getContext() { return &contextStorage; }
+    };
+
+    IRSpecializationState* createIRSpecializationState(
         EntryPointRequest*  entryPointRequest,
         ProgramLayout*      programLayout,
         CodeGenTarget       target,
         TargetRequest*      targetReq)
     {
+        IRSpecializationState* state = new IRSpecializationState();
+
+        state->programLayout = programLayout;
+        state->target = target;
+        state->targetReq = targetReq;
+
+
         auto compileRequest = entryPointRequest->compileRequest;
-        auto session = compileRequest->mSession;
         auto translationUnit = entryPointRequest->getTranslationUnit();
         auto originalIRModule = translationUnit->irModule;
-        if (!originalIRModule)
-        {
-            // We should already have emitted IR for the original
-            // translation unit, and it we don't have it, then
-            // we are now in trouble.
-            return nullptr;
-        }
-
-        // We now need to start cloning IR symbols from `originalIRModule`
-        // into a fresh IR module for this entry point. Along the way we need to:
-        //
-        // 1. Attach layout information from `programLayout` and/or `entryPointLayout`
-        //    onto the cloned IR symbols, to drive later code generation.
-        //
-        // 2. In cases where a function might have multiple target-specific definitions,
-        //    we need to pick the "best" one for the chosen code generation target.
-        //
-
-        IRSharedSpecContext sharedContextStorage;
 
+        auto sharedContext = state->getSharedContext();
         initializeSharedSpecContext(
-            &sharedContextStorage,
+            sharedContext,
             compileRequest->mSession,
             nullptr,
             originalIRModule);
+        state->irModule = sharedContext->module;
 
         // We also need to attach the IR definitions for symbols from
         // any loaded modules:
         for (auto loadedModule : compileRequest->loadedModulesList)
         {
-            insertGlobalValueSymbols(&sharedContextStorage, loadedModule->irModule);
+            insertGlobalValueSymbols(sharedContext, loadedModule->irModule);
         }
-        // any loaded modules
 
-
-        IRSpecContext contextStorage;
-        IRSpecContext*  context = &contextStorage;
-        context->shared = &sharedContextStorage;
-        context->builder = &sharedContextStorage.builderStorage;
+        auto context = state->getContext();
+        context->shared = sharedContext;
+        context->builder = &sharedContext->builderStorage;
         context->target = target;
 
         // Create the GlobalGenericParamSubstitution for substituting global generic types
@@ -3928,7 +4010,7 @@ namespace Slang
         // now specailize the program layout using the substitution
         RefPtr<ProgramLayout> newProgramLayout = specializeProgramLayout(targetReq, programLayout, globalParamSubst);
 
-        auto entryPointLayout = findEntryPointLayout(newProgramLayout, entryPointRequest);
+        state->newProgramLayout = newProgramLayout;
 
         // Next, we want to optimize lookup for layout infromation
         // associated with global declarations, so that we can 
@@ -3940,6 +4022,69 @@ namespace Slang
             context->globalVarLayouts.AddIfNotExists(mangledName, globalVarLayout);
         }
 
+        return state;
+    }
+
+    void destroyIRSpecializationState(IRSpecializationState* state)
+    {
+        delete state;
+    }
+
+    IRModule* getIRModule(IRSpecializationState* state)
+    {
+        return state->irModule;
+    }
+
+    IRGlobalValue* getSpecializedGlobalValueForDeclRef(
+        IRSpecializationState*  state,
+        DeclRef<Decl> const&    declRef)
+    {
+        // We will start be ensuring that we have code for
+        // the declaration itself.
+        auto decl = declRef.getDecl();
+        auto mangledDeclName = getMangledName(decl);
+
+        IRGlobalValue* irDeclVal = cloneGlobalValueWithMangledName(
+            state->getContext(),
+            mangledDeclName);
+        if(!irDeclVal)
+            return nullptr;
+
+        // Now we need to deal with specializing the given
+        // IR value based on the substitutions applied to
+        // our declaration reference.
+
+        if(!declRef.substitutions)
+            return irDeclVal;
+
+        SLANG_UNEXPECTED("unhandled");
+        UNREACHABLE_RETURN(nullptr);
+    }
+
+    void specializeIRForEntryPoint(
+        IRSpecializationState*  state,
+        EntryPointRequest*  entryPointRequest)
+    {
+        auto target = state->target;
+
+        auto compileRequest = entryPointRequest->compileRequest;
+        auto session = compileRequest->mSession;
+        auto translationUnit = entryPointRequest->getTranslationUnit();
+        auto originalIRModule = translationUnit->irModule;
+        if (!originalIRModule)
+        {
+            // We should already have emitted IR for the original
+            // translation unit, and it we don't have it, then
+            // we are now in trouble.
+            return;
+        }
+
+        auto context = state->getContext();
+        auto newProgramLayout = state->newProgramLayout;
+
+        auto entryPointLayout = findEntryPointLayout(newProgramLayout, entryPointRequest);
+
+
         // Next, we make sure to clone the global value for
         // the entry point function itself, and rely on
         // this step to recursively copy over anything else
@@ -3964,8 +4109,6 @@ namespace Slang
         default:
             break;
         }
-
-        return sharedContextStorage.module;
     }
 
     //