Add `sampleCount` parameter for MS textures. (#3001)

* Add `sampleCount` parameter for MS textures.

* Fix test.

---------

Co-authored-by: Yong He <yhe@nvidia.com>

9 files changed, 96 insertions, 69 deletions

diff --git a/source/slang/core.meta.slang b/source/slang/core.meta.slang
index 5342800be..08cbcd9d5 100644
--- a/source/slang/core.meta.slang
+++ b/source/slang/core.meta.slang
@@ -1329,9 +1329,12 @@ struct TextureTypeInfo
         flavor |= (access << 8);
 
         // emit a generic signature
-        // TODO: allow for multisample count to come in as well...
-        sb << "__generic<T = float4> ";
-
+        sb << "__generic<T = float4";
+        // Multi-sample rw texture types have an optional sampleCount parameter.
+        if (isMultisample && access == SLANG_RESOURCE_ACCESS_READ_WRITE)
+            sb << ", let sampleCount : int = 0";
+        sb << ">";
+            
         if(prefixInfo.combined)
         {
             sb << "__magic_type(TextureSampler," << int(flavor) << ")\n";
diff --git a/source/slang/slang-ast-type.h b/source/slang/slang-ast-type.h
index c5ec7e161..53544edfa 100644
--- a/source/slang/slang-ast-type.h
+++ b/source/slang/slang-ast-type.h
@@ -155,12 +155,17 @@ class TextureTypeBase : public ResourceType
 {
     SLANG_ABSTRACT_AST_CLASS(TextureTypeBase)
 
+    // The sampleCount parameter of a RWTexture*MS resource.
+    Val* sampleCount = nullptr;
 protected:
-    TextureTypeBase(TextureFlavor inFlavor, Type* inElementType)
+    TextureTypeBase(TextureFlavor inFlavor, Type* inElementType, Val* inSampleCount = nullptr)
     {
         elementType = inElementType;
         flavor = inFlavor;
+        sampleCount = inSampleCount;
     }
+
+    Val* getSampleCount() const { return sampleCount; }
 };
 
 
@@ -170,8 +175,8 @@ class TextureType : public TextureTypeBase
     SLANG_AST_CLASS(TextureType)
 
 protected:
-    TextureType(TextureFlavor flavor, Type* elementType)
-        : TextureTypeBase(flavor, elementType)
+    TextureType(TextureFlavor flavor, Type* elementType, Val* inSampleCount = nullptr)
+        : TextureTypeBase(flavor, elementType, inSampleCount)
     {}
 };
 
diff --git a/source/slang/slang-check-decl.cpp b/source/slang/slang-check-decl.cpp
index 5b7777154..4c1e967e3 100644
--- a/source/slang/slang-check-decl.cpp
+++ b/source/slang/slang-check-decl.cpp
@@ -270,6 +270,11 @@ namespace Slang
             checkVarDeclCommon(varDecl);
         }
 
