Several spirv intrinsics and small fix (#3307)

* spirv intrinsic for faceforward

* spirv intrinsic for fwidth

* spirv intrinsic for modf

* spirv intrinsic for nonuniformresourceindex

* spirv intrinsic for transpose

* Make sure address space matches for OpAccessChain

* Correct placement for OpDecorate instructions in spirv asm blocks

2 files changed, 83 insertions, 14 deletions

diff --git a/source/slang/hlsl.meta.slang b/source/slang/hlsl.meta.slang
index fdc2ce83d..43e40778c 100644
--- a/source/slang/hlsl.meta.slang
+++ b/source/slang/hlsl.meta.slang
@@ -2697,6 +2697,7 @@ vector<float16_t, N> f32tof16_(vector<float, N> value)
 __generic<T : __BuiltinFloatingPointType, let N : int>
 __target_intrinsic(hlsl)
 __target_intrinsic(glsl)
+__target_intrinsic(spirv, "OpExtInst resultType resultId glsl450 FaceForward _0 _1 _2")
 [__readNone]
 vector<T,N> faceforward(vector<T,N> n, vector<T,N> i, vector<T,N> ng)
 {
@@ -2918,11 +2919,15 @@ matrix<T, N, M> frexp(matrix<T, N, M> x, out matrix<int, N, M, L> exp)
 // Texture filter width
 __generic<T : __BuiltinFloatingPointType>
 [__readNone]
+__target_intrinsic(hlsl)
+__target_intrinsic(glsl)
+__target_intrinsic(spirv, "OpFWidth resultType resultId _0")
 T fwidth(T x);
 
 __generic<T : __BuiltinFloatingPointType, let N : int>
 __target_intrinsic(hlsl)
 __target_intrinsic(glsl)
+__target_intrinsic(spirv, "OpFWidth resultType resultId _0")
 [__readNone]
 vector<T, N> fwidth(vector<T, N> x)
 {
@@ -4058,6 +4063,9 @@ matrix<T,N,M> min(matrix<T,N,M> x, matrix<T,N,M> y)
 
 // split into integer and fractional parts (both with same sign)
 __generic<T : __BuiltinFloatingPointType>
+__target_intrinsic(hlsl)
+__target_intrinsic(glsl)
+__target_intrinsic(spirv, "OpExtInst resultType resultId glsl450 Modf _0 _1")
 [__readNone]
 T modf(T x, out T ip);
 
@@ -4352,22 +4360,48 @@ __generic<let N : int> float noise(vector<float, N> x)
 /// to this function as necessary in output code, rather than make this
 /// the user's responsibility, so that the default behavior of the language
 /// is more semantically "correct."
-__target_intrinsic(hlsl)
-__target_intrinsic(glsl, nonuniformEXT)
 __glsl_extension(GL_EXT_nonuniform_qualifier)
 [__readNone]
 uint NonUniformResourceIndex(uint index)
 {
-    return index;
+    __target_switch
+    {
+    case hlsl:
+        __intrinsic_asm "NonUniformResourceIndex";
+    case glsl:
+        __intrinsic_asm "nonuniformEXT";
+    case spirv:
+        spirv_asm
+        {
+            OpCapability ShaderNonUniform;
+            OpDecorate $index NonUniform;
+        };
+        return index;
+    default:
+        return index;
+    }
 }
 
-__target_intrinsic(hlsl)
-__target_intrinsic(glsl, nonuniformEXT)
 __glsl_extension(GL_EXT_nonuniform_qualifier)
 [__readNone]
 int NonUniformResourceIndex(int index)
 {
-    return index;
+    __target_switch
+    {
+    case hlsl:
+        __intrinsic_asm "NonUniformResourceIndex";
+    case glsl:
+        __intrinsic_asm "nonuniformEXT";
+    case spirv:
+        spirv_asm
+        {
+            OpCapability ShaderNonUniform;
+            OpDecorate $index NonUniform;
+        };
+        return index;
+    default:
+        return index;
+    }
 }
 
 /// HLSL allows NonUniformResourceIndex around non int/uint types.
@@ -5049,6 +5083,7 @@ matrix<T,N,M> tanh(matrix<T,N,M> x)
 __generic<T : __BuiltinFloatingPointType, let N : int, let M : int>
 __target_intrinsic(hlsl)
 __target_intrinsic(glsl)
+__target_intrinsic(spirv, "OpTranspose resultType resultId _0")
 [__readNone]
 [PreferRecompute]
 matrix<T, M, N> transpose(matrix<T, N, M> x)
@@ -5062,6 +5097,7 @@ matrix<T, M, N> transpose(matrix<T, N, M> x)
 __generic<T : __BuiltinIntegerType, let N : int, let M : int>
 __target_intrinsic(hlsl)
 __target_intrinsic(glsl)
+__target_intrinsic(spirv, "OpTranspose resultType resultId _0")
 [__readNone]
 [PreferRecompute]
 matrix<T, M, N> transpose(matrix<T, N, M> x)
@@ -5075,6 +5111,7 @@ matrix<T, M, N> transpose(matrix<T, N, M> x)
 __generic<T : __BuiltinLogicalType, let N : int, let M : int>
 __target_intrinsic(hlsl)
 __target_intrinsic(glsl)
+__target_intrinsic(spirv, "OpTranspose resultType resultId _0")
 [__readNone]
 [PreferRecompute]
 matrix<T, M, N> transpose(matrix<T, N, M> x)
diff --git a/source/slang/slang-emit-spirv.cpp b/source/slang/slang-emit-spirv.cpp
index 0eae87a10..f26340969 100644
--- a/source/slang/slang-emit-spirv.cpp
+++ b/source/slang/slang-emit-spirv.cpp
@@ -3732,19 +3732,25 @@ struct SPIRVEmitContext
 
     SpvInst* emitGetElementPtr(SpvInstParent* parent, IRGetElementPtr* inst)
     {
+        IRBuilder builder(m_irModule);
         auto base = inst->getBase();
-        SpvWord baseId = 0;
-        // Only used in debug build, but we don't want a warning/error for an unused initialized variable
+        const SpvWord baseId = getID(ensureInst(base));
+
+        // We might replace resultType with a different storage class equivalent
+        auto resultType = as<IRPtrTypeBase>(inst->getFullType());
+        SLANG_ASSERT(resultType);
 
-        if (as<IRPointerLikeType>(base->getDataType()) || as<IRPtrTypeBase>(base->getDataType()))
+        if(const auto basePtrType = as<IRPtrTypeBase>(base->getDataType()))
         {
-            baseId = getID(ensureInst(base));
+            // If the base pointer has a specific address space and the
+            // expected result type doesn't, then make sure they match.
+            // It's invalid spir-v if they don't match
+            resultType = getPtrTypeWithAddressSpace(cast<IRPtrTypeBase>(inst->getDataType()), basePtrType->getAddressSpace());
         }
         else
         {
-            SLANG_ASSERT(!"invalid IR: base of getElementPtr must be a pointer.");
+            SLANG_ASSERT(as<IRPointerLikeType>(base->getDataType()) || !"invalid IR: base of getElementPtr must be a pointer.");
         }
-        SLANG_ASSERT(as<IRPtrTypeBase>(inst->getFullType()));
         return emitOpAccessChain(
             parent,
             inst,
@@ -3835,12 +3841,36 @@ struct SPIRVEmitContext
         return emitOpVectorShuffle(parent, inst, inst->getFullType(), inst->getBase(), inst->getSource(), shuffleIndices.getArrayView());
     }
 
+    IRPtrTypeBase* getPtrTypeWithAddressSpace(IRPtrTypeBase* ptrTypeWithNoAddressSpace, IRIntegerValue addressSpace)
+    {
+        // If it's already ok, return as is
+        if(ptrTypeWithNoAddressSpace->getAddressSpace() == addressSpace)
+            return ptrTypeWithNoAddressSpace;
+
+        // It has an address space, but it doesn't match fail, this indicates a
+        // problem with whatever's creating these types
+        SLANG_ASSERT(!ptrTypeWithNoAddressSpace->hasAddressSpace());
+
+        IRBuilder builder(ptrTypeWithNoAddressSpace);
+        return builder.getPtrType(
+            ptrTypeWithNoAddressSpace->getOp(),
+            ptrTypeWithNoAddressSpace->getValueType(),
+            addressSpace
+        );
+    }
+
     SpvInst* emitStructuredBufferGetElementPtr(SpvInstParent* parent, IRInst* inst)
     {
         //"%addr = OpAccessChain resultType*StorageBuffer resultId _0 const(int, 0) _1;"
         IRBuilder builder(inst);
-        auto addr = emitInst(parent, inst, SpvOpAccessChain, inst->getDataType(), kResultID, inst->getOperand(0), emitIntConstant(0, builder.getIntType()), inst->getOperand(1));
-        return addr;
+        return emitOpAccessChain(
+            parent,
+            inst,
+            // Make sure the resulting pointer has the correct storage class
+            getPtrTypeWithAddressSpace(cast<IRPtrTypeBase>(inst->getDataType()), SpvStorageClassStorageBuffer),
+            inst->getOperand(0),
+            makeArray(emitIntConstant(0, builder.getIntType()), ensureInst(inst->getOperand(1)))
+        );
     }
 
     SpvInst* emitStructuredBufferGetDimensions(SpvInstParent* parent, IRInst* inst)
@@ -4442,6 +4472,8 @@ struct SPIRVEmitContext
                                 return getSection(SpvLogicalSectionID::Extensions);
                             case SpvOpExecutionMode:
                                 return getSection(SpvLogicalSectionID::ExecutionModes);
+                            case SpvOpDecorate:
+                                return getSection(SpvLogicalSectionID::Annotations);
                             default:
                                 return defaultParent;