[MLIR][OpenMP] Add MLIR operation for OpenMP teams

[skatrak]

Summary

This patch adds an operation definition for the OpenMP 'teams' construct.

[tschuett]

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

232

To be picky. No! For modern OpenMP, I use the teams construct in shared memory.

[skatrak]

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

232

Thank you for pointing this out, so I understand then that this is not a device-only construct? I'll work on removing this restriction and updating the operation description.

[tschuett]

Changes from 4.5 to 5.0:

The teams construct (see Section 10.2) was extended to support execution on the host device
without an enclosing target construct (see Section 13.8).

[kiranchandramohan]

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

227

I think this will force the teams region to have a single block only. I guess that is not the intention here.

232

The standard says the following. What is your usage @tschuett ?
"A teams region can only be strictly nested within the implicit parallel region or a target region. If a teams construct is nested within a target construct, that target construct must contain no statements, declarations or directives outside of the teams construct."

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

882–884

You can use the AllTypesMatch trait for this. For reference omp.wsloop op.

[tschuett]

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

232

I read the same sentence from the standard.

[skatrak]

Address reviewer's comments. Lift restriction of being nested inside of omp.target.

[skatrak]

Thank you for the comments, I just pushed some changes and posted a couple questions inline.

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

227

You're right, thank you. I misinterpreted here "single block" to mean "single region".

232

I focused on the part about being strictly nested within a target region, but didn't quite get the "implicit parallel region that surrounds the whole OpenMP program" when I read that paragraph prior to implementing this.

I made some changes to check that the enclosing target region doesn't have any other operations, but I'm not sure what exactly to check to make sure it's nested within the implicit global parallel region. Should we just check that omp.teams does not have an omp.parallel parent or is having any OpenMP dialect parent operation (except from omp.target) already unsupported use of omp.teams?

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

882–884

Unfortunately I can't seem to be able to do this, since both arguments are optional. If any of them is not specified and I set that trait, the TeamsOp::verifyInvariantsImpl() segfaults due to trying to access the type of these undefined arguments.

Is there a known workaround for this?

[kiranchandramohan]

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

232

I think, we can have a simple check here. If the immediate OpenMP dialect parent is any operation other than omp.target then it is an error. Otherwise do not signal any error.

I made some changes to check that the enclosing target region doesn't have any other operations

Would that exclude the following?

omp.target {
  %c1 = arith.constant 1 : i32
  %c10 = arith.constant 10: i32
  omp.teams num_teams(%c1 : i32 to %c10 : i32) {
    omp.terminator
  }
}

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

882–884

OK then leave it as is. It probably requires the operands to be always present together.

[skatrak]

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

232

I think, we can have a simple check here. If the immediate OpenMP dialect parent is any operation other than omp.target then it is an error. Otherwise do not signal any error.

That is what I implemented the first time around, but as @tschuett pointed out, the teams construct can appear outside of target regions, by binding to the implicit parallel region surrounding the program. I think we could check that its direct parent is an omp.target or it doesn't have any direct or indirect parent from the OpenMP dialect. I think that second check follows the definition of being strictly nested within that implicit parallel region. The other alternative is to just check the first, as you say, and put a TODO mentioning the other case where omp.teams should be eventually allowed.

Would that exclude the following?

omp.target {
  %c1 = arith.constant 1 : i32
  %c10 = arith.constant 10: i32
  omp.teams num_teams(%c1 : i32 to %c10 : i32) {
    omp.terminator
  }
}

Yes, it would. Now I realize the issue with that. The spec (version 5.2, section 10.2) does state "If a teams region is nested inside a target region, the corresponding target construct must not contain any statements, declarations or directives outside of the corresponding teams construct." But obviously the arguments for the construct's clauses must come from somewhere, and I understood that the region inside of an omp.target couldn't access what was outside of it.

I suppose that the way this works is it should have access to device-mapped variables, but I'm not sure how these can be accessed in MLIR. If it is by their SSA value outside of the target region, then I suppose that means values from outside omp.target can actually be used inside (they would only fail at runtime if they weren't mapped properly), so the constants in your example could be hoisted out and the restriction kept, assuming constants are always mapped to all devices.

