[RFC] Clause-based representation of OpenMP dialect operations

MLIR
1

Background

OpenMP clauses can be attached to different constructs, but their representation in the spec is generally the same regardless of the construct to which they are attached. In the OpenMP MLIR dialect, the representation of a clause corresponds to a list of arguments and a declarative assembly format string to define their textual representation.

The current problem is that all operations for which a certain clause applies must explicitly define the relevant arguments and assembly format associated to that clause, which leads to redundancy in the dialect and makes it very easy to end up with inconsistent representations for the same clause on different operations. Also, any modifications made to a given clause result in having to search for all instances of it across all operations, which is error-prone.

Proposal

It would be better to have a single source of truth about OpenMP clauses rather than have this information spread and duplicated across multiple OpenMP operations, which should make the dialect simpler to understand and contribute to. To this end, the following stack of PRs has been created:

The idea of this patch stack is to make OpenMP clauses independent top level definitions which can then be passed in a template argument list to OpenMP operation definitions, just as it’s already done for traits. OpenMP_Clauses are passed as a list template argument to OpenMP_Ops and they can define certain properties, which are joined together in order to make a default initialization for the fields of the same name of the OpenMP operation:

  • traits: Optional. They are added to the list of traits of the operation. This is a useful feature for clauses such as map, which must always come together with a MapClauseOwningOpInterface.
  • arguments: Mandatory. It defines how the clause is represented, through a dag(ins) field.
  • assemblyFormat: Optional. It should almost always be defined, but there are clauses for which the string representation is always interlinked with the op they appear in, such as collapse. This is the declarative definition of the printer / parser for the arguments. How these are combined depends on whether the clause is marked as required or optional.
  • description: Optional. It’s used to populate a clausesDescription field, so each operation definition must still define a description itself. That field is intended to be appended to the end of the OpenMP_Op’s description.
  • extraClassDeclaration: Optional. It can define some C++ code to be added to every OpenMP operation that includes that clause.

In order to give operation definitions fine-grained control over features of a certain clause might need to be inhibited, the OpenMP_Clause class takes “skipTraits”, “skipArguments”, “skipAssemblyFormat”, “skipDescription” and “skipExtraClassDeclaration” bit template arguments. These are intended to be used very sparingly for cases where some of the clauses might collide in some way otherwise.

For operations where clause-populated fields must be overriden and augmented with further operation-specific information, the clausesArgs, clausesAssemblyFormat and clausesExtraClassDeclaration fields contain what would otherwise be the default values for arguments, assemblyFormat and extraClassDeclaration, respectively.

Example

To illustrate how this change would impact the definition of operations in the OpenMP dialect, below is a simplified version of the current definition of omp.teams:

def TeamsOp : OpenMP_Op<"teams", [AttrSizedOperandSegments, ...]> {
  let summary = "teams construct";
  let description = [{
    The teams construct defines a...

    The optional `num_teams_upper` and `num_teams_lower` specify...

    The `allocators_vars` and `allocate_vars` parameters are...
  }];

  let arguments = (ins Optional<AnyInteger>:$num_teams_lower,
                       Optional<AnyInteger>:$num_teams_upper,
                       Variadic<AnyType>:$allocate_vars,
                       Variadic<AnyType>:$allocators_vars);

  let regions = (region AnyRegion:$region);

  let builders = [
    OpBuilder<(ins CArg<"const TeamsClauseOps &">:$clauses)>
  ];

  let assemblyFormat = [{
    oilist(
      `num_teams` `(` ( $num_teams_lower^ `:` type($num_teams_lower) )? `to`
                        $num_teams_upper `:` type($num_teams_upper) `)`
    | `allocate` `(`
        custom<AllocateAndAllocator>(
          $allocate_vars, type($allocate_vars),
          $allocators_vars, type($allocators_vars)
        ) `)`
    ) $region attr-dict
  }];

  let hasVerifier = 1;
}

The equivalent representation (resulting in the same outputs from tablegen) following the proposed approach would instead define two OpenMP_Clause instances, for the num_teams and allocate clauses, and an OpenMP_Op for omp.teams, as follows:

