Fix output of `groupshared` with IR type system (#492)

The basic problem was that the lowering logic was constructing (more or less) `Ptr<@GroupShared X>` instead of `@GroupShared Ptr<X>`.
There were also problems with passes not propagating through rates that should have been (e.g., legalization).

I've added a test case to actually validate `groupshared` support.

6 files changed, 114 insertions, 57 deletions

diff --git a/source/slang/emit.cpp b/source/slang/emit.cpp
index 008ddb270..e93dfbbb5 100644
--- a/source/slang/emit.cpp
+++ b/source/slang/emit.cpp
@@ -1243,6 +1243,7 @@ struct EmitVisitor
                 auto rateQualifiedType = cast<IRRateQualifiedType>(type);
                 emitTypeImpl(rateQualifiedType->getValueType(), declarator);
             }
+            break;
 
         case kIROp_ArrayType:
             emitArrayTypeImpl(cast<IRArrayType>(type), declarator);
@@ -3249,6 +3250,7 @@ struct EmitVisitor
                 auto valType = ptrType->getValueType();
 
                 auto name = getIRName(inst);
+                emitIRRateQualifiers(ctx, inst);
                 emitIRType(ctx, valType, name);
                 emit(";\n");
             }
@@ -4155,7 +4157,7 @@ struct EmitVisitor
         {
             for (auto pp = bb->getFirstParam(); pp; pp = pp->getNextParam())
             {
-                emitIRType(ctx, pp->getDataType(), getIRName(pp));
+                emitIRType(ctx, pp->getFullType(), getIRName(pp));
                 emit(";\n");
             }
         }
@@ -5165,6 +5167,7 @@ struct EmitVisitor
 
         emitIRVarModifiers(ctx, layout, varDecl, varType);
 
+        emitIRRateQualifiers(ctx, varDecl);
         emitIRType(ctx, varType, getIRName(varDecl));
 
         emitIRSemantics(ctx, varDecl);
@@ -5212,6 +5215,7 @@ struct EmitVisitor
             emit("static ");
         }
         emit("const ");
+        emitIRRateQualifiers(ctx, valDecl);
         emitIRType(ctx, valType, getIRName(valDecl));
 
         if (valDecl->getFirstBlock())
diff --git a/source/slang/ir-legalize-types.cpp b/source/slang/ir-legalize-types.cpp
index 20efc02b1..b9b553e15 100644
--- a/source/slang/ir-legalize-types.cpp
+++ b/source/slang/ir-legalize-types.cpp
@@ -600,6 +600,8 @@ static LegalVal legalizeLocalVar(
         context,
         irLocalVar->getDataType()->getValueType());
 
+    auto originalRate = irLocalVar->getRate();
+
     RefPtr<VarLayout> varLayout = findVarLayout(irLocalVar);
     RefPtr<TypeLayout> typeLayout = varLayout ? varLayout->typeLayout : nullptr;
 
