Implement texture_storage_Xd in WGSL (#5158)

* Implement texture_storage_Xd in WGSL

This commit implements `texture_storage_Xd` in WGSL, which is similar
to RWTextureXD in HLSL.

It is intresting that `texture_storage_Xd` doesn't take the shader
type as its input argument at all.
Instead, it takes "texel format" enum value as its first template
parameter, which can be found here:
https://www.w3.org/TR/WGSL/#storage-texel-formats

As an example, `texture_storage_2d<rg32uint, read_write>` expects
vec4<u32> as a value type for `Load` and `Store`, where Z-component
will be ignored and treated as zero and W-component will be treated
always as 1. The type `u32` is inferred from the enum value `rg32uint`.
Note that the number of component is always fixed to 4 regardless how
many components are actually stored.

3 files changed, 221 insertions, 17 deletions

diff --git a/source/slang/hlsl.meta.slang b/source/slang/hlsl.meta.slang
index c646f5cb7..fd5c2ebb7 100644
--- a/source/slang/hlsl.meta.slang
+++ b/source/slang/hlsl.meta.slang
@@ -3300,7 +3300,7 @@ extension __TextureImpl<T,Shape,isArray,0,sampleCount,$(access),isShadow, 0,form
                     || Shape.flavor == $(SLANG_TEXTURE_3D)
                     , "WGSL doesn't supports textureLoad for Cube texture.");
                 static_assert(isArray == 0 || Shape.flavor == $(SLANG_TEXTURE_2D)
-                    , "WGSL supports textureLoad for texture_storage_2d_array but not for array of 1D, 3D or Cube.");
+                    , "WGSL supports textureLoad for 2d_array but not for array of 1D, 3D or Cube.");
                 static_assert(isShadow == 0 || T is float
                     , "WGSL supports only f32 depth textures");
 
@@ -3477,21 +3477,43 @@ extension __TextureImpl<T,Shape,isArray,0,sampleCount,$(access),isShadow, 0,form
                         break;
                     }
             case wgsl:
-                static_assert(T is float || T is vector<float,2> || T is vector<float,3> || T is vector<float,4>
-                    , "WGSL supports only f32 type textures");
+                static_assert(Shape.flavor == $(SLANG_TEXTURE_1D)
+                    || Shape.flavor == $(SLANG_TEXTURE_2D)
+                    || Shape.flavor == $(SLANG_TEXTURE_3D)
+                    , "WGSL doesn't supports textureStore for Cube texture.");
+                static_assert(isArray == 0 || Shape.flavor == $(SLANG_TEXTURE_2D)
+                    , "WGSL supports textureStore for texture_store_2d_array but not for array of 1D, 3D or Cube.");
+
+                // WGSL requires the value type to be always `vec4`
                 if (isArray == 1)
                 {
-                    switch (Shape.flavor)
-                    {
-                    case $(SLANG_TEXTURE_1D):
-                        __intrinsic_asm "textureStore($0, ($1).x, i32(($1).y), $2)$z";
-                    case $(SLANG_TEXTURE_2D):
-                        __intrinsic_asm "textureStore($0, ($1).xy, i32(($1).z), $2)$z";
-                    case $(SLANG_TEXTURE_CUBE):
-                        __intrinsic_asm "textureStore($0, ($1).xyz, i32(($1).w), $2)$z";
-                    }
+                    if (T is int32_t   || T is int16_t  || T is int8_t)  __intrinsic_asm "textureStore($0, ($1).xy, i32(($1).z), vec4<i32>($2, 0, 0, 1))";
+                    if (T is int32_t2  || T is int16_t  || T is int8_t)  __intrinsic_asm "textureStore($0, ($1).xy, i32(($1).z), vec4<i32>($2, 0, 1))";
+                    if (T is int32_t3  || T is int16_t  || T is int8_t)  __intrinsic_asm "textureStore($0, ($1).xy, i32(($1).z), vec4<i32>($2, 1))";
+                    if (T is uint32_t  || T is uint16_t || T is uint8_t) __intrinsic_asm "textureStore($0, ($1).xy, i32(($1).z), vec4<u32>($2, 0, 0, 1))";
+                    if (T is uint32_t2 || T is uint16_t || T is uint8_t) __intrinsic_asm "textureStore($0, ($1).xy, i32(($1).z), vec4<u32>($2, 0, 1))";
+                    if (T is uint32_t3 || T is uint16_t || T is uint8_t) __intrinsic_asm "textureStore($0, ($1).xy, i32(($1).z), vec4<u32>($2, 1))";
+                    if (T is half)   __intrinsic_asm "textureStore($0, ($1).xy, i32(($1).z), vec4<f16>($2, 0, 0, 1))";
+                    if (T is half2)  __intrinsic_asm "textureStore($0, ($1).xy, i32(($1).z), vec4<f16>($2, 0, 1))";
+                    if (T is half3)  __intrinsic_asm "textureStore($0, ($1).xy, i32(($1).z), vec4<f16>($2, 1))";
+                    if (T is float)  __intrinsic_asm "textureStore($0, ($1).xy, i32(($1).z), vec4<f32>($2, 0, 0, 1))";
+                    if (T is float2) __intrinsic_asm "textureStore($0, ($1).xy, i32(($1).z), vec4<f32>($2, 0, 1))";
+                    if (T is float3) __intrinsic_asm "textureStore($0, ($1).xy, i32(($1).z), vec4<f32>($2, 1))";
+                    __intrinsic_asm "textureStore($0, ($1).xy, i32(($1).z), $2)";
                 }
