From 2da28c50d9c3699692eccde4b86d0b8d2323e55c Mon Sep 17 00:00:00 2001
From: Jay Kwak <82421531+jkwak-work@users.noreply.github.com>
Date: Fri, 19 Apr 2024 16:43:21 -0700
Subject: Support arithmetics on generic arguments (#3968)

Resovles an issue #3935

Slang had to fold the generic arguments after specialization.
---
 source/slang/slang-ir-insts.h        |  1 +
 source/slang/slang-ir-specialize.cpp | 26 ++++++++++++++++++++++++--
 2 files changed, 25 insertions(+), 2 deletions(-)

(limited to 'source')

diff --git a/source/slang/slang-ir-insts.h b/source/slang/slang-ir-insts.h
index d1615e89c..2d2fd4511 100644
--- a/source/slang/slang-ir-insts.h
+++ b/source/slang/slang-ir-insts.h
@@ -1234,6 +1234,7 @@ struct IRSpecialize : IRInst
     // after the generic value come the arguments
     UInt getArgCount() { return getOperandCount() - 1; }
     IRInst* getArg(UInt index) { return getOperand(index + 1); }
+    IRUse* getArgOperand(Index i) { return getOperands() + 1 + i; }
 
     IR_LEAF_ISA(Specialize)
 };
diff --git a/source/slang/slang-ir-specialize.cpp b/source/slang/slang-ir-specialize.cpp
index a4985eee4..72e39590b 100644
--- a/source/slang/slang-ir-specialize.cpp
+++ b/source/slang/slang-ir-specialize.cpp
@@ -7,6 +7,7 @@
 #include "slang-ir-ssa-simplification.h"
 #include "slang-ir-lower-witness-lookup.h"
 #include "slang-ir-dce.h"
+#include "slang-ir-sccp.h"
 #include "../core/slang-performance-profiler.h"
 
 namespace Slang
@@ -211,7 +212,29 @@ struct SpecializationContext
         IRGeneric* genericVal,
         IRSpecialize* specializeInst)
     {
-        // First, we want to see if an existing specialization
+        // We need to fold the generic arguments here in order to uniquely identify
+        // which specializations need to be generated.
+        // The folding of the generic-arguments are deferred until the specialization
+        // step here, because the exact value of the generic-arguments can be unknown
+        // until the specialization. The exact values of the generic-arguments may come
+        // from the other modules, and they will be unknown until linking the modules.
+        //
+        UInt argCount = specializeInst->getArgCount();
+        {
+            IRBuilder builder(module);
+            for (UInt ii = 0; ii < argCount; ++ii)
+            {
+                IRUse* argUse = specializeInst->getArgOperand(ii);
+                auto originalArg = argUse->get();
+                IRInst* foldedArg = tryConstantFoldInst(module, originalArg);
+                if (foldedArg == originalArg)
+                    continue;
+
+                specializeInst = as<IRSpecialize>(builder.replaceOperand(argUse, foldedArg));
+            }
+        }
+
+        // We want to see if an existing specialization
         // has already been made. To do that we will construct a key
         // for lookup in the generic specialization context.
         //
@@ -223,7 +246,6 @@ struct SpecializationContext
         //
         Key key;
         key.vals.add(specializeInst->getBase());
-        UInt argCount = specializeInst->getArgCount();
         for (UInt ii = 0; ii < argCount; ++ii)
         {
             key.vals.add(specializeInst->getArg(ii));
-- 
cgit v1.2.3