Add tooling support for SPV_KHR_maximal_reconvergence (KhronosGroup#5542)

* Validation for SPV_KHR_maximal_reconvergence
* Add pass to add/remove maximal reconvergence execution mode
---------

Co-authored-by: David Neto <[email protected]>

Author: alan-baker

Android.mk

@@ -157,6 +157,7 @@ SPVTOOLS_OPT_SRC_FILES := \
 		source/opt/loop_utils.cpp \
 		source/opt/mem_pass.cpp \
 		source/opt/merge_return_pass.cpp \
+		source/opt/modify_maximal_reconvergence.cpp \
 		source/opt/module.cpp \
 		source/opt/optimizer.cpp \
 		source/opt/pass.cpp \

BUILD.gn

@@ -737,6 +737,8 @@ static_library("spvtools_opt") {
     "source/opt/mem_pass.h",
     "source/opt/merge_return_pass.cpp",
     "source/opt/merge_return_pass.h",
+    "source/opt/modify_maximal_reconvergence.cpp",
+    "source/opt/modify_maximal_reconvergence.h",
     "source/opt/module.cpp",
     "source/opt/module.h",
     "source/opt/null_pass.h",

DEPS

@@ -13,7 +13,7 @@ vars = {
   'protobuf_revision': 'v21.12',
 
   're2_revision': '264e71e88e1c8a4b5ec326e70e9cf1d476f58a58',
-  'spirv_headers_revision': '7b0309708da5126b89e4ce6f19835f36dc912f2f',
+  'spirv_headers_revision': 'ae6a8b39717523d96683bc0d20b541944e28072f',
 }
 
 deps = {

include/spirv-tools/optimizer.hpp

@@ -1003,6 +1003,10 @@ Optimizer::PassToken CreateSwitchDescriptorSetPass(uint32_t ds_from,
 // OpBeginInterlockInvocationEXT and one OpEndInterlockInvocationEXT, in that
 // order.
 Optimizer::PassToken CreateInvocationInterlockPlacementPass();
+
+// Creates a pass to add/remove maximal reconvergence execution mode.
+// This pass either adds or removes maximal reconvergence from all entry points.
+Optimizer::PassToken CreateModifyMaximalReconvergencePass(bool add);
 }  // namespace spvtools
 
 #endif  // INCLUDE_SPIRV_TOOLS_OPTIMIZER_HPP_

source/opt/CMakeLists.txt

@@ -94,6 +94,7 @@ set(SPIRV_TOOLS_OPT_SOURCES
   loop_unswitch_pass.h
   mem_pass.h
   merge_return_pass.h
+  modify_maximal_reconvergence.h
   module.h
   null_pass.h
   passes.h
@@ -214,6 +215,7 @@ set(SPIRV_TOOLS_OPT_SOURCES
   loop_unswitch_pass.cpp
   mem_pass.cpp
   merge_return_pass.cpp
+  modify_maximal_reconvergence.cpp
   module.cpp
   optimizer.cpp
   pass.cpp

source/opt/modify_maximal_reconvergence.cpp

@@ -0,0 +1,103 @@
+// Copyright (c) 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "modify_maximal_reconvergence.h"
+
+#include "source/opt/ir_context.h"
+#include "source/util/make_unique.h"
+
+namespace spvtools {
+namespace opt {
+
+Pass::Status ModifyMaximalReconvergence::Process() {
+  bool changed = false;
+  if (add_) {
+    changed = AddMaximalReconvergence();
+  } else {
+    changed = RemoveMaximalReconvergence();
+  }
+  return changed ? Pass::Status::SuccessWithChange
+                 : Pass::Status::SuccessWithoutChange;
+}
+
+bool ModifyMaximalReconvergence::AddMaximalReconvergence() {
+  bool changed = false;
+  bool has_extension = false;
+  bool has_shader =
+      context()->get_feature_mgr()->HasCapability(spv::Capability::Shader);
+  for (auto extension : context()->extensions()) {
+    if (extension.GetOperand(0).AsString() == "SPV_KHR_maximal_reconvergence") {
+      has_extension = true;
+      break;
+    }
+  }
+
+  std::unordered_set<uint32_t> entry_points_with_mode;
+  for (auto mode : get_module()->execution_modes()) {
+    if (spv::ExecutionMode(mode.GetSingleWordInOperand(1)) ==
+        spv::ExecutionMode::MaximallyReconvergesKHR) {
+      entry_points_with_mode.insert(mode.GetSingleWordInOperand(0));
+    }
+  }
+
+  for (auto entry_point : get_module()->entry_points()) {
+    const uint32_t id = entry_point.GetSingleWordInOperand(1);
+    if (!entry_points_with_mode.count(id)) {
+      changed = true;
+      if (!has_extension) {
+        context()->AddExtension("SPV_KHR_maximal_reconvergence");
+        has_extension = true;
+      }
+      if (!has_shader) {
+        context()->AddCapability(spv::Capability::Shader);
+        has_shader = true;
+      }
+      context()->AddExecutionMode(MakeUnique<Instruction>(
+          context(), spv::Op::OpExecutionMode, 0, 0,
+          std::initializer_list<Operand>{
+              {SPV_OPERAND_TYPE_ID, {id}},
+              {SPV_OPERAND_TYPE_EXECUTION_MODE,
+               {static_cast<uint32_t>(
+                   spv::ExecutionMode::MaximallyReconvergesKHR)}}}));
+      entry_points_with_mode.insert(id);
+    }
+  }
+
+  return changed;
+}
+
+bool ModifyMaximalReconvergence::RemoveMaximalReconvergence() {
+  bool changed = false;
+  std::vector<Instruction*> to_remove;
+  Instruction* mode = &*get_module()->execution_mode_begin();
+  while (mode) {
+    if (mode->opcode() != spv::Op::OpExecutionMode &&
+        mode->opcode() != spv::Op::OpExecutionModeId) {
+      break;
+    }
+    if (spv::ExecutionMode(mode->GetSingleWordInOperand(1)) ==
+        spv::ExecutionMode::MaximallyReconvergesKHR) {
+      mode = context()->KillInst(mode);
+      changed = true;
+    } else {
+      mode = mode->NextNode();
+    }
+  }
+
+  changed |=
+      context()->RemoveExtension(Extension::kSPV_KHR_maximal_reconvergence);
+  return changed;
+}
+}  // namespace opt
+}  // namespace spvtools

source/opt/modify_maximal_reconvergence.h

@@ -0,0 +1,53 @@
+// Copyright (c) 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef LIBSPIRV_OPT_MODIFY_MAXIMAL_RECONVERGENCE_H_
+#define LIBSPIRV_OPT_MODIFY_MAXIMAL_RECONVERGENCE_H_
+
+#include "pass.h"
+
+namespace spvtools {
+namespace opt {
+
+// Modifies entry points to either add or remove MaximallyReconvergesKHR
+//
+// This pass will either add or remove MaximallyReconvergesKHR to all entry
+// points in the module. When adding the execution mode, it does not attempt to
+// determine whether any ray tracing invocation repack instructions might be
+// executed because it is a runtime restriction. That is left to the user.
+class ModifyMaximalReconvergence : public Pass {
+ public:
+  const char* name() const override { return "modify-maximal-reconvergence"; }
+  Status Process() override;
+
+  explicit ModifyMaximalReconvergence(bool add = true) : Pass(), add_(add) {}
+
+  IRContext::Analysis GetPreservedAnalyses() override {
+    return IRContext::kAnalysisDefUse |
+           IRContext::kAnalysisInstrToBlockMapping |
+           IRContext::kAnalysisDecorations | IRContext::kAnalysisCombinators |
+           IRContext::kAnalysisCFG | IRContext::kAnalysisNameMap |
+           IRContext::kAnalysisConstants | IRContext::kAnalysisTypes;
+  }
+
+ private:
+  bool AddMaximalReconvergence();
+  bool RemoveMaximalReconvergence();
+
+  bool add_;
+};
+}  // namespace opt
+}  // namespace spvtools
+
+#endif  // LIBSPIRV_OPT_MODIFY_MAXIMAL_RECONVERGENCE_H_

source/opt/optimizer.cpp

@@ -606,6 +606,23 @@ bool Optimizer::RegisterPassFromFlag(const std::string& flag,
       return false;
     }
     RegisterPass(CreateSwitchDescriptorSetPass(from_set, to_set));
+  } else if (pass_name == "modify-maximal-reconvergence") {
+    if (pass_args.size() == 0) {
+      Error(consumer(), nullptr, {},
+            "--modify-maximal-reconvergence requires an argument");
+      return false;
+    }
+    if (pass_args == "add") {
+      RegisterPass(CreateModifyMaximalReconvergencePass(true));
+    } else if (pass_args == "remove") {
+      RegisterPass(CreateModifyMaximalReconvergencePass(false));
+    } else {
+      Errorf(consumer(), nullptr, {},
+             "Invalid argument for --modify-maximal-reconvergence: %s (must be "
+             "'add' or 'remove')",
+             pass_args.c_str());
+      return false;
+    }
   } else {
     Errorf(consumer(), nullptr, {},
            "Unknown flag '--%s'. Use --help for a list of valid flags",
@@ -1141,6 +1158,11 @@ Optimizer::PassToken CreateInvocationInterlockPlacementPass() {
   return MakeUnique<Optimizer::PassToken::Impl>(
       MakeUnique<opt::InvocationInterlockPlacementPass>());
 }
+
+Optimizer::PassToken CreateModifyMaximalReconvergencePass(bool add) {
+  return MakeUnique<Optimizer::PassToken::Impl>(
+      MakeUnique<opt::ModifyMaximalReconvergence>(add));
+}
 }  // namespace spvtools
 
 extern "C" {

source/opt/passes.h

@@ -65,6 +65,7 @@
 #include "source/opt/loop_unroller.h"
 #include "source/opt/loop_unswitch_pass.h"
 #include "source/opt/merge_return_pass.h"
+#include "source/opt/modify_maximal_reconvergence.h"
 #include "source/opt/null_pass.h"
 #include "source/opt/private_to_local_pass.h"
 #include "source/opt/reduce_load_size.h"

source/val/validate_cfg.cpp

@@ -190,6 +190,8 @@ spv_result_t ValidateBranchConditional(ValidationState_t& _,
               "ID of an OpLabel instruction";
   }
 
+  // A similar requirement for SPV_KHR_maximal_reconvergence is deferred until
+  // entry point call trees have been reconrded.
   if (_.version() >= SPV_SPIRV_VERSION_WORD(1, 6) && true_id == false_id) {
     return _.diag(SPV_ERROR_INVALID_ID, inst)
            << "In SPIR-V 1.6 or later, True Label and False Label must be "
@@ -875,6 +877,95 @@ spv_result_t StructuredControlFlowChecks(
   return SPV_SUCCESS;
 }
 
+spv_result_t MaximalReconvergenceChecks(ValidationState_t& _) {
+  // Find all the entry points with the MaximallyReconvergencesKHR execution
+  // mode.
+  std::unordered_set<uint32_t> maximal_funcs;
+  std::unordered_set<uint32_t> maximal_entry_points;
+  for (auto entry_point : _.entry_points()) {
+    const auto* exec_modes = _.GetExecutionModes(entry_point);
+    if (exec_modes &&
+        exec_modes->count(spv::ExecutionMode::MaximallyReconvergesKHR)) {
+      maximal_entry_points.insert(entry_point);
+      maximal_funcs.insert(entry_point);
+    }
+  }
+
+  if (maximal_entry_points.empty()) {
+    return SPV_SUCCESS;
+  }
+
+  // Find all the functions reachable from a maximal reconvergence entry point.
+  for (const auto& func : _.functions()) {
+    const auto& entry_points = _.EntryPointReferences(func.id());
+    for (auto id : entry_points) {
+      if (maximal_entry_points.count(id)) {
+        maximal_funcs.insert(func.id());
+        break;
+      }
+    }
+  }
+
+  // Check for conditional branches with the same true and false targets.
+  for (const auto& inst : _.ordered_instructions()) {
+    if (inst.opcode() == spv::Op::OpBranchConditional) {
+      const auto true_id = inst.GetOperandAs<uint32_t>(1);
+      const auto false_id = inst.GetOperandAs<uint32_t>(2);
+      if (true_id == false_id && maximal_funcs.count(inst.function()->id())) {
+        return _.diag(SPV_ERROR_INVALID_ID, &inst)
+               << "In entry points using the MaximallyReconvergesKHR execution "
+                  "mode, True Label and False Label must be different labels";
+      }
+    }
+  }
+
+  // Check for invalid multiple predecessors. Only loop headers, continue
+  // targets, merge targets or switch targets or defaults may have multiple
+  // unique predecessors.
+  for (const auto& func : _.functions()) {
+    if (!maximal_funcs.count(func.id())) continue;
+
+    for (const auto* block : func.ordered_blocks()) {
+      std::unordered_set<uint32_t> unique_preds;
+      const auto* preds = block->predecessors();
+      if (!preds) continue;
+
+      for (const auto* pred : *preds) {
+        unique_preds.insert(pred->id());
+      }
+      if (unique_preds.size() < 2) continue;
+
+      const auto* terminator = block->terminator();
+      const auto index = terminator - &_.ordered_instructions()[0];
+      const auto* pre_terminator = &_.ordered_instructions()[index - 1];
+      if (pre_terminator->opcode() == spv::Op::OpLoopMerge) continue;
+
+      const auto* label = _.FindDef(block->id());
+      bool ok = false;
+      for (const auto& pair : label->uses()) {
+        const auto* use_inst = pair.first;
+        switch (use_inst->opcode()) {
+          case spv::Op::OpSelectionMerge:
+          case spv::Op::OpLoopMerge:
+          case spv::Op::OpSwitch:
+            ok = true;
+            break;
+          default:
+            break;
+        }
+      }
+      if (!ok) {
+        return _.diag(SPV_ERROR_INVALID_CFG, label)
+               << "In entry points using the MaximallyReconvergesKHR "
+                  "execution mode, this basic block must not have multiple "
+                  "unique predecessors";
+      }
+    }
+  }
+
+  return SPV_SUCCESS;
+}
+
 spv_result_t PerformCfgChecks(ValidationState_t& _) {
   for (auto& function : _.functions()) {
     // Check all referenced blocks are defined within a function
@@ -999,6 +1090,11 @@ spv_result_t PerformCfgChecks(ValidationState_t& _) {
         return error;
     }
   }
+
+  if (auto error = MaximalReconvergenceChecks(_)) {
+    return error;
+  }
+
   return SPV_SUCCESS;
 }