-                __intrinsic_asm "textureStore($0, $1, $2)$z";
+                if (T is int32_t   || T is int16_t  || T is int8_t)  __intrinsic_asm "textureStore($0, $1, vec4<i32>($2, 0, 0, 1))";
+                if (T is int32_t2  || T is int16_t  || T is int8_t)  __intrinsic_asm "textureStore($0, $1, vec4<i32>($2, 0, 1))";
+                if (T is int32_t3  || T is int16_t  || T is int8_t)  __intrinsic_asm "textureStore($0, $1, vec4<i32>($2, 1))";
+                if (T is uint32_t  || T is uint16_t || T is uint8_t) __intrinsic_asm "textureStore($0, $1, vec4<u32>($2, 0, 0, 1))";
+                if (T is uint32_t2 || T is uint16_t || T is uint8_t) __intrinsic_asm "textureStore($0, $1, vec4<u32>($2, 0, 1))";
+                if (T is uint32_t3 || T is uint16_t || T is uint8_t) __intrinsic_asm "textureStore($0, $1, vec4<u32>($2, 1))";
+                if (T is half)   __intrinsic_asm "textureStore($0, $1, vec4<f16>($2, 0, 0, 1))";
+                if (T is half2)  __intrinsic_asm "textureStore($0, $1, vec4<f16>($2, 0, 1))";
+                if (T is half3)  __intrinsic_asm "textureStore($0, $1, vec4<f16>($2, 1))";
+                if (T is float)  __intrinsic_asm "textureStore($0, $1, vec4<f32>($2, 0, 0, 1))";
+                if (T is float2) __intrinsic_asm "textureStore($0, $1, vec4<f32>($2, 0, 1))";
+                if (T is float3) __intrinsic_asm "textureStore($0, $1, vec4<f32>($2, 1))";
+                __intrinsic_asm "textureStore($0, $1, $2)";
             }
         }
 
diff --git a/source/slang/slang-emit-wgsl.cpp b/source/slang/slang-emit-wgsl.cpp
index 051cdb820..66e802571 100644
--- a/source/slang/slang-emit-wgsl.cpp
+++ b/source/slang/slang-emit-wgsl.cpp
@@ -298,6 +298,39 @@ void WGSLSourceEmitter::emit(const AddressSpace addressSpace)
     }
 }
 
