Add support for vk::image_format attribute (#2945)

5 files changed, 66 insertions, 7 deletions

diff --git a/source/slang/core.meta.slang b/source/slang/core.meta.slang
index c30fdf403..1677d6e12 100644
--- a/source/slang/core.meta.slang
+++ b/source/slang/core.meta.slang
@@ -3093,6 +3093,9 @@ attribute_syntax [__AttributeUsage(target : _AttributeTargets)] : AttributeUsage
 __attributeTarget(VarDeclBase)
 attribute_syntax [format(format : String)] : FormatAttribute;
 
+__attributeTarget(VarDeclBase)
+attribute_syntax [vk_image_format(format : String)] : FormatAttribute;
+
 __attributeTarget(Decl)
 attribute_syntax [allow(diagnostic: String)] : AllowAttribute;
 
diff --git a/source/slang/slang-ast-support-types.h b/source/slang/slang-ast-support-types.h
index 7e65dd02e..c366a4c4e 100644
--- a/source/slang/slang-ast-support-types.h
+++ b/source/slang/slang-ast-support-types.h
@@ -158,7 +158,9 @@ namespace Slang
 
     const ImageFormatInfo& getImageFormatInfo(ImageFormat format);
 
-    bool findImageFormatByName(char const* name, ImageFormat* outFormat);
+    bool findImageFormatByName(const UnownedStringSlice& name, ImageFormat* outFormat);
+    bool findVkImageFormatByName(const UnownedStringSlice& name, ImageFormat* outFormat);
+
     char const* getGLSLNameForImageFormat(ImageFormat format);
 
     // TODO(tfoley): We should ditch this enumeration
diff --git a/source/slang/slang-check-modifier.cpp b/source/slang/slang-check-modifier.cpp
index 2bc914a65..ce80d0002 100644
--- a/source/slang/slang-check-modifier.cpp
+++ b/source/slang/slang-check-modifier.cpp
@@ -591,9 +591,20 @@ namespace Slang
             }
 
             ImageFormat format = ImageFormat::unknown;
-            if(!findImageFormatByName(formatName.getBuffer(), &format))
+
+            if (attr->keywordName->text.getUnownedSlice() == toSlice("image"))
+            {
+                if(!findImageFormatByName(formatName.getUnownedSlice(), &format))
+                {
+                    getSink()->diagnose(attr->args[0], Diagnostics::unknownImageFormatName, formatName);
+                }
+            }
+            else
             {
-                getSink()->diagnose(attr->args[0], Diagnostics::unknownImageFormatName, formatName);
+                if (!findVkImageFormatByName(formatName.getUnownedSlice(), &format))
+                {
+                    getSink()->diagnose(attr->args[0], Diagnostics::unknownImageFormatName, formatName);
+                }
             }
 
             formatAttr->format = format;
diff --git a/source/slang/slang-image-format-defs.h b/source/slang/slang-image-format-defs.h
index 25dbae16a..97c41270b 100644
--- a/source/slang/slang-image-format-defs.h
+++ b/source/slang/slang-image-format-defs.h
@@ -43,5 +43,7 @@ FORMAT(rg8ui,           (UINT8,   2, sizeof(uint8_t) * 2))
 FORMAT(r32ui,           (UINT32,  1, sizeof(uint32_t)))
 FORMAT(r16ui,           (UINT16,  1, sizeof(uint16_t)))
 FORMAT(r8ui,            (UINT8,   1, sizeof(uint8_t)))
+FORMAT(r64ui,           (UINT64,   1, sizeof(uint64_t)))
+FORMAT(r64i,            (INT64,   1, sizeof(int64_t)))
 
 #undef FORMAT
diff --git a/source/slang/slang-syntax.cpp b/source/slang/slang-syntax.cpp
index 606c77096..53e77e646 100644
--- a/source/slang/slang-syntax.cpp
+++ b/source/slang/slang-syntax.cpp
@@ -1434,10 +1434,8 @@ static const ImageFormatInfo kImageFormatInfos[] =
 #undef SLANG_IMAGE_FORMAT_INFO
 };
 
-bool findImageFormatByName(char const* inName, ImageFormat* outFormat)
+bool findImageFormatByName(const UnownedStringSlice& name, ImageFormat* outFormat)
 {
-    const UnownedStringSlice name(inName);
-
     for (Index i = 0; i < SLANG_COUNT_OF(kImageFormatInfos); ++i)
     {
         const auto& info = kImageFormatInfos[i];
@@ -1450,12 +1448,55 @@ bool findImageFormatByName(char const* inName, ImageFormat* outFormat)
     return false;
 }
 
+// https://github.com/microsoft/DirectXShaderCompiler/blob/main/docs/SPIR-V.rst#id71
+#define SLANG_VK_TO_IMAGE_FORMAT(x) \
+    x(r11g11b10f, r11f_g11f_b10f) \
+    x(rgb10a2, rgb10_a2) \
+    x(rgb10a2ui, rgb10_a2ui) 
+
+struct VkImageFormatInfo
+{
+    UnownedStringSlice name;
+    ImageFormat format;
+};
+static const VkImageFormatInfo kVkImageFormatInfos[] =
+{
+#define SLANG_VK_IMAGE_FORMAT_INFO(name, format) { toSlice(#name), ImageFormat::format },
+    SLANG_VK_TO_IMAGE_FORMAT(SLANG_VK_IMAGE_FORMAT_INFO)
+};
+
+static const auto kSNorm = UnownedStringSlice::fromLiteral("snorm");
+
+bool findVkImageFormatByName(const UnownedStringSlice& name, ImageFormat* outFormat)
+{
+    // Handle names ending in snorm
+    if (name.endsWith(kSNorm))
+    {
+        StringBuilder buf;
+        //  format names end with snormal after a '_', so replace with that
+        buf << name.head(name.getLength() - kSNorm.getLength()) << "_" << kSNorm;
+        return findImageFormatByName(buf.getUnownedSlice(), outFormat);
+    }
+    
+    // Handle the special cases
+    for (const auto& vkInfo : kVkImageFormatInfos)
+    {
+        if (vkInfo.name == name)
+        {
+            *outFormat = vkInfo.format;
+            return true;
+        }
+    }
+
+    // Default to the regular lookup mechanism for everything else
+    return findImageFormatByName(name, outFormat);
+}
+
 char const* getGLSLNameForImageFormat(ImageFormat format)
 {
     return kImageFormatInfos[Index(format)].name.begin();
 }
 
-
 const ImageFormatInfo& getImageFormatInfo(ImageFormat format)
 {
     return kImageFormatInfos[Index(format)];