From 5fa35fcce532267a2ae5779dee9ff4d07fab6bf4 Mon Sep 17 00:00:00 2001
From: Anders Leino <aleino@nvidia.com>
Date: Fri, 11 Oct 2024 10:05:15 +0300
Subject: WGSL: Enable load & store from byte-addressible buffers (#5252)

---
 source/slang/slang-emit-wgsl.cpp | 40 ++++++++++++++++++++++++++++++++++++++--
 1 file changed, 38 insertions(+), 2 deletions(-)

(limited to 'source/slang/slang-emit-wgsl.cpp')

diff --git a/source/slang/slang-emit-wgsl.cpp b/source/slang/slang-emit-wgsl.cpp
index a05e03afc..b1a723dc5 100644
--- a/source/slang/slang-emit-wgsl.cpp
+++ b/source/slang/slang-emit-wgsl.cpp
@@ -348,6 +348,13 @@ void WGSLSourceEmitter::emitSimpleTypeImpl(IRType* type)
     }
     break;
 
+    case kIROp_HLSLByteAddressBufferType:
+    case kIROp_HLSLRWByteAddressBufferType:
+    {
+        m_writer->emit("array<u32>");
+    }
+    break;
+
     case kIROp_VoidType:
     {
         // There is no void type in WGSL.
@@ -590,13 +597,15 @@ void WGSLSourceEmitter::emitVarKeywordImpl(IRType * type, IRInst* varDecl)
         m_writer->emit("<workgroup>");
     }
     else if (type->getOp() == kIROp_HLSLRWStructuredBufferType ||
-        type->getOp() == kIROp_HLSLRasterizerOrderedStructuredBufferType)
+        type->getOp() == kIROp_HLSLRasterizerOrderedStructuredBufferType ||
+        type->getOp() == kIROp_HLSLRWByteAddressBufferType)
     {
         m_writer->emit("<");
         m_writer->emit("storage, read_write");
         m_writer->emit(">");
     }
-    else if (type->getOp() == kIROp_HLSLStructuredBufferType)
+    else if (type->getOp() == kIROp_HLSLStructuredBufferType ||
+        type->getOp() == kIROp_HLSLByteAddressBufferType)
     {
         m_writer->emit("<");
         m_writer->emit("storage, read");
@@ -1178,6 +1187,33 @@ bool WGSLSourceEmitter::tryEmitInstExprImpl(IRInst* inst, const EmitOpInfo& inOu
         }
     }
     break;
+
+    case kIROp_ByteAddressBufferLoad:
+    {
+        // Indices in Slang code count bytes, but in WASM they count u32's since
+        // byte address buffers translate to array<u32> in WASM, so divide by 4.
+        emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+        m_writer->emit("[(");
+        emitOperand(inst->getOperand(1), getInfo(EmitOp::General));
+        m_writer->emit(")/4]");
+        return true;
+    }
+    break;
+
+    case kIROp_ByteAddressBufferStore:
+    {
+        // Indices in Slang code count bytes, but in WASM they count u32's since
+        // byte address buffers translate to array<u32> in WASM, so divide by 4.
+        auto base = inst->getOperand(0);
+        emitOperand(base, EmitOpInfo());
+        m_writer->emit("[(");
+        emitOperand(inst->getOperand(1), getInfo(EmitOp::General));
+        m_writer->emit(")/4] = ");
+        emitOperand(inst->getOperand(inst->getOperandCount() - 1), getInfo(EmitOp::General));
+        return true;
+    }
+    break;
+
     }
 
     return false;
-- 
cgit v1.2.3