[mlir][openmp] Custom syntax for `omp.target` operation

[alexbatashev]

Summary

Add a custom parser and printer for omp.target operation.

[bondhugula]

This is useful support. Looks good - a comment below.

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

983

Use AttrSizedOperandSegments::getOperandSegmentSizeAttr().

[bondhugula]

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

403

Update doc above: I think there isn't a need to define these other than mentioning that they have the same meaning as in the OpenMP standard?

[alexbatashev]

address reviewers feedback

[alexbatashev]

Thanks @bondhugula! I've switched to AttrSizedOperandSegments::getOperandSegmentSizeAttr() instead of referencing the attribute by name and updated docs to reflect that new arguments have the same meaning as in OpenMP standard.

[alexbatashev]

@bondhugula can you please submit this patch on my behalf? I don't have merge rights to LLVM. Please, use these credentials: Alexander Batashev <[e-mail]>

[alexbatashev]

Some changes were missing from this differential. Return them back.

[alexbatashev]

FWIW, I couldn't see changes from my first changeset in this differential, so I squashed both commits into one and uploaded them again.

[bondhugula]

It'll be good if someone closely working on the omp dialect also reviews this. I've added a few of them.

[kiranchandramohan]

Thanks for this patch. In general, we are a bit reluctant to add the privatisation clauses to more operations since the discussion for handling privatisation (https://llvm.discourse.group/t/rfc-privatisation-in-openmp-dialect/3526) in the MLIR dialect has not concluded. Flang which uses the OpenMP dialect deals with privatisation when it lowers to MLIR. We should be back to discussing privatisation soon.

We can take the other changes.

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

837

Nit: num_threads -> thread_limit.

[alexbatashev]

Removed private and firstprivate clauses

[alexbatashev]

@kiranchandramohan thanks for letting me know and pointing to that interesting discussion!

[alexbatashev]

Fix typo

[kiranchandramohan]

LGTM.

[alexbatashev]

@kiranchandramohan can you please submit this patch on my behalf? I don't have merge rights to LLVM. Please, use these credentials: Alexander Batashev <[e-mail]>

[kiranchandramohan]

I have submitted the patch for you @alexbatashev. It seems that you have a few patches submitted to llvm, I think you can ask for commit access so that you can merge patches yourself.
You can find the details for obtaining commit access from the following link.
https://llvm.org/docs/DeveloperPolicy.html#obtaining-commit-access

[alexbatashev]

Thank you, @kiranchandramohan!

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

let description = [{
 
     The optional $device parameter specifies the device number for the target region.
 
     The optional $thread_limit specifies the limit on the number of threads
 
     The optional $nowait elliminates the implicit barrier so the parent task can make progress
     even if the target task is not yet completed.
 
-    TODO:  private, map, is_device_ptr, firstprivate, depend, defaultmap, in_reduction
+    TODO:  map, is_device_ptr, depend, defaultmap, in_reduction
 
   }];
 
   let arguments = (ins Optional<I1>:$if_expr,
                        Optional<AnyInteger>:$device,
                        Optional<AnyInteger>:$thread_limit,
+                       Variadic<AnyType>:$private_vars,
+                       Variadic<AnyType>:$firstprivate_vars,
                        UnitAttr:$nowait);
 
   let regions = (region AnyRegion:$region);
+
+  let parser = [{ return parseTargetOp(parser, result); }];
+  let printer = [{ return printTargetOp(p, *this); }];
 }
 
 
 //===----------------------------------------------------------------------===//
 // 2.16 master Construct
 //===----------------------------------------------------------------------===//
 def MasterOp : OpenMP_Op<"master"> {
   let summary = "master construct";

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

if (auto def = op.default_val())
     p << "default(" << def->drop_front(3) << ") ";
 
   if (auto bind = op.proc_bind_val())
     p << "proc_bind(" << bind << ") ";
 
   p.printRegion(op.getRegion());
 }
 
+static void printTargetOp(OpAsmPrinter &p, TargetOp op) {
+  p << " ";
+  if (auto ifCond = op.if_expr())
+    p << "if(" << ifCond << " : " << ifCond.getType() << ") ";
+
+  if (auto device = op.device())
+    p << "device(" << device << " : " << device.getType() << ") ";
+
+  if (auto threads = op.thread_limit())
+    p << "thread_limit(" << threads << " : " << threads.getType() << ") ";
+
+  printDataVars(p, op.private_vars(), "private");
+  printDataVars(p, op.firstprivate_vars(), "firstprivate");
+
+  if (op.nowait()) {
+    p << "nowait ";
+  }
+
+  p.printRegion(op.getRegion());
+}
+
 //===----------------------------------------------------------------------===//
 // Parser and printer for Linear Clause
 //===----------------------------------------------------------------------===//
 
 /// linear ::= `linear` `(` linear-list `)`
 /// linear-list := linear-val | linear-val linear-list
 /// linear-val := ssa-id-and-type `=` ssa-id-and-type
 static ParseResult
