Legalize array return types. (#2463)

3 files changed, 117 insertions, 0 deletions

diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index bab33e79d..896171f32 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -23,6 +23,7 @@
 #include "slang-ir-glsl-legalize.h"
 #include "slang-ir-insts.h"
 #include "slang-ir-inline.h"
+#include "slang-ir-legalize-array-return-type.h"
 #include "slang-ir-legalize-varying-params.h"
 #include "slang-ir-link.h"
 #include "slang-ir-com-interface.h"
@@ -525,6 +526,10 @@ Result linkAndOptimizeIR(
         simplifyIR(irModule);
     }
 
+    // Rewrite functions that return arrays to return them via `out` parameter,
+    // since our target languages doesn't allow returning arrays.
+    legalizeArrayReturnType(irModule);
+
 #if 0
     dumpIRIfEnabled(codeGenContext, irModule, "AFTER RESOURCE SPECIALIZATION");
 #endif
diff --git a/source/slang/slang-ir-legalize-array-return-type.cpp b/source/slang/slang-ir-legalize-array-return-type.cpp
new file mode 100644
index 000000000..711b21f1c
--- /dev/null
+++ b/source/slang/slang-ir-legalize-array-return-type.cpp
@@ -0,0 +1,99 @@
+#include "slang-ir-legalize-array-return-type.h"
+#include "slang-ir.h"
+#include "slang-ir-insts.h"
+#include "slang-ir-clone.h"
+
+namespace Slang
+{
+
+void makeFuncReturnViaOutParam(IRBuilder& builder, IRFunc* func)
+{
+    auto funcType = as<IRFuncType>(func->getFullType());
+    if (!funcType)
+        return;
+    auto arrayType = funcType->getResultType();
+    builder.setInsertBefore(funcType);
+    List<IRType*> paramTypes;
+    for (UInt i = 0; i < funcType->getParamCount(); i++)
+    {
+        paramTypes.add(funcType->getParamType(i));
+    }
+    auto outParamType = builder.getOutType(funcType->getResultType());
+    paramTypes.add(outParamType);
+
+    auto newFuncType = builder.getFuncType(paramTypes, builder.getVoidType());
+    func->setFullType(newFuncType);
+    auto firstBlock = func->getFirstBlock();
+    builder.setInsertInto(firstBlock);
+    auto outParam = builder.emitParam(outParamType);
+
+    // Collect return insts.
+    List<IRReturn*> returnInsts;
+    for (auto block : func->getBlocks())
+    {
+        for (auto inst : block->getChildren())
+        {
+            if (inst->getOp() == kIROp_Return)
+            {
+                returnInsts.add(as<IRReturn>(inst));
+            }
+        }
+    }
+
+    // Rewrite return inst into a store + return void.
+    for (auto returnInst : returnInsts)
+    {
+        builder.setInsertBefore(returnInst);
+        builder.emitStore(outParam, returnInst->getVal());
+        builder.emitReturn();
+        SLANG_RELEASE_ASSERT(!returnInst->hasUses());
+        returnInst->removeAndDeallocate();
+    }
+
+    // Rewrite call sites.
+    List<IRCall*> callSites;
+    for (auto use = func->firstUse; use; use = use->nextUse)
+    {
+        if (auto call = as<IRCall>(use->getUser()))
+        {
+            if (call->getCallee() == func)
+                callSites.add(call);
+        }
+    }
+    for (auto call : callSites)
+    {
+        builder.setInsertBefore(call);
+        auto tmpVar = builder.emitVar(arrayType);
+        List<IRInst*> args;
+        for (UInt i = 0; i < call->getArgCount(); i++)
+        {
+            args.add(call->getArg(i));
+        }
+        args.add(tmpVar);
+        builder.emitCallInst(builder.getVoidType(), func, args);
+        auto load = builder.emitLoad(tmpVar);
+        call->replaceUsesWith(load);
+        call->removeAndDeallocate();
+    }
+
+}
+
+void legalizeArrayReturnType(IRModule* module)
+{
+    SharedIRBuilder builderStorage;
+    builderStorage.init(module);
+    builderStorage.deduplicateAndRebuildGlobalNumberingMap();
+    IRBuilder builder(&builderStorage);
+
+    for (auto inst : module->getGlobalInsts())
+    {
+        if (auto func = as<IRFunc>(inst))
+        {
+            if (func->getResultType()->getOp() == kIROp_ArrayType)
+            {
+                makeFuncReturnViaOutParam(builder, func);
+            }
+        }
+    }
+}
+}
diff --git a/source/slang/slang-ir-legalize-array-return-type.h b/source/slang/slang-ir-legalize-array-return-type.h
new file mode 100644
index 000000000..0086f7827
--- /dev/null
+++ b/source/slang/slang-ir-legalize-array-return-type.h
@@ -0,0 +1,13 @@
+// slang-ir-legalize-array-return-type.h
+#pragma once
+
+#include "slang-ir-insts.h"
+
+namespace Slang
+{
+    struct IRModule;
+
+    // Turn array-typed return values into `out` parameters for backends that does not
+    // support arrays in return values.
+    void legalizeArrayReturnType(IRModule* module);
+}