+        void visitGenericValueParamDecl(GenericValueParamDecl* genValDecl)
+        {
+            checkVarDeclCommon(genValDecl);
+        }
+
         void visitGlobalGenericValueParamDecl(GlobalGenericValueParamDecl* decl)
         {
             checkVarDeclCommon(decl);
diff --git a/source/slang/slang-check-overload.cpp b/source/slang/slang-check-overload.cpp
index 1d2f82ac7..2fa13f3fa 100644
--- a/source/slang/slang-check-overload.cpp
+++ b/source/slang/slang-check-overload.cpp
@@ -196,6 +196,14 @@ namespace Slang
     {
         auto genericDeclRef = candidate.item.declRef.as<GenericDecl>();
 
+        // Only allow constructing a PartialGenericAppExpr when referencing a callable decl.
+        // Other types of generic decls must be fully specified.
+        bool allowPartialGenericApp = false;
+        if (as<CallableDecl>(genericDeclRef.getDecl()->inner))
+        {
+            allowPartialGenericApp = true;
+        }
+
         // The basic idea here is that we need to check that the
         // arguments to a generic application (e.g., `F<A1, A2, ...>`)
         // have the right "type," which in this context means
@@ -235,13 +243,27 @@ namespace Slang
             {
                 if (aa >= context.argCount)
                 {
-                    // If we have run out of arguments, then we don't
-                    // apply any more checks at this step. We will instead
-                    // attempt to *infer* an argument at this position
-                    // at a later stage.
-                    //
-                    candidate.flags |= OverloadCandidate::Flag::IsPartiallyAppliedGeneric;
-                    break;
+                    if (allowPartialGenericApp)
+                    {
+                        // If we have run out of arguments, and the referenced decl
+                        // allows partially applied specialization (i.e. a callable
+                        // decl) then we don't apply any more checks at this step.
+                        // We will instead attempt to *infer* an argument at this
+                        // position at a later stage.
+                        //
+                        candidate.flags |= OverloadCandidate::Flag::IsPartiallyAppliedGeneric;
+                        break;
+                    }
+                    else
+                    {
+                        // Otherwise, the generic decl had better provide a default value
+                        // or this reference is ill-formed.
+                        auto substType = typeParamRef.substitute(m_astBuilder, typeParamRef.getDecl()->initType.type);
+                        if (!substType)
+                            return false;
+                        checkedArgs.add(substType);
+                        continue;
+                    }
                 }
 
                 // We have a type parameter, and we expect to find
@@ -287,13 +309,29 @@ namespace Slang
             {
                 if (aa >= context.argCount)
                 {
-                    // If we have run out of arguments, then we don't
-                    // apply any more checks at this step. We will instead
-                    // attempt to *infer* an argument at this position
-                    // at a later stage.
-                    //
-                    candidate.flags |= OverloadCandidate::Flag::IsPartiallyAppliedGeneric;
-                    break;
+                    if (allowPartialGenericApp)
+                    {
+                        // If we have run out of arguments and the decl allows
+                        // partial specialization, then we don't apply any more
+                        // checks at this step. We will instead attempt to
+                        // *infer* an argument at this position at a later
+                        // stage.
+                        //
+                        candidate.flags |= OverloadCandidate::Flag::IsPartiallyAppliedGeneric;
+                        break;
+                    }
+                    else
+                    {
+                        // Otherwise, the generic decl had better provide a default value
+                        // or this reference is ill-formed.
+                        ensureDecl(valParamRef, DeclCheckState::Checked);
+                        ConstantFoldingCircularityInfo newCircularityInfo(valParamRef.getDecl(), nullptr);
+                        auto defaultVal = tryConstantFoldExpr(valParamRef.substitute(m_astBuilder, valParamRef.getDecl()->initExpr), &newCircularityInfo);
+                        if (!defaultVal)
+                            return false;
+                        checkedArgs.add(defaultVal);
+                        continue;
+                    }
                 }
 
                 // The case for a generic value parameter is similar to that
diff --git a/source/slang/slang-check-type.cpp b/source/slang/slang-check-type.cpp
index fd4c218fa..cee54388f 100644
--- a/source/slang/slang-check-type.cpp
+++ b/source/slang/slang-check-type.cpp
@@ -108,7 +108,7 @@ namespace Slang
         {
             return expr;
         }
-
+        
         getSink()->diagnose(expr, Diagnostics::expectedAType, expr->type);
         return CreateErrorExpr(expr);
     }
diff --git a/source/slang/slang-emit-hlsl.cpp b/source/slang/slang-emit-hlsl.cpp
index bc6b7a159..30de45773 100644
--- a/source/slang/slang-emit-hlsl.cpp
+++ b/source/slang/slang-emit-hlsl.cpp
@@ -287,6 +287,15 @@ void HLSLSourceEmitter::_emitHLSLTextureType(IRTextureTypeBase* texType)
     }
     m_writer->emit("<");
     emitType(texType->getElementType());
+    if (texType->getOperandCount() == 2)
+    {
+        auto sampleCount = as<IRIntLit>(texType->getSampleCount());
+        if (sampleCount->getValue() != 0)
+        {
+            m_writer->emit(", ");
+            m_writer->emit(sampleCount->getValue());
+        }
+    }
     m_writer->emit(" >");
 }
 
diff --git a/source/slang/slang-ir.h b/source/slang/slang-ir.h
index 2edda0f4f..e7cc0d6c8 100644
--- a/source/slang/slang-ir.h
+++ b/source/slang/slang-ir.h
@@ -1336,6 +1336,7 @@ struct IRResourceTypeBase : IRType
 struct IRResourceType : IRResourceTypeBase
 {
     IRType* getElementType() { return (IRType*)getOperand(0); }
+    IRType* getSampleCount() { return (IRType*)getOperand(1); }
 
     IR_PARENT_ISA(ResourceType)
 };
diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index 63dc39a20..e49da00c2 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -2063,45 +2063,9 @@ struct ValLoweringVisitor : ValVisitor<ValLoweringVisitor, LoweredValInfo, Lower
             operands.getBuffer());
     }
 
