[mlir][OpenMP] Lower simd if clause to LLVM IR

[domada]

Summary

Patch: https://reviews.llvm.org/D128940 added support for simd if clause for MLIR dialect. This patch adds support for lowering MLIR simd if clause to LLVM IR.

[kiranchandramohan]

Other options include,
-> Adding the if option to applySimd in the OpenMPIRBuilder. CreateParallel takes in the ifClause as an argument. But it can be argued that it does this since ifClause for parallel requires the insertion of runtime calls. https://github.com/llvm/llvm-project/blob/6858a17f66f6d609e2d801c443da951c391636a0/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp#L833
-> Handling the if clause in the OpenMPToLLVMConversion pass. https://github.com/llvm/llvm-project/blob/main/mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp
-> Or adding an OpenMPtoSCF pass like what OpenACC does and handling the if clause there. https://github.com/llvm/llvm-project/blob/main/mlir/lib/Conversion/OpenACCToSCF/OpenACCToSCF.cpp

[clementval]

flang/test/Lower/OpenMP/simd.f90

4

It might be better to not mix test like that. Normally here fir is tested. You can use tco to test fir to LLVM IR translation.

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

994

This (and the following lines) is better if it's done in the OpenMPIRBuilder.

[raghavendhra]

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

973

Nit : Please use spaces instead of Tab

[domada]

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

994

Hi,
I would like to ask you why you opt for modification of OpenMPIRBuilder instead of adding an OpenMPtoSCF pass like what OpenACC does and handling the if clause there. https://github.com/llvm/llvm-project/blob/main/mlir/lib/Conversion/OpenACCToSCF/OpenACCToSCF.cpp ?

[clementval]

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

994

I'm not sure I follow your question here. I personally did not opt for anything in this patch since it is yours. In OpenACC, for standalone data movement ops it's straight forward to do it in MLIR directly.

[shraiysh]

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

994

I think that can be done here too (like Kiran suggested). In that case, these lines would be eliminated. There are pros and cons to both the approaches 1 and 3 in Kiran's reply. If we follow the OpenACC approach, this same thing will have to be done inside clang, but it should be straightforward for OpenMP Dialect. If we follow the OpenMP IR Builder, it will probably be a bit trickier, but it will be a nice abstraction for both OpenMP Dialect and Clang.

