Legalization of non-struct when function expects struct, resolves #3840 (#3880)

* Legalization of non-struct when expects struct.

`__forceVarIntoStructTemporarily()` solves the issue of passing "non-struct type's" into a parameter that only accepts "struct type's".
The intrinsic solves the issue through checking the parameter of the intrinsic:
If the parameter is a "struct type"
* Return a reference to the parameter
else
* a "struct type" Temporary variable is made and the "non struct type" parameter is copied to a member of this struct. This struct is then returned by `__forceVarIntoStructTemporarily()`. Optionally if the use location of this call is a argument which can have side effects (out, inout, ref, etc.) the temporary struct variable is copied into the original "non struct type" parameter.

Testing code has "addComplexity" functions to avoid optimizations through forcing side effects so we can predict the code output.

* Address review comments

- ForceInline ray functions
- fix testing
- adjust how we replace operands in senarios to avoid unexpected side effects of replacing operands without any explicit checks

* Adjust nv test slightly and remove .glsl file

* Remove implicit LOD sampling & test additions

- Implicit LOD sampling is not allowed in a raygen. Implicit LOD sampling requires depth (from a fragment shader) to sample. Raygen does not have the depth, so this function was replaced.
- Changed other tests for correctness/clarity

* Test if Falcor breaks through use of ForceInline

* Add back force inline

may need to look at how Falcor wrote its slang shaders. This will be done if ForceInline causes issues since ForceInline should not affect code gen in an impactable way.

5 files changed, 200 insertions, 5 deletions

diff --git a/source/slang/hlsl.meta.slang b/source/slang/hlsl.meta.slang
index cd6de61b5..fa2a7ccb1 100644
--- a/source/slang/hlsl.meta.slang
+++ b/source/slang/hlsl.meta.slang
@@ -9572,6 +9572,26 @@ void CallShader(uint shaderIndex, inout Payload payload)
 
 // 10.3.2
 
+// Some functions only accept a "struct type" parameter. The 
+// following function addresses this issue by transforming non-struct
+// parameters into a struct. 
+// side effect typed use locations (`inout`,`out`, etc.) are managed.
+__generic<T>
+__intrinsic_op($(kIROp_ForceVarIntoStructTemporarily))
+Ref<T> __forceVarIntoStructTemporarily(inout T maybeStruct);
+
+__target_intrinsic(hlsl, "TraceRay")
+[require(hlsl, raytracing)]
+__generic<payload_t>
+void __traceRayHLSL(
+        RaytracingAccelerationStructure AccelerationStructure,
+        uint RayFlags,
+        uint InstanceInclusionMask,
+        uint RayContributionToHitGroupIndex,
+        uint MultiplierForGeometryContributionToHitGroupIndex,
+        uint MissShaderIndex,
+        RayDesc Ray,
+        inout payload_t Payload);
 
 __target_intrinsic(_GL_EXT_ray_tracing, "traceRayEXT")
 [require(glsl, raytracing)]
@@ -9599,6 +9619,7 @@ __generic<Payload>
 __intrinsic_op($(kIROp_GetVulkanRayTracingPayloadLocation))
 int __rayPayloadLocation(__ref Payload payload);
 
+[ForceInline]
 [require(glsl, raytracing)]
 [require(spirv, raytracing)]
 [require(hlsl, raytracing)]
