[Mlir-commits] [mlir] [MLIR][OpenMP] Fix MLIR->LLVM value matching in privatization logic (PR #103718)

[Leandro Lupori]

================
@@ -1424,35 +1424,107 @@ convertOmpParallel(omp::ParallelOp opInst, llvm::IRBuilderBase &builder,
     }
   };
 
-  SmallVector<omp::PrivateClauseOp> privatizerClones;
-  SmallVector<llvm::Value *> privateVariables;
+  SmallVector<omp::PrivateClauseOp> mlirPrivatizerClones;
+  SmallVector<llvm::Value *> llvmPrivateVars;
 
   // TODO: Perform appropriate actions according to the data-sharing
   // attribute (shared, private, firstprivate, ...) of variables.
   // Currently shared and private are supported.
   auto privCB = [&](InsertPointTy allocaIP, InsertPointTy codeGenIP,
-                    llvm::Value &, llvm::Value &vPtr,
-                    llvm::Value *&replacementValue) -> InsertPointTy {
-    replacementValue = &vPtr;
+                    llvm::Value &, llvm::Value &llvmOmpRegionInputPtr,
+                    llvm::Value *&llvmReplacementValue) -> InsertPointTy {
+    llvmReplacementValue = &llvmOmpRegionInputPtr;
 
     // If this is a private value, this lambda will return the corresponding
     // mlir value and its `PrivateClauseOp`. Otherwise, empty values are
     // returned.
-    auto [privVar, privatizerClone] =
+    auto [mlirPrivVar, mlirPrivatizerClone] =
         [&]() -> std::pair<mlir::Value, omp::PrivateClauseOp> {
       if (!opInst.getPrivateVars().empty()) {
-        auto privateVars = opInst.getPrivateVars();
-        auto privateSyms = opInst.getPrivateSyms();
+        auto mlirPrivVars = opInst.getPrivateVars();
+        auto mlirPrivSyms = opInst.getPrivateSyms();
 
-        for (auto [privVar, privatizerAttr] :
-             llvm::zip_equal(privateVars, *privateSyms)) {
+        // Try to find a privatizer that corresponds to the LLVM value being
+        // prvatized.
+        for (auto [mlirPrivVar, mlirPrivatizerAttr] :
+             llvm::zip_equal(mlirPrivVars, *mlirPrivSyms)) {
           // Find the MLIR private variable corresponding to the LLVM value
           // being privatized.
-          llvm::Value *llvmPrivVar = moduleTranslation.lookupValue(privVar);
-          if (llvmPrivVar != &vPtr)
+          llvm::Value *mlirToLLVMPrivVar =
+              moduleTranslation.lookupValue(mlirPrivVar);
+
+          // Check if the LLVM value being privatized matches the LLVM value
+          // mapped to privVar. In some cases, this is not trivial ...
+          auto isMatch = [&]() {
+            if (mlirToLLVMPrivVar == nullptr)
+              return false;
+
+            // If both values are trivially equal, we found a match.
+            if (mlirToLLVMPrivVar == &llvmOmpRegionInputPtr)
+              return true;
+
+            // Otherwise, we check if both llvmOmpRegionInputPtr and
+            // mlirToLLVMPrivVar refer to the same memory (through a load/store
+            // pair). This happens if a struct (i.e. multi-field value) is being
+            // privatized.
+            //
+            // For example, if the privatized value is defined by:
+            // ```
+            //   %priv_val = alloca { ptr, i64 }, align 8
+            // ```
+            //
+            // Initialized this value (outside the omp region) will be something
----------------
luporl wrote:

```suggestion
            // The initialization of this value (outside the omp region) will be something
```

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