shader-slang / slang

Making it easier to work with shaders
MIT License
1.78k stars 159 forks source link

Add generic descriptor indexing intrinsic #4389

Open dubiousconst282 opened 2 weeks ago

dubiousconst282 commented 2 weeks ago

Add generic descriptor indexing intrinsic

Closes #4351


This seems to be working nicely as far as I can tell, but I have a few questions:

csyonghe commented 2 weeks ago

I think you can just create an ArrayTypeLayout for the type, with the size being unsized and element type being a special "DynamicResourceEntry" type. The TypeLayout for DynamicResourceEntry type should be a SimpleTypeLayout that has one DescriptorTableSlot entry.

csyonghe commented 2 weeks ago

I am not sure why this change need to interact with specialization, they should be orthogonal?

jkwak-work commented 2 weeks ago

I am not sure why this change need to interact with specialization, they should be orthogonal?

I think the specialization side code is trying to add global uniform variable for each types used on "get" method. If I borrow the example from the proposed design, it will looks like this,

static const int Resource = 0;
static const int Sampler = 1;
interface IBuiltinResourceType<let kind: int>
{}

extension Texture2D : IBuiltinResourceType<Resource>
{}

extension Texture1D : IBuiltinResourceType<Resource>
{}

extension Sampler2D : IBuiltinResourceType<Sampler>
{}

struct ResourceArray<let kind:int>
{
    T get<T:IBuiltinResourceType<kind>>(int index);
}

// Specialization code adds two following global variables,
// because "get" method was used with two different types:
// uniform Texture1D resArray_0[];
// uniform Texture2D resArray_1[];
ResourceArray<Resource> resArray;

void main()
{
    Texture1D t1D = resArray.get<Texture1D>(1);
    Texture2D t2D = resArray.get<Texture2D>(2);
}

I think this should happen in the glsl-legalization rather than specialization.

jkwak-work commented 2 weeks ago

Here is my answers to the questions in the description

Type layout is created based on a temporary StructuredBuffer[] type. Reflection will still report the type as the ResourceArray struct, but I'm not sure if this would cause any issues or if it should be done in some other way.

I think the reflection side needs to be handled properly but I don't see any code change related to that. We will need to start from writing a test case for it and clarify what kind of changes will be required.

Specialization will report an error if the ResourceArray variable is used by anything other than the get intrinsic. Should this be done elsewhere or is fine to keep it as is?

As I mentioned on my previous comment, I think that should be handled in the legalization part not specialization part.

Not sure about the naming, and whether it is better to have an enum instead of global constants like discussed in the original issue.

I think enum is better.

csyonghe commented 2 weeks ago

Yes, we should just create those aliased global variables during glsl legalization pass. Basically, convert getResourceElementInst into getElement(aliasedGlobalArrayForResourceType).

dubiousconst282 commented 1 week ago

Thank you both for the review.

I have moved the aliased parameter generation code to legalization. At first I had implemented it in the SPIRV legalization pass, but then moved it to specialization so it could be shared with the GLSL target, without realizing it would do the wrong thing for the other targets.

I have not fully addressed the comments involving type layout and reflection yet, because I'm thinking that having just a "DynamicResource" type would be more general and maybe require less changes on that side (considering that it would be needed for proxying through the ArrayType anyway). Then the revised proposal would be something like:

struct __DynamicResource<let kind : __DynamicResourceKind = __DynamicResourceKind.General> 
{
    T as<T : __IDynamicResourceCastable<kind>>();
}
enum __DynamicResourceKind
{
    General = 0, // CBV_SRV_UAV
    Sampler = 1
}
interface __IDynamicResourceCastable<let kind : __DynamicResourceKind = __DynamicResourceKind.General> { }

__generic<T, Shape : __ITextureShape, let isArray : int, let isMS : int, let sampleCount : int, let access : int, let isShadow : int, let isCombined : int, let format : int>
extension __TextureImpl<T, Shape, isArray, isMS, sampleCount, access, isShadow, isCombined, format> : __IDynamicResourceCastable<__DynamicResourceKind.General> {
    __implicit_conversion(1)
    __init(__DynamicResource res);
}

