Fix handling of arrays of resources in type legalization (#896)

Before type legalization we might have code like: (using pseudo-Slang-IR):

    struct P { ... Texture2D<float>[] t; }
    global_param p : ParameterBlock<P>;
    ...

    // p.t[someIndex].Load(...);
    //
    let ptrToArrayOfTextures = getFieldPtr(p, "t")                            : Ptr<Texture2D<float>[]>;
    let ptrToTexture         = getElementPtr(ptrToArrayOfTextures, someIndex) : Ptr<Texture2D<float>>;
    let texture              = load(ptrToTexture)                             : Texture2D<float>;
    let result               = call(loadFunc, texture, ...)                   : float;

Legalization needs to move the `t` array there out of the `p` parameter block, so the global declarations become something like:

    struct P_Ordinary { ... }; // no more "t" field
    global_param p_ordinary : ParameterBlock<p_ordinary);
    global_param p_t : Texture2D<float>[];

In terms of the code to access `p.t[someIndex]` the problem is that `p_t` has one less level of indirection than `p.t` had. We solve this in the type legalization pass using "pseudo-types" and "pseudo-values," where one of the cases is `implicitIndirect` which holds a value of type `T`, but indicates that it should act like a value of type `T*`.

We then use some basic rules for dealing with `implicitIndirect` values, such as:

    load(implicitDeref(x)) : T => x : T
    getFieldPtr(implicitDeref(s), f) => implicitDeref(getField(s, f))
    getElementPtr(implicitDeref(a), i) => implicitDeref(getElement(a, i))

The bug here was that for the `getFieldPtr` and `getElementPtr` cases, we weren't computing the type of the `getField` or `getElement` instruction correctly. We were copying the type from the `getFieldPtr` or `getElementPtr` operation over directly, but those will be *pointer* types and we need the type of whatever they point to.

