Fixes for Shader Execution Reordering on VK (#2929)

* Fixes for Shader Execution Reordering on VK

There are some mismatches between the way that hit objects are
handled between the current NVAPI/HLSL and proposed GLSL extensions
for shader execution reordering. These mismatches create complications
for generating valid GLSL/SPIR-V code from input Slang.

Many of the problems that apply to `HitObject` also apply to the
existing `RayQuery<>` type used for "inline" ray tracing.

In the case of `RayQuery<>` we have that for *both* HLSL and
GLSL/SPIR-V:

* A `RayQuery` (or `rayQueryEXT`) is an opaque handle to underlying
  mutable storage

* The storage that backs a `RayQuery` is allocated as part of the
  "defualt constructor" for a local variable declared with type
  `RayQuery`.

* The `RayQuery` API provides numerous operations that mutate the
  storage referred to by the opaque handle.

The key difference between HLSL and GLSL/SPIR-V for the case of a
`RayQuery` amounts to:

* In HLSL, local variables of type `RayQuery` can be assigned to,
  and assignment has by-reference semantics. It is possible to create
  multiple aliased handles to the same underlying storage.

* In GLSL/SPIR-V, local variables of type `rayQueryEXT` cannot be
  assigned to, returned from functions, etc. It is impossible to
  create multiple aliased handles to the same underlying storage.

The case for `HitObject`s is signicantly *more* messy, because:

* In NVAPI/HLSL a `HitObject` is effectively a "value type" in that
  it only exposes constructors, and there is no way to mutate the
  state of a `HitObject` other than by assignment to a variable of that
  type. It makes no semantic difference whether a `HitObject` directly
  stores the value(s), or if it is a handle, since there is no way
  to introduce aliasing of mutable state. Assignment of `HitObject`s
  semantically creates a copy.

* In GLSL/SPIR-V, a `hitObjectNV` is, like a `rayQueryEXT`, a handle
  to underlying mutable state. These handles cannot be assigned,
  returned from functions, etc. There is no way to make a copy of
  a hit object.

This change includes several changes to how *both* `RayQuery<>` and
`HitObject` are implemented, with the intention of getting more cases
to work correctly when compiling for GLSL/SPIR-V, and to set up a
more clear mental model for the semantics we want to give to these
types in Slang, and how those semantics can/should map to our targets.

An overview of important changes:

* Marked a few operations on `RayQuery` as `[mutating]` that
  realistically should have already been that way.

* Marked the `HitObject` type as being non-copyable (an attribute we
  do not currently enforce), and marked the various GLSL operations that
  construct a hit object as having an `out` parameter of the `HitObject`
  type (even if they are nominally specified in GLSL as not writing
  to the correspondign parameter).

* Added a distinct IR opcode (`allocateOpaqueHandle`) to represent the
  implicit allocation that happens when declaring a variable of type
  `HitObject` or `RayQuery`, and made the "implicit constructor" for
  those types map to the new op. This operation took a lot of tweaking
  to get emitting in a reasonable way, and I'm still not 100% sure that
  all of the emission-related logic for it is strictly required
  (or correct).

* Added new IR instructions for `HitObject` and `RayQuery` types, and
  made the stdlib types map to those IR instructions.

* Treat `HitObject` and `RayQuery` as resource types for the purpose
  of our existing pass that specializes calls to functions that have
  outputs of resource type

* Added a new test case that includes a function that returns a
  `HitObject` as its result.

* Many test cases saw slight changes in their output (especially around
  the relative ordering of declarations of `HitObject`s and `RayQuery`s
  with other instructions)

* Remove debugging logic

1 files changed, 90 insertions, 27 deletions

diff --git a/source/slang/slang-emit-c-like.cpp b/source/slang/slang-emit-c-like.cpp
index 75d7b3cf8..c7ea3a09a 100644
--- a/source/slang/slang-emit-c-like.cpp
+++ b/source/slang/slang-emit-c-like.cpp
@@ -1239,10 +1239,12 @@ bool CLikeSourceEmitter::shouldFoldInstIntoUseSites(IRInst* inst)
         {
             return true;
         }
+        if (as<IRHitObjectType>(type))
+        {
+            return true;
+        }
     }
 
-
-
     // If the instruction is at global scope, then it might represent
     // a constant (e.g., the value of an enum case).
     //
@@ -1292,6 +1294,7 @@ bool CLikeSourceEmitter::shouldFoldInstIntoUseSites(IRInst* inst)
     // definition for certain types on certain targets (e.g. `out TriangleStream<T>`
     // for GLSL), so we check this only after all those special cases are
     // considered.
+    //
     if (inst->getOp() == kIROp_undefined)
         return false;
 
@@ -1382,6 +1385,7 @@ bool CLikeSourceEmitter::shouldFoldInstIntoUseSites(IRInst* inst)
     if(inst->getParent() != user->getParent())
         return false;
 
+
     // Now let's look at all the instructions between this instruction
     // and the user. If any of them might have side effects, then lets
     // bail out now.
@@ -1410,6 +1414,25 @@ bool CLikeSourceEmitter::shouldFoldInstIntoUseSites(IRInst* inst)
     if(as<IRUnconditionalBranch>(user))
         return false;
 
+    // HACK: As a special case, an `allocateOpaqueHandle` operation should
+    // only be folded in if its only use is as the operand of a `store`
+    // that will *itself* get peephole merged in as the initial-value expression
+    // of a `var`:
+    //
+    if (inst->getOp() == kIROp_AllocateOpaqueHandle)
+    {
+        auto store = as<IRStore>(user);
+        if (!store) return false;
+        if (store->getVal() != inst) return false;
+
+        auto var = as<IRVar>(store->getPtr());
+        if (!var) return false;
+
+        if(var->getNextInst() != store) return false;
+
+        return true;
+    }
+
     // Okay, if we reach this point then the user comes later in
     // the same block, and there are no instructions with side
     // effects in between, so it seems safe to fold things in.
