Merge pull request shader-slang#324 from tfoleyNV/generic-struct-specialization

Support generic `struct` types during IR-based emit

Author: csyonghe

source/slang/emit.cpp

@@ -116,6 +116,10 @@ struct SharedEmitContext
 
     // Map used to tell AST lowering what decls are represented by IR.
     HashSet<Decl*>* irDeclSetForAST = nullptr;
+
+    // Are we doing IR-only emit, so that everything should get
+    // its mangled name?
+    bool isFullIRMode = false;
 };
 
 struct EmitContext
@@ -2996,10 +3000,21 @@ struct EmitVisitor
         }
     }
 
+    bool isBuiltinDecl(Decl* decl)
+    {
+        for (auto dd = decl; dd; dd = dd->ParentDecl)
+        {
+            if (dd->FindModifier<FromStdLibModifier>())
+                return true;
+        }
+        return false;
+    }
+
     void EmitDeclRef(DeclRef<Decl> declRef)
     {
         // Are we emitting an AST in a context where some declarations
         // are actually stored as IR code?
+
         if(auto irDeclSet = context->shared->irDeclSetForAST)
         {
             Decl* decl = declRef.getDecl();
@@ -3010,6 +3025,17 @@ struct EmitVisitor
             }
         }
 
+        if (context->shared->isFullIRMode)
+        {
+            // Don't apply this to builting declarations
+            if (!isBuiltinDecl(declRef.getDecl()))
+            {
+                emit(getIRName(declRef));
+                return;
+            }
+        }
+
+
 
         // TODO: need to qualify a declaration name based on parent scopes/declarations
 
@@ -4453,26 +4479,34 @@ emitDeclImpl(decl, nullptr);
 
     String getIRName(DeclRefBase const& declRef)
     {
-        // It is a bit ugly, but we need a deterministic way
-        // to get a name for things when emitting from the IR
-        // that won't conflict with any keywords, builtins, etc.
-        // in the target language.
+        // In general, when referring to a declaration that has been lowered
+        // via the IR, we want to use its mangled name.
+        //
+        // There are two main exceptions to this:
+        //
+        // 1. For debugging, we accept the `-no-mangle` flag which basically
+        //    instructs us to try to use the original name of all declarations,
+        //    to make the output more like what is expected to come out of
+        //    fxc pass-through. This case should get deprecated some day.
         //
-        // Eventually we should accomplish this by using
-        // mangled names everywhere, but that complicates things
-        // when we are also using direct comparison to fxc/glslang
-        // output for some of our tests.
+        // 2. It is really annoying to have the fields of a `struct` type
+        //    get ridiculously lengthy mangled names, and this also messes
+        //    up stuff like specialization (since the mangled name of a field
+        //    would then include the mangled name of the outer type).
         //
 
         String name;
         if (context->shared->entryPoint->compileRequest->compileFlags & SLANG_COMPILE_FLAG_NO_MANGLING)
         {
+            // Special case (1):
             name.append(getText(declRef.GetName()));
+            return name;
         }
-        else
-        {
-            name.append(getMangledName(declRef));
-        }
+
+        // Special case (2): not implemented yet.
+
+        // General case:
+        name.append(getMangledName(declRef));
         return name;
     }
 
@@ -7040,6 +7074,24 @@ emitDeclImpl(decl, nullptr);
         {}
     }
 
+    void emitIRUsedTypesForGlobalValueWithCode(
+        EmitContext*            ctx,
+        IRGlobalValueWithCode*  value)
+    {
+        for( auto bb = value->getFirstBlock(); bb; bb = bb->getNextBlock() )
+        {
+            for( auto pp = bb->getFirstParam(); pp; pp = pp->getNextParam() )
+            {
+                emitIRUsedTypesForValue(ctx, pp);
+            }
+
+            for( auto ii = bb->getFirstInst(); ii; ii = ii->getNextInst() )
+            {
+                emitIRUsedTypesForValue(ctx, ii);
+            }
+        }
+    }
+
     void emitIRUsedTypesForValue(
         EmitContext*    ctx,
         IRValue*        value)
