[flang] add helper to create descriptor with new base address

[jeanPerier]

There is currently no helper to create a descriptor for a copy of a Fortran entity based on the descriptor of the original entity and the base address of the copy (most places that are doing this currently are also doing allocation of the copy at the same time or using the runtime).
Add a helper for this with a unit test.

[llvmbot]

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

Author: None (jeanPerier)

Changes

There is currently no helper to create a descriptor for a copy of a Fortran entity based on the descriptor of the original entity and the base address of the copy (most places that are doing this currently are also doing allocation of the copy at the same time or using the runtime).
Add a helper for this with a unit test.

---

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

3 Files Affected:

• (modified) flang/include/flang/Optimizer/Builder/FIRBuilder.h (+9)
• (modified) flang/lib/Optimizer/Builder/FIRBuilder.cpp (+22)
• (modified) flang/unittests/Optimizer/Builder/FIRBuilderTest.cpp (+47)

diff --git a/flang/include/flang/Optimizer/Builder/FIRBuilder.h b/flang/include/flang/Optimizer/Builder/FIRBuilder.h
index 4b3087ed45788..d3af3bafbf279 100644
--- a/flang/include/flang/Optimizer/Builder/FIRBuilder.h
+++ b/flang/include/flang/Optimizer/Builder/FIRBuilder.h
@@ -959,6 +959,15 @@ mlir::Value genLifetimeStart(mlir::OpBuilder &builder, mlir::Location loc,
 void genLifetimeEnd(mlir::OpBuilder &builder, mlir::Location loc,
                     mlir::Value mem);
 
+/// Given a fir.box or fir.class \p box describing an entity and a raw address
+/// \p newAddr for an entity with the same Fortran properties (rank, dynamic
+/// type, length parameters and bounds) and attributes (POINTER or ALLOCATABLE),
+/// create a box for \p newAddr with the same type as \p box. This assumes \p
+/// newAddr is for contiguous storage (\p box does not have to be contiguous).
+mlir::Value getDescriptorWithNewBaseAddress(fir::FirOpBuilder &builder,
+                                            mlir::Location loc, mlir::Value box,
+                                            mlir::Value newAddr);
+
 } // namespace fir::factory
 
 #endif // FORTRAN_OPTIMIZER_BUILDER_FIRBUILDER_H
diff --git a/flang/lib/Optimizer/Builder/FIRBuilder.cpp b/flang/lib/Optimizer/Builder/FIRBuilder.cpp
index b6501fd530992..5e6e20861fd85 100644
--- a/flang/lib/Optimizer/Builder/FIRBuilder.cpp
+++ b/flang/lib/Optimizer/Builder/FIRBuilder.cpp
@@ -1974,3 +1974,25 @@ void fir::factory::genLifetimeEnd(mlir::OpBuilder &builder, mlir::Location loc,
                                   mlir::Value cast) {
   mlir::LLVM::LifetimeEndOp::create(builder, loc, cast);
 }
+
+mlir::Value fir::factory::getDescriptorWithNewBaseAddress(
+    fir::FirOpBuilder &builder, mlir::Location loc, mlir::Value box,
+    mlir::Value newAddr) {
+  auto boxType = llvm::dyn_cast<fir::BaseBoxType>(box.getType());
+  assert(boxType &&
+         "expected a box type input in getDescriptorWithNewBaseAddress");
+  if (boxType.isAssumedRank())
+    TODO(loc, "changing descriptor base address for an assumed rank entity");
+  llvm::SmallVector<mlir::Value> lbounds;
+  fir::factory::genDimInfoFromBox(builder, loc, box, &lbounds,
+                                  /*extents=*/nullptr, /*strides=*/nullptr);
+  fir::BoxValue inputBoxValue(box, lbounds, /*explicitParams=*/{});
+  fir::ExtendedValue openedInput =
+      fir::factory::readBoxValue(builder, loc, inputBoxValue);
+  mlir::Value shape = fir::isArray(openedInput)
+                          ? builder.createShape(loc, openedInput)
+                          : mlir::Value{};
+  mlir::Value typeMold = fir::isPolymorphicType(boxType) ? box : mlir::Value{};
+  return builder.createBox(loc, boxType, newAddr, shape, /*slice=*/{},
+                           fir::getTypeParams(openedInput), typeMold);
+}
diff --git a/flang/unittests/Optimizer/Builder/FIRBuilderTest.cpp b/flang/unittests/Optimizer/Builder/FIRBuilderTest.cpp
index e3e364720af67..fffd4ab5446ca 100644
--- a/flang/unittests/Optimizer/Builder/FIRBuilderTest.cpp
+++ b/flang/unittests/Optimizer/Builder/FIRBuilderTest.cpp
@@ -644,3 +644,50 @@ TEST_F(FIRBuilderTest, genArithIntegerOverflow) {
   auto op4_ioff = op4_iofi.getOverflowAttr().getValue();
   EXPECT_EQ(op4_ioff, nsw);
 }
