Support for dynamic array of textures access for Texture footprint. (#6392)

* Support for dynamic array of textures access for Texture footprint.

* Do a check for DXIL for dynamic array/footprint.

* format code (#40)

---------

Co-authored-by: jsmall-nvidia <jsmall@nvidia.com>
Co-authored-by: slangbot <ellieh+slangbot@nvidia.com>
Co-authored-by: Yong He <yonghe@outlook.com>

2 files changed, 188 insertions, 1 deletions

diff --git a/source/slang/slang-ir-lower-binding-query.cpp b/source/slang/slang-ir-lower-binding-query.cpp
index 0935f1e54..541640aaf 100644
--- a/source/slang/slang-ir-lower-binding-query.cpp
+++ b/source/slang/slang-ir-lower-binding-query.cpp
@@ -288,7 +288,118 @@ struct BindingQueryLoweringContext : public WorkListPass
     //
     OpaqueValueInfo computeOpaqueValueInfo(IRInst* opaqueValue)
     {
-        if (auto globalParam = as<IRGlobalParam>(opaqueValue))
+        if (auto getElement = as<IRGetElement>(opaqueValue))
+        {
+            IRInst* baseInst = getElement->getBase();
+            IRInst* indexInst = getElement->getIndex();
+
+            IRInst* elementType = getElement->getDataType();
+
+            // TODO(JS): This a hack to make this work for arrays of resource type.
+            // It won't work in the general case as it stands because we would need
+            // to propogate layout kind types needed at usage sites.
+            // Without knowing the resource kind that is being processed it's not possible
+            // to accumulate the calculation.
+            //
+            // So presumably we need to request a binding query for a specific resource kind.
+            // We could do this by making the type of the binding query hold the type.
+
+            // We need to add instructions which will work out the binding for the base
+            OpaqueValueInfo baseInfo = findOrComputeOpaqueValueInfo(baseInst);
+
+            // If we couldn't find it we are done
+            if (baseInfo.registerIndex == nullptr || baseInfo.registerSpace == nullptr)
+            {
+                return baseInfo;
+            }
+
+
+            LayoutResourceKind kind = LayoutResourceKind::None;
+            Index stride = 1;
+
+            if (auto resourceType = as<IRResourceType>(elementType))
+            {
+                const auto shape = resourceType->getShape();
+
+                switch (shape)
+                {
+                case SLANG_TEXTURE_1D:
+                case SLANG_TEXTURE_2D:
+                case SLANG_TEXTURE_3D:
+                case SLANG_TEXTURE_CUBE:
+                case SLANG_STRUCTURED_BUFFER:
+                case SLANG_BYTE_ADDRESS_BUFFER:
+                case SLANG_TEXTURE_BUFFER:
+                    {
+                        const auto access = resourceType->getAccess();
+                        bool isReadOnly = (access == SLANG_RESOURCE_ACCESS_READ);
+
+                        kind = isReadOnly ? LayoutResourceKind::ShaderResource
+                                          : LayoutResourceKind::UnorderedAccess;
+                        break;
+                    }
+                default:
+                    break;
+                }
+            }
+            else if (as<IRSamplerStateTypeBase>(elementType))
+            {
+                kind = LayoutResourceKind::SamplerState;
+            }
+            else if (as<IRConstantBufferType>(elementType))
+            {
+                kind = LayoutResourceKind::ConstantBuffer;
+            }
+
+            if (kind == LayoutResourceKind::None)
+            {
+                // Can't determine the kind
+                return OpaqueValueInfo();
+            }
+
+            // If the element type has type layout we can try and use that
+            if (auto layoutDecoration = elementType->findDecoration<IRLayoutDecoration>())
+            {
+                // We have to calculate
+                if (auto elementTypeLayout = as<IRTypeLayout>(layoutDecoration->getLayout()))
+                {
+                    IRTypeSizeAttr* sizeAttr = elementTypeLayout->findSizeAttr(kind);
+                    sizeAttr = sizeAttr ? sizeAttr
+                                        : elementTypeLayout->findSizeAttr(
+                                              LayoutResourceKind::DescriptorTableSlot);
+
+                    if (!sizeAttr)
+                    {
+                        // Couldn't work it out
+                        return OpaqueValueInfo();
+                    }
+
+                    // TODO(JS): Perhaps we have to do something else if not finite?
+                    stride = sizeAttr->getFiniteSize();
+                }
+            }
+
+            SLANG_UNUSED(indexInst);
+
+            // Okay we need to create an instruction which is
+            // base + stride * index
+
+            IRBuilder builder(module);
+
+            builder.setInsertBefore(opaqueValue);
+
+            auto calcRegisterInst = builder.emitAdd(
+                indexType,
+                builder.emitMul(indexType, builder.getIntValue(indexType, stride), indexInst),
+                baseInfo.registerIndex);
+
+            OpaqueValueInfo finalInfo;
+            finalInfo.registerIndex = calcRegisterInst;
+            finalInfo.registerSpace = baseInfo.registerSpace;
+
+            return finalInfo;
+        }
+        else if (auto globalParam = as<IRGlobalParam>(opaqueValue))
         {
             // The simple/base case is when we have a global shader
             // parameter that has layout information attached.
diff --git a/tests/gpu-feature/texture/query/footprint/nv-shader-texture-footprint-array.slang b/tests/gpu-feature/texture/query/footprint/nv-shader-texture-footprint-array.slang
new file mode 100644
index 000000000..dc60cf263
--- /dev/null
+++ b/tests/gpu-feature/texture/query/footprint/nv-shader-texture-footprint-array.slang
@@ -0,0 +1,76 @@
+//DISABLE_TEST:SIMPLE:-target dxil-assembly -entry fragmentMain -profile sm_6_5 -stage fragment -DNV_SHADER_EXTN_SLOT=u0 
+
+//TEST:SIMPLE(filecheck=SPIRV):-target spirv-assembly -entry fragmentMain -stage fragment
+//TEST:SIMPLE(filecheck=DXIL):-target dxil-assembly -entry fragmentMain -profile sm_6_5 -stage fragment -DNV_SHADER_EXTN_SLOT=u0 
+//TEST:SIMPLE(filecheck=HLSL):-target hlsl -entry fragmentMain -stage fragment
+
+//DISABLED_TEST:SIMPLE:-target spirv-assembly -entry fragmentMain -stage fragment
+//DISABLED_TEST:SIMPLE:-target dxil-assembly -entry fragmentMain -stage fragment
+//DISABLED_TEST:SIMPLE:-target hlsl -entry fragmentMain -stage fragment
+
+uniform Texture2D textures[] : register(t2, space10);
+uniform SamplerState sampler;
+uniform RWStructuredBuffer<uint> outputBuffer;
+
+static Texture2D _getBindlessTexture2d(uint texIdx)
+{
+    return textures[NonUniformResourceIndex(texIdx)];
+}
+
+void accumulate(inout uint r, uint u)
+{
+    r = r ^ u;
+}
+
+void accumulate(inout uint r, bool b)
+{
+    accumulate(r, uint(b));
+}
+
+void accumulate(inout uint r, uint2 u)
+{
+    accumulate(r, u.x);
+    accumulate(r, u.y);
+}
+
+void accumulate(inout uint r, uint3 u)
+{
+    accumulate(r, u.x);
+    accumulate(r, u.y);
+    accumulate(r, u.z);
+}
+
+void accumulate(inout uint r, TextureFootprint2D f)
+{
+    accumulate(r, f.anchor);
+    accumulate(r, f.offset);
+    accumulate(r, f.mask);
+    accumulate(r, f.lod);
+    accumulate(r, f.granularity);
+    accumulate(r, f.isSingleLevel);
+}
+
+cbuffer Uniforms
+{
+    uniform float2 coords;
+    uniform uint granularity;
+};
+
+void fragmentMain(
+    float v : VARYING)
+{
+    uint index = uint(v);
+    uint r = 0;
+
+    accumulate(r, _getBindlessTexture2d(index).queryFootprintCoarse(granularity, sampler, coords));
+
+// SPIRV: Extension  "SPV_NV_shader_image_footprint"
+// SPIRV: ImageSampleFootprintNV
+
+// DXIL: struct struct.NvShaderExtnStruct
+
+// HLSL: NvFootprintCoarse
+
+    outputBuffer[index] = r;
+}
+