From ab567491e15aa5c94d6e64956f46d7896583f321 Mon Sep 17 00:00:00 2001
From: comex <comexk@gmail.com>
Date: Thu, 26 Nov 2020 14:39:27 -0500
Subject: [PATCH] Add support for `OpGroupNonUniform{All,Any,AllEqual,Ballot}`,
 and fix `OpGroupNonUniformShuffleXor`.

These Vulkan 1.1 operations can be used in place of
`OpSubgroup{All,Any,AllEqual,Ballot}KHR`, among other things.

For `OpGroupNonUniformShuffleXor`, which was already implemented, turns
out the scope argument needs to be encoded not as an immediate, but as
an id that points to a constant integer.
---
 include/sirit/sirit.h      | 24 +++++++++++++++++++++++-
 src/instructions/group.cpp | 22 +++++++++++++++++++++-
 2 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/include/sirit/sirit.h b/include/sirit/sirit.h
index 7e1fd22..d7dfc78 100644
--- a/include/sirit/sirit.h
+++ b/include/sirit/sirit.h
@@ -1017,7 +1017,29 @@ public:
 
     /// Return the value of the invocation identified by the current invocation's id within the
     /// group xor'ed with mask.
-    Id OpGroupNonUniformShuffleXor(Id result_type, spv::Scope scope, Id value, Id mask);
+    Id OpGroupNonUniformShuffleXor(Id result_type, Id scope, Id value, Id mask);
+
+    /// Evaluates a predicate for all active invocations in the group, resulting in
+    /// true if predicate evaluates to true for all active invocations in the
+    /// group, otherwise the result is false.
+    Id OpGroupNonUniformAll(Id result_type, Id scope, Id predicate);
+
+    /// Evaluates a predicate for all active invocations in the group,
+    /// resulting in true if predicate evaluates to true for any active
+    /// invocation in the group, otherwise the result is false.
+    Id OpGroupNonUniformAny(Id result_type, Id scope, Id predicate);
+
+    /// Evaluates a value for all active invocations in the group. The result
+    /// is true if Value is equal for all active invocations in the group.
+    /// Otherwise, the result is false.
+    Id OpGroupNonUniformAllEqual(Id result_type, Id scope, Id value);
+
+    /// Result is a bitfield value combining the Predicate value from all
+    /// invocations in the group that execute the same dynamic instance of this
+    /// instruction. The bit is set to one if the corresponding invocation is
+    /// active and the Predicate for that invocation evaluated to true;
+    /// otherwise, it is set to zero.
+    Id OpGroupNonUniformBallot(Id result_type, Id scope, Id predicate);
 
     // Atomic
 
diff --git a/src/instructions/group.cpp b/src/instructions/group.cpp
index 605a17f..274b5b6 100644
--- a/src/instructions/group.cpp
+++ b/src/instructions/group.cpp
@@ -36,10 +36,30 @@ Id Module::OpSubgroupAllEqualKHR(Id result_type, Id predicate) {
     return *code << OpId{spv::Op::OpSubgroupAllEqualKHR, result_type} << predicate << EndOp{};
 }
 
-Id Module::OpGroupNonUniformShuffleXor(Id result_type, spv::Scope scope, Id value, Id mask) {
+Id Module::OpGroupNonUniformShuffleXor(Id result_type, Id scope, Id value, Id mask) {
     code->Reserve(6);
     return *code << OpId{spv::Op::OpGroupNonUniformShuffleXor, result_type} << scope << value
                  << mask << EndOp{};
 }
 
+Id Module::OpGroupNonUniformAll(Id result_type, Id scope, Id predicate) {
+   code->Reserve(5);
+   return *code << OpId{spv::Op::OpGroupNonUniformAll, result_type} << scope << predicate << EndOp{};
+}
+
+Id Module::OpGroupNonUniformAny(Id result_type, Id scope, Id predicate) {
+   code->Reserve(5);
+   return *code << OpId{spv::Op::OpGroupNonUniformAny, result_type} << scope << predicate << EndOp{};
+}
+
+Id Module::OpGroupNonUniformAllEqual(Id result_type, Id scope, Id value) {
+   code->Reserve(5);
+   return *code << OpId{spv::Op::OpGroupNonUniformAllEqual, result_type} << scope << value << EndOp{};
+}
+
+Id Module::OpGroupNonUniformBallot(Id result_type, Id scope, Id predicate) {
+   code->Reserve(5);
+   return *code << OpId{spv::Op::OpGroupNonUniformBallot, result_type} << scope << predicate << EndOp{};
+}
+
 } // namespace Sirit