Respect matrix layout in uniform and in/out parameters for HLSL target. (#5013)

* Respect matrix layout in uniform and in/out parameters for HLSL target.

* Update test.

* Fix test.

* fix test.

* Fix metal layout calculation.

* Fix compile error.

* Fix compiler error.

---------

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

12 files changed, 109 insertions, 57 deletions

diff --git a/source/slang/slang-check-decl.cpp b/source/slang/slang-check-decl.cpp
index ec59469d3..ad3f94fc3 100644
--- a/source/slang/slang-check-decl.cpp
+++ b/source/slang/slang-check-decl.cpp
@@ -93,7 +93,7 @@ namespace Slang
         void checkDerivativeMemberAttributeParent(VarDeclBase* varDecl, DerivativeMemberAttribute* attr);
         void checkExtensionExternVarAttribute(VarDeclBase* varDecl, ExtensionExternVarModifier* m);
         void checkMeshOutputDecl(VarDeclBase* varDecl);
-
+        void maybeApplyMatrixLayoutModifier(VarDeclBase* varDecl);
         void checkVarDeclCommon(VarDeclBase* varDecl);
 
         void visitVarDecl(VarDecl* varDecl)
@@ -1516,6 +1516,22 @@ namespace Slang
             }
         }
     }
+    void SemanticsDeclHeaderVisitor::maybeApplyMatrixLayoutModifier(VarDeclBase* varDecl)
+    {
+        if (auto matrixType = as<MatrixExpressionType>(varDecl->type.type))
+        {
+            if (auto matrixLayoutModifier = varDecl->findModifier<MatrixLayoutModifier>())
+            {
+                auto matrixLayout = as<ColumnMajorLayoutModifier>(matrixLayoutModifier) ? SLANG_MATRIX_LAYOUT_COLUMN_MAJOR : SLANG_MATRIX_LAYOUT_ROW_MAJOR;
+                auto newMatrixType = getASTBuilder()->getMatrixType(
+                    matrixType->getElementType(),
+                    matrixType->getRowCount(),
+                    matrixType->getColumnCount(),
+                    getASTBuilder()->getIntVal(getASTBuilder()->getIntType(), matrixLayout));
+                varDecl->type.type = newMatrixType;
+            }
+        }
+    }
 
     void SemanticsDeclHeaderVisitor::checkVarDeclCommon(VarDeclBase* varDecl)
     {
@@ -1598,19 +1614,7 @@ namespace Slang
         }
 
         // If there is a matrix layout modifier, we will modify the matrix type now.
-        if (auto matrixType = as<MatrixExpressionType>(varDecl->type.type))
-        {
-            if (auto matrixLayoutModifier = varDecl->findModifier<MatrixLayoutModifier>())
-            {
-                auto matrixLayout = as<ColumnMajorLayoutModifier>(matrixLayoutModifier) ? SLANG_MATRIX_LAYOUT_COLUMN_MAJOR : SLANG_MATRIX_LAYOUT_ROW_MAJOR;
-                auto newMatrixType = getASTBuilder()->getMatrixType(
-                    matrixType->getElementType(),
-                    matrixType->getRowCount(),
-                    matrixType->getColumnCount(),
-                    getASTBuilder()->getIntVal(getASTBuilder()->getIntType(), matrixLayout));
-                varDecl->type.type = newMatrixType;
-            }
-        }
+        maybeApplyMatrixLayoutModifier(varDecl);
 
         if (varDecl->initExpr)
         {
@@ -7676,6 +7680,8 @@ namespace Slang
             }
         }
 
+        maybeApplyMatrixLayoutModifier(paramDecl);
+
         // Only texture types are allowed to have memory qualifiers on parameters
         if(!paramDecl->type || paramDecl->type->astNodeType != ASTNodeType::TextureType)
         {
diff --git a/source/slang/slang-emit-c-like.cpp b/source/slang/slang-emit-c-like.cpp
index ad365581a..1893929f8 100644
--- a/source/slang/slang-emit-c-like.cpp
+++ b/source/slang/slang-emit-c-like.cpp
@@ -3520,7 +3520,7 @@ void CLikeSourceEmitter::emitParamTypeImpl(IRType* type, String const& name)
     {
         type = constRefType->getValueType();
     }
-
+    emitParamTypeModifier(type);
     emitType(type, name);
 }
 
