Lower all ByteAddressBuffer uses for SPIRV. (#3143)

Co-authored-by: Yong He <yhe@nvidia.com>

11 files changed, 253 insertions, 37 deletions

diff --git a/source/slang/hlsl.meta.slang b/source/slang/hlsl.meta.slang
index cfc2bc3a5..ec7809b90 100644
--- a/source/slang/hlsl.meta.slang
+++ b/source/slang/hlsl.meta.slang
@@ -14,6 +14,15 @@ __generic<T>
 __intrinsic_op($(kIROp_StructuredBufferGetDimensions))
 uint2 __structuredBufferGetDimensions(ConsumeStructuredBuffer<T> buffer);
 
+__intrinsic_op($(kIROp_StructuredBufferGetDimensions))
+uint2 __structuredBufferGetDimensions<T>(StructuredBuffer<T> buffer);
+
+__intrinsic_op($(kIROp_StructuredBufferGetDimensions))
+uint2 __structuredBufferGetDimensions<T>(RWStructuredBuffer<T> buffer);
+
+__intrinsic_op($(kIROp_StructuredBufferGetDimensions))
+uint2 __structuredBufferGetDimensions<T>(RasterizerOrderedStructuredBuffer<T> buffer);
+
 __generic<T>
 __magic_type(HLSLAppendStructuredBufferType)
 __intrinsic_type($(kIROp_HLSLAppendStructuredBufferType))
@@ -37,35 +46,61 @@ __magic_type(HLSLByteAddressBufferType)
 __intrinsic_type($(kIROp_HLSLByteAddressBufferType))
 struct ByteAddressBuffer
 {
-    __target_intrinsic(glsl, "$1 = $0._data.length() * 4")
     [__readNone]
-    void GetDimensions(
-        out uint dim);
+    __target_intrinsic(hlsl)
+    __target_intrinsic(cpp)
+    __target_intrinsic(cuda)
+    [__unsafeForceInlineEarly]
+    void GetDimensions(out uint dim);
+
+    [__unsafeForceInlineEarly]
+    __specialized_for_target(spirv)
+    __specialized_for_target(glsl)
+    void GetDimensions(out uint dim)
+    {
+        dim = __structuredBufferGetDimensions(__getEquivalentStructuredBuffer<uint>(this)).x*4;
+    }
 
-    __target_intrinsic(glsl, "$0._data[$1/4]")
     [__readNone]
-    uint Load(int location);
+    [ForceInline]
+    __target_intrinsic(hlsl)
+    uint Load(int location)
+    {
+        return __byteAddressBufferLoad<uint>(this, location);
+    }
 
     [__readNone]
     uint Load(int location, out uint status);
 
-    __target_intrinsic(glsl, "uvec2($0._data[$1/4], $0._data[$1/4+1])")
     [__readNone]
-    uint2 Load2(int location);
+    [ForceInline]
+    __target_intrinsic(hlsl)
+    uint2 Load2(int location)
+    {
+        return __byteAddressBufferLoad<uint2>(this, location);
+    }
 
     [__readNone]
     uint2 Load2(int location, out uint status);
 
-    __target_intrinsic(glsl, "uvec3($0._data[$1/4], $0._data[$1/4+1], $0._data[$1/4+2])")
     [__readNone]
-    uint3 Load3(int location);
+    [ForceInline]
+    __target_intrinsic(hlsl)
+    uint3 Load3(int location)
+    {
+        return __byteAddressBufferLoad<uint3>(this, location);
+    }
 
     [__readNone]
     uint3 Load3(int location, out uint status);
 
-    __target_intrinsic(glsl, "uvec4($0._data[$1/4], $0._data[$1/4+1], $0._data[$1/4+2], $0._data[$1/4+3])")
     [__readNone]
-    uint4 Load4(int location);
+    [ForceInline]
+    __target_intrinsic(hlsl)
+    uint4 Load4(int location)
+    {
+        return __byteAddressBufferLoad<uint4>(this, location);
+    }
 
     [__readNone]
     uint4 Load4(int location, out uint status);
@@ -244,11 +279,16 @@ __magic_type(HLSLStructuredBufferType)
 __intrinsic_type($(kIROp_HLSLStructuredBufferType))
 struct StructuredBuffer
 {
-    __target_intrinsic(glsl, "$1 = $0._data.length(); $2 = 0")
     [__readNone]
+    [__unsafeForceInlineEarly]
     void GetDimensions(
         out uint numStructs,
-        out uint stride);
+        out uint stride)
+    {
+        let rs = __structuredBufferGetDimensions(this);
+        numStructs = rs.x;
+        stride = rs.y;
+    }
 
     __intrinsic_op($(kIROp_StructuredBufferLoad))
     __target_intrinsic(glsl, "$0._data[$1]")
@@ -321,34 +361,56 @@ struct $(item.name)
     // Note(tfoley): supports all operations from `ByteAddressBuffer`
     // TODO(tfoley): can this be made a sub-type?
 
-    __target_intrinsic(glsl, "$1 = $0._data.length() * 4")
-    void GetDimensions(
-        out uint dim);
+    __target_intrinsic(hlsl)
+    __target_intrinsic(cpp)
+    __target_intrinsic(cuda)
+    [__unsafeForceInlineEarly]
+    void GetDimensions(out uint dim);
 
-    __target_intrinsic(glsl, "$0._data[$1/4]")
+    [__unsafeForceInlineEarly]
+    __specialized_for_target(spirv)
+    __specialized_for_target(glsl)
+    void GetDimensions(out uint dim)
+    {
+        dim = __structuredBufferGetDimensions(__getEquivalentStructuredBuffer<uint>(this)).x*4;
+    }
+
+    __target_intrinsic(hlsl)
     [__NoSideEffect]
-    uint Load(int location);
+    uint Load(int location)
+    {
+        return __byteAddressBufferLoad<uint>(this, location);
+    }
 
     [__NoSideEffect]
     uint Load(int location, out uint status);
 
-    __target_intrinsic(glsl, "uvec2($0._data[$1/4], $0._data[$1/4+1])")
+    __target_intrinsic(hlsl)
     [__NoSideEffect]
-    uint2 Load2(int location);
+    uint2 Load2(int location)
+    {
+        return __byteAddressBufferLoad<uint2>(this, location);
+    }
 
     [__NoSideEffect]
     uint2 Load2(int location, out uint status);
 
-    __target_intrinsic(glsl, "uvec3($0._data[$1/4], $0._data[$1/4+1], $0._data[$1/4+2])")
+    __target_intrinsic(hlsl)
     [__NoSideEffect]
-    uint3 Load3(int location);
+    uint3 Load3(int location)
+    {
+        return __byteAddressBufferLoad<uint3>(this, location);
+    }
 
     [__NoSideEffect]
     uint3 Load3(int location, out uint status);
 
-    __target_intrinsic(glsl, "uvec4($0._data[$1/4], $0._data[$1/4+1], $0._data[$1/4+2], $0._data[$1/4+3])")
+    __target_intrinsic(hlsl)
     [__NoSideEffect]
-    uint4 Load4(int location);
+    uint4 Load4(int location)
+    {
+        return __byteAddressBufferLoad<uint4>(this, location);
+    }
 
     [__NoSideEffect]
     uint4 Load4(int location, out uint status);
@@ -689,25 +751,39 @@ ${{{{
         UINT dest,
         UINT value);
 
-    __target_intrinsic(glsl, "$0._data[$1/4] = $2")
+    __target_intrinsic(hlsl)
+    [ForceInline]
     void Store(
         uint address,
-        uint value);
+        uint value)
+    {
+        __byteAddressBufferStore(this, address, value);
+    }
 
-    __target_intrinsic(glsl, "$0._data[$1/4] = $2.x, $0._data[$1/4+1] = $2.y")
-    void Store2(
-        uint address,
-        uint2 value);
+    __target_intrinsic(hlsl)
+    [ForceInline]
+    void Store2(uint address, uint2 value)
+    {
+        __byteAddressBufferStore(this, address, value);
+    }
 
-    __target_intrinsic(glsl, "$0._data[$1/4] = $2.x, $0._data[$1/4+1] = $2.y, $0._data[$1/4+2] = $2.z")
+    __target_intrinsic(hlsl)
+    [ForceInline]
     void Store3(
         uint address,
-        uint3 value);
+        uint3 value)
+    {
+        __byteAddressBufferStore(this, address, value);
+    }
 
-    __target_intrinsic(glsl, "$0._data[$1/4] = $2.x, $0._data[$1/4+1] = $2.y, $0._data[$1/4+2] = $2.z, $0._data[$1/4+3] = $2.w")
+    __target_intrinsic(hlsl)
+    [ForceInline]
     void Store4(
         uint address,
-        uint4 value);
+        uint4 value)
+    {
+        __byteAddressBufferStore(this, address, value);
+    }
 
     void Store<T>(int offset, T value)
     {
@@ -738,10 +814,17 @@ struct $(item.name)
 {
     uint DecrementCounter();
 
-    __target_intrinsic(glsl, "$1 = $0._data.length(); $2 = 0")
+    [__readNone]
+    [__unsafeForceInlineEarly]
+    __target_intrinsic(hlsl)
     void GetDimensions(
         out uint numStructs,
-        out uint stride);
+        out uint stride)
+    {
+        let rs = __structuredBufferGetDimensions(this);
+        numStructs = rs.x;
+        stride = rs.y;
+    }
 
     uint IncrementCounter();
 
@@ -1850,6 +1933,12 @@ __generic<T : __BuiltinFloatingPointType> vector<T,4> dst(vector<T,4> x, vector<
 __intrinsic_op($(kIROp_GetEquivalentStructuredBuffer))
 RWStructuredBuffer<T> __getEquivalentStructuredBuffer<T>(RWByteAddressBuffer b);
 
+__intrinsic_op($(kIROp_GetEquivalentStructuredBuffer))
+StructuredBuffer<T> __getEquivalentStructuredBuffer<T>(ByteAddressBuffer b);
+
+__intrinsic_op($(kIROp_GetEquivalentStructuredBuffer))
+RasterizerOrderedStructuredBuffer<T> __getEquivalentStructuredBuffer<T>(RasterizerOrderedByteAddressBuffer b);
+
 // Error message
 
 // void errorf( string format, ... );
diff --git a/source/slang/slang-emit-cpp.cpp b/source/slang/slang-emit-cpp.cpp
index dee8e7197..0b01bdf4e 100644
--- a/source/slang/slang-emit-cpp.cpp
+++ b/source/slang/slang-emit-cpp.cpp
@@ -1192,6 +1192,40 @@ const UnownedStringSlice* CPPSourceEmitter::getVectorElementNames(IRVectorType*
     return getVectorElementNames(basicType->getBaseType(), elemCount);
 }
 
+bool CPPSourceEmitter::tryEmitInstStmtImpl(IRInst* inst)
+{
+    switch (inst->getOp())
+    {
+    case kIROp_StructuredBufferGetDimensions:
+    {
+        auto count = _generateUniqueName(UnownedStringSlice("_elementCount"));
+        auto stride = _generateUniqueName(UnownedStringSlice("_stride"));
+
+        m_writer->emit("uint ");
+        m_writer->emit(count);
+        m_writer->emit(";\n");
+        m_writer->emit("uint ");
+        m_writer->emit(stride);
+        m_writer->emit(";\n");
+        emitOperand(inst->getOperand(0), leftSide(getInfo(EmitOp::General), getInfo(EmitOp::Postfix)));
+        m_writer->emit(".GetDimensions(&");
+        m_writer->emit(count);
+        m_writer->emit(", &");
+        m_writer->emit(stride);
+        m_writer->emit(");\n");
+        emitInstResultDecl(inst);
+        m_writer->emit("uint2(");
+        m_writer->emit(count);
+        m_writer->emit(", ");
+        m_writer->emit(stride);
+        m_writer->emit(");\n");
+        return true;
+    }
+    default:
+        return false;
+    }
+}
+
 bool CPPSourceEmitter::tryEmitInstExprImpl(IRInst* inst, const EmitOpInfo& inOuterPrec)
 {
     switch (inst->getOp())
diff --git a/source/slang/slang-emit-cpp.h b/source/slang/slang-emit-cpp.h
index 89978e68a..73a0f864e 100644
--- a/source/slang/slang-emit-cpp.h
+++ b/source/slang/slang-emit-cpp.h
@@ -53,6 +53,8 @@ protected:
     virtual void _emitType(IRType* type, DeclaratorInfo* declarator) SLANG_OVERRIDE;
     virtual void emitVectorTypeNameImpl(IRType* elementType, IRIntegerValue elementCount) SLANG_OVERRIDE;
     virtual bool tryEmitInstExprImpl(IRInst* inst, const EmitOpInfo& inOuterPrec) SLANG_OVERRIDE;
+    virtual bool tryEmitInstStmtImpl(IRInst* inst) SLANG_OVERRIDE;
+
     virtual void emitPreModuleImpl() SLANG_OVERRIDE;
     virtual void emitSimpleValueImpl(IRInst* value) SLANG_OVERRIDE;
     virtual void emitSimpleFuncParamImpl(IRParam* param) SLANG_OVERRIDE;
diff --git a/source/slang/slang-emit-cuda.cpp b/source/slang/slang-emit-cuda.cpp
index 345aa3168..3e974f60e 100644
--- a/source/slang/slang-emit-cuda.cpp
+++ b/source/slang/slang-emit-cuda.cpp
@@ -449,6 +449,40 @@ void CUDASourceEmitter::emitIntrinsicCallExprImpl(IRCall* inst, IRTargetIntrinsi
     Super::emitIntrinsicCallExprImpl(inst, targetIntrinsic, inOuterPrec);
 }
 
+bool CUDASourceEmitter::tryEmitInstStmtImpl(IRInst* inst)
+{
+    switch (inst->getOp())
+    {
+    case kIROp_StructuredBufferGetDimensions:
+    {
+        auto count = _generateUniqueName(UnownedStringSlice("_elementCount"));
+        auto stride = _generateUniqueName(UnownedStringSlice("_stride"));
+
+        m_writer->emit("uint ");
+        m_writer->emit(count);
+        m_writer->emit(";\n");
+        m_writer->emit("uint ");
+        m_writer->emit(stride);
+        m_writer->emit(";\n");
+        emitOperand(inst->getOperand(0), leftSide(getInfo(EmitOp::General), getInfo(EmitOp::Postfix)));
+        m_writer->emit(".GetDimensions(&");
+        m_writer->emit(count);
+        m_writer->emit(", &");
+        m_writer->emit(stride);
+        m_writer->emit(");\n");
+        emitInstResultDecl(inst);
+        m_writer->emit("make_uint2(");
+        m_writer->emit(count);
+        m_writer->emit(", ");
+        m_writer->emit(stride);
+        m_writer->emit(");\n");
+        return true;
+    }
+    default:
+        return false;
+    }
+}
+
 bool CUDASourceEmitter::tryEmitInstExprImpl(IRInst* inst, const EmitOpInfo& inOuterPrec)
 {
     switch(inst->getOp())
diff --git a/source/slang/slang-emit-cuda.h b/source/slang/slang-emit-cuda.h
index b067e0010..a08afd862 100644
--- a/source/slang/slang-emit-cuda.h
+++ b/source/slang/slang-emit-cuda.h
@@ -92,6 +92,7 @@ protected:
 
     virtual bool tryEmitGlobalParamImpl(IRGlobalParam* varDecl, IRType* varType) SLANG_OVERRIDE;
     virtual bool tryEmitInstExprImpl(IRInst* inst, const EmitOpInfo& inOuterPrec) SLANG_OVERRIDE;
+    virtual bool tryEmitInstStmtImpl(IRInst* inst) SLANG_OVERRIDE;
     virtual void emitIntrinsicCallExprImpl(IRCall* inst, IRTargetIntrinsicDecoration* targetIntrinsic, EmitOpInfo const& inOuterPrec) SLANG_OVERRIDE;
     
     virtual void emitModuleImpl(IRModule* module, DiagnosticSink* sink) SLANG_OVERRIDE;
diff --git a/source/slang/slang-emit-glsl.cpp b/source/slang/slang-emit-glsl.cpp
index e1f74f70d..e40a5ceca 100644
--- a/source/slang/slang-emit-glsl.cpp
+++ b/source/slang/slang-emit-glsl.cpp
@@ -8,6 +8,7 @@
 #include "slang-mangled-lexer.h"
 
 #include "slang-legalize-types.h"
+#include "slang-ir-layout.h"
 
 #include <assert.h>
 
@@ -2036,6 +2037,22 @@ bool GLSLSourceEmitter::tryEmitInstStmtImpl(IRInst* inst)
             m_writer->emit(", -1);\n");
             return true;
         }
+    case kIROp_StructuredBufferGetDimensions:
+        {
+            emitInstResultDecl(inst);
+            m_writer->emit("uvec2(");
+            emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+            m_writer->emit("._data.length(), ");
+            auto elementType = as<IRHLSLStructuredBufferTypeBase>(inst->getOperand(0)->getDataType())->getElementType();
+            IRIntegerValue stride = 0;
+            if (auto sizeDecor = elementType->findDecoration<IRSizeAndAlignmentDecoration>())
+            {
+                stride = align(sizeDecor->getSize(), (int)sizeDecor->getAlignment());
+            }
+            m_writer->emit(stride);
+            m_writer->emit(");\n");
+            return true;
+        }
     default:
         return false;
     }
diff --git a/source/slang/slang-emit-spirv-ops.h b/source/slang/slang-emit-spirv-ops.h
index d91afad0c..253f90c22 100644
--- a/source/slang/slang-emit-spirv-ops.h
+++ b/source/slang/slang-emit-spirv-ops.h
@@ -1044,6 +1044,20 @@ SpvInst* emitOpCompositeConstruct(
     return emitInst(parent, inst, SpvOpCompositeConstruct, idResultType, kResultID, constituents);
 }
 
+// https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#OpCompositeConstruct
+template<typename T, typename T1, typename T2>
+SpvInst* emitOpCompositeConstruct(
+    SpvInstParent* parent,
+    IRInst* inst,
+    const T& idResultType,
+    const T1& constituent1,
+    const T2& constituent2
+)
+{
+    static_assert(isSingular<T>);
+    return emitInst(parent, inst, SpvOpCompositeConstruct, idResultType, kResultID, constituent1, constituent2);
+}
+
 // https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#OpCompositeExtract
 template<typename T1, typename T2, Index N>
 SpvInst* emitOpCompositeExtract(
diff --git a/source/slang/slang-emit-spirv.cpp b/source/slang/slang-emit-spirv.cpp
index 9a7f5ad31..9d8c4d89a 100644
--- a/source/slang/slang-emit-spirv.cpp
+++ b/source/slang/slang-emit-spirv.cpp
@@ -1822,6 +1822,8 @@ struct SPIRVEmitContext
             return emitStore(parent, as<IRStore>(inst));
         case kIROp_RWStructuredBufferGetElementPtr:
             return emitStructuredBufferGetElementPtr(parent, inst);
+        case kIROp_StructuredBufferGetDimensions:
+            return emitStructuredBufferGetDimensions(parent, inst);
         case kIROp_swizzle:
             return emitSwizzle(parent, as<IRSwizzle>(inst));
         case kIROp_IntCast:
@@ -2915,6 +2917,21 @@ struct SPIRVEmitContext
         return addr;
     }
 
+    SpvInst* emitStructuredBufferGetDimensions(SpvInstParent* parent, IRInst* inst)
+    {
+        IRBuilder builder(inst);
+        auto arrayLength = emitInst(parent, nullptr, SpvOpArrayLength, builder.getUIntType(), kResultID, inst->getOperand(0), SpvLiteralInteger::from32(0));
+        auto elementType = as<IRPtrType>(inst->getOperand(0)->getDataType())->getValueType();
+        IRIntegerValue stride = 0;
+        if (auto sizeDecor = elementType->findDecoration<IRSizeAndAlignmentDecoration>())
+        {
+            stride = align(sizeDecor->getSize(), (int)sizeDecor->getAlignment());
+        }
+        auto strideOperand = emitIntConstant(stride, builder.getUIntType());
+        auto result = emitOpCompositeConstruct(parent, inst, inst->getDataType(), arrayLength, strideOperand);
+        return result;
+    }
+
     SpvInst* emitSwizzle(SpvInstParent* parent, IRSwizzle* inst)
     {
         if (inst->getElementCount() == 1)
diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index 03d62b540..8ee641acc 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -650,6 +650,8 @@ Result linkAndOptimizeIR(
             break;
 
         case CodeGenTarget::GLSL:
+        case CodeGenTarget::SPIRV:
+        case CodeGenTarget::SPIRVAssembly:
             // For GLSL targets, we want to translate the vector load/store
             // operations into scalar ops. This is in part as a simplification,
             // but it also ensures that our generated code respects the lax
@@ -793,6 +795,7 @@ Result linkAndOptimizeIR(
     switch (target)
     {
     case CodeGenTarget::GLSL:
+    case CodeGenTarget::SPIRV:
         {
             legalizeImageSubscriptForGLSL(irModule);
             legalizeConstantBufferLoadForGLSL(irModule);
diff --git a/source/slang/slang-ir-autodiff.cpp b/source/slang/slang-ir-autodiff.cpp
index 645662caa..4c5608341 100644
--- a/source/slang/slang-ir-autodiff.cpp
+++ b/source/slang/slang-ir-autodiff.cpp
@@ -709,7 +709,6 @@ IRFunc *DifferentiableTypeConformanceContext::getOrCreateExistentialDAddMethod()
     builder.addNameHintDecoration(existentialDAddFunc, UnownedStringSlice("__existential_dadd"));
     builder.addBackwardDifferentiableDecoration(existentialDAddFunc);
     
-    this->existentialDAddFunc = existentialDAddFunc;
     return existentialDAddFunc;
 }
 
diff --git a/source/slang/slang-ir-lower-buffer-element-type.cpp b/source/slang/slang-ir-lower-buffer-element-type.cpp
index 0472e44df..909ffea83 100644
--- a/source/slang/slang-ir-lower-buffer-element-type.cpp
+++ b/source/slang/slang-ir-lower-buffer-element-type.cpp
@@ -500,6 +500,7 @@ namespace Slang
                             case kIROp_StructuredBufferLoadStatus:
                             case kIROp_RWStructuredBufferLoad:
                             case kIROp_RWStructuredBufferLoadStatus:
+                            case kIROp_StructuredBufferConsume:
                                 {
                                     IRCloneEnv cloneEnv = {};
                                     builder.setInsertBefore(user);
@@ -512,6 +513,7 @@ namespace Slang
                                 }
                             case kIROp_Store:
                             case kIROp_RWStructuredBufferStore:
+                            case kIROp_StructuredBufferAppend:
                                 {
                                     // Use must be the dest operand of the store inst.
                                     if (use != user->getOperands() + 0)
@@ -524,6 +526,8 @@ namespace Slang
                                         store->val.set(packedVal);
                                     else if (auto sbStore = as<IRRWStructuredBufferStore>(user))
                                         sbStore->setOperand(2, packedVal);
+                                    else if (auto sbAppend = as<IRStructuredBufferAppend>(user))
+                                        sbAppend->setOperand(1, packedVal);
                                     else
                                         SLANG_UNREACHABLE("unhandled store type");
                                     break;
@@ -565,6 +569,8 @@ namespace Slang
                             case kIROp_RWStructuredBufferGetElementPtr:
                                 ptrValsWorkList.add(user);
                                 break;
+                            case kIROp_StructuredBufferGetDimensions:
+                                break;
                             default:
                                 SLANG_UNREACHABLE("unhandled inst of a buffer/pointer value that needs storage lowering.");
                                 break;