[Mlir-commits] [mlir] [OpenMP][MLIR] Add omp.distribute op to the OMP dialect (PR #67720)

[Sergio Afonso]

skatrak wrote:

I like the idea of potentially collecting all information on how a loop may be split up across different levels of parallelism in the `omp.wsloop` MLIR operation, just because it makes it simpler to call the right runtime function. In my mind, that approach taken to what I think to be its full realization would probably at least impact the `distribute`, `do` and `simd` directives, though maybe also `loop` or others. Assuming the "worksharing type" attribute can be printed/parsed how I envision it, these various OpenMP constructs could be represented as follows:
```mlir
// !$omp target teams distribute parallel do
omp.target {
  omp.teams {
    omp.parallel {
      omp.wsloop distribute for (%i) : i32 (%lb) to (%ub) step (%step) {
      }
    }
  }
}

// !$omp distribute
omp.wsloop distribute (%i) : i32 (%lb) to (%ub) step (%step) {
}

// !$omp distribute parallel do
omp.parallel {
  omp.wsloop distribute for (%i) : i32 (%lb) to (%ub) step (%step) {
  }
}

// !$omp distribute parallel do simd
omp.parallel {
  omp.wsloop distribute for simd (%i) : i32 (%lb) to (%ub) step (%step) {
  }
}

// !$omp distribute simd
omp.wsloop distribute simd (%i) : i32 (%lb) to (%ub) step (%step) {
}

// !$omp do
omp.wsloop for (%i) : i32 (%lb) to (%ub) step (%step) {
}

// !$omp do simd
omp.wsloop for simd (%i) : i32 (%lb) to (%ub) step (%step) {
}

// !$omp simd
omp.wsloop simd (%i) : i32 (%lb) to (%ub) step (%step) {
}

// !$omp parallel do
omp.parallel {
  omp.wsloop for (%i) : i32 (%lb) to (%ub) step (%step) {
  }
}

// !$omp parallel do simd
omp.parallel {
  omp.wsloop for simd (%i) : i32 (%lb) to (%ub) step (%step) {
  }
}
```
This ignores for the time being the introduction of `omp.canonical_loop`. I assume that whichever approach is taken there would need to somehow interact with or take over `omp.wsloop`, so in principle that fact wouldn't change by this addition. The main shortcoming of this approach, as I see it, is that certain clauses may be applicable to multiple of these directives, but we would be trying to store them all in a single MLIR operation, so we may possibly need to disambiguate by adding multiple variants to certain clauses. For example:

```mlir
// !$omp distribute private(a)
// !$omp parallel do private(b)
omp.parallel {
  omp.wsloop private(%1, %2) distribute for ... {}
  // or
  omp.wsloop distribute_private(%1) for_private(%2) distribute for ... {}
}
// Would something like this be valid?
// !$omp distribute collapse(1)
// !$omp parallel do collapse(2)
omp.parallel {
  omp.wsloop collapse(%whichever_takes_precedence) distribute for ... {}
  // or
  omp.wsloop distribute_collapse(%1) for_collapse(%2) distribute for ... {}
}
```
This is the list of clauses, according to the 5.2 spec document, accepted by the `distribute`, `do` and `simd` constructs. There are several clauses that apply to multiple or all of these directives, so it may be the case for some that semantically it's not the same attaching them to one of these constructs as doing it to all of them, or it may be allowed for them to have different values for each directive.

| Clause        | distribute | do      | simd    |
|:--------------|------------|---------|---------|
| allocate      | ✓    | ✓ |         |
| aligned       |            |         | ✓ |
| collapse      | ✓    | ✓ | ✓ |
| dist_schedule | ✓    |         |         |
| firstprivate  | ✓    | ✓ |         |
| if            |            |         | ✓ |
| lastprivate   | ✓    | ✓ | ✓ |
| linear        |            | ✓ | ✓ |
| nontemporal   |            |         | ✓ |
| nowait        |            | ✓ |         |
| order         | ✓    | ✓ | ✓ |
| ordered       |            | ✓ |         |
| private       | ✓    | ✓ | ✓ |
| reduction     |            | ✓ | ✓ |
| schedule      |            | ✓ |         |
| safelen       |            |         | ✓ |
| simdlen       |            |         | ✓ |

https://github.com/llvm/llvm-project/pull/67720