968 lines · cpp
1//===- OptimizedBufferization.cpp - special cases for bufferization -------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8// In some special cases we can bufferize hlfir expressions in a more optimal9// way so as to avoid creating temporaries. This pass handles these. It should10// be run before the catch-all bufferization pass.11//12// This requires constant subexpression elimination to have already been run.13//===----------------------------------------------------------------------===//14 15#include "flang/Optimizer/Analysis/AliasAnalysis.h"16#include "flang/Optimizer/Builder/FIRBuilder.h"17#include "flang/Optimizer/Builder/HLFIRTools.h"18#include "flang/Optimizer/Dialect/FIROps.h"19#include "flang/Optimizer/Dialect/FIRType.h"20#include "flang/Optimizer/HLFIR/HLFIRDialect.h"21#include "flang/Optimizer/HLFIR/HLFIROps.h"22#include "flang/Optimizer/HLFIR/Passes.h"23#include "flang/Optimizer/OpenMP/Passes.h"24#include "flang/Optimizer/Support/Utils.h"25#include "flang/Optimizer/Transforms/Utils.h"26#include "mlir/Dialect/Func/IR/FuncOps.h"27#include "mlir/IR/Dominance.h"28#include "mlir/IR/PatternMatch.h"29#include "mlir/Interfaces/SideEffectInterfaces.h"30#include "mlir/Pass/Pass.h"31#include "mlir/Support/LLVM.h"32#include "mlir/Transforms/GreedyPatternRewriteDriver.h"33#include "llvm/ADT/TypeSwitch.h"34#include <iterator>35#include <memory>36#include <mlir/Analysis/AliasAnalysis.h>37#include <optional>38 39namespace hlfir {40#define GEN_PASS_DEF_OPTIMIZEDBUFFERIZATION41#include "flang/Optimizer/HLFIR/Passes.h.inc"42} // namespace hlfir43 44#define DEBUG_TYPE "opt-bufferization"45 46namespace {47 48/// This transformation should match in place modification of arrays.49/// It should match code of the form50/// %array = some.operation // array has shape %shape51/// %expr = hlfir.elemental %shape : [...] {52/// bb0(%arg0: index)53/// %0 = hlfir.designate %array(%arg0)54/// [...] // no other reads or writes to %array55/// hlfir.yield_element %element56/// }57/// hlfir.assign %expr to %array58/// hlfir.destroy %expr59///60/// Or61///62/// %read_array = some.operation // shape %shape63/// %expr = hlfir.elemental %shape : [...] {64/// bb0(%arg0: index)65/// %0 = hlfir.designate %read_array(%arg0)66/// [...]67/// hlfir.yield_element %element68/// }69/// %write_array = some.operation // with shape %shape70/// [...] // operations which don't effect write_array71/// hlfir.assign %expr to %write_array72/// hlfir.destroy %expr73///74/// In these cases, it is safe to turn the elemental into a do loop and modify75/// elements of %array in place without creating an extra temporary for the76/// elemental. We must check that there are no reads from the array at indexes77/// which might conflict with the assignment or any writes. For now we will keep78/// that strict and say that all reads must be at the elemental index (it is79/// probably safe to read from higher indices if lowering to an ordered loop).80class ElementalAssignBufferization81 : public mlir::OpRewritePattern<hlfir::ElementalOp> {82private:83 struct MatchInfo {84 mlir::Value array;85 hlfir::AssignOp assign;86 hlfir::DestroyOp destroy;87 };88 /// determines if the transformation can be applied to this elemental89 static std::optional<MatchInfo> findMatch(hlfir::ElementalOp elemental);90 91 /// Returns the array indices for the given hlfir.designate.92 /// It recognizes the computations used to transform the one-based indices93 /// into the array's lb-based indices, and returns the one-based indices94 /// in these cases.95 static llvm::SmallVector<mlir::Value>96 getDesignatorIndices(hlfir::DesignateOp designate);97 98public:99 using mlir::OpRewritePattern<hlfir::ElementalOp>::OpRewritePattern;100 101 llvm::LogicalResult102 matchAndRewrite(hlfir::ElementalOp elemental,103 mlir::PatternRewriter &rewriter) const override;104};105 106/// recursively collect all effects between start and end (including start, not107/// including end) start must properly dominate end, start and end must be in108/// the same block. If any operations with unknown effects are found,109/// std::nullopt is returned110static std::optional<mlir::SmallVector<mlir::MemoryEffects::EffectInstance>>111getEffectsBetween(mlir::Operation *start, mlir::Operation *end) {112 mlir::SmallVector<mlir::MemoryEffects::EffectInstance> ret;113 if (start == end)114 return ret;115 assert(start->getBlock() && end->getBlock() && "TODO: block arguments");116 assert(start->getBlock() == end->getBlock());117 assert(mlir::DominanceInfo{}.properlyDominates(start, end));118 119 mlir::Operation *nextOp = start;120 while (nextOp && nextOp != end) {121 std::optional<mlir::SmallVector<mlir::MemoryEffects::EffectInstance>>122 effects = mlir::getEffectsRecursively(nextOp);123 if (!effects)124 return std::nullopt;125 ret.append(*effects);126 nextOp = nextOp->getNextNode();127 }128 return ret;129}130 131/// If effect is a read or write on val, return whether it aliases.132/// Otherwise return mlir::AliasResult::NoAlias133static mlir::AliasResult134containsReadOrWriteEffectOn(const mlir::MemoryEffects::EffectInstance &effect,135 mlir::Value val) {136 fir::AliasAnalysis aliasAnalysis;137 138 if (mlir::isa<mlir::MemoryEffects::Read, mlir::MemoryEffects::Write>(139 effect.getEffect())) {140 mlir::Value accessedVal = effect.getValue();141 if (mlir::isa<fir::DebuggingResource>(effect.getResource()))142 return mlir::AliasResult::NoAlias;143 if (!accessedVal)144 return mlir::AliasResult::MayAlias;145 if (accessedVal == val)146 return mlir::AliasResult::MustAlias;147 148 // if the accessed value might alias val149 mlir::AliasResult res = aliasAnalysis.alias(val, accessedVal);150 if (!res.isNo())151 return res;152 153 // FIXME: alias analysis of fir.load154 // follow this common pattern:155 // %ref = hlfir.designate %array(%index)156 // %val = fir.load $ref157 if (auto designate = accessedVal.getDefiningOp<hlfir::DesignateOp>()) {158 if (designate.getMemref() == val)159 return mlir::AliasResult::MustAlias;160 161 // if the designate is into an array that might alias val162 res = aliasAnalysis.alias(val, designate.getMemref());163 if (!res.isNo())164 return res;165 }166 }167 return mlir::AliasResult::NoAlias;168}169 170// Helper class for analyzing two array slices represented171// by two hlfir.designate operations.172class ArraySectionAnalyzer {173public:174 // The result of the analyzis is one of the values below.175 enum class SlicesOverlapKind {176 // Slices overlap is unknown.177 Unknown,178 // Slices are definitely identical.179 DefinitelyIdentical,180 // Slices are definitely disjoint.181 DefinitelyDisjoint,182 // Slices may be either disjoint or identical,183 // i.e. there is definitely no partial overlap.184 EitherIdenticalOrDisjoint185 };186 187 // Analyzes two hlfir.designate results and returns the overlap kind.188 // The callers may use this method when the alias analysis reports189 // an alias of some kind, so that we can run Fortran specific analysis190 // on the array slices to see if they are identical or disjoint.191 // Note that the alias analysis are not able to give such an answer192 // about the references.193 static SlicesOverlapKind analyze(mlir::Value ref1, mlir::Value ref2);194 195private:196 struct SectionDesc {197 // An array section is described by <lb, ub, stride> tuple.198 // If the designator's subscript is not a triple, then199 // the section descriptor is constructed as <lb, nullptr, nullptr>.200 mlir::Value lb, ub, stride;201 202 SectionDesc(mlir::Value lb, mlir::Value ub, mlir::Value stride)203 : lb(lb), ub(ub), stride(stride) {204 assert(lb && "lower bound or index must be specified");205 normalize();206 }207 208 // Normalize the section descriptor:209 // 1. If UB is nullptr, then it is set to LB.210 // 2. If LB==UB, then stride does not matter,211 // so it is reset to nullptr.212 // 3. If STRIDE==1, then it is reset to nullptr.213 void normalize() {214 if (!ub)215 ub = lb;216 if (lb == ub)217 stride = nullptr;218 if (stride)219 if (auto val = fir::getIntIfConstant(stride))220 if (*val == 1)221 stride = nullptr;222 }223 224 bool operator==(const SectionDesc &other) const {225 return lb == other.lb && ub == other.ub && stride == other.stride;226 }227 };228 229 // Given an operand_iterator over the indices operands,230 // read the subscript values and return them as SectionDesc231 // updating the iterator. If isTriplet is true,232 // the subscript is a triplet, and the result is <lb, ub, stride>.233 // Otherwise, the subscript is a scalar index, and the result234 // is <index, nullptr, nullptr>.235 static SectionDesc readSectionDesc(mlir::Operation::operand_iterator &it,236 bool isTriplet) {237 if (isTriplet)238 return {*it++, *it++, *it++};239 return {*it++, nullptr, nullptr};240 }241 242 // Return the ordered lower and upper bounds of the section.243 // If stride is known to be non-negative, then the ordered244 // bounds match the <lb, ub> of the descriptor.245 // If stride is known to be negative, then the ordered246 // bounds are <ub, lb> of the descriptor.247 // If stride is unknown, we cannot deduce any order,248 // so the result is <nullptr, nullptr>249 static std::pair<mlir::Value, mlir::Value>250 getOrderedBounds(const SectionDesc &desc) {251 mlir::Value stride = desc.stride;252 // Null stride means stride=1.253 if (!stride)254 return {desc.lb, desc.ub};255 // Reverse the bounds, if stride is negative.256 if (auto val = fir::getIntIfConstant(stride)) {257 if (*val >= 0)258 return {desc.lb, desc.ub};259 else260 return {desc.ub, desc.lb};261 }262 263 return {nullptr, nullptr};264 }265 266 // Given two array sections <lb1, ub1, stride1> and267 // <lb2, ub2, stride2>, return true only if the sections268 // are known to be disjoint.269 //270 // For example, for any positive constant C:271 // X:Y does not overlap with (Y+C):Z272 // X:Y does not overlap with Z:(X-C)273 static bool areDisjointSections(const SectionDesc &desc1,274 const SectionDesc &desc2) {275 auto [lb1, ub1] = getOrderedBounds(desc1);276 auto [lb2, ub2] = getOrderedBounds(desc2);277 if (!lb1 || !lb2)278 return false;279 // Note that this comparison must be made on the ordered bounds,280 // otherwise 'a(x:y:1) = a(z:x-1:-1) + 1' may be incorrectly treated281 // as not overlapping (x=2, y=10, z=9).282 if (isLess(ub1, lb2) || isLess(ub2, lb1))283 return true;284 return false;285 }286 287 // Given two array sections <lb1, ub1, stride1> and288 // <lb2, ub2, stride2>, return true only if the sections289 // are known to be identical.290 //291 // For example:292 // <x, x, stride>293 // <x, nullptr, nullptr>294 //295 // These sections are identical, from the point of which array296 // elements are being addresses, even though the shape297 // of the array slices might be different.298 static bool areIdenticalSections(const SectionDesc &desc1,299 const SectionDesc &desc2) {300 if (desc1 == desc2)301 return true;302 return false;303 }304 305 // Return true, if v1 is known to be less than v2.306 static bool isLess(mlir::Value v1, mlir::Value v2);307};308 309ArraySectionAnalyzer::SlicesOverlapKind310ArraySectionAnalyzer::analyze(mlir::Value ref1, mlir::Value ref2) {311 if (ref1 == ref2)312 return SlicesOverlapKind::DefinitelyIdentical;313 314 auto des1 = ref1.getDefiningOp<hlfir::DesignateOp>();315 auto des2 = ref2.getDefiningOp<hlfir::DesignateOp>();316 // We only support a pair of designators right now.317 if (!des1 || !des2)318 return SlicesOverlapKind::Unknown;319 320 if (des1.getMemref() != des2.getMemref()) {321 // If the bases are different, then there is unknown overlap.322 LLVM_DEBUG(llvm::dbgs() << "No identical base for:\n"323 << des1 << "and:\n"324 << des2 << "\n");325 return SlicesOverlapKind::Unknown;326 }327 328 // Require all components of the designators to be the same.329 // It might be too strict, e.g. we may probably allow for330 // different type parameters.331 if (des1.getComponent() != des2.getComponent() ||332 des1.getComponentShape() != des2.getComponentShape() ||333 des1.getSubstring() != des2.getSubstring() ||334 des1.getComplexPart() != des2.getComplexPart() ||335 des1.getTypeparams() != des2.getTypeparams()) {336 LLVM_DEBUG(llvm::dbgs() << "Different designator specs for:\n"337 << des1 << "and:\n"338 << des2 << "\n");339 return SlicesOverlapKind::Unknown;340 }341 342 // Analyze the subscripts.343 auto des1It = des1.getIndices().begin();344 auto des2It = des2.getIndices().begin();345 bool identicalTriplets = true;346 bool identicalIndices = true;347 for (auto [isTriplet1, isTriplet2] :348 llvm::zip(des1.getIsTriplet(), des2.getIsTriplet())) {349 SectionDesc desc1 = readSectionDesc(des1It, isTriplet1);350 SectionDesc desc2 = readSectionDesc(des2It, isTriplet2);351 352 // See if we can prove that any of the sections do not overlap.353 // This is mostly a Polyhedron/nf performance hack that looks for354 // particular relations between the lower and upper bounds355 // of the array sections, e.g. for any positive constant C:356 // X:Y does not overlap with (Y+C):Z357 // X:Y does not overlap with Z:(X-C)358 if (areDisjointSections(desc1, desc2))359 return SlicesOverlapKind::DefinitelyDisjoint;360 361 if (!areIdenticalSections(desc1, desc2)) {362 if (isTriplet1 || isTriplet2) {363 // For example:364 // hlfir.designate %6#0 (%c2:%c7999:%c1, %c1:%c120:%c1, %0)365 // hlfir.designate %6#0 (%c2:%c7999:%c1, %c1:%c120:%c1, %1)366 //367 // If all the triplets (section speficiers) are the same, then368 // we do not care if %0 is equal to %1 - the slices are either369 // identical or completely disjoint.370 //371 // Also, treat these as identical sections:372 // hlfir.designate %6#0 (%c2:%c2:%c1)373 // hlfir.designate %6#0 (%c2)374 identicalTriplets = false;375 LLVM_DEBUG(llvm::dbgs() << "Triplet mismatch for:\n"376 << des1 << "and:\n"377 << des2 << "\n");378 } else {379 identicalIndices = false;380 LLVM_DEBUG(llvm::dbgs() << "Indices mismatch for:\n"381 << des1 << "and:\n"382 << des2 << "\n");383 }384 }385 }386 387 if (identicalTriplets) {388 if (identicalIndices)389 return SlicesOverlapKind::DefinitelyIdentical;390 else391 return SlicesOverlapKind::EitherIdenticalOrDisjoint;392 }393 394 LLVM_DEBUG(llvm::dbgs() << "Different sections for:\n"395 << des1 << "and:\n"396 << des2 << "\n");397 return SlicesOverlapKind::Unknown;398}399 400bool ArraySectionAnalyzer::isLess(mlir::Value v1, mlir::Value v2) {401 auto removeConvert = [](mlir::Value v) -> mlir::Operation * {402 auto *op = v.getDefiningOp();403 while (auto conv = mlir::dyn_cast_or_null<fir::ConvertOp>(op))404 op = conv.getValue().getDefiningOp();405 return op;406 };407 408 auto isPositiveConstant = [](mlir::Value v) -> bool {409 if (auto val = fir::getIntIfConstant(v))410 return *val > 0;411 return false;412 };413 414 auto *op1 = removeConvert(v1);415 auto *op2 = removeConvert(v2);416 if (!op1 || !op2)417 return false;418 419 // Check if they are both constants.420 if (auto val1 = fir::getIntIfConstant(op1->getResult(0)))421 if (auto val2 = fir::getIntIfConstant(op2->getResult(0)))422 return *val1 < *val2;423 424 // Handle some variable cases (C > 0):425 // v2 = v1 + C426 // v2 = C + v1427 // v1 = v2 - C428 if (auto addi = mlir::dyn_cast<mlir::arith::AddIOp>(op2))429 if ((addi.getLhs().getDefiningOp() == op1 &&430 isPositiveConstant(addi.getRhs())) ||431 (addi.getRhs().getDefiningOp() == op1 &&432 isPositiveConstant(addi.getLhs())))433 return true;434 if (auto subi = mlir::dyn_cast<mlir::arith::SubIOp>(op1))435 if (subi.getLhs().getDefiningOp() == op2 &&436 isPositiveConstant(subi.getRhs()))437 return true;438 return false;439}440 441llvm::SmallVector<mlir::Value>442ElementalAssignBufferization::getDesignatorIndices(443 hlfir::DesignateOp designate) {444 mlir::Value memref = designate.getMemref();445 446 // If the object is a box, then the indices may be adjusted447 // according to the box's lower bound(s). Scan through448 // the computations to try to find the one-based indices.449 if (mlir::isa<fir::BaseBoxType>(memref.getType())) {450 // Look for the following pattern:451 // %13 = fir.load %12 : !fir.ref<!fir.box<...>452 // %14:3 = fir.box_dims %13, %c0 : (!fir.box<...>, index) -> ...453 // %17 = arith.subi %14#0, %c1 : index454 // %18 = arith.addi %arg2, %17 : index455 // %19 = hlfir.designate %13 (%18) : (!fir.box<...>, index) -> ...456 //457 // %arg2 is a one-based index.458 459 auto isNormalizedLb = [memref](mlir::Value v, unsigned dim) {460 // Return true, if v and dim are such that:461 // %14:3 = fir.box_dims %13, %dim : (!fir.box<...>, index) -> ...462 // %17 = arith.subi %14#0, %c1 : index463 // %19 = hlfir.designate %13 (...) : (!fir.box<...>, index) -> ...464 if (auto subOp =465 mlir::dyn_cast_or_null<mlir::arith::SubIOp>(v.getDefiningOp())) {466 auto cst = fir::getIntIfConstant(subOp.getRhs());467 if (!cst || *cst != 1)468 return false;469 if (auto dimsOp = mlir::dyn_cast_or_null<fir::BoxDimsOp>(470 subOp.getLhs().getDefiningOp())) {471 if (memref != dimsOp.getVal() ||472 dimsOp.getResult(0) != subOp.getLhs())473 return false;474 auto dimsOpDim = fir::getIntIfConstant(dimsOp.getDim());475 return dimsOpDim && dimsOpDim == dim;476 }477 }478 return false;479 };480 481 llvm::SmallVector<mlir::Value> newIndices;482 for (auto index : llvm::enumerate(designate.getIndices())) {483 if (auto addOp = mlir::dyn_cast_or_null<mlir::arith::AddIOp>(484 index.value().getDefiningOp())) {485 for (unsigned opNum = 0; opNum < 2; ++opNum)486 if (isNormalizedLb(addOp->getOperand(opNum), index.index())) {487 newIndices.push_back(addOp->getOperand((opNum + 1) % 2));488 break;489 }490 491 // If new one-based index was not added, exit early.492 if (newIndices.size() <= index.index())493 break;494 }495 }496 497 // If any of the indices is not adjusted to the array's lb,498 // then return the original designator indices.499 if (newIndices.size() != designate.getIndices().size())500 return designate.getIndices();501 502 return newIndices;503 }504 505 return designate.getIndices();506}507 508std::optional<ElementalAssignBufferization::MatchInfo>509ElementalAssignBufferization::findMatch(hlfir::ElementalOp elemental) {510 mlir::Operation::user_range users = elemental->getUsers();511 // the only uses of the elemental should be the assignment and the destroy512 if (std::distance(users.begin(), users.end()) != 2) {513 LLVM_DEBUG(llvm::dbgs() << "Too many uses of the elemental\n");514 return std::nullopt;515 }516 517 // If the ElementalOp must produce a temporary (e.g. for518 // finalization purposes), then we cannot inline it.519 if (hlfir::elementalOpMustProduceTemp(elemental)) {520 LLVM_DEBUG(llvm::dbgs() << "ElementalOp must produce a temp\n");521 return std::nullopt;522 }523 524 MatchInfo match;525 for (mlir::Operation *user : users)526 mlir::TypeSwitch<mlir::Operation *, void>(user)527 .Case([&](hlfir::AssignOp op) { match.assign = op; })528 .Case([&](hlfir::DestroyOp op) { match.destroy = op; });529 530 if (!match.assign || !match.destroy) {531 LLVM_DEBUG(llvm::dbgs() << "Couldn't find assign or destroy\n");532 return std::nullopt;533 }534 535 // the array is what the elemental is assigned into536 // TODO: this could be extended to also allow hlfir.expr by first bufferizing537 // the incoming expression538 match.array = match.assign.getLhs();539 mlir::Type arrayType = mlir::dyn_cast<fir::SequenceType>(540 fir::unwrapPassByRefType(match.array.getType()));541 if (!arrayType) {542 LLVM_DEBUG(llvm::dbgs() << "AssignOp's result is not an array\n");543 return std::nullopt;544 }545 546 // require that the array elements are trivial547 // TODO: this is just to make the pass easier to think about. Not an inherent548 // limitation549 mlir::Type eleTy = hlfir::getFortranElementType(arrayType);550 if (!fir::isa_trivial(eleTy)) {551 LLVM_DEBUG(llvm::dbgs() << "AssignOp's data type is not trivial\n");552 return std::nullopt;553 }554 555 // The array must have the same shape as the elemental.556 //557 // f2018 10.2.1.2 (3) requires the lhs and rhs of an assignment to be558 // conformable unless the lhs is an allocatable array. In HLFIR we can559 // see this from the presence or absence of the realloc attribute on560 // hlfir.assign. If it is not a realloc assignment, we can trust that561 // the shapes do conform.562 //563 // TODO: the lhs's shape is dynamic, so it is hard to prove that564 // there is no reallocation of the lhs due to the assignment.565 // We can probably try generating multiple versions of the code566 // with checking for the shape match, length parameters match, etc.567 if (match.assign.isAllocatableAssignment()) {568 LLVM_DEBUG(llvm::dbgs() << "AssignOp may involve (re)allocation of LHS\n");569 return std::nullopt;570 }571 572 // the transformation wants to apply the elemental in a do-loop at the573 // hlfir.assign, check there are no effects which make this unsafe574 575 // keep track of any values written to in the elemental, as these can't be576 // read from or written to between the elemental and the assignment577 mlir::SmallVector<mlir::Value, 1> notToBeAccessedBeforeAssign;578 // likewise, values read in the elemental cannot be written to between the579 // elemental and the assign580 mlir::SmallVector<mlir::Value, 1> notToBeWrittenBeforeAssign;581 582 // 1) side effects in the elemental body - it isn't sufficient to just look583 // for ordered elementals because we also cannot support out of order reads584 std::optional<mlir::SmallVector<mlir::MemoryEffects::EffectInstance>>585 effects = getEffectsBetween(&elemental.getBody()->front(),586 elemental.getBody()->getTerminator());587 if (!effects) {588 LLVM_DEBUG(llvm::dbgs()589 << "operation with unknown effects inside elemental\n");590 return std::nullopt;591 }592 for (const mlir::MemoryEffects::EffectInstance &effect : *effects) {593 mlir::AliasResult res = containsReadOrWriteEffectOn(effect, match.array);594 if (res.isNo()) {595 if (effect.getValue()) {596 if (mlir::isa<mlir::MemoryEffects::Write>(effect.getEffect()))597 notToBeAccessedBeforeAssign.push_back(effect.getValue());598 else if (mlir::isa<mlir::MemoryEffects::Read>(effect.getEffect()))599 notToBeWrittenBeforeAssign.push_back(effect.getValue());600 }601 602 // this is safe in the elemental603 continue;604 }605 606 // don't allow any aliasing writes in the elemental607 if (mlir::isa<mlir::MemoryEffects::Write>(effect.getEffect())) {608 LLVM_DEBUG(llvm::dbgs() << "write inside the elemental body\n");609 return std::nullopt;610 }611 612 if (effect.getValue() == nullptr) {613 LLVM_DEBUG(llvm::dbgs()614 << "side-effect with no value, cannot analyze further\n");615 return std::nullopt;616 }617 618 // allow if and only if the reads are from the elemental indices, in order619 // => each iteration doesn't read values written by other iterations620 // don't allow reads from a different value which may alias: fir alias621 // analysis isn't precise enough to tell us if two aliasing arrays overlap622 // exactly or only partially. If they overlap partially, a designate at the623 // elemental indices could be accessing different elements: e.g. we could624 // designate two slices of the same array at different start indexes. These625 // two MustAlias but index 1 of one array isn't the same element as index 1626 // of the other array.627 if (!res.isPartial()) {628 if (auto designate =629 effect.getValue().getDefiningOp<hlfir::DesignateOp>()) {630 ArraySectionAnalyzer::SlicesOverlapKind overlap =631 ArraySectionAnalyzer::analyze(match.array, designate.getMemref());632 if (overlap ==633 ArraySectionAnalyzer::SlicesOverlapKind::DefinitelyDisjoint)634 continue;635 636 if (overlap == ArraySectionAnalyzer::SlicesOverlapKind::Unknown) {637 LLVM_DEBUG(llvm::dbgs() << "possible read conflict: " << designate638 << " at " << elemental.getLoc() << "\n");639 return std::nullopt;640 }641 auto indices = getDesignatorIndices(designate);642 auto elementalIndices = elemental.getIndices();643 if (indices.size() == elementalIndices.size() &&644 std::equal(indices.begin(), indices.end(), elementalIndices.begin(),645 elementalIndices.end()))646 continue;647 648 LLVM_DEBUG(llvm::dbgs() << "possible read conflict: " << designate649 << " at " << elemental.getLoc() << "\n");650 return std::nullopt;651 }652 }653 LLVM_DEBUG(llvm::dbgs() << "disallowed side-effect: " << effect.getValue()654 << " for " << elemental.getLoc() << "\n");655 return std::nullopt;656 }657 658 // 2) look for conflicting effects between the elemental and the assignment659 effects = getEffectsBetween(elemental->getNextNode(), match.assign);660 if (!effects) {661 LLVM_DEBUG(662 llvm::dbgs()663 << "operation with unknown effects between elemental and assign\n");664 return std::nullopt;665 }666 for (const mlir::MemoryEffects::EffectInstance &effect : *effects) {667 // not safe to access anything written in the elemental as this write668 // will be moved to the assignment669 for (mlir::Value val : notToBeAccessedBeforeAssign) {670 mlir::AliasResult res = containsReadOrWriteEffectOn(effect, val);671 if (!res.isNo()) {672 LLVM_DEBUG(llvm::dbgs()673 << "disallowed side-effect: " << effect.getValue() << " for "674 << elemental.getLoc() << "\n");675 return std::nullopt;676 }677 }678 // Anything that is read inside the elemental can only be safely read679 // between the elemental and the assignment.680 for (mlir::Value val : notToBeWrittenBeforeAssign) {681 mlir::AliasResult res = containsReadOrWriteEffectOn(effect, val);682 if (!res.isNo() &&683 !mlir::isa<mlir::MemoryEffects::Read>(effect.getEffect())) {684 LLVM_DEBUG(llvm::dbgs()685 << "disallowed non-read side-effect: " << effect.getValue()686 << " for " << elemental.getLoc() << "\n");687 return std::nullopt;688 }689 }690 }691 692 return match;693}694 695llvm::LogicalResult ElementalAssignBufferization::matchAndRewrite(696 hlfir::ElementalOp elemental, mlir::PatternRewriter &rewriter) const {697 std::optional<MatchInfo> match = findMatch(elemental);698 if (!match)699 return rewriter.notifyMatchFailure(700 elemental, "cannot prove safety of ElementalAssignBufferization");701 702 mlir::Location loc = elemental->getLoc();703 fir::FirOpBuilder builder(rewriter, elemental.getOperation());704 auto rhsExtents = hlfir::getIndexExtents(loc, builder, elemental.getShape());705 706 // create the loop at the assignment707 builder.setInsertionPoint(match->assign);708 hlfir::Entity lhs{match->array};709 lhs = hlfir::derefPointersAndAllocatables(loc, builder, lhs);710 mlir::Value lhsShape = hlfir::genShape(loc, builder, lhs);711 llvm::SmallVector<mlir::Value> lhsExtents =712 hlfir::getIndexExtents(loc, builder, lhsShape);713 llvm::SmallVector<mlir::Value> extents =714 fir::factory::deduceOptimalExtents(rhsExtents, lhsExtents);715 716 // Generate a loop nest looping around the hlfir.elemental shape and clone717 // hlfir.elemental region inside the inner loop718 hlfir::LoopNest loopNest =719 hlfir::genLoopNest(loc, builder, extents, !elemental.isOrdered(),720 flangomp::shouldUseWorkshareLowering(elemental));721 builder.setInsertionPointToStart(loopNest.body);722 auto yield = hlfir::inlineElementalOp(loc, builder, elemental,723 loopNest.oneBasedIndices);724 hlfir::Entity elementValue{yield.getElementValue()};725 rewriter.eraseOp(yield);726 727 // Assign the element value to the array element for this iteration.728 auto arrayElement =729 hlfir::getElementAt(loc, builder, lhs, loopNest.oneBasedIndices);730 hlfir::AssignOp::create(731 builder, loc, elementValue, arrayElement, /*realloc=*/false,732 /*keep_lhs_length_if_realloc=*/false, match->assign.getTemporaryLhs());733 734 rewriter.eraseOp(match->assign);735 rewriter.eraseOp(match->destroy);736 rewriter.eraseOp(elemental);737 return mlir::success();738}739 740/// Expand hlfir.assign of a scalar RHS to array LHS into a loop nest741/// of element-by-element assignments:742/// hlfir.assign %cst to %0 : f32, !fir.ref<!fir.array<6x6xf32>>743/// into:744/// fir.do_loop %arg0 = %c1 to %c6 step %c1 unordered {745/// fir.do_loop %arg1 = %c1 to %c6 step %c1 unordered {746/// %1 = hlfir.designate %0 (%arg1, %arg0) :747/// (!fir.ref<!fir.array<6x6xf32>>, index, index) -> !fir.ref<f32>748/// hlfir.assign %cst to %1 : f32, !fir.ref<f32>749/// }750/// }751class BroadcastAssignBufferization752 : public mlir::OpRewritePattern<hlfir::AssignOp> {753private:754public:755 using mlir::OpRewritePattern<hlfir::AssignOp>::OpRewritePattern;756 757 llvm::LogicalResult758 matchAndRewrite(hlfir::AssignOp assign,759 mlir::PatternRewriter &rewriter) const override;760};761 762llvm::LogicalResult BroadcastAssignBufferization::matchAndRewrite(763 hlfir::AssignOp assign, mlir::PatternRewriter &rewriter) const {764 // Since RHS is a scalar and LHS is an array, LHS must be allocated765 // in a conforming Fortran program, and LHS cannot be reallocated766 // as a result of the assignment. So we can ignore isAllocatableAssignment767 // and do the transformation always.768 mlir::Value rhs = assign.getRhs();769 if (!fir::isa_trivial(rhs.getType()))770 return rewriter.notifyMatchFailure(771 assign, "AssignOp's RHS is not a trivial scalar");772 773 hlfir::Entity lhs{assign.getLhs()};774 if (!lhs.isArray())775 return rewriter.notifyMatchFailure(assign,776 "AssignOp's LHS is not an array");777 778 mlir::Type eleTy = lhs.getFortranElementType();779 if (!fir::isa_trivial(eleTy))780 return rewriter.notifyMatchFailure(781 assign, "AssignOp's LHS data type is not trivial");782 783 mlir::Location loc = assign->getLoc();784 fir::FirOpBuilder builder(rewriter, assign.getOperation());785 builder.setInsertionPoint(assign);786 lhs = hlfir::derefPointersAndAllocatables(loc, builder, lhs);787 mlir::Value shape = hlfir::genShape(loc, builder, lhs);788 llvm::SmallVector<mlir::Value> extents =789 hlfir::getIndexExtents(loc, builder, shape);790 791 if (lhs.isSimplyContiguous() && extents.size() > 1) {792 // Flatten the array to use a single assign loop, that can be better793 // optimized.794 mlir::Value n = extents[0];795 for (size_t i = 1; i < extents.size(); ++i)796 n = mlir::arith::MulIOp::create(builder, loc, n, extents[i]);797 llvm::SmallVector<mlir::Value> flatExtents = {n};798 799 mlir::Type flatArrayType;800 mlir::Value flatArray = lhs.getBase();801 if (mlir::isa<fir::BoxType>(lhs.getType())) {802 shape = builder.genShape(loc, flatExtents);803 flatArrayType = fir::BoxType::get(fir::SequenceType::get(eleTy, 1));804 flatArray = fir::ReboxOp::create(builder, loc, flatArrayType, flatArray,805 shape, /*slice=*/mlir::Value{});806 } else {807 // Array references must have fixed shape, when used in assignments.808 auto seqTy =809 mlir::cast<fir::SequenceType>(fir::unwrapRefType(lhs.getType()));810 llvm::ArrayRef<int64_t> fixedShape = seqTy.getShape();811 int64_t flatExtent = 1;812 for (int64_t extent : fixedShape)813 flatExtent *= extent;814 flatArrayType =815 fir::ReferenceType::get(fir::SequenceType::get({flatExtent}, eleTy));816 flatArray = builder.createConvert(loc, flatArrayType, flatArray);817 }818 819 hlfir::LoopNest loopNest =820 hlfir::genLoopNest(loc, builder, flatExtents, /*isUnordered=*/true,821 flangomp::shouldUseWorkshareLowering(assign));822 builder.setInsertionPointToStart(loopNest.body);823 824 mlir::Value arrayElement =825 hlfir::DesignateOp::create(builder, loc, fir::ReferenceType::get(eleTy),826 flatArray, loopNest.oneBasedIndices);827 hlfir::AssignOp::create(builder, loc, rhs, arrayElement);828 } else {829 hlfir::LoopNest loopNest =830 hlfir::genLoopNest(loc, builder, extents, /*isUnordered=*/true,831 flangomp::shouldUseWorkshareLowering(assign));832 builder.setInsertionPointToStart(loopNest.body);833 auto arrayElement =834 hlfir::getElementAt(loc, builder, lhs, loopNest.oneBasedIndices);835 hlfir::AssignOp::create(builder, loc, rhs, arrayElement);836 }837 838 rewriter.eraseOp(assign);839 return mlir::success();840}841 842class EvaluateIntoMemoryAssignBufferization843 : public mlir::OpRewritePattern<hlfir::EvaluateInMemoryOp> {844 845public:846 using mlir::OpRewritePattern<hlfir::EvaluateInMemoryOp>::OpRewritePattern;847 848 llvm::LogicalResult849 matchAndRewrite(hlfir::EvaluateInMemoryOp,850 mlir::PatternRewriter &rewriter) const override;851};852 853static llvm::LogicalResult854tryUsingAssignLhsDirectly(hlfir::EvaluateInMemoryOp evalInMem,855 mlir::PatternRewriter &rewriter) {856 mlir::Location loc = evalInMem.getLoc();857 hlfir::DestroyOp destroy;858 hlfir::AssignOp assign;859 for (auto user : llvm::enumerate(evalInMem->getUsers())) {860 if (user.index() > 2)861 return mlir::failure();862 mlir::TypeSwitch<mlir::Operation *, void>(user.value())863 .Case([&](hlfir::AssignOp op) { assign = op; })864 .Case([&](hlfir::DestroyOp op) { destroy = op; });865 }866 if (!assign || !destroy || destroy.mustFinalizeExpr() ||867 assign.isAllocatableAssignment())868 return mlir::failure();869 870 hlfir::Entity lhs{assign.getLhs()};871 // EvaluateInMemoryOp memory is contiguous, so in general, it can only be872 // replace by the LHS if the LHS is contiguous.873 if (!lhs.isSimplyContiguous())874 return mlir::failure();875 // Character assignment may involves truncation/padding, so the LHS876 // cannot be used to evaluate RHS in place without proving the LHS and877 // RHS lengths are the same.878 if (lhs.isCharacter())879 return mlir::failure();880 fir::AliasAnalysis aliasAnalysis;881 // The region must not read or write the LHS.882 // Note that getModRef is used instead of mlir::MemoryEffects because883 // EvaluateInMemoryOp is typically expected to hold fir.calls and that884 // Fortran calls cannot be modeled in a useful way with mlir::MemoryEffects:885 // it is hard/impossible to list all the read/written SSA values in a call,886 // but it is often possible to tell that an SSA value cannot be accessed,887 // hence getModRef is needed here and below. Also note that getModRef uses888 // mlir::MemoryEffects for operations that do not have special handling in889 // getModRef.890 if (aliasAnalysis.getModRef(evalInMem.getBody(), lhs).isModOrRef())891 return mlir::failure();892 // Any variables affected between the hlfir.evalInMem and assignment must not893 // be read or written inside the region since it will be moved at the894 // assignment insertion point.895 auto effects = getEffectsBetween(evalInMem->getNextNode(), assign);896 if (!effects) {897 LLVM_DEBUG(898 llvm::dbgs()899 << "operation with unknown effects between eval_in_mem and assign\n");900 return mlir::failure();901 }902 for (const mlir::MemoryEffects::EffectInstance &effect : *effects) {903 mlir::Value affected = effect.getValue();904 if (!affected ||905 aliasAnalysis.getModRef(evalInMem.getBody(), affected).isModOrRef())906 return mlir::failure();907 }908 909 rewriter.setInsertionPoint(assign);910 fir::FirOpBuilder builder(rewriter, evalInMem.getOperation());911 mlir::Value rawLhs = hlfir::genVariableRawAddress(loc, builder, lhs);912 hlfir::computeEvaluateOpIn(loc, builder, evalInMem, rawLhs);913 rewriter.eraseOp(assign);914 rewriter.eraseOp(destroy);915 rewriter.eraseOp(evalInMem);916 return mlir::success();917}918 919llvm::LogicalResult EvaluateIntoMemoryAssignBufferization::matchAndRewrite(920 hlfir::EvaluateInMemoryOp evalInMem,921 mlir::PatternRewriter &rewriter) const {922 if (mlir::succeeded(tryUsingAssignLhsDirectly(evalInMem, rewriter)))923 return mlir::success();924 // Rewrite to temp + as_expr here so that the assign + as_expr pattern can925 // kick-in for simple types and at least implement the assignment inline926 // instead of call Assign runtime.927 fir::FirOpBuilder builder(rewriter, evalInMem.getOperation());928 mlir::Location loc = evalInMem.getLoc();929 auto [temp, isHeapAllocated] = hlfir::computeEvaluateOpInNewTemp(930 loc, builder, evalInMem, evalInMem.getShape(), evalInMem.getTypeparams());931 rewriter.replaceOpWithNewOp<hlfir::AsExprOp>(932 evalInMem, temp, /*mustFree=*/builder.createBool(loc, isHeapAllocated));933 return mlir::success();934}935 936class OptimizedBufferizationPass937 : public hlfir::impl::OptimizedBufferizationBase<938 OptimizedBufferizationPass> {939public:940 void runOnOperation() override {941 mlir::MLIRContext *context = &getContext();942 943 mlir::GreedyRewriteConfig config;944 // Prevent the pattern driver from merging blocks945 config.setRegionSimplificationLevel(946 mlir::GreedySimplifyRegionLevel::Disabled);947 948 mlir::RewritePatternSet patterns(context);949 // TODO: right now the patterns are non-conflicting,950 // but it might be better to run this pass on hlfir.assign951 // operations and decide which transformation to apply952 // at one place (e.g. we may use some heuristics and953 // choose different optimization strategies).954 // This requires small code reordering in ElementalAssignBufferization.955 patterns.insert<ElementalAssignBufferization>(context);956 patterns.insert<BroadcastAssignBufferization>(context);957 patterns.insert<EvaluateIntoMemoryAssignBufferization>(context);958 959 if (mlir::failed(mlir::applyPatternsGreedily(960 getOperation(), std::move(patterns), config))) {961 mlir::emitError(getOperation()->getLoc(),962 "failure in HLFIR optimized bufferization");963 signalPassFailure();964 }965 }966};967} // namespace968