[mlir][OpenMP] Allow composite SIMD REDUCTION and IF (#147568)

Reduction support: https://github.com/llvm/llvm-project/pull/146671
If Support is fixed in this PR

The problem for the IF clause in composite constructs was that wsloop
and simd both operate on the same CanonicalLoopInfo structure: with the
SIMD processed first, followed by the wsloop. Previously the IF clause
generated code like
```
if (cond) {
  while (...) {
    simd_loop_body;
  }
} else {
  while (...) {
    nonsimd_loop_body;
  }
}
```
The problem with this is that this invalidates the CanonicalLoopInfo
structure to be processed by the wsloop later. To avoid this, in this
patch I preserve the original loop, moving the IF clause inside of the
loop:
```
while (...) {
  if (cond) {
    simd_loop_body;
  } else {
    non_simd_loop_body;
  }
}
```
On simple examples I tried LLVM was able to hoist the if condition
outside of the loop at -O3.

The disadvantage of this is that we cannot add the
llvm.loop.vectorize.enable attribute on either the SIMD or non-SIMD
loops because they both share a loop back edge. There's no way of
solving this without keeping the old design of having two different
loops: which cannot be represented using only one CanonicalLoopInfo
structure. I don't think the presence or absence of this attribute makes
much difference. In my testing it is the llvm.loop.parallel_access
metadata which makes the difference to vectorization. LLVM will
vectorize if legal whether or not this attribute is there in the TRUE
branch. In the FALSE branch this means the loop might be vectorized even
when the condition is false: but I think this is still standards
compliant: OpenMP 6.0 says that when the if clause is false that should
be treated like the SIMDLEN clause is one. The SIMDLEN clause is defined
as a "hint". For the same reason, SIMDLEN and SAFELEN clauses are
silently ignored when SIMD IF is used.

I think it is better to implement SIMD IF and ignore SIMDLEN and SAFELEN
and some vectorization encouragement metadata when combined with IF than
to ignore IF because IF could have correctness consequences whereas the
rest are optimiztion hints. For example, the user might use the IF
clause to disable SIMD programatically when it is known not safe to
vectorize the loop. In this case it is not at all safe to add the
parallel access or SAFELEN metadata.

1 files changed, 22 insertions, 11 deletions

diff --git a/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp b/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp
index be98be2..d019ade 100644
--- a/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp
+++ b/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp
@@ -2242,23 +2242,34 @@ TEST_F(OpenMPIRBuilderTest, ApplySimdIf) {
   PB.registerFunctionAnalyses(FAM);
   LoopInfo &LI = FAM.getResult<LoopAnalysis>(*F);
 
-  // Check if there are two loops (one with enabled vectorization)
+  // Check if there is one loop containing branches with and without
+  // vectorization
   const std::vector<Loop *> &TopLvl = LI.getTopLevelLoops();
-  EXPECT_EQ(TopLvl.size(), 2u);
+  EXPECT_EQ(TopLvl.size(), 1u);
 
   Loop *L = TopLvl[0];
   EXPECT_TRUE(findStringMetadataForLoop(L, "llvm.loop.parallel_accesses"));
-  EXPECT_TRUE(getBooleanLoopAttribute(L, "llvm.loop.vectorize.enable"));
-  EXPECT_EQ(getIntLoopAttribute(L, "llvm.loop.vectorize.width"), 3);
-
-  // The second loop should have disabled vectorization
-  L = TopLvl[1];
-  EXPECT_FALSE(findStringMetadataForLoop(L, "llvm.loop.parallel_accesses"));
+  // These attributes cannot not be set because the loop is shared between simd
+  // and non-simd versions
   EXPECT_FALSE(getBooleanLoopAttribute(L, "llvm.loop.vectorize.enable"));
-  // Check for llvm.access.group metadata attached to the printf
-  // function in the loop body.
+  EXPECT_EQ(getIntLoopAttribute(L, "llvm.loop.vectorize.width"), 0);
+
+  // Check for if condition
   BasicBlock *LoopBody = CLI->getBody();
-  EXPECT_TRUE(any_of(*LoopBody, [](Instruction &I) {
+  BranchInst *IfCond = cast<BranchInst>(LoopBody->getTerminator());
+  EXPECT_EQ(IfCond->getCondition(), IfCmp);
+  BasicBlock *TrueBranch = IfCond->getSuccessor(0);
+  BasicBlock *FalseBranch = IfCond->getSuccessor(1)->getUniqueSuccessor();
+
+  // Check for llvm.access.group metadata attached to the printf
+  // function in the true body.
+  EXPECT_TRUE(any_of(*TrueBranch, [](Instruction &I) {
+    return I.getMetadata("llvm.access.group") != nullptr;
+  }));
+
+  // Check for llvm.access.group metadata attached to the printf
+  // function in the false body.
+  EXPECT_FALSE(any_of(*FalseBranch, [](Instruction &I) {
     return I.getMetadata("llvm.access.group") != nullptr;
   }));
 }