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.

1 files changed, 50 insertions, 4 deletions

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(">");
             }
         }