@@ -9615,7 +9636,17 @@ void TraceRay(
 {
     __target_switch
     {
-    case hlsl: __intrinsic_asm "TraceRay";
+    case hlsl:
+        __traceRayHLSL(
+            AccelerationStructure,
+            RayFlags,
+            InstanceInclusionMask,
+            RayContributionToHitGroupIndex,
+            MultiplierForGeometryContributionToHitGroupIndex,
+            MissShaderIndex,
+            Ray,
+            __forceVarIntoStructTemporarily(Payload));
+        return;
     case cuda: __intrinsic_asm "traceOptiXRay";
     case glsl:
     {
@@ -9673,6 +9704,20 @@ void TraceRay(
 //
 // https://github.com/KhronosGroup/GLSL/blob/master/extensions/nv/GLSL_NV_ray_tracing_motion_blur.txt
 
+__target_intrinsic(hlsl, "TraceMotionRay")
+[require(hlsl, raytracing_motionblur)]
+__generic<payload_t>
+void __traceMotionRayHLSL(
+        RaytracingAccelerationStructure AccelerationStructure,
+        uint RayFlags,
+        uint InstanceInclusionMask,
+        uint RayContributionToHitGroupIndex,
+        uint MultiplierForGeometryContributionToHitGroupIndex,
+        uint MissShaderIndex,
+        RayDesc Ray,
+        float CurrentTime,
+        inout payload_t Payload);
+
 __glsl_extension(GL_NV_ray_tracing_motion_blur)
 __target_intrinsic(glsl, "traceRayMotionNV")
 [require(glsl, raytracing_motionblur)]
@@ -9690,6 +9735,7 @@ void __traceMotionRay(
     float                           CurrentTime,
     int                             PayloadLocation);
 
+[ForceInline]
 [require(glsl, raytracing_motionblur)]
 [require(spirv, raytracing_motionblur)]
 [require(hlsl, raytracing_motionblur)]
@@ -9707,7 +9753,18 @@ void TraceMotionRay(
 {
     __target_switch
     {
-    case hlsl: __intrinsic_asm "TraceMotionRay";
+    case hlsl:
+        __traceMotionRayHLSL(
+            AccelerationStructure,
+            RayFlags,
+            InstanceInclusionMask,
+            RayContributionToHitGroupIndex,
+            MultiplierForGeometryContributionToHitGroupIndex,
+            MissShaderIndex,
+            Ray,
+            CurrentTime,
+            __forceVarIntoStructTemporarily(Payload));
+        return;
     case glsl:
     {
         [__vulkanRayPayload]
@@ -11555,7 +11612,7 @@ struct HitObject
                     MultiplierForGeometryContributionToHitGroupIndex, 
                     MissShaderIndex, 
                     Ray, 
-                    Payload,
+                    __forceVarIntoStructTemporarily(Payload),
                     hitObj);
                 return hitObj;
             }
@@ -11641,7 +11698,16 @@ struct HitObject
         __target_switch
         {
         case hlsl:
-            __intrinsic_asm "TraceMotionRay";
+            __traceMotionRayHLSL(
+                AccelerationStructure,
+                RayFlags,
+                InstanceInclusionMask,
+                RayContributionToHitGroupIndex,
+                MultiplierForGeometryContributionToHitGroupIndex,
+                MissShaderIndex,
+                Ray,
+                CurrentTime,
+                __forceVarIntoStructTemporarily(Payload));
         case glsl:
             {
                 [__vulkanRayPayload]
@@ -12143,6 +12209,14 @@ struct HitObject
         }
     }
 
+    __target_intrinsic(hlsl, "NvInvokeHitObject")
+    [require(hlsl, ser)]
+    __generic<payload_t>
+    static void __InvokeHLSL(
+        RaytracingAccelerationStructure AccelerationStructure,
+        HitObject HitOrMiss,
+        inout payload_t Payload);
+
         /// Invokes closesthit or miss shading for the specified hit object. In case of a NOP HitObject, no
         /// shader is invoked.
     [__requiresNVAPI]
@@ -12157,7 +12231,11 @@ struct HitObject
     {
         __target_switch
         {
-        case hlsl: __intrinsic_asm "NvInvokeHitObject";
+        case hlsl:
+            __InvokeHLSL(
+                AccelerationStructure,
+                HitOrMiss,
+                __forceVarIntoStructTemporarily(Payload));
         case glsl:
             {
                 [__vulkanRayPayload]
diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index 70dddb4f6..f724b1941 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -28,6 +28,7 @@
 #include "slang-ir-explicit-global-init.h"
 #include "slang-ir-fuse-satcoop.h"
 #include "slang-ir-glsl-legalize.h"
+#include "slang-ir-hlsl-legalize.h"
 #include "slang-ir-insts.h"
 #include "slang-ir-inline.h"
 #include "slang-ir-legalize-array-return-type.h"
@@ -839,6 +840,10 @@ Result linkAndOptimizeIR(
         break;
     }
 
+    // Legalize non struct parameters that are expected to be structs for HLSL. 
+    if(isD3DTarget(targetRequest))
+        legalizeNonStructParameterToStructForHLSL(irModule);
+
     // Legalize `ImageSubscript` and constant buffer loads for GLSL.
     switch (target)
     {
diff --git a/source/slang/slang-ir-hlsl-legalize.cpp b/source/slang/slang-ir-hlsl-legalize.cpp
new file mode 100644
index 000000000..ccd480d8c
--- /dev/null
+++ b/source/slang/slang-ir-hlsl-legalize.cpp
@@ -0,0 +1,93 @@
+// slang-ir-hlsl-legalize.cpp
+#include "slang-ir-hlsl-legalize.h"
+
+#include <functional>
+
+#include "slang-ir.h"
+#include "slang-ir-insts.h"
+#include "slang-ir-inst-pass-base.h"
+#include "slang-ir-specialize-function-call.h"
+#include "slang-ir-util.h"
+
+namespace Slang
+{
+
+void searchChildrenForForceVarIntoStructTemporarily(IRModule* module, IRInst* inst)
+{
+    for(auto child : inst->getChildren())
+    {
+        switch(child->getOp())
+        {
+        case kIROp_Block:
+        {
+            searchChildrenForForceVarIntoStructTemporarily(module, child);
+            break;
+        }
+        case kIROp_Call:
+        {
+            auto call = as<IRCall>(child);
+            for(UInt i = 0; i < call->getArgCount(); i++)
+            {
+                auto arg = call->getArg(i);
+                if(arg->getOp() != kIROp_ForceVarIntoStructTemporarily)
+                    continue;
+                auto forceStructArg = arg->getOperand(0);
+                auto forceStructBaseType = as<IRType>(forceStructArg->getDataType()->getOperand(0));
+                if(forceStructBaseType->getOp() == kIROp_StructType)
+                {
+                    call->setArg(i, arg->getOperand(0));
+                    continue;
+                }
+
+                // When `__forceVarIntoStructTemporarily` is called with a non-struct type parameter, 
+                // we create a temporary struct and copy the parameter into the struct.
+                // This struct is then subsituted for the return of `__forceVarIntoStructTemporarily`.
+                // Optionally, if `__forceVarIntoStructTemporarily` is a parameter to a side effect type 
+                // (`ref`, `out`, `inout`) we copy the struct back into our original non-struct parameter.
+                IRBuilder builder(call);
+
+                builder.setInsertBefore(call->getCallee());
+                auto structType = builder.createStructType();
+                StringBuilder structName;
+                builder.addNameHintDecoration(structType, UnownedStringSlice("ForceVarIntoStructTemporarily_t"));
+
+                auto elementBufferKey = builder.createStructKey();
+                builder.addNameHintDecoration(elementBufferKey, UnownedStringSlice("data"));
+                auto _dataField = builder.createStructField(structType, elementBufferKey, forceStructBaseType);
+
+                builder.setInsertBefore(call);
+                auto structVar = builder.emitVar(structType);
+                builder.addNameHintDecoration(structVar, UnownedStringSlice("forceVarIntoStructTemporarily"));
+                builder.emitStore(
+                    builder.emitFieldAddress(builder.getPtrType(_dataField->getFieldType()), structVar, _dataField->getKey()),
+                    builder.emitLoad(forceStructArg));
+
+                arg->replaceUsesWith(structVar);
+                arg->removeAndDeallocate();
+
+                auto argType = call->getCallee()->getDataType()->getOperand(i+1);
+                if (!isPtrLikeOrHandleType(argType))
+                    continue;
+
+                builder.setInsertAfter(call);
+                builder.emitStore(
+                    forceStructArg,
+                    builder.emitFieldAddress(builder.getPtrType(_dataField->getFieldType()), structVar, _dataField->getKey()));
+            }
+            break;
+        }
+        }
+    }
+}
+
+void legalizeNonStructParameterToStructForHLSL(IRModule* module)
+{
+    for(auto globalInst : module->getGlobalInsts())
+    {
+        if (globalInst->getOp() != kIROp_Func)
+            continue;
+        searchChildrenForForceVarIntoStructTemporarily(module, globalInst);
+    }
+}
+
+} // namespace Slang
diff --git a/source/slang/slang-ir-hlsl-legalize.h b/source/slang/slang-ir-hlsl-legalize.h
new file mode 100644
index 000000000..3970dc364
--- /dev/null
+++ b/source/slang/slang-ir-hlsl-legalize.h
@@ -0,0 +1,17 @@
+// slang-ir-hlsl-legalize.h
+#pragma once
+#include"../core/slang-list.h"
+#include "slang-compiler.h"
+
+namespace Slang
+{
+
+class DiagnosticSink;
+class Session;
+
+struct IRFunc;
+struct IRModule;
+
+void legalizeNonStructParameterToStructForHLSL(IRModule* module);
+
+}
\ No newline at end of file
diff --git a/source/slang/slang-ir-inst-defs.h b/source/slang/slang-ir-inst-defs.h
index e51f6dc95..25f331708 100644
--- a/source/slang/slang-ir-inst-defs.h
+++ b/source/slang/slang-ir-inst-defs.h
@@ -664,6 +664,8 @@ INST(GetVulkanRayTracingPayloadLocation, GetVulkanRayTracingPayloadLocation, 1,
 
 INST(GetLegalizedSPIRVGlobalParamAddr, kIROp_GetLegalizedSPIRVGlobalParamAddr, 1, 0)
 
+INST(ForceVarIntoStructTemporarily, ForceVarIntoStructTemporarily, 1, 0)
+
 INST(MakeArrayList, makeArrayList, 0, 0)
 INST(MakeTensorView, makeTensorView, 0, 0)
 INST(AllocateTorchTensor, allocTorchTensor, 0, 0)