[Mlir-commits] [flang] [mlir] [flang][OpenMP] Delayed privatization for variables with `equivalence` association (PR #100531)

Kareem Ergawy llvmlistbot at llvm.org
Thu Jul 25 02:05:12 PDT 2024


https://github.com/ergawy updated https://github.com/llvm/llvm-project/pull/100531

>From a8a797f65d32aba191aa645ff3d0d51b6e1501a7 Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Thu, 25 Jul 2024 02:45:30 -0500
Subject: [PATCH] [flang][OpenMP] Delayed privatization for variable with
 `equivalence` association

Handles variables that are storage associated via `equivalence`. The
problem is that these variables are declared as `fir.ptr`s while their
privatized storage is declared as `fir.ref` which was triggering a
validation error in the OpenMP dialect. So far, I just disabled a small
part of the validation logic until we find a way to properly check of
compatible types without leaking the FIR dialect specifics in the OpenMP
dialect validation logic.
---
 .../lib/Lower/OpenMP/DataSharingProcessor.cpp |  9 +++--
 .../DelayedPrivatization/equivalence.f90      | 36 +++++++++++++++++++
 mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp  |  6 +++-
 mlir/test/Dialect/OpenMP/invalid.mlir         | 15 ++++----
 mlir/test/Target/LLVMIR/openmp-private.mlir   | 35 ++++++++++++++++++
 5 files changed, 90 insertions(+), 11 deletions(-)
 create mode 100644 flang/test/Lower/OpenMP/DelayedPrivatization/equivalence.f90

diff --git a/flang/lib/Lower/OpenMP/DataSharingProcessor.cpp b/flang/lib/Lower/OpenMP/DataSharingProcessor.cpp
index 7e76a81e0df92..a8fb4c44f3614 100644
--- a/flang/lib/Lower/OpenMP/DataSharingProcessor.cpp
+++ b/flang/lib/Lower/OpenMP/DataSharingProcessor.cpp
@@ -523,6 +523,7 @@ void DataSharingProcessor::doPrivatize(const semantics::Symbol *sym,
     fir::ExtendedValue symExV = converter.getSymbolExtendedValue(*sym);
 
     symTable->pushScope();
+    mlir::Type yieldedType;
 
     // Populate the `alloc` region.
     {
@@ -542,9 +543,10 @@ void DataSharingProcessor::doPrivatize(const semantics::Symbol *sym,
       symTable->addSymbol(*sym, localExV);
       symTable->pushScope();
       cloneSymbol(sym);
-      firOpBuilder.create<mlir::omp::YieldOp>(
+      auto yieldOp = firOpBuilder.create<mlir::omp::YieldOp>(
           hsb.getAddr().getLoc(),
           symTable->shallowLookupSymbol(*sym).getAddr());
+      yieldedType = yieldOp.getOperand(0).getType();
       symTable->popScope();
     }
 
@@ -553,8 +555,9 @@ void DataSharingProcessor::doPrivatize(const semantics::Symbol *sym,
       mlir::Region &copyRegion = result.getCopyRegion();
       // First block argument corresponding to the original/host value while
       // second block argument corresponding to the privatized value.
-      mlir::Block *copyEntryBlock = firOpBuilder.createBlock(
-          &copyRegion, /*insertPt=*/{}, {symType, symType}, {symLoc, symLoc});
+      mlir::Block *copyEntryBlock =
+          firOpBuilder.createBlock(&copyRegion, /*insertPt=*/{},
+                                   {symType, yieldedType}, {symLoc, symLoc});
       firOpBuilder.setInsertionPointToEnd(copyEntryBlock);
 
       auto addSymbol = [&](unsigned argIdx, bool force = false) {
diff --git a/flang/test/Lower/OpenMP/DelayedPrivatization/equivalence.f90 b/flang/test/Lower/OpenMP/DelayedPrivatization/equivalence.f90
new file mode 100644
index 0000000000000..23a3b3f96e3e5
--- /dev/null
+++ b/flang/test/Lower/OpenMP/DelayedPrivatization/equivalence.f90
@@ -0,0 +1,36 @@
+! Test delayed privatization for variables that are storage associated via `EQUIVALENCE`.
+
+! RUN: %flang_fc1 -emit-hlfir -fopenmp -mmlir --openmp-enable-delayed-privatization \
+! RUN:   -o - %s 2>&1 | FileCheck %s
+! RUN: bbc -emit-hlfir -fopenmp --openmp-enable-delayed-privatization -o - %s 2>&1 \
+! RUN:   | FileCheck %s
+
+subroutine private_common
+  real x, y
+  equivalence (x,y)
+  !$omp parallel firstprivate(x)
+    x = 3.14
+  !$omp end parallel
+end subroutine
+
+! CHECK:  omp.private {type = firstprivate} @[[X_PRIVATIZER:.*]] : ![[X_TYPE:fir.ptr<f32>]] alloc {
+! CHECK:  ^bb0(%{{.*}}: ![[X_TYPE]]):
+! CHECK:    %[[PRIV_ALLOC:.*]] = fir.alloca f32 {bindc_name = "x", {{.*}}}
+! CHECK:    %[[PRIV_DECL:.*]]:2 = hlfir.declare %[[PRIV_ALLOC]] {{{.*}}} : (![[PRIV_TYPE:fir.ref<f32>]]) -> ({{.*}})
+! CHECK:    omp.yield(%[[PRIV_DECL]]#0 : ![[PRIV_TYPE]])
+! CHECK:  } copy {
+! CHECK:  ^bb0(%[[ORIG_PTR:.*]]: ![[X_TYPE]], %[[PRIV_REF:.*]]: ![[PRIV_TYPE]]):
+! CHECK:    %[[ORIG_VAL:.*]] = fir.load %[[ORIG_PTR]] : !fir.ptr<f32>
+! CHECK:    hlfir.assign %[[ORIG_VAL]] to %[[PRIV_REF]] temporary_lhs : f32, ![[PRIV_TYPE]]
+! CHECK:    omp.yield(%[[PRIV_REF]] : ![[PRIV_TYPE]])
+! CHECK:  }
+
+! CHECK:  func.func @_QPprivate_common() {
+! CHECK:    omp.parallel private(@[[X_PRIVATIZER]] %{{.*}}#0 -> %[[PRIV_ARG:.*]] : ![[X_TYPE]]) {
+! CHECK:      %[[REG_DECL:.*]]:2 = hlfir.declare %[[PRIV_ARG]] {{{.*}}} : (![[X_TYPE]]) -> ({{.*}})
+! CHECK:      %[[CST:.*]] = arith.constant {{.*}}
+! CHECK:      hlfir.assign %[[CST]] to %[[REG_DECL]]#0 : {{.*}}
+! CHECK:      omp.terminator
+! CHECK:    }
+! CHECK:    return
+! CHECK:  }
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index f5ec5a476ad8f..ffcd3b11a12f3 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -2505,7 +2505,11 @@ LogicalResult PrivateClauseOp::verify() {
              << "Did not expect any values to be yielded.";
     }
 
-    if (yieldedTypes.size() == 1 && yieldedTypes.front() == symType)
+    // TODO How can we check that 2 types are compatible without leaking
+    // dialect-specific (e.g. FIR) information? The problem is that in some
+    // cases, we would get the input type of the symbol as, e.g. `fir.ptr<i32>`
+    // while the allocated private memory is of type `fir.ref<i32>`.
+    if (yieldedTypes.size() == 1 /*&& yieldedTypes.front() == symType*/)
       return success();
 
     auto error = mlir::emitError(yieldOp.getLoc())
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index 9977dd57e3023..d511c063c8d9c 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -2207,13 +2207,14 @@ func.func @omp_distribute_unconstrained_order() -> () {
   }
   return
 }
-// -----
-omp.private {type = private} @x.privatizer : i32 alloc {
-^bb0(%arg0: i32):
-  %0 = arith.constant 0.0 : f32
-  // expected-error @below {{Invalid yielded value. Expected type: 'i32', got: 'f32'}}
-  omp.yield(%0 : f32)
-}
+
+// Disabled for now. See relevant TODO in PrivateClauseOp::verify().
+//omp.private {type = private} @x.privatizer : i32 alloc {
+//^bb0(%arg0: i32):
+//  %0 = arith.constant 0.0 : f32
+//  // _expected-error_ @below {{Invalid yielded value. Expected type: 'i32', got: 'f32'}}
+//  omp.yield(%0 : f32)
+//}
 
 // -----
 
diff --git a/mlir/test/Target/LLVMIR/openmp-private.mlir b/mlir/test/Target/LLVMIR/openmp-private.mlir
index 5924377581db6..6bc73fc6f102b 100644
--- a/mlir/test/Target/LLVMIR/openmp-private.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-private.mlir
@@ -234,3 +234,38 @@ omp.declare_reduction @reducer.part : !llvm.ptr init {
 ^bb0(%arg0: !llvm.ptr):
   omp.yield
 }
+
+// -----
+
+llvm.func @_QPequivalence() {
+  %0 = llvm.mlir.constant(1 : i64) : i64
+  %1 = llvm.alloca %0 x !llvm.array<4 x i8> : (i64) -> !llvm.ptr
+  %2 = llvm.mlir.constant(0 : index) : i64
+  %3 = llvm.getelementptr %1[0, %2] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.array<4 x i8>
+  omp.parallel private(@_QFequivalenceEx_firstprivate_ptr_f32 %3 -> %arg0 : !llvm.ptr) {
+    %4 = llvm.mlir.constant(3.140000e+00 : f32) : f32
+    llvm.store %4, %arg0 : f32, !llvm.ptr
+    omp.terminator
+  }
+  llvm.return
+}
+
+omp.private {type = firstprivate} @_QFequivalenceEx_firstprivate_ptr_f32 : !llvm.ptr alloc {
+^bb0(%arg0: !llvm.ptr):
+  %0 = llvm.mlir.constant(1 : i64) : i64
+  %1 = llvm.alloca %0 x f32 {bindc_name = "x", pinned} : (i64) -> !llvm.ptr
+  omp.yield(%1 : !llvm.ptr)
+} copy {
+^bb0(%arg0: !llvm.ptr, %arg1: !llvm.ptr):
+  %0 = llvm.load %arg0 : !llvm.ptr -> f32
+  llvm.store %0, %arg1 : f32, !llvm.ptr
+  omp.yield(%arg1 : !llvm.ptr)
+}
+
+// CHECK: define internal void @_QPequivalence..omp_par
+// CHECK:   %[[PRIV_ALLOC:.*]] = alloca float, i64 1, align 4
+// CHECK:   %[[HOST_VAL:.*]] = load float, ptr %{{.*}}, align 4
+// Test that we initialzie the firstprivate variable.
+// CHECK:   store float %[[HOST_VAL]], ptr %[[PRIV_ALLOC]], align 4
+// Test that we inlined the body of the parallel region.
+// CHECK:   store float 0x{{.*}}, ptr %[[PRIV_ALLOC]], align 4



More information about the Mlir-commits mailing list