def OpenMP_NumTeamsClause : OpenMP_Clause<
    /*isRequired=*/false, /*skipTraits=*/false, /*skipArguments=*/false,
    /*skipAssemblyFormat=*/false, /*skipDescription=*/false,
    /*skipExtraClassDeclaration=*/false> {
  let arguments = (ins
    Optional<AnyInteger>:$num_teams_lower,
    Optional<AnyInteger>:$num_teams_upper
  );

  let assemblyFormat = [{
    `num_teams` `(` ( $num_teams_lower^ `:` type($num_teams_lower) )? `to`
                      $num_teams_upper `:` type($num_teams_upper) `)`
  }];

  let description = [{
    The optional `num_teams_upper` and `num_teams_lower` arguments specify...
  }];
}

def OpenMP_AllocateClause : OpenMP_Clause<
    /*isRequired=*/false, /*skipTraits=*/false, /*skipArguments=*/false,
    /*skipAssemblyFormat=*/false, /*skipDescription=*/false,
    /*skipExtraClassDeclaration=*/false> {
  let arguments = (ins
    Variadic<AnyType>:$allocate_vars,
    Variadic<AnyType>:$allocators_vars
  );

  let assemblyFormat = [{
    `allocate` `(`
      custom<AllocateAndAllocator>($allocate_vars, type($allocate_vars),
                                   $allocators_vars, type($allocators_vars)) `)`
  }];

  let description = [{
    The `allocators_vars` and `allocate_vars` parameters are...
  }];
}

def TeamsOp : OpenMP_Op<"teams", [AttrSizedOperandSegments...], [
    OpenMP_NumTeamsClause, OpenMP_AllocateClause
  ], singleRegion = true> {
  let summary = "teams construct";
  let description = [{
    The teams construct defines a...
  }] # clausesDescription;

  let builders = [
    OpBuilder<(ins CArg<"const TeamsClauseOps &">:$clauses)>
  ];

  let hasVerifier = 1;
}
2 Likes
2

Thanks @skatrak for the proposal. This looks great to me. It probably matches the standards description better than the current approach.

I see that the region is also a template argument. Could you describe the different kinds of regions possible (no region, single block, single region, multiple regions) and how it will be passed to the operation. Also, some operands have relations to the block argument of the region, now that you are planning to make the region a template argument, how does this work?

The arguments (operands, attributes) of an operation were previously immediately clear by looking at the definition of the operation. Now these are spread across multiple template arguments, we will have to look at all of them to understand all the arguments of an operation.

How is the ordering of the arguments of the operation finally decided. We had oilist before that made it order-independent. Is that still the case.

Have you looked at other dialects? Do any of them have this approach where arguments of the operation are declared through template arguments? Or are we the first to take this approach? May be @mehdi_amini or @ftynse can advise whether there are any issues from an MLIR modelling point of view.

3

I see that the region is also a template argument. Could you describe the different kinds of regions possible (no region, single block, single region, multiple regions) and how it will be passed to the operation. Also, some operands have relations to the block argument of the region, now that you are planning to make the region a template argument, how does this work?

This initial implementation only takes a single bit hasSingleRegion argument. It’s used to represent the most common 2 types of OpenMP operations: the ones without a region (argument is 0) and the ones with a single AnyRegion named $region (argument is 1). This is helpful for avoiding the need for operations to define an assemblyFormat, since the default clause-populated string will end with “$region attr-dict” or “attr-dict” depending on this flag.

It doesn’t do anything about simplifying the handling of entry block arguments or representing other types / amounts of regions. The regions and assemblyFormat fields can be overriden by operation definitions to support other kinds of regions. At the moment, I’m on the edge between thinking this feature helps with readability and thinking this may obscure the op’s definition more than it helps. I’m open to suggestions on this. The same example above, if we remove that argument, would be slightly longer:

def TeamsOp : OpenMP_Op<"teams", [AttrSizedOperandSegments...], [
    OpenMP_NumTeamsClause, OpenMP_AllocateClause
  ]> {
  let summary = "teams construct";
  let description = [{
    The teams construct defines a...
  }] # clausesDescription;

  let builders = [
    OpBuilder<(ins CArg<"const TeamsClauseOps &">:$clauses)>
  ];

  let hasVerifier = 1;
  // Changes:
  let regions = (region AnyRegion:$region);
  let assemblyFormat = clausesAssemblyFormat # " $region attr-dict";
}

The arguments (operands, attributes) of an operation were previously immediately clear by looking at the definition of the operation. Now these are spread across multiple template arguments, we will have to look at all of them to understand all the arguments of an operation.

