[Mlir-commits] [mlir] [mlir][OpenMP] - MLIR to LLVMIR translation support for delayed privatization in `omp.target` ops. (PR #109668)

[Pranav Bhandarkar]

================
@@ -3435,6 +3450,59 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
       const auto &arg = targetRegion.front().getArgument(argIndex);
       moduleTranslation.mapValue(arg, mapOpValue);
     }
+
+    // Do privatization after moduleTranslation has already recorded
+    // mapped values.
+    if (!targetOp.getPrivateVars().empty()) {
+      auto oldIP = builder.saveIP();
+      builder.restoreIP(allocaIP);
+
+      OperandRange privateVars = targetOp.getPrivateVars();
+      std::optional<ArrayAttr> privateSyms = targetOp.getPrivateSyms();
+      unsigned numMapVars = targetOp.getMapVars().size();
+      Block &firstTargetBlock = targetRegion.front();
+      BlockArgument *blockArgsStart = firstTargetBlock.getArguments().begin();
+      BlockArgument *privArgsStart = blockArgsStart + numMapVars;
+      BlockArgument *privArgsEnd =
+          privArgsStart + targetOp.getPrivateVars().size();
+      MutableArrayRef privateBlockArgs(privArgsStart, privArgsEnd);
+
+      for (auto [privVar, privatizerNameAttr, privBlockArg] :
+           llvm::zip_equal(privateVars, *privateSyms, privateBlockArgs)) {
+
+        SymbolRefAttr privSym = llvm::cast<SymbolRefAttr>(privatizerNameAttr);
+        omp::PrivateClauseOp privatizer = findPrivatizer(&opInst, privSym);
+        if (privatizer.getDataSharingType() ==
+                omp::DataSharingClauseType::FirstPrivate ||
+            !privatizer.getDeallocRegion().empty()) {
+          opInst.emitError("Translation of omp.target from MLIR to LLVMIR "
+                           "failed because translation of firstprivate and "
+                           " private allocatables is not supported yet");
+          bodyGenStatus = failure();
+        } else {
+          omp::PrivateClauseOp clonedPrivatizer =
+              clonePrivatizer(moduleTranslation, privatizer, &opInst);
+          Region &allocRegion = clonedPrivatizer.getAllocRegion();
+          BlockArgument allocRegionArg = allocRegion.getArgument(0);
+          replaceAllUsesInRegionWith(allocRegionArg, privVar, allocRegion);
+          SmallVector<llvm::Value *, 1> yieldedValues;
+          if (failed(inlineConvertOmpRegions(
+                  allocRegion, "omp.targetop.privatizer", builder,
+                  moduleTranslation, &yieldedValues))) {
+            opInst.emitError(
+                "failed to inline `alloc` region of an `omp.private` "
+                "op in the target region");
+            bodyGenStatus = failure();
+          } else {
+            assert(yieldedValues.size() == 1);
+            moduleTranslation.mapValue(privBlockArg, yieldedValues.front());
+          }
+          clonedPrivatizer.erase();
+          builder.restoreIP(oldIP);
+        }
+      }
+    }
+    builder.restoreIP(codeGenIP);
----------------
bhandarkar-pranav wrote:

Thank you for this suggestion @tblah and @ergawy. Indeed, a multi-block `alloc` region in the privatizer exposed a bug in my code. I have fixed it now. Please take a look. Mind you that the the test I have added has an `alloc` region is artificially multi-block (It needn't be). We'll have real multi-block privatizers when we deal with allocatables which I am working on right now (I have it working using the old approach of inlining in MLIR before converting to LLVMIR. Need to change that code to this newer approach now).

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