@@ -614,14 +616,25 @@ static LegalVal legalizeLocalVar(
     switch (maybeSimpleType.flavor)
     {
     case LegalType::Flavor::simple:
-        // Easy case: the type is usable as-is, and we
-        // should just do that.
-        irLocalVar->setFullType(context->builder->getPtrType(
-            maybeSimpleType.getSimple()));
-        return LegalVal::simple(irLocalVar);
+        {
+            // Easy case: the type is usable as-is, and we
+            // should just do that.
+            auto type = maybeSimpleType.getSimple();
+            type = context->builder->getPtrType(type);
+            if( originalRate )
+            {
+                type = context->builder->getRateQualifiedType(
+                    originalRate,
+                    type);
+            }
+            irLocalVar->setFullType(type);
+            return LegalVal::simple(irLocalVar);
+        }
 
     default:
     {
+        // TODO: We don't handle rates in this path.
+
         context->insertBeforeLocalVar = irLocalVar;
 
         LegalVarChain* varChain = nullptr;
diff --git a/source/slang/ir.cpp b/source/slang/ir.cpp
index 2615c1c07..397726b0a 100644
--- a/source/slang/ir.cpp
+++ b/source/slang/ir.cpp
@@ -548,6 +548,9 @@ namespace Slang
     //
     IRParentInst* mergeCandidateParentsForHoistableInst(IRParentInst* left, IRParentInst* right)
     {
+        // If the candidates are both the same, then who cares?
+        if(left == right) return left;
+
         // If either `left` or `right` is a block, then we need to be
         // a bit careful, because blocks can see other values just using
         // the dominance relationship, without a direct parent-child relationship.
@@ -4805,6 +4808,27 @@ namespace Slang
         IRGlobalValueWithCode*  clonedValue,
         IRGlobalValueWithCode*  originalValue);
 
+    IRRate* cloneRate(
+        IRSpecContextBase*  context,
+        IRRate*             rate)
+    {
+        return (IRRate*) cloneType(context, rate);
+    }
+
+    void maybeSetClonedRate(
+        IRSpecContextBase*  context,
+        IRBuilder*          builder,
+        IRInst*             clonedValue,
+        IRInst*             originalValue)
+    {
+        if(auto rate = originalValue->getRate() )
+        {
+            clonedValue->setFullType(builder->getRateQualifiedType(
+                cloneRate(context, rate),
+                clonedValue->getFullType()));
+        }
+    }
+
     IRGlobalVar* cloneGlobalVarImpl(
         IRSpecContextBase*              context,
         IRBuilder*                      builder,
@@ -4814,11 +4838,7 @@ namespace Slang
         auto clonedVar = builder->createGlobalVar(
             cloneType(context, originalVar->getDataType()->getValueType()));
 
-        if(auto rate = originalVar->getRate() )
-        {
-            clonedVar->setFullType(builder->getRateQualifiedType(
-                rate, clonedVar->getFullType()));
-        }
+        maybeSetClonedRate(context, builder, clonedVar, originalVar);
 
         registerClonedValue(context, clonedVar, originalValues);
 
diff --git a/source/slang/lower-to-ir.cpp b/source/slang/lower-to-ir.cpp
index 4f5e8bceb..3953490fe 100644
--- a/source/slang/lower-to-ir.cpp
+++ b/source/slang/lower-to-ir.cpp
@@ -1195,6 +1195,25 @@ void addVarDecorations(
     }
 }
 
+/// If `decl` has a modifier that should turn into a
+/// rate qualifier, then apply it to `inst`.
+void maybeSetRate(
+    IRGenContext*   context,
+    IRInst*         inst,
+    Decl*           decl)
+{
+    auto builder = context->irBuilder;
+
+    if (decl->HasModifier<HLSLGroupSharedModifier>())
+    {
+        inst->setFullType(builder->getRateQualifiedType(
+            builder->getGroupSharedRate(),
+            inst->getFullType()));
+    }
+}
+
+
+
 LoweredValInfo createVar(
     IRGenContext*   context,
     IRType*         type,
@@ -1205,6 +1224,8 @@ LoweredValInfo createVar(
 
     if (decl)
     {
+        maybeSetRate(context, irAlloc, decl);
+
         addVarDecorations(context, irAlloc, decl);
 
         builder->addHighLevelDeclDecoration(irAlloc, decl);
@@ -3192,22 +3213,6 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
     {
         IRType* varType = lowerType(context, decl->getType());
 
-        if (decl->HasModifier<HLSLGroupSharedModifier>())
-        {
-            // TODO: here we are applying the rate qualifier to
-            // the *data type* of the variable, when we really
-            // should be applying the rate to the variable itself.
-            //
-            // This ends up making a distinction between
-            // `Ptr<@GroupShared X>` and `@GroupShared Ptr<X>`.
-            // The latter is more technically correct, but the
-            // code generation logic currently looks for the former.
-
-            varType = getBuilder()->getRateQualifiedType(
-                getBuilder()->getGroupSharedRate(),
-                varType);
-        }
-
         auto builder = getBuilder();
 
         IRGlobalValueWithCode* irGlobal = nullptr;
@@ -3226,6 +3231,8 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
         }
         irGlobal->mangledName = context->getSession()->getNameObj(getMangledName(decl));
 
+        maybeSetRate(context, irGlobal, decl);
+
         if (decl)
         {
             builder->addHighLevelDeclDecoration(irGlobal, decl);
@@ -3300,36 +3307,6 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
         // initializer expression a bit carefully (it should only
         // be initialized on-demand at its first use).
 
-        // Some qualifiers on a variable will change how we allocate it,
-        // so we need to reflect that somehow. The first example
-        // we run into is the `groupshared` qualifier, which marks
-        // a variable in a compute shader as having per-group allocation
-        // rather than the traditional per-thread (or rather per-thread
-        // per-activation-record) allocation.
-        //
-        // Options include:
-        //
-        //  - Use a distinct allocation opration, so that the type
-        //    of the variable address/value is unchanged.
-        //
-        //  - Add a notion of an "address space" to pointer types,
-        //    so that we can allocate things in distinct spaces.
-        //
-        //  - Add a notion of a "rate" so that we can declare a
-        //    variable with a distinct rate.
-        //
-        // For now we might do the expedient thing and handle this
-        // via a notion of an "address space."
-
-        if (decl->HasModifier<HLSLGroupSharedModifier>())
-        {
-            // TODO: This logic is duplicated with the global-variable
-            // case. We should seek to share it.
-            varType = getBuilder()->getRateQualifiedType(
-                getBuilder()->getGroupSharedRate(),
-                varType);
-        }
-
         LoweredValInfo varVal = createVar(context, varType, decl);
 
         if( auto initExpr = decl->initExpr )
diff --git a/tests/compute/groupshared.slang b/tests/compute/groupshared.slang
new file mode 100644
index 000000000..e471e6148
--- /dev/null
+++ b/tests/compute/groupshared.slang
@@ -0,0 +1,39 @@
+// groupshared.slang
+
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -dx12
+//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):dxbinding(0),glbinding(0),out
+
+#define THREAD_COUNT 4
+
+groupshared int gA[THREAD_COUNT];
+int test(int val)
+{
+	gA[val] = val;
+	GroupMemoryBarrierWithGroupSync();
+	val = gA[val ^ 1];
+
+/* TODO: once function-scope `static` works
+	static groupshared int gB[THREAD_COUNT];
+
+	gB[val] = val;
+	GroupMemoryBarrierWithGroupSync();
+	val = gB[val ^ 2];
+*/
+
+	return val;
+}
+
+RWStructuredBuffer<int> gBuffer;
+
+[numthreads(THREAD_COUNT, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+	uint tid = dispatchThreadID.x;
+
+	int val = int(tid);
+	val = test(val);
+	gBuffer[tid] = val;
+}
\ No newline at end of file
diff --git a/tests/compute/groupshared.slang.expected.txt b/tests/compute/groupshared.slang.expected.txt
new file mode 100644
index 000000000..8a5d99f5b
--- /dev/null
+++ b/tests/compute/groupshared.slang.expected.txt
@@ -0,0 +1,4 @@
+1
+0
+3
+2