Direct SPIRV for RayTracing Inline. (#3201)

* Direct SPIRV for RayTracing Inline.

* Fix.

* Cleanup.

* Fix compile error.

---------

Co-authored-by: Yong He <yhe@nvidia.com>

8 files changed, 156 insertions, 59 deletions

diff --git a/source/slang/hlsl.meta.slang b/source/slang/hlsl.meta.slang
index 752c99a5b..dce44c9cc 100644
--- a/source/slang/hlsl.meta.slang
+++ b/source/slang/hlsl.meta.slang
@@ -7765,26 +7765,7 @@ struct RayQuery <let rayFlagsGeneric : RAY_FLAG = RAY_FLAG_NONE>
     __intrinsic_op($(kIROp_AllocateOpaqueHandle))
     __init();
 
-    // Initialize a ray-tracing query.
-    //
-    // This method may be called on a "fresh" ray query, or
-    // on one that is already tracing a ray. In the latter
-    // case any state related to the ray previously being
-    // traced is overwritten.
-    //
-    // The `rayFlags` here will be bitwise ORed with
-    // the `rayFlags` passed as a generic argument to
-    // `RayQuery` to get the effective ray flags, which
-    // must obey any API-imposed restrictions.
-    //
-    __target_intrinsic(hlsl)
-    [mutating]
-    void TraceRayInline(
-        RaytracingAccelerationStructure accelerationStructure,
-        RAY_FLAG                        rayFlags,
-        uint                            instanceInclusionMask,
-        RayDesc                         ray);
-
+    
     __target_intrinsic(glsl, "rayQueryInitializeEXT($0, $1, $2, $3, $4, $5, $6, $7)")
     __glsl_extension(GL_EXT_ray_query)
     __glsl_version(460)
@@ -7796,10 +7777,31 @@ struct RayQuery <let rayFlagsGeneric : RAY_FLAG = RAY_FLAG_NONE>
         float3                          origin,
         float                           tMin,
         float3                          direction,
-        float                           tMax);
+        float                           tMax)
+    {
+        __target_switch
+        {
+        case glsl: __intrinsic_asm "rayQueryInitializeEXT($0, $1, $2, $3, $4, $5, $6, $7)";
+        case spirv:
+            spirv_asm {
+                OpRayQueryInitializeKHR &this $accelerationStructure $rayFlags $instanceInclusionMask $origin $tMin $direction $tMax;
+            };
+        }
+    }
 
+    // Initialize a ray-tracing query.
+    //
+    // This method may be called on a "fresh" ray query, or
+    // on one that is already tracing a ray. In the latter
+    // case any state related to the ray previously being
+    // traced is overwritten.
+    //
+    // The `rayFlags` here will be bitwise ORed with
+    // the `rayFlags` passed as a generic argument to
+    // `RayQuery` to get the effective ray flags, which
+    // must obey any API-imposed restrictions.
+    //
     [__unsafeForceInlineEarly]
-    __specialized_for_target(glsl)
     [mutating]
     void TraceRayInline(
         RaytracingAccelerationStructure accelerationStructure,
@@ -7807,16 +7809,22 @@ struct RayQuery <let rayFlagsGeneric : RAY_FLAG = RAY_FLAG_NONE>
         uint                            instanceInclusionMask,
         RayDesc                         ray)
     {
-        __rayQueryInitializeEXT(
-            accelerationStructure,
-            rayFlags | rayFlagsGeneric,
-            instanceInclusionMask,
-            ray.Origin,
-            ray.TMin,
-            ray.Direction,
-            ray.TMax);
+        __target_switch
+        {
+        case hlsl: __intrinsic_asm ".TraceRayInline";
+        case glsl:
+        case spirv:
+            __rayQueryInitializeEXT(
+                accelerationStructure,
+                rayFlags | rayFlagsGeneric,
+                instanceInclusionMask,
+                ray.Origin,
+                ray.TMin,
+                ray.Direction,
+                ray.TMax);
+        }
     }
-
+    
     // Resume the ray query coroutine.
     //
     // If the coroutine suspends because of encountering
@@ -7831,11 +7839,21 @@ struct RayQuery <let rayFlagsGeneric : RAY_FLAG = RAY_FLAG_NONE>
     // functions to appropriately handle the closest hit (it any)
     // that was found.
     //