+
+TEST_F(FIRBuilderTest, getDescriptorWithNewBaseAddress) {
+  auto builder = getBuilder();
+  auto loc = builder.getUnknownLoc();
+
+  // Build an input fir.box for a 1-D array of i64 with constant extent 10.
+  auto i64Ty = builder.getI64Type();
+  auto seqTy = fir::SequenceType::get({10}, i64Ty);
+  auto refArrTy = fir::ReferenceType::get(seqTy);
+  auto ptrTy = fir::PointerType::get(seqTy);
+  auto boxTy = fir::BoxType::get(ptrTy);
+  // Create an undef box descriptor value (descriptor contents are unspecified).
+  mlir::Value inputBox = fir::UndefOp::create(builder, loc, boxTy);
+
+  // New base address (same element type and properties).
+  mlir::Value addr2 = fir::UndefOp::create(builder, loc, refArrTy);
+
+  mlir::Value newBox = fir::factory::getDescriptorWithNewBaseAddress(
+      builder, loc, inputBox, addr2);
+
+  // The returned descriptor must have the same type as the input box.
+  EXPECT_EQ(newBox.getType(), inputBox.getType());
+
+  // It must be constructed by an embox using the new base address.
+  ASSERT_TRUE(llvm::isa_and_nonnull<fir::EmboxOp>(newBox.getDefiningOp()));
+  auto embox = llvm::dyn_cast<fir::EmboxOp>(newBox.getDefiningOp());
+  EXPECT_EQ(embox.getMemref(), addr2);
+
+  // The shape should be derived from the input box; expect a fir.shape with one
+  // extent that comes from a fir.box_dims reading from the original input box.
+  mlir::Value shape = embox.getShape();
+  ASSERT_TRUE(shape);
+  ASSERT_TRUE(llvm::isa_and_nonnull<fir::ShapeShiftOp>(shape.getDefiningOp()));
+  auto shapeOp = llvm::dyn_cast<fir::ShapeShiftOp>(shape.getDefiningOp());
+  ASSERT_EQ(shapeOp.getExtents().size(), 1u);
+  mlir::Value extent0 = shapeOp.getExtents()[0];
+  ASSERT_TRUE(llvm::isa_and_nonnull<fir::BoxDimsOp>(extent0.getDefiningOp()));
+  auto dimOp = llvm::dyn_cast<fir::BoxDimsOp>(extent0.getDefiningOp());
+  EXPECT_EQ(dimOp.getVal(), inputBox);
+
+  // Also verify the origin comes from a BoxDims on the same input box.
+  ASSERT_EQ(shapeOp.getOrigins().size(), 1u);
+  mlir::Value origin0 = shapeOp.getOrigins()[0];
+  ASSERT_TRUE(llvm::isa_and_nonnull<fir::BoxDimsOp>(origin0.getDefiningOp()));
+  auto lbOp = llvm::dyn_cast<fir::BoxDimsOp>(origin0.getDefiningOp());
+  EXPECT_EQ(lbOp.getVal(), inputBox);
+}

[vzakhari]

LGTM