[flang][flang-driver][mlir][OpenMP] atomic control support

clang/include/clang/Driver/Options.td

@@ -2303,21 +2303,21 @@ def fsymbol_partition_EQ : Joined<["-"], "fsymbol-partition=">, Group<f_Group>,
 
 defm atomic_remote_memory : BoolFOption<"atomic-remote-memory",
   LangOpts<"AtomicRemoteMemory">, DefaultFalse,
-  PosFlag<SetTrue, [], [ClangOption, CC1Option], "May have">,
-  NegFlag<SetFalse, [], [ClangOption], "Assume no">,
-  BothFlags<[], [ClangOption], " atomic operations on remote memory">>;
+  PosFlag<SetTrue, [], [ClangOption, CC1Option, FlangOption, FC1Option], "May have">,
+  NegFlag<SetFalse, [], [ClangOption, FlangOption], "Assume no">,
+  BothFlags<[], [ClangOption, FlangOption], " atomic operations on remote memory">>;
 
 defm atomic_fine_grained_memory : BoolFOption<"atomic-fine-grained-memory",
   LangOpts<"AtomicFineGrainedMemory">, DefaultFalse,
-  PosFlag<SetTrue, [], [ClangOption, CC1Option], "May have">,
-  NegFlag<SetFalse, [], [ClangOption], "Assume no">,
-  BothFlags<[], [ClangOption], " atomic operations on fine-grained memory">>;
+  PosFlag<SetTrue, [], [ClangOption, CC1Option, FlangOption, FC1Option], "May have">,
+  NegFlag<SetFalse, [], [ClangOption, FlangOption], "Assume no">,
+  BothFlags<[], [ClangOption, FlangOption], " atomic operations on fine-grained memory">>;
 
 defm atomic_ignore_denormal_mode : BoolFOption<"atomic-ignore-denormal-mode",
   LangOpts<"AtomicIgnoreDenormalMode">, DefaultFalse,
-  PosFlag<SetTrue, [], [ClangOption, CC1Option], "Allow">,
-  NegFlag<SetFalse, [], [ClangOption], "Disallow">,
-  BothFlags<[], [ClangOption], " atomic operations to ignore denormal mode">>;
+  PosFlag<SetTrue, [], [ClangOption, CC1Option, FlangOption, FC1Option], "Allow">,
+  NegFlag<SetFalse, [], [ClangOption, FlangOption], "Disallow">,
+  BothFlags<[], [ClangOption, FlangOption], " atomic operations to ignore denormal mode">>;
 
 defm memory_profile : OptInCC1FFlag<"memory-profile", "Enable", "Disable", " heap memory profiling">;
 def fmemory_profile_EQ : Joined<["-"], "fmemory-profile=">,
@@ -5314,9 +5314,9 @@ defm amdgpu_precise_memory_op
                   " precise memory mode (AMDGPU only)">;
 
 def munsafe_fp_atomics : Flag<["-"], "munsafe-fp-atomics">,
-  Visibility<[ClangOption, CC1Option]>, Alias<fatomic_ignore_denormal_mode>;
+  Visibility<[ClangOption, CC1Option, FlangOption, FC1Option]>, Alias<fatomic_ignore_denormal_mode>;
 def mno_unsafe_fp_atomics : Flag<["-"], "mno-unsafe-fp-atomics">,
-  Visibility<[ClangOption]>, Alias<fno_atomic_ignore_denormal_mode>;
+  Visibility<[ClangOption, FlangOption]>, Alias<fno_atomic_ignore_denormal_mode>;
 
 def faltivec : Flag<["-"], "faltivec">, Group<f_Group>;
 def fno_altivec : Flag<["-"], "fno-altivec">, Group<f_Group>;

flang/include/flang/Frontend/TargetOptions.h

@@ -53,6 +53,11 @@ class TargetOptions {
 
   /// Print verbose assembly
   bool asmVerbose = false;
+
+  /// Atomic control options
+  bool atomicIgnoreDenormalMode = false;
+  bool atomicRemoteMemory = false;
+  bool atomicFineGrainedMemory = false;
 };
 
 } // end namespace Fortran::frontend

