Misc SPIRV fixes. (#3220)

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

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
     );