[Mlir-commits] [llvm] [mlir] [mlir][OpenMP] Allow multi-block omp.iterator regions (PR #229198)

[llvmlistbot at llvm.org]

================
@@ -3050,16 +3049,41 @@ convertIteratorRegion(llvm::Value *linearIV, IteratorInfo &iterInfo,
         iterInfo.getLowerBounds()[d],
         builder.CreateMul(idx, iterInfo.getSteps()[d]), "omp.it.phys_iv");
 
-    moduleTranslation.mapValue(iteratorRegionBlock.getArgument(d), physIV);
+    moduleTranslation.mapValue(entryBlock.getArgument(d), physIV);
   }
 
   // Translate the iterator region into the loop body.
-  moduleTranslation.mapBlock(&iteratorRegionBlock, builder.GetInsertBlock());
-  if (mlir::failed(moduleTranslation.convertBlock(iteratorRegionBlock,
-                                                  /*ignoreArguments=*/true,
-                                                  builder))) {
-    return mlir::failure();
+  if (iteratorRegion.hasOneBlock()) {
+    moduleTranslation.mapBlock(&entryBlock, builder.GetInsertBlock());
+    return moduleTranslation.convertBlock(entryBlock, /*ignoreArguments=*/true,
+                                          builder);
+  }
+
+  // Unlike convertOmpOpRegions, do not forward the yielded value through a PHI:
+  // it may have no LLVM type (e.g. `!omp.affinity_entry_ty`). The unique
+  // `omp.yield` block is the continuation's only predecessor, so callers can
+  // look up the yielded value directly.
+  llvm::BasicBlock *continuation =
+      splitBB(builder, true, "omp.iterator.region.cont");
+  llvm::Instruction *sourceBranch = builder.GetInsertBlock()->getTerminator();
+  for (mlir::Block &block : iteratorRegion)
+    moduleTranslation.mapBlock(
+        &block,
+        llvm::BasicBlock::Create(builder.getContext(), "omp.iterator.region",
+                                 continuation->getParent(), continuation));
+  sourceBranch->setSuccessor(0, moduleTranslation.lookupBlock(&entryBlock));
+
+  for (mlir::Block *block : getBlocksSortedByDominance(iteratorRegion)) {
+    llvm::IRBuilderBase::InsertPointGuard guard(builder);
+    if (failed(moduleTranslation.convertBlock(*block, block->isEntryBlock(),
+                                              builder)))
+      return mlir::failure();
+    if (isa<omp::YieldOp>(block->getTerminator()))
+      builder.CreateBr(continuation);
   }
+  LLVM::detail::connectPHINodes(iteratorRegion, moduleTranslation);
----------------
MattPD wrote:

The new internal PHI receives an i64 value for this declared i32 argument. `mlir-opt repro.mlir` accepts the input, but `mlir-translate --mlir-to-llvmir repro.mlir` aborts with `All operands to PHI node must be the same type as the PHI node!`. Before the patch, the verifier rejected the multi-block region.

```mlir
llvm.func @narrow(%x: !llvm.ptr) {
  %lo = llvm.mlir.constant(1 : i32) : i32
  %hi = llvm.mlir.constant(2 : i32) : i32
  %it = omp.iterator(%i: i32) = (%lo to %hi step %lo) {
    llvm.br ^next(%i : i32)
  ^next(%v: i32):
    llvm.store %v, %x : i32, !llvm.ptr
    omp.yield(%x : !llvm.ptr)
  } -> !omp.iterated<!llvm.ptr>
  omp.taskwait depend(taskdependin -> %it : !omp.iterated<!llvm.ptr>)
  llvm.return
}
```

`IteratorInfo` still converts the bounds to i64, and reconstruction maps that i64 value directly to the entry argument. Do you think you could preserve the declared induction type in this prerequisite, rather than relying on the later changes in https://github.com/llvm/llvm-project/pull/227454? Converting the i64 value back to the declared type would still lose values outside the i64 range.

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