[mlir][OpenMP] Add if clause to OpenMP simd construct

[domada]

Summary

Added if clause support for simd loop construct in definition of OpenMP MLIR dialect.

[arnamoy10]

Thanks for working on this patch. I think we need to add a testcase in flang/test/Lower/OpenMP/simd.f90

[shraiysh]

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

441

1. Is this different from wsloop control? If not, can we reuse it as LoopControl?
2. Without a token like for before custom<LoopControl>, the error messages for wrong syntax (like missing loop control) will be very confusing. I would suggest adding for before the custom assembly format to have legible error messages. See omp.wsloop's assemblyFormat for reference.

[domada]

Added test case for simd directive with if clause in file: flang/test/Lower/OpenMP/simd.f90

[arnamoy10]

flang/test/Lower/OpenMP/simd.f90

25 ↗

(On Diff #442026)

It will be nice to see an example with more realistic looking test case e.g. if (n > threshold), where n and threshold are both arguments, and n is the loop upper bound.

[domada]

Addressed review remarks

[domada]

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

441

Ad 1. wsloop control can contain optional inclusive keyword. Is this keyword also allowed for simd loop? I assumed that inclusive keyword is not allowed for simd construct. Moreover, the list of attributes for WsLoop is different than for SIMD loop. That's why I used different method.
Ad 2. Good point -> I will add it.

441

Ad 1. Done -> Added inclusive keyword as well
Ad 2. Done

[shraiysh]

LGTM. Thank you!

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

419

nit: this has to wrap at 80 columns. it looks like it has wrapped much earlier.

[domada]

Fixed layout of simd description

[peixin]

The OMPIRBuilder and lowering to LLVMIR for if clause of SIMD construct is not supported yet. If you merge this before implementing those, users will get unexpected results for a fortran code since it behaves as if there is no if clause.

Usually, we implement the OMPIRBuilder and lowering to LLVMIR first in case some addition process needs to be done in lowering to MLIR. I would recommend you split this patch into two: 1. MLIR operation def (can be merged now) 2. lowering to MLIR (delayed after supporting OMPIRBuilder and lowering to LLVMIR). It's better to verify that the generated MLIR for if clause of SIMD construct in this patch works OK in OMPIRBuilder.

Also, I noticed there should be TODO for other clauses of SIMD construct. I think the lowering of clauses really should be refactored. @shraiysh Do you plan to continue that work? If yes, please add TODO for clauses of SIMD construct.

flang/lib/Lower/OpenMP.cpp

831–839

[shraiysh]

@shraiysh Do you plan to continue that work?

Not at the moment no. I am working on other tasks at the moment and might not be able to spend time on it for a couple weeks. If there is an urgent requirement, it's okay if someone else commandeers that revision - D121583.

[domada]

flang/lib/Lower/OpenMP.cpp

831–839

Do you propose to move this if statement inside for loop: for (const auto &clause : loopOpClauseList.v) ?

[kiranchandramohan]

In D128940#3629981, @peixin wrote:

The OMPIRBuilder and lowering to LLVMIR for if clause of SIMD construct is not supported yet. If you merge this before implementing those, users will get unexpected results for a fortran code since it behaves as if there is no if clause.

Usually, we implement the OMPIRBuilder and lowering to LLVMIR first in case some addition process needs to be done in lowering to MLIR. I would recommend you split this patch into two: 1. MLIR operation def (can be merged now) 2. lowering to MLIR (delayed after supporting OMPIRBuilder and lowering to LLVMIR). It's better to verify that the generated MLIR for if clause of SIMD construct in this patch works OK in OMPIRBuilder.

Is it OK to add a return failure in the translation of SIMD op if there is an if clause operand? Links below to the translation place of SIMD op and an existing return failure.
https://github.com/llvm/llvm-project/blob/e6ff553979e850eeb7f0bbe77deab1c88fc764b3/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp#L903
https://github.com/llvm/llvm-project/blob/e6ff553979e850eeb7f0bbe77deab1c88fc764b3/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp#L570

Also, I noticed there should be TODO for other clauses of SIMD construct. I think the lowering of clauses really should be refactored. @shraiysh Do you plan to continue that work? If yes, please add TODO for clauses of SIMD construct.

For now, may be @domada can add a static function that returns an MLIR value given an If clause. This static function can be later moved to the refactored class.

[peixin]

Is it OK to add a return failure in the translation of SIMD op if there is an if clause operand? Links below to the translation place of SIMD op and an existing return failure.
https://github.com/llvm/llvm-project/blob/e6ff553979e850eeb7f0bbe77deab1c88fc764b3/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp#L903
https://github.com/llvm/llvm-project/blob/e6ff553979e850eeb7f0bbe77deab1c88fc764b3/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp#L570

Good point. This is also one option.

Also, I noticed there should be TODO for other clauses of SIMD construct. I think the lowering of clauses really should be refactored. @shraiysh Do you plan to continue that work? If yes, please add TODO for clauses of SIMD construct.

For now, may be @domada can add a static function that returns an MLIR value given an If clause. This static function can be later moved to the refactored class.

Agree.

flang/lib/Lower/OpenMP.cpp

831–839

Sorry, I mean the following:

for (const auto &clause : loopOpClauseList.v) {
  if (const auto &scheduleClause =
          std::get_if<Fortran::parser::OmpClause::Schedule>(&clause.u)) {
    if (const auto &chunkExpr =
            std::get<std::optional<Fortran::parser::ScalarIntExpr>>(
                scheduleClause->v.t)) {
      if (const auto *expr = Fortran::semantics::GetExpr(*chunkExpr)) {
        scheduleChunkClauseOperand =
            fir::getBase(converter.genExprValue(*expr, stmtCtx));
      }
    }
  } else if (const auto &ifClause = std::get_if<Fortran::parser::OmpClause::If>(&clause.u)) {
    auto &expr = std::get<Fortran::parser::ScalarLogicalExpr>(ifClause->v.t);
    mlir::Value ifVal = fir::getBase(
        converter.genExprValue(*Fortran::semantics::GetExpr(expr), stmtCtx));
    ifClauseOperand = firOpBuilder.createConvert(
        currentLocation, firOpBuilder.getI1Type(), ifVal);
  }
}

[domada]

Added check for unsupported lowering if clause to LLVM IR.
Refactor Flang-MLIR lowering. Extracted function for handling if clause.

[peixin]

LGTM

[kiranchandramohan]

LGTM. Thanks and Welcome to the OpenMP for Flang team!

Nit: You can also probably mention in the commit message that,
-> the patch also lowers the clause from the Flang parse-tree to MLIR, the translation to LLVM IR will be added in a later patch.
-> Reuses the custom handling for LoopControl in the assembly-format.
-> Includes the inclusive attribute.

flang/lib/Lower/OpenMP.cpp

 static void genOMP(Fortran::lower::AbstractConverter &converter,
                    Fortran::lower::pft::Evaluation &eval,
                    const Fortran::parser::OpenMPLoopConstruct &loopConstruct) {
 
   fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
   mlir::Location currentLocation = converter.getCurrentLocation();
   llvm::SmallVector<mlir::Value> lowerBound, upperBound, step, linearVars,
       linearStepVars, reductionVars;
-  mlir::Value scheduleChunkClauseOperand;
+  mlir::Value scheduleChunkClauseOperand, ifClauseOperand;
   mlir::Attribute scheduleClauseOperand, noWaitClauseOperand,
       orderedClauseOperand, orderClauseOperand;
+  Fortran::lower::StatementContext stmtCtx;
   const auto &loopOpClauseList = std::get<Fortran::parser::OmpClauseList>(
       std::get<Fortran::parser::OmpBeginLoopDirective>(loopConstruct.t).t);
 
   const auto ompDirective =
       std::get<Fortran::parser::OmpLoopDirective>(
           std::get<Fortran::parser::OmpBeginLoopDirective>(loopConstruct.t).t)
           .v;
   if (llvm::omp::OMPD_parallel_do == ompDirective) {
static void genOMP(Fortran::lower::AbstractConverter &converter,
 
   for (const auto &clause : loopOpClauseList.v) {
     if (const auto &scheduleClause =
             std::get_if<Fortran::parser::OmpClause::Schedule>(&clause.u)) {
       if (const auto &chunkExpr =
               std::get<std::optional<Fortran::parser::ScalarIntExpr>>(
                   scheduleClause->v.t)) {
         if (const auto *expr = Fortran::semantics::GetExpr(*chunkExpr)) {
-          Fortran::lower::StatementContext stmtCtx;
           scheduleChunkClauseOperand =
               fir::getBase(converter.genExprValue(*expr, stmtCtx));
         }
       }
     }
+    if (const auto &ifClause =
+            std::get_if<Fortran::parser::OmpClause::If>(&clause.u)) {
+      auto &expr = std::get<Fortran::parser::ScalarLogicalExpr>(ifClause->v.t);
+      mlir::Value ifVal = fir::getBase(
+          converter.genExprValue(*Fortran::semantics::GetExpr(expr), stmtCtx));
+      ifClauseOperand = firOpBuilder.createConvert(
+          currentLocation, firOpBuilder.getI1Type(), ifVal);
+    }
   }
 
   // The types of lower bound, upper bound, and step are converted into the
   // type of the loop variable if necessary.
   mlir::Type loopVarType = getLoopVarType(converter, loopVarTypeSize);
   for (unsigned it = 0; it < (unsigned)lowerBound.size(); it++) {
     lowerBound[it] = firOpBuilder.createConvert(currentLocation, loopVarType,
                                                 lowerBound[it]);
     upperBound[it] = firOpBuilder.createConvert(currentLocation, loopVarType,
                                                 upperBound[it]);
     step[it] =
         firOpBuilder.createConvert(currentLocation, loopVarType, step[it]);
   }
 
   // 2.9.3.1 SIMD construct
   // TODO: Support all the clauses
   if (llvm::omp::OMPD_simd == ompDirective) {
     TypeRange resultType;
     auto SimdLoopOp = firOpBuilder.create<mlir::omp::SimdLoopOp>(
-        currentLocation, resultType, lowerBound, upperBound, step);
+        currentLocation, resultType, lowerBound, upperBound, step,
+        ifClauseOperand);
     createBodyOfOp<omp::SimdLoopOp>(SimdLoopOp, converter, currentLocation,
                                     eval, &loopOpClauseList, iv);
     return;
   }
 
   // FIXME: Add support for following clauses:
   // 1. linear
   // 2. order

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

let description = [{
     The simd construct can be applied to a loop to indicate that the loop can be 
     transformed into a SIMD loop (that is, multiple iterations of the loop can 
     be executed concurrently using SIMD instructions).. The lower and upper 
     bounds specify a half-open range: the range includes the lower bound but 
     does not include the upper bound.
 
     The body region can contain any number of blocks. The region is terminated
     by "omp.yield" instruction without operands.
+
+    When an if clause is present and evaluates to false,
+    the preferred number of iterations to be executed concurrently is one,
+    regardless of whether a simdlen clause is specified.
     ```
     omp.simdloop (%i1, %i2) : index = (%c0, %c0) to (%c10, %c10) 
                                       step (%c1, %c1) {
       // block operations
       omp.yield
     }
     ``` 
   }];
 
   // TODO: Add other clauses
   let arguments = (ins Variadic<IntLikeType>:$lowerBound,
              Variadic<IntLikeType>:$upperBound,
-             Variadic<IntLikeType>:$step);
+             Variadic<IntLikeType>:$step,
+             Optional<I1>:$if_expr
+     );
  
   let regions = (region AnyRegion:$region);
+  let assemblyFormat = [{
+    oilist(`if` `(` $if_expr `)`
+    ) custom<SimdLoopControl>($region, $lowerBound, $upperBound, $step,
+                                  type($step)) attr-dict
+  }];
 
   let extraClassDeclaration = [{
     /// Returns the number of loops in the simd loop nest.
     unsigned getNumLoops() { return lowerBound().size(); }
 
   }];
  
   let hasCustomAssemblyFormat = 1;

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

 /// Parses an OpenMP Simd construct [2.9.3.1]
 ///
 /// simdloop ::= `omp.simdloop` loop-control clause-list
 /// loop-control ::= `(` ssa-id-list `)` `:` type `=`  loop-bounds
 /// loop-bounds := `(` ssa-id-list `)` to `(` ssa-id-list `)` steps
 /// steps := `step` `(`ssa-id-list`)`
 /// clause-list ::= clause clause-list | empty
 /// clause ::= TODO
-ParseResult SimdLoopOp::parse(OpAsmParser &parser, OperationState &result) {
+ParseResult parseSimdLoopControl(
+    OpAsmParser &parser, Region &region,
+    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &lowerBound,
+    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &upperBound,
+    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &steps,
+    SmallVectorImpl<Type> &loopVarTypes) {
   // Parse an opening `(` followed by induction variables followed by `)`
   SmallVector<OpAsmParser::Argument> ivs;
   Type loopVarType;
-  SmallVector<OpAsmParser::UnresolvedOperand> lower, upper, steps;
   if (parser.parseArgumentList(ivs, OpAsmParser::Delimiter::Paren) ||
       parser.parseColonType(loopVarType) ||
       // Parse loop bounds.
       parser.parseEqual() ||
-      parser.parseOperandList(lower, ivs.size(),
+      parser.parseOperandList(lowerBound, ivs.size(),
                               OpAsmParser::Delimiter::Paren) ||
-      parser.resolveOperands(lower, loopVarType, result.operands) ||
       parser.parseKeyword("to") ||
-      parser.parseOperandList(upper, ivs.size(),
-                              OpAsmParser::Delimiter::Paren) ||
-      parser.resolveOperands(upper, loopVarType, result.operands) ||
-      // Parse step values.
-      parser.parseKeyword("step") ||
-      parser.parseOperandList(steps, ivs.size(),
-                              OpAsmParser::Delimiter::Paren) ||
-      parser.resolveOperands(steps, loopVarType, result.operands))
+      parser.parseOperandList(upperBound, ivs.size(),
+                              OpAsmParser::Delimiter::Paren))
     return failure();
 
-  int numIVs = static_cast<int>(ivs.size());
-  SmallVector<int> segments{numIVs, numIVs, numIVs};
-  // TODO: Add parseClauses() when we support clauses
-  result.addAttribute("operand_segment_sizes",
-                      parser.getBuilder().getI32VectorAttr(segments));
+  // Parse step values.
+  if (parser.parseKeyword("step") ||
+      parser.parseOperandList(steps, ivs.size(), OpAsmParser::Delimiter::Paren))
+    return failure();
 
   // Now parse the body.
-  Region *body = result.addRegion();
+  loopVarTypes = SmallVector<Type>(ivs.size(), loopVarType);
   for (auto &iv : ivs)
     iv.type = loopVarType;
-  return parser.parseRegion(*body, ivs);
+  return parser.parseRegion(region, ivs);
 }
 
-void SimdLoopOp::print(OpAsmPrinter &p) {
-  auto args = getRegion().front().getArguments();
-  p << " (" << args << ") : " << args[0].getType() << " = (" << lowerBound()
-    << ") to (" << upperBound() << ") ";
-  p << "step (" << step() << ") ";
-
-  p.printRegion(region(), /*printEntryBlockArgs=*/false);
+void printSimdLoopControl(OpAsmPrinter &p, Operation *op, Region &region,
+                          ValueRange lowerBound, ValueRange upperBound,
+                          ValueRange steps, TypeRange loopVarTypes) {
+  auto args = region.front().getArguments();
+  p << " (" << args << ") : " << args[0].getType() << " = (" << lowerBound
+    << ") to (" << upperBound << ") ";
+  p << "step (" << steps << ") ";
+  p.printRegion(region, /*printEntryBlockArgs=*/false);
 }
 
 //===----------------------------------------------------------------------===//
 // Verifier for Simd construct [2.9.3.1]
 //===----------------------------------------------------------------------===//
 
 LogicalResult SimdLoopOp::verify() {
   if (this->lowerBound().empty()) {

mlir/test/Dialect/OpenMP/invalid.mlir

 
 // -----
 
 func.func @omp_simdloop(%lb : index, %ub : index, %step : i32) -> () {
   // expected-error @below {{op failed to verify that all of {lowerBound, upperBound, step} have same type}}
   "omp.simdloop" (%lb, %ub, %step) ({
     ^bb0(%iv: index):
       omp.yield
-  }) {operand_segment_sizes = dense<[1,1,1]> : vector<3xi32>} :
+  }) {operand_segment_sizes = dense<[1,1,1,0]> : vector<4xi32>} :
     (index, index, i32) -> () 
 
   return
 }
 
 // -----
 
 // expected-error @below {{op expects initializer region with one argument of the reduction type}}

mlir/test/Dialect/OpenMP/ops.mlir

 }
 
 // CHECK-LABEL: omp_simdloop
 func.func @omp_simdloop(%lb : index, %ub : index, %step : index) -> () {
   // CHECK: omp.simdloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
   "omp.simdloop" (%lb, %ub, %step) ({
     ^bb0(%iv: index):
       omp.yield
-  }) {operand_segment_sizes = dense<[1,1,1]> : vector<3xi32>} :
+  }) {operand_segment_sizes = dense<[1,1,1,0]> : vector<4xi32>} :
     (index, index, index) -> () 
 
   return
 }
 
-
 // CHECK-LABEL: omp_simdloop_pretty
 func.func @omp_simdloop_pretty(%lb : index, %ub : index, %step : index) -> () {
   // CHECK: omp.simdloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
   omp.simdloop (%iv) : index = (%lb) to (%ub) step (%step) {
     omp.yield
   }
   return
 }
 
+// CHECK-LABEL: omp_simdloop_pretty_if
+func.func @omp_simdloop_pretty_if(%lb : index, %ub : index, %step : index, %if_cond : i1) -> () {
+  // CHECK: omp.simdloop if(%{{.*}}) (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
+  omp.simdloop if(%if_cond) (%iv): index = (%lb) to (%ub) step (%step) {
+    omp.yield
+  }
+  return
+}
+
 // CHECK-LABEL: omp_simdloop_pretty_multiple
 func.func @omp_simdloop_pretty_multiple(%lb1 : index, %ub1 : index, %step1 : index, %lb2 : index, %ub2 : index, %step2 : index) -> () {
   // CHECK: omp.simdloop (%{{.*}}, %{{.*}}) : index = (%{{.*}}, %{{.*}}) to (%{{.*}}, %{{.*}}) step (%{{.*}}, %{{.*}})
   omp.simdloop (%iv1, %iv2) : index = (%lb1, %lb2) to (%ub1, %ub2) step (%step1, %step2) {
     omp.yield
   }
   return
 }

mlir/test/Target/LLVMIR/openmp-llvm.mlir

"omp.simdloop" (%lb, %ub, %step) ({
     ^bb0(%iv: i64):
       %3 = llvm.mlir.constant(2.000000e+00 : f32) : f32
       // The form of the emitted IR is controlled by OpenMPIRBuilder and
       // tested there. Just check that the right metadata is added.
       // CHECK: llvm.access.group
       %4 = llvm.getelementptr %arg0[%iv] : (!llvm.ptr<f32>, i64) -> !llvm.ptr<f32>
       llvm.store %3, %4 : !llvm.ptr<f32>
       omp.yield
-  }) {operand_segment_sizes = dense<[1,1,1]> : vector<3xi32>} :
+  }) {operand_segment_sizes = dense<[1,1,1,0]> : vector<4xi32>} :
     (i64, i64, i64) -> () 
 
   llvm.return
 }
 // CHECK: llvm.loop.parallel_accesses
 // CHECK-NEXT: llvm.loop.vectorize.enable
 
 // -----