From 7c2ff54758d26b73074fd14143ecd843ba685e0d Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Mon, 4 Nov 2024 17:37:50 -0800
Subject: Various WGSL fixes. (#5490)

* [WGSL] make sure switch has a default label.

* Various WGSL fixes.

* Update rhi submodule commit

* format code

* Remove unnecessary DISABLE_TEST directive on not applicable test.

* Matrix comp mul + `select`.

* Legalize binary ops for wgsl.

---------

Co-authored-by: slangbot <186143334+slangbot@users.noreply.github.com>
---
 source/slang/slang-emit-wgsl.cpp | 126 +++++++++++++++++++++++++++++++++++++--
 1 file changed, 121 insertions(+), 5 deletions(-)

(limited to 'source/slang/slang-emit-wgsl.cpp')

diff --git a/source/slang/slang-emit-wgsl.cpp b/source/slang/slang-emit-wgsl.cpp
index 4aca03a61..d8ec01776 100644
--- a/source/slang/slang-emit-wgsl.cpp
+++ b/source/slang/slang-emit-wgsl.cpp
@@ -497,6 +497,34 @@ void WGSLSourceEmitter::emitLayoutQualifiersImpl(IRVarLayout* layout)
     }
 }
 
+static bool isStaticConst(IRInst* inst)
+{
+    if (inst->getParent()->getOp() == kIROp_Module)
+    {
+        return true;
+    }
+    switch (inst->getOp())
+    {
+    case kIROp_MakeVector:
+    case kIROp_swizzle:
+    case kIROp_swizzleSet:
+    case kIROp_IntCast:
+    case kIROp_FloatCast:
+    case kIROp_CastFloatToInt:
+    case kIROp_CastIntToFloat:
+    case kIROp_BitCast:
+        {
+            for (UInt i = 0; i < inst->getOperandCount(); i++)
+            {
+                if (!isStaticConst(inst->getOperand(i)))
+                    return false;
+            }
+            return true;
+        }
+    }
+    return false;
+}
+
 void WGSLSourceEmitter::emitVarKeywordImpl(IRType* type, IRInst* varDecl)
 {
     switch (varDecl->getOp())
@@ -505,14 +533,10 @@ void WGSLSourceEmitter::emitVarKeywordImpl(IRType* type, IRInst* varDecl)
     case kIROp_GlobalVar:
     case kIROp_Var:         m_writer->emit("var"); break;
     default:
-        if (as<IRModuleInst>(varDecl->getParent()))
-        {
+        if (isStaticConst(varDecl))
             m_writer->emit("const");
-        }
         else
-        {
             m_writer->emit("var");
-        }
         break;
     }
 
@@ -977,6 +1001,33 @@ void WGSLSourceEmitter::emitCallArg(IRInst* inst)
     }
 }
 
+bool WGSLSourceEmitter::shouldFoldInstIntoUseSites(IRInst* inst)
+{
+    bool result = CLikeSourceEmitter::shouldFoldInstIntoUseSites(inst);
+    if (result)
+    {
+        // If inst is a matrix, and is used in a component-wise multiply,
+        // we need to not fold it.
+        if (as<IRMatrixType>(inst->getDataType()))
+        {
+            for (auto use = inst->firstUse; use; use = use->nextUse)
+            {
+                auto user = use->getUser();
+                if (user->getOp() == kIROp_Mul)
+                {
+                    if (as<IRMatrixType>(user->getOperand(0)->getDataType()) &&
+                        as<IRMatrixType>(user->getOperand(1)->getDataType()))
+                    {
+                        return false;
+                    }
+                }
+            }
+        }
+    }
+    return result;
+}
+
+
 bool WGSLSourceEmitter::tryEmitInstExprImpl(IRInst* inst, const EmitOpInfo& inOuterPrec)
 {
     EmitOpInfo outerPrec = inOuterPrec;
@@ -1126,6 +1177,71 @@ bool WGSLSourceEmitter::tryEmitInstExprImpl(IRInst* inst, const EmitOpInfo& inOu
             return true;
         }
         break;
+
+    case kIROp_GetStringHash:
+        {
+            auto getStringHashInst = as<IRGetStringHash>(inst);
+            auto stringLit = getStringHashInst->getStringLit();
+
+            if (stringLit)
+            {
+                auto slice = stringLit->getStringSlice();
+                emitType(inst->getDataType());
+                m_writer->emit("(");
+                m_writer->emit((int)getStableHashCode32(slice.begin(), slice.getLength()).hash);
+                m_writer->emit(")");
+            }
+            else
+            {
+                // Couldn't handle
+                diagnoseUnhandledInst(inst);
+            }
+            return true;
+        }
+
+    case kIROp_Mul:
+        {
+            if (!as<IRMatrixType>(inst->getOperand(0)->getDataType()) ||
+                !as<IRMatrixType>(inst->getOperand(1)->getDataType()))
+            {
+                return false;
+            }
+            // Mul(m1, m2) should be translated to component-wise multiplication in WGSL.
+            auto matrixType = as<IRMatrixType>(inst->getDataType());
+            auto rowCount = getIntVal(matrixType->getRowCount());
+            emitType(inst->getDataType());
+            m_writer->emit("(");
+            for (IRIntegerValue i = 0; i < rowCount; i++)
+            {
+                if (i != 0)
+                {
+                    m_writer->emit(", ");
+                }
+                emitOperand(inst->getOperand(0), getInfo(EmitOp::Postfix));
+                m_writer->emit("[");
+                m_writer->emit(i);
+                m_writer->emit("] * ");
+                emitOperand(inst->getOperand(1), getInfo(EmitOp::Postfix));
+                m_writer->emit("[");
+                m_writer->emit(i);
+                m_writer->emit("]");
+            }
+            m_writer->emit(")");
+
+            return true;
+        }
+
+    case kIROp_Select:
+        {
+            m_writer->emit("select(");
+            emitOperand(inst->getOperand(2), getInfo(EmitOp::General));
+            m_writer->emit(", ");
+            emitOperand(inst->getOperand(1), getInfo(EmitOp::General));
+            m_writer->emit(", ");
+            emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+            m_writer->emit(")");
+            return true;
+        }
     }
 
     return false;
-- 
cgit v1.2.3