+static const char* getWgslImageFormat(IRTextureTypeBase* type)
+{
+    // You can find the supported WGSL texel format from the URL:
+    // https://www.w3.org/TR/WGSL/#storage-texel-formats
+    //
+    ImageFormat imageFormat = type->hasFormat() ? (ImageFormat)type->getFormat() : ImageFormat::unknown;
+    switch (imageFormat)
+    {
+    case ImageFormat::rgba8:       return "rgba8unorm";
+    case ImageFormat::rgba8_snorm: return "rgba8snorm";
+    case ImageFormat::rgba8ui:     return "rgba8uint";
+    case ImageFormat::rgba8i:      return "rgba8sint";
+    case ImageFormat::rgba16ui:    return "rgba16uint";
+    case ImageFormat::rgba16i:     return "rgba16sint";
+    case ImageFormat::rgba16f:     return "rgba16float";
+    case ImageFormat::r32ui:       return "r32uint";
+    case ImageFormat::r32i:        return "r32sint";
+    case ImageFormat::r32f:        return "r32float";
+    case ImageFormat::rg32ui:      return "rg32uint";
+    case ImageFormat::rg32i:       return "rg32sint";
+    case ImageFormat::rg32f:       return "rg32float";
+    case ImageFormat::rgba32ui:    return "rgba32uint";
+    case ImageFormat::rgba32i:     return "rgba32sint";
+    case ImageFormat::rgba32f:     return "rgba32float";
+    case ImageFormat::unknown:
+        // Unlike SPIR-V, WGSL doesn't have a texel format for "unknown".
+        return "rgba32float";
+    default:
+        // We may need to print a warning for types WGSL doesn't support
+        return "rgba32float";
+    }
+}
+
 void WGSLSourceEmitter::emitSimpleTypeImpl(IRType* type)
 {
     switch (type->getOp())
@@ -467,11 +500,24 @@ void WGSLSourceEmitter::emitSimpleTypeImpl(IRType* type)
             if (!texType->isShadow())
             {
                 m_writer->emit("<");
+
                 auto elemType = texType->getElementType();
-                if (auto vecElemType = as<IRVectorType>(elemType))
-                    emitSimpleType(vecElemType->getElementType());
-                else
-                    emitType(elemType);
+
+                switch (texType->getAccess())
+                {
+                case SLANG_RESOURCE_ACCESS_READ_WRITE:
+                    m_writer->emit(getWgslImageFormat(texType));
+                    m_writer->emit(", read_write");
+                    break;
+
+                default:
+                    if (auto vecElemType = as<IRVectorType>(elemType))
+                        emitSimpleType(vecElemType->getElementType());
+                    else
+                        emitType(elemType);
+                    break;
+                }
+
                 m_writer->emit(">");
             }
         }