@@ -1843,6 +1866,7 @@ void CLikeSourceEmitter::defaultEmitInstExpr(IRInst* inst, const EmitOpInfo& inO
 
     case kIROp_undefined:
     case kIROp_DefaultConstruct:
+    case kIROp_AllocateOpaqueHandle:
         m_writer->emit(getName(inst));
         break;
 
@@ -2441,19 +2465,13 @@ void CLikeSourceEmitter::_emitInst(IRInst* inst)
     case kIROp_DefaultConstruct:
         {
             auto type = inst->getDataType();
-            emitType(type, getName(inst));
+            _emitInstAsDefaultInitializedVar(inst, type);
+        }
+        break;
 
-            // On targets that support empty initializers, we will emit it.
-            switch (this->getTarget())
-            {
-            case CodeGenTarget::CPPSource:
-            case CodeGenTarget::HostCPPSource:
-            case CodeGenTarget::PyTorchCppBinding:
-            case CodeGenTarget::CUDASource:
-                m_writer->emit(" = {}");
-                break;
-            }
-            m_writer->emit(";\n");
+    case kIROp_AllocateOpaqueHandle:
+        {
+            _emitAllocateOpaqueHandleImpl(inst);
         }
         break;
 
@@ -2466,17 +2484,8 @@ void CLikeSourceEmitter::_emitInst(IRInst* inst)
 
     case kIROp_Store:
         {
-            if (inst->getPrevInst() == inst->getOperand(0) && inst->getOperand(0)->getOp() == kIROp_Var)
-            {
-                // If we are storing into a var that is defined right before the store, we have
-                // already folded the store in the initialization of the var, so we can skip here.
-                break;
-            }
-            auto prec = getInfo(EmitOp::Assign);
-            emitDereferenceOperand(inst->getOperand(0), leftSide(getInfo(EmitOp::General), prec));
-            m_writer->emit(" = ");
-            emitOperand(inst->getOperand(1), rightSide(prec, getInfo(EmitOp::General)));
-            m_writer->emit(";\n");
+            auto store = cast<IRStore>(inst);
+            emitStore(store);
         }
         break;
 
@@ -2620,6 +2629,51 @@ void CLikeSourceEmitter::_emitInst(IRInst* inst)
     }
 }
 
+void CLikeSourceEmitter::emitStore(IRStore* store)
+{
+    if (store->getPrevInst() == store->getOperand(0) && store->getOperand(0)->getOp() == kIROp_Var)
+    {
+        // If we are storing into a `var` that is defined right before the store, we have
+        // already folded the store in the initialization of the `var`, so we can skip here.
+        //
+        return;
+    }
+    _emitStoreImpl(store);
+}
+
+void CLikeSourceEmitter::_emitStoreImpl(IRStore* store)
+{
+    auto srcVal = store->getVal();
+    auto dstPtr = store->getPtr();
+    auto prec = getInfo(EmitOp::Assign);
+    emitDereferenceOperand(dstPtr, leftSide(getInfo(EmitOp::General), prec));
+    m_writer->emit(" = ");
+    emitOperand(srcVal, rightSide(prec, getInfo(EmitOp::General)));
+    m_writer->emit(";\n");
+}
+
+void CLikeSourceEmitter::_emitInstAsDefaultInitializedVar(IRInst* inst, IRType* type)
+{
+    emitType(type, getName(inst));
+
+    // On targets that support empty initializers, we will emit it.
+    switch (this->getTarget())
+    {
+    case CodeGenTarget::CPPSource:
+    case CodeGenTarget::HostCPPSource:
+    case CodeGenTarget::PyTorchCppBinding:
+    case CodeGenTarget::CUDASource:
+        m_writer->emit(" = {}");
+        break;
+    }
+    m_writer->emit(";\n");
+}
+
+void CLikeSourceEmitter::_emitAllocateOpaqueHandleImpl(IRInst* allocateInst)
+{
+    _emitInstAsDefaultInitializedVar(allocateInst, allocateInst->getDataType());
+}
+
 void CLikeSourceEmitter::emitSemanticsUsingVarLayout(IRVarLayout* varLayout)
 {
     if(auto semanticAttr = varLayout->findAttr<IRSemanticAttr>())
@@ -3415,18 +3469,27 @@ void CLikeSourceEmitter::emitVar(IRVar* varDecl)
 
     emitLayoutSemantics(varDecl);
 
+    // TODO: ideally this logic should scan ahead to see if it can find a `store`
+    // instruction that writes to the `var`, within the same block, such that all
+    // of the intervening instructions are safe to fold.
+    //
     if (auto store = as<IRStore>(varDecl->getNextInst()))
     {
         if (store->getPtr() == varDecl)
         {
-            m_writer->emit(" = ");
-            emitOperand(store->getVal(), getInfo(EmitOp::General));
+            _emitInstAsVarInitializerImpl(store->getVal());
         }
     }
 
     m_writer->emit(";\n");
 }
 
+void CLikeSourceEmitter::_emitInstAsVarInitializerImpl(IRInst* inst)
+{
+    m_writer->emit(" = ");
+    emitOperand(inst, getInfo(EmitOp::General));
+}
+
 void CLikeSourceEmitter::emitGlobalVar(IRGlobalVar* varDecl)
 {
     auto allocatedType = varDecl->getDataType();