@@ -3862,7 +3862,7 @@ void CLikeSourceEmitter::emitVarModifiers(IRVarLayout* layout, IRInst* varDecl,
     if (!layout)
         return;
 
-    emitMatrixLayoutModifiersImpl(layout);
+    emitMatrixLayoutModifiersImpl(varType);
 
     // Target specific modifier output
     emitImageFormatModifierImpl(varDecl, varType);
diff --git a/source/slang/slang-emit-c-like.h b/source/slang/slang-emit-c-like.h
index 756b59913..00ad156d1 100644
--- a/source/slang/slang-emit-c-like.h
+++ b/source/slang/slang-emit-c-like.h
@@ -497,7 +497,7 @@ public:
     virtual void emitMeshShaderModifiersImpl(IRInst* varInst) { SLANG_UNUSED(varInst) }
     virtual void emitSimpleTypeImpl(IRType* type) = 0;
     virtual void emitVarDecorationsImpl(IRInst* varDecl) { SLANG_UNUSED(varDecl);  }
-    virtual void emitMatrixLayoutModifiersImpl(IRVarLayout* layout) { SLANG_UNUSED(layout);  }
+    virtual void emitMatrixLayoutModifiersImpl(IRType* varType) { SLANG_UNUSED(varType);  }
     virtual void emitTypeImpl(IRType* type, const StringSliceLoc* nameLoc);
     virtual void emitSimpleValueImpl(IRInst* inst);
     virtual void emitModuleImpl(IRModule* module, DiagnosticSink* sink);
@@ -505,6 +505,7 @@ public:
     virtual void emitVarExpr(IRInst* inst, EmitOpInfo const& outerPrec);
     virtual void emitOperandImpl(IRInst* inst, EmitOpInfo const& outerPrec);
     virtual void emitParamTypeImpl(IRType* type, String const& name);
+    virtual void emitParamTypeModifier(IRType* type) { SLANG_UNUSED(type); }
     virtual void emitIntrinsicCallExprImpl(IRCall* inst, UnownedStringSlice intrinsicDefinition, IRInst* intrinsicInst, EmitOpInfo const& inOuterPrec);
     virtual void emitFunctionPreambleImpl(IRInst* inst) { SLANG_UNUSED(inst); }
     virtual void emitLoopControlDecorationImpl(IRLoopControlDecoration* decl) { SLANG_UNUSED(decl); }
diff --git a/source/slang/slang-emit-cuda.cpp b/source/slang/slang-emit-cuda.cpp
index ab7ee8f58..05485855c 100644
--- a/source/slang/slang-emit-cuda.cpp
+++ b/source/slang/slang-emit-cuda.cpp
@@ -810,9 +810,9 @@ void CUDASourceEmitter::emitVarDecorationsImpl(IRInst* varDecl)
     Super::emitVarDecorationsImpl(varDecl);
 }
 
-void CUDASourceEmitter::emitMatrixLayoutModifiersImpl(IRVarLayout* layout)
+void CUDASourceEmitter::emitMatrixLayoutModifiersImpl(IRType* varType)
 {
-    Super::emitMatrixLayoutModifiersImpl(layout);
+    Super::emitMatrixLayoutModifiersImpl(varType);
 }
 
 bool CUDASourceEmitter::tryEmitGlobalParamImpl(IRGlobalParam* varDecl, IRType* varType)
diff --git a/source/slang/slang-emit-cuda.h b/source/slang/slang-emit-cuda.h
index ac7a33302..8b9c98c41 100644
--- a/source/slang/slang-emit-cuda.h
+++ b/source/slang/slang-emit-cuda.h
@@ -76,7 +76,7 @@ protected:
     virtual void emitSimpleValueImpl(IRInst* inst) SLANG_OVERRIDE;
     virtual void emitVectorTypeNameImpl(IRType* elementType, IRIntegerValue elementCount) SLANG_OVERRIDE;
     virtual void emitVarDecorationsImpl(IRInst* varDecl) SLANG_OVERRIDE;
-    virtual void emitMatrixLayoutModifiersImpl(IRVarLayout* layout) SLANG_OVERRIDE;
+    virtual void emitMatrixLayoutModifiersImpl(IRType* varType) SLANG_OVERRIDE;
     virtual void emitFunctionPreambleImpl(IRInst* inst) SLANG_OVERRIDE;
     virtual void _maybeEmitExportLike(IRInst* inst) SLANG_OVERRIDE { SLANG_UNUSED(inst); }
 