@@ -7050,19 +7102,23 @@ emitDeclImpl(decl, nullptr);
         case kIROp_Func:
             {
                 auto irFunc = (IRFunc*) value;
+
+                // Don't emit anything for a generic function,
+                // since we only care about the types used by
+                // the actual specializations.
+                if (irFunc->genericDecl)
+                    return;
+
                 emitIRUsedType(ctx, irFunc->getResultType());
-                for( auto bb = irFunc->getFirstBlock(); bb; bb = bb->getNextBlock() )
-                {
-                    for( auto pp = bb->getFirstParam(); pp; pp = pp->getNextParam() )
-                    {
-                        emitIRUsedTypesForValue(ctx, pp);
-                    }
 
-                    for( auto ii = bb->getFirstInst(); ii; ii = ii->getNextInst() )
-                    {
-                        emitIRUsedTypesForValue(ctx, ii);
-                    }
-                }
+                emitIRUsedTypesForGlobalValueWithCode(ctx, irFunc);
+            }
+            break;
+
+        case kIROp_global_var:
+            {
+                auto irGlobal = (IRGlobalVar*) value;
+                emitIRUsedTypesForGlobalValueWithCode(ctx, irGlobal);
             }
             break;
 
@@ -7372,6 +7428,8 @@ String emitEntryPoint(
             // compilation work. We thus start by cloning any code needed
             // by the entry point over to our fresh IR module.
 
+            sharedContext.isFullIRMode = true;
+
             specializeIRForEntryPoint(
                 irSpecializationState,
                 entryPoint);

source/slang/ir.cpp

@@ -3079,7 +3079,20 @@ namespace Slang
     {
         if(!originalValue)
             return;
-        context->getClonedValues().Add(originalValue, clonedValue);
+
+        // Note: setting the entry direclty here rather than
+        // using `Add` or `AddIfNotExists` because we can conceivably
+        // clone the same value (e.g., a basic block inside a generic
+        // function) multiple times, and that is okay, and we really
+        // just need to keep track of the most recent value.
+
+        // TODO: The same thing could potentially be handled more
+        // cleanly by having a notion of scoping for these cloned-value
+        // mappings, so that we register cloned values for things
+        // inside of a function to a temporary mapping that we
+        // throw away after the function is done.
+
+        context->getClonedValues()[originalValue] = clonedValue;
     }
 
     // Information on values to use when registering a cloned value

source/slang/mangle.cpp

@@ -184,6 +184,18 @@ namespace Slang
         }
     }
 
+    // TODO: this needs to be centralized
+    RefPtr<GenericSubstitution> getOutermostGenericSubst(
+        RefPtr<Substitutions> inSubst)
+    {
+        for (auto subst = inSubst; subst; subst = subst->outer)
+        {
+            if (auto genericSubst = subst.As<GenericSubstitution>())
+                return genericSubst;
+        }
+        return nullptr;
+    }
+
     void emitQualifiedName(
         ManglingContext*    context,
         DeclRef<Decl>       declRef)
@@ -221,7 +233,7 @@ namespace Slang
             // There are two cases here: either we have specializations
             // in place for the parent generic declaration, or we don't.
 
-            auto subst = declRef.substitutions.As<GenericSubstitution>();
+            auto subst = getOutermostGenericSubst(declRef.substitutions);
             if( subst && subst->genericDecl == parentGenericDeclRef.getDecl() )
             {
                 // This is the case where we *do* have substitutions.

tests/compute/generic-struct.slang

@@ -0,0 +1,37 @@
+//TEST(compute):COMPARE_COMPUTE:-xslang -use-ir
+//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):dxbinding(0),glbinding(0),out
+
+// Check that user code can declare and use a generic
+// `struct` type.
+
+RWStructuredBuffer<int> outputBuffer;
+
+__generic<T>
+struct GenStruct
+{
+	T x;
+	T y;
+};
+
+__generic<T>
+T test(T val)
+{
+	GenStruct<T> gs;
+	gs.x = val;
+	gs.y = val;
+	return gs.x;
+}
+
+
+[numthreads(4, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+	uint tid = dispatchThreadID.x;
+
+	float outVal = 0;
+
+	outVal += test<uint>(tid);
+	outVal += test<float>(tid * 16.0f);
+
+	outputBuffer[tid] = int(outVal);
+}

tests/compute/generic-struct.slang.expected.txt

@@ -0,0 +1,4 @@
+0
+11
+22
+33