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(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 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(newBox.getDefiningOp())); + auto embox = llvm::dyn_cast(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(shape.getDefiningOp())); + auto shapeOp = llvm::dyn_cast(shape.getDefiningOp()); + ASSERT_EQ(shapeOp.getExtents().size(), 1u); + mlir::Value extent0 = shapeOp.getExtents()[0]; + ASSERT_TRUE(llvm::isa_and_nonnull(extent0.getDefiningOp())); + auto dimOp = llvm::dyn_cast(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(origin0.getDefiningOp())); + auto lbOp = llvm::dyn_cast(origin0.getDefiningOp()); + EXPECT_EQ(lbOp.getVal(), inputBox); +}