[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
+            // like this:
+            //
+            // clang-format off
+            // ```
+            //   %partially_init_priv_val = insertvalue { ptr, i64 } undef,
+            //                              ptr %some_ptr, 0
+            //   %fully_init_priv_val = insertvalue { ptr, i64 } %partially_init_priv_val,
+            //                          i64 %some_i64, 1
+            //   ...
+            //   store { ptr, i64 } %fully_init_priv_val, ptr %priv_val, align 8
+            // ```
+            // clang-format on
+            //
+            // Now, we enter the OMP region, in order to access this privatized
+            // value, we need to load from the allocated memory:
+            // ```
+            // omp.par.entry:
+            //   %priv_val_load = load { ptr, i64 }, ptr %priv_val, align 8
+            // ```
+            //
+            // The 2 LLVM values tracked here map as follows:
+            // - `mlirToLLVMPrivVar`     -> `%fully_init_priv_val`
+            // - `llvmOmpRegionInputPtr` -> `priv_val_load`
----------------
luporl wrote:

```suggestion
            // - `llvmOmpRegionInputPtr` -> `%priv_val_load`
```

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