[flang-commits] [flang] [flang] Canonicalize fir.array_coor by pulling in embox/rebox. (PR #92858)

[via flang-commits]

================
@@ -417,6 +423,216 @@ mlir::LogicalResult fir::ArrayCoorOp::verify() {
   return mlir::success();
 }
 
+// Pull in fir.embox and fir.rebox into fir.array_coor when possible.
+struct SimplifyArrayCoorOp : public mlir::OpRewritePattern<fir::ArrayCoorOp> {
+  using mlir::OpRewritePattern<fir::ArrayCoorOp>::OpRewritePattern;
+  mlir::LogicalResult
+  matchAndRewrite(fir::ArrayCoorOp op,
+                  mlir::PatternRewriter &rewriter) const override {
+    mlir::Value memref = op.getMemref();
+    if (!mlir::isa<fir::BaseBoxType>(memref.getType()))
+      return mlir::failure();
+
+    mlir::Value boxedMemref, boxedShape, boxedSlice;
+    if (auto emboxOp =
+            mlir::dyn_cast_or_null<fir::EmboxOp>(memref.getDefiningOp())) {
+      boxedMemref = emboxOp.getMemref();
+      boxedShape = emboxOp.getShape();
+      boxedSlice = emboxOp.getSlice();
+      // If any of operands, that are not currently supported for migration
+      // to ArrayCoorOp, is present, don't rewrite.
+      if (!emboxOp.getTypeparams().empty() || emboxOp.getSourceBox() ||
+          emboxOp.getAccessMap())
+        return mlir::failure();
+    } else if (auto reboxOp = mlir::dyn_cast_or_null<fir::ReboxOp>(
+                   memref.getDefiningOp())) {
+      boxedMemref = reboxOp.getBox();
+      boxedShape = reboxOp.getShape();
+      boxedSlice = reboxOp.getSlice();
+    } else {
+      return mlir::failure();
+    }
+
+    // Slices changing the number of dimensions are not supported
+    // for array_coor yet.
+    unsigned origBoxRank;
+    if (mlir::isa<fir::BaseBoxType>(boxedMemref.getType()))
+      origBoxRank = fir::getBoxRank(boxedMemref.getType());
+    else if (auto arrTy = mlir::dyn_cast<fir::SequenceType>(
+                 fir::unwrapRefType(boxedMemref.getType())))
+      origBoxRank = arrTy.getDimension();
+    else
+      return mlir::failure();
+
+    if (fir::getBoxRank(memref.getType()) != origBoxRank)
+      return mlir::failure();
+
+    // Slices with substring are not supported by array_coor.
+    if (boxedSlice)
+      if (auto sliceOp =
+              mlir::dyn_cast_or_null<fir::SliceOp>(boxedSlice.getDefiningOp()))
+        if (!sliceOp.getSubstr().empty())
+          return mlir::failure();
+
+    // If embox/rebox and array_coor have conflicting shapes or slices,
+    // do nothing.
+    if (op.getShape() && boxedShape && boxedShape != op.getShape())
+      return mlir::failure();
+    if (op.getSlice() && boxedSlice && boxedSlice != op.getSlice())
+      return mlir::failure();
+
+    // The embox/rebox and array_coor either have compatible
+    // shape/slice at this point or shape/slice is null
+    // in one of them but not in the other.
+    // The compatibility means they are equal or both null.
+    if (!op.getShape()) {
+      if (boxedShape) {
+        if (op.getSlice()) {
+          if (!boxedSlice) {
+            // %0 = fir.rebox %arg(%shape)
+            // %1 = fir.array_coor %0 [%slice] %idx
+            // Both the slice indices and %idx are 1-based, so the rebox
+            // may be pulled in as:
+            // %1 = fir.array_coor %arg [%slice] %idx
+            boxedShape = nullptr;
+          } else {
+            // %0 = fir.rebox %arg(%shape) [%slice]
+            // %1 = fir.array_coor %0 [%slice] %idx
+            // I believe this FIR may only be valid if the shape specifies
+            // that all lower bounds are 1s and the slice's start indices
+            // are all 1s.
+            // We could pull in the rebox as:
+            // %1 = fir.array_coor %arg [%slice] %idx
+            // Do not do anything for the time being.
+            return mlir::failure();
+          }
+        } else { // !op.getSlice()
+          if (!boxedSlice) {
+            // %0 = fir.rebox %arg(%shape)
+            // %1 = fir.array_coor %0 %idx
+            // Pull in as:
+            // %1 = fir.array_coor %arg %idx
+            boxedShape = nullptr;
+          } else {
+            // %0 = fir.rebox %arg(%shape) [%slice]
+            // %1 = fir.array_coor %0 %idx
+            // Pull in as:
+            // %1 = fir.array_coor %arg(%shape) %[slice] %idx
+          }
+        }
+      } else { // !boxedShape
+        if (op.getSlice()) {
+          if (!boxedSlice) {
+            // %0 = fir.rebox %arg
+            // %1 = fir.array_coor %0 [%slice] %idx
+            // Pull in as:
+            // %1 = fir.array_coor %arg [%slice] %idx
+          } else {
+            // %0 = fir.rebox %arg [%slice]
+            // %1 = fir.array_coor %0 [%slice] %idx
+            // Pull in as:
+            // %1 = fir.array_coor %arg [%slice] %idx
----------------
jeanPerier wrote:

Good catch Renaud, yes this is a actually changing the semantics of fir.array_coor and is not a bug fix as I also thought it was. Lowering without HLFIR used to generate indices that were also based on the shift.

```
subroutine foo(x)
  real :: x(10:19)
  x(12:15) = 0.
end subroutine
```
Was lowered by `bbc -emit-fir --hlfir=false -o - | fir-opt --array-value-copy -canonicalize -o -` to
```
    %0 = fir.shape_shift %c10, %c10 : (index, index) -> !fir.shapeshift<1>
    %1 = fir.slice %c12, %c15, %c1 : (index, index, index) -> !fir.slice<1>
    %2 = fir.undefined !fir.array<10xf32>
    %3 = fir.do_loop %arg1 = %c0 to %c3 step %c1 unordered iter_args(%arg2 = %2) -> (!fir.array<10xf32>) {
      %4 = arith.addi %arg1, %c10 : index
      %5 = fir.array_coor %arg0(%0) [%1] %4 : (!fir.ref<!fir.array<10xf32>>, !fir.shapeshift<1>, !fir.slice<1>, index) -> !fir.ref<f32>
      fir.store %cst to %5 : !fir.ref<f32>
      fir.result %2 : !fir.array<10xf32>
    }
    return
  }
```

It is a bit dumb since this requires an extra addi in FIR that will just later get combined with the subi in codegen. Fortran semantics for array sections is to be one based (`LBOUND(x(2:10)) = 1`).

If we keep the current fir.array_coor semantics, that would imply that the patterns should actually change the indices into "old_indices + shift_lb - 1". (the only case where you would not need to do that was if the pattern was `%box = fir.rebox %shift  %slice .... fir.array_coor %box %shift %indices`, but that is never generated by lowering).

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