Issue/legalize resource (#4769)

* Fix the issue that NonUniformResourceIndex is ignored

Fix the issue that after `specializeFunctionCalls`,
`NonUniformResourceIndex` is ignored in the generated specialized
function.

The reason is that if the function has a non-uniform resource parameter,
we will legalize it by replacing the resource parameter with a index,
and indexing of the resource will be moved inside the specialized function.

e.g.

```
void func(ResourceType resource) { ... }
func(resource[NonUniformResourceIndex(0)])
```
will be specialized into

```
void func(int index) { resource[index]; }
func(0);
```

In this case, inside the function, we will loose the information about
whether the resource is a non-uniform. So we add the handling for this
corner case by adding insert a `NonUniformResourceIndex` into the
specialized function:

```
void func(int index) {
int nonUniformIdx = NonUniformResourceIndex(index);
resource[nonUniformIdx];
}
```

* Fix the issue that arguments mismatch after specilization callsite

specializeCall() call could cause arguments mismatch with the parameters
of the specialized function.

For example, if the function parameter contains a resource type

```
void func(ResourceType res) { ... }

int index = ...
func(resources[index]);
```

This will be specialized into

```
void func(int index) { resources[index] }

int index = ...
func(index);
```

However, if we have more than 1 call sites, and the other call site
doesn't use `int` as the index, e.g.

```
uint index = ...
func(resources[index]);
```
this call site will be specialized into

```
uint index = ...
func(index);
```

this will be invalid, because the argument doesn't match the parameter.
so we just add the data type of the new arguments into the function key such that

For the uniformity info, we add a new attribute "IROp_NonUniformAttr",
so we will form a IRAttributedType that encodes both uniformity and data
type, and use it as the key of call info. So if there is call site using the different
data type for the resource index, we will specialize a new function for this.


* Handle the intCast and uintCast operation

Since after intCast/uintCast of nonuniformIndex, it's still a
nonuniformIndex. So we will handle this case as well.

Also, add a new test to cover this.

2 files changed, 92 insertions, 2 deletions

diff --git a/source/slang/slang-ir-inst-defs.h b/source/slang/slang-ir-inst-defs.h
index 84ee634a1..39de083f0 100644
--- a/source/slang/slang-ir-inst-defs.h
+++ b/source/slang/slang-ir-inst-defs.h
@@ -1173,6 +1173,7 @@ INST_RANGE(Layout, VarLayout, EntryPointLayout)
     INST(UNormAttr, unorm, 0, HOISTABLE)
     INST(SNormAttr, snorm, 0, HOISTABLE)
     INST(NoDiffAttr, no_diff, 0, HOISTABLE)
+    INST(NonUniformAttr, nonuniform, 0, HOISTABLE)
 
     /* SemanticAttr */
         INST(UserSemanticAttr, userSemantic, 2, HOISTABLE)
diff --git a/source/slang/slang-ir-specialize-function-call.cpp b/source/slang/slang-ir-specialize-function-call.cpp
index a41ca1e99..98aba0fae 100644
--- a/source/slang/slang-ir-specialize-function-call.cpp
+++ b/source/slang/slang-ir-specialize-function-call.cpp
@@ -5,6 +5,7 @@
 #include "slang-ir-clone.h"
 #include "slang-ir-insts.h"
 #include "slang-ir-ssa-simplification.h"
+#include "slang-ir-util.h"
 
 namespace Slang
 {
@@ -363,7 +364,7 @@ struct FunctionParameterSpecializationContext
             // a new callee function based on the original
             // function and the information we gathered.
             //
-            newFunc = generateSpecializedFunc(oldFunc, funcInfo);
+            newFunc = generateSpecializedFunc(oldFunc, funcInfo, callInfo);
             specializedFuncs.add(callInfo.key, newFunc);
         }
 
@@ -381,6 +382,7 @@ struct FunctionParameterSpecializationContext
         newCall->insertBefore(oldCall);
         oldCall->replaceUsesWith(newCall);
         oldCall->removeAndDeallocate();
+
     }
 
     // Before diving into the details on how we gather information
@@ -559,6 +561,21 @@ struct FunctionParameterSpecializationContext
             // the arguments at the new call site, and
             // don't add anything to the specialization key.
             //
