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.

2 files changed, 85 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(">");
             }
         }