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

[Slava Zakharin 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
----------------
vzakhari wrote:

Thanks!

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