[MLIR][OpenMP] Add support for basic SIMD construct

[arnamoy10]

Summary

Patch adds a new operation for the SIMD construct. The op is designed to be very similar to the existing wsloop operation, so that the CanonicalLoopInfo of OpenMPIRBuilder can be used.

[peixin]

You may need one test case in mlir/test/Target/LLVMIR/openmp-llvm.mlir.

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

348

2.9.3 is for simd directives. I think you are implementing simd contruct [2.9.3.1]. No construct called simd loop contruct. In addition, simd loop is kind of confusing since someone may think of it as worksharing-loop simd contruct.

[shraiysh]

Thanks for the patch @arnamoy10. Requesting a few changes.

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

379

The current head has removed such definitions of parsers/printers in favor of fields hasCustomAssemblyFormat and hasVerifier (It is a minor change). Can you please rebase and use them?

Another point is, can we use assemblyFormat for this? We can have a custom parser/printer for loop bounds if assembly format does not allow separating region arguments from the region itself. (@kiranchandramohan @clementval FYI, I was planning on doing the same with WSLoop assemblyFormat.) Thoughts?

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

896

There are no test for this. The tests should be added to mlir/test/Target/LLVMIR/openmp-llvm.mlir.

900

Can we instead move it to the verifier right away instead of adding a TODO.

903–908

There is a constructor for LocationDescription that just takes an IRBuilderBase. That can be used here.

[shraiysh]

Also, please run git-clang-format HEAD~ before submitting it to eliminate formatting related errors.

[arnamoy10]

Addressing reviewers' comments.

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

348

I think SimdLoopOp is OK. for worksharing-loop simd contruct, we can call the Op WsSimdLoopOp

379

Thanks, did the mods. For Assemblyformat addition I will wait until your patch for the wsloop

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

900

Thanks, added.

903–908

Thanks for the suggestion. However, I will be needing the diLoc in later code, therefore I need to keep the subprogram and diLoc variables.

[peixin]

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

348

OK. Agree.

[shraiysh]

Thanks for addressing the comments @arnamoy10. Generally looks good. A few nits and comments.

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

359

Nit: This line has over 80 characters.

385

Nit: This line has over 80 characters.

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

1372

Nit: Can we please put this closer to the parser/printer?

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

900

I think we can safely remove this from here now, as it will be checked in verifier.

903–908

You can use ompLoc.DL for that as done here.

mlir/test/Dialect/OpenMP/ops.mlir

345

Can we please have one test with multiple arguments in this syntax? More than one (%lb, %ub and %step).

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

668 ↗