Once the types are fixed, we can properly generate legalized IR for `p.t[someIndex].Load(...) that looks like:

    let arrayOfTextures = p_t                                    : Texture2D<float>[];
    let texture         = getElement(arrayOfTextures, someIndex) : Texture2D<float>;
    let result           = call(loadFunc, texture, ...)          : float;

The old was giving the `texture` intermediate a type of `Ptr<Texure2D<float>>`. That didn't actually trip up too many things, because we mostly just went on to emit code from something with slightly incorrect types for intermediates that never show up in the generated HLSL/GLSL.

Where this caused a problem is for some of the intrinsic function definitions for the GLSL/Vulkan back-end, because those do things that inspect operand types. In particular the `$z` opcode in our intrinsic function strings triggers logic that looks at a texture operand, and uses its type to try to find the appropriate swizzle to get from a 4-component vector to the appropriate type for the operation (e.g., for a load from a `Texture2D<float>` we need to swizzle with `.x` to get a single scalar out of the matching GLSL texture fetch operation).

The main fix in this change is thus to make `getElementPtr` and `getFieldPtr` legalization properly account for the fact that when switching to `getElement` or `getField` we need a result type that is the "pointee" of the original result.

There was already logic to extract the pointed-to type from a pointer in `ir-specialize.cpp`, so I extracted that to a re-usable function in the IR as `tryGetPointedToType` (returns null if the type isn't actually a pointer).
This logic needed to be extended for type legalization, to deal with the various "pseudo-type" cases.

There is another fix in this change which is marking the `NonUniformResourceIndex` function as `[__readNone]`, which enables it to be more aggressively folded into use sites. Without that fix, we risk emitting code like:

```glsl
int tmp = nonUniformEXT(someIndex);
vec4 result = texelFetch(arrayOfTextures[tmp], ...);
```

The problem with that code is that (at least by my reading of the spec), assigning to the variable `tmp` that isn't declared with the `nonUniformEXT` qualifier effectively loses that qualifier, and drivers are free to assume that `tmp` is uniform when used to index into `arrayOfTextures`.

Marking the `NonUniformResourceIndex` function as `[__readNone]` indicates that it has no side effects, which should mean that our emit logic no longer needs to emit it was its own line of code to be safe.
The effects of this change are confirmed by both the new test case added, and the existing `non-uniform-indexing` test.

9 files changed, 190 insertions, 72 deletions

diff --git a/source/slang/hlsl.meta.slang b/source/slang/hlsl.meta.slang
index 4044600aa..a12e45c36 100644
--- a/source/slang/hlsl.meta.slang
+++ b/source/slang/hlsl.meta.slang
@@ -1036,10 +1036,12 @@ __generic<let N : int> float noise(vector<float, N> x);
 /// is more semantically "correct."
 __target_intrinsic(glsl, nonuniformEXT)
 __glsl_extension(GL_EXT_nonuniform_qualifier)
+[__readNone]
 uint NonUniformResourceIndex(uint index);
 
 __target_intrinsic(glsl, nonuniformEXT)
 __glsl_extension(GL_EXT_nonuniform_qualifier)
+[__readNone]
 int NonUniformResourceIndex(int index);
 
 // Normalize a vector
diff --git a/source/slang/hlsl.meta.slang.h b/source/slang/hlsl.meta.slang.h
index 1342badb8..a82a733ac 100644
--- a/source/slang/hlsl.meta.slang.h
+++ b/source/slang/hlsl.meta.slang.h
@@ -1112,10 +1112,12 @@ SLANG_RAW("/// the user's responsibility, so that the default behavior of the la
 SLANG_RAW("/// is more semantically \"correct.\"\n")
 SLANG_RAW("__target_intrinsic(glsl, nonuniformEXT)\n")
 SLANG_RAW("__glsl_extension(GL_EXT_nonuniform_qualifier)\n")
+SLANG_RAW("[__readNone]\n")
 SLANG_RAW("uint NonUniformResourceIndex(uint index);\n")
 SLANG_RAW("\n")
 SLANG_RAW("__target_intrinsic(glsl, nonuniformEXT)\n")
 SLANG_RAW("__glsl_extension(GL_EXT_nonuniform_qualifier)\n")
+SLANG_RAW("[__readNone]\n")
 SLANG_RAW("int NonUniformResourceIndex(int index);\n")
 SLANG_RAW("\n")
 SLANG_RAW("// Normalize a vector\n")
@@ -1529,7 +1531,7 @@ for (int aa = 0; aa < kBaseBufferAccessLevelCount; ++aa)
 
     sb << "};\n";
 }
-SLANG_RAW("#line 1456 \"hlsl.meta.slang\"")
+SLANG_RAW("#line 1458 \"hlsl.meta.slang\"")
 SLANG_RAW("\n")
 SLANG_RAW("\n")
 SLANG_RAW("\n")
diff --git a/source/slang/ir-legalize-types.cpp b/source/slang/ir-legalize-types.cpp
index 0c1465fb4..00b8109ee 100644
--- a/source/slang/ir-legalize-types.cpp
+++ b/source/slang/ir-legalize-types.cpp
@@ -560,6 +560,60 @@ static LegalVal wrapBufferValue(
     }
 }
 
+static IRType* getPointedToType(
+    IRTypeLegalizationContext*  context,
+    IRType*                     ptrType)
+{
+    auto valueType = tryGetPointedToType(context->builder, ptrType);
+    if( !valueType )
+    {
+        SLANG_UNEXPECTED("expected a pointer type during type legalization");
+    }
+    return valueType;
+}
+
+static LegalType getPointedToType(
+    IRTypeLegalizationContext*  context,
+    LegalType                   type)
+{
+    switch( type.flavor )
+    {
+    case LegalType::Flavor::none:
+        return LegalType();
+
+    case LegalType::Flavor::simple:
+        return LegalType::simple(getPointedToType(context, type.getSimple()));
+
+    case LegalType::Flavor::implicitDeref:
+        return type.getImplicitDeref()->valueType;
+
+    case LegalType::Flavor::pair:
+        {
+            auto pairType = type.getPair();
+            auto ordinary = getPointedToType(context, pairType->ordinaryType);
+            auto special = getPointedToType(context, pairType->specialType);
+            return LegalType::pair(ordinary, special, pairType->pairInfo);
+        }
+
+    case LegalType::Flavor::tuple:
+        {
+            auto tupleType = type.getTuple();
+            RefPtr<TuplePseudoType> resultTuple = new TuplePseudoType();
+            for( auto ee : tupleType->elements )
+            {
+                TuplePseudoType::Element resultElement;
+                resultElement.key = ee.key;
+                resultElement.type = getPointedToType(context, ee.type);
+            }
+            return LegalType::tuple(resultTuple);
+        }
+
+    default:
+        SLANG_UNEXPECTED("unhandled case in type legalization");
+        UNREACHABLE_RETURN(LegalType());
+    }
+}
+
 static LegalVal legalizeFieldAddress(
     IRTypeLegalizationContext*    context,
     LegalType                   type,
@@ -678,8 +732,8 @@ static LegalVal legalizeFieldAddress(
             // adding an implicit dereference on top of that.
             //
             auto implicitDerefVal = legalPtrOperand.getImplicitDeref();
-
-            return LegalVal::implicitDeref(legalizeFieldExtract(context,type, implicitDerefVal, fieldKey));
+            auto valueType = getPointedToType(context, type);
+            return LegalVal::implicitDeref(legalizeFieldExtract(context, valueType, implicitDerefVal, fieldKey));
         }
 
     default:
@@ -816,7 +870,6 @@ static LegalVal legalizeGetElement(
         indexOperand);
 }
 
-
 static LegalVal legalizeGetElementPtr(
     IRTypeLegalizationContext*  context,
     LegalType                   type,
@@ -914,10 +967,15 @@ static LegalVal legalizeGetElementPtr(
             // instead (and then wrap the element value with an
             // implicit dereference).
             //
+            // The result type for our `getElement` instruction needs
+            // to be the type *pointed to* by `type`, and not `type.
+            //
+            auto valueType = getPointedToType(context, type);
+
             auto implicitDerefVal = legalPtrOperand.getImplicitDeref();
             return LegalVal::implicitDeref(legalizeGetElement(
                 context,
-                type,
+                valueType,
                 implicitDerefVal,
                 indexOperand));
         }
