Skip to content

[flang] Inline hlfir.reshape as hlfir.elemental. #124683

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
vzakhari merged 3 commits into from
Jan 30, 2025
Merged

[flang] Inline hlfir.reshape as hlfir.elemental. #124683

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
5b6b201
[flang] Inline hlfir.reshape as hlfir.elemental.
vzakhari Jan 28, 2025
30e193f
Fixed handling of dynamically optional PAD.
vzakhari Jan 28, 2025
b1adf5d
Added negative tests and a comment.
vzakhari Jan 28, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
208 changes: 208 additions & 0 deletions flang/lib/Optimizer/HLFIR/Transforms/SimplifyHLFIRIntrinsics.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -951,6 +951,213 @@ class DotProductConversion
}
};

class ReshapeAsElementalConversion
: public mlir::OpRewritePattern<hlfir::ReshapeOp> {
public:
using mlir::OpRewritePattern<hlfir::ReshapeOp>::OpRewritePattern;

llvm::LogicalResult
matchAndRewrite(hlfir::ReshapeOp reshape,
mlir::PatternRewriter &rewriter) const override {
// Do not inline RESHAPE with ORDER yet. The runtime implementation
// may be good enough, unless the temporary creation overhead
// is high.
// TODO: If ORDER is constant, then we can still easily inline.
// TODO: If the result's rank is 1, then we can assume ORDER == (/1/).
if (reshape.getOrder())
return rewriter.notifyMatchFailure(reshape,
"RESHAPE with ORDER argument");

// Verify that the element types of ARRAY, PAD and the result
// match before doing any transformations.
hlfir::Entity result = hlfir::Entity{reshape};
hlfir::Entity array = hlfir::Entity{reshape.getArray()};
mlir::Type elementType = array.getFortranElementType();
if (result.getFortranElementType() != elementType)
return rewriter.notifyMatchFailure(
reshape, "ARRAY and result have different types");
mlir::Value pad = reshape.getPad();
Copy link
Contributor

@jeanPerier jeanPerier Jan 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

PAD is dynamically optional. If its actual argument is an OPTIONAL/POINTER/ALLOCATABLE, its presence should be checked at runtime.

You probably need to do something about that here (or at least to detect and do not do the transformation for now).

Copy link
Contributor Author

@vzakhari vzakhari Jan 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks! I moved the reads from PAD under the check of whether we have to read from it or not.

if (pad && hlfir::getFortranElementType(pad.getType()) != elementType)
return rewriter.notifyMatchFailure(reshape,
"ARRAY and PAD have different types");

// TODO: selecting between ARRAY and PAD of non-trivial element types
// requires more work. We have to select between two references
// to elements in ARRAY and PAD. This requires conditional
// bufferization of the element, if ARRAY/PAD is an expression.
if (pad && !fir::isa_trivial(elementType))
return rewriter.notifyMatchFailure(reshape,
"PAD present with non-trivial type");

mlir::Location loc = reshape.getLoc();
fir::FirOpBuilder builder{rewriter, reshape.getOperation()};
// Assume that all the indices arithmetic does not overflow
// the IndexType.
builder.setIntegerOverflowFlags(mlir::arith::IntegerOverflowFlags::nuw);

llvm::SmallVector<mlir::Value, 1> typeParams;
hlfir::genLengthParameters(loc, builder, array, typeParams);

// Fetch the extents of ARRAY, PAD and result beforehand.
llvm::SmallVector<mlir::Value, Fortran::common::maxRank> arrayExtents =
hlfir::genExtentsVector(loc, builder, array);

mlir::Value arraySize, padSize;
llvm::SmallVector<mlir::Value, Fortran::common::maxRank> padExtents;
if (pad) {
// If PAD is present, we have to use array size to start taking
// elements from the PAD array.
arraySize = computeArraySize(loc, builder, arrayExtents);

padExtents = hlfir::genExtentsVector(loc, builder, hlfir::Entity{pad});
// PAD size is needed to wrap around the linear index addressing
// the PAD array.
padSize = computeArraySize(loc, builder, padExtents);
}
hlfir::Entity shape = hlfir::Entity{reshape.getShape()};
llvm::SmallVector<mlir::Value, Fortran::common::maxRank> resultExtents;
mlir::Type indexType = builder.getIndexType();
for (int idx = 0; idx < result.getRank(); ++idx)
resultExtents.push_back(hlfir::loadElementAt(
loc, builder, shape,
builder.createIntegerConstant(loc, indexType, idx + 1)));
auto resultShape = builder.create<fir::ShapeOp>(loc, resultExtents);

auto genKernel = [&](mlir::Location loc, fir::FirOpBuilder &builder,
mlir::ValueRange inputIndices) -> hlfir::Entity {
mlir::Value linearIndex =
computeLinearIndex(loc, builder, resultExtents, inputIndices);
fir::IfOp ifOp;
if (pad) {
// PAD is present. Check if this element comes from the PAD array.
mlir::Value isInsideArray = builder.create<mlir::arith::CmpIOp>(
loc, mlir::arith::CmpIPredicate::ult, linearIndex, arraySize);
ifOp = builder.create<fir::IfOp>(loc, elementType, isInsideArray,
/*withElseRegion=*/true);

// In the 'else' block, return an element from the PAD.
builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
// Subtract the ARRAY size from the zero-based linear index
// to get the zero-based linear index into PAD.
mlir::Value padLinearIndex =
builder.create<mlir::arith::SubIOp>(loc, linearIndex, arraySize);
// PAD wraps around, when additional elements are needed.
padLinearIndex =
builder.create<mlir::arith::RemUIOp>(loc, padLinearIndex, padSize);
llvm::SmallVector<mlir::Value, Fortran::common::maxRank> padIndices =
delinearizeIndex(loc, builder, padExtents, padLinearIndex);
mlir::Value padElement =
hlfir::loadElementAt(loc, builder, hlfir::Entity{pad}, padIndices);
builder.create<fir::ResultOp>(loc, padElement);

// In the 'then' block, return an element from the ARRAY.
builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
}

llvm::SmallVector<mlir::Value, Fortran::common::maxRank> arrayIndices =
delinearizeIndex(loc, builder, arrayExtents, linearIndex);
mlir::Value arrayElement =
hlfir::loadElementAt(loc, builder, array, arrayIndices);

if (ifOp) {
builder.create<fir::ResultOp>(loc, arrayElement);
builder.setInsertionPointAfter(ifOp);
arrayElement = ifOp.getResult(0);
}

return hlfir::Entity{arrayElement};
};
hlfir::ElementalOp elementalOp = hlfir::genElementalOp(
loc, builder, elementType, resultShape, typeParams, genKernel,
/*isUnordered=*/true,
/*polymorphicMold=*/result.isPolymorphic() ? array : mlir::Value{},
reshape.getResult().getType());
assert(elementalOp.getResult().getType() == reshape.getResult().getType());
rewriter.replaceOp(reshape, elementalOp);
return mlir::success();
}

private:
/// Compute zero-based linear index given an array extents
/// and one-based indices:
/// \p extents: [e0, e1, ..., en]
/// \p indices: [i0, i1, ..., in]
///
/// linear-index :=
/// (...((in-1)*e(n-1)+(i(n-1)-1))*e(n-2)+...)*e0+(i0-1)
static mlir::Value computeLinearIndex(mlir::Location loc,
fir::FirOpBuilder &builder,
mlir::ValueRange extents,
mlir::ValueRange indices) {
std::size_t rank = extents.size();
assert(rank = indices.size());
Copy link
Collaborator

@davemgreen davemgreen Jan 30, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: Should be ==? GCC was giving a warning.

Copy link
Contributor Author

@vzakhari vzakhari Jan 31, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you! == indeed.

Copy link
Contributor Author

@vzakhari vzakhari Jan 31, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed by 381416a

mlir::Type indexType = builder.getIndexType();
mlir::Value zero = builder.createIntegerConstant(loc, indexType, 0);
mlir::Value one = builder.createIntegerConstant(loc, indexType, 1);
mlir::Value linearIndex = zero;
for (auto idx : llvm::enumerate(llvm::reverse(indices))) {
mlir::Value tmp = builder.create<mlir::arith::SubIOp>(
loc, builder.createConvert(loc, indexType, idx.value()), one);
tmp = builder.create<mlir::arith::AddIOp>(loc, linearIndex, tmp);
if (idx.index() + 1 < rank)
tmp = builder.create<mlir::arith::MulIOp>(
loc, tmp,
builder.createConvert(loc, indexType,
extents[rank - idx.index() - 2]));

linearIndex = tmp;
}
return linearIndex;
}

/// Compute one-based array indices from the given zero-based \p linearIndex
/// and the array \p extents [e0, e1, ..., en].
/// i0 := linearIndex % e0 + 1
/// linearIndex := linearIndex / e0
/// i1 := linearIndex % e1 + 1
/// linearIndex := linearIndex / e1
/// ...
/// i(n-1) := linearIndex % e(n-1) + 1
/// linearIndex := linearIndex / e(n-1)
/// in := linearIndex + 1
static llvm::SmallVector<mlir::Value, Fortran::common::maxRank>
delinearizeIndex(mlir::Location loc, fir::FirOpBuilder &builder,
mlir::ValueRange extents, mlir::Value linearIndex) {
llvm::SmallVector<mlir::Value, Fortran::common::maxRank> indices;
mlir::Type indexType = builder.getIndexType();
mlir::Value one = builder.createIntegerConstant(loc, indexType, 1);
linearIndex = builder.createConvert(loc, indexType, linearIndex);

for (std::size_t dim = 0; dim < extents.size(); ++dim) {
mlir::Value currentIndex;
if (dim == extents.size() - 1) {
currentIndex = linearIndex;
} else {
mlir::Value extent =
builder.createConvert(loc, indexType, extents[dim]);
currentIndex =
builder.create<mlir::arith::RemUIOp>(loc, linearIndex, extent);
linearIndex =
builder.create<mlir::arith::DivUIOp>(loc, linearIndex, extent);
}
indices.push_back(
builder.create<mlir::arith::AddIOp>(loc, currentIndex, one));
}
return indices;
}

static mlir::Value computeArraySize(mlir::Location loc,
fir::FirOpBuilder &builder,
mlir::ValueRange extents) {
mlir::Type indexType = builder.getIndexType();
mlir::Value size = builder.createIntegerConstant(loc, indexType, 1);
for (auto extent : extents)
size = builder.create<mlir::arith::MulIOp>(
loc, size, builder.createConvert(loc, indexType, extent));
return size;
}
};

class SimplifyHLFIRIntrinsics
: public hlfir::impl::SimplifyHLFIRIntrinsicsBase<SimplifyHLFIRIntrinsics> {
public:
Expand Down Expand Up @@ -987,6 +1194,7 @@ class SimplifyHLFIRIntrinsics
patterns.insert<MatmulConversion<hlfir::MatmulOp>>(context);

patterns.insert<DotProductConversion>(context);
patterns.insert<ReshapeAsElementalConversion>(context);

if (mlir::failed(mlir::applyPatternsGreedily(
getOperation(), std::move(patterns), config))) {
Expand Down
Loading
Loading