I believe that with implicit casts, this would have full parity with the HLSL syntax in terms of usage:

[vk::binding(0, 0)]
__DynamicResource g_TextureHeap[];

[numthreads(1)]
void computeMain(int2 tid: SV_DispatchThreadID) {
    RWTexture2D tex = g_TextureHeap[0];

    tex[tid] = 1.0;
}

Please let if I should proceed with this or keep the current interface.

csyonghe commented 1 week ago

Sure, the new approach looks good to me too. In this case is the g_textureHeap just a normal unsized array of DynamicResource?

dubiousconst282 commented 1 week ago

Yes, it is user defined like before. (will edit previous comment to clarify.)

dubiousconst282 commented 1 week ago

Okay I think this is ready for another review.

The only thing that I can tell that is missing is support for casting to buffer types like StructuredBuffer<T> and others that don't derive from ITextureImpl, although with pointers I personally feel like this isn't as critical and can be added later, maybe with the template generator thingy.

I forgot to split my commits, so here's a few notes:

csyonghe commented 6 days ago

I want to bring this up explicitly here that the current implementation cannot handle cases where the global dynamic resource parameter is not referenced directly in use sites and instead are passed around with local variables or function parameters. For example, the following case will not work:

DynamicResource gRes;

void f(DynamicResource t)
{
    Texture2D tt = t;
    tt.Sample(...);
}

void main()
{
    f(gRes);
}

Which is fine as a limitation for now especially when this type will be used only internally in stdlib when we implement the full Bindless support. But I want to make sure the limitation is well understood.

dubiousconst282 commented 6 days ago

I think this case is handled correctly by specialization (it is covered by the test), e.g.:

__DynamicResource g_SingleRes;
RWStructuredBuffer<float4> outputBuffer;

[noinline]
float4 IndirectAccess(__DynamicResource resource, int2 pos)
{
    Texture2D texture = resource;
    return texture[pos];
}

[numthreads(1)]
void computeMain(int2 tid: SV_DispatchThreadID) {
    outputBuffer[tid.x + tid.y * 128] = IndirectAccess(g_SingleRes, tid);
}

On the GLSL targets compiles to:

#version 460
#extension GL_EXT_scalar_block_layout : require
#extension GL_EXT_samplerless_texture_functions : require
layout(row_major) uniform;
layout(row_major) buffer;

#line 73 0
layout(scalar, binding = 1) buffer StructuredBuffer_vectorx3Cfloatx2C4x3E_t_0 {
    vec4 _data[];
} outputBuffer_0;

#line 2364 1
layout(binding = 0)
uniform texture2D g_SingleRes_0;

#line 84 0
vec4 IndirectAccess_0(ivec2 _S1)
{

#line 79
    return (texelFetch((g_SingleRes_0), ((ivec2(uvec2(_S1)))), 0));
}

