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-exception
6//
7//===----------------------------------------------------------------------===//
8// In some special cases we can bufferize hlfir expressions in a more optimal
9// way so as to avoid creating temporaries. This pass handles these. It should
10// 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_OPTIMIZEDBUFFERIZATION
41#include "flang/Optimizer/HLFIR/Passes.h.inc"
42} // namespace hlfir
43
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 form
50/// %array = some.operation // array has shape %shape
51/// %expr = hlfir.elemental %shape : [...] {
52/// bb0(%arg0: index)
53/// %0 = hlfir.designate %array(%arg0)
54/// [...] // no other reads or writes to %array
55/// hlfir.yield_element %element
56/// }
57/// hlfir.assign %expr to %array
58/// hlfir.destroy %expr
59///
60/// Or
61///
62/// %read_array = some.operation // shape %shape
63/// %expr = hlfir.elemental %shape : [...] {
64/// bb0(%arg0: index)
65/// %0 = hlfir.designate %read_array(%arg0)
66/// [...]
67/// hlfir.yield_element %element
68/// }
69/// %write_array = some.operation // with shape %shape
70/// [...] // operations which don't effect write_array
71/// hlfir.assign %expr to %write_array
72/// hlfir.destroy %expr
73///
74/// In these cases, it is safe to turn the elemental into a do loop and modify
75/// elements of %array in place without creating an extra temporary for the
76/// elemental. We must check that there are no reads from the array at indexes
77/// which might conflict with the assignment or any writes. For now we will keep
78/// that strict and say that all reads must be at the elemental index (it is
79/// probably safe to read from higher indices if lowering to an ordered loop).
80class ElementalAssignBufferization
81 : 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 elemental
89 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 indices
93 /// into the array's lb-based indices, and returns the one-based indices
94 /// 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::LogicalResult
102 matchAndRewrite(hlfir::ElementalOp elemental,
103 mlir::PatternRewriter &rewriter) const override;
104};
105
106/// recursively collect all effects between start and end (including start, not
107/// including end) start must properly dominate end, start and end must be in
108/// the same block. If any operations with unknown effects are found,
109/// std::nullopt is returned
110static 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::NoAlias
133static mlir::AliasResult
134containsReadOrWriteEffectOn(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 val
149 mlir::AliasResult res = aliasAnalysis.alias(val, accessedVal);
150 if (!res.isNo())
151 return res;
152
153 // FIXME: alias analysis of fir.load
154 // follow this common pattern:
155 // %ref = hlfir.designate %array(%index)
156 // %val = fir.load $ref
157 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 val
162 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 represented
171// 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 EitherIdenticalOrDisjoint
185 };
186
187 // Analyzes two hlfir.designate results and returns the overlap kind.
188 // The callers may use this method when the alias analysis reports
189 // an alias of some kind, so that we can run Fortran specific analysis
190 // 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 answer
192 // 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, then
199 // 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 SectionDesc
231 // 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 result
234 // 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 ordered
244 // bounds match the <lb, ub> of the descriptor.
245 // If stride is known to be negative, then the ordered
246 // 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 else
260 return {desc.ub, desc.lb};
261 }
262
263 return {nullptr, nullptr};
264 }
265
266 // Given two array sections <lb1, ub1, stride1> and
267 // <lb2, ub2, stride2>, return true only if the sections
268 // are known to be disjoint.
269 //
270 // For example, for any positive constant C:
271 // X:Y does not overlap with (Y+C):Z
272 // 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 treated
281 // 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> and
288 // <lb2, ub2, stride2>, return true only if the sections
289 // 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 array
296 // elements are being addresses, even though the shape
297 // 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::SlicesOverlapKind
310ArraySectionAnalyzer::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 for
330 // 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 for
354 // particular relations between the lower and upper bounds
355 // of the array sections, e.g. for any positive constant C:
356 // X:Y does not overlap with (Y+C):Z
357 // 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, then
368 // we do not care if %0 is equal to %1 - the slices are either
369 // 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 else
391 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 + C
426 // v2 = C + v1
427 // v1 = v2 - C
428 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 adjusted
447 // according to the box's lower bound(s). Scan through
448 // 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 : index
454 // %18 = arith.addi %arg2, %17 : index
455 // %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 : index
463 // %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 destroy
512 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. for
518 // 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 into
536 // TODO: this could be extended to also allow hlfir.expr by first bufferizing
537 // the incoming expression
538 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 trivial
547 // TODO: this is just to make the pass easier to think about. Not an inherent
548 // limitation
549 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 be
558 // conformable unless the lhs is an allocatable array. In HLFIR we can
559 // see this from the presence or absence of the realloc attribute on
560 // hlfir.assign. If it is not a realloc assignment, we can trust that
561 // the shapes do conform.
562 //
563 // TODO: the lhs's shape is dynamic, so it is hard to prove that
564 // there is no reallocation of the lhs due to the assignment.
565 // We can probably try generating multiple versions of the code
566 // 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 the
573 // hlfir.assign, check there are no effects which make this unsafe
574
575 // keep track of any values written to in the elemental, as these can't be
576 // read from or written to between the elemental and the assignment
577 mlir::SmallVector<mlir::Value, 1> notToBeAccessedBeforeAssign;
578 // likewise, values read in the elemental cannot be written to between the
579 // elemental and the assign
580 mlir::SmallVector<mlir::Value, 1> notToBeWrittenBeforeAssign;
581
582 // 1) side effects in the elemental body - it isn't sufficient to just look
583 // for ordered elementals because we also cannot support out of order reads
584 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 elemental
603 continue;
604 }
605
606 // don't allow any aliasing writes in the elemental
607 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 order
619 // => each iteration doesn't read values written by other iterations
620 // don't allow reads from a different value which may alias: fir alias
621 // analysis isn't precise enough to tell us if two aliasing arrays overlap
622 // exactly or only partially. If they overlap partially, a designate at the
623 // elemental indices could be accessing different elements: e.g. we could
624 // designate two slices of the same array at different start indexes. These
625 // two MustAlias but index 1 of one array isn't the same element as index 1
626 // 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: " << designate
638 << " 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: " << designate
649 << " 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 assignment
659 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 write
668 // will be moved to the assignment
669 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 read
679 // 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 assignment
707 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 clone
717 // hlfir.elemental region inside the inner loop
718 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 builder.create<hlfir::AssignOp>(
731 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 nest
741/// 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 BroadcastAssignBufferization
752 : public mlir::OpRewritePattern<hlfir::AssignOp> {
753private:
754public:
755 using mlir::OpRewritePattern<hlfir::AssignOp>::OpRewritePattern;
756
757 llvm::LogicalResult
758 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 allocated
765 // in a conforming Fortran program, and LHS cannot be reallocated
766 // as a result of the assignment. So we can ignore isAllocatableAssignment
767 // 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 better
793 // optimized.
794 mlir::Value n = extents[0];
795 for (size_t i = 1; i < extents.size(); ++i)
796 n = builder.create<mlir::arith::MulIOp>(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 = builder.create<fir::ReboxOp>(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 builder.create<hlfir::DesignateOp>(loc, fir::ReferenceType::get(eleTy),
826 flatArray, loopNest.oneBasedIndices);
827 builder.create<hlfir::AssignOp>(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 builder.create<hlfir::AssignOp>(loc, rhs, arrayElement);
836 }
837
838 rewriter.eraseOp(assign);
839 return mlir::success();
840}
841
842class EvaluateIntoMemoryAssignBufferization
843 : public mlir::OpRewritePattern<hlfir::EvaluateInMemoryOp> {
844
845public:
846 using mlir::OpRewritePattern<hlfir::EvaluateInMemoryOp>::OpRewritePattern;
847
848 llvm::LogicalResult
849 matchAndRewrite(hlfir::EvaluateInMemoryOp,
850 mlir::PatternRewriter &rewriter) const override;
851};
852
853static llvm::LogicalResult
854tryUsingAssignLhsDirectly(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 be
872 // 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 LHS
876 // cannot be used to evaluate RHS in place without proving the LHS and
877 // 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 because
883 // EvaluateInMemoryOp is typically expected to hold fir.calls and that
884 // 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 uses
888 // mlir::MemoryEffects for operations that do not have special handling in
889 // getModRef.
890 if (aliasAnalysis.getModRef(evalInMem.getBody(), lhs).isModOrRef())
891 return mlir::failure();
892 // Any variables affected between the hlfir.evalInMem and assignment must not
893 // be read or written inside the region since it will be moved at the
894 // 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 can
925 // kick-in for simple types and at least implement the assignment inline
926 // 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 OptimizedBufferizationPass
937 : 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 blocks
945 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.assign
951 // operations and decide which transformation to apply
952 // at one place (e.g. we may use some heuristics and
953 // 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} // namespace
968

source code of flang/lib/Optimizer/HLFIR/Transforms/OptimizedBufferization.cpp