-    // Lower a type where the type declaration being referenced is assumed
-    // to be an intrinsic type with a single generic type parameter, and
-    // which can thus be lowered to a simple IR type with the appropriate opcode.
-    IRType* lowerGenericIntrinsicType(DeclRefType* type, Type* elementType)
-    {
-        auto intrinsicTypeModifier = type->declRef.getDecl()->findModifier<IntrinsicTypeModifier>();
-        SLANG_ASSERT(intrinsicTypeModifier);
-        IROp op = IROp(intrinsicTypeModifier->irOp);
-        IRInst* irElementType = lowerType(context, elementType);
-        return getBuilder()->getType(
-            op,
-            1,
-            &irElementType);
-    }
-
-    IRType* lowerGenericIntrinsicType(DeclRefType* type, Type* elementType, IntVal* count)
-    {
-        auto intrinsicTypeModifier = type->declRef.getDecl()->findModifier<IntrinsicTypeModifier>();
-        SLANG_ASSERT(intrinsicTypeModifier);
-        IROp op = IROp(intrinsicTypeModifier->irOp);
-        IRInst* irElementType = lowerType(context, elementType);
-
-        IRInst* irCount = lowerSimpleVal(context, count);
-
-        IRInst* const operands[2] =
-        {
-            irElementType,
-            irCount,
-        };
-
-        return getBuilder()->getType(
-            op,
-            SLANG_COUNT_OF(operands),
-            operands);
-    }
-
     IRType* visitResourceType(ResourceType* type)
     {
-        return lowerGenericIntrinsicType(type, type->elementType);
+        return lowerSimpleIntrinsicType(type);
     }
 
     IRType* visitSamplerStateType(SamplerStateType* type)
@@ -2111,7 +2075,7 @@ struct ValLoweringVisitor : ValVisitor<ValLoweringVisitor, LoweredValInfo, Lower
 
     IRType* visitBuiltinGenericType(BuiltinGenericType* type)
     {
-        return lowerGenericIntrinsicType(type, type->elementType);
+        return lowerSimpleIntrinsicType(type);
     }
 
     IRType* visitUntypedBufferResourceType(UntypedBufferResourceType* type)
@@ -2121,18 +2085,12 @@ struct ValLoweringVisitor : ValVisitor<ValLoweringVisitor, LoweredValInfo, Lower
 
     IRType* visitHLSLPatchType(HLSLPatchType* type)
     {
-        Type* elementType = type->getElementType();
-        IntVal* count = type->getElementCount();
-
-        return lowerGenericIntrinsicType(type, elementType, count);
+        return lowerSimpleIntrinsicType(type);
     }
 
     IRType* visitMeshOutputType(MeshOutputType* type)
     {
-        Type* elementType = type->getElementType();
-        IntVal* count = type->getMaxElementCount();
-
-        return lowerGenericIntrinsicType(type, elementType, count);
+        return lowerSimpleIntrinsicType(type);
     }
 
     IRType* visitExtractExistentialType(ExtractExistentialType* type)
diff --git a/source/slang/slang-syntax.cpp b/source/slang/slang-syntax.cpp
index f70d2bd85..249969729 100644
--- a/source/slang/slang-syntax.cpp
+++ b/source/slang/slang-syntax.cpp
@@ -573,9 +573,17 @@ Index getFilterCountImpl(const ReflectClassInfo& clsInfo, MemberFilterStyle filt
             else if (magicMod->magicName == "Texture")
             {
                 SLANG_ASSERT(subst && subst->getArgs().getCount() >= 1);
+                auto textureTag = TextureFlavor(magicMod->tag);
+                Val* sampleCount = nullptr;
+                if (textureTag.isMultisample() && textureTag.getAccess() == SLANG_RESOURCE_ACCESS_READ_WRITE)
+                {
+                    if (subst->getArgs().getCount() >= 2)
+                        sampleCount = ExtractGenericArgInteger(subst->getArgs().getLast());
+                }
                 auto textureType = astBuilder->getOrCreate<TextureType>(
-                    TextureFlavor(magicMod->tag),
-                    ExtractGenericArgType(subst->getArgs()[0]));
+                    textureTag,
+                    ExtractGenericArgType(subst->getArgs()[0]),
+                    sampleCount);
                 textureType->declRef = declRef;
                 return textureType;
             }