diff --git a/source/slang/ir-specialize.cpp b/source/slang/ir-specialize.cpp
index 7a129a03b..e491eaa31 100644
--- a/source/slang/ir-specialize.cpp
+++ b/source/slang/ir-specialize.cpp
@@ -1153,54 +1153,6 @@ struct SpecializationContext
         }
     }
 
-        /// Given a type being used as pointer, try to determine the type it points to.
-    IRType* tryGetPointedToType(
-        IRBuilder*  builder,
-        IRType*     type)
-    {
-        // The "true" pointers and the pointer-like stdlib types are the easy cases.
-        if( auto ptrType = as<IRPtrTypeBase>(type) )
-        {
-            return ptrType->getValueType();
-        }
-        else if( auto ptrLikeType = as<IRPointerLikeType>(type) )
-        {
-            return ptrLikeType->getElementType();
-        }
-        //
-        // A more interesting case arises when we have a `BindExistentials<P<T>, ...>`
-        // where `P<T>` is a pointer(-like) type.
-        //
-        else if( auto bindExistentials = as<IRBindExistentialsType>(type) )
-        {
-            // We know that `BindExistentials` won't introduce its own
-            // existential type parameters, nor will any of the pointer(-like)
-            // type constructors `P`.
-            //
-            // Thus we know that the type that is pointed to should be
-            // the same as `BindExistentials<T, ...>`.
-            //
-            auto baseType = bindExistentials->getBaseType();
-            if( auto baseElementType = tryGetPointedToType(builder, baseType) )
-            {
-                UInt existentialArgCount = bindExistentials->getExistentialArgCount();
-                List<IRInst*> existentialArgs;
-                for( UInt ii = 0; ii < existentialArgCount; ++ii )
-                {
-                    existentialArgs.Add(bindExistentials->getExistentialArg(ii));
-                }
-                return builder->getBindExistentialsType(
-                    baseElementType,
-                    existentialArgCount,
-                    existentialArgs.Buffer());
-            }
-        }
-
-        // TODO: We may need to handle other cases here.
-
-        return nullptr;
-    }
-
     void maybeSpecializeLoad(IRLoad* inst)
     {
         auto ptrArg = inst->ptr.get();
diff --git a/source/slang/ir.cpp b/source/slang/ir.cpp
index 677429e32..3f6f8e0a3 100644
--- a/source/slang/ir.cpp
+++ b/source/slang/ir.cpp
@@ -236,6 +236,60 @@ namespace Slang
         return nullptr;
     }
 
