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

* Defer immutable buffer loads when emitting spirv.

* Fix.

* Fix.

* Fix.

* Fix tests.

* Fix test.

11 files changed, 141 insertions, 35 deletions

diff --git a/source/slang/slang-check-decl.cpp b/source/slang/slang-check-decl.cpp
index 6d8788b4f..9c77303b9 100644
--- a/source/slang/slang-check-decl.cpp
+++ b/source/slang/slang-check-decl.cpp
@@ -14471,6 +14471,15 @@ VarDeclBase* getTrailingUnsizedArrayElement(
     return nullptr;
 }
 
+bool isImmutableBufferType(Type* type)
+{
+    if (as<UniformParameterGroupType>(type))
+        return true;
+    if (auto resourceType = as<ResourceType>(type))
+        return resourceType->getAccess() == SLANG_RESOURCE_ACCESS_READ;
+    return false;
+}
+
 bool isOpaqueHandleType(Type* type)
 {
     while (auto modifiedType = as<ModifiedType>(type))
diff --git a/source/slang/slang-check-expr.cpp b/source/slang/slang-check-expr.cpp
index c3238bd9b..b28b458da 100644
--- a/source/slang/slang-check-expr.cpp
+++ b/source/slang/slang-check-expr.cpp
@@ -538,6 +538,11 @@ Expr* SemanticsVisitor::constructDerefExpr(Expr* base, QualType elementType, Sou
     {
         derefExpr->type.isLeftValue = true;
     }
+    else if (isImmutableBufferType(base->type))
+    {
+        derefExpr->type.isLeftValue = false;
+        derefExpr->type.isWriteOnly = false;
+    }
     else
     {
         derefExpr->type.isLeftValue = base->type.isLeftValue;
diff --git a/source/slang/slang-check-impl.h b/source/slang/slang-check-impl.h
index e73f65d75..7ddec20fb 100644
--- a/source/slang/slang-check-impl.h
+++ b/source/slang/slang-check-impl.h
@@ -3145,6 +3145,8 @@ bool isUnsizedArrayType(Type* type);
 
 bool isInterfaceType(Type* type);
 
+bool isImmutableBufferType(Type* type);
+
 // Check if `type` is nullable. An `Optional<T>` will occupy the same space as `T`, if `T`
 // is nullable.
 bool isNullableType(Type* type);
diff --git a/source/slang/slang-emit-c-like.cpp b/source/slang/slang-emit-c-like.cpp
index 24acea42f..af4345942 100644
--- a/source/slang/slang-emit-c-like.cpp
+++ b/source/slang/slang-emit-c-like.cpp
@@ -2707,7 +2707,7 @@ void CLikeSourceEmitter::defaultEmitInstExpr(IRInst* inst, const EmitOpInfo& inO
     case kIROp_NonUniformResourceIndex:
         emitOperand(
             inst->getOperand(0),
-            getInfo(EmitOp::General)); // Directly emit NonUniformResourceIndex Operand0;
+            outerPrec); // Directly emit NonUniformResourceIndex Operand0;
         break;
 
     case kIROp_GetNativeStr:
diff --git a/source/slang/slang-emit-spirv-ops.h b/source/slang/slang-emit-spirv-ops.h
index 4d33e045e..a5e4d730a 100644
--- a/source/slang/slang-emit-spirv-ops.h
+++ b/source/slang/slang-emit-spirv-ops.h
@@ -672,6 +672,21 @@ SpvInst* emitOpAccessChain(
     return emitInst(parent, inst, SpvOpAccessChain, idResultType, kResultID, base, indexes);
 }
 
+// https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#OpInBoundsAccessChain
+template<typename T1, typename T2, typename Ts>
+SpvInst* emitOpInBoundsAccessChain(
+    SpvInstParent* parent,
+    IRInst* inst,
+    const T1& idResultType,
+    const T2& base,
+    const Ts& indexes)
+{
+    static_assert(isSingular<T1>);
+    static_assert(isSingular<T2>);
+    static_assert(isPlural<Ts>);
+    return emitInst(parent, inst, SpvOpInBoundsAccessChain, idResultType, kResultID, base, indexes);
+}
+
 
 // https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#OpPtrAccessChain
 template<typename T1, typename T2, typename T3>
diff --git a/source/slang/slang-emit-spirv.cpp b/source/slang/slang-emit-spirv.cpp
index 1e6f27e7f..62c667de1 100644
--- a/source/slang/slang-emit-spirv.cpp
+++ b/source/slang/slang-emit-spirv.cpp
@@ -6824,7 +6824,7 @@ struct SPIRVEmitContext : public SourceEmitterBase, public SPIRVEmitSharedContex
             getStructFieldId(baseStructType, as<IRStructKey>(fieldAddress->getField())),
             builder.getIntType());
         SLANG_ASSERT(as<IRPtrTypeBase>(fieldAddress->getFullType()));
-        return emitOpAccessChain(
+        return emitOpInBoundsAccessChain(
             parent,
             fieldAddress,
             fieldAddress->getFullType(),
@@ -6869,7 +6869,6 @@ struct SPIRVEmitContext : public SourceEmitterBase, public SPIRVEmitSharedContex
         // We might replace resultType with a different storage class equivalent
         auto resultType = as<IRPtrTypeBase>(inst->getDataType());
         SLANG_ASSERT(resultType);
-
         if (const auto basePtrType = as<IRPtrTypeBase>(base->getDataType()))
         {
             // If the base pointer has a specific address space and the
diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index 49b27383d..db8a9ba61 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -1227,11 +1227,6 @@ Result linkAndOptimizeIR(
     // Inline calls to any functions marked with [__unsafeInlineEarly] or [ForceInline].
     performForceInlining(irModule);
 
-    // Push `structuredBufferLoad` to the end of access chain to avoid loading unnecessary data.
-    if (isKhronosTarget(targetRequest) || isMetalTarget(targetRequest) ||
-        isWGPUTarget(targetRequest))
-        deferBufferLoad(irModule);
-
     // Specialization can introduce dead code that could trip
     // up downstream passes like type legalization, so we
     // will run a DCE pass to clean up after the specialization.
@@ -1386,6 +1381,11 @@ Result linkAndOptimizeIR(
     specializeResourceUsage(codeGenContext, irModule);
     specializeFuncsForBufferLoadArgs(codeGenContext, irModule);
 
+    // Push `structuredBufferLoad` to the end of access chain to avoid loading unnecessary data.
+    if (isKhronosTarget(targetRequest) || isMetalTarget(targetRequest) ||
+        isWGPUTarget(targetRequest))
+        deferBufferLoad(irModule);
+
     // We also want to specialize calls to functions that
     // takes unsized array parameters if possible.
     // Moreover, for Khronos targets, we also want to specialize calls to functions
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);
                 }
diff --git a/source/slang/slang-ir-insts.h b/source/slang/slang-ir-insts.h
index 8ce4e46cd..afe06f9a1 100644
--- a/source/slang/slang-ir-insts.h
+++ b/source/slang/slang-ir-insts.h
@@ -4310,6 +4310,7 @@ public:
 
     IRInst* emitLoad(IRType* type, IRInst* ptr);
     IRInst* emitLoad(IRType* type, IRInst* ptr, IRInst* align);
+    IRInst* emitLoad(IRType* type, IRInst* ptr, IRAlignedAttr* align);
     IRInst* emitLoad(IRInst* ptr);
 
     IRInst* emitLoadReverseGradient(IRType* type, IRInst* diffValue);
diff --git a/source/slang/slang-ir-specialize-function-call.cpp b/source/slang/slang-ir-specialize-function-call.cpp
index 7a9fc5f6f..c03e644de 100644
--- a/source/slang/slang-ir-specialize-function-call.cpp
+++ b/source/slang/slang-ir-specialize-function-call.cpp
@@ -338,7 +338,9 @@ struct FunctionParameterSpecializationContext
         // correctly check the preconditions.
         //
         auto oldFunc = as<IRFunc>(oldCall->getCallee());
-        SLANG_ASSERT(oldFunc);
+        if (!oldFunc)
+            return;
+
         SLANG_ASSERT(oldFunc->isDefinition());
 
         // Our first information-gathering pass will
@@ -390,6 +392,14 @@ struct FunctionParameterSpecializationContext
         newCall->insertBefore(oldCall);
         oldCall->replaceUsesWith(newCall);
         oldCall->removeAndDeallocate();
+
+        // If old func is no longer used after the specialization,
+        // remove it.
+        if (!oldFunc->hasUses())
+        {
+            if (!shouldInstBeLiveIfParentIsLive(oldFunc, IRDeadCodeEliminationOptions{}))
+                oldFunc->removeAndDeallocate();
+        }
     }
 
     // Before diving into the details on how we gather information
diff --git a/source/slang/slang-ir.cpp b/source/slang/slang-ir.cpp
index dc25b1489..7c687f4f1 100644
--- a/source/slang/slang-ir.cpp
+++ b/source/slang/slang-ir.cpp
@@ -5107,6 +5107,20 @@ IRInst* IRBuilder::emitLoad(IRType* type, IRInst* ptr, IRInst* align)
     return inst;
 }
 
+IRInst* IRBuilder::emitLoad(IRType* type, IRInst* ptr, IRAlignedAttr* align)
+{
+    if (align)
+    {
+        auto inst = createInst<IRLoad>(this, kIROp_Load, type, ptr, align);
+        addInst(inst);
+        return inst;
+    }
+    else
+    {
+        return emitLoad(type, ptr);
+    }
+}
+
 IRInst* IRBuilder::emitLoad(IRInst* ptr)
 {
     // Note: a `load` operation does not consider the rate