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

[Kareem Ergawy via flang-commits]

================
@@ -3316,7 +3316,51 @@ createDeviceArgumentAccessor(MapInfoData &mapData, llvm::Argument &arg,
 
   return builder.saveIP();
 }
+static bool privatizerNeedsMap(omp::PrivateClauseOp &privatizer) {
+  Region &allocRegion = privatizer.getAllocRegion();
+  Value blockArg0 = allocRegion.getArgument(0);
+  return !blockArg0.use_empty();
+}
+
+// Return the llvm::Value * corresponding to the privateVar that
+// is being privatized. It isn't always as simple as looking up
+// moduleTranslation with privateVar. For instance, in case of
+// an allocatable, the descriptor for the allocatable is privatized.
+// This descriptor is mapped using an MapInfoOp. So, this function
+// will return a pointer to the llvm::Value corresponding to the
+// block argument for the mapped descriptor.
+static llvm::Value *
+findHostAssociatedValue(Value privateVar, omp::TargetOp targetOp,
+                        llvm::DenseMap<Value, int> &mappedPrivateVars,
+                        llvm::IRBuilderBase &builder,
+                        LLVM::ModuleTranslation &moduleTranslation) {
+  if (mappedPrivateVars.contains(privateVar)) {
+    int blockArgIndex = mappedPrivateVars[privateVar];
+    Value blockArg = targetOp.getRegion().getArgument(blockArgIndex);
+    mlir::Type privVarType = privateVar.getType();
+    mlir::Type blockArgType = blockArg.getType();
+    assert(isa<LLVM::LLVMPointerType>(blockArgType) &&
+           "A block argument corresponding to a mapped var should have "
+           "!llvm.ptr type");
+
+    if (privVarType == blockArg.getType()) {
+      llvm::Value *v = moduleTranslation.lookupValue(blockArg);
+      return v;
+    }
 
+    if (!isa<LLVM::LLVMPointerType>(privVarType)) {
+      // This typically happens when the privatized type is lowered from
+      // boxchar<KIND> and gets lowered to !llvm.struct<(ptr, i64)>. That is the
+      // struct/pair is passed by value. But, mapped values are passed only as
+      // pointers, so before we privatize, we must load the pointer.
+      llvm::Value *load =
+          builder.CreateLoad(moduleTranslation.convertType(privVarType),
+                             moduleTranslation.lookupValue(blockArg));
+      return load;
+    }
+  }
+  return moduleTranslation.lookupValue(privateVar);
----------------
ergawy wrote:

nit: Early return, just to follow the code style.
```suggestion
  if (!mappedPrivateVars.contains(privateVar))
    return moduleTranslation.lookupValue(privateVar);
    
    int blockArgIndex = mappedPrivateVars[privateVar];
    Value blockArg = targetOp.getRegion().getArgument(blockArgIndex);
    mlir::Type privVarType = privateVar.getType();
    mlir::Type blockArgType = blockArg.getType();
    assert(isa<LLVM::LLVMPointerType>(blockArgType) &&
           "A block argument corresponding to a mapped var should have "
           "!llvm.ptr type");

    if (privVarType == blockArg.getType()) {
      llvm::Value *v = moduleTranslation.lookupValue(blockArg);
      return v;
    }

    if (!isa<LLVM::LLVMPointerType>(privVarType)) {
      // This typically happens when the privatized type is lowered from
      // boxchar<KIND> and gets lowered to !llvm.struct<(ptr, i64)>. That is the
      // struct/pair is passed by value. But, mapped values are passed only as
      // pointers, so before we privatize, we must load the pointer.
      llvm::Value *load =
          builder.CreateLoad(moduleTranslation.convertType(privVarType),
                             moduleTranslation.lookupValue(blockArg));
      return load;
    }
```

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