GPU Foreach Loop (#1498)

* GPU Foreach Loop

This PR introduces the completed GPU foreach loop and updates the
heterogeneous-hello-world example to use it.  This PR builds on the
previous introduction of the GPU Foreach loop parsing and semantic
checking PR (#1482) by introducing IR lowering and emmitting.  THe new
feature can be used by having a GPU_Foreach loop interacting with a
named non-CPP entry point, and using the -heterogeneous flag.

* Fix to path

Co-authored-by: Tim Foley <tfoleyNV@users.noreply.github.com>

7 files changed, 113 insertions, 18 deletions

diff --git a/source/slang/slang-emit-c-like.cpp b/source/slang/slang-emit-c-like.cpp
index 0708bc9a8..9e9117dcd 100644
--- a/source/slang/slang-emit-c-like.cpp
+++ b/source/slang/slang-emit-c-like.cpp
@@ -2425,6 +2425,31 @@ void CLikeSourceEmitter::defaultEmitInstExpr(IRInst* inst, const EmitOpInfo& inO
         m_writer->emit(")");
         break;
     }
+    case kIROp_GpuForeach:
+    {
+        auto operand = inst->getOperand(2);
+        if (as<IRFunc>(operand))
+        {
+            //emitOperand(operand->findDecoration<IREntryPointDecoration>(), getInfo(EmitOp::General));
+            emitOperand(operand, getInfo(EmitOp::General));
+        }
+        else
+        {
+            SLANG_UNEXPECTED("Expected 3rd operand to be a function");
+        }
+        m_writer->emit("_wrapper(");
+        emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+        m_writer->emit(", ");
+        emitOperand(inst->getOperand(1), getInfo(EmitOp::General));
+        UInt argCount = inst->getOperandCount();
+        for (UInt aa = 3; aa < argCount; ++aa)
+        {
+            m_writer->emit(", ");
+            emitOperand(inst->getOperand(aa), getInfo(EmitOp::General));
+        }
+        m_writer->emit(")");
+        break;
+    }
     default:
         diagnoseUnhandledInst(inst);
         break;
diff --git a/source/slang/slang-emit-cpp.cpp b/source/slang/slang-emit-cpp.cpp
index ebb6bda2d..0c2b4cd93 100644
--- a/source/slang/slang-emit-cpp.cpp
+++ b/source/slang/slang-emit-cpp.cpp
@@ -2277,8 +2277,8 @@ void CPPSourceEmitter::emitPreprocessorDirectivesImpl()
     {
         // Put all into an anonymous namespace
         // This includes any generated types, and generated intrinsics
-
-        m_writer->emit("namespace { // anonymous \n\n");
+        if (!m_compileRequest->getLinkage()->m_heterogeneous)
+            m_writer->emit("namespace { // anonymous \n\n");
         m_writer->emit("#ifdef SLANG_PRELUDE_NAMESPACE\n");
         m_writer->emit("using namespace SLANG_PRELUDE_NAMESPACE;\n");
         m_writer->emit("#endif\n\n");
@@ -2556,11 +2556,11 @@ void CPPSourceEmitter::emitModuleImpl(IRModule* module)
                 continue;
             if (auto entryPointDecoration = func->findDecoration<IREntryPointDecoration>())
             {
-                String someName = entryPointDecoration->getName()->getStringSlice();
+                String entryPointName  = entryPointDecoration->getName()->getStringSlice();
                 for (int index = 0; index < program->getEntryPointCount(); index++)
                 {
                     auto entryPoint = program->getEntryPoint(index);
-                    if (someName == entryPoint->getName()->text)
+                    if (entryPointName  == entryPoint->getName()->text)
                     {
                         for (auto targetRequest : linkage->targets)
                         {
@@ -2586,13 +2586,13 @@ void CPPSourceEmitter::emitModuleImpl(IRModule* module)
                                 auto ptr = (const unsigned char*)blob->getBufferPointer();
 
                                 m_writer->emit("size_t __");
-                                m_writer->emit(someName);
+                                m_writer->emit(entryPointName );
                                 m_writer->emit("Size = ");
                                 m_writer->emitInt64(blob->getBufferSize());
                                 m_writer->emit(";\n");
 
                                 m_writer->emit("unsigned char __");
-                                m_writer->emit(someName);
+                                m_writer->emit(entryPointName );
                                 m_writer->emit("[] = {");
                                 for (unsigned int i = 0; i < blob->getBufferSize() - 1; i++) {
                                     m_writer->emitUInt64(ptr[i]);
@@ -2602,6 +2602,26 @@ void CPPSourceEmitter::emitModuleImpl(IRModule* module)
                                 m_writer->emit("};\n");
                             }
                         }
+                        // Emit a wrapper function for calling the shader blob
+                        m_writer->emit("void ");
+                        m_writer->emit(entryPointName);
+                        m_writer->emit("_wrapper(gfx_Renderer_0* renderer, Vector<uint32_t, 3> gridDims, \n");
+                        m_writer->emit("\tRWStructuredBuffer<float> buffer)\n{");
+                        m_writer->emit("\n\tgfx_ShaderProgram_0* shaderProgram = loadShaderProgram_0(renderer, __");
+                        m_writer->emit(entryPointName);
+                        m_writer->emit(", __");
+                        m_writer->emit(entryPointName);
+                        m_writer->emit("Size);");
+                        m_writer->emit("\n\tgfx_DescriptorSetLayout_0* setLayout = buildDescriptorSetLayout_0(renderer);");
+                        m_writer->emit("\n\tgfx_PipelineLayout_0* pipelineLayout = buildPipeline_0(renderer, setLayout);");
+                        m_writer->emit("\n\tgfx_DescriptorSet_0* descriptorSet = ");
+                        m_writer->emit("buildDescriptorSet_0(renderer, setLayout, unconvertBuffer_0(buffer));");
+                        m_writer->emit("\n\tgfx_PipelineState_0* pipelineState = ");
+                        m_writer->emit("buildPipelineState_0(shaderProgram, renderer, pipelineLayout);");
+
+                        m_writer->emit("\n\tdispatchComputation_0(renderer, pipelineState, pipelineLayout, ");
+                        m_writer->emit("descriptorSet, gridDims.x, gridDims.y, gridDims.z);");
+                        m_writer->emit("\n}\n");
                     }
                 }
             }
@@ -2645,8 +2665,8 @@ void CPPSourceEmitter::emitModuleImpl(IRModule* module)
     if (m_target == CodeGenTarget::CPPSource)
     {
         // Need to close the anonymous namespace when outputting for C++
-
-        m_writer->emit("} // anonymous\n\n");
+        if (!linkage->m_heterogeneous)
+            m_writer->emit("} // anonymous\n\n");
     }
 
      // Finally we need to output dll entry points
diff --git a/source/slang/slang-ir-inst-defs.h b/source/slang/slang-ir-inst-defs.h
index 3761828b9..74bd15531 100644
--- a/source/slang/slang-ir-inst-defs.h
+++ b/source/slang/slang-ir-inst-defs.h
@@ -467,6 +467,9 @@ INST(SampleGrad, sampleGrad, 4, 0)
 
 INST(GroupMemoryBarrierWithGroupSync, GroupMemoryBarrierWithGroupSync, 0, 0)
 
+// GPU_FOREACH loop of the form 
+INST(GpuForeach, gpuForeach, 3, 0)
+
 /* Decoration */
 
 INST(HighLevelDeclDecoration,               highLevelDecl,          1, 0)
diff --git a/source/slang/slang-ir-insts.h b/source/slang/slang-ir-insts.h
index 55af7db78..a66ca8f7a 100644
--- a/source/slang/slang-ir-insts.h
+++ b/source/slang/slang-ir-insts.h
@@ -1946,6 +1946,8 @@ struct IRBuilder
         return emitWrapExistential(type, value, slotArgCount, slotArgVals.getBuffer());
     }
 
+    IRInst* emitGpuForeach(List<IRInst*> args);
+
     IRUndefined* emitUndefined(IRType* type);
 
     IRInst* findOrAddInst(
diff --git a/source/slang/slang-ir-link.cpp b/source/slang/slang-ir-link.cpp
index 9de2b5d1c..274fbf6d9 100644
--- a/source/slang/slang-ir-link.cpp
+++ b/source/slang/slang-ir-link.cpp
@@ -1491,19 +1491,21 @@ LinkedIR linkIR(
             cloneValue(context, bindInst);
         }
     }
-
-    for (IRModule* irModule : irModules)
+    if (target == CodeGenTarget::CPPSource)
     {
-        for (auto inst : irModule->getGlobalInsts())
+        for (IRModule* irModule : irModules)
         {
-            auto hasPublic = inst->findDecoration<IRPublicDecoration>();
-            if (!hasPublic)
-                continue;
-
-            auto cloned = cloneValue(context, inst);
-            if (!cloned->findDecorationImpl(kIROp_KeepAliveDecoration))
+            for (auto inst : irModule->getGlobalInsts())
             {
-                context->builder->addKeepAliveDecoration(cloned);
+                auto hasPublic = inst->findDecoration<IRPublicDecoration>();
+                if (!hasPublic)
+                    continue;
+
+                auto cloned = cloneValue(context, inst);
+                if (!cloned->findDecorationImpl(kIROp_KeepAliveDecoration))
+                {
+                    context->builder->addKeepAliveDecoration(cloned);
+                }
             }
         }
     }
diff --git a/source/slang/slang-ir.cpp b/source/slang/slang-ir.cpp
index 13840a84a..2a1db6310 100644
--- a/source/slang/slang-ir.cpp
+++ b/source/slang/slang-ir.cpp
@@ -3752,6 +3752,17 @@ namespace Slang
         return inst;
     }
 
+    IRInst* IRBuilder::emitGpuForeach(List<IRInst*> args)
+    {
+        auto inst = createInst<IRInst>(
+            this,
+            kIROp_GpuForeach,
+            getVoidType(),
+            args.getCount(),
+            args.getBuffer());
+        addInst(inst);
+        return inst;
+    }
 
     //
     // Decorations
diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index 9c4808f31..864491f7e 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -4116,7 +4116,39 @@ struct StmtLoweringVisitor : StmtVisitor<StmtLoweringVisitor>
 
     void visitGpuForeachStmt(GpuForeachStmt* stmt)
     {
+        auto builder = getBuilder();
         startBlockIfNeeded(stmt);
+
+        auto renderer = getSimpleVal(context, lowerRValueExpr(context, stmt->renderer));
+        auto gridDims = getSimpleVal(context, lowerRValueExpr(context, stmt->gridDims));
+
+        List<IRInst*> irArgs;
+        if (auto callExpr = as<InvokeExpr>(stmt->kernelCall))
+        {
+            irArgs.add(renderer);
+            irArgs.add(gridDims);
+            auto fref = getSimpleVal(context, lowerRValueExpr(context, callExpr->functionExpr));
+            irArgs.add(fref);
+            for (auto arg : callExpr->arguments)
+            {
+                // if a reference to dispatchThreadID, don't emit
+                if (auto declRefExpr = as<DeclRefExpr>(arg))
+                {
+                    if (declRefExpr->declRef.getDecl() == stmt->dispatchThreadID)
+                    {
+                        continue;
+                    }
+                }
+                auto irArg = getSimpleVal(context, lowerRValueExpr(context, arg));
+                irArgs.add(irArg);
+            }
+        }
+        else
+        {
+            SLANG_UNEXPECTED("GPUForeach parsing produced an invalid result");
+        }
+
+        builder->emitGpuForeach(irArgs);
         return;
     }