NonUniformResourceIndex not working when resource assigned from ResourceDescriptorHeap is returned from function

[djmoth]

Description
I have tried to make a type-safe resource handle system using templates:

template<typename T>
struct StructuredBufferHandle
{
    uint Index;
};

template<typename T>
struct StructuredBufferOffset
{
    static const uint HandleBits = 12;
    static const uint OffsetBits = sizeof(uint) * 8 - HandleBits;
    static const uint MaxHandleIndex = (1u << HandleBits) - 1;
    static const uint MaxOffset = (1u << OffsetBits) - 1;

    uint PackedOffset;

    StructuredBufferHandle<T> GetHandle()
    {
        return (StructuredBufferHandle<T>)(PackedOffset & MaxHandleIndex);
    }

    uint GetOffset()
    {
        return PackedOffset >> HandleBits;
    }

    bool operator==(StructuredBufferOffset<T> other)
    {
        return PackedOffset == other.PackedOffset;
    }
};

template<typename T>
struct SubStructuredBuffer
{
    StructuredBuffer<T> Buffer;
    uint Offset;
    uint Length;

    T operator[](uint index)
    {
        return Buffer[Offset + index];
    }

    T Load(uint index)
    {
        return Buffer[Offset + index];
    }
};

template<typename T>
StructuredBuffer<T> GetStructuredBuffer(StructuredBufferHandle<T> handle)
{
    return ResourceDescriptorHeap[handle.Index];
}

template<typename T>
StructuredBuffer<T> GetStructuredBufferNonUniform(StructuredBufferHandle<T> handle)
{
    return ResourceDescriptorHeap[NonUniformResourceIndex(handle.Index)];
}

template<typename T>
SubStructuredBuffer<T> GetSubStructuredBufferNonUniform(StructuredBufferOffset<T> offset)
{
    SubStructuredBuffer<T> result;
    result.Buffer = GetStructuredBufferNonUniform(offset.GetHandle());
    result.Offset = offset.GetOffset();
    result.Length = 0;
    return result;
}

The uniform version works fine on both AMD and NVIDIA cards. But the non-uniform version does not work on AMD and results in rendering artifacts. These can be fixed with Radeon Developer Tool Suite as described in https://asawicki.info/news_1788_debugging_amd-specific_issues_with_driver_experiments_tool. This leads be to believe that the issue is with the NonUniformResourceIndex not being respected somewhere down the line. It specifically doesn't work in GetSubStructuredBufferNonUniform. I have tried replacing GetStructuredBufferNonUniform with a macro:

#define GetStructuredBufferNonUniform(handle) ResourceDescriptorHeap[NonUniformResourceIndex((handle).Index)];

I tried did this because https://microsoft.github.io/DirectX-Specs/d3d/HLSL_SM_6_6_DynamicResources.html#resourcedescriptorheap-and-samplerdescriptorheap specifically says:

ResourceDescriptorHeap[index] must be used to assign a local or global CBV, SRV, or UAV resource variable or function call argument.

Which could be interpreted as the not being allowed as a return value of a function. But this made no difference.
When analyzing my code in PIX, which is a large raytracing compute shader, PIX finds multiple places in which NonUniformResourceIndex is missing with the Dr. PIX query NonUniformResourceIndex. Changing all instances of ResourceDescriptorHeap[handle.Index] to ResourceDescriptorHeap[NonUniformResourceIndex(handle.Index)] made no difference in the number of instances. This shows that it is not just due to me not calling the right GetStructuredBufferNonUniform function.

Steps to Reproduce
https://godbolt.org/z/heq67bT8c
I hope this snippet can reproduce the issue, but I have not had the opportunity to test it on AMD hardware. Perhaps the bug can be seen by analyzing the annotations of the handles in the DXIL, but I am not able to verify it myself. The handle macros at the top of the file are used to tightly pack handles in cbuffers as uints, as structs are 16-byte aligned.

Actual Behavior
A StructuredBuffer returned by ResourceDescriptorHeap[NonUniformResourceIndex(handle.Index)] and assigned to a field of a struct returned by a function does not retain its NonUniformResourceIndex annotation throughout its lifetime.

Environment

• DXC version: 1.8.2505.1
• Host Operating System: Windows 11