I realize there are many considerations here, so maybe @jsjodin, @agozillon, @domada, @TIFitis have other concerns they'd like to bring up related to the visibility of SSA values across target boundaries and data mapping. For the time being, I don't mind lifting the restriction and allowing other operations to coexist with omp.teams inside of omp.target.

[kiranchandramohan]

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

232

For the time being, I don't mind lifting the restriction and allowing other operations to coexist with omp.teams inside of omp.target.

This is alrite for now.

The first check can also be simple. Only if the immediate OpenMP dialect parent is not an omp.target operation then error out. If on the device side, there is not omp.target operation as per the new design then this check might have to be skipped.

[skatrak]

Update verifier to follow reviewer's recommendations.

[skatrak]

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

232

The first check can also be simple. Only if the immediate OpenMP dialect parent is not an omp.target operation then error out. If on the device side, there is not omp.target operation as per the new design then this check might have to be skipped.

Done. After speaking with @jsjodin, the current plan to represent outlined target regions will keep omp.target operations, so it won't need special handling here.

[kiranchandramohan]

LG.

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

879–880

The alternative is to only issue this error if the parent operation is an OpenMP dialect operation and not a TargetOp.

883–884

Nit: Can this type be spelled out?

[skatrak]

Address comments.

[skatrak]

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

879–880

Just made some changes to make that check in place of the related TODO comment. I'll wait for your thumbs up before landing this

[kiranchandramohan]

LG.

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

let description = [{
     operation.  These regions are not expected to return any value so the
     terminator takes no operands. The terminator op returns control to the
     enclosing op.
   }];
 
   let assemblyFormat = "attr-dict";
 }
 
+//===----------------------------------------------------------------------===//
+// 2.7 teams Construct
+//===----------------------------------------------------------------------===//
+def TeamsOp : OpenMP_Op<"teams", [
+              AttrSizedOperandSegments, RecursiveMemoryEffects,
+              ReductionClauseInterface]> {
+  let summary = "teams construct";
+  let description = [{
+    The teams construct defines a region of code that triggers the creation of a
+    league of teams. Once created, the number of teams remains constant for the
+    duration of its code region.
+
+    The optional $num_teams_upper and $num_teams_lower specify the limit on the
+    number of teams to be created. If only the upper bound is specified, it acts
+    as if the lower bound was set to the same value. It is not supported to set
+    $num_teams_lower if $num_teams_upper is not specified. They define a closed
+    range, where both the lower and upper bounds are included.
+
+    If the $if_expr is present and it evaluates to `false`, the number of teams
+    created is one.
+
+    The optional $thread_limit specifies the limit on the number of threads.
+
+    The $allocators_vars and $allocate_vars parameters are a variadic list of
+    values that specify the memory allocator to be used to obtain storage for
+    private values.
+  }];
+
+  let arguments = (ins Optional<AnyInteger>:$num_teams_lower,
+                       Optional<AnyInteger>:$num_teams_upper,
+                       Optional<I1>:$if_expr,
+                       Optional<AnyInteger>:$thread_limit,
+                       Variadic<AnyType>:$allocate_vars,
+                       Variadic<AnyType>:$allocators_vars,
+                       Variadic<OpenMP_PointerLikeType>:$reduction_vars,
+                       OptionalAttr<SymbolRefArrayAttr>:$reductions);
+
+  let regions = (region AnyRegion:$region);
+
+  let assemblyFormat = [{
+    oilist(
+      `num_teams` `(` ( $num_teams_lower^ `:` type($num_teams_lower) )? `to`
+                        $num_teams_upper `:` type($num_teams_upper) `)`
+    | `if` `(` $if_expr `)`
+    | `thread_limit` `(` $thread_limit `:` type($thread_limit) `)`
+    | `reduction` `(`
+        custom<ReductionVarList>(
+          $reduction_vars, type($reduction_vars), $reductions
+        ) `)`
+    | `allocate` `(`
+        custom<AllocateAndAllocator>(
+          $allocate_vars, type($allocate_vars),
+          $allocators_vars, type($allocators_vars)
+        ) `)`
+    ) $region attr-dict
+  }];
+
+  let hasVerifier = 1;
+}
+
 def OMP_ScheduleModNone         : I32EnumAttrCase<"none", 0>;
 def OMP_ScheduleModMonotonic    : I32EnumAttrCase<"monotonic", 1>;
 def OMP_ScheduleModNonmonotonic : I32EnumAttrCase<"nonmonotonic", 2>;
 // FIXME: remove this value for the modifier because this is handled using a
 // separate attribute
 def OMP_ScheduleModSIMD         : I32EnumAttrCase<"simd", 3>;
 
 def ScheduleModifier

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

 LogicalResult ParallelOp::verify() {
   if (getAllocateVars().size() != getAllocatorsVars().size())
     return emitError(
         "expected equal sizes for allocate and allocator variables");
   return verifyReductionVarList(*this, getReductions(), getReductionVars());
 }
 
 //===----------------------------------------------------------------------===//
