Partially implement vk_buffer_ref proposal. (#3814)

6 files changed, 147 insertions, 0 deletions

diff --git a/source/slang/hlsl.meta.slang b/source/slang/hlsl.meta.slang
index b660437d9..d2dfe6267 100644
--- a/source/slang/hlsl.meta.slang
+++ b/source/slang/hlsl.meta.slang
@@ -13912,6 +13912,41 @@ extension __TextureImpl<float, Shape, 0, 0, 0, $(kStdlibResourceAccessReadWrite)
 
 // Buffer Pointer
 
+namespace vk
+{
+    // Partial implementation of the vk::buffer_ref proposal:
+    // https://github.com/microsoft/hlsl-specs/blob/main/proposals/0010-vk-buffer-ref.md
+    struct BufferPointer<T, let Alignment : int = 0>
+    {
+        T *_ptr;
+        [ForceInline] __init(T *ptr) { _ptr = ptr; }
+        [ForceInline] __init(uint64_t val) { _ptr = (T *)val; }
+        [ForceInline] Ref<T> Get() { return *_ptr; }
+        [ForceInline] T *getPtr() { return _ptr;}
+    }
+    [ForceInline]
+    BufferPointer<U, alignment> static_pointer_cast<U, let alignment : int = 0, T, let a : int>(BufferPointer<T, a> src)
+    {
+        return BufferPointer<U, alignment>((U*)(src.getPtr()));
+    }
+    [ForceInline]
+    BufferPointer<U, alignment> reinterpret_pointer_cast<U, let alignment : int = 0, T, let a : int>(BufferPointer<T, a> src)
+    {
+        return BufferPointer<U, alignment>((U *)(src.getPtr()));
+    }
+}
+
+attribute_syntax[vk_aliased_pointer] : VkAliasedPointerAttribute;
+attribute_syntax[vk_restrict_pointer] : VkRestrictPointerAttribute;
+
+extension uint64_t
+{
+    __init<T, let alignment : int>(vk::BufferPointer<T, alignment> ptr)
+    {
+        this = (uint64_t)ptr._ptr;
+    }
+}
+
 __generic<T, let Alignment : int = 16>
 __intrinsic_type($(kIROp_HLSLConstBufferPointerType))
 __glsl_extension(GL_EXT_buffer_reference)
diff --git a/source/slang/slang-ast-modifier.h b/source/slang/slang-ast-modifier.h
index 5ab7eae3c..87057b694 100644
--- a/source/slang/slang-ast-modifier.h
+++ b/source/slang/slang-ast-modifier.h
@@ -741,6 +741,15 @@ class GLSLBindingAttribute : public Attribute
     int32_t set = 0;
 };
 
