[Mlir-commits] [flang] [mlir] [MLIR][OpenMP] Support basic materialization for `omp.private` ops (PR #81715)

[Kareem Ergawy]

================
@@ -1000,6 +1000,26 @@ convertOmpWsLoop(Operation &opInst, llvm::IRBuilderBase &builder,
   return success();
 }
 
+/// Replace the region arguments of the parallel op (which correspond to private
+/// variables) with the actual private varibles they correspond to. This
+/// prepares the parallel op so that it matches what is expected by the
+/// OMPIRBuilder.
+static void prepareOmpParallelForPrivatization(omp::ParallelOp opInst) {
----------------
ergawy wrote:

> Could you elaborate on why this is necessary to match what the OMPIRBuilder expects?

The `OpenMPIRBuilder::createParallel(...)` [collects the inputs and outputs of the parallel region](https://github.com/llvm/llvm-project/blob/main/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp#L1502). The [inputs are then iterated](https://github.com/llvm/llvm-project/blob/main/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp#L1580) to invoke the private/firstprivate/shared callback.

The input as defined by the `CodeExtractor` is [any value defined outside the region and used inside it](https://github.com/llvm/llvm-project/blob/main/llvm/lib/Transforms/Utils/CodeExtractor.cpp#L649).

Therefore, if `OpenMPIRBuilder::createParallel` is invoked with the block arguments still used inside the parallel region, the IR builder won't detect these as inputs and won't invoke the `PrivCB` for private/firstprivate variables.

That's why `prepareOmpParallelForPrivatization` "rewires" the uses of private vars inside the region from the block args to the original SSA values.

> it also seems like it is making the omp.parallel operation invalid

Indeed, this does invalidate the op. But my understanding is that the op is at the end of its life at this stage and there is no need to keep it valid. However, to make this transformation less damaging, I removed the line that erases the region arguments. That way the region arguments are still there but just not used anymore. The op then is at a transitional state where original SSA values for private variables are used inside the region but the privatization logic is not inlined yet.

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