flang/include/flang/Optimizer/Dialect/Support/FIRContext.h

@@ -58,6 +58,22 @@ void setTargetCPU(mlir::ModuleOp mod, llvm::StringRef cpu);
 /// Get the target CPU string from the Module or return a null reference.
 llvm::StringRef getTargetCPU(mlir::ModuleOp mod);
 
+/// Sets that Denormal Mode can be ignored for lowering of floating point atomic
+/// operations.
+void setAtomicIgnoreDenormalMode(mlir::ModuleOp mod, bool value);
+/// Gets whether Denormal Mode can be ignored or not for lowering of floating
+/// point atomic operations.
+bool getAtomicIgnoreDenormalMode(mlir::ModuleOp mod);
+/// Sets that fine grained memory is used for lowering of atomic operations.
+void setAtomicFineGrainedMemory(mlir::ModuleOp mod, bool value);
+/// Gets whether fine grained memory is used or not for lowering of atomic
+/// operations.
+bool getAtomicFineGrainedMemory(mlir::ModuleOp mod);
+/// Sets that remote memory is used for lowering of atomic operations.
+void setAtomicRemoteMemory(mlir::ModuleOp mod, bool value);
+/// Gets whether remote memory is used or not for lowering of atomic operations.
+bool getAtomicRemoteMemory(mlir::ModuleOp mod);
+
 /// Set the tune CPU for the module. `cpu` must not be deallocated while
 /// module `mod` is still live.
 void setTuneCPU(mlir::ModuleOp mod, llvm::StringRef cpu);

flang/lib/Frontend/CompilerInvocation.cpp

@@ -496,6 +496,16 @@ static void parseTargetArgs(TargetOptions &opts, llvm::opt::ArgList &args) {
           args.getLastArg(clang::driver::options::OPT_triple))
     opts.triple = a->getValue();
 
+  opts.atomicIgnoreDenormalMode = args.hasFlag(
+      clang::driver::options::OPT_fatomic_ignore_denormal_mode,
+      clang::driver::options::OPT_fno_atomic_ignore_denormal_mode, false);
+  opts.atomicFineGrainedMemory = args.hasFlag(
+      clang::driver::options::OPT_fatomic_fine_grained_memory,
+      clang::driver::options::OPT_fno_atomic_fine_grained_memory, false);
+  opts.atomicRemoteMemory =
+      args.hasFlag(clang::driver::options::OPT_fatomic_remote_memory,
+                   clang::driver::options::OPT_fno_atomic_remote_memory, false);
+
   if (const llvm::opt::Arg *a =
           args.getLastArg(clang::driver::options::OPT_target_cpu))
     opts.cpu = a->getValue();

flang/lib/Lower/Bridge.cpp

@@ -6666,6 +6666,10 @@ Fortran::lower::LoweringBridge::LoweringBridge(
   fir::setKindMapping(*module, kindMap);
   fir::setTargetCPU(*module, targetMachine.getTargetCPU());
   fir::setTuneCPU(*module, targetOpts.cpuToTuneFor);
+  fir::setAtomicIgnoreDenormalMode(*module,
+                                   targetOpts.atomicIgnoreDenormalMode);
+  fir::setAtomicFineGrainedMemory(*module, targetOpts.atomicFineGrainedMemory);
+  fir::setAtomicRemoteMemory(*module, targetOpts.atomicRemoteMemory);
   fir::setTargetFeatures(*module, targetMachine.getTargetFeatureString());
   fir::support::setMLIRDataLayout(*module, targetMachine.createDataLayout());
   fir::setIdent(*module, Fortran::common::getFlangFullVersion());

flang/lib/Lower/OpenMP/Atomic.cpp

@@ -634,9 +634,16 @@ genAtomicUpdate(lower::AbstractConverter &converter,
     }
   }
 