diff --git a/source/slang/slang-emit-glsl.cpp b/source/slang/slang-emit-glsl.cpp
index fff1cddbe..56113409d 100644
--- a/source/slang/slang-emit-glsl.cpp
+++ b/source/slang/slang-emit-glsl.cpp
@@ -2850,28 +2850,28 @@ void GLSLSourceEmitter::emitVarDecorationsImpl(IRInst* varDecl)
 
 }
 
-void GLSLSourceEmitter::emitMatrixLayoutModifiersImpl(IRVarLayout* layout)
+void GLSLSourceEmitter::emitMatrixLayoutModifiersImpl(IRType* varType)
 {
     // When a variable has a matrix type, we want to emit an explicit
     // layout qualifier based on what the layout has been computed to be.
     //
 
-    auto typeLayout = layout->getTypeLayout()->unwrapArray();
+    auto matrixType = as<IRMatrixType>(unwrapArray(varType));
 
-    if (auto matrixTypeLayout = as<IRMatrixTypeLayout>(typeLayout))
+    if (matrixType)
     {
         // Reminder: the meaning of row/column major layout
         // in our semantics is the *opposite* of what GLSL
         // calls them, because what they call "columns"
         // are what we call "rows."
         //
-        switch (matrixTypeLayout->getMode())
+        switch (getIntVal(matrixType->getLayout()))
         {
-            case kMatrixLayoutMode_ColumnMajor:
+            case SLANG_MATRIX_LAYOUT_COLUMN_MAJOR:
                 m_writer->emit("layout(row_major)\n");
                 break;
 
-            case kMatrixLayoutMode_RowMajor:
+            case SLANG_MATRIX_LAYOUT_ROW_MAJOR:
                 m_writer->emit("layout(column_major)\n");
                 break;
         }
diff --git a/source/slang/slang-emit-glsl.h b/source/slang/slang-emit-glsl.h
index cc8a4b02d..8958c7608 100644
--- a/source/slang/slang-emit-glsl.h
+++ b/source/slang/slang-emit-glsl.h
@@ -42,7 +42,7 @@ protected:
     virtual void emitSimpleTypeImpl(IRType* type) SLANG_OVERRIDE;
     virtual void emitVectorTypeNameImpl(IRType* elementType, IRIntegerValue elementCount) SLANG_OVERRIDE;
     virtual void emitVarDecorationsImpl(IRInst* varDecl) SLANG_OVERRIDE;
-    virtual void emitMatrixLayoutModifiersImpl(IRVarLayout* layout) SLANG_OVERRIDE;
+    virtual void emitMatrixLayoutModifiersImpl(IRType* varType) SLANG_OVERRIDE;
     virtual void emitTypeImpl(IRType* type, const StringSliceLoc* nameAndLoc) SLANG_OVERRIDE;
     virtual void emitParamTypeImpl(IRType* type, String const& name) SLANG_OVERRIDE;
     virtual void emitFuncDecorationImpl(IRDecoration* decoration) SLANG_OVERRIDE;
diff --git a/source/slang/slang-emit-hlsl.cpp b/source/slang/slang-emit-hlsl.cpp
index a867d4c06..f9765a555 100644
--- a/source/slang/slang-emit-hlsl.cpp
+++ b/source/slang/slang-emit-hlsl.cpp
@@ -959,15 +959,37 @@ void HLSLSourceEmitter::emitSimpleTypeImpl(IRType* type)
         case kIROp_MatrixType:
         {
             auto matType = (IRMatrixType*)type;
-
-            // TODO(tfoley): should really emit these with sugar
-            m_writer->emit("matrix<");
-            emitType(matType->getElementType());
-            m_writer->emit(",");
-            emitVal(matType->getRowCount(), getInfo(EmitOp::General));
-            m_writer->emit(",");
-            emitVal(matType->getColumnCount(), getInfo(EmitOp::General));
-            m_writer->emit("> ");           
+            bool canUseSugar = true;
+            switch (matType->getElementType()->getOp())
+            {
+            case kIROp_IntType:
+            case kIROp_UIntType:
+            case kIROp_FloatType:
+                canUseSugar = true;
+                break;
+            default:
+                canUseSugar = false;
+                break;
+            }
+            if (!as<IRIntLit>(matType->getRowCount()) || !as<IRIntLit>(matType->getColumnCount()))
+                canUseSugar = false;
+            if (canUseSugar)
+            {
+                emitType(matType->getElementType());
+                m_writer->emitInt64(getIntVal(matType->getRowCount()));
+                m_writer->emit("x");
+                m_writer->emitInt64(getIntVal(matType->getColumnCount()));
+            }
+            else
+            {
+                m_writer->emit("matrix<");
+                emitType(matType->getElementType());
+                m_writer->emit(",");
+                emitVal(matType->getRowCount(), getInfo(EmitOp::General));
+                m_writer->emit(",");
+                emitVal(matType->getColumnCount(), getInfo(EmitOp::General));
+                m_writer->emit("> ");
+            }
             return;
         }
         case kIROp_SamplerStateType:
@@ -1306,25 +1328,24 @@ void HLSLSourceEmitter::emitVarDecorationsImpl(IRInst* varDecl)
     }
 }
 
