[flang-commits] [clang] [flang] [flang][OpenMP] Add -f[no]-openmp-simd (PR #150269)

[Tom Eccles via flang-commits]

================
@@ -0,0 +1,360 @@
+#include "flang/Optimizer/Builder/FIRBuilder.h"
+#include "flang/Optimizer/Transforms/Utils.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include <llvm/Support/Debug.h>
+#include <mlir/IR/MLIRContext.h>
+#include <mlir/IR/Operation.h>
+#include <mlir/IR/PatternMatch.h>
+#include <mlir/Support/LLVM.h>
+
+namespace flangomp {
+#define GEN_PASS_DEF_SIMDONLYPASS
+#include "flang/Optimizer/OpenMP/Passes.h.inc"
+} // namespace flangomp
+
+namespace {
+
+#define DEBUG_TYPE "omp-simd-only-pass"
+
+class SimdOnlyConversionPattern : public mlir::RewritePattern {
+public:
+  SimdOnlyConversionPattern(mlir::MLIRContext *ctx)
+      : mlir::RewritePattern(MatchAnyOpTypeTag{}, 1, ctx) {}
+
+  mlir::LogicalResult
+  matchAndRewrite(mlir::Operation *op,
+                  mlir::PatternRewriter &rewriter) const override {
+    if (op->getDialect()->getNamespace() !=
+        mlir::omp::OpenMPDialect::getDialectNamespace())
+      return rewriter.notifyMatchFailure(op, "Not an OpenMP op");
+
+    if (auto simdOp = mlir::dyn_cast<mlir::omp::SimdOp>(op)) {
+      // Remove the composite attr given that the op will no longer be composite
+      if (simdOp.isComposite()) {
+        simdOp.setComposite(false);
+        return mlir::success();
+      }
+
+      return rewriter.notifyMatchFailure(op, "Op is a plain SimdOp");
+    }
+
+    if (op->getParentOfType<mlir::omp::SimdOp>())
+      return rewriter.notifyMatchFailure(op, "Op is nested under a SimdOp");
+
+    if (!mlir::isa<mlir::func::FuncOp>(op->getParentOp()) &&
+        (mlir::isa<mlir::omp::TerminatorOp>(op) ||
+         mlir::isa<mlir::omp::YieldOp>(op)))
+      return rewriter.notifyMatchFailure(op,
+                                         "Non top-level yield or terminator");
+
+    // SectionOp overrides its BlockArgInterface based on the parent SectionsOp.
+    // We need to make sure we only rewrite omp.sections once all omp.section
+    // ops inside it have been rewritten, otherwise the individual omp.section
+    // ops will not be able to access their argument values.
+    if (auto sectionsOp = mlir::dyn_cast<mlir::omp::SectionsOp>(op)) {
+      for (auto &opInSections : sectionsOp.getRegion().getOps())
+        if (mlir::isa<mlir::omp::SectionOp>(opInSections))
+          return rewriter.notifyMatchFailure(
+              op, "SectionsOp still contains individual sections");
+    }
+
+    LLVM_DEBUG(llvm::dbgs() << "SimdOnlyPass matched OpenMP op:\n");
+    LLVM_DEBUG(op->dump());
+
+    // Erase ops that don't need any special handling
+    if (mlir::isa<mlir::omp::BarrierOp>(op) ||
+        mlir::isa<mlir::omp::FlushOp>(op) ||
+        mlir::isa<mlir::omp::TaskyieldOp>(op) ||
+        mlir::isa<mlir::omp::MapBoundsOp>(op) ||
+        mlir::isa<mlir::omp::TargetEnterDataOp>(op) ||
+        mlir::isa<mlir::omp::TargetExitDataOp>(op) ||
+        mlir::isa<mlir::omp::TargetUpdateOp>(op) ||
+        mlir::isa<mlir::omp::OrderedOp>(op) ||
+        mlir::isa<mlir::omp::CancelOp>(op) ||
+        mlir::isa<mlir::omp::CancellationPointOp>(op) ||
+        mlir::isa<mlir::omp::ScanOp>(op) ||
+        mlir::isa<mlir::omp::TaskwaitOp>(op)) {
+      rewriter.eraseOp(op);
+      return mlir::success();
+    }
+
+    fir::FirOpBuilder builder(rewriter, op);
+    mlir::Location loc = op->getLoc();
+
+    auto inlineSimpleOp = [&](mlir::Operation *ompOp) -> bool {
+      if (!ompOp)
+        return false;
+
+      llvm::SmallVector<std::pair<mlir::Value, mlir::BlockArgument>>
+          blockArgsPairs;
+      if (auto iface =
+              mlir::dyn_cast<mlir::omp::BlockArgOpenMPOpInterface>(op)) {
+        iface.getBlockArgsPairs(blockArgsPairs);
+        for (auto [value, argument] : blockArgsPairs)
+          rewriter.replaceAllUsesWith(argument, value);
+      }
+
+      if (ompOp->getRegion(0).getBlocks().size() == 1) {
+        auto &block = *ompOp->getRegion(0).getBlocks().begin();
+        // This block is about to be removed so any arguments should have been
+        // replaced by now.
+        block.eraseArguments(0, block.getNumArguments());
+        if (auto terminatorOp =
+                mlir::dyn_cast<mlir::omp::TerminatorOp>(block.back())) {
+          rewriter.eraseOp(terminatorOp);
+        }
+        rewriter.inlineBlockBefore(&block, op, {});
----------------
tblah wrote:

ultra-nit
```suggestion
        rewriter.inlineBlockBefore(&block, ompOp, {});
```

https://github.com/llvm/llvm-project/pull/150269