+class VkAliasedPointerAttribute : public Attribute
+{
+    SLANG_AST_CLASS(VkAliasedPointerAttribute)
+};
+
+class VkRestrictPointerAttribute : public Attribute
+{
+    SLANG_AST_CLASS(VkRestrictPointerAttribute)
+};
 
 class GLSLOffsetLayoutAttribute : public Attribute 
 {
diff --git a/source/slang/slang-ir-spirv-legalize.cpp b/source/slang/slang-ir-spirv-legalize.cpp
index 2b7e86f47..ca3729abc 100644
--- a/source/slang/slang-ir-spirv-legalize.cpp
+++ b/source/slang/slang-ir-spirv-legalize.cpp
@@ -1669,6 +1669,15 @@ struct SPIRVLegalizationContext : public SourceEmitterBase
     void translatePtrResultType(IRInst* inst)
     {
         auto ptrType = as<IRPtrType>(inst->getDataType());
+        if (!ptrType)
+        {
+            if (auto refType = as<IRRefType>(inst->getDataType()))
+            {
+                // Functions that return ref type should be treated as returning a pointer.
+                IRBuilder builder(inst);
+                ptrType = builder.getPtrType(refType->getValueType());
+            }
+        }
         auto newPtrType = translateToStorageBufferPointer(ptrType);
         if (newPtrType == ptrType)
             return;
diff --git a/source/slang/slang-preprocessor.cpp b/source/slang/slang-preprocessor.cpp
index b0986f64c..702259381 100644
--- a/source/slang/slang-preprocessor.cpp
+++ b/source/slang/slang-preprocessor.cpp
@@ -2570,7 +2570,40 @@ static PreprocessorExpressionValue ParseAndEvaluateUnaryExpression(PreprocessorD
 
                 return LookupMacro(context, name) != NULL;
             }
+            else if (token.getContent() == "__has_feature")
+            {
+                // handle `defined(someName)`
+
+                // Possibly parse a `(`
+                Token leftParen;
+                if (PeekRawTokenType(context) == TokenType::LParent)
+                {
+                    leftParen = AdvanceRawToken(context);
+                }
 
+                // Expect an identifier
+                Token nameToken;
+                if (!ExpectRaw(context, TokenType::Identifier, Diagnostics::expectedTokenInDefinedExpression, &nameToken))
+                {
+                    return 0;
+                }
+                
+                // If we saw an opening `(`, then expect one to close
+                if (leftParen.type != TokenType::Unknown)
+                {
+                    if (!ExpectRaw(context, TokenType::RParent, Diagnostics::expectedTokenInDefinedExpression))
+                    {
+                        GetSink(context)->diagnose(leftParen.loc, Diagnostics::seeOpeningToken, leftParen);
+                        return 0;
+                    }
+                }
+
+                if (nameToken.getContent() == "hlsl_vk_buffer_pointer")
+                {
+                    return 1;
+                }
+                return 0;
+            }
             // An identifier here means it was not defined as a macro (or
             // it is defined, but as a function-like macro. These should
             // just evaluate to zero (possibly with a warning)
diff --git a/tests/spirv/vk-buffer-pointer-1.slang b/tests/spirv/vk-buffer-pointer-1.slang
new file mode 100644
index 000000000..7d4e5dc1c
--- /dev/null
+++ b/tests/spirv/vk-buffer-pointer-1.slang
@@ -0,0 +1,34 @@
+//TEST:SIMPLE(filecheck=CHECK): -target spirv -emit-spirv-directly -O0
+
+// Note: spirv-opt will crash on the generated spirv, so use -O0 to disable it.
+
+// CHECK: OpEntryPoint
+// CHECK: OpTypeForwardPointer
+typedef vk::BufferPointer<block_s> block_p;
+
+struct block_s
+{
+      float4 x;
+      block_p next;
+};
+
+struct TestPushConstant_t
+{
+      block_p root;
+};
+
+[[vk::push_constant]] TestPushConstant_t g_PushConstants;
+
+[shader("fragment")]
+float4 MainPs(void) : SV_Target0
+{
+#if __has_feature(hlsl_vk_buffer_pointer)
+      block_p g_p = g_PushConstants.root;
+      g_p = g_p.Get().next;
+      if ((uint64_t)g_p == 0) // Null pointer test
+          return float4(0.0, 0.0, 0.0, 0.0);
+      return g_p.Get().x;
+#else
+    return float4(0.0);
+#endif
+}
\ No newline at end of file
diff --git a/tests/spirv/vk-buffer-pointer.slang b/tests/spirv/vk-buffer-pointer.slang
new file mode 100644
index 000000000..9fcc4deb2
--- /dev/null
+++ b/tests/spirv/vk-buffer-pointer.slang
@@ -0,0 +1,27 @@
+//TEST:SIMPLE(filecheck=CHECK): -target spirv -emit-spirv-directly
+
+// CHECK: OpEntryPoint
+
+struct Globals_s
+{
+      float4 g_vSomeConstantA;
+      float4 g_vTestFloat4;
+      float4 g_vSomeConstantB;
+};
+
+typedef vk::BufferPointer<Globals_s> Globals_p;
+
+struct TestPushConstant_t
+{
+      Globals_p m_nBufferDeviceAddress;
+};
+
+[[vk::push_constant]] TestPushConstant_t g_PushConstants;
+
+[shader("fragment")]
+float4 MainPs(void) : SV_Target0
+{
+      float4 vTest = float4(1.0,0.0,0.0,0.0);
+      g_PushConstants.m_nBufferDeviceAddress.Get().g_vTestFloat4 = vTest;
+      return vTest;
+}
\ No newline at end of file