Defer immutable buffer loads when emitting spirv. (#7579)

* Defer immutable buffer loads when emitting spirv.

* Fix.

* Fix.

* Fix.

* Fix tests.

* Fix test.

1 files changed, 77 insertions, 26 deletions

diff --git a/source/slang/slang-ir-defer-buffer-load.cpp b/source/slang/slang-ir-defer-buffer-load.cpp
index bd3b78f9e..e71892fe5 100644
--- a/source/slang/slang-ir-defer-buffer-load.cpp
+++ b/source/slang/slang-ir-defer-buffer-load.cpp
@@ -18,7 +18,6 @@ struct DeferBufferLoadContext
     Dictionary<IRInst*, IRInst*> mapPtrToValue;
 
     IRFunc* currentFunc = nullptr;
-    IRDominatorTree* dominatorTree = nullptr;
 
     // Ensure that for an original SSA value, we have formed a pointer that can be used to load the
     // value.
@@ -57,6 +56,9 @@ struct DeferBufferLoadContext
                 result = b.emitFieldAddress(ptr, valueInst->getOperand(1));
                 break;
             }
+        case kIROp_Load:
+            result = valueInst->getOperand(0);
+            break;
         }
         if (result)
         {
@@ -65,7 +67,39 @@ struct DeferBufferLoadContext
         return result;
     }
 
-    static bool isStructuredBufferLoad(IRInst* inst)
+    static bool isImmutableLocation(IRInst* loc)
+    {
+        switch (loc->getOp())
+        {
+        case kIROp_GetStructuredBufferPtr:
+        case kIROp_ImageSubscript:
+            return isImmutableLocation(loc->getOperand(0));
+        default:
+            break;
+        }
+
+        auto type = loc->getDataType();
+        if (!type)
+            return false;
+
+        switch (type->getOp())
+        {
+        case kIROp_HLSLStructuredBufferType:
+        case kIROp_HLSLByteAddressBufferType:
+        case kIROp_ConstantBufferType:
+        case kIROp_ParameterBlockType:
+            return true;
+        default:
+            break;
+        }
+
+        if (auto textureType = as<IRTextureType>(type))
+            return textureType->getAccess() == SLANG_RESOURCE_ACCESS_READ;
+
+        return false;
+    }
+
+    static bool isImmutableBufferLoad(IRInst* inst)
     {
         // Note: we cannot defer loads from RWStructuredBuffer because there can be other
         // instructions that modify the buffer.
@@ -74,6 +108,11 @@ struct DeferBufferLoadContext
         case kIROp_StructuredBufferLoad:
         case kIROp_StructuredBufferLoadStatus:
             return true;
+        case kIROp_Load:
+            {
+                auto rootAddr = getRootAddr(inst->getOperand(0));
+                return isImmutableLocation(rootAddr);
+            }
         default:
             return false;
         }
@@ -88,33 +127,49 @@ struct DeferBufferLoadContext
         IRInst* result = nullptr;
         if (mapPtrToValue.tryGetValue(ptr, result))
             return result;
-        builder.setInsertAfter(ptr);
-        result = builder.emitLoad(ptr);
-        mapPtrToValue[ptr] = result;
+        IRAlignedAttr* align = nullptr;
+        if (auto load = as<IRLoad>(loadInst))
+            align = load->findAttr<IRAlignedAttr>();
+        if (!as<IRModuleInst>(ptr->getParent()))
+        {
+            builder.setInsertAfter(ptr);
+            IRType* valueType = tryGetPointedToType(&builder, ptr->getFullType());
+            result = builder.emitLoad(valueType, ptr, align);
+            mapPtrToValue[ptr] = result;
+        }
+        else
+        {
+            builder.setInsertBefore(loadInst);
+            IRType* valueType = tryGetPointedToType(&builder, ptr->getFullType());
+            result = builder.emitLoad(valueType, ptr, align);
+            // Since we are inserting the load in a local scope, we can't register
+            // the mapping to the pointer, since the global pointer needs to be
+            // loaded once per function.
+        }
         return result;
     }
 
     static bool isSimpleType(IRInst* type)
     {
-        if (as<IRBasicType>(type))
-            return true;
-        if (as<IRVectorType>(type))
-            return true;
-        if (as<IRMatrixType>(type))
-            return true;
-        return false;
+        if (auto modType = as<IRRateQualifiedType>(type))
+            type = modType->getValueType();
+        if (as<IRStructType>(type))
+            return false;
+        if (as<IRTupleType>(type))
+            return false;
+        if (as<IRArrayTypeBase>(type))
+            return false;
+        return true;
     }
 
     void deferBufferLoadInst(IRBuilder& builder, List<IRInst*>& workList, IRInst* loadInst)
     {
         // Don't defer the load anymore if the type is simple.
-        if (isSimpleType(loadInst->getDataType()))
+        if (isSimpleType(loadInst->getDataType()) || loadInst->findAttr<IRAlignedAttr>())
         {
-            if (!isStructuredBufferLoad(loadInst))
-            {
-                auto materializedVal = materializePointer(builder, loadInst);
-                loadInst->replaceUsesWith(materializedVal);
-            }
+            auto materializedVal = materializePointer(builder, loadInst);
+            loadInst->transferDecorationsTo(materializedVal);
+            loadInst->replaceUsesWith(materializedVal);
             return;
         }
 
@@ -141,18 +196,15 @@ struct DeferBufferLoadContext
                     }
                     break;
                 default:
-                    if (!isStructuredBufferLoad(loadInst))
-                    {
-                        needMaterialize = true;
-                        return;
-                    }
-                    break;
+                    needMaterialize = true;
+                    return;
                 }
             });
 
         if (needMaterialize)
         {
             auto val = materializePointer(builder, loadInst);
+            loadInst->transferDecorationsTo(val);
             loadInst->replaceUsesWith(val);
             loadInst->removeAndDeallocate();
         }
@@ -170,7 +222,6 @@ struct DeferBufferLoadContext
         removeRedundancyInFunc(func, false);
 
         currentFunc = func;
-        dominatorTree = func->getModule()->findOrCreateDominatorTree(func);
 
         List<IRInst*> workList;
 
@@ -178,7 +229,7 @@ struct DeferBufferLoadContext
         {
             for (auto inst : block->getChildren())
             {
-                if (isStructuredBufferLoad(inst))
+                if (isImmutableBufferLoad(inst))
                 {
                     workList.add(inst);
                 }