Fix issue of infinity float literal (#5489)

* Fix issue of infinity float literal

* add parameters for the test

* Correct the way to construct inf and nan

In WGSL, expression of "1.0/0.0" is not allowed, it will report compile error,
so to construct infinity or nan, we have to assign the float literal to a variable
and then use it to bypass the compile error.

By doing so, we add getInfinity and getNan functions to the builtin
prelude to wgsl.

---------

Co-authored-by: Yong He <yonghe@outlook.com>

2 files changed, 63 insertions, 2 deletions

diff --git a/source/slang/slang-emit-wgsl.cpp b/source/slang/slang-emit-wgsl.cpp
index 52ff790d7..dea95c6ec 100644
--- a/source/slang/slang-emit-wgsl.cpp
+++ b/source/slang/slang-emit-wgsl.cpp
@@ -26,6 +26,39 @@
 namespace Slang
 {
 
+// In WGSL, expression of "1.0/0.0" is not allowed, it will report compile error,
+// so to construct infinity or nan, we have to assign the float literal to a variable
+// and then use it to bypass the compile error.
+static const char* kWGSLBuiltinPreludeGetInfinity = R"(
+fn _slang_getInfinity(positive: bool) -> f32
+{
+    let a = select(f32(-1.0), f32(1.0), positive);
+    let b = f32(0.0);
+    return a / b;
+}
+)";
+
+static const char* kWGSLBuiltinPreludeGetNan = R"(
+fn _slang_getNan() -> f32
+{
+    let a = f32(0.0);
+    let b = f32(0.0);
+    return a / b;
+}
+)";
+
+void WGSLSourceEmitter::ensurePrelude(const char* preludeText)
+{
+    IRStringLit* stringLit;
+    if (!m_builtinPreludes.tryGetValue(preludeText, stringLit))
+    {
+        IRBuilder builder(m_irModule);
+        stringLit = builder.getStringValue(UnownedStringSlice(preludeText));
+        m_builtinPreludes[preludeText] = stringLit;
+    }
+    m_requiredPreludes.add(stringLit);
+}
+
 void WGSLSourceEmitter::emitSwitchCaseSelectorsImpl(
     const SwitchRegion::Case* const currentCase,
     const bool isDefault)
@@ -878,8 +911,32 @@ void WGSLSourceEmitter::emitSimpleValueImpl(IRInst* inst)
 
                 case BaseType::Float:
                     {
-                        m_writer->emit(litInst->value.floatVal);
-                        m_writer->emit("f");
+                        IRConstant::FloatKind kind = litInst->getFloatKind();
+                        switch (kind)
+                        {
+                        case IRConstant::FloatKind::Nan:
+                            {
+                                ensurePrelude(kWGSLBuiltinPreludeGetNan);
+                                m_writer->emit("_slang_getNan()");
+                                break;
+                            }
+                        case IRConstant::FloatKind::PositiveInfinity:
+                            {
+                                ensurePrelude(kWGSLBuiltinPreludeGetInfinity);
+                                m_writer->emit("_slang_getInfinity(true)");
+                                break;
+                            }
+                        case IRConstant::FloatKind::NegativeInfinity:
+                            {
+                                ensurePrelude(kWGSLBuiltinPreludeGetInfinity);
+                                m_writer->emit("_slang_getInfinity(false)");
+                                break;
+                            }
+                        default:
+                            m_writer->emit(litInst->value.floatVal);
+                            m_writer->emit("f");
+                            break;
+                        }
                     }
                     break;
 
diff --git a/source/slang/slang-emit-wgsl.h b/source/slang/slang-emit-wgsl.h
index 1a8ec2fd5..f178d8f66 100644
--- a/source/slang/slang-emit-wgsl.h
+++ b/source/slang/slang-emit-wgsl.h
@@ -51,6 +51,10 @@ public:
     void emit(const AddressSpace addressSpace);
 
     virtual bool shouldFoldInstIntoUseSites(IRInst* inst) SLANG_OVERRIDE;
+    Dictionary<const char*, IRStringLit*> m_builtinPreludes;
+
+protected:
+    void ensurePrelude(const char* preludeText);
 
 private:
     // Emit the matrix type with 'rowCountWGSL' WGSL-rows and 'colCountWGSL' WGSL-columns