-    __target_intrinsic(glsl, rayQueryProceedEXT)
     __glsl_extension(GL_EXT_ray_query)
     __glsl_version(460)
     [mutating]
-    bool Proceed();
+    bool Proceed()
+    {
+        __target_switch
+        {
+        case hlsl: __intrinsic_asm ".Proceed";
+        case glsl: __intrinsic_asm "rayQueryProceedEXT";
+        case spirv: return spirv_asm
+            {
+                result:$$bool = OpRayQueryProceedKHR &this
+            };
+        }
+    }
 
     // Causes the ray query to terminate.
     //
@@ -7843,11 +7861,18 @@ struct RayQuery <let rayFlagsGeneric : RAY_FLAG = RAY_FLAG_NONE>
     // traversal has terminated, so that subsequent
     // `Proceed()` calls will return `false`.
     //
-    __target_intrinsic(glsl, rayQueryTerminateEXT)
     __glsl_extension(GL_EXT_ray_query)
     __glsl_version(460)
     [mutating]
-    void Abort();
+    void Abort()
+    {
+        __target_switch
+        {
+        case hlsl: __intrinsic_asm ".Abort";
+        case glsl: __intrinsic_asm "rayQueryTerminateEXT";
+        case spirv: spirv_asm { OpRayQueryTerminateKHR &this };
+        }
+    }
 
     // Get the type of candidate hit being considered.
     //
@@ -7858,11 +7883,22 @@ struct RayQuery <let rayFlagsGeneric : RAY_FLAG = RAY_FLAG_NONE>
     // the kind of candidate hit that must be resolved by
     // user code.
     //
-    __target_intrinsic(glsl, "rayQueryGetIntersectionTypeEXT($0, false)")
     __glsl_extension(GL_EXT_ray_query)
     __glsl_version(460)
     [__readNone]
-    CANDIDATE_TYPE CandidateType();
+    CANDIDATE_TYPE CandidateType()
+    {
+        __target_switch
+        {
+        case hlsl: __intrinsic_asm ".CandidateType";
+        case glsl: __intrinsic_asm "rayQueryGetIntersectionTypeEXT($0, false)";
+        case spirv:
+            uint RayQueryCandidateIntersectionKHR = 0;
+            return spirv_asm {
+                result:$$CANDIDATE_TYPE = OpRayQueryGetIntersectionTypeKHR &this $RayQueryCandidateIntersectionKHR;
+            };
+        }
+    }
 
     // Access properties of a candidate hit.
 
@@ -7980,11 +8016,23 @@ struct RayQuery <let rayFlagsGeneric : RAY_FLAG = RAY_FLAG_NONE>
     void CommitProceduralPrimitiveHit(float t);
 
     // Get the status of the committed (closest) hit, if any.
-    __target_intrinsic(glsl, "rayQueryGetIntersectionTypeEXT($0, true)")
     __glsl_extension(GL_EXT_ray_query)
     __glsl_version(460)
     [__readNone]
-    COMMITTED_STATUS CommittedStatus();
+    COMMITTED_STATUS CommittedStatus()
+    {
+        __target_switch
+        {
+        case hlsl: __intrinsic_asm ".CommittedStatus";
+        case glsl: __intrinsic_asm "rayQueryGetIntersectionTypeEXT($0, true)";
+        case spirv:
+            uint RayQueryCommittedIntersectionKHR = 1;
+            return spirv_asm
+            {
+                result:$$COMMITTED_STATUS = OpRayQueryGetIntersectionTypeKHR &this $RayQueryCommittedIntersectionKHR;
+            };
+        }
+    }
 
     // Access properties of the committed hit.
     //
@@ -8012,11 +8060,23 @@ struct RayQuery <let rayFlagsGeneric : RAY_FLAG = RAY_FLAG_NONE>
     [__readNone]
     float4x3 CommittedWorldToObject4x3();
 
-    __target_intrinsic(glsl, "rayQueryGetIntersectionTEXT($0, true)")
     __glsl_extension(GL_EXT_ray_query)
     __glsl_version(460)
     [__readNone]