diff --git a/tests/wgsl/texture-storage.slang b/tests/wgsl/texture-storage.slang
new file mode 100644
index 000000000..84921870b
--- /dev/null
+++ b/tests/wgsl/texture-storage.slang
@@ -0,0 +1,136 @@
+//TEST:SIMPLE(filecheck=WGSL): -stage fragment -entry fragMain -target wgsl
+
+// In WGSL, `textureSample` and other variants work only for f32 type.
+// But textureLoad works for i32, u32 and f32.
+
+//TEST_INPUT: ubuffer(data=[0], stride=4):out,name outputBuffer
+RWStructuredBuffer<int> outputBuffer;
+
+// f32 types
+
+//TEST_INPUT: RWTexture1D(format=R32G32_FLOAT, size=4, content = zero):name w1D_f32v2
+//TEST_INPUT: RWTexture2D(format=R32G32_FLOAT, size=4, content = zero):name w2D_f32v2
+//TEST_INPUT: RWTexture3D(format=R32G32_FLOAT, size=4, content = zero):name w3D_f32v2
+//TEST_INPUT: RWTexture2D(format=R32G32_FLOAT, size=4, content = zero, arrayLength=2):name w2DArray_f32v2
+// WGSL: var w1D_f32v2{{[^:]*}}: texture_storage_1d<rg32float, read_write>
+[format("rg32f")] RWTexture1D<float2>      w1D_f32v2;
+[format("rg32f")] RWTexture2D<float2>      w2D_f32v2;
+[format("rg32f")] RWTexture3D<float2>      w3D_f32v2;
+[format("rg32f")] RWTexture2DArray<float2> w2DArray_f32v2;
+
+//TEST_INPUT: RWTexture1D(format=R8G8B8A8_UNORM, size=4, content = zero):name w1D_f32v4
+//TEST_INPUT: RWTexture2D(format=R8G8B8A8_UNORM, size=4, content = zero):name w2D_f32v4
+//TEST_INPUT: RWTexture3D(format=R8G8B8A8_UNORM, size=4, content = zero):name w3D_f32v4
+//TEST_INPUT: RWTexture2D(format=R8G8B8A8_UNORM, size=4, content = zero, arrayLength=2):name w2DArray_f32v4
+// WGSL: var w1D_f32v4{{[^:]*}}: texture_storage_1d<rgba8unorm, read_write>
+[format("rgba8")] RWTexture1D<float4>      w1D_f32v4;
+[format("rgba8")] RWTexture2D<float4>      w2D_f32v4;
+[format("rgba8")] RWTexture3D<float4>      w3D_f32v4;
+[format("rgba8")] RWTexture2DArray<float4> w2DArray_f32v4;
+
+// i32 types
+
+//TEST_INPUT: RWTexture1D(format=R32G32_SINT, size=4, content = zero):name w1D_i32v2
+//TEST_INPUT: RWTexture2D(format=R32G32_SINT, size=4, content = zero):name w2D_i32v2
+//TEST_INPUT: RWTexture3D(format=R32G32_SINT, size=4, content = zero):name w3D_i32v2
+//TEST_INPUT: RWTexture2D(format=R32G32_SINT, size=4, content = zero, arrayLength=2):name w2DArray_i32v2
+// WGSL: var w1D_i32v2{{[^:]*}}: texture_storage_1d<rg32sint, read_write>
+[format("rg32i")] RWTexture1D<int2>      w1D_i32v2;
+[format("rg32i")] RWTexture2D<int2>      w2D_i32v2;
+[format("rg32i")] RWTexture3D<int2>      w3D_i32v2;
+[format("rg32i")] RWTexture2DArray<int2> w2DArray_i32v2;
+
+//TEST_INPUT: RWTexture1D(format=R32G32B32A32_SINT, size=4, content = zero):name w1D_i32v4
+//TEST_INPUT: RWTexture2D(format=R32G32B32A32_SINT, size=4, content = zero):name w2D_i32v4
+//TEST_INPUT: RWTexture3D(format=R32G32B32A32_SINT, size=4, content = zero):name w3D_i32v4
+//TEST_INPUT: RWTexture2D(format=R32G32B32A32_SINT, size=4, content = zero, arrayLength=2):name w2DArray_i32v4
+// WGSL: var w1D_i32v4{{[^:]*}}: texture_storage_1d<rgba32sint, read_write>
+[format("rgba32i")] RWTexture1D<int4>      w1D_i32v4;
+[format("rgba32i")] RWTexture2D<int4>      w2D_i32v4;
+[format("rgba32i")] RWTexture3D<int4>      w3D_i32v4;
+[format("rgba32i")] RWTexture2DArray<int4> w2DArray_i32v4;
+
+// u32 types
+
+//TEST_INPUT: RWTexture1D(format=R32G32_UINT, size=4, content = zero):name w1D_u32v2
+//TEST_INPUT: RWTexture2D(format=R32G32_UINT, size=4, content = zero):name w2D_u32v2
+//TEST_INPUT: RWTexture3D(format=R32G32_UINT, size=4, content = zero):name w3D_u32v2
+//TEST_INPUT: RWTexture2D(format=R32G32_UINT, size=4, content = zero, arrayLength=2):name w2DArray_u32v2
+// WGSL: var w1D_u32v2{{[^:]*}}: texture_storage_1d<rg32uint, read_write>
+[format("rg32ui")] RWTexture1D<uint2>      w1D_u32v2;
+[format("rg32ui")] RWTexture2D<uint2>      w2D_u32v2;
+[format("rg32ui")] RWTexture3D<uint2>      w3D_u32v2;
+[format("rg32ui")] RWTexture2DArray<uint2> w2DArray_u32v2;
+
+//TEST_INPUT: RWTexture1D(format=R16G16B16A16_UINT, size=4, content = zero):name w1D_u32v4
+//TEST_INPUT: RWTexture2D(format=R16G16B16A16_UINT, size=4, content = zero):name w2D_u32v4
+//TEST_INPUT: RWTexture3D(format=R16G16B16A16_UINT, size=4, content = zero):name w3D_u32v4
+//TEST_INPUT: RWTexture2D(format=R16G16B16A16_UINT, size=4, content = zero, arrayLength=2):name w2DArray_u32v4
+// WGSL: var w1D_u32v4{{[^:]*}}: texture_storage_1d<rgba16uint, read_write>
+[format("rgba16ui")] RWTexture1D<uint4>      w1D_u32v4;
+[format("rgba16ui")] RWTexture2D<uint4>      w2D_u32v4;
+[format("rgba16ui")] RWTexture3D<uint4>      w3D_u32v4;
+[format("rgba16ui")] RWTexture2DArray<uint4> w2DArray_u32v4;
+
+
+__generic<T:__BuiltinArithmeticType, let N:int, let sampleIndex:int, let format:int>
+[ForceInline] // Workaround for a WGSL requirement that `texture_storage_Xd` must always be a global variable
+bool TEST_textureStorage_StoreLoad(
+    RWTexture1D<vector<T,N>, sampleIndex, format> w1D,
+    RWTexture2D<vector<T,N>, sampleIndex, format> w2D,
+    RWTexture3D<vector<T,N>, sampleIndex, format> w3D,
+    RWTexture2DArray<vector<T,N>, sampleIndex, format> w2DArray)
+{
+    typealias Tvn = vector<T,N>;
+
+    // ===================
+    // o[i] = v;
+    // https://www.w3.org/TR/WGSL/#texturestore
+    // ===================
+
+    // TODO: store before load
+    // WGSL: textureStore({{\(*}}w1D
+    w1D[0] = Tvn(T(1));
+
+    // WGSL: textureStore({{\(*}}w2D
+    w2D[0] = Tvn(T(1));
+
+    // WGSL: textureStore({{\(*}}w3D
+    w3D[0] = Tvn(T(1));
+
+    // WGSL: textureStore({{\(*}}w2DArray
+    w2DArray[0] = Tvn(T(1));
+
+    return true
+        // ===================
+        // T Load()
+        // https://www.w3.org/TR/WGSL/#textureload
+        // ===================
+
+        // WGSL: textureLoad({{\(*}}w1D
+        && all(Tvn(T(0)) == w1D.Load(0))
+
+        // WGSL: textureLoad({{\(*}}w2D
+        && all(Tvn(T(0)) == w2D.Load(int2(0, 0)))
+
+        // WGSL: textureLoad({{\(*}}w3D
+        && all(Tvn(T(0)) == w3D.Load(int3(0, 0, 0)))
+
+        // WGSL: textureLoad({{\(*}}w2DArray
+        && all(Tvn(T(0)) == w2DArray.Load(int3(0, 0, 0)))
+        ;
+}
+
+void fragMain()
+{
+    bool result = true
+        && TEST_textureStorage_StoreLoad<float,    2>(w1D_f32v2, w2D_f32v2, w3D_f32v2, w2DArray_f32v2)
+        && TEST_textureStorage_StoreLoad<float,    4>(w1D_f32v4, w2D_f32v4, w3D_f32v4, w2DArray_f32v4)
+        && TEST_textureStorage_StoreLoad<int32_t,  2>(w1D_i32v2, w2D_i32v2, w3D_i32v2, w2DArray_i32v2)
+        && TEST_textureStorage_StoreLoad<int32_t,  4>(w1D_i32v4, w2D_i32v4, w3D_i32v4, w2DArray_i32v4)
+        && TEST_textureStorage_StoreLoad<uint32_t, 2>(w1D_u32v2, w2D_u32v2, w3D_u32v2, w2DArray_u32v2)
+        && TEST_textureStorage_StoreLoad<uint32_t, 4>(w1D_u32v4, w2D_u32v4, w3D_u32v4, w2DArray_u32v4)
+        ;
+
+    outputBuffer[0] = int(result);
+}