[mlir][OpenMP] Added omp.atomic.write lowering to LLVM IR

[shraiysh]

Summary

This patch adds omp.atomic.write lowering to LLVM IR.
Also, changed the syntax to have equal symbol instead of the comma to
make it more intuitive.

[kiranchandramohan]

Look OK. Have a few nit comments.

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

926

Nit: Use similar case style for AO (we have used ao above).

929

Nit: Can you add comments for the two false fields?

mlir/test/Target/LLVMIR/openmp-llvm.mlir

755

Nit: Remove the debug check.

[shraiysh]

Thanks for the review @kiranchandramohan. Addressed comments.

[kiranchandramohan]

LGTM.

[peixin]

LGTM.

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

931

forget to change this AO to ao?

mlir/test/Target/LLVMIR/openmp-llvm.mlir

756

What will be generated in IR for hint clause?

[shraiysh]

Thanks for the review @kiranchandramohan and @peixin. Addressed comments and rebase with main.

[shraiysh]

mlir/test/Target/LLVMIR/openmp-llvm.mlir

756

Hint clause is ignored at the moment - removed the testcase.

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

 /// value ::= operand `:` type
 static ParseResult parseAtomicWriteOp(OpAsmParser &parser,
                                       OperationState &result) {
   OpAsmParser::OperandType address, value;
   Type addrType, valueType;
   SmallVector<ClauseType> clauses = {memoryOrderClause, hintClause};
   SmallVector<int> segments;
 
-  if (parser.parseOperand(address) || parser.parseComma() ||
+  if (parser.parseOperand(address) || parser.parseEqual() ||
       parser.parseOperand(value) ||
       parseClauses(parser, result, clauses, segments) ||
       parser.parseColonType(addrType) || parser.parseComma() ||
       parser.parseType(valueType) ||
       parser.resolveOperand(address, addrType, result.operands) ||
       parser.resolveOperand(value, valueType, result.operands))
     return failure();
   return success();
 }
 
 /// Printer for AtomicWriteOp
 static void printAtomicWriteOp(OpAsmPrinter &p, AtomicWriteOp op) {
-  p << " " << op.address() << ", " << op.value() << " ";
+  p << " " << op.address() << " = " << op.value() << " ";
   if (op.memory_order())
     p << "memory_order(" << op.memory_order() << ") ";
   if (op.hintAttr())
     printSynchronizationHint(p, op, op.hintAttr());
   p << ": " << op.address().getType() << ", " << op.value().getType();
 }
 
 /// Verifier for AtomicWriteOp

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

convertOmpAtomicRead(Operation &opInst, llvm::IRBuilderBase &builder,
   // Restore the IP and insert Atomic Read.
   builder.restoreIP(currentIP);
   llvm::OpenMPIRBuilder::AtomicOpValue atomicV = {v, false, false};
   llvm::OpenMPIRBuilder::AtomicOpValue x = {address, false, false};
   builder.restoreIP(ompBuilder->createAtomicRead(ompLoc, x, atomicV, ao));
   return success();
 }
 
+/// Converts an omp.atomic.write operation to LLVM IR.
+static LogicalResult
+convertOmpAtomicWrite(Operation &opInst, llvm::IRBuilderBase &builder,
+                      LLVM::ModuleTranslation &moduleTranslation) {
+  auto writeOp = cast<omp::AtomicWriteOp>(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(writeOp.memory_order());
+  llvm::Value *expr = moduleTranslation.lookupValue(writeOp.value());
+  llvm::Value *dest = moduleTranslation.lookupValue(writeOp.address());
+  llvm::OpenMPIRBuilder::AtomicOpValue x = {dest, /*isSigned=*/false,
+                                            /*isVolatile=*/false};
+  builder.restoreIP(ompBuilder->createAtomicWrite(ompLoc, x, expr, 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
 /// reduction within WsLoopOp, but can be easily extended.
 static LogicalResult
 convertOmpReductionOp(omp::ReductionOp reductionOp,
                       llvm::IRBuilderBase &builder,
                       LLVM::ModuleTranslation &moduleTranslation) {
return llvm::TypeSwitch<Operation *, LogicalResult>(op)
         return convertOmpOrdered(*op, builder, moduleTranslation);
       })
       .Case([&](omp::WsLoopOp) {
         return convertOmpWsLoop(*op, builder, moduleTranslation);
       })
       .Case([&](omp::AtomicReadOp) {
         return convertOmpAtomicRead(*op, builder, moduleTranslation);
       })
+      .Case([&](omp::AtomicWriteOp) {
+        return convertOmpAtomicWrite(*op, builder, moduleTranslation);
+      })
       .Case([&](omp::SectionsOp) {
         return convertOmpSections(*op, builder, moduleTranslation);
       })
       .Case<omp::YieldOp, omp::TerminatorOp, omp::ReductionDeclareOp,
             omp::CriticalDeclareOp>([](auto op) {
         // `yield` and `terminator` can be just omitted. The block structure
         // was created in the region that handles their parent operation.
         // `reduction.declare` will be used by reductions and is not

mlir/test/Dialect/OpenMP/invalid.mlir

func @omp_atomic_read6(%addr : memref<i32>) {
   %1 = omp.atomic.read  %addr hint(speculative) hint(contended) : memref<i32> -> i32
   return
 }
 
 // -----
 
 func @omp_atomic_write1(%addr : memref<i32>, %val : i32) {
   // expected-error @below {{the hints omp_sync_hint_uncontended and omp_sync_hint_contended cannot be combined}}
-  omp.atomic.write  %addr, %val hint(contended, uncontended) : memref<i32>, i32
+  omp.atomic.write  %addr = %val hint(contended, uncontended) : memref<i32>, i32
   return
 }
 
 // -----
 
 func @omp_atomic_write2(%addr : memref<i32>, %val : i32) {
   // expected-error @below {{memory-order must not be acq_rel or acquire for atomic writes}}
-  omp.atomic.write  %addr, %val memory_order(acq_rel) : memref<i32>, i32
+  omp.atomic.write  %addr = %val memory_order(acq_rel) : memref<i32>, i32
   return
 }
 
 // -----
 
 func @omp_atomic_write3(%addr : memref<i32>, %val : i32) {
   // expected-error @below {{memory-order must not be acq_rel or acquire for atomic writes}}
-  omp.atomic.write  %addr, %val memory_order(acquire) : memref<i32>, i32
+  omp.atomic.write  %addr = %val memory_order(acquire) : memref<i32>, i32
   return
 }
 
 // -----
 
 func @omp_atomic_write4(%addr : memref<i32>, %val : i32) {
   // expected-error @below {{at most one memory_order clause can appear on the omp.atomic.write operation}}
-  omp.atomic.write  %addr, %val memory_order(release) memory_order(seq_cst) : memref<i32>, i32
+  omp.atomic.write  %addr = %val memory_order(release) memory_order(seq_cst) : memref<i32>, i32
   return
 }
 
 // -----
 
 func @omp_atomic_write5(%addr : memref<i32>, %val : i32) {
   // expected-error @below {{at most one hint clause can appear on the omp.atomic.write operation}}
-  omp.atomic.write  %addr, %val hint(contended) hint(speculative) : memref<i32>, i32
+  omp.atomic.write  %addr = %val hint(contended) hint(speculative) : memref<i32>, i32
   return
 }
 
 // -----
 
 func @omp_atomic_write6(%addr : memref<i32>, %val : i32) {
   // expected-error @below {{attribute 'memory_order' failed to satisfy constraint: MemoryOrderKind Clause}}
-  omp.atomic.write  %addr, %val memory_order(xyz) : memref<i32>, i32
+  omp.atomic.write  %addr = %val memory_order(xyz) : memref<i32>, i32
   return
 }
 
 // -----
 
 func @omp_atomic_update1(%x: memref<i32>, %expr: i32, %foo: memref<i32>) {
   // expected-error @below {{atomic update variable %x not found in the RHS of the assignment statement in an atomic.update operation}}
   omp.atomic.update %x = %foo add %expr : memref<i32>, i32

mlir/test/Dialect/OpenMP/ops.mlir

func @omp_atomic_read(%addr : memref<i32>) {
   // CHECK: %{{.*}} = omp.atomic.read %[[ADDR]] memory_order(seq_cst) hint(contended, speculative) : memref<i32> -> i32
   %8 = omp.atomic.read %addr hint(speculative, contended) memory_order(seq_cst) : memref<i32> -> i32
   return
 }
 
 // CHECK-LABEL: omp_atomic_write
 // CHECK-SAME: (%[[ADDR:.*]]: memref<i32>, %[[VAL:.*]]: i32)
 func @omp_atomic_write(%addr : memref<i32>, %val : i32) {
-  // CHECK: omp.atomic.write %[[ADDR]], %[[VAL]] : memref<i32>, i32
-  omp.atomic.write %addr, %val : memref<i32>, i32
-  // CHECK: omp.atomic.write %[[ADDR]], %[[VAL]] memory_order(seq_cst) : memref<i32>, i32
-  omp.atomic.write %addr, %val memory_order(seq_cst) : memref<i32>, i32
-  // CHECK: omp.atomic.write %[[ADDR]], %[[VAL]] memory_order(release) : memref<i32>, i32
-  omp.atomic.write %addr, %val memory_order(release) : memref<i32>, i32
-  // CHECK: omp.atomic.write %[[ADDR]], %[[VAL]] memory_order(relaxed) : memref<i32>, i32
-  omp.atomic.write %addr, %val memory_order(relaxed) : memref<i32>, i32
-  // CHECK: omp.atomic.write %[[ADDR]], %[[VAL]] hint(uncontended, speculative) : memref<i32>, i32
-  omp.atomic.write %addr, %val hint(speculative, uncontended) : memref<i32>, i32
+  // CHECK: omp.atomic.write %[[ADDR]] = %[[VAL]] : memref<i32>, i32
+  omp.atomic.write %addr = %val : memref<i32>, i32
+  // CHECK: omp.atomic.write %[[ADDR]] = %[[VAL]] memory_order(seq_cst) : memref<i32>, i32
+  omp.atomic.write %addr = %val memory_order(seq_cst) : memref<i32>, i32
+  // CHECK: omp.atomic.write %[[ADDR]] = %[[VAL]] memory_order(release) : memref<i32>, i32
+  omp.atomic.write %addr = %val memory_order(release) : memref<i32>, i32
+  // CHECK: omp.atomic.write %[[ADDR]] = %[[VAL]] memory_order(relaxed) : memref<i32>, i32
+  omp.atomic.write %addr = %val memory_order(relaxed) : memref<i32>, i32
+  // CHECK: omp.atomic.write %[[ADDR]] = %[[VAL]] hint(uncontended, speculative) : memref<i32>, i32
+  omp.atomic.write %addr = %val hint(speculative, uncontended) : memref<i32>, i32
   return
 }
 
 // CHECK-LABEL: omp_atomic_update
 // CHECK-SAME: (%[[X:.*]]: memref<i32>, %[[EXPR:.*]]: i32, %[[XBOOL:.*]]: memref<i1>, %[[EXPRBOOL:.*]]: i1)
 func @omp_atomic_update(%x : memref<i32>, %expr : i32, %xBool : memref<i1>, %exprBool : i1) {
   // CHECK: omp.atomic.update %[[X]] = %[[X]] add %[[EXPR]] : memref<i32>, i32
   omp.atomic.update %x = %x add %expr : memref<i32>, i32

mlir/test/Target/LLVMIR/openmp-llvm.mlir

llvm.func @omp_atomic_read(%arg0 : !llvm.ptr<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
 }
 
 // -----
 
+// CHECK-LABEL: @omp_atomic_write
+// CHECK-SAME: (i32* %[[x:.*]], i32 %[[expr:.*]])
+llvm.func @omp_atomic_write(%x: !llvm.ptr<i32>, %expr: i32) -> () {
+  // CHECK: store atomic i32 %[[expr]], i32* %[[x]] monotonic, align 4
+  omp.atomic.write %x = %expr : !llvm.ptr<i32>, i32
+  // CHECK: store atomic i32 %[[expr]], i32* %[[x]] seq_cst, align 4
+  // CHECK: call void @__kmpc_flush(%struct.ident_t* @{{.*}})
+  omp.atomic.write %x = %expr memory_order(seq_cst) : !llvm.ptr<i32>, i32
+  // CHECK: store atomic i32 %[[expr]], i32* %[[x]] release, align 4
+  // CHECK: call void @__kmpc_flush(%struct.ident_t* @{{.*}})
+  omp.atomic.write %x = %expr memory_order(release) : !llvm.ptr<i32>, i32
+  // CHECK: store atomic i32 %[[expr]], i32* %[[x]] monotonic, align 4
+  omp.atomic.write %x = %expr memory_order(relaxed) : !llvm.ptr<i32>, i32
+  llvm.return
+}
+
+// -----
+
 // CHECK-LABEL: @omp_sections_empty
 llvm.func @omp_sections_empty() -> () {
   omp.sections {
     omp.terminator
   }
   // CHECK-NEXT: ret void
   llvm.return
 }