[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();
----------------
jeanPerier wrote:

For more safety, you could return here if the rebox has a shape that is not  a fir.shift, because that means there is some reshaping happening, and it would be incorrect to replace the fir.array_coor operand by the rebox input.
This pattern is likely not generated directly by lowering, but could happen after store/load forwarding of fir.box in:

```
subroutine foo(x)
  real, pointer :: p(:, :)
  real, contiguous, optional, target :: x(:, :)
  p(1:10, 1:20) => x
  call bar(p(3,4))
end subroutine
```
bbc -emit-fir:
````
    %9 = fir.shape_shift %c1, %c10, %c1, %c20 : (index, index, index, index) -> !fir.shapeshift<2>
    %10 = fir.rebox %6(%9) : (!fir.box<!fir.array<?x?xf32>>, !fir.shapeshift<2>) -> !fir.box<!fir.ptr<!fir.array<?x?xf32>>>
    fir.store %10 to %5 : !fir.ref<!fir.box<!fir.ptr<!fir.array<?x?xf32>>>>
    %11 = fir.load %5 : !fir.ref<!fir.box<!fir.ptr<!fir.array<?x?xf32>>>>
    %c0_0 = arith.constant 0 : index
    %12:3 = fir.box_dims %11, %c0_0 : (!fir.box<!fir.ptr<!fir.array<?x?xf32>>>, index) -> (index, index, index)
    %c1_1 = arith.constant 1 : index
    %13:3 = fir.box_dims %11, %c1_1 : (!fir.box<!fir.ptr<!fir.array<?x?xf32>>>, index) -> (index, index, index)
    %14 = fir.shift %12#0, %13#0 : (index, index) -> !fir.shift<2>
    %15 = fir.array_coor %11(%14) %c3, %c4 : (!fir.box<!fir.ptr<!fir.array<?x?xf32>>>, !fir.shift<2>, index, index) -> !fir.ref<f32>
````

Your requirement that the shape operand must be the same if present would prevent any issues with the above example, but one can imagine similar FIR where the array_coor would have no shape operand, so I would advise rejecting rebox with shape<>/shapeshift<> operands.

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