+    // IRPtrTypeBase
+
+    IRType* tryGetPointedToType(
+        IRBuilder*  builder,
+        IRType*     type)
+    {
+        if( auto rateQualType = as<IRRateQualifiedType>(type) )
+        {
+            type = rateQualType->getDataType();
+        }
+
+        // The "true" pointers and the pointer-like stdlib types are the easy cases.
+        if( auto ptrType = as<IRPtrTypeBase>(type) )
+        {
+            return ptrType->getValueType();
+        }
+        else if( auto ptrLikeType = as<IRPointerLikeType>(type) )
+        {
+            return ptrLikeType->getElementType();
+        }
+        //
+        // A more interesting case arises when we have a `BindExistentials<P<T>, ...>`
+        // where `P<T>` is a pointer(-like) type.
+        //
+        else if( auto bindExistentials = as<IRBindExistentialsType>(type) )
+        {
+            // We know that `BindExistentials` won't introduce its own
+            // existential type parameters, nor will any of the pointer(-like)
+            // type constructors `P`.
+            //
+            // Thus we know that the type that is pointed to should be
+            // the same as `BindExistentials<T, ...>`.
+            //
+            auto baseType = bindExistentials->getBaseType();
+            if( auto baseElementType = tryGetPointedToType(builder, baseType) )
+            {
+                UInt existentialArgCount = bindExistentials->getExistentialArgCount();
+                List<IRInst*> existentialArgs;
+                for( UInt ii = 0; ii < existentialArgCount; ++ii )
+                {
+                    existentialArgs.Add(bindExistentials->getExistentialArg(ii));
+                }
+                return builder->getBindExistentialsType(
+                    baseElementType,
+                    existentialArgCount,
+                    existentialArgs.Buffer());
+            }
+        }
+
+        // TODO: We may need to handle other cases here.
+
+        return nullptr;
+    }
+
 
     // IRBlock
 
@@ -2395,21 +2449,8 @@ namespace Slang
         // results) at the "default" rate of the parent function,
         // unless a subsequent analysis pass constraints it.
 
-        IRType* valueType = nullptr;
-        if(auto ptrType = as<IRPtrTypeBase>(ptr->getDataType()))
-        {
-            valueType = ptrType->getValueType();
-        }
-        else if(auto ptrLikeType = as<IRPointerLikeType>(ptr->getDataType()))
-        {
-            valueType = ptrLikeType->getElementType();
-        }
-        else
-        {
-            // Bad!
-            SLANG_ASSERT(ptrType);
-            return nullptr;
-        }
+        IRType* valueType = tryGetPointedToType(this, ptr->getFullType());
+        SLANG_ASSERT(valueType);
 
         // Ugly special case: if the front-end created a variable with
         // type `Ptr<@R T>` instead of `@R Ptr<T>`, then the above
