Support arithmetics on generic arguments (#3968)

Resovles an issue #3935

Slang had to fold the generic arguments after specialization.

3 files changed, 60 insertions, 2 deletions

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));
diff --git a/tests/bugs/gh-3935.slang b/tests/bugs/gh-3935.slang
new file mode 100644
index 000000000..9c8e9f19c
--- /dev/null
+++ b/tests/bugs/gh-3935.slang
@@ -0,0 +1,35 @@
+// Test if generic arguments using arithmetics are folded properly.
+
+//TEST:SIMPLE(filecheck=HLSL): -stage compute -entry computeMain -target hlsl
+
+RWStructuredBuffer<float> outputBuffer;
+
+__generic<let ArraySize:int>
+struct MyStruct
+{
+  int elems[ArraySize];
+};
+
+__generic<let ArraySize:int>
+MyStruct<ArraySize / 2> Reduce(MyStruct<ArraySize> o)
+{
+  MyStruct<ArraySize / 2> result;
+  //HLSL:[[ReturnType:MyStruct_[0-9]*]] Reduce
+  //HLSL:[[ReturnType]] result
+
+  result.elems[0] = o.elems[0];
+
+  // Error happened here as the return type differs
+  // from the type of "result".
+  return result;
+}
+
+[numthreads(1, 1, 1)]
+void computeMain(int3 dispatchThreadID : SV_DispatchThreadID)
+{
+  MyStruct<4> a;
+  a.elems[0] = 2;
+
+  let result = Reduce(a);
+  outputBuffer[0] = result.elems[0];
+}