(On Diff #412106)

Same here. Can we please have one test with multiple arguments in this syntax? More than one (%lb, %ub and %step).

[arnamoy10]

Addressing reviewers' comments

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

359

No, just 78. It is also clang-formatted

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

1372

Done, thank you

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

900

Done, thank you

903–908

Thanks, done

mlir/test/Dialect/OpenMP/ops.mlir

345

Thanks done

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

668 ↗

(On Diff #412106)

Thanks, done

[shraiysh]

LGTM.

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

359

I still see it as 82 characters (78 text + 4 spaces at the beginning). Unfortunately, clang-format does not format tablegen files so we have to do this manually.

[arnamoy10]

Rebase on main and fix line character count

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

let extraClassDeclaration = [{
     /// Returns the number of reduction variables.
     unsigned getNumReductionVars() { return reduction_vars().size(); }
   }];
   let parser = [{ return parseWsLoopOp(parser, result); }];
   let printer = [{ return printWsLoopOp(p, *this); }];
   let verifier = [{ return ::verifyWsLoopOp(*this); }];
 }
 
+//===----------------------------------------------------------------------===//
+// 2.9.3  SIMD Directive
+//===----------------------------------------------------------------------===//
+
+def SimdLoopOp : OpenMP_Op<"simdloop", [AttrSizedOperandSegments,
+                         AllTypesMatch<["lowerBound", "upperBound", "step"]>]> {
+ let summary = "simd loop construct";
+  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.
+    ```
+    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);
+ 
+  let regions = (region AnyRegion:$region);
+
+  let extraClassDeclaration = [{
+    /// Returns the number of loops in the simd loop nest.
+    unsigned getNumLoops() { return lowerBound().size(); }
+
+  }];
+ 
+  let parser = [{ return parseSimdLoopOp(parser, result); }];
+  let printer = [{ return printSimdLoopOp(p, *this); }];
+}
+
 def YieldOp : OpenMP_Op<"yield",
     [NoSideEffect, ReturnLike, Terminator,
-     ParentOneOf<["WsLoopOp", "ReductionDeclareOp"]>]> {
+     ParentOneOf<["WsLoopOp", "ReductionDeclareOp", "SimdLoopOp"]>]> {
   let summary = "loop yield and termination operation";
   let description = [{
     "omp.yield" yields SSA values from the OpenMP dialect op region and
     terminates the region. The semantics of how the values are yielded is
     defined by the parent operation.
     If "omp.yield" has any operands, the operands must match the parent
     operation's results.
   }];

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

static void printWsLoopOp(OpAsmPrinter &p, WsLoopOp op) {
   if (!op.reduction_vars().empty())
     printReductionVarList(p, op.reductions(), op.reduction_vars());
 
   p << ' ';
   p.printRegion(op.region(), /*printEntryBlockArgs=*/false);
 }
 
 //===----------------------------------------------------------------------===//
+// SimdLoopOp
+//===----------------------------------------------------------------------===//
+/// Parses an OpenMP Simd Loop operation
+///
+/// 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
+static ParseResult parseSimdLoopOp(OpAsmParser &parser,
+                                   OperationState &result) {
+
+  // Parse an opening `(` followed by induction variables followed by `)`
+  SmallVector<OpAsmParser::OperandType> ivs;
+  if (parser.parseRegionArgumentList(ivs, /*requiredOperandCount=*/-1,
+                                     OpAsmParser::Delimiter::Paren))
+    return failure();
+  int numIVs = static_cast<int>(ivs.size());
+  Type loopVarType;
+  if (parser.parseColonType(loopVarType))
+    return failure();
+  // Parse loop bounds.
+  SmallVector<OpAsmParser::OperandType> lower;
+  if (parser.parseEqual() ||
+      parser.parseOperandList(lower, numIVs, OpAsmParser::Delimiter::Paren) ||
+      parser.resolveOperands(lower, loopVarType, result.operands))
+    return failure();
+  SmallVector<OpAsmParser::OperandType> upper;
+  if (parser.parseKeyword("to") ||
+      parser.parseOperandList(upper, numIVs, OpAsmParser::Delimiter::Paren) ||
+      parser.resolveOperands(upper, loopVarType, result.operands))
+    return failure();
+
+  // Parse step values.
+  SmallVector<OpAsmParser::OperandType> steps;
+  if (parser.parseKeyword("step") ||
+      parser.parseOperandList(steps, numIVs, OpAsmParser::Delimiter::Paren) ||
+      parser.resolveOperands(steps, loopVarType, result.operands))
+    return failure();
+
+  SmallVector<int> segments{numIVs, numIVs, numIVs};
+  // TODO: Add parseClauses() when we support clauses
+  result.addAttribute("operand_segment_sizes",
+                      parser.getBuilder().getI32VectorAttr(segments));
+
+  // Now parse the body.
+  Region *body = result.addRegion();
+  SmallVector<Type> ivTypes(numIVs, loopVarType);
+  SmallVector<OpAsmParser::OperandType> blockArgs(ivs);
+  if (parser.parseRegion(*body, blockArgs, ivTypes))
+    return failure();
+  return success();
+}
+
+static void printSimdLoopOp(OpAsmPrinter &p, SimdLoopOp op) {
+  auto args = op.getRegion().front().getArguments();
+  p << " (" << args << ") : " << args[0].getType() << " = (" << op.lowerBound()
+    << ") to (" << op.upperBound() << ") ";
+  p << "step (" << op.step() << ") ";
+
+  p.printRegion(op.region(), /*printEntryBlockArgs=*/false);
+}
+
+//===----------------------------------------------------------------------===//
 // ReductionOp
 //===----------------------------------------------------------------------===//
 
 static ParseResult parseAtomicReductionRegion(OpAsmParser &parser,
                                               Region &region) {
   if (parser.parseOptionalKeyword("atomic"))
     return success();
   return parser.parseRegion(region);

 static LogicalResult verifyWsLoopOp(WsLoopOp op) {
   return verifyReductionVarList(op, op.reductions(), op.reduction_vars());
 }
 
 //===----------------------------------------------------------------------===//
 // Verifier for critical construct (2.17.1)
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verifyCriticalDeclareOp(CriticalDeclareOp op) {
   return verifySynchronizationHint(op, op.hint());
 }
 
 static LogicalResult verifyCriticalOp(CriticalOp op) {
 
   if (op.nameAttr()) {
     auto symbolRef = op.nameAttr().cast<SymbolRefAttr>();
     auto decl =

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

convertOmpWsLoop(Operation &opInst, llvm::IRBuilderBase &builder,
   auto nextInsertionPoint =
       ompBuilder->createBarrier(contInsertPoint, llvm::omp::OMPD_for);
   tempTerminator->eraseFromParent();
   builder.restoreIP(nextInsertionPoint);
 
   return success();
 }
 
+/// Converts an OpenMP simd loop into LLVM IR using OpenMPIRBuilder.
+static LogicalResult
+convertOmpSimdLoop(Operation &opInst, llvm::IRBuilderBase &builder,
+                   LLVM::ModuleTranslation &moduleTranslation) {
+  auto loop = cast<omp::SimdLoopOp>(opInst);
+  // TODO: this should be in the op verifier instead.
+  if (loop.lowerBound().empty())
+    return failure();
+
+  llvm::DISubprogram *subprogram =
+      builder.GetInsertBlock()->getParent()->getSubprogram();
+  const llvm::DILocation *diLoc =
+      moduleTranslation.translateLoc(opInst.getLoc(), subprogram);
+  llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder.saveIP(),
+                                                    llvm::DebugLoc(diLoc));
+
+  // 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();
+  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.
+    llvm::BasicBlock *entryBlock = ip.getBlock();
+    llvm::BasicBlock *exitBlock =
+        entryBlock->splitBasicBlock(ip.getPoint(), "omp.simdloop.exit");
+    convertOmpOpRegions(loop.region(), "omp.simdloop.region", *entryBlock,
+                        *exitBlock, 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) {
+      loc = llvm::OpenMPIRBuilder::LocationDescription(bodyInsertPoints.back(),
+                                                       llvm::DebugLoc(diLoc));
+      computeIP = loopInfos.front()->getPreheaderIP();
+    }
+    loopInfos.push_back(ompBuilder->createCanonicalLoop(
+        loc, bodyGen, lowerBound, upperBound, step,
+        /*IsSigned=*/true, /*Inclusive=*/true, computeIP));
+
+    if (failed(bodyGenStatus))
+      return failure();
+  }
+
+  // Collapse loops.
+  llvm::IRBuilderBase::InsertPoint afterIP = loopInfos.front()->getAfterIP();
+  llvm::CanonicalLoopInfo *loopInfo =
+      ompBuilder->collapseLoops(diLoc, loopInfos, {});
+
+  ompBuilder->applySimd(ompLoc.DL, loopInfo);
+
+  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::TypeSwitch<Operation *, LogicalResult>(op)
         return convertOmpOrderedRegion(*op, builder, moduleTranslation);
       })
       .Case([&](omp::OrderedOp) {
         return convertOmpOrdered(*op, builder, moduleTranslation);
       })
       .Case([&](omp::WsLoopOp) {
         return convertOmpWsLoop(*op, builder, moduleTranslation);
       })
+      .Case([&](omp::SimdLoopOp) {
+        return convertOmpSimdLoop(*op, builder, moduleTranslation);
+      })
       .Case([&](omp::AtomicReadOp) {
         return convertOmpAtomicRead(*op, builder, moduleTranslation);
       })
       .Case([&](omp::AtomicWriteOp) {
         return convertOmpAtomicWrite(*op, builder, moduleTranslation);
       })
       .Case([&](omp::SectionsOp) {
         return convertOmpSections(*op, builder, moduleTranslation);

mlir/test/Dialect/OpenMP/invalid.mlir

func @proc_bind_not_allowed(%lb : index, %ub : index, %step : index) {
   // expected-error @below {{proc_bind is not a valid clause for the omp.wsloop operation}}
   omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) proc_bind(close) {
     omp.yield
   }
 }
 
 // -----
 
+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>} :
+    (index, index, i32) -> () 
+
+  return
+}
+
+// -----
+
 // expected-error @below {{op expects initializer region with one argument of the reduction type}}
 omp.reduction.declare @add_f32 : f64
 init {
 ^bb0(%arg: f32):
   %0 = arith.constant 0.0 : f32
   omp.yield (%0 : f32)
 }
 combiner {

mlir/test/Dialect/OpenMP/ops.mlir

omp.wsloop (%iv1, %iv2) : i32 = (%lb1, %lb2) to (%ub1, %ub2) step (%step1, %step2) {
     memref.store %iv1, %data1[%1] : memref<?xi32>
     memref.store %iv2, %data1[%2] : memref<?xi32>
     omp.yield
   }
 
   return
 }
 
+// CHECK-LABEL: omp_simdloop
+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>} :
+    (index, index, index) -> () 
+
+  return
+}
+
+// CHECK-LABEL: omp_simdloop_pretty
+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_target
 func @omp_target(%if_cond : i1, %device : si32,  %num_threads : si32) -> () {
 
     // Test with optional operands; if_expr, device, thread_limit, and nowait.
     // CHECK: omp.target
     "omp.target"(%if_cond, %device, %num_threads) ({
        // CHECK: omp.terminator
        omp.terminator