[flang] Allow embox's source_box to be a !fir.box.

[vzakhari]

In order to create temporary copies of assumed-type arrays
(e.g. for -frepack-arrays), we have to allow the source_box
to be a !fir.box.

This patch replaces #147618.

[llvmbot]

@llvm/pr-subscribers-flang-fir-hlfir

Author: Slava Zakharin (vzakhari)

Changes

In order to create temporary copies of assumed-type arrays
(e.g. for -frepack-arrays), we have to allow the source_box
to be a !fir.box.

This patch replaces #147618.

---

Full diff: https://github.com/llvm/llvm-project/pull/148305.diff

4 Files Affected:

• (modified) flang/include/flang/Optimizer/Dialect/FIRCG/CGOps.td (+2-1)
• (modified) flang/include/flang/Optimizer/Dialect/FIROps.td (+5-9)
• (modified) flang/include/flang/Optimizer/Dialect/FIRTypes.td (+4-2)
• (modified) flang/test/Fir/fir-ops.fir (+10)

diff --git a/flang/include/flang/Optimizer/Dialect/FIRCG/CGOps.td b/flang/include/flang/Optimizer/Dialect/FIRCG/CGOps.td
index d48caf1a8cb3b..04f839386498c 100644
--- a/flang/include/flang/Optimizer/Dialect/FIRCG/CGOps.td
+++ b/flang/include/flang/Optimizer/Dialect/FIRCG/CGOps.td
@@ -60,7 +60,8 @@ def fircg_XEmboxOp : fircg_Op<"ext_embox", [AttrSizedOperandSegments]> {
       Variadic<AnyIntegerType>:$slice,
       Variadic<AnyCoordinateType>:$subcomponent,
       Variadic<AnyIntegerType>:$substr, Variadic<AnyIntegerType>:$lenParams,
-      Optional<fir_ClassType>:$sourceBox, OptionalAttr<I32Attr>:$allocator_idx);
+      Optional<BoxOrClassType>:$sourceBox,
+      OptionalAttr<I32Attr>:$allocator_idx);
   let results = (outs BoxOrClassType);
 
   let assemblyFormat = [{
diff --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td
index e3f5c4403002a..99b5105ab365e 100644
--- a/flang/include/flang/Optimizer/Dialect/FIROps.td
+++ b/flang/include/flang/Optimizer/Dialect/FIROps.td
@@ -834,15 +834,11 @@ def fir_EmboxOp : fir_Op<"embox", [NoMemoryEffect, AttrSizedOperandSegments]> {
         - allocator_idx: specify special allocator to use.
   }];
 
-  let arguments = (ins
-    AnyReferenceLike:$memref,
-    Optional<AnyShapeType>:$shape,
-    Optional<fir_SliceType>:$slice,
-    Variadic<AnyIntegerType>:$typeparams,
-    Optional<fir_ClassType>:$sourceBox,
-    OptionalAttr<AffineMapAttr>:$accessMap,
-    OptionalAttr<I32Attr>:$allocator_idx
-  );
+  let arguments = (ins AnyReferenceLike:$memref, Optional<AnyShapeType>:$shape,
+      Optional<fir_SliceType>:$slice, Variadic<AnyIntegerType>:$typeparams,
+      Optional<BoxOrClassType>:$sourceBox,
+      OptionalAttr<AffineMapAttr>:$accessMap,
+      OptionalAttr<I32Attr>:$allocator_idx);
 
   let results = (outs BoxOrClassType);
 
diff --git a/flang/include/flang/Optimizer/Dialect/FIRTypes.td b/flang/include/flang/Optimizer/Dialect/FIRTypes.td
index 0ead54df3ca97..2fdc9a96804da 100644
--- a/flang/include/flang/Optimizer/Dialect/FIRTypes.td
+++ b/flang/include/flang/Optimizer/Dialect/FIRTypes.td
@@ -626,8 +626,10 @@ def AnyBoxLike : TypeConstraint<Or<[fir_BoxType.predicate,
     fir_BoxCharType.predicate, fir_BoxProcType.predicate,
     fir_ClassType.predicate]>, "any box">;
 
-def BoxOrClassType : TypeConstraint<Or<[fir_BoxType.predicate,
-    fir_ClassType.predicate]>, "box or class">;
+def BoxOrClassLike
+    : TypeConstraint<Or<[fir_BoxType.predicate, fir_ClassType.predicate]>,
+                     "box or class">;
+def BoxOrClassType : Type<BoxOrClassLike.predicate, "box or class">;
 
 def AnyRefOrBoxLike : TypeConstraint<Or<[AnyReferenceLike.predicate,
     AnyBoxLike.predicate, FunctionType.predicate]>,
diff --git a/flang/test/Fir/fir-ops.fir b/flang/test/Fir/fir-ops.fir
index 3585bf9efca3e..0892eb9fa0de8 100644
--- a/flang/test/Fir/fir-ops.fir
+++ b/flang/test/Fir/fir-ops.fir
@@ -843,6 +843,16 @@ func.func @embox_tdesc(%arg0: !fir.class<!fir.array<10x!fir.type<derived_poly{a:
   return
 }
 
+// CHECK-LABEL:   func.func @embox_with_source_box(
+// CHECK-SAME:      %[[ARG0:.*]]: !fir.heap<!fir.array<?xnone>>,
+// CHECK-SAME:      %[[ARG1:.*]]: !fir.shape<1>,
+// CHECK-SAME:      %[[ARG2:.*]]: !fir.box<!fir.array<?xnone>>) -> !fir.class<!fir.heap<!fir.array<?xnone>>> {
+func.func @embox_with_source_box(%arg0: !fir.heap<!fir.array<?xnone>>, %arg1: !fir.shape<1>, %arg2: !fir.box<!fir.array<?xnone>>) -> !fir.class<!fir.heap<!fir.array<?xnone>>> {
+// CHECK:           %[[VAL_0:.*]] = fir.embox %[[ARG0]](%[[ARG1]]) source_box %[[ARG2]] : (!fir.heap<!fir.array<?xnone>>, !fir.shape<1>, !fir.box<!fir.array<?xnone>>) -> !fir.class<!fir.heap<!fir.array<?xnone>>>
+  %0 = fir.embox %arg0(%arg1) source_box %arg2 : (!fir.heap<!fir.array<?xnone>>, !fir.shape<1>, !fir.box<!fir.array<?xnone>>) -> !fir.class<!fir.heap<!fir.array<?xnone>>>
+  return %0 : !fir.class<!fir.heap<!fir.array<?xnone>>>
+}
+
 func.func @test_fortran_var_attrs() {
   %0 = fir.alloca !fir.heap<f32> {fortran_attrs = #fir.var_attrs<allocatable>}
   %1 = fir.alloca f32 {fortran_attrs = #fir.var_attrs<asynchronous>}

[clementval]

LGTM