-void HLSLSourceEmitter::emitMatrixLayoutModifiersImpl(IRVarLayout* layout)
+void HLSLSourceEmitter::emitMatrixLayoutModifiersImpl(IRType* type)
 {
-    // When a variable has a matrix type, we want to emit an explicit
-    // layout qualifier based on what the layout has been computed to be.
-    //
-
-    auto typeLayout = layout->getTypeLayout()->unwrapArray();
-
-    if (auto matrixTypeLayout = as<IRMatrixTypeLayout>(typeLayout))
+    auto matType = as<IRMatrixType>(type);
+    if (!matType)
+        return;
+    auto matrixLayout = getIntVal(matType->getLayout());
+    if (getTargetProgram()->getOptionSet().getMatrixLayoutMode() != (MatrixLayoutMode)matrixLayout)
     {
-        switch (matrixTypeLayout->getMode())
+        switch (matrixLayout)
         {
-            case kMatrixLayoutMode_ColumnMajor:
-                m_writer->emit("column_major ");
-                break;
-
-            case kMatrixLayoutMode_RowMajor:
-                m_writer->emit("row_major ");
-                break;
+        case SLANG_MATRIX_LAYOUT_COLUMN_MAJOR:
+            m_writer->emit("column_major ");
+            break;
+        case SLANG_MATRIX_LAYOUT_ROW_MAJOR:
+            m_writer->emit("row_major ");
+            break;
+        default:
+            break;
         }
     }
 }
diff --git a/source/slang/slang-emit-hlsl.h b/source/slang/slang-emit-hlsl.h
index 2137a6d3d..4d721a1ac 100644
--- a/source/slang/slang-emit-hlsl.h
+++ b/source/slang/slang-emit-hlsl.h
@@ -45,7 +45,8 @@ protected:
     virtual void emitSimpleTypeImpl(IRType* type) SLANG_OVERRIDE;
     virtual void emitVectorTypeNameImpl(IRType* elementType, IRIntegerValue elementCount) SLANG_OVERRIDE;
     virtual void emitVarDecorationsImpl(IRInst* varDecl) SLANG_OVERRIDE;
-    virtual void emitMatrixLayoutModifiersImpl(IRVarLayout* layout) SLANG_OVERRIDE;
+    virtual void emitMatrixLayoutModifiersImpl(IRType* varType) SLANG_OVERRIDE;
+    virtual void emitParamTypeModifier(IRType* type) SLANG_OVERRIDE { emitMatrixLayoutModifiersImpl(type); }
 
     virtual bool tryEmitInstExprImpl(IRInst* inst, const EmitOpInfo& inOuterPrec) SLANG_OVERRIDE;
     virtual void emitSimpleValueImpl(IRInst* inst) SLANG_OVERRIDE;
diff --git a/source/slang/slang-emit-metal.cpp b/source/slang/slang-emit-metal.cpp
index 61e310401..a0fe220e5 100644
--- a/source/slang/slang-emit-metal.cpp
+++ b/source/slang/slang-emit-metal.cpp
@@ -1091,7 +1091,7 @@ void MetalSourceEmitter::emitVarDecorationsImpl(IRInst* varInst)
     SLANG_UNUSED(varInst);
 }
 