layout(local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
void main()
{

#line 83
    ivec2 _S2 = ivec2(gl_GlobalInvocationID.xy);
    outputBuffer_0._data[_S2.x + _S2.y * 128] = IndirectAccess_0(_S2);
    return;
}

Although I agree this would not be a significant limitation since this type is intended mainly for internal/library use.

csyonghe commented 6 days ago

Ah, I see that you added a case for dynamic resource to be treated as specializable parameter type. That will work then.

csyonghe commented 6 days ago

But generally, DynamicResource isn't a first class type that users can expect to work everywhere and we may not want to expose it directly. If a user defines a DynamicResource typed field in a struct, it is not going to work well.

My plan is that we should instead provide a __getResourceFromHandle(uint64 index) function, so that we automatically define and index into that global array. This means that the user will never need to deal with dynamic resource type themselves.

With this, we can also provide a Bindless type where T is a resource type and Bindless is a simple struct wrapping a uint64 that provides a get method to return T.

dubiousconst282 commented 6 days ago

Can you clarify on how __getResourceFromHandle() would work? The thing with DynamicResource is that it allows users to specify the binding points directly, which I think maps more closely to how Vulkan's descriptor indexing works.

My intent for this type was that it would allow users to define their own Bindless type, since an earlier proposal had uncertainties on the representation and was closed (https://github.com/shader-slang/slang/issues/768).

csyonghe commented 6 days ago

The user will need to specify the binding for the implicit heap via a compiler option, e.g. -fvk-dynamic-descriptor-set 10.

The reason to make the global array system defined instead of user defined is to avoid conflicts and allow reusing code written by different developers.

If library A and library B are written by two developers and they all want to use bindless model, how can one consume both A and B without worrying about conflicts?

The only way to facilitate seamless sharing is for the compiler to own this definition so there will always be one declaration of such array and there could never be conflicts.

csyonghe commented 6 days ago

Internally, we will introduce a kIROp_GetGlobalResourceHeap inst that is hoistable so it will be globally dedupplicated after linking.

Then we make __getGlobalResourceHeap() intrinsic map to this op code.

Then we can have __getResourceFromHandle implemented by calling this intrinsic.

The binding for the global heap can be provide explicitly with a compiler option, or implicitly by the system where we just allocate the next available space after all other parameters have been assigned a binding. We will need the reflection API to be able to tell the user which space is allocated for the implicit heap.

When emitting code for spirv, we will have a pass to rewrite kIROp_GetGlobalResourceHeap into a global DynamicResource array, before the specialization pass, so all the code provided in this PR will run and turn them into aliases of concrete resource arrays.

csyonghe commented 6 days ago

With this, the bindless type can simply be implemented in stdlib as an ordinary struct type:

struct Bindless<T: __BuiltinResourceType>
{
    uint64 handle;

    [ForceInline]
    T get() { return __getGlobalResourceHeap<T.ResourceKind>().get<T>(handle); }
} 
dubiousconst282 commented 6 days ago

From what I can tell, normal descriptors already suffer with mostly the same conflict issues, although I can see why having a more strict solution like you described is preferred.

I think having the binding points be defined via a compiler option instead instead of implicitly would be better.

Ideally the finer resource types that Vulkan defines compared to HLSL should be considered as well. Although VK_EXT_mutable_descriptor_type makes this less of an issue, it is still more limiting in some ways (not as widely supported as plain descriptor indexing, can't bind a texture as Image and SampledImage at the same time).

csyonghe commented 6 days ago

I still want to have this PR merged in. No matter what we do for Bindless they will be based on the changes here.

dubiousconst282 commented 6 days ago

Ah, sorry, I missed that __getGlobalResourceHeap will reuse what's implemented here.

dubiousconst282 commented 5 days ago

The tests incomplete-type.slang and constant-buffer-unsized.slang fail when the stdlib declares an extension for ConstantBuffer like attached below. I've had no luck debugging so far, maybe you have some hint for me to look for?

__generic<T>
extension ConstantBuffer<T> : __IDynamicResourceCastable<__DynamicResourceKind.General>
{
    __intrinsic_op($(kIROp_CastDynamicResource))
    __implicit_conversion($(kConversionCost_GenericParamUpcast))
    __init(__DynamicResource res);
}
csyonghe commented 5 days ago

What kind of failure are you seeing? ConstantBuffer has an implicit conversion to T, adding another conversion here may be changing the diagnostic message that leads to a false failure.

dubiousconst282 commented 5 days ago

Both tests are supposed to trigger a diagnostic about invalid usage of ConstantBuffer, but this never happens when the extensions (or really any other __init function) are present.


It turns out that this diagnostic check happens in the same method where default variable initialization is handled, and it will early return if there is no appliable initializer function for the variable type:

https://github.com/shader-slang/slang/blob/323fdff4544ac3b640bdc617d7f7d64ab56fbf5a/source/slang/slang-check-decl.cpp#L2082-L2086 https://github.com/shader-slang/slang/blob/323fdff4544ac3b640bdc617d7f7d64ab56fbf5a/source/slang/slang-check-decl.cpp#L2136-L2146

I don't think the early returns are correct, so I've changed them to a else block. I couldn't get all of the tests to run on my laptop so I'll wait to see what CI thinks