From 00474731a90add7c0aaffc3a76d713ae2cecdbbe Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Thu, 21 Mar 2024 21:10:48 -0700
Subject: Partially implement vk_buffer_ref proposal. (#3814)

---
 source/slang/hlsl.meta.slang             | 35 ++++++++++++++++++++++++++++++++
 source/slang/slang-ast-modifier.h        |  9 ++++++++
 source/slang/slang-ir-spirv-legalize.cpp |  9 ++++++++
 source/slang/slang-preprocessor.cpp      | 33 ++++++++++++++++++++++++++++++
 4 files changed, 86 insertions(+)

(limited to 'source')

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)
-- 
cgit v1.2.3