-    float CommittedRayT();
+    float CommittedRayT()
+    {
+        __target_switch
+        {
+        case glsl: __intrinsic_asm "rayQueryGetIntersectionTEXT($0, true)";
+        case hlsl: __intrinsic_asm ".CommittedRayT";
+        case spirv:
+            uint RayQueryCommittedIntersectionKHR = 1;
+            return spirv_asm
+            {
+                result:$$float = OpRayQueryGetIntersectionTKHR &this $RayQueryCommittedIntersectionKHR
+            };
+        }
+    }
 
     __target_intrinsic(glsl, "rayQueryGetIntersectionInstanceIdEXT($0, true)")
     __glsl_extension(GL_EXT_ray_query)
diff --git a/source/slang/slang-check-expr.cpp b/source/slang/slang-check-expr.cpp
index 75fd5177e..5ea897448 100644
--- a/source/slang/slang-check-expr.cpp
+++ b/source/slang/slang-check-expr.cpp
@@ -670,6 +670,8 @@ namespace Slang
                         ThisExpr* expr = m_astBuilder->create<ThisExpr>();
                         expr->type.type = thisType;
                         expr->loc = loc;
+                        if (auto declRefExpr = as<DeclRefExpr>(originalExpr))
+                            expr->scope = declRefExpr->scope;
 
                         // Whether or not the implicit `this` is mutable depends
                         // on the context in which it is used, and the lookup
diff --git a/source/slang/slang-emit-spirv-ops.h b/source/slang/slang-emit-spirv-ops.h
index 57c074751..42a2023fc 100644
--- a/source/slang/slang-emit-spirv-ops.h
+++ b/source/slang/slang-emit-spirv-ops.h
@@ -229,6 +229,20 @@ SpvInst* emitOpTypeSampler(IRInst* inst)
     );
 }
 
+SpvInst* emitOpTypeAccelerationStructure(IRInst* inst)
+{
+    return emitInstMemoized(
+        getSection(SpvLogicalSectionID::ConstantsAndTypes), inst, SpvOpTypeAccelerationStructureKHR, kResultID
+    );
+}
+
+SpvInst* emitOpTypeRayQuery(IRInst* inst)
+{
+    return emitInstMemoized(
+        getSection(SpvLogicalSectionID::ConstantsAndTypes), inst, SpvOpTypeRayQueryKHR, kResultID
+    );
+}
+
 // https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#OpTypeArray
 template<typename T1, typename T2>
 SpvInst* emitOpTypeArray(IRInst* inst, const T1& elementType, const T2& length)
diff --git a/source/slang/slang-emit-spirv.cpp b/source/slang/slang-emit-spirv.cpp
index 07f1b2aee..88fb8b0fe 100644
--- a/source/slang/slang-emit-spirv.cpp
+++ b/source/slang/slang-emit-spirv.cpp
@@ -1430,11 +1430,17 @@ struct SPIRVEmitContext
                 );
             }
         case kIROp_SamplerStateType:
-                return emitOpTypeSampler(inst);
-        // > OpTypeArray
-        // > OpTypeRuntimeArray
-        // > OpTypeOpaque
-        // > OpTypePointer
+            return emitOpTypeSampler(inst);
+
+        case kIROp_RaytracingAccelerationStructureType:
+            requireSPIRVCapability(SpvCapabilityRayTracingKHR);
+            ensureExtensionDeclaration(UnownedStringSlice("SPV_KHR_ray_tracing"));
+            return emitOpTypeAccelerationStructure(inst);
+
+        case kIROp_RayQueryType:
+            ensureExtensionDeclaration(UnownedStringSlice("SPV_KHR_ray_query"));
+            requireSPIRVCapability(SpvCapabilityRayQueryKHR);
+            return emitOpTypeRayQuery(inst);
 
         case kIROp_FuncType:
             // > OpTypeFunction
@@ -1544,6 +1550,8 @@ struct SPIRVEmitContext
             }
         case kIROp_GetStringHash:
             return emitGetStringHash(inst);
