From 5cb960a72449dad36594e8ad2bfa899f5aaa11be Mon Sep 17 00:00:00 2001
From: Sai Praveen Bangaru <31557731+saipraveenb25@users.noreply.github.com>
Date: Wed, 13 Nov 2024 19:50:52 -0500
Subject: Fix WGSL emit for '&' and add bindings for thread group size (#5557)

Co-authored-by: Yong He <yonghe@outlook.com>
---
 source/slang-wasm/slang-wasm-bindings.cpp | 14 ++++++++++++++
 source/slang-wasm/slang-wasm.cpp          | 12 ++++++++++++
 source/slang-wasm/slang-wasm.h            | 16 ++++++++++++++++
 source/slang/slang-emit-wgsl.cpp          |  1 +
 4 files changed, 43 insertions(+)

diff --git a/source/slang-wasm/slang-wasm-bindings.cpp b/source/slang-wasm/slang-wasm-bindings.cpp
index 8c8257795..43a516ccc 100644
--- a/source/slang-wasm/slang-wasm-bindings.cpp
+++ b/source/slang-wasm/slang-wasm-bindings.cpp
@@ -53,6 +53,16 @@ EMSCRIPTEN_BINDINGS(slang)
             allow_raw_pointers())
         .function("getBindingIndex", &slang::wgsl::VariableLayoutReflection::getBindingIndex);
 
+    class_<slang::wgsl::EntryPointReflection>("EntryPointReflection")
+        .function(
+            "getComputeThreadGroupSize",
+            &slang::wgsl::EntryPointReflection::getComputeThreadGroupSize);
+
+    class_<slang::wgsl::EntryPointReflection::ThreadGroupSize>("ThreadGroupSize")
+        .property("x", &slang::wgsl::EntryPointReflection::ThreadGroupSize::x)
+        .property("y", &slang::wgsl::EntryPointReflection::ThreadGroupSize::y)
+        .property("z", &slang::wgsl::EntryPointReflection::ThreadGroupSize::z);
+
     class_<slang::wgsl::ProgramLayout>("ProgramLayout")
         .function("toJsonObject", &slang::wgsl::ProgramLayout::toJsonObject)
         .function("getParameterCount", &slang::wgsl::ProgramLayout::getParameterCount)
@@ -63,6 +73,10 @@ EMSCRIPTEN_BINDINGS(slang)
         .function(
             "getGlobalParamsTypeLayout",
             &slang::wgsl::ProgramLayout::getGlobalParamsTypeLayout,
+            allow_raw_pointers())
+        .function(
+            "findEntryPointByName",
+            &slang::wgsl::ProgramLayout::findEntryPointByName,
             allow_raw_pointers());
 
     enum_<slang::BindingType>("BindingType")
diff --git a/source/slang-wasm/slang-wasm.cpp b/source/slang-wasm/slang-wasm.cpp
index 10fc86029..a1ce2920d 100644
--- a/source/slang-wasm/slang-wasm.cpp
+++ b/source/slang-wasm/slang-wasm.cpp
@@ -442,6 +442,18 @@ TypeLayoutReflection* ProgramLayout::getGlobalParamsTypeLayout()
     return (slang::wgsl::TypeLayoutReflection*)(interface()->getGlobalParamsTypeLayout());
 }
 
+EntryPointReflection* ProgramLayout::findEntryPointByName(std::string name)
+{
+    return (slang::wgsl::EntryPointReflection*)(interface()->findEntryPointByName(name.c_str()));
+}
+
+EntryPointReflection::ThreadGroupSize EntryPointReflection::getComputeThreadGroupSize()
+{
+    SlangUInt size[3];
+    interface()->getComputeThreadGroupSize(3, size);
+    return {size[0], size[1], size[2]};
+}
+
 BindingType TypeLayoutReflection::getDescriptorSetDescriptorRangeType(
     unsigned int setIndex,
     unsigned int rangeIndex)
diff --git a/source/slang-wasm/slang-wasm.h b/source/slang-wasm/slang-wasm.h
index 4288aed18..f6b14117a 100644
--- a/source/slang-wasm/slang-wasm.h
+++ b/source/slang-wasm/slang-wasm.h
@@ -72,6 +72,20 @@ public:
     }
 };
 
+class EntryPointReflection
+{
+
+public:
+    struct ThreadGroupSize
+    {
+        unsigned int x;
+        unsigned int y;
+        unsigned int z;
+    };
+
+    ThreadGroupSize getComputeThreadGroupSize();
+    slang::EntryPointReflection* interface() const { return (slang::EntryPointReflection*)this; }
+};
 
 class ProgramLayout
 {
@@ -81,6 +95,8 @@ public:
 
     slang::wgsl::TypeLayoutReflection* getGlobalParamsTypeLayout();
 
+    slang::wgsl::EntryPointReflection* findEntryPointByName(std::string name);
+
     slang::ProgramLayout* interface() const { return (slang::ProgramLayout*)this; }
 
     emscripten::val toJsonObject();
diff --git a/source/slang/slang-emit-wgsl.cpp b/source/slang/slang-emit-wgsl.cpp
index b5ed7b9d2..fcc4b615f 100644
--- a/source/slang/slang-emit-wgsl.cpp
+++ b/source/slang/slang-emit-wgsl.cpp
@@ -1341,6 +1341,7 @@ bool WGSLSourceEmitter::tryEmitInstExprImpl(IRInst* inst, const EmitOpInfo& inOu
         }
     case kIROp_BitXor:
     case kIROp_BitOr:
+    case kIROp_BitAnd:
         {
             // Emit bitwise operators with paranthesis to avoid precedence issues
             const auto emitOp = getEmitOpForOp(inst->getOp());
-- 
cgit v1.2.3