+            // We should also add 2 more things such that our specialization
+            // can handle the corner cases that if the oldBase is a nonuniform
+            // resource and also the data type of oldIndex will be handled correctly.
+            // By doing so, we form an IRAttributedType to include both information
+            // and add it to the key of call info.
+
+            List<IRAttr*> irAttrs;
+            if (findNonuniformIndexInst(oldIndex))
+            {
+                IRAttr* attr = getBuilder()->getAttr(kIROp_NonUniformAttr);
+                irAttrs.add(attr);
+            }
+            auto irType = getBuilder()->getAttributedType(oldIndex->getDataType(), irAttrs);
+            ioInfo.key.vals.add(irType);
+
             ioInfo.newArgs.add(oldIndex);
         }
         else if (oldArg->getOp() == kIROp_Load)
@@ -577,6 +594,27 @@ struct FunctionParameterSpecializationContext
         }
     }
 
+    IRInst* findNonuniformIndexInst(IRInst* inst)
+    {
+        while(1)
+        {
+            if (inst == nullptr)
+                return nullptr;
+
+            if (inst->getOp() == kIROp_NonUniformResourceIndex)
+                return inst;
+
+            if (inst->getOp() == kIROp_IntCast)
+            {
+                inst = inst->getOperand(0);
+            }
+            else
+            {
+                return nullptr;
+            }
+        }
+    }
+
     // The remaining information we've discussed is only
     // gathered once we decide we want to generate a
     // specialized function, but it follows much the same flow.
@@ -803,7 +841,8 @@ struct FunctionParameterSpecializationContext
     //
     IRFunc* generateSpecializedFunc(
         IRFunc*                         oldFunc,
-        FuncSpecializationInfo const&   funcInfo)
+        FuncSpecializationInfo const&   funcInfo,
+        CallSpecializationInfo const&   callInfo)
     {
         // We will make use of the infrastructure for cloning
         // IR code, that is defined in `ir-clone.{h,cpp}`.
@@ -933,6 +972,18 @@ struct FunctionParameterSpecializationContext
             newBodyInst->insertBefore(newFirstOrdinary);
         }
 
+        // We need to handle a corner case where the new argument of
+        // the callee of this specialized function could be a use of
+        // NonUniformResourceIndex(), in such case, any indexing operation
+        // on the global buffer by using this new argument should be
+        // decorated with NonUniformDecoration. However, inside the new
+        // specialized function, we don't have that information anymore.
+        // Therefore, we will need to scan the new argument list to find out
+        // this case, and insert the NonUniformResourceIndex() instruction
+        // on the corresponding parameter of the new specialized function.
+        maybeInsertNonUniformResourceIndex(newFunc, funcInfo, callInfo);
+
+
         // At this point we've created a new specialized function,
         // and as such it may contain call sites that were not
         // covered when we built our initial work list.
@@ -964,6 +1015,44 @@ struct FunctionParameterSpecializationContext
 
         return newFunc;
     }
+
+    void maybeInsertNonUniformResourceIndex(
+        IRFunc* newFunc,
+        FuncSpecializationInfo const&   funcInfo,
+        CallSpecializationInfo const&   callInfo)
+    {
+        auto builder = getBuilder();
+        uint32_t paramIndex = 0;
+
+        SLANG_ASSERT(callInfo.newArgs.getCount() == funcInfo.newParams.getCount());
+
+        // Iterate over the new arguments, new parameters pair, and
+        // find out if there is any use of NonUniformResourceIndex()
+        // in the new arguments.
+        for (auto newArg : callInfo.newArgs)
+        {
+            if (auto nonuniformIdxInst = findNonuniformIndexInst(newArg))
+            {
+                auto firstOrdinary = newFunc->getFirstOrdinaryInst();
+
+                IRCloneEnv cloneEnv;
+                auto newParam = funcInfo.newParams[paramIndex];
+
+                // Clone the NonUniformResourceIndex call and insert it at beginning
+                // of the function. Then replace every use of the parameter with the
+                // NonUniformResourceIndex.
+                auto clonedInst = cloneInstAndOperands(&cloneEnv, builder, nonuniformIdxInst);
+                clonedInst->insertBefore(firstOrdinary);
+                newParam->replaceUsesWith(clonedInst);
+
+                // At last, set the operand of the NonUniformResourceIndex to the new parameter
+                // because we haven't done it yet during inst clone.
+                clonedInst->setOperand(0, newParam);
+            }
+            paramIndex++;
+        }
+
+    }
 };
 
 // The top-level function for invoking the specialization pass