I think what Roger meant was why choose one over another? Or, in other words - which one of these should we do here? (I don't have a strong preference for one or another)

[shraiysh]

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

994

Apologies for taking the wrong name.
I think what Roger Dominik meant was why choose one over another? Or, in other words - which one of these should we do here? (I don't have a strong preference for one or another)

[kiranchandramohan]

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

994

In general, we should consider lowering the clauses or operation in translation as the last possible option. Translation is intended to be a light-weight transformation and indeed it is for LLVM dialect. For using the OpenMPIRBuilder we do some more stuff in translation but we should try to keep it as lightweight as possible and implement things in the OpenMPIRBuilder or if that is not possible, in a conversion pass. So preference will be,

1. OpenMPIRBuilder
2. OpenMPToSCF
3. OpenMPToLLVM
4. Translation

Also, it will be good to have things implemented in a uniform fashion for various OpenMP Ops and also among OpenACC and OpenMP.

[domada]

Applied remarks. Moved generation of if clause to OMPIRBuilder.

[domada]

Fix build failures

[Meinersbur]

llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

341

[serious] This is really ugly from a software architecture point of view. Where previously all could expect a CanonicalLoopInfo has a preheader has to exist, it now any loop-handling code has to check for it.

The description of CanonicalLoopInfo needs to be updated.

616–617

What this function does should be achievable with 3 elementary steps:

1. Code versioning (IfCond)
2. collapseLoops
3. applySimd(CanonicalLoopInfo, SimLen)

Step 1 does not yet exist as an OpenMPIRBuilder method but but maybe should be made available independent of applySimd to by usable for anything that implements an if clause.

I would not object to have such a helper function doing all 3 steps together, but the description should reflect that.

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

2566–2571

[nit] unrelated

2874–2887

Please avoid copy&paste code. The common code can be extracted into a separate function.

However, I don't think a version taking a Loop info is even necessary. See comment below.

2949–2959

[serious] As far as I understand, the only reason this instantiates a pass manger is to use the cloneLoopWithPreheader function. However, it is sufficient to clone all block from CanonicalLoopInfo::getHeader() to CanonicalLoopInfo::getAfter(). Since the ExitBlock is guaranteed to be the only exit out of the loop (at least until I can introduce breaks; D124823 is a step into that direction), all blocks in the loop can be simply enumerated and cloned without cloneLoopWithPreheader.

This functionality is probably useful for implementing any loop versioning for if-clauses and should probably made its own method.

[domada]

llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

341

Hi,
thanks for the remarks.

I decided to modify createCanonicalLoop interface, because omp loop trip count was calculated outside the loop preheader.

I needed to calculate trip count inside loop preheader because I used cloneLoopWithPreheader function. This function needs a BasicBlock which dominates. If we create loop trip count outside loop preheader then we cannot use loop preheader as dominator block.

According to your comments we can omit cloneLoopWithPreheader function call. In consequence there is no need to modify createCanonicalLoop function.

[domada]

Implemented simd if clause- applying review remarks:

1. Use simplified cloning mechanism (without cloneLoopWithPreheader).
2. Removed addLoopMetadata(Loop *L, Metadata) function.
3. Added more tests.

[domada]

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

2949–2959

I do not use cloneLoopWithPreheader function any more, but I still need pass manager to define Loop blocks. Original applySimd function also uses pass manager for defining loop blocks. Please let me know if you have better idea, how to extract all body blocks.

[Meinersbur]

llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

609

As mentioned, I think this can be integrated into the existing applySimd methods. All it does need is the IfCond argument, and if it is set, implement the loop versioning. After that, use the same code to attach the metadata.

collapseLoops can be called by the translation layer to get a single CanonicalLoopInfo handle to be passed to this function.

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

2949–2959

Enumerating loop block without LoopInfo will be possible after D124823. Instantiation of an entire pass manager infrastructure does not scale well and should not be a long-term solution. It runs DominatorTree, LoopInfo etc on the entire function. n loops in the function have a complexity of at O(n³), because DominatorTree is already O(n²).

3012

3039

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

974–980

[domada]

Applied remarks:
Do not collapse loops inside function responsible for applying SIMD metadirectives
Added special function for loop versioning

[Meinersbur]

LGTM, this turned out really nice. Thank you!

Some nitpicks left. Please apply before committing.

llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

622

This could be made private as front-ends are encouraged to use the IfCond argument for the relevant instead that may implement specialized behavior (like passing the value to the runtime).

It still might be useful for front-ends (eg implementing dynamic metadirective), but I'd wait to such use cases to occur before adding this to the public API.

Edit: I noticed it is tested in the unit test, which requires it to be public.

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

2949–2959

Could you add a TODO/FIXME comment mentioning that we should not depend on a pass manger?

3031

operator[] may modify the underlaying hashtable.

The name MappedLatch is more meaningful than MappedValue.

3034

[domada]

llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

622

I moved it to private and I removed corresponding unit test. This functionality is also tested by applySimdLoopIf unit test.

flang/test/Lower/OpenMP/simd.f90

 ! Tests for 2.9.3.1 Simd
 
 ! RUN: bbc -fopenmp -emit-fir %s -o - | FileCheck %s
+! RUN: bbc -fopenmp -emit-fir %s -o - | fir-opt --fir-to-llvm-ir | FileCheck %s --check-prefixes="LLVMIRDialect"
+! RUN: bbc -fopenmp -emit-fir %s -o - | fir-opt --fir-to-llvm-ir | tco | FileCheck %s --check-prefix=LLVMIR
 
 !CHECK-LABEL: func @_QPsimdloop()
 subroutine simdloop
 integer :: i
   !$OMP SIMD
   ! CHECK: %[[LB:.*]] = arith.constant 1 : i32
   ! CHECK-NEXT: %[[UB:.*]] = arith.constant 9 : i32
   ! CHECK-NEXT: %[[STEP:.*]] = arith.constant 1 : i32

 !CHECK-LABEL: func @_QPsimdloop_with_if_clause
 subroutine simdloop_with_if_clause(n, threshold)
 integer :: i, n, threshold
   !$OMP SIMD IF( n .GE. threshold )
   ! CHECK: %[[LB:.*]] = arith.constant 1 : i32
   ! CHECK: %[[UB:.*]] = fir.load %arg0
   ! CHECK: %[[STEP:.*]] = arith.constant 1 : i32
   ! CHECK: %[[COND:.*]] = arith.cmpi sge
+  ! LLVMIRDialect: %[[LB_LLVMIR1:.*]] = llvm.mlir.constant(1 : i32) : i32
+  ! LLVMIRDialect: %[[UB_LLVMIR1:.*]] = llvm.load %arg0 : !llvm.ptr<i32>
+  ! LLVMIRDialect: %[[STEP_LLVMIR1:.*]] = llvm.mlir.constant(1 : i32) : i32
+  ! LLVMIRDialect: %[[COND_LLVMIR1:.*]] = llvm.icmp "sge"
   ! CHECK: omp.simdloop if(%[[COND:.*]]) for (%[[I:.*]]) : i32 = (%[[LB]]) to (%[[UB]]) inclusive  step (%[[STEP]]) {
+  ! LLVMIRDialect: omp.simdloop if(%[[COND_LLVMIR1:.*]]) for (%[[I:.*]]) : i32 = (%[[LB_LLVMIR1]]) to (%[[UB_LLVMIR1]]) inclusive  step (%[[STEP_LLVMIR1]]) {
+  ! LLVMIR: "llvm.loop.parallel_accesses"
+  ! LLVMIR: !{!"llvm.loop.vectorize.enable", i1 true}
+  ! LLVMIR: !{!"llvm.loop.vectorize.enable", i1 false}
   do i = 1, n
     ! CHECK: fir.store %[[I]] to %[[LOCAL:.*]] : !fir.ref<i32>
     ! CHECK: %[[LD:.*]] = fir.load %[[LOCAL]] : !fir.ref<i32>
     ! CHECK: fir.call @_FortranAioOutputInteger32({{.*}}, %[[LD]]) : (!fir.ref<i8>, i32) -> i1
     print*, i
   end do
   !$OMP END SIMD
 end subroutine

llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

public:
   /// loop body.
   ///
   /// \param Loc       The insert and source location description. The insert
   ///                  location can be between two instructions or the end of a
   ///                  degenerate block (e.g. a BB under construction).
   /// \param BodyGenCB Callback that will generate the loop body code.
   /// \param TripCount Number of iterations the loop body is executed.
   /// \param Name      Base name used to derive BB and instruction names.
+  /// \param Preheader If not null, place trip count values in that block.
   ///
   /// \returns An object representing the created control flow structure which
   ///          can be used for loop-associated directives.
   CanonicalLoopInfo *createCanonicalLoop(const LocationDescription &Loc,
                                          LoopBodyGenCallbackTy BodyGenCB,
                                          Value *TripCount,
-                                         const Twine &Name = "loop");
+                                         const Twine &Name = "loop",
+                                         BasicBlock *Preheader = nullptr);
 
   /// Generator for the control flow structure of an OpenMP canonical loop.
   ///
   /// Instead of a logical iteration space, this allows specifying user-defined
   /// loop counter values using increment, upper- and lower bounds. To
   /// disambiguate the terminology when counting downwards, instead of lower
   /// bounds we use \p Start for the loop counter value in the first body
   /// iteration.
public:
   /// \param IsSigned  Whether Start, Stop and Step are signed integers.
   /// \param InclusiveStop Whether \p Stop itself is a valid value for the loop
   ///                      counter.
   /// \param ComputeIP Insertion point for instructions computing the trip
   ///                  count. Can be used to ensure the trip count is available
   ///                  at the outermost loop of a loop nest. If not set,
   ///                  defaults to the preheader of the generated loop.
   /// \param Name      Base name used to derive BB and instruction names.
+  /// \param CreatePreheader If true, this function creates preheader block.
+  ///                  Trip count values should be placed in that block.
   ///
   /// \returns An object representing the created control flow structure which
   ///          can be used for loop-associated directives.
   CanonicalLoopInfo *createCanonicalLoop(const LocationDescription &Loc,
                                          LoopBodyGenCallbackTy BodyGenCB,
                                          Value *Start, Value *Stop, Value *Step,
                                          bool IsSigned, bool InclusiveStop,
                                          InsertPointTy ComputeIP = {},
-                                         const Twine &Name = "loop");
+                                         const Twine &Name = "loop",
+                                         bool CreatePreheader = false);
 
   /// Collapse a loop nest into a single loop.
   ///
   /// Merges loops of a loop nest into a single CanonicalLoopNest representation
   /// that has the same number of innermost loop iterations as the origin loop
   /// nest. The induction variables of the input loops are derived from the
   /// collapsed loop's induction variable. This is intended to be used to
   /// implement OpenMP's collapse clause. Before applying a directive,
public:
   ///                    to defer unrolling to the LoopUnrollPass.
   void unrollLoopPartial(DebugLoc DL, CanonicalLoopInfo *Loop, int32_t Factor,
                          CanonicalLoopInfo **UnrolledCLI);
 
   /// Add metadata to simd-ize a loop.
   ///
   /// \param Loop    The loop to simd-ize.
   /// \param Simdlen The Simdlen length to apply to the simd loop.
   void applySimd(CanonicalLoopInfo *Loop, ConstantInt *Simdlen);
 
+  /// Add metadata to simd-ize a loops.
+  ///
+  /// \param DL     Debug location for instructions added by unrolling.
+  /// \param Loops  The set of loops annotated with pragma simd directive
+  /// \param IfCond The optional if clause
+  void applySimd(DebugLoc DL, ArrayRef<CanonicalLoopInfo *> Loops,
+                 Value *IfCond = nullptr, ConstantInt *Simdlen = nullptr);
+
   /// Generator for '#omp flush'
   ///
   /// \param Loc The location where the flush directive was encountered
   void createFlush(const LocationDescription &Loc);
 
   /// Generator for '#omp taskwait'
   ///
   /// \param Loc The location where the taskwait directive was encountered.
   void createTaskwait(const LocationDescription &Loc);
 
   /// Generator for '#omp taskyield'
   ///
public:
   /// \param DL        DebugLoc used for the instructions in the skeleton.
   /// \param TripCount Value to be used for the trip count.
   /// \param F         Function in which to insert the BasicBlocks.
   /// \param PreInsertBefore  Where to insert BBs that execute before the body,
   ///                         typically the body itself.
   /// \param PostInsertBefore Where to insert BBs that execute after the body.
   /// \param Name      Base name used to derive BB
   ///                  and instruction names.
+  /// \param Preheader Preheader block. If it is not null, no new preheader
+  ///                  is created.
   ///
   /// \returns The CanonicalLoopInfo that represents the emitted loop.
-  CanonicalLoopInfo *createLoopSkeleton(DebugLoc DL, Value *TripCount,
-                                        Function *F,
-                                        BasicBlock *PreInsertBefore,
-                                        BasicBlock *PostInsertBefore,
-                                        const Twine &Name = {});
+  CanonicalLoopInfo *
+  createLoopSkeleton(DebugLoc DL, Value *TripCount, Function *F,
+                     BasicBlock *PreInsertBefore, BasicBlock *PostInsertBefore,
+                     const Twine &Name = {}, BasicBlock *Preheader = nullptr);
 };
 
 /// Class to represented the control flow structure of an OpenMP canonical loop.
 ///
 /// The control-flow structure is standardized for easy consumption by
 /// directives associated with loops. For instance, the worksharing-loop
 /// construct may change this control flow such that each loop iteration is
 /// executed on only one thread. The constraints of a canonical loop in brief

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

 #include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/IR/Value.h"
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Target/TargetOptions.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/CodeExtractor.h"
 #include "llvm/Transforms/Utils/LoopPeel.h"
 #include "llvm/Transforms/Utils/UnrollLoop.h"
 
 #include <cstdint>
 
 #define DEBUG_TYPE "openmp-ir-builder"
 
OpenMPIRBuilder::createMasked(const LocationDescription &Loc,
   Instruction *ExitCall = Builder.CreateCall(ExitRTLFn, ArgsEnd);
 
   return EmitOMPInlinedRegion(OMPD, EntryCall, ExitCall, BodyGenCB, FiniCB,
                               /*Conditional*/ true, /*hasFinalize*/ true);
 }
 
 CanonicalLoopInfo *OpenMPIRBuilder::createLoopSkeleton(
     DebugLoc DL, Value *TripCount, Function *F, BasicBlock *PreInsertBefore,
-    BasicBlock *PostInsertBefore, const Twine &Name) {
+    BasicBlock *PostInsertBefore, const Twine &Name, BasicBlock *Preheader) {
   Module *M = F->getParent();
   LLVMContext &Ctx = M->getContext();
   Type *IndVarTy = TripCount->getType();
+  bool PreheaderCreated = false;
 
   // Create the basic block structure.
-  BasicBlock *Preheader =
-      BasicBlock::Create(Ctx, "omp_" + Name + ".preheader", F, PreInsertBefore);
+  if (!Preheader) {
+    PreheaderCreated = true;
+    Preheader = BasicBlock::Create(Ctx, "omp_" + Name + ".preheader", F,
+                                   PreInsertBefore);
+  }
   BasicBlock *Header =
       BasicBlock::Create(Ctx, "omp_" + Name + ".header", F, PreInsertBefore);
   BasicBlock *Cond =
       BasicBlock::Create(Ctx, "omp_" + Name + ".cond", F, PreInsertBefore);
   BasicBlock *Body =
       BasicBlock::Create(Ctx, "omp_" + Name + ".body", F, PreInsertBefore);
   BasicBlock *Latch =
       BasicBlock::Create(Ctx, "omp_" + Name + ".inc", F, PostInsertBefore);
   BasicBlock *Exit =
       BasicBlock::Create(Ctx, "omp_" + Name + ".exit", F, PostInsertBefore);
   BasicBlock *After =
       BasicBlock::Create(Ctx, "omp_" + Name + ".after", F, PostInsertBefore);
 
   // Use specified DebugLoc for new instructions.
   Builder.SetCurrentDebugLocation(DL);
-
+  if (PreheaderCreated)
     Builder.SetInsertPoint(Preheader);
   Builder.CreateBr(Header);
 
   Builder.SetInsertPoint(Header);
   PHINode *IndVarPHI = Builder.CreatePHI(IndVarTy, 2, "omp_" + Name + ".iv");
   IndVarPHI->addIncoming(ConstantInt::get(IndVarTy, 0), Preheader);
   Builder.CreateBr(Cond);
 
   Builder.SetInsertPoint(Cond);
CanonicalLoopInfo *OpenMPIRBuilder::createLoopSkeleton(
   CL->Exit = Exit;
 
 #ifndef NDEBUG
   CL->assertOK();
 #endif
   return CL;
 }
 
-CanonicalLoopInfo *
-OpenMPIRBuilder::createCanonicalLoop(const LocationDescription &Loc,
-                                     LoopBodyGenCallbackTy BodyGenCB,
-                                     Value *TripCount, const Twine &Name) {
+CanonicalLoopInfo *OpenMPIRBuilder::createCanonicalLoop(
+    const LocationDescription &Loc, LoopBodyGenCallbackTy BodyGenCB,
+    Value *TripCount, const Twine &Name, BasicBlock *Preheader) {
   BasicBlock *BB = Loc.IP.getBlock();
   BasicBlock *NextBB = BB->getNextNode();
+  CanonicalLoopInfo *CL;
+  if (Preheader)
+    CL = createLoopSkeleton(Loc.DL, TripCount, BB->getParent(),
+                            Preheader->getNextNode(), Preheader->getNextNode(),
+                            Name, Preheader);
+  else
+    CL = createLoopSkeleton(Loc.DL, TripCount, BB->getParent(), NextBB, NextBB,
+                            Name);
 
-  CanonicalLoopInfo *CL = createLoopSkeleton(Loc.DL, TripCount, BB->getParent(),
-                                             NextBB, NextBB, Name);
   BasicBlock *After = CL->getAfter();
 
   // If location is not set, don't connect the loop.
   if (updateToLocation(Loc)) {
     // Split the loop at the insertion point: Branch to the preheader and move
     // every following instruction to after the loop (the After BB). Also, the
     // new successor is the loop's after block.
     spliceBB(Builder, After, /*CreateBranch=*/false);
     Builder.CreateBr(CL->getPreheader());
   }
 
   // Emit the body content. We do it after connecting the loop to the CFG to
   // avoid that the callback encounters degenerate BBs.
   BodyGenCB(CL->getBodyIP(), CL->getIndVar());
 
 #ifndef NDEBUG
   CL->assertOK();
 #endif
   return CL;
 }
 
 CanonicalLoopInfo *OpenMPIRBuilder::createCanonicalLoop(
     const LocationDescription &Loc, LoopBodyGenCallbackTy BodyGenCB,
     Value *Start, Value *Stop, Value *Step, bool IsSigned, bool InclusiveStop,
-    InsertPointTy ComputeIP, const Twine &Name) {
+    InsertPointTy ComputeIP, const Twine &Name, bool CreatePreheader) {
 
   // Consider the following difficulties (assuming 8-bit signed integers):
   //  * Adding \p Step to the loop counter which passes \p Stop may overflow:
   //      DO I = 1, 100, 50
   ///  * A \p Step of INT_MIN cannot not be normalized to a positive direction:
   //      DO I = 100, 0, -128
 
   // Start, Stop and Step must be of the same integer type.
   auto *IndVarTy = cast<IntegerType>(Start->getType());
   assert(IndVarTy == Stop->getType() && "Stop type mismatch");
   assert(IndVarTy == Step->getType() && "Step type mismatch");
-
+  BasicBlock *Preheader = nullptr;
+  if (CreatePreheader) {
+    Module *M = Loc.IP.getBlock()->getParent()->getParent();
+    Function *F = Loc.IP.getBlock()->getParent();
+    LLVMContext &Ctx = M->getContext();
+    Preheader = BasicBlock::Create(Ctx, "omp_" + Name + ".preheader", F,
+                                   Loc.IP.getBlock()->getNextNode());
+  }
   LocationDescription ComputeLoc =
       ComputeIP.isSet() ? LocationDescription(ComputeIP, Loc.DL) : Loc;
   updateToLocation(ComputeLoc);
 
   ConstantInt *Zero = ConstantInt::get(IndVarTy, 0);
   ConstantInt *One = ConstantInt::get(IndVarTy, 1);
+  if (CreatePreheader)
+    Builder.SetInsertPoint(Preheader);
 
   // Like Step, but always positive.
   Value *Incr = Step;
 
   // Distance between Start and Stop; always positive.
   Value *Span;
 
   // Condition whether there are no iterations are executed at all, e.g. because
CanonicalLoopInfo *OpenMPIRBuilder::createCanonicalLoop(
 
   auto BodyGen = [=](InsertPointTy CodeGenIP, Value *IV) {
     Builder.restoreIP(CodeGenIP);
     Value *Span = Builder.CreateMul(IV, Step);
     Value *IndVar = Builder.CreateAdd(Span, Start);
     BodyGenCB(Builder.saveIP(), IndVar);
   };
   LocationDescription LoopLoc = ComputeIP.isSet() ? Loc.IP : Builder.saveIP();
-  return createCanonicalLoop(LoopLoc, BodyGen, TripCount, Name);
+  return createCanonicalLoop(LoopLoc, BodyGen, TripCount, Name, Preheader);
 }
 
 // Returns an LLVM function to call for initializing loop bounds using OpenMP
 // static scheduling depending on `type`. Only i32 and i64 are supported by the
 // runtime. Always interpret integers as unsigned similarly to
 // CanonicalLoopInfo.
 static FunctionCallee getKmpcForStaticInitForType(Type *Ty, Module &M,
                                                   OpenMPIRBuilder &OMPBuilder) {
OpenMPIRBuilder::collapseLoops(DebugLoc DL, ArrayRef<CanonicalLoopInfo *> Loops,
   // Loop control blocks that may become orphaned later.
   SmallVector<BasicBlock *, 12> OldControlBBs;
   OldControlBBs.reserve(6 * Loops.size());
   for (CanonicalLoopInfo *Loop : Loops)
     Loop->collectControlBlocks(OldControlBBs);
 
   // Setup the IRBuilder for inserting the trip count computation.
   Builder.SetCurrentDebugLocation(DL);
-  if (ComputeIP.isSet())
+  if (ComputeIP.isSet()) {
     Builder.restoreIP(ComputeIP);
-  else
+  } else {
     Builder.restoreIP(Outermost->getPreheaderIP());
+  }
 
   // Derive the collapsed' loop trip count.
   // TODO: Find common/largest indvar type.
   Value *CollapsedTripCount = nullptr;
   for (CanonicalLoopInfo *L : Loops) {
     assert(L->isValid() &&
            "All loops to collapse must be valid canonical loops");
     Value *OrigTripCount = L->getTripCount();
     if (!CollapsedTripCount) {
#ifndef NDEBUG
   for (CanonicalLoopInfo *GenL : Result)
     GenL->assertOK();
 #endif
   return Result;
 }
 
 /// Attach loop metadata \p Properties to the loop described by \p Loop. If the
 /// loop already has metadata, the loop properties are appended.
+static void addLoopMetadata(Loop *Loop, ArrayRef<Metadata *> Properties) {
+  assert(Loop->isValid() && "Expecting a valid CanonicalLoopInfo");
+
+  // Nothing to do if no property to attach.
+  if (Properties.empty())
+    return;
+
+  // If the loop already has metadata, prepend it to the new metadata.
+  BasicBlock *Latch = Loop->getLoopLatch();
+
+  LLVMContext &Ctx = Latch->getParent()->getContext();
+  SmallVector<Metadata *> NewLoopProperties;
+  NewLoopProperties.push_back(nullptr);
+
+  assert(Latch && "A valid CanonicalLoopInfo must have a unique latch");
+  MDNode *Existing = Latch->getTerminator()->getMetadata(LLVMContext::MD_loop);
+  if (Existing)
+    append_range(NewLoopProperties, drop_begin(Existing->operands(), 1));
+
+  append_range(NewLoopProperties, Properties);
+  MDNode *LoopID = MDNode::getDistinct(Ctx, NewLoopProperties);
+  LoopID->replaceOperandWith(0, LoopID);
+
+  Latch->getTerminator()->setMetadata(LLVMContext::MD_loop, LoopID);
+}
+
+/// Attach loop metadata \p Properties to the loop described by \p Loop. If the
+/// loop already has metadata, the loop properties are appended.
 static void addLoopMetadata(CanonicalLoopInfo *Loop,
                             ArrayRef<Metadata *> Properties) {
   assert(Loop->isValid() && "Expecting a valid CanonicalLoopInfo");
 
   // Nothing to do if no property to attach.
   if (Properties.empty())
     return;
 

 void OpenMPIRBuilder::unrollLoopHeuristic(DebugLoc, CanonicalLoopInfo *Loop) {
   LLVMContext &Ctx = Builder.getContext();
   addLoopMetadata(
       Loop, {
                 MDNode::get(Ctx, MDString::get(Ctx, "llvm.loop.unroll.enable")),
             });
 }
 
+void OpenMPIRBuilder::applySimd(DebugLoc DL,
+                                ArrayRef<CanonicalLoopInfo *> Loops,
+                                Value *IfCond, ConstantInt *Simdlen) {
+  LLVMContext &Ctx = Builder.getContext();
+  FunctionAnalysisManager FAM;
+  Function *F = Loops.front()->getFunction();
+  FAM.registerPass([]() { return DominatorTreeAnalysis(); });
+  FAM.registerPass([]() { return LoopAnalysis(); });
+  FAM.registerPass([]() { return PassInstrumentationAnalysis(); });
+
+  LoopAnalysis LIA;
+  LoopInfo &&LI = LIA.run(*F, FAM);
+  DominatorTreeAnalysis DTA;
+  DominatorTree &&DT = DTA.run(*F, FAM);
+
+  CanonicalLoopInfo *CollapsedLoop = collapseLoops(DL, Loops, {});
+  Loop *L = LI.getLoopFor(CollapsedLoop->getHeader());
+  if (IfCond) {
+    Instruction *SplitBefore;
+    if (Instruction::classof(IfCond)) {
+      SplitBefore = dyn_cast<Instruction>(IfCond);
+    } else {
+      CanonicalLoopInfo *FrontLoop = Loops.front();
+      SplitBefore = FrontLoop->getPreheader()->getTerminator();
+    }
+    BasicBlock *Head = SplitBefore->getParent();
+
+    Instruction *HeadOldTerm = Head->getTerminator();
+    llvm::LLVMContext &C = Head->getContext();
+    llvm::BasicBlock *ThenBlock = llvm::BasicBlock::Create(
+        C, "simd.if.then", Head->getParent(), Head->getNextNode());
+    llvm::BasicBlock *ElseBlock = llvm::BasicBlock::Create(
+        C, "simd.if.else", Head->getParent(), Loops.front()->getExit());
+    SmallVector<BasicBlock *, 8> PostLoopBlocks;
+    Loop *PostLoop;
+    ValueToValueMapTy VMap;
+
+    // Create if condition branch.
+    Builder.SetInsertPoint(HeadOldTerm);
+    Instruction *BrInstr =
+        Builder.CreateCondBr(IfCond, ThenBlock, /*ifFalse*/ ElseBlock);
+    InsertPointTy IP{BrInstr->getParent(), ++BrInstr->getIterator()};
+    // Then block contains branch to omp loop which needs to be vectorized
+    spliceBB(IP, ThenBlock, false);
+    ThenBlock->replaceSuccessorsPhiUsesWith(Head, ThenBlock);
+
+    // Else block contains branch to cloned omp loop which is not vectorized
+    PostLoop = cloneLoopWithPreheader(CollapsedLoop->getExit(),
+                                      CollapsedLoop->getPreheader(), L, VMap,
+                                      ".else", &LI, &DT, PostLoopBlocks);
+    remapInstructionsInBlocks(PostLoopBlocks, VMap);
+    Builder.SetInsertPoint(ElseBlock);
+    Builder.CreateBr(PostLoop->getLoopPreheader());
+    ConstantAsMetadata *BoolConst =
+        ConstantAsMetadata::get(ConstantInt::getFalse(Type::getInt1Ty(Ctx)));
+    // Disable vectorization if the simd condition is not met
+    addLoopMetadata(
+        PostLoop,
+        {MDNode::get(Ctx, {MDString::get(Ctx, "llvm.loop.vectorize.enable"),
+                           BoolConst})});
+  }
+  // Add simd metadata to the loop
+  applySimd(CollapsedLoop, Simdlen);
+}
+
 void OpenMPIRBuilder::applySimd(CanonicalLoopInfo *CanonicalLoop,
                                 ConstantInt *Simdlen) {
   LLVMContext &Ctx = Builder.getContext();
 
   Function *F = CanonicalLoop->getFunction();
 
   FunctionAnalysisManager FAM;
   FAM.registerPass([]() { return DominatorTreeAnalysis(); });
   FAM.registerPass([]() { return LoopAnalysis(); });
   FAM.registerPass([]() { return PassInstrumentationAnalysis(); });
 
   LoopAnalysis LIA;
   LoopInfo &&LI = LIA.run(*F, FAM);
 
   Loop *L = LI.getLoopFor(CanonicalLoop->getHeader());
 
   SmallSet<BasicBlock *, 8> Reachable;
 
   // Get the basic blocks from the loop in which memref instructions
   // can be found.
   // TODO: Generalize getting all blocks inside a CanonicalizeLoopInfo,
   // preferably without running any passes.
   for (BasicBlock *Block : L->getBlocks()) {
     if (Block == CanonicalLoop->getCond() ||
         Block == CanonicalLoop->getHeader())
       continue;
     Reachable.insert(Block);
   }
 
   // Add access group metadata to memory-access instructions.
   MDNode *AccessGroup = MDNode::getDistinct(Ctx, {});
   for (BasicBlock *BB : Reachable)
     addSimdMetadata(BB, AccessGroup, LI);
 
   // Use the above access group metadata to create loop level
   // metadata, which should be distinct for each loop.
   ConstantAsMetadata *BoolConst =

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

 #include "mlir/Support/LLVM.h"
 #include "mlir/Target/LLVMIR/ModuleTranslation.h"
 
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Frontend/OpenMP/OMPIRBuilder.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/IRBuilder.h"
+#include "llvm/Transforms/Utils/BasicBlockUtils.h"
 
 using namespace mlir;
 
 namespace {
 static llvm::omp::ScheduleKind
 convertToScheduleKind(Optional<omp::ClauseScheduleKind> schedKind) {
   if (!schedKind.has_value())
     return llvm::omp::OMP_SCHEDULE_Default;
convertOmpSimdLoop(Operation &opInst, llvm::IRBuilderBase &builder,
   llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder);
 
   // Generator of the canonical loop body.
   // TODO: support error propagation in OpenMPIRBuilder and use it instead of
   // relying on captured variables.
   SmallVector<llvm::CanonicalLoopInfo *> loopInfos;
   SmallVector<llvm::OpenMPIRBuilder::InsertPointTy> bodyInsertPoints;
   LogicalResult bodyGenStatus = success();
-
-  // TODO: The code generation for if clause is not supported yet.
-  if (loop.if_expr())
-    return failure();
-
   auto bodyGen = [&](llvm::OpenMPIRBuilder::InsertPointTy ip, llvm::Value *iv) {
     // Make sure further conversions know about the induction variable.
     moduleTranslation.mapValue(
         loop.getRegion().front().getArgument(loopInfos.size()), iv);
 
     // Capture the body insertion point for use in nested loops. BodyIP of the
     // CanonicalLoopInfo always points to the beginning of the entry block of
     // the body.
     bodyInsertPoints.push_back(ip);
 
     if (loopInfos.size() != loop.getNumLoops() - 1)
       return;
 
     // Convert the body of the loop.
     builder.restoreIP(ip);
     convertOmpOpRegions(loop.region(), "omp.simdloop.region", builder,
                         moduleTranslation, bodyGenStatus);
   };
-
   // Delegate actual loop construction to the OpenMP IRBuilder.
   // TODO: this currently assumes SimdLoop is semantically similar to SCF loop,
   // i.e. it has a positive step, uses signed integer semantics. Reconsider
   // this code when SimdLoop clearly supports more cases.
   llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
   for (unsigned i = 0, e = loop.getNumLoops(); i < e; ++i) {
     llvm::Value *lowerBound =
         moduleTranslation.lookupValue(loop.lowerBound()[i]);
     llvm::Value *upperBound =
         moduleTranslation.lookupValue(loop.upperBound()[i]);
     llvm::Value *step = moduleTranslation.lookupValue(loop.step()[i]);
 
     // Make sure loop trip count are emitted in the preheader of the outermost
     // loop at the latest so that they are all available for the new collapsed
     // loop will be created below.
     llvm::OpenMPIRBuilder::LocationDescription loc = ompLoc;
     llvm::OpenMPIRBuilder::InsertPointTy computeIP = ompLoc.IP;
-    if (i != 0) {
+    if (i == 0) {
+      // OMPBuilder should create loop preheader for the most outermost loop
+      // All trip count values will be placed in that block
+      loopInfos.push_back(ompBuilder->createCanonicalLoop(
+          loc, bodyGen, lowerBound, upperBound, step,
+          /*IsSigned=*/true, /*Inclusive=*/true, computeIP, "loop",
+          /*CreatePreheader=*/true));
+    } else {
+      // Preheader block for the most outermost loop has been created.
       loc = llvm::OpenMPIRBuilder::LocationDescription(bodyInsertPoints.back(),
                                                        ompLoc.DL);
       computeIP = loopInfos.front()->getPreheaderIP();
-    }
       loopInfos.push_back(ompBuilder->createCanonicalLoop(
           loc, bodyGen, lowerBound, upperBound, step,
-        /*IsSigned=*/true, /*Inclusive=*/true, computeIP));
-
+          /*IsSigned=*/true, /*Inclusive=*/true, computeIP, "loop",
+          /*CreatePreheader=*/false));
+    }
     if (failed(bodyGenStatus))
       return failure();
   }
 
   // Collapse loops.
-  llvm::IRBuilderBase::InsertPoint afterIP = loopInfos.front()->getAfterIP();
-  llvm::CanonicalLoopInfo *loopInfo =
-      ompBuilder->collapseLoops(ompLoc.DL, loopInfos, {});
-
-  ompBuilder->applySimd(loopInfo, nullptr);
-
-  builder.restoreIP(afterIP);
+  llvm::IRBuilderBase::InsertPoint AfterIP = loopInfos.front()->getAfterIP();
+  if (loop.if_expr()) {
+    ompBuilder->applySimd(ompLoc.DL, loopInfos,
+                          moduleTranslation.lookupValue(loop.if_expr()),
+                          nullptr);
+  } else {
+    ompBuilder->applySimd(ompLoc.DL, loopInfos);
+  }
+  builder.restoreIP(AfterIP);
   return success();
 }
 
 /// Convert an Atomic Ordering attribute to llvm::AtomicOrdering.
 llvm::AtomicOrdering
 convertAtomicOrdering(Optional<omp::ClauseMemoryOrderKind> ao) {
   if (!ao)
     return llvm::AtomicOrdering::Monotonic; // Default Memory Ordering
 
   switch (*ao) {
   case omp::ClauseMemoryOrderKind::Seq_cst:
     return llvm::AtomicOrdering::SequentiallyConsistent;
   case omp::ClauseMemoryOrderKind::Acq_rel:
     return llvm::AtomicOrdering::AcquireRelease;
   case omp::ClauseMemoryOrderKind::Acquire:
     return llvm::AtomicOrdering::Acquire;
   case omp::ClauseMemoryOrderKind::Release:
     return llvm::AtomicOrdering::Release;
   case omp::ClauseMemoryOrderKind::Relaxed:
     return llvm::AtomicOrdering::Monotonic;
   }