[Mlir-commits] [mlir] [MLIR][OpenMP] Add omp.simd operation (PR #79843)

[Sergio Afonso]

skatrak wrote:

I've gone through all the legal OpenMP combined and composite constructs (I hope I didn't miss any) to see which additional operations would be needed in order to represent composite constructs independently, and to also compare how they would look in MLIR compared to "wrapper" ops similar what's been proposed for `omp.distribute`.

The table below lists all these combined/composite constructs as they would be represented using composite ops or wrapper ops (taking some liberties with ops that aren't in the dialect yet), but as a summary what I have found out is that as of OpenMP 5.2, we'd only need to add 5 of these operations: `omp.wssimdloop`, `omp.distsimdloop`, `omp.distparwsloop`, `omp.distparwssimdloop`, `omp.tasksimdloop` (and converting `omp.distribute` into `omp.distloop`, which would represent the loop itself, to follow the same pattern).

The main problem I see with these is redundancy. The set of accepted arguments for these operations would be the union of all accepted arguments for their leaf constructs, and also they all turned out to represent loops (with indices, ranges, etc). I suppose the second concern would be addressed by keeping all loop-related information into the upcoming `omp.canonical_loop` and have these ops just have the corresponding CLI as an argument or something similar. Maybe there are ways to avoid repeating the same lists of arguments in TableGen if my first concern turns out to be a real maintainability issue.

So I think that it might be a good idea to indeed create these composite MLIR operations, since they represent distinct behaviors and their shortcomings seem solvable in principle. With regards to the options of de-composing operations in an MLIR pass or leaving that to the MLIR to LLVM IR translation stage, I guess we could figure that out over time. Probably the second approach would be easiest to start.

It'd be interesting to hear if @mjklemm, @jsjodin, @DominikAdamski, @kparzysz or @agozillon  have any comments on this.

P.S. Sorry for the massive table, couldn't make it smaller 😅.

<table>
<tr>
<td>

**Construct**

</td>
<td>

**Composite ops**

</td>
<td>

**Wrapper ops**

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp do simd`

</td>
<td>

```mlir
omp.wssimdloop %i... {
  BODY
}
```

</td>
<td>

```mlir
omp.wsloop %i... {
  omp.simd {
    BODY
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp distribute simd`

</td>
<td>

```mlir
omp.distsimdloop %i... {
  BODY
}
```

</td>
<td>

```mlir
omp.distribute {
  omp.simdloop %i... {
    BODY
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp distribute parallel do`

</td>
<td>

```mlir
omp.distparwsloop %i... {
  BODY
}
```

</td>
<td>

```mlir
omp.distribute {
  omp.parallel {
    omp.wsloop %i... {
      BODY
    }
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp distribute parallel do simd`

</td>
<td>

```mlir
omp.distparwssimdloop %i... {
  BODY
}
```

</td>
<td>

```mlir
omp.distribute {
  omp.parallel {
    omp.wsloop %i... {
      omp.simd {
        BODY
      }
    }
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp taskloop simd`

</td>
<td>

```mlir
omp.tasksimdloop %i... {
  BODY
}
```

</td>
<td>

```mlir
omp.taskloop %i... {
  omp.simd {
    BODY
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp parallel do`

</td>
<td>

```mlir
omp.parallel {
  omp.wsloop %i... {
    BODY
  }
}
```

</td>
<td>

```mlir
omp.parallel {
  omp.wsloop %i... {
    BODY
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp parallel loop`

</td>
<td>

```mlir
omp.parallel {
  omp.loop %i... {
    BODY
  }
}
```

</td>
<td>

```mlir
omp.parallel {
  omp.loop %i... {
    BODY
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp parallel sections`

</td>
<td>

```mlir
omp.parallel {
  omp.sections {
    BODY
  }
}
```

</td>
<td>

```mlir
omp.parallel {
  omp.sections {
    BODY
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp parallel workshare`

</td>
<td>

```mlir
omp.parallel {
  omp.workshare {
    BODY
  }
}
```

</td>
<td>

```mlir
omp.parallel {
  omp.workshare {
    BODY
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp parallel do simd`

</td>
<td>

```mlir
omp.parallel {
  omp.wssimdloop %i... {
    BODY
  }
}
```

</td>
<td>

```mlir
omp.parallel {
  omp.wsloop %i... {
    omp.simd {
      BODY
    }
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp parallel masked`

</td>
<td>

```mlir
omp.parallel {
  omp.masked {
    BODY
  }
}
```

</td>
<td>

```mlir
omp.parallel {
  omp.masked {
    BODY
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp masked taskloop`

</td>
<td>

```mlir
omp.masked {
  omp.taskloop %i... {
    BODY
  }
}
```

</td>
<td>

```mlir
omp.masked {
  omp.taskloop %i... {
    BODY
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp masked taskloop simd`

</td>
<td>

```mlir
omp.masked {
  omp.tasksimdloop %i... {
    BODY
  }
}
```

</td>
<td>

```mlir
omp.masked {
  omp.taskloop %i... {
    omp.simd {
      BODY
    }
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp parallel masked taskloop`

</td>
<td>

```mlir
omp.parallel {
  omp.masked {
    omp.taskloop %i... {
      BODY
    }
  }
}
```

</td>
<td>

```mlir
omp.parallel {
  omp.masked {
    omp.taskloop %i... {
      BODY
    }
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp parallel masked taskloop simd`

</td>
<td>

```mlir
omp.parallel {
  omp.masked {
    omp.tasksimdloop %i... {
      BODY
    }
  }
}
```

</td>
<td>

```mlir
omp.parallel {
  omp.masked {
    omp.taskloop %i... {
      omp.simd {
        BODY
      }
    }
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp teams distribute`

</td>
<td>

```mlir
omp.teams {
  omp.distloop %i... {
    BODY
  }
}
```

</td>
<td>

```mlir
omp.teams {
  omp.distribute {
    omp.wsloop %i... {
      BODY
    }
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp teams distribute simd`

</td>
<td>

```mlir
omp.teams {
  omp.distsimdloop %i... {
    BODY
  }
}
```

</td>
<td>

```mlir
omp.teams {
  omp.distribute {
    omp.simdloop %i... {
      BODY
    }
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp teams distribute parallel do`

</td>
<td>

```mlir
omp.teams {
  omp.distparwsloop %i... {
    BODY
  }
}
```

</td>
<td>

```mlir
omp.teams {
  omp.distribute {
    omp.parallel {
      omp.wsloop %i... {
        BODY
      }
    }
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp teams distribute parallel do simd`

</td>
<td>

```mlir
omp.teams {
  omp.distparwssimdloop %i... {
    BODY
  }
}
```

</td>
<td>

```mlir
omp.teams {
  omp.distribute {
    omp.parallel {
      omp.wsloop %i... {
        omp.simd {
          BODY
        }
      }
    }
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp teams loop`

</td>
<td>

```mlir
omp.teams {
  omp.loop %i... {
    BODY
  }
}
```

</td>
<td>

```mlir
omp.teams {
  omp.loop %i... {
    BODY
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp target parallel`

</td>
<td>

```mlir
omp.target {
  omp.parallel {
    BODY
  }
}
```

</td>
<td>

```mlir
omp.target {
  omp.parallel {
    BODY
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp target parallel do`

</td>
<td>

```mlir
omp.target {
  omp.parallel {
    omp.wsloop %i... {
      BODY
    }
  }
}
```

</td>
<td>

```mlir
omp.target {
  omp.parallel {
    omp.wsloop %i... {
      BODY
    }
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp target parallel do simd`

</td>
<td>

```mlir
omp.target {
  omp.parallel {
    omp.wssimdloop %i... {
      BODY
    }
  }
}
```

</td>
<td>

```mlir
omp.target {
  omp.parallel {
    omp.wsloop %i... {
      omp.simd {
        BODY
      }
    }
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp target parallel loop`

</td>
<td>

```mlir
omp.target {
  omp.parallel {
    omp.loop %i... {
      BODY
    }
  }
}
```

</td>
<td>

```mlir
omp.target {
  omp.parallel {
    omp.loop %i... {
      BODY
    }
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp target simd`

</td>
<td>

```mlir
omp.target {
  omp.simdloop %i... {
    BODY
  }
}
```

</td>
<td>

```mlir
omp.target {
  omp.simdloop %i... {
    BODY
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp target teams`

</td>
<td>

```mlir
omp.target {
  omp.teams {
    BODY
  }
}
```

</td>
<td>

```mlir
omp.target {
  omp.teams {
    BODY
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp target teams distribute`

</td>
<td>

```mlir
omp.target {
  omp.teams {
    omp.distloop %i... {
      BODY
    }
  }
}
```

</td>
<td>

```mlir
omp.target {
  omp.teams {
    omp.distribute {
      omp.wsloop %i... {
        BODY
      }
    }
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp target teams distribute simd`

</td>
<td>

```mlir
omp.target {
  omp.teams {
    omp.distsimdloop %i... {
      BODY
    }
  }
}
```

</td>
<td>

```mlir
omp.target {
  omp.teams {
    omp.distribute {
      omp.simdloop %i... {
        BODY
      }
    }
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp target teams loop`

</td>
<td>

```mlir
omp.target {
  omp.teams {
    omp.loop %i... {
      BODY
    }
  }
}
```

</td>
<td>

```mlir
omp.target {
  omp.teams {
    omp.loop %i... {
      BODY
    }
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp target teams distribute parallel do`

</td>
<td>

```mlir
omp.target {
  omp.teams {
    omp.distparwsloop %i... {
      BODY
    }
  }
}
```

</td>
<td>

```mlir
omp.target {
  omp.teams {
    omp.distribute {
      omp.parallel {
        omp.wsloop %i... {
          BODY
        }
      }
    }
  }
}
```

</td>
</tr>

<!-- ROW -->
<tr>
<td>

`!$omp target teams distribute parallel do simd`

</td>
<td>

```mlir
omp.target {
  omp.teams {
    omp.distparwssimdloop %i... {
      BODY
    }
  }
}
```

</td>
<td>

```mlir
omp.target {
  omp.teams {
    omp.distribute {
      omp.parallel {
        omp.wsloop %i... {
          omp.simd {
            BODY
          }
        }
      }
    }
  }
}
```

</td>
</tr>

</table>


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