[WIP] PoC Lowering of dialect attribute on builtin.module via DialectConversion

[skatrak]

Summary

This is just a proof-of-concept showing one way to do codegen for OpenMP
dialect discardable attributes attached to builtin.module in MLIR without
relying on LLVMTranslationDialectInterface::amendOperation(). Downstream
users of the OpenMP dialect interested on doing their own lowering within MLIR
could implement something similar to this, although there could be other better
ways to achieve the same. This is related to a discussion in patch D147219, and
not intended to be merged.

It can be tested in this way:

$ flang-new -fc1 -triple x86_64-unknown-linux-gnu -emit-llvm-bc -fopenmp \
            -mrelocation-model static -target-cpu x86-64 -mmlir \
            -mlir-print-ir-after-all -o requires-host.bc -x f95-cpp-input \
            llvm-project/flang/test/Lower/OpenMP/requires.f90 &> host.mlir
$ flang-new -fc1 -triple amdgcn-amd-amdhsa -emit-llvm-bc -fopenmp \
            -mrelocation-model pic -pic-level 2 -target-cpu gfx1030 \
            -fopenmp-host-ir-file-path requires-host.bc -fopenmp-is-device \
            -mmlir -mlir-print-ir-after-all -o requires-device.bc -x \
            f95-cpp-input llvm-project/flang/test/Lower/OpenMP/requires.f90 \
            &> device.mlir
$ llvm-dis requires-host.bc
$ llvm-dis requires-device.bc

Depends on D147218

[kiranchandramohan]

Thanks @skatrak.

flang/lib/Optimizer/CodeGen/CodeGen.cpp

void runOnOperation() override final {
     // which contains regions it is legal if the region contains only the
     // LLVM dialect. Add OpenMP dialect as a legal dialect for conversion and
     // legalize conversion of OpenMP operations without regions.
     mlir::configureOpenMPToLLVMConversionLegality(target, typeConverter);
     target.addLegalDialect<mlir::omp::OpenMPDialect>();
     target.addLegalDialect<mlir::acc::OpenACCDialect>();
 
     // required NOPs for applying a full conversion
-    target.addLegalOp<mlir::ModuleOp>();
+    target.addDynamicallyLegalOp<mlir::ModuleOp>([](mlir::ModuleOp op) {
+      // TODO Check for other attributes that need lowering
+      return !op->hasAttr("omp.requires");
+    });
 
     // If we're on Windows, we might need to rename some libm calls.
     bool isMSVC = fir::getTargetTriple(mod).isOSMSVCRT();
     if (isMSVC) {
       pattern.insert<RenameMSVCLibmCallees, RenameMSVCLibmFuncs>(context);
 
       target.addDynamicallyLegalOp<mlir::LLVM::CallOp>(
           [](mlir::LLVM::CallOp op) {

mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp

for (unsigned idx = 0; idx < curOp.getNumRegions(); idx++) {
                                              *this->getTypeConverter())))
         return failure();
     }
 
     rewriter.eraseOp(curOp);
     return success();
   }
 };
