Legalization of function parameter types.

This commit addresses issue #275

This commit includes following changes:

1. legalize function parameter IRParam instructions
2. legalize function parameter types in IRFuncType
3. legalize call sites (IRCall) with proper arguments
4. legalize local vars that has a mixed resource type.

5 files changed, 317 insertions, 61 deletions

diff --git a/source/slang/emit.cpp b/source/slang/emit.cpp
index a8cd11af4..c72adea38 100644
--- a/source/slang/emit.cpp
+++ b/source/slang/emit.cpp
@@ -97,8 +97,6 @@ struct SharedEmitContext
     Dictionary<IRValue*, UInt> mapIRValueToID;
 
     HashSet<Decl*> irDeclsVisited;
-
-    Dictionary<IRBlock*, IRBlock*> irMapContinueTargetToLoopHead;
 };
 
 struct EmitContext
@@ -5566,10 +5564,10 @@ emitDeclImpl(decl, nullptr);
 
                         emit("for(;;)\n{\n");
 
-                        // Register information so that `continue` sites
-                        // can do the right thing:
-                        ctx->shared->irMapContinueTargetToLoopHead.Add(continueBlock, targetBlock);
-
+                        // TODO: Okay, we *said* we'd do this special
+                        // handling of the `continue` sites, but
+                        // we aren't actually setting anything up here...
+                        //
 
                         emitIRStmtsForBlocks(
                             ctx,
@@ -5590,28 +5588,7 @@ emitDeclImpl(decl, nullptr);
                 return;
 
             case kIROp_continue:
-                // With out current strategy for outputting loops,
-                // just outputting an AST-level `continue` here won't
-                // actually execute the statements in the continue block.
-                //
-                // Instead, we have to manually output those statements
-                // directly here, and *then* do an AST-level `continue`.
-                //
-                // This leads to code duplication when we have multiple
-                // `continue` sites in the original program, but it avoids
-                // introducing additional temporaries for control flow.
-                {
-                    auto continueInst = (IRContinue*) terminator;
-                    auto targetBlock = continueInst->getTargetBlock();
-                    IRBlock* loopHead = nullptr;
-                    ctx->shared->irMapContinueTargetToLoopHead.TryGetValue(targetBlock, loopHead);
-                    SLANG_ASSERT(loopHead);
-                    emitIRStmtsForBlocks(
-                        ctx,
-                        targetBlock,
-                        loopHead);
-                    emit("continue;\n");
-                }
+                emit("continue;\n");
                 return;
 
             case kIROp_loopTest:
diff --git a/source/slang/ir-legalize-types.cpp b/source/slang/ir-legalize-types.cpp
index 5b08acee8..2719f336b 100644
--- a/source/slang/ir-legalize-types.cpp
+++ b/source/slang/ir-legalize-types.cpp
@@ -210,7 +210,15 @@ struct TypeLegalizationContext
     // When inserting new globals, put them before this one.
     IRGlobalValue* insertBeforeGlobal = nullptr;
 
+    // When inserting new parameters, put them before this one.
+    IRParam* insertBeforeParam = nullptr;
+
     Dictionary<IRValue*, LegalVal> mapValToLegalVal;
+
+    IRVar* insertBeforeLocalVar = nullptr;
+    // store local var instructions that have been replaced here, so we can free them
+    // when legalization has done
+    List<IRInst*> oldLocalVars;
 };
 
 static void registerLegalizedValue(
@@ -353,6 +361,22 @@ static LegalType legalizeType(
     return LegalType::simple(type);
 }
 
+// Represents the "chain" of declarations that
+// were followed to get to a variable that we
+// are now declaring as a leaf variable.
+struct LegalVarChain
+{
+    LegalVarChain*  next;
+    VarLayout*      varLayout;
+};
+
+static LegalVal declareVars(
+    TypeLegalizationContext*    context,
+    IROp                        op,
+    LegalType                   type,
+    TypeLayout*                 typeLayout,
+    LegalVarChain*              varChain);
+
 // Legalize a type, and then expect it to
 // result in a simple type.
 static RefPtr<Type> legalizeSimpleType(
@@ -388,6 +412,42 @@ static LegalVal legalizeOperand(
     return LegalVal::simple(irValue);
 }
 
+static void getArgumentValues(
+    List<IRValue*> & instArgs, 
+    LegalVal val)
+{
+    switch (val.flavor)
+    {
+    case LegalVal::Flavor::simple:
+        instArgs.Add(val.getSimple());
+        break;
+    case LegalVal::Flavor::implicitDeref:
+        getArgumentValues(instArgs, val.getImplicitDeref());
+        break;
+    case LegalVal::Flavor::tuple:
+        {
+            for (auto elem : val.getTuple()->elements)
+                getArgumentValues(instArgs, elem.val);
+        }
+        break;
+    }
+}
+
+static LegalVal legalizeCall(
+    TypeLegalizationContext*    context,
+    IRCall* callInst)
+{
+    // TODO: implement legalization of non-simple return types
+    auto retType = legalizeType(context, callInst->type);
+    SLANG_ASSERT(retType.flavor == LegalType::Flavor::simple);
+    
+    List<IRValue*> instArgs;
+    for (auto i = 1u; i < callInst->argCount; i++)
+        getArgumentValues(instArgs, legalizeOperand(context, callInst->getArg(i)));
+
+    return LegalVal::simple(context->builder->emitCallInst(callInst->type, callInst->func.usedValue, instArgs.Count(), instArgs.Buffer()));
+}
+
 static LegalVal legalizeLoad(
     TypeLegalizationContext*    context,
     LegalVal                    legalPtrVal)
@@ -431,6 +491,48 @@ static LegalVal legalizeLoad(
     }
 }
 
+static LegalVal legalizeStore(
+    TypeLegalizationContext*    context,
+    LegalVal                    legalPtrVal,
+    LegalVal                    legalVal)
+{
+    switch (legalPtrVal.flavor)
+    {
+    case LegalVal::Flavor::simple:
+    {
+        context->builder->emitStore(legalPtrVal.getSimple(), legalVal.getSimple());
+        return legalVal;
+    }
+    break;
+
+    case LegalVal::Flavor::implicitDeref:
+        // TODO: what is the right behavior here?
+        if (legalVal.flavor == LegalVal::Flavor::implicitDeref)
+            return legalizeStore(context, legalPtrVal.getImplicitDeref(), legalVal.getImplicitDeref());
+        else
+            return legalizeStore(context, legalPtrVal.getImplicitDeref(), legalVal);
+
+    case LegalVal::Flavor::tuple:
+    {
+        // We need to emit a store for each element of
+        // the tuple.
+        auto destTuple = legalPtrVal.getTuple();
+        auto valTuple = legalVal.getTuple();
+        SLANG_ASSERT(destTuple->elements.Count() == valTuple->elements.Count());
+        for (UInt i = 0; i < valTuple->elements.Count(); i++)
+        {
+            legalizeStore(context, destTuple->elements[i].val, valTuple->elements[i].val);
+        }
+        return legalVal;
+    }
+    break;
+
+    default:
+        SLANG_UNEXPECTED("unhandled case");
+        break;
+    }
+}
+
 static LegalVal legalizeFieldAddress(
     TypeLegalizationContext*    context,
     LegalType                   type,
@@ -492,6 +594,12 @@ static LegalVal legalizeInst(
     case kIROp_FieldAddress:
         return legalizeFieldAddress(context, type, args[0], args[1]);
 
+    case kIROp_Store:
+        return legalizeStore(context, args[0], args[1]);
+
+    case kIROp_Call:
+        return legalizeCall(context, (IRCall*)inst);
+
     default:
         // TODO: produce a user-visible diagnostic here
         SLANG_UNEXPECTED("non-simple operand(s)!");
@@ -499,10 +607,74 @@ static LegalVal legalizeInst(
     }
 }
 
+RefPtr<VarLayout> findVarLayout(IRValue* value)
+{
+    if (auto layoutDecoration = value->findDecoration<IRLayoutDecoration>())
+        return layoutDecoration->layout.As<VarLayout>();
+    return nullptr;
+}
+
+static LegalVal legalizeLocalVar(
+    TypeLegalizationContext*    context,
+    IRVar*                irLocalVar)
+{
+    // Legalize the type for the variable's value
+    auto legalValueType = legalizeType(
+        context,
+        irLocalVar->getType()->getValueType());
+
+    RefPtr<VarLayout> varLayout = findVarLayout(irLocalVar);
+    RefPtr<TypeLayout> typeLayout = varLayout ? varLayout->typeLayout : nullptr;
+
+    // If we've decided to do implicit deref on the type,
+    // then go ahead and declare a value of the pointed-to type.
+    LegalType maybeSimpleType = legalValueType;
+    while (maybeSimpleType.flavor == LegalType::Flavor::implicitDeref)
+    {
+        maybeSimpleType = maybeSimpleType.getImplicitDeref()->valueType;
+    }
+
+    switch (maybeSimpleType.flavor)
+    {
+    case LegalType::Flavor::simple:
+        // Easy case: the type is usable as-is, and we
+        // should just do that.
+        irLocalVar->type = context->session->getPtrType(
+            maybeSimpleType.getSimple());
+        return LegalVal::simple(irLocalVar);
+
+    default:
+    {
+        context->insertBeforeLocalVar = irLocalVar;
+
+        LegalVarChain* varChain = nullptr;
+        LegalVarChain varChainStorage;
+        if (varLayout)
+        {
+            varChainStorage.next = nullptr;
+            varChainStorage.varLayout = varLayout;
+            varChain = &varChainStorage;
+        }
+
+        LegalVal newVal = declareVars(context, kIROp_Var, legalValueType, typeLayout, varChain);
+
+        // Remove the old local var.
+        irLocalVar->removeFromParent();
+        // add old local var to list
+        context->oldLocalVars.Add(irLocalVar);
+        return newVal;
+    }
+    break;
+    }
+}
+
 static LegalVal legalizeInst(
     TypeLegalizationContext*    context,
     IRInst*                     inst)
 {
+    if (inst->op == kIROp_Var)
+        return legalizeLocalVar(context, (IRVar*)inst);
+
     // Need to legalize all the operands.
     auto argCount = inst->getArgCount();
     List<LegalVal> legalArgs;
@@ -567,40 +739,87 @@ static LegalVal legalizeInst(
     return legalVal;
 }
 
+static void addParamType(IRFuncType * ftype, LegalType t)
+{
+    switch (t.flavor)
+    {
+    case LegalType::Flavor::simple:
+        ftype->paramTypes.Add(t.obj.As<Type>());
+        break;
+    case LegalType::Flavor::implicitDeref:
+    {
+        auto imp = t.obj.As<ImplicitDerefType>();
+        addParamType(ftype, imp->valueType);
+        break;
+    }
+    case LegalType::Flavor::tuple:
+    {
+        auto tup = t.obj.As<TupleType>();
+        for (auto & elem : tup->elements)
+            addParamType(ftype, elem.type);
+    }
+    break;
+    default:
+        SLANG_ASSERT(false);
+    }
+}
+
 static void legalizeFunc(
     TypeLegalizationContext*    context,
     IRFunc*                     irFunc)
 {
     // Overwrite the function's type with
     // the result of legalization.
-    irFunc->type = legalizeSimpleType(context, irFunc->type);
-
+    auto newFuncType = new IRFuncType();
+    newFuncType->setSession(context->session);
+    auto oldFuncType = irFunc->type.As<IRFuncType>();
+    newFuncType->resultType = legalizeSimpleType(context, oldFuncType->resultType);
+    for (auto & paramType : oldFuncType->paramTypes)
+    {
+        auto legalParamType = legalizeType(context, paramType);
+        addParamType(newFuncType, legalParamType);
+    }
+    irFunc->type = newFuncType;
+    List<LegalVal> paramVals;
+    List<IRValue*> oldParams;
+
+    // we use this list to store replaced local var insts.
+    // these old instructions will be freed when we are done.
+    context->oldLocalVars.Clear();
+    
     // Go through the blocks of the function
     for (auto bb = irFunc->getFirstBlock(); bb; bb = bb->getNextBlock())
     {
         // Legalize the parameters of the block, which may
         // involve increasing the number of parameters
-        for (auto pp = bb->getFirstParam(); pp; pp = pp->getNextParam())
+        for (auto pp = bb->getFirstParam(); pp; pp = pp->nextParam)
         {
             auto legalParamType = legalizeType(context, pp->getType());
-
-            switch (legalParamType.flavor)
+            if (legalParamType.flavor != LegalType::Flavor::simple)
             {
-            case LegalType::Flavor::simple:
-                // The type is simple, so we can just rewrite it in place
-                pp->type = legalParamType.getSimple();
-                break;
+                context->insertBeforeParam = pp;
+                context->builder->curBlock = nullptr;
 
-            default:
-                // We have something like a tuple, and will need
-                // to expand into multiple parameters now.
-                SLANG_UNEXPECTED("need to handle it!");
-                break;
+                auto paramVal = declareVars(context, kIROp_Param, legalParamType, nullptr, nullptr);
+                paramVals.Add(paramVal);
+                if (pp == bb->getFirstParam())
+                {
+                    bb->firstParam = pp;
+                    while (bb->firstParam->prevParam)
+                        bb->firstParam = bb->firstParam->prevParam;
+                }
+                bb->lastParam = pp->prevParam;
+                if (pp->prevParam)
+                    pp->prevParam->nextParam = pp->nextParam;
+                if (pp->nextParam)
+                    pp->nextParam->prevParam = pp->prevParam;
+                auto oldParam = pp;
+                oldParams.Add(oldParam);
+                registerLegalizedValue(context, oldParam, paramVal);
             }
-
+           
         }
 
-
         // Now legalize the instructions inside the block
         IRInst* nextInst = nullptr;
         for (auto ii = bb->getFirstInst(); ii; ii = nextInst)
@@ -611,18 +830,17 @@ static void legalizeFunc(
 
             registerLegalizedValue(context, ii, legalVal);
         }
+
+    }
+    for (auto & op : oldParams)
+    {
+        SLANG_ASSERT(op->firstUse == nullptr || op->firstUse->nextUse == nullptr);
+        op->deallocate();
     }
+    for (auto & lv : context->oldLocalVars)
+        lv->deallocate();
 }
 
-// Represents the "chain" of declarations that
-// were followed to get to a variable that we
-// are now declaring as a leaf variable.
-struct LegalVarChain
-{
-    LegalVarChain*  next;
-    VarLayout*      varLayout;
-};
-
 static LegalVal declareSimpleVar(
     TypeLegalizationContext*    context,
     IROp                        op,
@@ -696,7 +914,37 @@ static LegalVal declareSimpleVar(
             return LegalVal::simple(globalVar);
         }
         break;
+    case kIROp_Var:
+    {
+        IRBuilder* builder = context->builder;
+
+        auto localVar = builder->emitVar(type);
+        localVar->removeFromParent();
+        localVar->insertBefore(context->insertBeforeLocalVar);
+        if (varLayout)
+        {
+            builder->addLayoutDecoration(localVar, varLayout);
+        }
+        return LegalVal::simple(localVar);
+    }
+    break;
+    case kIROp_Param:
+        {
+            IRBuilder* builder = context->builder;
+            auto param = builder->emitParam(type);
+            if (context->insertBeforeParam->prevParam)
+                context->insertBeforeParam->prevParam->nextParam = param;
+            param->prevParam = context->insertBeforeParam->prevParam;
+            param->nextParam = context->insertBeforeParam;
+            context->insertBeforeParam->prevParam = param;
+            if (varLayout)
+            {
+                builder->addLayoutDecoration(param, varLayout);
+            }
 
+            return LegalVal::simple(param);
+        }
+        break;
     default:
         SLANG_UNEXPECTED("unexpected IR opcode");
         break;
@@ -808,13 +1056,6 @@ static LegalVal declareVars(
     }
 }
 
-RefPtr<VarLayout> findVarLayout(IRValue* value)
-{
-    if (auto layoutDecoration = value->findDecoration<IRLayoutDecoration>())
-        return layoutDecoration->layout.As<VarLayout>();
-    return nullptr;
-}
-
 static void legalizeGlobalVar(
     TypeLegalizationContext*    context,
     IRGlobalVar*                irGlobalVar)
diff --git a/source/slang/lower-to-ir.cpp b/source/slang/lower-to-ir.cpp
index bbbc1812b..759c386cd 100644
--- a/source/slang/lower-to-ir.cpp
+++ b/source/slang/lower-to-ir.cpp
@@ -2850,7 +2850,6 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
             IRType* irParamType = irResultType;
             paramTypes.Add(irParamType);
             subBuilder->emitParam(irParamType);
-
             // TODO: we need some way to wire this up to the `newValue`
             // or whatever name we give for that parameter inside
             // the setter body.
diff --git a/tests/compute/func-param-legalize.slang b/tests/compute/func-param-legalize.slang
new file mode 100644
index 000000000..285fcfbb7
--- /dev/null
+++ b/tests/compute/func-param-legalize.slang
@@ -0,0 +1,35 @@
+//TEST(compute):COMPARE_COMPUTE:-xslang -use-ir
+//TEST_INPUT:Texture2D(size=4, content = one) : dxbinding(0),glbinding(0)
+//TEST_INPUT: Sampler : dxbinding(0),glbinding(0,1,2,3,4,5,6)
+//TEST_INPUT: ubuffer(data=[0 0 0 0], stride=4):dxbinding(0),glbinding(0),out
+
+struct Param
+{
+    Texture2D tex;
+    SamplerState samplerState;
+    float base;
+};
+
+Texture2D diffuseMap;
+SamplerState samplerState;
+RWStructuredBuffer<float> outputBuffer;
+
+float4 run(Param p)
+{
+    return p.tex.SampleLevel(p.samplerState, float2(0.0), 0) + p.base;
+}
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    Param p;
+    p.tex = diffuseMap;
+    p.samplerState = samplerState;
+    p.base = -0.5;
+	float4 outVal = run(p);
+	
+    outputBuffer[0] = outVal.x;
+	outputBuffer[1] = outVal.y;
+	outputBuffer[2] = outVal.z;
+	outputBuffer[3] = outVal.w;
+}
\ No newline at end of file
diff --git a/tests/compute/func-param-legalize.slang.expected.txt b/tests/compute/func-param-legalize.slang.expected.txt
new file mode 100644
index 000000000..e4e4c642a
--- /dev/null
+++ b/tests/compute/func-param-legalize.slang.expected.txt
@@ -0,0 +1,4 @@
+3F000000
+3F000000
+3F000000
+3F000000