if (nonspeculative && speculative)
     return op->emitOpError() << "the hints omp_sync_hint_nonspeculative and "
                                 "omp_sync_hint_speculative cannot be combined.";
   return success();
 }
 
 enum ClauseType {
   ifClause,
   numThreadsClause,
+  deviceClause,
+  threadLimitClause,
   privateClause,
   firstprivateClause,
   lastprivateClause,
   sharedClause,
   copyinClause,
   allocateClause,
   defaultClause,
   procBindClause,
static ParseResult parseClauses(OpAsmParser &parser, OperationState &result,
 
   // Stores the last parsed clause keyword
   StringRef clauseKeyword;
   StringRef opName = result.name.getStringRef();
 
   // Containers for storing operands, types and attributes for various clauses
   std::pair<OpAsmParser::OperandType, Type> ifCond;
   std::pair<OpAsmParser::OperandType, Type> numThreads;
+  std::pair<OpAsmParser::OperandType, Type> device;
+  std::pair<OpAsmParser::OperandType, Type> threadLimit;
 
   SmallVector<OpAsmParser::OperandType> privates, firstprivates, lastprivates,
       shareds, copyins;
   SmallVector<Type> privateTypes, firstprivateTypes, lastprivateTypes,
       sharedTypes, copyinTypes;
 
   SmallVector<OpAsmParser::OperandType> allocates, allocators;
   SmallVector<Type> allocateTypes, allocatorTypes;
if (clauseKeyword == "if") {
         return failure();
       clauseSegments[pos[ifClause]] = 1;
     } else if (clauseKeyword == "num_threads") {
       if (checkAllowed(numThreadsClause) || parser.parseLParen() ||
           parser.parseOperand(numThreads.first) ||
           parser.parseColonType(numThreads.second) || parser.parseRParen())
         return failure();
       clauseSegments[pos[numThreadsClause]] = 1;
+    } else if (clauseKeyword == "device") {
+      if (checkAllowed(deviceClause) || parser.parseLParen() ||
+          parser.parseOperand(device.first) ||
+          parser.parseColonType(device.second) || parser.parseRParen())
+        return failure();
+      clauseSegments[pos[deviceClause]] = 1;
+    } else if (clauseKeyword == "thread_limit") {
+      if (checkAllowed(threadLimitClause) || parser.parseLParen() ||
+          parser.parseOperand(threadLimit.first) ||
+          parser.parseColonType(threadLimit.second) || parser.parseRParen())
+        return failure();
+      clauseSegments[pos[threadLimitClause]] = 1;
     } else if (clauseKeyword == "private") {
       if (checkAllowed(privateClause) ||
           parseOperandAndTypeList(parser, privates, privateTypes))
         return failure();
       clauseSegments[pos[privateClause]] = privates.size();
     } else if (clauseKeyword == "firstprivate") {
       if (checkAllowed(firstprivateClause) ||
           parseOperandAndTypeList(parser, firstprivates, firstprivateTypes))
if (done[ifClause] && clauseSegments[pos[ifClause]] &&
     return failure();
 
   // Add num_threads parameter.
   if (done[numThreadsClause] && clauseSegments[pos[numThreadsClause]] &&
       failed(parser.resolveOperand(numThreads.first, numThreads.second,
                                    result.operands)))
     return failure();
 
+  // Add device parameter.
+  if (done[deviceClause] && clauseSegments[pos[deviceClause]] &&
+      failed(
+          parser.resolveOperand(device.first, device.second, result.operands)))
+    return failure();
+
+  // Add num_threads parameter.
+  if (done[threadLimitClause] && clauseSegments[pos[threadLimitClause]] &&
+      failed(parser.resolveOperand(threadLimit.first, threadLimit.second,
+                                   result.operands)))
+    return failure();
+
   // Add private parameters.
   if (done[privateClause] && clauseSegments[pos[privateClause]] &&
       failed(parser.resolveOperands(privates, privateTypes,
                                     privates[0].location, result.operands)))
     return failure();
 
   // Add firstprivate parameters.
   if (done[firstprivateClause] && clauseSegments[pos[firstprivateClause]] &&
static ParseResult parseParallelOp(OpAsmParser &parser,
   Region *body = result.addRegion();
   SmallVector<OpAsmParser::OperandType> regionArgs;
   SmallVector<Type> regionArgTypes;
   if (parser.parseRegion(*body, regionArgs, regionArgTypes))
     return failure();
   return success();
 }
 
+/// Parses a target operation.
+///
+/// operation ::= `omp.target` clause-list
+/// clause-list ::= clause | clause clause-list
+/// clause ::= if | device | private | firstprivate | nowait
+///
+static ParseResult parseTargetOp(OpAsmParser &parser, OperationState &result) {
+  SmallVector<ClauseType> clauses = {ifClause,           deviceClause,
+                                     threadLimitClause,  privateClause,
+                                     firstprivateClause, nowaitClause};
+
+  SmallVector<int> segments;
+
+  if (failed(parseClauses(parser, result, clauses, segments)))
+    return failure();
+
+  result.addAttribute("operand_segment_sizes",
+                      parser.getBuilder().getI32VectorAttr(segments));
+
+  Region *body = result.addRegion();
+  SmallVector<OpAsmParser::OperandType> regionArgs;
+  SmallVector<Type> regionArgTypes;
+  if (parser.parseRegion(*body, regionArgs, regionArgTypes))
+    return failure();
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // Parser, printer and verifier for SectionsOp
 //===----------------------------------------------------------------------===//
 
 /// Parses an OpenMP Sections operation
 ///
 /// sections ::= `omp.sections` clause-list
 /// clause-list ::= clause clause-list | empty

mlir/test/Dialect/OpenMP/ops.mlir

omp.wsloop (%iv1, %iv2) : i32 = (%lb1, %lb2) to (%ub1, %ub2) step (%step1, %step2) {
     memref.store %iv2, %data1[%2] : memref<?xi32>
     omp.yield
   }
 
   return
 }
 
 // CHECK-LABEL: omp_target
-func @omp_target(%if_cond : i1, %device : si32,  %num_threads : si32) -> () {
+func @omp_target(%if_cond : i1, %device : si32,  %num_threads : si32, %data_var : memref<i32>) -> () {
 
-    // Test with optional operands; if_expr, device, thread_limit, and nowait.
-    // CHECK: omp.target
-    "omp.target"(%if_cond, %device, %num_threads) ({
+    // Test with optional operands; if_expr, device, thread_limit, private, firstprivate and nowait.
+    // CHECK: omp.target if({{.*}}) device({{.*}}) thread_limit({{.*}}) private({{.*}}) firstprivate({{.*}}) nowait
+    "omp.target"(%if_cond, %device, %num_threads, %data_var, %data_var) ({
        // CHECK: omp.terminator
        omp.terminator
-    }) {operand_segment_sizes = dense<[1,1,1]>: vector<3xi32>, nowait } : ( i1, si32, si32  ) -> ()
+    }) {operand_segment_sizes = dense<[1,1,1,1,1]>: vector<5xi32>, nowait } : ( i1, si32, si32, memref<i32>, memref<i32> ) -> ()
 
     // CHECK: omp.barrier
     omp.barrier
 
     return
 }
 
+// CHECK-LABEL: omp_target_pretty
+func @omp_target_pretty(%if_cond : i1, %device : si32,  %num_threads : si32, %data_var : memref<i32>) -> () {
+    // CHECK: omp.target if({{.*}}) device({{.*}}) firstprivate({{.*}})
+    omp.target if(%if_cond : i1) device(%device : si32) thread_limit(%num_threads : si32) firstprivate(%data_var : memref<i32>) {
+      omp.terminator
+    }
+
+    // CHECK: omp.target if({{.*}}) device({{.*}}) private({{.*}}) nowait
+    omp.target if(%if_cond : i1) device(%device : si32) thread_limit(%num_threads : si32) private(%data_var : memref<i32>) nowait {
+      omp.terminator
+    }
+
+    return
+}
+
 // CHECK: omp.reduction.declare
 // CHECK-LABEL: @add_f32
 // CHECK: : f32
 // CHECK: init
 // CHECK: ^{{.+}}(%{{.+}}: f32):
 // CHECK:   omp.yield
 // CHECK: combiner
 // CHECK: ^{{.+}}(%{{.+}}: f32, %{{.+}}: f32):