+// TeamsOp
+//===----------------------------------------------------------------------===//
+
+static bool opInGlobalImplicitParallelRegion(Operation *op) {
+  while ((op = op->getParentOp()))
+    if (isa<OpenMPDialect>(op->getDialect()))
+      return false;
+  return true;
+}
+
+LogicalResult TeamsOp::verify() {
+  // Check parent region
+  // TODO If nested inside of a target region, also check that it does not
+  // contain any statements, declarations or directives other than this
+  // omp.teams construct. The issue is how to support the initialization of
+  // this operation's own arguments (allow SSA values across omp.target?).
+  Operation *op = getOperation();
+  if (!isa<TargetOp>(op->getParentOp()) &&
+      !opInGlobalImplicitParallelRegion(op))
+    return emitError("expected to be nested inside of omp.target or not nested "
+                     "in any OpenMP dialect operations");
+
+  // Check for num_teams clause restrictions
+  if (auto numTeamsLowerBound = getNumTeamsLower()) {
+    auto numTeamsUpperBound = getNumTeamsUpper();
+    if (!numTeamsUpperBound)
+      return emitError("expected num_teams upper bound to be defined if the "
+                       "lower bound is defined");
+    if (numTeamsLowerBound.getType() != numTeamsUpperBound.getType())
+      return emitError(
+          "expected num_teams upper bound and lower bound to be the same type");
+  }
+
+  // Check for allocate clause restrictions
+  if (getAllocateVars().size() != getAllocatorsVars().size())
+    return emitError(
+        "expected equal sizes for allocate and allocator variables");
+
+  return verifyReductionVarList(*this, getReductions(), getReductionVars());
+}
+
+//===----------------------------------------------------------------------===//
 // Verifier for SectionsOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult SectionsOp::verify() {
   if (getAllocateVars().size() != getAllocatorsVars().size())
     return emitError(
         "expected equal sizes for allocate and allocator variables");
 

mlir/test/Dialect/OpenMP/invalid.mlir

omp.atomic.capture {
     }
     omp.atomic.read %v = %x memory_order(seq_cst) : memref<i32>, i32
   }
   return
 }
 
 // -----
 
+func.func @omp_teams_parent() {
+  omp.parallel {
+    // expected-error @below {{expected to be nested inside of omp.target or not nested in any OpenMP dialect operations}}
+    omp.teams {
+      omp.terminator
+    }
+    omp.terminator
+  }
+  return
+}
+
+// -----
+
+func.func @omp_teams_allocate(%data_var : memref<i32>) {
+  omp.target {
+    // expected-error @below {{expected equal sizes for allocate and allocator variables}}
+    "omp.teams" (%data_var) ({
+      omp.terminator
+    }) {operand_segment_sizes = array<i32: 0,0,0,0,1,0,0>} : (memref<i32>) -> ()
+    omp.terminator
+  }
+  return
+}
+
+// -----
+
+func.func @omp_teams_num_teams1(%lb : i32) {
+  omp.target {
+    // expected-error @below {{expected num_teams upper bound to be defined if the lower bound is defined}}
+    "omp.teams" (%lb) ({
+      omp.terminator
+    }) {operand_segment_sizes = array<i32: 1,0,0,0,0,0,0>} : (i32) -> ()
+    omp.terminator
+  }
+  return
+}
+
+// -----
+
+func.func @omp_teams_num_teams2(%lb : i32, %ub : i16) {
+  omp.target {
+    // expected-error @below {{expected num_teams upper bound and lower bound to be the same type}}
+    omp.teams num_teams(%lb : i32 to %ub : i16) {
+      omp.terminator
+    }
+    omp.terminator
+  }
+  return
+}
+
+// -----
+
 func.func @omp_sections(%data_var : memref<i32>) -> () {
   // expected-error @below {{expected equal sizes for allocate and allocator variables}}
   "omp.sections" (%data_var) ({
     omp.terminator
   }) {operand_segment_sizes = array<i32: 0,1,0>} : (memref<i32>) -> ()
   return
 }
 

