[MLIR][OpenMP] Teams Construct Operation

[abidmalikwaterloo]

Summary

Added the OpenMP Teams Construct operation definition for MLIR.

[clementval]

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

100

From another patch?

142

You might want to align OptionalAttr with the Optional above. I think other operations in this file are like that.

299–300

What's the reason to remove this?

325

From another patch?

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

799 ↗

(On Diff #357044)

Should not be removed?

mlir/test/Dialect/OpenMP/ops.mlir

299

Should not be removed?

[abidmalikwaterloo]

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

100

I think I need to update my repo.

[abidmalikwaterloo]

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

100

This is from the previous patch.

299–300

Its accidentally removed because of moving code between branches. I. will fix all these mistakes.

325

yes. It is from another patch.

mlir/test/Dialect/OpenMP/ops.mlir

299

I branched for teams construct. Do we need it for this patch?

[abidmalikwaterloo]

Let me submit it again. Most of the stuff is still on my machine but did not appear or seemed removed.

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

799 ↗

(On Diff #357044)

This is from another patch that is still under review.

[abidmalikwaterloo]

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

142

Should we keep it in line with the parallel operation?

mlir/test/Dialect/OpenMP/ops.mlir

299

This is from another patch.

[clementval]

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

142

142

Just like the suggestion above.

[clementval]

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

120

Just a simple question here: I'm not sure we need a 1-to-1 mapping between construct and dialect operation. Did you think through how omp.target and omp.teams operation will interact?

[abidmalikwaterloo]

Here is a simple example of interaction:

func @omp_compute(%x: memref<1024xf32>, %y: memref<1024xf32>,
  %n: index, %a: f32, %N:i32) -> memref<1024xf32> {
  %c0 = constant 0 : index
  %c1 = constant 1 : index
  %c0 = constant 0 : index
  %c1 = constant 1 : index

  omp.target  {
    omp.teams num_teams(2) thread_limit(%N) {
       omp.distribute{
             loop.for %arg0 = %c0 to %n step %c1 {
             %xi = load %x[%arg0] : memref<1024xf32>
             %yi = load %y[%arg0] : memref<1024xf32>
             %ax = mulf %a, %xi : f32
             %yy = addf %ax, %yi : f32
             store %yy, %y[%arg0] : memref<1024xf32>
		}
            }
      }
  }
 return %y : memref<1024xf32>
}

[abidmalikwaterloo]

Resumbitting the patch as the previous one was missing some codes

[clementval]

I think your patch is still mixing with previous one.

[abidmalikwaterloo]

In D105581#2909226, @clementval wrote:

I think your patch is still mixing with the previous one.

@kiranchandramohan
What is the solution to remove this confusionn? This patch does not need the basic parsing patch for target operation. I feel my branching was not correctly done which created this confusion.

[clementval]

In D105581#2977048, @abidmalikwaterloo wrote:

In D105581#2909226, @clementval wrote:

I think your patch is still mixing with the previous one.

@kiranchandramohan
What is the solution to remove this confusionn? This patch does not need the basic parsing patch for target operation. I feel my branching was not correctly done which created this confusion.

You likely did your patch from a branch other than main so you have remaining of the previous patches work.
Apply your patch on a fresh branch from main and fix the problematic diffs

[abidmalikwaterloo]

In D105581#2978930, @clementval wrote:

In D105581#2977048, @abidmalikwaterloo wrote:

In D105581#2909226, @clementval wrote:

I think your patch is still mixing with the previous one.

@kiranchandramohan
What is the solution to remove this confusionn? This patch does not need the basic parsing patch for target operation. I feel my branching was not correctly done which created this confusion.

You likely did your patch from a branch other than main so you have remaining of the previous patches work.
Apply your patch on a fresh branch from main and fix the problematic diffs

OK

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

120

Are you asking with respect to the discussion [[ https://lists.llvm.org/pipermail/llvm-dev/2020-February/139136.html | [llvm-dev] About OpenMP dialect in MLIR]]? Here how I see the interaction:

omp.target {
      omp.teams {
          loop.for{
          }
     }
}

[abidmalikwaterloo]

rebase the branch to main to remove the redundent code

[clementval]

Please just update the NIT comment before proceeding.

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

123

Please split the line at 80.

[abidmalikwaterloo]

Made changes to the length of comments as per advised by the reviewer.

[peixin]

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

141

Why num_threads? num_teams and thread_limit clauses are in teams construct.

146

copyin clause is not allowed in teams construct.

reduction and defualt clauses are missed.

[skatrak]

I think this patch can be closed now, since D154441 landed with support for that operation.

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

def ParallelOp : OpenMP_Op<"parallel", [AttrSizedOperandSegments]> {
 
   let regions = (region AnyRegion:$region);
 
   let builders = [
     OpBuilder<(ins CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)>
   ];
   let parser = [{ return parseParallelOp(parser, result); }];
   let printer = [{ return printParallelOp(p, *this); }];
   let verifier = [{ return ::verifyParallelOp(*this); }];
   
 }
 
 def TerminatorOp : OpenMP_Op<"terminator", [Terminator]> {
   let summary = "terminator for OpenMP regions";
   let description = [{
     A terminator operation for regions that appear in the body of OpenMP
     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]> {
+  let summary = "teams construct";
+  let description = [{
+    The teams construct creates a league of initial teams and the initial thread 
+    in each team executes the region.
+
+    The optional $num_threads_var parameter specifies the number of threads which
+    should be used to execute the parallel region.
+
+    The optional $default_val attribute specifies the default data sharing attribute
+    of values used in the  region that are not passed explicitly as parameters
+
+    The $private_vars, $firstprivate_vars, $shared_vars and $copyin_vars parameters
+    are a variadic list of values that specify the data sharing attribute of
+    those values.
+
+    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<AnyType>:$num_threads_var,
+             OptionalAttr<ClauseDefault>:$default_val,
+             Variadic<AnyType>:$private_vars,
+             Variadic<AnyType>:$firstprivate_vars,
+             Variadic<AnyType>:$shared_vars,
+             Variadic<AnyType>:$copyin_vars,
+             Variadic<AnyType>:$allocate_vars,
+             Variadic<AnyType>:$allocators_vars);
+
+  let regions = (region AnyRegion:$region);
+}
+
 //===----------------------------------------------------------------------===//
 // 2.9.2 Workshare Loop Construct
 //===----------------------------------------------------------------------===//
 
 def WsLoopOp : OpenMP_Op<"wsloop", [AttrSizedOperandSegments,
                          AllTypesMatch<["lowerBound", "upperBound", "step"]>]> {
   let summary = "workshare loop construct";
   let description = [{
let description = [{
     enforces an order on the memory operations of the variables explicitly
     specified or implied.
   }];
 
   let arguments = (ins Variadic<AnyType>:$varList);
 
   let assemblyFormat = [{ ( `(` $varList^ `:` type($varList) `)` )? attr-dict}];
 }
+
 //===----------------------------------------------------------------------===//
 // 2.14.5 target construct
 //===----------------------------------------------------------------------===//
 
 def TargetOp : OpenMP_Op<"target",[AttrSizedOperandSegments]> {
   let summary = "target construct";
   let description = [{
     The target construct includes a region of code which is to be executed
let description = [{
     
     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
     
   }];
 
   let arguments = (ins Optional<I1>:$if_expr,
                        Optional<AnyInteger>:$device,
                        Optional<AnyInteger>:$thread_limit,
                        UnitAttr:$nowait);
 
   let regions = (region AnyRegion:$region);
   
   let builders = [
        OpBuilder<(ins "IntegerAttr":$if_expr,
-                       "IntegerAttr":$device,
-                       "IntegerAttr":$thread_limit,
-                       "UnitAttr":$nowait)>
   ];
 }
 
 
 //===----------------------------------------------------------------------===//
 // 2.16 master Construct
 //===----------------------------------------------------------------------===//
 def MasterOp : OpenMP_Op<"master"> {
   let summary = "master construct";
   let description = [{

mlir/test/Dialect/OpenMP/ops.mlir

"omp.target"(%if_cond, %device, %num_threads) ({
        // CHECK: omp.terminator
        omp.terminator
     }) {operand_segment_sizes = dense<[1,1,1]>: vector<3xi32>, nowait } : ( i1, si32, si32  ) -> ()
 
     // CHECK: omp.barrier
     omp.barrier
 
     return
 }
+
+
+//  CHECK-LABEL: omp_teams
+func @omp_teams(%data_var : memref<i32>, %num_threads : si32) -> () {
+
+    // CHECK: omp.teams
+    "omp.teams"(%num_threads, %data_var, %data_var, %data_var, %data_var) ({
+      // CHECK: omp.terminator
+      omp.terminator
+    }) {operand_segment_sizes = dense<[1,1,1,1,1,0,0]>: vector<7xi32>, default_val = "defshared"} : (si32, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
+
+  // CHECK: omp.barrier
+    omp.barrier
+
+  return
+}
+