Unverified Commit d6061d29 authored by Sergio Afonso's avatar Sergio Afonso Committed by GitHub
Browse files

[Flang][OpenMP] Add pass to replace allocas with device shared memory (#161863) 

This patch introduces a new Flang OpenMP MLIR pass, only ran for target
device modules, that identifies `fir.alloca` operations that should use
device shared memory and replaces them with pairs of
`omp.alloc_shared_mem` and `omp.free_shared_mem` operations.

This works in conjunction to the MLIR to LLVM IR translation pass'
handling of privatization, mapping and reductions in the OpenMP dialect
to properly select the right memory space for allocations based on where
they are made and where they are used.

This pass, in particular, handles explicit stack allocations in MLIR,
whereas the aforementioned translation pass takes care of implicit ones
represented by entry block arguments.
parent 568e1ad6

flang/include/flang/Optimizer/Support/InitFIR.h

+2 −0
Original line number Diff line number Diff line
@@ -34,6 +34,7 @@ 
#include "mlir/Dialect/Math/IR/Math.h"

#include "mlir/Dialect/OpenACC/OpenACC.h"

#include "mlir/Dialect/OpenACC/Transforms/Passes.h"

#include "mlir/Dialect/OpenMP/Transforms/Passes.h"

#include "mlir/Dialect/SCF/IR/SCF.h"

#include "mlir/Dialect/SCF/Transforms/Passes.h"

#include "mlir/InitAllDialects.h"
@@ -106,6 +107,7 @@ inline void loadDialects(mlir::MLIRContext &context) { 
/// but is a smaller set since we aren't using many of the passes found there.

inline void registerMLIRPassesForFortranTools() {

  mlir::acc::registerOpenACCPasses();

  mlir::omp::registerOpenMPPasses();

  mlir::registerCanonicalizerPass();

  mlir::registerCSEPass();

  mlir::affine::registerAffineLoopFusionPass();

flang/lib/Optimizer/Passes/Pipelines.cpp

+7 −0
Original line number Diff line number Diff line
@@ -13,6 +13,7 @@ 
#include "flang/Optimizer/OpenACC/Passes.h"

#include "mlir/Conversion/Passes.h"

#include "mlir/Dialect/LLVMIR/Transforms/Passes.h"

#include "mlir/Dialect/OpenMP/Transforms/Passes.h"

#include "llvm/Support/CommandLine.h"



/// Force setting the no-alias attribute on fuction arguments when possible.
@@ -440,6 +441,12 @@ void createDefaultFIRCodeGenPassPipeline(mlir::PassManager &pm, 
  }



  fir::addFIRToLLVMPass(pm, config);



  // Convert applicable OpenMP stack allocations to shared memory allocations

  // for GPU targets. This pass must run after any alloca-generating passes to

  // ensure all are adequately accounted for.

  if (config.EnableOpenMP && !config.EnableOpenMPSimd)

    pm.addPass(mlir::omp::createStackToSharedPass());

}



/// Create a pass pipeline for lowering from MLIR to LLVM IR

flang/test/Fir/basic-program.fir

+2 −0
Original line number Diff line number Diff line
@@ -178,5 +178,7 @@ func.func @_QQmain() { 
// PASSES-NEXT:  LowerNontemporalPass

// PASSES-NEXT: FIRToLLVMLowering

// PASSES-NEXT: ReconcileUnrealizedCasts

// PASSES-NEXT: 'llvm.func' Pipeline

// PASSES-NEXT:  StackToSharedPass

// PASSES-NEXT: PrepareForOMPOffloadPrivatizationPass

// PASSES-NEXT: LLVMIRLoweringPass

mlir/docs/Passes.md

+4 −0
Original line number Diff line number Diff line
@@ -76,6 +76,10 @@ This document describes the available MLIR passes and their contracts. 


[include "MemRefPasses.md"]



## 'omp' Dialect Passes



[include "OpenMPPasses.md"]



## 'shard' Dialect Passes



[include "ShardPasses.md"]

mlir/include/mlir/Dialect/OpenMP/Transforms/Passes.h

+5 −1
Original line number Diff line number Diff line
@@ -13,6 +13,10 @@ 


namespace mlir {



namespace LLVM {

class LLVMFuncOp;

} // namespace LLVM



namespace omp {



/// Generate the code for registering conversion passes.
@@ -23,4 +27,4 @@ namespace omp { 
} // namespace omp

} // namespace mlir



#endif // MLIR_DIALECT_LLVMIR_TRANSFORMS_PASSES_H
 
#endif // MLIR_DIALECT_OPENMP_TRANSFORMS_PASSES_H