mlir/test/Dialect/OpenMP/ops.mlir

omp.wsloop for (%iv) : index = (%lb) to (%ub) step (%step) {
       omp.yield
     }
     // CHECK: omp.terminator
     omp.terminator
   }
   return
 }
 
+// CHECK-LABEL: omp_teams
+func.func @omp_teams(%lb : i32, %ub : i32, %if_cond : i1, %num_threads : i32,
+                     %data_var : memref<i32>) -> () {
+  // Test nesting inside of omp.target
+  omp.target {
+    // CHECK: omp.teams
+    omp.teams {
+      // CHECK: omp.terminator
+      omp.terminator
+    }
+    // CHECK: omp.terminator
+    omp.terminator
+  }
+
+  // CHECK: omp.teams
+  omp.teams {
+    %0 = arith.constant 1 : i32
+    // CHECK: omp.terminator
+    omp.terminator
+  }
+
+  // Test num teams.
+  // CHECK: omp.teams num_teams(%{{.+}} : i32 to %{{.+}} : i32)
+  omp.teams num_teams(%lb : i32 to %ub : i32) {
+    // CHECK: omp.terminator
+    omp.terminator
+  }
+
+  // CHECK: omp.teams num_teams( to %{{.+}} : i32)
+  omp.teams num_teams(to %ub : i32) {
+    // CHECK: omp.terminator
+    omp.terminator
+  }
+
+  // Test if.
+  // CHECK: omp.teams if(%{{.+}})
+  omp.teams if(%if_cond) {
+    // CHECK: omp.terminator
+    omp.terminator
+  }
+
+  // Test thread limit.
+  // CHECK: omp.teams thread_limit(%{{.+}} : i32)
+  omp.teams thread_limit(%num_threads : i32) {
+    // CHECK: omp.terminator
+    omp.terminator
+  }
+
+  // Test reduction.
+  %c1 = arith.constant 1 : i32
+  %0 = llvm.alloca %c1 x i32 : (i32) -> !llvm.ptr<f32>
+  // CHECK: omp.teams reduction(@add_f32 -> %{{.+}} : !llvm.ptr<f32>) {
+  omp.teams reduction(@add_f32 -> %0 : !llvm.ptr<f32>) {
+    %1 = arith.constant 2.0 : f32
+    // CHECK: omp.reduction %{{.+}}, %{{.+}}
+    omp.reduction %1, %0 : f32, !llvm.ptr<f32>
+    // CHECK: omp.terminator
+    omp.terminator
+  }
+
+  // Test allocate.
+  // CHECK: omp.teams allocate(%{{.+}} : memref<i32> -> %{{.+}} : memref<i32>)
+  omp.teams allocate(%data_var : memref<i32> -> %data_var : memref<i32>) {
+    // CHECK: omp.terminator
+    omp.terminator
+  }
+
+  return
+}
+
 // CHECK-LABEL: func @sections_reduction
 func.func @sections_reduction() {
   %c1 = arith.constant 1 : i32
   %0 = llvm.alloca %c1 x i32 : (i32) -> !llvm.ptr<f32>
   // CHECK: omp.sections reduction(@add_f32 -> {{.+}} : !llvm.ptr<f32>)
   omp.sections reduction(@add_f32 -> %0 : !llvm.ptr<f32>) {
     // CHECK: omp.section
     omp.section {