[flang][fir] Add affine optimization pass pipeline.

Author: NexMing

flang/include/flang/Optimizer/Passes/CommandLineOpts.h

@@ -42,6 +42,7 @@ extern llvm::cl::opt<bool> disableCfgConversion;
 extern llvm::cl::opt<bool> disableFirAvc;
 extern llvm::cl::opt<bool> disableFirMao;
 
+extern llvm::cl::opt<bool> enableAffineOpt;
 extern llvm::cl::opt<bool> disableFirAliasTags;
 extern llvm::cl::opt<bool> useOldAliasTags;
 

flang/include/flang/Optimizer/Passes/Pipelines.h

@@ -18,8 +18,11 @@
 #include "flang/Optimizer/Passes/CommandLineOpts.h"
 #include "flang/Optimizer/Transforms/Passes.h"
 #include "flang/Tools/CrossToolHelpers.h"
+#include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
 #include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h"
 #include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
+#include "mlir/Conversion/SCFToOpenMP/SCFToOpenMP.h"
+#include "mlir/Dialect/Affine/Passes.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMAttrs.h"
 #include "mlir/Pass/PassManager.h"

flang/lib/Optimizer/Passes/CMakeLists.txt

@@ -21,6 +21,7 @@ add_flang_library(flangPasses
   MLIRPass
   MLIRReconcileUnrealizedCasts
   MLIRSCFToControlFlow
+  MLIRSCFToOpenMP
   MLIRSupport
   MLIRTransforms
 )

flang/lib/Optimizer/Passes/CommandLineOpts.cpp

@@ -55,6 +55,7 @@ cl::opt<bool> useOldAliasTags(
     cl::desc("Use a single TBAA tree for all functions and do not use "
              "the FIR alias tags pass"),
     cl::init(false), cl::Hidden);
+EnableOption(AffineOpt, "affine-opt", "affine optimization");
 
 /// CodeGen Passes
 DisableOption(CodeGenRewrite, "codegen-rewrite", "rewrite FIR for codegen");

flang/lib/Optimizer/Passes/Pipelines.cpp

@@ -209,6 +209,26 @@ void createDefaultFIROptimizerPassPipeline(mlir::PassManager &pm,
   if (pc.AliasAnalysis && !disableFirAliasTags && !useOldAliasTags)
     pm.addPass(fir::createAddAliasTags());
 
+  // We can first convert the FIR dialect to the Affine dialect, perform
+  // optimizations on top of it, and then lower it to the FIR dialect.
+  // TODO: These optimization passes (like PromoteToAffinePass) are currently
+  // experimental, so it's important to actively identify and address issues.
+  if (enableAffineOpt && pc.OptLevel.isOptimizingForSpeed()) {
+    pm.addPass(fir::createPromoteToAffinePass());
+    pm.addPass(mlir::createCSEPass());
+    pm.addPass(mlir::affine::createAffineLoopInvariantCodeMotionPass());
+    pm.addPass(mlir::affine::createAffineLoopNormalizePass());
+    pm.addPass(mlir::affine::createSimplifyAffineStructuresPass());
+    pm.addPass(mlir::affine::createAffineParallelize(
+        mlir::affine::AffineParallelizeOptions{1, false}));
+    pm.addPass(fir::createAffineDemotionPass());
+    pm.addPass(mlir::createLowerAffinePass());
+    if (pc.EnableOpenMP) {
+      pm.addPass(mlir::createConvertSCFToOpenMPPass());
+      pm.addPass(mlir::createCanonicalizerPass());
+    }
+  }
+
   addNestedPassToAllTopLevelOperations<PassConstructor>(
       pm, fir::createStackReclaim);
   // convert control flow to CFG form

flang/test/Driver/mlir-pass-pipeline.f90

@@ -4,6 +4,7 @@
 ! -O0 is the default:
 ! RUN: %flang_fc1 -S -mmlir --mlir-pass-statistics -mmlir --mlir-pass-statistics-display=pipeline %s -O0 -o /dev/null 2>&1 | FileCheck --check-prefixes=ALL %s
 ! RUN: %flang_fc1 -S -mmlir --mlir-pass-statistics -mmlir --mlir-pass-statistics-display=pipeline %s -O2 -o /dev/null 2>&1 | FileCheck --check-prefixes=ALL,O2 %s
+! RUN: %flang_fc1 -S -mmlir --mlir-pass-statistics -mmlir --mlir-pass-statistics-display=pipeline -mllvm --enable-affine-opt %s -O2 -o /dev/null 2>&1 | FileCheck --check-prefixes=ALL,O2,AFFINE %s
 
 ! REQUIRES: asserts
 
@@ -105,6 +106,19 @@
 ! ALL-NEXT: SimplifyFIROperations
 ! O2-NEXT:  AddAliasTags
 
+! AFFINE-NEXT: 'func.func' Pipeline
+! AFFINE-NEXT:   AffineDialectPromotion
+! AFFINE-NEXT: CSE
+! AFFINE-NEXT:   (S) 0 num-cse'd - Number of operations CSE'd
+! AFFINE-NEXT:   (S) 0 num-dce'd - Number of operations DCE'd
+! AFFINE-NEXT: 'func.func' Pipeline
+! AFFINE-NEXT:   AffineLoopInvariantCodeMotion
+! AFFINE-NEXT:   AffineLoopNormalize
+! AFFINE-NEXT:   SimplifyAffineStructures
+! AFFINE-NEXT:   AffineParallelize
+! AFFINE-NEXT:   AffineDialectDemotion
+! AFFINE-NEXT: LowerAffinePass
+
 ! ALL-NEXT: Pipeline Collection : ['fir.global', 'func.func', 'omp.declare_reduction', 'omp.private']
 ! ALL-NEXT:    'fir.global' Pipeline
 ! ALL-NEXT:      StackReclaim

flang/test/Integration/OpenMP/auto-omp.f90

@@ -0,0 +1,10 @@
+! RUN: %flang_fc1 -O1 -mllvm --enable-affine-opt -emit-llvm -fopenmp -o - %s \
+! RUN: | FileCheck %s
+
+!CHECK-LABEL: define void @foo_(ptr captures(none) %0) {{.*}} {
+!CHECK: call void{{.*}}@__kmpc_fork_call{{.*}}@[[OMP_OUTLINED_FN_1:.*]])
+
+subroutine foo(a)
+  integer, dimension(100, 100), intent(out) :: a
+  a = 1
+end subroutine foo