diff --git a/flang/test/Lower/OpenMP/requires.f90 b/flang/test/Lower/OpenMP/requires.f90 --- a/flang/test/Lower/OpenMP/requires.f90 +++ b/flang/test/Lower/OpenMP/requires.f90 @@ -2,10 +2,10 @@ ! This test checks the lowering of requires into MLIR -! MLIR attributes -!MLIR: module attributes { -!MLIR-SAME: omp.atomic_default_mem_order = #omp -!MLIR-SAME: omp.requires = #omp +! Module attributes +! MLIR: module attributes { +! MLIR-SAME: omp.atomic_default_mem_order = #omp +! MLIR-SAME: omp.requires = #omp program requires !$omp requires unified_shared_memory reverse_offload atomic_default_mem_order(seq_cst) diff --git a/mlir/lib/Target/LLVMIR/CMakeLists.txt b/mlir/lib/Target/LLVMIR/CMakeLists.txt --- a/mlir/lib/Target/LLVMIR/CMakeLists.txt +++ b/mlir/lib/Target/LLVMIR/CMakeLists.txt @@ -37,6 +37,7 @@ MLIRDLTIDialect MLIRLLVMDialect MLIRLLVMIRTransforms + MLIROpenMPDialect MLIRTranslateLib ) diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp --- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp +++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp @@ -15,6 +15,7 @@ #include "mlir/IR/IRMapping.h" #include "mlir/IR/Operation.h" #include "mlir/Support/LLVM.h" +#include "mlir/Support/LogicalResult.h" #include "mlir/Target/LLVMIR/Dialect/OpenMPCommon.h" #include "mlir/Target/LLVMIR/ModuleTranslation.h" @@ -24,6 +25,8 @@ #include "llvm/Frontend/OpenMP/OMPIRBuilder.h" #include "llvm/IR/DebugInfoMetadata.h" #include "llvm/IR/IRBuilder.h" +#include "llvm/Transforms/Utils/ModuleUtils.h" +#include using namespace mlir; @@ -1031,11 +1034,24 @@ return success(); } +static std::optional +getAtomicDefaultMemOrder(Operation &opInst) { + // Try to get the omp.atomic_default_mem_order attribute, if present + if (auto offloadModule = + opInst.getParentOfType()) + return offloadModule.getAtomicDefaultMemOrder(); + + return std::nullopt; +} + /// Convert an Atomic Ordering attribute to llvm::AtomicOrdering. -llvm::AtomicOrdering -convertAtomicOrdering(std::optional ao) { +static llvm::AtomicOrdering +convertAtomicOrdering(std::optional ao, + std::optional defaultAo) { + // If not specified, try using the default atomic ordering gathered from a + // requires atomic_mem_default_order clause, if present if (!ao) - return llvm::AtomicOrdering::Monotonic; // Default Memory Ordering + ao = defaultAo.value_or(omp::ClauseMemoryOrderKind::Relaxed); switch (*ao) { case omp::ClauseMemoryOrderKind::Seq_cst: @@ -1056,13 +1072,14 @@ static LogicalResult convertOmpAtomicRead(Operation &opInst, llvm::IRBuilderBase &builder, LLVM::ModuleTranslation &moduleTranslation) { - auto readOp = cast(opInst); llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder(); llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder); - llvm::AtomicOrdering AO = convertAtomicOrdering(readOp.getMemoryOrderVal()); + auto defaultAO = getAtomicDefaultMemOrder(opInst); + llvm::AtomicOrdering AO = + convertAtomicOrdering(readOp.getMemoryOrderVal(), defaultAO); llvm::Value *x = moduleTranslation.lookupValue(readOp.getX()); llvm::Value *v = moduleTranslation.lookupValue(readOp.getV()); @@ -1083,7 +1100,9 @@ llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder(); llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder); - llvm::AtomicOrdering ao = convertAtomicOrdering(writeOp.getMemoryOrderVal()); + auto defaultAO = getAtomicDefaultMemOrder(opInst); + llvm::AtomicOrdering ao = + convertAtomicOrdering(writeOp.getMemoryOrderVal(), defaultAO); llvm::Value *expr = moduleTranslation.lookupValue(writeOp.getValue()); llvm::Value *dest = moduleTranslation.lookupValue(writeOp.getAddress()); llvm::Type *ty = moduleTranslation.convertType(writeOp.getValue().getType()); @@ -1147,8 +1166,9 @@ /*isSigned=*/false, /*isVolatile=*/false}; + auto defaultAO = getAtomicDefaultMemOrder(*opInst.getOperation()); llvm::AtomicOrdering atomicOrdering = - convertAtomicOrdering(opInst.getMemoryOrderVal()); + convertAtomicOrdering(opInst.getMemoryOrderVal(), defaultAO); // Generate update code. LogicalResult updateGenStatus = success(); @@ -1236,8 +1256,9 @@ /*isSigned=*/false, /*isVolatile=*/false}; + auto defaultAO = getAtomicDefaultMemOrder(*atomicCaptureOp.getOperation()); llvm::AtomicOrdering atomicOrdering = - convertAtomicOrdering(atomicCaptureOp.getMemoryOrderVal()); + convertAtomicOrdering(atomicCaptureOp.getMemoryOrderVal(), defaultAO); LogicalResult updateGenStatus = success(); auto updateFn = [&](llvm::Value *atomicx, @@ -1542,6 +1563,25 @@ return bodyGenStatus; } +/// Converts the module-level set of OpenMP requires clauses into LLVM IR using +/// OpenMPIRBuilder. +static LogicalResult +convertRequires(Operation &op, omp::ClauseRequiresAttr requiresAttr, + LLVM::ModuleTranslation &moduleTranslation) { + auto *ompBuilder = moduleTranslation.getOpenMPBuilder(); + + // No need to read requiresAttr here, because it has already been done in + // translateModuleToLLVMIR(). There, flags are stored in the + // OpenMPIRBuilderConfig object, available to the OpenMPIRBuilder. + auto *regFn = + ompBuilder->createRegisterRequires(ompBuilder->createPlatformSpecificName( + {"omp_offloading", "requires_reg"})); + + // Add registration function as global constructor + llvm::appendToGlobalCtors(ompBuilder->M, regFn, /* Priority = */ 0); + return success(); +} + namespace { /// Implementation of the dialect interface that converts operations belonging @@ -1556,6 +1596,12 @@ LogicalResult convertOperation(Operation *op, llvm::IRBuilderBase &builder, LLVM::ModuleTranslation &moduleTranslation) const final; + + /// Processes omp dialect attributes attached to operations from a different + /// dialect. + LogicalResult + amendOperation(Operation *op, NamedAttribute attribute, + LLVM::ModuleTranslation &moduleTranslation) const final; }; } // namespace @@ -1665,6 +1711,21 @@ }); } +LogicalResult OpenMPDialectLLVMIRTranslationInterface::amendOperation( + Operation *op, NamedAttribute attribute, + LLVM::ModuleTranslation &moduleTranslation) const { + return TypeSwitch(attribute.getValue()) + .Case([&](omp::ClauseRequiresAttr requiresAttr) { + return convertRequires(*op, requiresAttr, moduleTranslation); + }) + .Default([](Attribute) { + // The omp.atomic_default_mem_order attribute does not need lowering + // in this switch, as it is read directly during OpenMP atomic ops + // lowering + return success(); + }); +} + void mlir::registerOpenMPDialectTranslation(DialectRegistry ®istry) { registry.insert(); registry.addExtension(+[](MLIRContext *ctx, omp::OpenMPDialect *dialect) { diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp --- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp +++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp @@ -1311,6 +1311,25 @@ LLVM::ensureDistinctSuccessors(module); ModuleTranslation translator(module, std::move(llvmModule)); + + // Set OpenMP IR Builder configuration + if (auto offloadMod = dyn_cast(module)) { + llvm::OpenMPIRBuilderConfig config; + config.setIsEmbedded(offloadMod.getIsDevice()); + config.setIsTargetCodegen(false); + + const auto requiresFlags = offloadMod.getRequires(); + config.setHasRequiresReverseOffload(bitEnumContainsAll( + requiresFlags, omp::ClauseRequires::reverse_offload)); + config.setHasRequiresUnifiedAddress(bitEnumContainsAll( + requiresFlags, omp::ClauseRequires::unified_address)); + config.setHasRequiresUnifiedSharedMemory(bitEnumContainsAll( + requiresFlags, omp::ClauseRequires::unified_shared_memory)); + config.setHasRequiresDynamicAllocators(bitEnumContainsAll( + requiresFlags, omp::ClauseRequires::dynamic_allocators)); + translator.getOpenMPBuilder()->setConfig(config); + } + if (failed(translator.convertFunctionSignatures())) return nullptr; if (failed(translator.convertGlobals())) diff --git a/mlir/test/Target/LLVMIR/openmp-llvm.mlir b/mlir/test/Target/LLVMIR/openmp-llvm.mlir --- a/mlir/test/Target/LLVMIR/openmp-llvm.mlir +++ b/mlir/test/Target/LLVMIR/openmp-llvm.mlir @@ -2454,3 +2454,54 @@ } llvm.return } + +// ----- + +// Check that the atomic default memory order is picked up by atomic operations. +module attributes { + omp.atomic_default_mem_order = #omp +} { + // CHECK-LABEL: @omp_atomic_default_mem_order + // CHECK-SAME: (ptr %[[ARG0:.*]], ptr %[[ARG1:.*]], i32 %[[EXPR:.*]]) + llvm.func @omp_atomic_default_mem_order(%arg0 : !llvm.ptr, + %arg1 : !llvm.ptr, + %expr : i32) -> () { + + // CHECK: %[[X1:.*]] = load atomic i32, ptr %[[ARG0]] seq_cst, align 4 + // CHECK: store i32 %[[X1]], ptr %[[ARG1]], align 4 + omp.atomic.read %arg1 = %arg0 : !llvm.ptr, i32 + + // CHECK: store atomic i32 %[[EXPR]], ptr %[[ARG1]] seq_cst, align 4 + // CHECK: call void @__kmpc_flush(ptr @{{.*}}) + omp.atomic.write %arg1 = %expr : !llvm.ptr, i32 + + // CHECK: atomicrmw add ptr %[[ARG1]], i32 %[[EXPR]] seq_cst + omp.atomic.update %arg1 : !llvm.ptr { + ^bb0(%xval: i32): + %newval = llvm.add %xval, %expr : i32 + omp.yield(%newval : i32) + } + + // CHECK: %[[xval:.*]] = atomicrmw xchg ptr %[[ARG0]], i32 %[[EXPR]] seq_cst + // CHECK: store i32 %[[xval]], ptr %[[ARG1]] + omp.atomic.capture { + omp.atomic.read %arg1 = %arg0 : !llvm.ptr, i32 + omp.atomic.write %arg0 = %expr : !llvm.ptr, i32 + } + + llvm.return + } +} + +// ----- + +// Check that OpenMP requires flags are registered by a global constructor. +module attributes { + omp.requires = #omp +} { + // CHECK: @llvm.global_ctors = appending global [1 x { i32, ptr, ptr }] + // CHECK-SAME: [{ i32, ptr, ptr } { i32 0, ptr @[[REG_FN:.*]], ptr null }] + // CHECK: define {{.*}} @[[REG_FN]]({{.*}}) + // CHECK-NOT: } + // CHECK: call void @__tgt_register_requires(i64 10) +}