+  mlir::ModuleOp module = builder.getModule();
+  mlir::omp::AtomicControlAttr atomicControlAttr =
+      mlir::omp::AtomicControlAttr::get(
+          builder.getContext(), fir::getAtomicIgnoreDenormalMode(module),
+          fir::getAtomicFineGrainedMemory(module),
+          fir::getAtomicRemoteMemory(module));
   builder.restoreInsertionPoint(atomicAt);
   auto updateOp = builder.create<mlir::omp::AtomicUpdateOp>(
-      loc, atomAddr, hint, makeMemOrderAttr(converter, memOrder));
+      loc, atomAddr, atomicControlAttr, hint,
+      makeMemOrderAttr(converter, memOrder));
 
   mlir::Region &region = updateOp->getRegion(0);
   mlir::Block *block = builder.createBlock(&region, {}, {atomType}, {loc});

flang/lib/Optimizer/Dialect/Support/FIRContext.cpp

@@ -88,6 +88,53 @@ void fir::setTuneCPU(mlir::ModuleOp mod, llvm::StringRef cpu) {
   mod->setAttr(tuneCpuName, mlir::StringAttr::get(ctx, cpu));
 }
 
+static constexpr const char *atomicIgnoreDenormalModeName =
+    "fir.atomic_ignore_denormal_mode";
+void fir::setAtomicIgnoreDenormalMode(mlir::ModuleOp mod, bool value) {
+  if (value) {
+    auto *ctx = mod.getContext();
+    mod->setAttr(atomicIgnoreDenormalModeName, mlir::UnitAttr::get(ctx));
+  } else {
+    if (mod->hasAttr(atomicIgnoreDenormalModeName))
+      mod->removeAttr(atomicIgnoreDenormalModeName);
+  }
+}
+
+bool fir::getAtomicIgnoreDenormalMode(mlir::ModuleOp mod) {
+  return mod->hasAttrOfType<mlir::UnitAttr>(atomicIgnoreDenormalModeName);
+}
+
+static constexpr const char *atomicFineGrainedMemoryName =
+    "fir.atomic_fine_grained_memory";
+void fir::setAtomicFineGrainedMemory(mlir::ModuleOp mod, bool value) {
+  if (value) {
+    auto *ctx = mod.getContext();
+    mod->setAttr(atomicFineGrainedMemoryName, mlir::UnitAttr::get(ctx));
+  } else {
+    if (mod->hasAttr(atomicFineGrainedMemoryName))
+      mod->removeAttr(atomicFineGrainedMemoryName);
+  }
+}
+
+bool fir::getAtomicFineGrainedMemory(mlir::ModuleOp mod) {
+  return mod->hasAttrOfType<mlir::UnitAttr>(atomicFineGrainedMemoryName);
+}
+static constexpr const char *atomicRemoteMemoryName =
+    "fir.atomic_remote_memory";
+void fir::setAtomicRemoteMemory(mlir::ModuleOp mod, bool value) {
+  if (value) {
+    auto *ctx = mod.getContext();
+    mod->setAttr(atomicRemoteMemoryName, mlir::UnitAttr::get(ctx));
+  } else {
+    if (mod->hasAttr(atomicRemoteMemoryName))
+      mod->removeAttr(atomicRemoteMemoryName);
+  }
+}
+
+bool fir::getAtomicRemoteMemory(mlir::ModuleOp mod) {
+  return mod->hasAttrOfType<mlir::UnitAttr>(atomicRemoteMemoryName);
+}
+
 llvm::StringRef fir::getTuneCPU(mlir::ModuleOp mod) {
   if (auto attr = mod->getAttrOfType<mlir::StringAttr>(tuneCpuName))
     return attr.getValue();

flang/test/Fir/convert-to-llvm-openmp-and-fir.fir

@@ -619,7 +619,7 @@ func.func @_QPs() {
     %c1_i32 = arith.constant 1 : i32
     %1 = arith.addi %arg0, %c1_i32 : i32
     omp.yield(%1 : i32)
-  }
+  } {atomic_control = #omp.atomic_control<>}
   return
 }
 fir.global internal @_QFsEc : i32 {
@@ -634,7 +634,7 @@ fir.global internal @_QFsEc : i32 {
 // CHECK:      %[[CONST_1:.*]] = llvm.mlir.constant(1 : i32) : i32
 // CHECK:      %[[OUT_VAL:.*]] = llvm.add %[[IN_VAL]], %[[CONST_1]]  : i32
 // CHECK:      omp.yield(%[[OUT_VAL]] : i32)
-// CHECK:    }
+// CHECK:    } {atomic_control = #omp.atomic_control<>}
 // CHECK:    llvm.return
 // CHECK:  }
 // CHECK:  llvm.mlir.global internal @[[GLOBAL]]() {{.*}} : i32 {

flang/test/Lower/OpenMP/atomic-control-options.f90

@@ -0,0 +1,37 @@
+! RUN: %flang_fc1 -emit-hlfir -triple amdgcn-amd-amdhsa -fopenmp -fopenmp-is-device -munsafe-fp-atomics %s -o - | FileCheck -check-prefix=UNSAFE-FP-ATOMICS %s
+! RUN: %flang_fc1 -emit-hlfir -triple amdgcn-amd-amdhsa -fopenmp -fopenmp-is-device -fatomic-ignore-denormal-mode %s -o - | FileCheck -check-prefix=IGNORE-DENORMAL %s
+! RUN: %flang_fc1 -emit-hlfir -triple amdgcn-amd-amdhsa -fopenmp -fopenmp-is-device -fatomic-fine-grained-memory %s -o - | FileCheck -check-prefix=FINE-GRAINED-MEMORY %s
+! RUN: %flang_fc1 -emit-hlfir -triple amdgcn-amd-amdhsa -fopenmp -fopenmp-is-device -fatomic-remote-memory %s -o - | FileCheck -check-prefix=REMOTE-MEMORY %s
+program test
+    implicit none
+    integer :: A, B, threads
+    threads = 128
+    A = 0
+    B = 0
+    !UNSAFE-FP-ATOMICS: omp.atomic.update %{{.*}} : !fir.ref<i32> {
+    !UNSAFE-FP-ATOMICS: } {atomic_control = #omp.atomic_control<ignore_denormal_mode = true>}
+    !IGNORE-DENORMAL: omp.atomic.update %{{.*}} : !fir.ref<i32> {
+    !IGNORE-DENORMAL: } {atomic_control = #omp.atomic_control<ignore_denormal_mode = true>}
+    !FINE-GRAINED-MEMORY: omp.atomic.update %{{.*}} : !fir.ref<i32> {
+    !FINE-GRAINED-MEMORY: } {atomic_control = #omp.atomic_control<fine_grained_memory = true>}
+    !REMOTE-MEMORY: omp.atomic.update %{{.*}} : !fir.ref<i32> {
+    !REMOTE-MEMORY: } {atomic_control = #omp.atomic_control<remote_memory = true>}
+    !$omp target parallel num_threads(threads)
+    !$omp atomic
+    A =  A + 1
+    !$omp end target parallel
+    !UNSAFE-FP-ATOMICS: omp.atomic.update %{{.*}} : !fir.ref<i32> {
+    !UNSAFE-FP-ATOMICS: } {atomic_control = #omp.atomic_control<ignore_denormal_mode = true>}
+    !IGNORE-DENORMAL: omp.atomic.update %{{.*}} : !fir.ref<i32> {
+    !IGNORE-DENORMAL: } {atomic_control = #omp.atomic_control<ignore_denormal_mode = true>}
+    !FINE-GRAINED-MEMORY: omp.atomic.update %{{.*}} : !fir.ref<i32> {
+    !FINE-GRAINED-MEMORY: } {atomic_control = #omp.atomic_control<fine_grained_memory = true>}
+    !REMOTE-MEMORY: omp.atomic.update %{{.*}} : !fir.ref<i32> {
+    !REMOTE-MEMORY: } {atomic_control = #omp.atomic_control<remote_memory = true>}
+    !$omp target parallel num_threads(threads)
+    !$omp atomic capture
+        A = A + B
+        B = A
+    !$omp end atomic
+    !$omp end target parallel
+end program test

mlir/include/mlir/Dialect/OpenMP/OpenMPAttrDefs.td

@@ -22,6 +22,21 @@ class OpenMP_Attr<string name, string attrMnemonic, list<Trait> traits = [],
   let mnemonic = attrMnemonic;
 }
 
+//===----------------------------------------------------------------------===//
+// AtomicControlAttr
+//===----------------------------------------------------------------------===//
+
+// Atomic control attributes hold information about architectural
+// characteristics which are required for lowering atomic operations.
+def AtomicControlAttr : OpenMP_Attr<"AtomicControl", "atomic_control"> {
+  let parameters =
+      (ins DefaultValuedParameter<"bool", "false">:$ignore_denormal_mode,
+          DefaultValuedParameter<"bool", "false">:$fine_grained_memory,
+          DefaultValuedParameter<"bool", "false">:$remote_memory);
+
+  let assemblyFormat = "`<` struct(params) `>`";
+}
+
 //===----------------------------------------------------------------------===//
 // DeclareTargetAttr
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

@@ -1543,9 +1543,11 @@ def AtomicUpdateOp : OpenMP_Op<"atomic.update", traits = [
     operations.
   }] # clausesDescription;
 
-  let arguments = !con((ins Arg<OpenMP_PointerLikeType,
-                                "Address of variable to be updated",
-                                [MemRead, MemWrite]>:$x), clausesArgs);
+  let arguments = !con(
+      (ins Arg<OpenMP_PointerLikeType,
+               "Address of variable to be updated", [MemRead, MemWrite]>:$x,
+          AtomicControlAttr:$atomic_control),
+      clausesArgs);
 
   // Override region definition.
   let regions = (region SizedRegion<1>:$region);

mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir

@@ -133,7 +133,7 @@ func.func @atomic_update() {
     %1 = arith.constant 1 : i32
     %2 = arith.addi %arg0, %1  : i32
     omp.yield(%2 : i32)
-  }
+  } {atomic_control = #omp.atomic_control<>}
   return
 }
 llvm.mlir.global internal @_QFsEc() : i32 {

mlir/test/Dialect/OpenMP/canonicalize.mlir

@@ -4,7 +4,7 @@ func.func @update_no_op(%x : memref<i32>) {
   omp.atomic.update %x : memref<i32> {
   ^bb0(%xval : i32):
     omp.yield(%xval : i32)
-  }
+  } {atomic_control = #omp.atomic_control<>}
   return
 }
 
@@ -17,7 +17,7 @@ func.func @update_write_op(%x : memref<i32>, %value: i32) {
   omp.atomic.update %x : memref<i32> {
   ^bb0(%xval : i32):
     omp.yield(%value : i32)
-  }
+  } {atomic_control = #omp.atomic_control<>}
   return
 }
 
@@ -33,7 +33,7 @@ func.func @update_normal(%x : memref<i32>, %value: i32) {
   ^bb0(%xval : i32):
     %newval = arith.addi %xval, %value : i32
     omp.yield(%newval : i32)
-  }
+  } {atomic_control = #omp.atomic_control<>}
   return
 }
 
@@ -50,7 +50,7 @@ func.func @update_unnecessary_computations(%x: memref<i32>) {
   ^bb0(%xval: i32):
     %newval = arith.addi %xval, %c0 : i32
     omp.yield(%newval: i32)
-  }
+  } {atomic_control = #omp.atomic_control<>}
   return
 }
 
@@ -65,7 +65,7 @@ func.func @update_unnecessary_computations(%x: memref<i32>) {
   ^bb0(%xval: i32):
     %newval = arith.muli %xval, %c0 : i32
     omp.yield(%newval: i32)
-  }
+  } {atomic_control = #omp.atomic_control<>}
   return
 }