+        case kIROp_AllocateOpaqueHandle:
+            return nullptr;
         default:
             {
                 if (as<IRSPIRVAsmOperand>(inst))
diff --git a/source/slang/slang-ir-spirv-legalize.cpp b/source/slang/slang-ir-spirv-legalize.cpp
index 63be9d19d..8d785bceb 100644
--- a/source/slang/slang-ir-spirv-legalize.cpp
+++ b/source/slang/slang-ir-spirv-legalize.cpp
@@ -202,6 +202,22 @@ struct SPIRVLegalizationContext : public SourceEmitterBase
                 i->removeAndDeallocate();
     }
 
+    bool isSpirvUniformConstantType(IRType* type)
+    {
+        if (as<IRTextureTypeBase>(type))
+            return true;
+        if (as<IRSamplerStateTypeBase>(type))
+            return true;
+        switch (type->getOp())
+        {
+        case kIROp_RaytracingAccelerationStructureType:
+        case kIROp_RayQueryType:
+            return true;
+        default:
+            return false;
+        }
+    }
+
     void processGlobalParam(IRGlobalParam* inst)
     {
         // If the global param is not a pointer type, make it so and insert explicit load insts.
@@ -237,8 +253,7 @@ struct SPIRVLegalizationContext : public SourceEmitterBase
             // Textures and Samplers can't be in Uniform for Vulkan, if they are
             // placed here then put them in UniformConstant instead
             if (storageClass == SpvStorageClassUniform
-                && (as<IRTextureTypeBase>(inst->getDataType())
-                    || as<IRSamplerStateTypeBase>(inst->getDataType())))
+                && isSpirvUniformConstantType(inst->getDataType()))
             {
                 storageClass = SpvStorageClassUniformConstant;
             }
diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index 106e5b5a3..8d1a86d4d 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -2537,12 +2537,18 @@ ParameterDirection getParameterDirection(VarDeclBase* paramDecl)
     ///
 ParameterDirection getThisParamDirection(Decl* parentDecl, ParameterDirection defaultDirection)
 {
+    auto parentParent = getParentDecl(parentDecl);
     // The `this` parameter for a `class` is always `in`.
-    if (as<ClassDecl>(parentDecl->parentDecl))
+    if (as<ClassDecl>(parentParent))
     {
         return kParameterDirection_In;
     }
 
+    if (parentParent->findModifier<NonCopyableTypeAttribute>())
+    {
+        return kParameterDirection_Ref;
+    }
+
     // Applications can opt in to a mutable `this` parameter,
     // by applying the `[mutating]` attribute to their
     // declaration.
diff --git a/source/slang/slang-parser.cpp b/source/slang/slang-parser.cpp
index 3f5267577..892edd328 100644
--- a/source/slang/slang-parser.cpp
+++ b/source/slang/slang-parser.cpp
@@ -6294,13 +6294,8 @@ namespace Slang
     static std::optional<SPIRVAsmOperand> parseSPIRVAsmOperand(Parser* parser)
     {
         const auto slangIdentOperand = [&](auto flavor){
-            const auto tok = parser->ReadToken(TokenType::Identifier);
-
-            VarExpr* varExpr = parser->astBuilder->create<VarExpr>();
-            varExpr->scope = parser->currentScope;
-            varExpr->loc = tok.getLoc();
-            varExpr->name = tok.getName();
-            return SPIRVAsmOperand{flavor, tok, varExpr};
+            auto token = parser->tokenReader.peekToken();
+            return SPIRVAsmOperand{flavor, token, parseAtomicExpr(parser)};
         };
 
         const auto slangTypeExprOperand = [&](auto flavor) {
diff --git a/tests/expected-failure.txt b/tests/expected-failure.txt
index 0acf97262..724f2588e 100644
--- a/tests/expected-failure.txt
+++ b/tests/expected-failure.txt
@@ -1,8 +1,5 @@
 tests/autodiff/global-param-hoisting.slang.1 (vk)
 tests/bugs/buffer-swizzle-store.slang.1 (vk)
-tests/bugs/gh-3075.slang.2 (vk)
-tests/bugs/ray-query-in-generic.slang.1 (vk)
-tests/compute/ray-tracing-inline.slang.1 (vk)
 tests/language-feature/constants/constexpr-loop.slang.1 (vk)
 tests/optimization/func-resource-result/func-resource-result-complex.slang.1 (vk)
 tests/type/texture-sampler/texture-sampler-2d.slang (vk)