[Mlir-commits] [mlir] 31cf42b - [mlir][OpenMP] Added omp.atomic.read lowering

[Shraiysh Vaishay]

Author: Shraiysh Vaishay
Date: 2021-12-07T11:17:30+05:30
New Revision: 31cf42bd9a71c53c7ddc6f0e3ec41fcea24f2f7c

URL: https://github.com/llvm/llvm-project/commit/31cf42bd9a71c53c7ddc6f0e3ec41fcea24f2f7c
DIFF: https://github.com/llvm/llvm-project/commit/31cf42bd9a71c53c7ddc6f0e3ec41fcea24f2f7c.diff

LOG: [mlir][OpenMP] Added omp.atomic.read lowering

This patch adds lowering from omp.atomic.read to LLVM IR along with the
memory ordering clause. Tests for the same are also added.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D115134

Added: 
    

Modified: 
    mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
    mlir/test/Target/LLVMIR/openmp-llvm.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 38f93f7faf920..d623dbb0c3cf8 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -786,6 +786,55 @@ convertOmpWsLoop(Operation &opInst, llvm::IRBuilderBase &builder,
   return success();
 }
 
+// Convert an Atomic Ordering attribute to llvm::AtomicOrdering.
+llvm::AtomicOrdering convertAtomicOrdering(Optional<StringRef> AOAttr) {
+  if (!AOAttr.hasValue())
+    return llvm::AtomicOrdering::Monotonic; // Default Memory Ordering
+
+  return StringSwitch<llvm::AtomicOrdering>(AOAttr.getValue())
+      .Case("seq_cst", llvm::AtomicOrdering::SequentiallyConsistent)
+      .Case("acq_rel", llvm::AtomicOrdering::AcquireRelease)
+      .Case("acquire", llvm::AtomicOrdering::Acquire)
+      .Case("release", llvm::AtomicOrdering::Release)
+      .Case("relaxed", llvm::AtomicOrdering::Monotonic)
+      .Default(llvm::AtomicOrdering::Monotonic);
+}
+
+// Convert omp.atomic.read operation to LLVM IR.
+static LogicalResult
+convertOmpAtomicRead(Operation &opInst, llvm::IRBuilderBase &builder,
+                     LLVM::ModuleTranslation &moduleTranslation) {
+
+  auto readOp = cast<omp::AtomicReadOp>(opInst);
+  llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
+
+  // Set up the source location value for OpenMP runtime.
+  llvm::DISubprogram *subprogram =
+      builder.GetInsertBlock()->getParent()->getSubprogram();
+  const llvm::DILocation *diLoc =
+      moduleTranslation.translateLoc(opInst.getLoc(), subprogram);
+  llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder.saveIP(),
+                                                    llvm::DebugLoc(diLoc));
+  llvm::AtomicOrdering AO = convertAtomicOrdering(readOp.memory_order());
+  llvm::Value *address = moduleTranslation.lookupValue(readOp.address());
+  llvm::OpenMPIRBuilder::InsertPointTy currentIP = builder.saveIP();
+
+  // Insert alloca for result.
+  llvm::OpenMPIRBuilder::InsertPointTy allocaIP =
+      findAllocaInsertPoint(builder, moduleTranslation);
+  builder.restoreIP(allocaIP);
+  llvm::Value *v = builder.CreateAlloca(
+      moduleTranslation.convertType(readOp.getResult().getType()));
+  moduleTranslation.mapValue(readOp.getResult(), v);
+
+  // Restore the IP and insert Atomic Read.
+  builder.restoreIP(currentIP);
+  llvm::OpenMPIRBuilder::AtomicOpValue V = {v, false, false};
+  llvm::OpenMPIRBuilder::AtomicOpValue X = {address, false, false};
+  builder.restoreIP(ompBuilder->createAtomicRead(ompLoc, X, V, AO));
+  return success();
+}
+
 /// Converts an OpenMP reduction operation using OpenMPIRBuilder. Expects the
 /// mapping between reduction variables and their private equivalents to have
 /// been stored on the ModuleTranslation stack. Currently only supports
@@ -907,6 +956,9 @@ LogicalResult OpenMPDialectLLVMIRTranslationInterface::convertOperation(
       .Case([&](omp::WsLoopOp) {
         return convertOmpWsLoop(*op, builder, moduleTranslation);
       })
+      .Case([&](omp::AtomicReadOp) {
+        return convertOmpAtomicRead(*op, builder, moduleTranslation);
+      })
       .Case<omp::YieldOp, omp::TerminatorOp, omp::ReductionDeclareOp,
             omp::CriticalDeclareOp>([](auto op) {
         // `yield` and `terminator` can be just omitted. The block structure

diff  --git a/mlir/test/Target/LLVMIR/openmp-llvm.mlir b/mlir/test/Target/LLVMIR/openmp-llvm.mlir
index d954fd180bb2e..77fd518c49cb3 100644
--- a/mlir/test/Target/LLVMIR/openmp-llvm.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-llvm.mlir
@@ -706,3 +706,33 @@ llvm.func @omp_ordered(%arg0 : i32, %arg1 : i32, %arg2 : i32, %arg3 : i64,
 
   llvm.return
 }
+
+// -----
+
+// CHECK-LABEL: @omp_atomic_read
+// CHECK-SAME: (i32* %[[ARG0:.*]])
+llvm.func @omp_atomic_read(%arg0 : !llvm.ptr<i32>) -> () {
+  // CHECK: %{{.*}} = alloca i32, align 4
+  // CHECK: %{{.*}} = alloca i32, align 4
+  // CHECK: %{{.*}} = alloca i32, align 4
+  // CHECK: %{{.*}} = alloca i32, align 4
+
+  // CHECK: %[[X1:.*]] = load atomic i32, i32* %[[ARG0]] monotonic, align 4
+  // CHECK: store i32 %[[X1]], i32* %{{.*}}, align 4
+  %x1 = omp.atomic.read %arg0 : !llvm.ptr<i32> -> i32
+
+  // CHECK: %[[X2:.*]] = load atomic i32, i32* %[[ARG0]] seq_cst, align 4
+  // CHECK: call void @__kmpc_flush(%{{.*}})
+  // CHECK: store i32 %[[X2]], i32* %{{.*}}, align 4
+  %x2 = omp.atomic.read %arg0 memory_order(seq_cst) : !llvm.ptr<i32> -> i32
+
+  // CHECK: %[[X3:.*]] = load atomic i32, i32* %[[ARG0]] acquire, align 4
+  // CHECK: call void @__kmpc_flush(%{{.*}})
+  // CHECK: store i32 %[[X3]], i32* %{{.*}}, align 4
+  %x3 = omp.atomic.read %arg0 memory_order(acquire) : !llvm.ptr<i32> -> i32
+
+  // CHECK: %[[X4:.*]] = load atomic i32, i32* %[[ARG0]] monotonic, align 4
+  // CHECK: store i32 %[[X4]], i32* %{{.*}}, align 4
+  %x4 = omp.atomic.read %arg0 memory_order(relaxed) : !llvm.ptr<i32> -> i32
+  llvm.return
+}