[Mlir-commits] [mlir] [mlir][spirv] Allow yielding values from selection regions (PR #133702)

[Jakub Kuderski]

================
@@ -2000,17 +2024,79 @@ LogicalResult ControlFlowStructurizer::structurize() {
   for (auto *block : constructBlocks)
     block->dropAllReferences();
 
+  // All internal uses should be removed from original blocks by now, so
+  // whatever is left is an outside use and will need to be yielded from
+  // the newly created selection region.
+  if (!isLoop) {
+    for (Block *block : constructBlocks) {
+      for (Operation &op : *block) {
+        if (!op.use_empty())
+          for (Value result : op.getResults()) {
+            valuesToYield.push_back(mapper.lookupOrNull(result));
+            outsideUses.push_back(result);
+          }
+      }
+      for (BlockArgument &arg : block->getArguments()) {
+        if (!arg.use_empty()) {
+          valuesToYield.push_back(mapper.lookupOrNull(arg));
+          outsideUses.push_back(arg);
+        }
+      }
+    }
+  }
+
+  assert(valuesToYield.size() == outsideUses.size());
+
+  // If we need to yield any values from the selection region we will take
+  // care of it here.
+  if (!isLoop && !valuesToYield.empty()) {
+    LLVM_DEBUG(logger.startLine()
+               << "[cf] yielding values from the selection region\n");
+
+    // Update `mlir.merge` with values to be yield.
+    auto mergeOps = body.back().getOps<spirv::MergeOp>();
+    assert(std::next(mergeOps.begin()) == mergeOps.end());
+    Operation *merge = *mergeOps.begin();
+    merge->setOperands(valuesToYield);
+
+    // MLIR does not allow changing the number of results of an operation, so
+    // we create a new SelectionOp with required list of results and move
+    // the region from the initial SelectionOp. The initial operation is then
+    // removed. Since we move the region to the new op all links between blocks
+    // and remapping we have previously done should be preserved.
+    builder.setInsertionPoint(&mergeBlock->front());
+    auto selectionOp = builder.create<spirv::SelectionOp>(
+        location, TypeRange(outsideUses),
+        static_cast<spirv::SelectionControl>(control));
+    selectionOp->getRegion(0).takeBody(body);
+
+    // Remove initial op and swap the pointer to the newly created one.
+    op->erase();
+    op = selectionOp;
+
+    // Update all outside uses to use results of the SelectionOp and remove
+    // block arguments from the original merge block.
+    for (size_t i = 0; i < outsideUses.size(); i++)
+      outsideUses[i].replaceAllUsesWith(selectionOp.getResult(i));
+    for (size_t i = 0; i < mergeBlock->getNumArguments(); ++i)
----------------
kuhar wrote:

nit: don't recalculate the end value

https://github.com/llvm/llvm-project/pull/133702