This is true, but at the same time we wouldn’t have to guess which argument(s) applied to each clause. Although we would be adding a level of indirection, it would still be relatively easy to figure out what arguments are part of the operation by checking the corresponding clause definitions. I think the fact that there are now *ClauseOps-based builders simplifies the most potentially problematic use case, since figuring out the complete sorted list of arguments to pass to a builder could take some looking around.

How is the ordering of the arguments of the operation finally decided. We had oilist before that made it order-independent. Is that still the case.

Basically, each clause definition is marked with a bit template argument to say whether it’s optional or required. The clause-based assembly format of an operation will contain the space-separated assembly formats of all required clauses followed by an oilist containing the assembly formats of optional clauses separated by “|”. This allows us to replicate the previous behavior without having to define it explicitly for each operation.

Have you looked at other dialects? Do any of them have this approach where arguments of the operation are declared through template arguments? Or are we the first to take this approach?

I took inspiration from the LLVM_MemAccessOpBase class, which defines a set of inputs (aliasAttrs) that is concatenated with other arguments by certain definitions (e.g. LLVM_LoadOp) to produce the final argument list of certain operations.

In this case, it’s a bit more involved because there are multiple fields created this way (e.g. arguments, assemblyFormat, …), and they are populated from multiple definitions provided as template arguments (the list of clauses) rather than having a default value defined by the parent class. My guess is that this isn’t a common pattern, but the general approach seems to me well-suited to the representation of OpenMP operations.

4

In the example you provide, the teams operation does not define any arguments of its own, how would it look like if it need to? Would you then have to expand the OpenMP_Op template?

5

It would be possible to override the clause-populated list of arguments (or any other field) by just setting the field to something else:

   let arguments = (ins SomethingElse:$x...);

That’s generally not something that we’d want to do, since it loses automatically-populated arguments. That’s what the clausesArgs field is intended to be used for. It’s for the cases when we want to add arguments, but still keep the ones that would otherwise be defined by default:

  let arguments = !con(clausesArgs, (ins SomethingElse:$x...));

The same feature is also available for all other clause-populated fields via the clausesAssemblyFormat and clausesExtraClassDeclaration fields. The example above could be extended with one extra argument as follows:

def TeamsOp : OpenMP_Op<"teams", [AttrSizedOperandSegments...], [
    OpenMP_NumTeamsClause, OpenMP_AllocateClause
  ], singleRegion = true> {
  let summary = "teams construct";
  let description = [{
    The teams construct defines a...
  }] # clausesDescription;

  let builders = [
    OpBuilder<(ins CArg<"const TeamsClauseOps &">:$clauses)>
  ];

  let hasVerifier = 1;

  // Changes:
  let arguments = !con(clausesArgs, (ins I32:$some_int));
  let assemblyFormat = "some_int `(` $some_int `)` " # clausesAssemblyFormat # " $region attr-dict";
}
6

(Sorry for the delay.)

I have used a similar approach, but can’t recall if it was upstream or in own of the downstreams. No particular complaints about this from me, tablegen readability is awful but not something we control.

Two recommendations:

  • Find a way to test this behavior, e.g., by defining dummy ops in some test-omp dialect and checking the raw output of mlir-tblgen (this just lists all definitions after instantiation and substitution happened).
  • Think of a way to produce understandable error messages when user forgot to do !con or string concatenation for some of the values coming from the base class. Assembly format will complain about some fields missing, but not obvious to the user because they are deep inside layers of tablegen. Something like op description will not complain at all. This may require writing a pseudo-backend for tablegen that does the checks and errors out on failure, and does nothing on success.

Otherwise, looks really good from the modeling perspective.

1 Like
7

Looks like I’m a bit delayed with my response as well :slight_smile:

Thank you for your recommendations, they are important aspects to consider. I just updated the first PR in the stack adding some tablegen tests, as you proposed in the first point.

After looking into how to deal with the second case, I agree it would be nice, from the perspective of someone adding/modifying OpenMP ops, to get some sort of reminder about considering “inherited” properties from the clause list when overriding operation fields that would otherwise be automatically populated. Otherwise it would be easy to lose optional information and errors on required information would be quite difficult to understand. I agree that adding some sort of mlir-tblgen -check-openmp-ops option that doesn’t actually output anything could possibly be the way to achieve this.

However, it wouldn’t be a trivial thing to implement, as tablegen generators are something I have no prior experience on. Since this is not a functional change but rather an improvement to error reporting, would you be ok with implementing it as a follow-up PR after the main stack lands?

8

No objection from me.