-void MetalSourceEmitter::emitMatrixLayoutModifiersImpl(IRVarLayout*)
+void MetalSourceEmitter::emitMatrixLayoutModifiersImpl(IRType*)
 {
     // Metal only supports column major layout, and we must have
     // already translated all matrix ops to assume column-major
diff --git a/source/slang/slang-emit-metal.h b/source/slang/slang-emit-metal.h
index 6460b076a..67aa0d506 100644
--- a/source/slang/slang-emit-metal.h
+++ b/source/slang/slang-emit-metal.h
@@ -44,7 +44,7 @@ protected:
     virtual void emitParamTypeImpl(IRType* type, String const& name) SLANG_OVERRIDE;
     virtual void emitVectorTypeNameImpl(IRType* elementType, IRIntegerValue elementCount) SLANG_OVERRIDE;
     virtual void emitVarDecorationsImpl(IRInst* varDecl) SLANG_OVERRIDE;
-    virtual void emitMatrixLayoutModifiersImpl(IRVarLayout* layout) SLANG_OVERRIDE;
+    virtual void emitMatrixLayoutModifiersImpl(IRType* varType) SLANG_OVERRIDE;
 
     virtual bool tryEmitInstExprImpl(IRInst* inst, const EmitOpInfo& inOuterPrec) SLANG_OVERRIDE;
     virtual bool tryEmitInstStmtImpl(IRInst* inst) SLANG_OVERRIDE;
diff --git a/source/slang/slang-type-layout.cpp b/source/slang/slang-type-layout.cpp
index 7d6d047d3..f654135a1 100644
--- a/source/slang/slang-type-layout.cpp
+++ b/source/slang/slang-type-layout.cpp
@@ -606,6 +606,28 @@ struct CUDALayoutRulesImpl : DefaultLayoutRulesImpl
     }
 };
 
+struct MetalLayoutRulesImpl : public CPULayoutRulesImpl
+{
+    SimpleLayoutInfo GetVectorLayout(BaseType elementType, SimpleLayoutInfo elementInfo, size_t elementCount) override
+    {
+        SLANG_UNUSED(elementType);
+
+        const auto elementSize = elementInfo.size.getFiniteValue();
+        auto alignedElementCount = 1 << Math::Log2Ceil((uint32_t)elementCount);
+
+        // Metal aligns vectors to 2/4 element boundaries.
+        size_t size = elementSize * elementCount;
+        size_t alignment = alignedElementCount * elementSize;
+
+        SimpleLayoutInfo vectorInfo;
+        vectorInfo.kind = elementInfo.kind;
+        vectorInfo.size = size;
+        vectorInfo.alignment = alignment;
+     
+        return vectorInfo;
+    }
+};
+
 struct HLSLStructuredBufferLayoutRulesImpl : DefaultLayoutRulesImpl
 {
     // HLSL structured buffers drop the restrictions added for constant buffers,
@@ -1696,6 +1718,7 @@ struct MetalArgumentBufferElementLayoutRulesImpl : ObjectLayoutRulesImpl, Defaul
 
 static MetalObjectLayoutRulesImpl kMetalObjectLayoutRulesImpl;
 static MetalArgumentBufferElementLayoutRulesImpl kMetalArgumentBufferElementLayoutRulesImpl;
+static MetalLayoutRulesImpl kMetalLayoutRulesImpl;
 
 LayoutRulesImpl kMetalAnyValueLayoutRulesImpl_ = {
     &kMetalLayoutRulesFamilyImpl,
@@ -1704,7 +1727,7 @@ LayoutRulesImpl kMetalAnyValueLayoutRulesImpl_ = {
 };
 
 LayoutRulesImpl kMetalConstantBufferLayoutRulesImpl_ = {
-    &kMetalLayoutRulesFamilyImpl, & kCPULayoutRulesImpl, &kMetalObjectLayoutRulesImpl,
+    &kMetalLayoutRulesFamilyImpl, & kMetalLayoutRulesImpl, &kMetalObjectLayoutRulesImpl,
 };
 
 LayoutRulesImpl kMetalParameterBlockLayoutRulesImpl_ = {
@@ -1712,7 +1735,7 @@ LayoutRulesImpl kMetalParameterBlockLayoutRulesImpl_ = {
 };
 
 LayoutRulesImpl kMetalStructuredBufferLayoutRulesImpl_ = {
-    &kMetalLayoutRulesFamilyImpl, &kCPULayoutRulesImpl, & kMetalObjectLayoutRulesImpl,
+    &kMetalLayoutRulesFamilyImpl, &kMetalLayoutRulesImpl, & kMetalObjectLayoutRulesImpl,
 };
 
 LayoutRulesImpl kMetalVaryingInputLayoutRulesImpl_ = {