From af8ce68e9fd7b6255b6e4e9e9524a285497116dc Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Thu, 21 Sep 2023 10:39:19 -0700
Subject: Misc SPIRV fixes. (#3220)

Co-authored-by: Yong He <yhe@nvidia.com>
---
 source/slang/slang-ir-lower-buffer-element-type.cpp |  2 +-
 source/slang/slang-ir-spirv-legalize.cpp            | 19 ++++++++++++++++++-
 source/slang/slang-stdlib-textures.cpp              | 19 +++++++++++++++----
 3 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/source/slang/slang-ir-lower-buffer-element-type.cpp b/source/slang/slang-ir-lower-buffer-element-type.cpp
index 6a4d639c9..1c02af82a 100644
--- a/source/slang/slang-ir-lower-buffer-element-type.cpp
+++ b/source/slang/slang-ir-lower-buffer-element-type.cpp
@@ -438,8 +438,8 @@ namespace Slang
                             auto result = builder.emitCast(info.loweredType, param);
                             builder.emitReturn(result);
                         }
+                        return info;
                     }
-                    break;
                 default:
                     break;
                 }
diff --git a/source/slang/slang-ir-spirv-legalize.cpp b/source/slang/slang-ir-spirv-legalize.cpp
index 1590f94c8..309fc5e38 100644
--- a/source/slang/slang-ir-spirv-legalize.cpp
+++ b/source/slang/slang-ir-spirv-legalize.cpp
@@ -955,6 +955,21 @@ struct SPIRVLegalizationContext : public SourceEmitterBase
         inst->insertAtEnd(m_module->getModuleInst());
     }
 
+    static bool isAsmInst(IRInst* inst)
+    {
+        return (as<IRSPIRVAsmInst>(inst) || as<IRSPIRVAsmOperand>(inst));
+    }
+
+    void processSPIRVAsm(IRSPIRVAsm* inst)
+    {
+        // Move anything that is not an spirv instruction to the outer parent.
+        for (auto child : inst->getModifiableChildren())
+        {
+            if (!isAsmInst(child))
+                child->insertBefore(inst);
+        }
+    }
+
     void processModule()
     {
         // Process global params before anything else, so we don't generate inefficient
@@ -1034,7 +1049,9 @@ struct SPIRVLegalizationContext : public SourceEmitterBase
             case kIROp_MakeOptionalNone:
                 processConstructor(inst);
                 break;
-
+            case kIROp_SPIRVAsm:
+                processSPIRVAsm(as<IRSPIRVAsm>(inst));
+                break;
             default:
                 for (auto child = inst->getLastChild(); child; child = child->getPrevInst())
                 {
diff --git a/source/slang/slang-stdlib-textures.cpp b/source/slang/slang-stdlib-textures.cpp
index bec3a7155..76ca51380 100644
--- a/source/slang/slang-stdlib-textures.cpp
+++ b/source/slang/slang-stdlib-textures.cpp
@@ -543,12 +543,23 @@ void TextureTypeInfo::writeQueryFunctions()
     }
 }
 
-static String spirvReadIntrinsic()
+static String spirvReadIntrinsic(SlangResourceAccess access)
 {
     StringBuilder spirvBuilder;
     const char* i = "                    ";
-    spirvBuilder << i << "%sampled : __sampledType(T) = OpImageRead $this $location;\n";
-    spirvBuilder << i << "__truncate $$T result __sampledType(T) %sampled;";
+    switch (access)
+    {
+    case SLANG_RESOURCE_ACCESS_NONE:
+    case SLANG_RESOURCE_ACCESS_READ:
+        spirvBuilder << i << "%sampled : __sampledType(T) = OpImageFetch $this $location;\n";
+        spirvBuilder << i << "__truncate $$T result __sampledType(T) %sampled;";
+        break;
+
+    default:
+        spirvBuilder << i << "%sampled : __sampledType(T) = OpImageRead $this $location;\n";
+        spirvBuilder << i << "__truncate $$T result __sampledType(T) %sampled;";
+        break;
+    }
     return spirvBuilder;
 }
 
@@ -663,7 +674,7 @@ void TextureTypeInfo::writeSubscriptFunctions()
         "operator[]",
         "get",
         glslBuilder,
-        spirvReadIntrinsic(),
+        spirvReadIntrinsic(access),
         cudaBuilder
     );
 
-- 
cgit v1.2.3