+
+/// Rewrite pattern that processes `builtin.module` operations if they contain
+/// the `omp.requires` attribute, removing it afterwards to mark the operation
+/// as already-processed.
+struct ModuleRequiresAttrConversion : public ConvertOpToLLVMPattern<ModuleOp> {
+  using ConvertOpToLLVMPattern<ModuleOp>::ConvertOpToLLVMPattern;
+  LogicalResult
+  matchAndRewrite(ModuleOp curOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    if (!curOp->hasAttr("omp.requires"))
+      return success();
+
+    rewriter.updateRootInPlace(curOp.getOperation(), [&] {
+      auto offloadMod = cast<omp::OffloadModuleInterface>(curOp.getOperation());
+      if (!offloadMod.getIsDevice()) {
+        // Ad-hoc definition of 'requires' flags expected by the OpenMP runtime.
+        enum {
+          REVERSE_OFFLOAD = 0x002,
+          UNIFIED_ADDRESS = 0x004,
+          UNIFIED_SHARED_MEMORY = 0x008,
+          DYNAMIC_ALLOCATORS = 0x010
+        };
+
+        // Populate flags from `omp.requires` attribute.
+        int32_t requiresFlags = 0;
+        auto requiresClauses = offloadMod.getRequires();
+        if (bitEnumContainsAll(requiresClauses,
+                               omp::ClauseRequires::reverse_offload))
+          requiresFlags |= REVERSE_OFFLOAD;
+        if (bitEnumContainsAll(requiresClauses,
+                               omp::ClauseRequires::unified_address))
+          requiresFlags |= UNIFIED_ADDRESS;
+        if (bitEnumContainsAll(requiresClauses,
+                               omp::ClauseRequires::unified_shared_memory))
+          requiresFlags |= UNIFIED_SHARED_MEMORY;
+        if (bitEnumContainsAll(requiresClauses,
+                               omp::ClauseRequires::dynamic_allocators))
+          requiresFlags |= DYNAMIC_ALLOCATORS;
+
+        OpBuilder::InsertionGuard guard(rewriter);
+        rewriter.setInsertionPointToStart(curOp.getBody());
+
+        // Declaration of `__tgt_register_requires(i32)`.
+        auto i32Type = IntegerType::get(getContext(), 32);
+        auto regFnType = LLVM::LLVMFunctionType::get(
+            LLVM::LLVMVoidType::get(getContext()), {i32Type});
+        auto regFn = rewriter.create<LLVM::LLVMFuncOp>(
+            curOp->getLoc(), "__tgt_register_requires", regFnType,
+            LLVM::linkage::Linkage::External);
+
+        // Definition of `.omp_offloading.requires_reg()`.
+        auto ctorFnType = LLVM::LLVMFunctionType::get(
+            LLVM::LLVMVoidType::get(getContext()), {});
+        auto ctorFn = rewriter.create<LLVM::LLVMFuncOp>(
+            curOp->getLoc(), ".omp_offloading.requires_reg", ctorFnType,
+            LLVM::linkage::Linkage::Internal);
+        ctorFn.addEntryBlock();
+
+        OpBuilder ctorFnBuilder(ctorFn.getBody());
+        auto flagsConstOp = ctorFnBuilder.create<LLVM::ConstantOp>(
+            ctorFn->getLoc(), i32Type,
+            IntegerAttr::get(i32Type, requiresFlags));
+        ctorFnBuilder.create<LLVM::CallOp>(ctorFn->getLoc(), TypeRange(),
+                                           SymbolRefAttr::get(regFn),
+                                           flagsConstOp.getResult());
+        ctorFnBuilder.create<LLVM::ReturnOp>(ctorFn->getLoc(), nullptr);
+
+        // Registration of `.omp_offloading.requires_reg()` as constructor.
+        auto ctors = rewriter.getArrayAttr(
+            {FlatSymbolRefAttr::get(getContext(), ctorFn.getSymName())});
+        auto priorities =
+            rewriter.getArrayAttr({rewriter.getI32IntegerAttr(0)});
+        rewriter.create<LLVM::GlobalCtorsOp>(curOp->getLoc(), ctors,
+                                             priorities);
+      }
+
+      curOp->removeAttr("omp.requires");
+    });
+    return success();
+  }
+};
 } // namespace
 
 void mlir::configureOpenMPToLLVMConversionLegality(
     ConversionTarget &target, LLVMTypeConverter &typeConverter) {
   target.addDynamicallyLegalOp<
       mlir::omp::AtomicUpdateOp, mlir::omp::CriticalOp, mlir::omp::TargetOp,
       mlir::omp::DataOp, mlir::omp::ParallelOp, mlir::omp::WsLoopOp,
       mlir::omp::SimdLoopOp, mlir::omp::MasterOp, mlir::omp::SectionOp,
patterns.add<
       RegionOpConversion<omp::DataOp>, RegionOpConversion<omp::TargetOp>,
       RegionLessOpWithVarOperandsConversion<omp::AtomicReadOp>,
       RegionLessOpWithVarOperandsConversion<omp::AtomicWriteOp>,
       RegionOpWithVarOperandsConversion<omp::AtomicUpdateOp>,
       RegionLessOpWithVarOperandsConversion<omp::FlushOp>,
       RegionLessOpWithVarOperandsConversion<omp::ThreadprivateOp>,
       RegionLessOpConversion<omp::YieldOp>,
       RegionLessOpConversion<omp::EnterDataOp>,
-      RegionLessOpConversion<omp::ExitDataOp>>(converter);
+      RegionLessOpConversion<omp::ExitDataOp>, ModuleRequiresAttrConversion>(
+      converter);
 }
 
 namespace {
 struct ConvertOpenMPToLLVMPass
     : public impl::ConvertOpenMPToLLVMPassBase<ConvertOpenMPToLLVMPass> {
   using Base::Base;
 
   void runOnOperation() override;