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

[Kareem Ergawy]

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

Fixes #102935

Updates matching logic for finding the LLVM value that corresponds to an MLIR value. We need that matching to find the delayed privatizer for an LLVM value being privatized.

The issue occures when there is an "indirect" correspondence between MLIR and LLVM values: in some cases the values we are trying to match stem from a pair of load/store ops that point to the same memref. This PR adds such matching logic.

>From c7eb7d375eb4be713331e49488dec07d60fb52f4 Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Wed, 14 Aug 2024 03:40:59 -0500
Subject: [PATCH] [MLIR][OpenMP] Fix MLIR->LLVM value matching in privatization
 logic

Fixes #102935

Updates matching logic for finding the LLVM value that corresponds to
an MLIR value. We need that matching to find the delayed privatizer for
an LLVM value being privatized.

The issue occures when there is an "indirect" correspondence between
MLIR and LLVM values: in some cases the values we are trying to match
stem from a pair of load/store ops that point to the same memref. This
PR adds such matching logic.
---
 .../OpenMP/OpenMPToLLVMIRTranslation.cpp      | 33 ++++++++++++++-
 .../Target/LLVMIR/openmp-firstprivate.mlir    | 42 +++++++++++++++++++
 2 files changed, 74 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 458d05d5059db7..75133ec82cc9f0 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -1444,12 +1444,43 @@ convertOmpParallel(omp::ParallelOp opInst, llvm::IRBuilderBase &builder,
         auto privateVars = opInst.getPrivateVars();
         auto privateSyms = opInst.getPrivateSyms();
 
+        // Try to find a privatizer that corresponds to the LLVM value being
+        // prvatized.
         for (auto [privVar, privatizerAttr] :
              llvm::zip_equal(privateVars, *privateSyms)) {
           // Find the MLIR private variable corresponding to the LLVM value
           // being privatized.
           llvm::Value *llvmPrivVar = moduleTranslation.lookupValue(privVar);
-          if (llvmPrivVar != &vPtr)
+
+          // Check if the LLVM value being privatized matches the LLVM value
+          // mapped to privVar. In some cases, this is not trivial ...
+          auto isMatch = [](llvm::Value *vPtr, llvm::Value *llvmPrivVar) {
+            // If both values are trivially equal, we found a match.
+            if (llvmPrivVar == nullptr)
+              return false;
+
+            if (llvmPrivVar == vPtr)
+              return true;
+
+            auto *vPtrLoad = llvm::dyn_cast_if_present<llvm::LoadInst>(vPtr);
+
+            if (vPtrLoad == nullptr)
+              return false;
+
+            // Otherwise, we check if both vPtr and llvmPrivVar refer to the
+            // same memory (through a load/store pair).
+            for (auto &use : llvmPrivVar->uses()) {
+              auto *llvmPrivVarStore =
+                  llvm::dyn_cast_if_present<llvm::StoreInst>(use.getUser());
+              if (llvmPrivVarStore && (vPtrLoad->getPointerOperand() ==
+                                       llvmPrivVarStore->getPointerOperand()))
+                return true;
+            }
+
+            return false;
+          };
+
+          if (!isMatch(&vPtr, llvmPrivVar))
             continue;
 
           SymbolRefAttr privSym = llvm::cast<SymbolRefAttr>(privatizerAttr);
diff --git a/mlir/test/Target/LLVMIR/openmp-firstprivate.mlir b/mlir/test/Target/LLVMIR/openmp-firstprivate.mlir
index 65ae98b2a74c6e..b06ad96f4592c5 100644
--- a/mlir/test/Target/LLVMIR/openmp-firstprivate.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-firstprivate.mlir
@@ -114,3 +114,45 @@ omp.private {type = firstprivate} @multi_block.privatizer : !llvm.ptr alloc {
   llvm.store %arg2, %arg3 : f32, !llvm.ptr
   omp.yield(%arg3 : !llvm.ptr)
 }
+
+// -----
+
+// Verifies fix for https://github.com/llvm/llvm-project/issues/102935.
+//
+// The issue happens since we previously failed to match MLIR values to their
+// corresponding LLVM values in some cases (e.g. char strings with non-const
+// length).
+llvm.func @non_const_len_char_test(%n: !llvm.ptr {fir.bindc_name = "n"}) {
+  %n_val = llvm.load %n : !llvm.ptr -> i64
+  %orig_alloc = llvm.alloca %n_val x i8 {bindc_name = "str"} : (i64) -> !llvm.ptr
+  %orig_val = llvm.mlir.undef : !llvm.struct<(ptr, i64)>
+  %orig_val_init = llvm.insertvalue %orig_alloc, %orig_val[0] : !llvm.struct<(ptr, i64)>
+  omp.parallel private(@non_const_len_char %orig_val_init -> %priv_arg : !llvm.struct<(ptr, i64)>) {
+    %dummy = llvm.extractvalue %priv_arg[0] : !llvm.struct<(ptr, i64)>
+    omp.terminator
+  }
+  llvm.return
+}
+
+omp.private {type = firstprivate} @non_const_len_char : !llvm.struct<(ptr, i64)> alloc {
+^bb0(%orig_val: !llvm.struct<(ptr, i64)>):
+  %str_len = llvm.extractvalue %orig_val[1] : !llvm.struct<(ptr, i64)>
+  %priv_alloc = llvm.alloca %str_len x i8 {bindc_name = "str", pinned} : (i64) -> !llvm.ptr
+  %priv_val = llvm.mlir.undef : !llvm.struct<(ptr, i64)>
+  %priv_val_init = llvm.insertvalue %priv_alloc, %priv_val[0] : !llvm.struct<(ptr, i64)>
+  omp.yield(%priv_val_init : !llvm.struct<(ptr, i64)>)
+} copy {
+^bb0(%orig_val: !llvm.struct<(ptr, i64)>, %priv_val: !llvm.struct<(ptr, i64)>):
+  llvm.call @foo() : () -> ()
+  omp.yield(%priv_val : !llvm.struct<(ptr, i64)>)
+}
+
+llvm.func @foo()
+
+// CHECK-LABEL: @non_const_len_char_test..omp_par({{.*}})
+// CHECK-NEXT:    omp.par.entry:
+// Verify that we found the privatizer by checking that we properly inlined the
+// bodies of the alloc and copy regions.
+// CHECK:         %[[STR_LEN:.*]] = extractvalue { ptr, i64 } %{{.*}}, 1
+// CHECK:         %{{.*}} = alloca i8, i64 %[[STR_LEN]], align 1
+// CHECK:         call void @foo()