From cb835b9596f78771dfd584aaa72aee7bc360b3d4 Mon Sep 17 00:00:00 2001
From: Sai Praveen Bangaru <31557731+saipraveenb25@users.noreply.github.com>
Date: Tue, 14 Jan 2025 17:48:17 -0500
Subject: Fix issue with specialization using arithmetic expressions (#6084)

---
 source/slang/slang-ir-specialize.cpp | 45 ++++++++++++++++++++++++++++++++++++
 source/slang/slang.cpp               |  3 ++-
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/source/slang/slang-ir-specialize.cpp b/source/slang/slang-ir-specialize.cpp
index 40cd40758..a9b0d4412 100644
--- a/source/slang/slang-ir-specialize.cpp
+++ b/source/slang/slang-ir-specialize.cpp
@@ -71,6 +71,42 @@ struct SpecializationContext
         module->getContainerPool().free(&cleanInsts);
     }
 
+    bool isUnsimplifiedArithmeticInst(IRInst* inst)
+    {
+        switch (inst->getOp())
+        {
+        case kIROp_Add:
+        case kIROp_Sub:
+        case kIROp_Mul:
+        case kIROp_Div:
+        case kIROp_Neg:
+        case kIROp_Not:
+        case kIROp_Eql:
+        case kIROp_Neq:
+        case kIROp_Leq:
+        case kIROp_Geq:
+        case kIROp_Less:
+        case kIROp_IRem:
+        case kIROp_FRem:
+        case kIROp_Greater:
+        case kIROp_Lsh:
+        case kIROp_Rsh:
+        case kIROp_BitAnd:
+        case kIROp_BitOr:
+        case kIROp_BitXor:
+        case kIROp_BitNot:
+        case kIROp_BitCast:
+        case kIROp_CastIntToFloat:
+        case kIROp_CastFloatToInt:
+        case kIROp_IntCast:
+        case kIROp_FloatCast:
+        case kIROp_Select:
+            return true;
+        default:
+            return false;
+        }
+    }
+
     // An instruction is then fully specialized if and only
     // if it is in our set.
     //
@@ -133,6 +169,14 @@ struct SpecializationContext
             return areAllOperandsFullySpecialized(inst);
         }
 
+        if (isUnsimplifiedArithmeticInst(inst))
+        {
+            // For arithmetic insts, we want to wait for simplification before specialization,
+            // since different insts can simplify to the same value.
+            //
+            return false;
+        }
+
         // The default case is that a global value is always specialized.
         if (inst->getParent() == module->getModuleInst())
         {
@@ -1092,6 +1136,7 @@ struct SpecializationContext
             {
                 this->changed = true;
                 eliminateDeadCode(module->getModuleInst());
+                applySparseConditionalConstantPropagationForGlobalScope(this->module, this->sink);
             }
 
             // Once the work list has gone dry, we should have the invariant
diff --git a/source/slang/slang.cpp b/source/slang/slang.cpp
index 5ec199658..efc1c6fd1 100644
--- a/source/slang/slang.cpp
+++ b/source/slang/slang.cpp
@@ -1493,7 +1493,8 @@ DeclRef<Decl> Linkage::specializeWithArgTypes(
     DiagnosticSink* sink)
 {
     SemanticsVisitor visitor(getSemanticsForReflection());
-    visitor = visitor.withSink(sink);
+    SemanticsVisitor::ExprLocalScope scope;
+    visitor = visitor.withSink(sink).withExprLocalScope(&scope);
 
     SLANG_AST_BUILDER_RAII(getASTBuilder());
 
-- 
cgit v1.2.3