Implement GetDimensions for structured buffers on WGSL (#6609)

* Implement structured buffer GetDimensions for WGSL

* byte address buffer fix

* enable wgsl test and minor fixups

* maybe fix get equivalent structured buffer

* remove unnecessary include

* clean up some code in meta file

3 files changed, 39 insertions, 11 deletions

diff --git a/source/slang/hlsl.meta.slang b/source/slang/hlsl.meta.slang
index c26a7613b..c7e4925be 100644
--- a/source/slang/hlsl.meta.slang
+++ b/source/slang/hlsl.meta.slang
@@ -67,15 +67,15 @@ __intrinsic_op($(kIROp_StructuredBufferGetDimensions))
 uint2 __structuredBufferGetDimensions(ConsumeStructuredBuffer<T,L> buffer);
 
 __intrinsic_op($(kIROp_StructuredBufferGetDimensions))
-[require(cpp_cuda_glsl_hlsl_metal_spirv, structuredbuffer)]
+[require(cpp_cuda_glsl_hlsl_metal_spirv_wgsl, structuredbuffer)]
 uint2 __structuredBufferGetDimensions<T,L:IBufferDataLayout>(StructuredBuffer<T,L> buffer);
 
 __intrinsic_op($(kIROp_StructuredBufferGetDimensions))
-[require(cpp_cuda_glsl_hlsl_metal_spirv, structuredbuffer_rw)]
+[require(cpp_cuda_glsl_hlsl_metal_spirv_wgsl, structuredbuffer_rw)]
 uint2 __structuredBufferGetDimensions<T,L:IBufferDataLayout>(RWStructuredBuffer<T,L> buffer);
 
 __intrinsic_op($(kIROp_StructuredBufferGetDimensions))
-[require(cpp_cuda_glsl_hlsl_metal_spirv, structuredbuffer_rw)]
+[require(cpp_cuda_glsl_hlsl_metal_spirv_wgsl, structuredbuffer_rw)]
 uint2 __structuredBufferGetDimensions<T,L:IBufferDataLayout>(RasterizerOrderedStructuredBuffer<T,L> buffer);
 
 //@public:
@@ -143,7 +143,7 @@ struct ByteAddressBuffer
     ///@param[out] dim The number of bytes in the buffer.
     [__readNone]
     [ForceInline]
-    [require(cpp_cuda_glsl_hlsl_metal_spirv, structuredbuffer)]
+    [require(cpp_cuda_glsl_hlsl_metal_spirv_wgsl, structuredbuffer)]
     void GetDimensions(out uint dim)
     {
         __target_switch
@@ -151,9 +151,7 @@ struct ByteAddressBuffer
         case cpp: __intrinsic_asm ".GetDimensions";
         case cuda: __intrinsic_asm ".GetDimensions";
         case hlsl: __intrinsic_asm ".GetDimensions";
-        case glsl:
-        case metal:
-        case spirv:
+        default:
             dim = __structuredBufferGetDimensions(__getEquivalentStructuredBuffer<uint>(this)).x*4;
         }
     }
@@ -4922,7 +4920,7 @@ struct $(item.name)
     /// Get the number of bytes in the buffer.
     ///@param[out] dim The number of bytes in the buffer.
     [ForceInline]
-    [require(cpp_cuda_glsl_hlsl_spirv, structuredbuffer_rw)]
+    [require(cpp_cuda_glsl_hlsl_spirv_wgsl, structuredbuffer_rw)]
     void GetDimensions(out uint dim)
     {
         __target_switch
@@ -4930,8 +4928,7 @@ struct $(item.name)
         case cpp: __intrinsic_asm ".GetDimensions";
         case cuda: __intrinsic_asm ".GetDimensions";
         case hlsl: __intrinsic_asm ".GetDimensions";
-        case glsl:
-        case spirv:
+        default:
             dim = __structuredBufferGetDimensions(__getEquivalentStructuredBuffer<uint>(this)).x*4;
         }
     }
@@ -5939,7 +5936,7 @@ struct $(item.name)
     /// @param stride The stride, in bytes, of each structure element.
     [__readNone]
     [ForceInline]
-    [require(cpp_cuda_glsl_hlsl_metal_spirv, structuredbuffer_rw)]
+    [require(cpp_cuda_glsl_hlsl_metal_spirv_wgsl, structuredbuffer_rw)]
     void GetDimensions(
         out uint numStructs,
         out uint stride)
diff --git a/source/slang/slang-emit-wgsl.cpp b/source/slang/slang-emit-wgsl.cpp
index d87cd06de..227c6cf75 100644
--- a/source/slang/slang-emit-wgsl.cpp
+++ b/source/slang/slang-emit-wgsl.cpp
@@ -1,5 +1,6 @@
 #include "slang-emit-wgsl.h"
 
+#include "slang-ir-layout.h"
 #include "slang-ir-util.h"
 
 // A note on row/column "terminology reversal".
@@ -1240,6 +1241,35 @@ bool WGSLSourceEmitter::tryEmitInstStmtImpl(IRInst* inst)
             m_writer->emit("(1));\n");
             return true;
         }
+    case kIROp_StructuredBufferGetDimensions:
+        {
+            IRIntegerValue strideValue;
+            auto dataType = inst->getOperand(0)->getDataType();
+            auto structuredBufferType = as<IRHLSLStructuredBufferTypeBase>(dataType);
+            if (structuredBufferType)
+            {
+                auto elementType = structuredBufferType->getElementType();
+                auto sizeDecor = elementType->findDecoration<IRSizeAndAlignmentDecoration>();
+                SLANG_ASSERT(sizeDecor);
+                strideValue = align(sizeDecor->getSize(), (int)sizeDecor->getAlignment());
+            }
+            else
+            {
+                SLANG_ASSERT(as<IRByteAddressBufferTypeBase>(dataType));
+                // ByteAddressBuffer(s) are an array of 32 bit integers, stride is 4 bytes.
+                strideValue = 4;
+            }
+
+            emitInstResultDecl(inst);
+            m_writer->emit("vec2<u32>(");
+            m_writer->emit("arrayLength(&");
+            emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+            m_writer->emit(")");
+            m_writer->emit(", ");
+            m_writer->emit(strideValue);
+            m_writer->emit(");\n");
+            return true;
+        }
     }
 }
 
diff --git a/tests/cross-compile/get-dimensions.slang b/tests/cross-compile/get-dimensions.slang
index 27cbb18de..8c1acd813 100644
--- a/tests/cross-compile/get-dimensions.slang
+++ b/tests/cross-compile/get-dimensions.slang
@@ -3,6 +3,7 @@
 //TEST(compute):COMPARE_COMPUTE_EX:-d3d12 -compute -shaderobj
 //TEST(compute):COMPARE_COMPUTE_EX:-d3d11 -compute -shaderobj
 //TEST(compute):COMPARE_COMPUTE_EX:-cpu -compute -shaderobj
+//TEST(compute):COMPARE_COMPUTE_EX:-wgpu -compute -shaderobj
 
 struct Thing
 {