diff --git a/source/slang/ir.h b/source/slang/ir.h
index 09013fe48..d4497ec91 100644
--- a/source/slang/ir.h
+++ b/source/slang/ir.h
@@ -25,6 +25,7 @@ class   Layout;
 class   Type;
 class   Session;
 class   Name;
+struct  IRBuilder;
 struct  IRFunc;
 struct  IRGlobalValueWithCode;
 struct  IRInst;
@@ -934,6 +935,13 @@ SIMPLE_IR_TYPE(OutType, OutTypeBase)
 SIMPLE_IR_TYPE(InOutType, OutTypeBase)
 SIMPLE_IR_TYPE(ExistentialBoxType, PtrTypeBase)
 
+    /// Get the type pointed to be `ptrType`, or `nullptr` if it is not a pointer(-like) type.
+    ///
+    /// The given IR `builder` will be used if new instructions need to be created.
+IRType* tryGetPointedToType(
+        IRBuilder*  builder,
+        IRType*     type);
+
 struct IRFuncType : IRType
 {
     IRType* getResultType() { return (IRType*) getOperand(0); }
diff --git a/tests/cross-compile/non-uniform-indexing.slang.glsl b/tests/cross-compile/non-uniform-indexing.slang.glsl
index 5ea5ed73a..96abe6bac 100644
--- a/tests/cross-compile/non-uniform-indexing.slang.glsl
+++ b/tests/cross-compile/non-uniform-indexing.slang.glsl
@@ -17,10 +17,12 @@ in vec3 _S2;
 
 void main()
 {
-    int _S3 = nonuniformEXT(int(_S2.z));
+    vec4 _S3 = texture(
+    	sampler2D(
+    		t_0[nonuniformEXT(int(_S2.z))],
+    		s_0),
+		_S2.xy);
 
-    vec4 _S4 = texture(sampler2D(t_0[_S3],s_0), _S2.xy);
-
-    _S1 = _S4;
+    _S1 = _S3;
     return;
 }
diff --git a/tests/cross-compile/vk-texture-indexing.slang b/tests/cross-compile/vk-texture-indexing.slang
new file mode 100644
index 000000000..9a086d5bd
--- /dev/null
+++ b/tests/cross-compile/vk-texture-indexing.slang
@@ -0,0 +1,24 @@
+// vk-texture-indexing.slang
+
+//TEST:CROSS_COMPILE:-target spirv-assembly -entry main -stage fragment
+
+struct Params
+{
+    Texture2D<float> textures[10];
+    SamplerState sampler;
+};
+
+ParameterBlock<Params> gParams;
+
+float fetchData(uint2 coords, uint index)
+{
+    return gParams.textures[NonUniformResourceIndex(index)][coords];
+}
+
+float4 main(
+	uint3 uv : UV)
+    : SV_Target
+{
+    float v = fetchData(uv.xy, uv.z);
+    return v;
+}
diff --git a/tests/cross-compile/vk-texture-indexing.slang.glsl b/tests/cross-compile/vk-texture-indexing.slang.glsl
new file mode 100644
index 000000000..069e6efc3
--- /dev/null
+++ b/tests/cross-compile/vk-texture-indexing.slang.glsl
@@ -0,0 +1,29 @@
+#version 450
+#extension GL_EXT_nonuniform_qualifier : require
+#extension GL_EXT_samplerless_texture_functions : require
+
+layout(binding = 0)
+uniform texture2D  gParams_textures_0[10];
+
+float fetchData_0(uvec2 coords_0, uint index_0)
+{
+    float _S1 = texelFetch(
+    	gParams_textures_0[nonuniformEXT(index_0)],
+    	ivec2(coords_0),
+    	0).x;
+
+    return _S1;
+}
+
+layout(location = 0)
+out vec4 _S2;
+
+flat layout(location = 0)
+in uvec3 _S3;
+
+void main()
+{
+    float v_0 = fetchData_0(_S3.xy, _S3.z);
+    _S2 = vec4(v_0);
+    return;
+}