/
LoopUtils.cpp
3082 lines (2700 loc) · 128 KB
/
LoopUtils.cpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
//===- LoopUtils.cpp ---- Misc utilities for loop transformation ----------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file implements miscellaneous loop transformation routines.
//
//===----------------------------------------------------------------------===//
#include "mlir/Transforms/LoopUtils.h"
#include "mlir/Analysis/AffineAnalysis.h"
#include "mlir/Analysis/LoopAnalysis.h"
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Analysis/Utils.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/IR/AffineValueMap.h"
#include "mlir/Dialect/SCF/SCF.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/BlockAndValueMapping.h"
#include "mlir/IR/Function.h"
#include "mlir/IR/IntegerSet.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Support/MathExtras.h"
#include "mlir/Transforms/RegionUtils.h"
#include "mlir/Transforms/Utils.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"
#define DEBUG_TYPE "LoopUtils"
using namespace mlir;
using llvm::SetVector;
using llvm::SmallMapVector;
namespace {
// This structure is to pass and return sets of loop parameters without
// confusing the order.
struct LoopParams {
Value lowerBound;
Value upperBound;
Value step;
};
} // namespace
/// Computes the cleanup loop lower bound of the loop being unrolled with
/// the specified unroll factor; this bound will also be upper bound of the main
/// part of the unrolled loop. Computes the bound as an AffineMap with its
/// operands or a null map when the trip count can't be expressed as an affine
/// expression.
static void getCleanupLoopLowerBound(AffineForOp forOp, unsigned unrollFactor,
AffineMap &map,
SmallVectorImpl<Value> &operands) {
auto lbMap = forOp.getLowerBoundMap();
// Single result lower bound map only.
if (lbMap.getNumResults() != 1) {
map = AffineMap();
return;
}
AffineMap tripCountMap;
SmallVector<Value, 4> tripCountOperands;
buildTripCountMapAndOperands(forOp, &tripCountMap, &tripCountOperands);
// Sometimes the trip count cannot be expressed as an affine expression.
if (!tripCountMap) {
map = AffineMap();
return;
}
OpBuilder b(forOp);
auto lb = b.create<AffineApplyOp>(forOp.getLoc(), lbMap,
forOp.getLowerBoundOperands());
// For each upper bound expr, get the range.
// Eg: affine.for %i = lb to min (ub1, ub2),
// where tripCountExprs yield (tr1, tr2), we create affine.apply's:
// lb + tr1 - tr1 % ufactor, lb + tr2 - tr2 % ufactor; the results of all
// these affine.apply's make up the cleanup loop lower bound.
SmallVector<AffineExpr, 4> bumpExprs(tripCountMap.getNumResults());
SmallVector<Value, 4> bumpValues(tripCountMap.getNumResults());
int64_t step = forOp.getStep();
for (unsigned i = 0, e = tripCountMap.getNumResults(); i < e; i++) {
auto tripCountExpr = tripCountMap.getResult(i);
bumpExprs[i] = (tripCountExpr - tripCountExpr % unrollFactor) * step;
auto bumpMap = AffineMap::get(tripCountMap.getNumDims(),
tripCountMap.getNumSymbols(), bumpExprs[i]);
bumpValues[i] =
b.create<AffineApplyOp>(forOp.getLoc(), bumpMap, tripCountOperands);
}
SmallVector<AffineExpr, 4> newUbExprs(tripCountMap.getNumResults());
for (unsigned i = 0, e = bumpExprs.size(); i < e; i++)
newUbExprs[i] = b.getAffineDimExpr(0) + b.getAffineDimExpr(i + 1);
operands.clear();
operands.push_back(lb);
operands.append(bumpValues.begin(), bumpValues.end());
map = AffineMap::get(1 + tripCountMap.getNumResults(), 0, newUbExprs,
b.getContext());
// Simplify the map + operands.
fullyComposeAffineMapAndOperands(&map, &operands);
map = simplifyAffineMap(map);
canonicalizeMapAndOperands(&map, &operands);
// Remove any affine.apply's that became dead from the simplification above.
for (auto v : bumpValues)
if (v.use_empty())
v.getDefiningOp()->erase();
if (lb.use_empty())
lb.erase();
}
// Build the IR that performs ceil division of a positive value by a constant:
// ceildiv(a, B) = divis(a + (B-1), B)
// where divis is rounding-to-zero division.
static Value ceilDivPositive(OpBuilder &builder, Location loc, Value dividend,
int64_t divisor) {
assert(divisor > 0 && "expected positive divisor");
assert(dividend.getType().isIndex() && "expected index-typed value");
Value divisorMinusOneCst = builder.create<ConstantIndexOp>(loc, divisor - 1);
Value divisorCst = builder.create<ConstantIndexOp>(loc, divisor);
Value sum = builder.create<AddIOp>(loc, dividend, divisorMinusOneCst);
return builder.create<SignedDivIOp>(loc, sum, divisorCst);
}
// Build the IR that performs ceil division of a positive value by another
// positive value:
// ceildiv(a, b) = divis(a + (b - 1), b)
// where divis is rounding-to-zero division.
static Value ceilDivPositive(OpBuilder &builder, Location loc, Value dividend,
Value divisor) {
assert(dividend.getType().isIndex() && "expected index-typed value");
Value cstOne = builder.create<ConstantIndexOp>(loc, 1);
Value divisorMinusOne = builder.create<SubIOp>(loc, divisor, cstOne);
Value sum = builder.create<AddIOp>(loc, dividend, divisorMinusOne);
return builder.create<SignedDivIOp>(loc, sum, divisor);
}
/// Promotes the loop body of a forOp to its containing block if the forOp
/// was known to have a single iteration.
// TODO: extend this for arbitrary affine bounds.
LogicalResult mlir::promoteIfSingleIteration(AffineForOp forOp) {
Optional<uint64_t> tripCount = getConstantTripCount(forOp);
if (!tripCount || tripCount.getValue() != 1)
return failure();
if (forOp.getLowerBoundMap().getNumResults() != 1)
return failure();
// Replaces all IV uses to its single iteration value.
auto iv = forOp.getInductionVar();
auto *parentBlock = forOp.getOperation()->getBlock();
if (!iv.use_empty()) {
if (forOp.hasConstantLowerBound()) {
OpBuilder topBuilder(forOp.getParentOfType<FuncOp>().getBody());
auto constOp = topBuilder.create<ConstantIndexOp>(
forOp.getLoc(), forOp.getConstantLowerBound());
iv.replaceAllUsesWith(constOp);
} else {
auto lbOperands = forOp.getLowerBoundOperands();
auto lbMap = forOp.getLowerBoundMap();
OpBuilder builder(parentBlock, Block::iterator(forOp));
if (lbMap == builder.getDimIdentityMap()) {
// No need of generating an affine.apply.
iv.replaceAllUsesWith(lbOperands[0]);
} else {
auto affineApplyOp =
builder.create<AffineApplyOp>(forOp.getLoc(), lbMap, lbOperands);
iv.replaceAllUsesWith(affineApplyOp);
}
}
}
// Move the loop body operations, except for its terminator, to the loop's
// containing block.
forOp.getBody()->back().erase();
parentBlock->getOperations().splice(Block::iterator(forOp),
forOp.getBody()->getOperations());
forOp.erase();
return success();
}
/// Promotes the loop body of a forOp to its containing block if the forOp
/// it can be determined that the loop has a single iteration.
LogicalResult mlir::promoteIfSingleIteration(scf::ForOp forOp) {
auto lbCstOp = forOp.lowerBound().getDefiningOp<ConstantIndexOp>();
auto ubCstOp = forOp.upperBound().getDefiningOp<ConstantIndexOp>();
auto stepCstOp = forOp.step().getDefiningOp<ConstantIndexOp>();
if (!lbCstOp || !ubCstOp || !stepCstOp || lbCstOp.getValue() < 0 ||
ubCstOp.getValue() < 0 || stepCstOp.getValue() < 0)
return failure();
int64_t tripCount = mlir::ceilDiv(ubCstOp.getValue() - lbCstOp.getValue(),
stepCstOp.getValue());
if (tripCount != 1)
return failure();
auto iv = forOp.getInductionVar();
iv.replaceAllUsesWith(lbCstOp);
// Move the loop body operations, except for its terminator, to the loop's
// containing block.
auto *parentBlock = forOp.getOperation()->getBlock();
forOp.getBody()->back().erase();
parentBlock->getOperations().splice(Block::iterator(forOp),
forOp.getBody()->getOperations());
forOp.erase();
return success();
}
/// Promotes all single iteration 'for' ops in `f`, i.e., moves
/// their body into the containing Block.
void mlir::promoteSingleIterationLoops(FuncOp f) {
// Gathers all innermost loops through a post order pruned walk.
f.walk([](Operation *op) {
if (auto forOp = dyn_cast<AffineForOp>(op))
promoteIfSingleIteration(forOp);
else if (auto forOp = dyn_cast<scf::ForOp>(op))
promoteIfSingleIteration(forOp);
});
}
/// Generates an affine.for op with the specified lower and upper bounds
/// while generating the right IV remappings to realize shifts for operations in
/// its body. The operations that go into the loop body are specified in
/// opGroupQueue starting from the specified offset, and in that order. The
/// first element of the pair specifies the shift applied to that group of
/// operations; the shift is multiplied by the loop step before being applied.
/// Returns nullptr if the generated loop simplifies to a single iteration one.
static AffineForOp generateShiftedLoop(
AffineMap lbMap, AffineMap ubMap,
const std::vector<std::pair<uint64_t, ArrayRef<Operation *>>> &opGroupQueue,
unsigned offset, AffineForOp srcForOp, OpBuilder b) {
auto lbOperands = srcForOp.getLowerBoundOperands();
auto ubOperands = srcForOp.getUpperBoundOperands();
assert(lbMap.getNumInputs() == lbOperands.size());
assert(ubMap.getNumInputs() == ubOperands.size());
auto loopChunk = b.create<AffineForOp>(srcForOp.getLoc(), lbOperands, lbMap,
ubOperands, ubMap, srcForOp.getStep());
auto loopChunkIV = loopChunk.getInductionVar();
auto srcIV = srcForOp.getInductionVar();
BlockAndValueMapping operandMap;
auto bodyBuilder = OpBuilder::atBlockTerminator(loopChunk.getBody());
for (auto it = opGroupQueue.begin() + offset, e = opGroupQueue.end(); it != e;
++it) {
uint64_t shift = it->first;
auto ops = it->second;
// All 'same shift' operations get added with their operands being
// remapped to results of cloned operations, and their IV used remapped.
// Generate the remapping if the shift is not zero: remappedIV = newIV -
// shift.
if (!srcIV.use_empty() && shift != 0) {
auto ivRemap = bodyBuilder.create<AffineApplyOp>(
srcForOp.getLoc(),
bodyBuilder.getSingleDimShiftAffineMap(
-static_cast<int64_t>(srcForOp.getStep() * shift)),
loopChunkIV);
operandMap.map(srcIV, ivRemap);
} else {
operandMap.map(srcIV, loopChunkIV);
}
for (auto *op : ops)
bodyBuilder.clone(*op, operandMap);
};
if (succeeded(promoteIfSingleIteration(loopChunk)))
return AffineForOp();
return loopChunk;
}
// The skewing of operations with respect to one another can be used for
// example to allow overlap of asynchronous operations (such as DMA
// communication) with computation, or just relative shifting of operations
// for better register reuse, locality or parallelism. As such, the shifts are
// typically expected to be at most of the order of the number of operations.
// This method should not be used as a substitute for loop distribution/fission.
// This method uses an algorithm// in time linear in the number of operations
// in the body of the for loop - (using the 'sweep line' paradigm). This method
// asserts preservation of SSA dominance. A check for that as well as that for
// memory-based dependence preservation check rests with the users of this
// method.
LogicalResult mlir::affineForOpBodySkew(AffineForOp forOp,
ArrayRef<uint64_t> shifts,
bool unrollPrologueEpilogue) {
assert(forOp.getBody()->getOperations().size() == shifts.size() &&
"too few/many shifts");
if (forOp.getBody()->begin() == std::prev(forOp.getBody()->end()))
return success();
// If the trip counts aren't constant, we would need versioning and
// conditional guards (or context information to prevent such versioning). The
// better way to pipeline for such loops is to first tile them and extract
// constant trip count "full tiles" before applying this.
auto mayBeConstTripCount = getConstantTripCount(forOp);
if (!mayBeConstTripCount.hasValue()) {
LLVM_DEBUG(forOp.emitRemark("non-constant trip count loop not handled"));
return success();
}
uint64_t tripCount = mayBeConstTripCount.getValue();
assert(isOpwiseShiftValid(forOp, shifts) &&
"shifts will lead to an invalid transformation\n");
int64_t step = forOp.getStep();
unsigned numChildOps = shifts.size();
// Do a linear time (counting) sort for the shifts.
uint64_t maxShift = *std::max_element(shifts.begin(), shifts.end());
if (maxShift >= numChildOps) {
// Large shifts are not the typical use case.
forOp.emitWarning("not shifting because shifts are unrealistically large");
return success();
}
// An array of operation groups sorted by shift amount; each group has all
// operations with the same shift in the order in which they appear in the
// body of the 'affine.for' op.
std::vector<std::vector<Operation *>> sortedOpGroups(maxShift + 1);
unsigned pos = 0;
for (auto &op : forOp.getBody()->without_terminator()) {
auto shift = shifts[pos++];
sortedOpGroups[shift].push_back(&op);
}
// Unless the shifts have a specific pattern (which actually would be the
// common use case), prologue and epilogue are not meaningfully defined.
// Nevertheless, if 'unrollPrologueEpilogue' is set, we will treat the first
// loop generated as the prologue and the last as epilogue and unroll these
// fully.
AffineForOp prologue, epilogue;
// Do a sweep over the sorted shifts while storing open groups in a
// vector, and generating loop portions as necessary during the sweep. A block
// of operations is paired with its shift.
std::vector<std::pair<uint64_t, ArrayRef<Operation *>>> opGroupQueue;
auto origLbMap = forOp.getLowerBoundMap();
uint64_t lbShift = 0;
OpBuilder b(forOp);
for (uint64_t d = 0, e = sortedOpGroups.size(); d < e; ++d) {
// If nothing is shifted by d, continue.
if (sortedOpGroups[d].empty())
continue;
if (!opGroupQueue.empty()) {
assert(d > 0 &&
"Queue expected to be empty when the first block is found");
// The interval for which the loop needs to be generated here is:
// [lbShift, min(lbShift + tripCount, d)) and the body of the
// loop needs to have all operations in opQueue in that order.
AffineForOp res;
if (lbShift + tripCount * step < d * step) {
res = generateShiftedLoop(
b.getShiftedAffineMap(origLbMap, lbShift),
b.getShiftedAffineMap(origLbMap, lbShift + tripCount * step),
opGroupQueue, /*offset=*/0, forOp, b);
// Entire loop for the queued op groups generated, empty it.
opGroupQueue.clear();
lbShift += tripCount * step;
} else {
res = generateShiftedLoop(b.getShiftedAffineMap(origLbMap, lbShift),
b.getShiftedAffineMap(origLbMap, d),
opGroupQueue, /*offset=*/0, forOp, b);
lbShift = d * step;
}
if (res) {
// Simplify/canonicalize the affine.for.
OwningRewritePatternList patterns;
AffineForOp::getCanonicalizationPatterns(patterns, res.getContext());
bool erased;
applyOpPatternsAndFold(res, std::move(patterns), &erased);
if (!erased && !prologue)
prologue = res;
if (!erased)
epilogue = res;
}
} else {
// Start of first interval.
lbShift = d * step;
}
// Augment the list of operations that get into the current open interval.
opGroupQueue.push_back({d, sortedOpGroups[d]});
}
// Those operations groups left in the queue now need to be processed (FIFO)
// and their loops completed.
for (unsigned i = 0, e = opGroupQueue.size(); i < e; ++i) {
uint64_t ubShift = (opGroupQueue[i].first + tripCount) * step;
epilogue = generateShiftedLoop(b.getShiftedAffineMap(origLbMap, lbShift),
b.getShiftedAffineMap(origLbMap, ubShift),
opGroupQueue, /*offset=*/i, forOp, b);
lbShift = ubShift;
if (!prologue)
prologue = epilogue;
}
// Erase the original for op.
forOp.erase();
if (unrollPrologueEpilogue && prologue)
loopUnrollFull(prologue);
if (unrollPrologueEpilogue && !epilogue && epilogue != prologue)
loopUnrollFull(epilogue);
return success();
}
/// Checks the legality of tiling of a hyper-rectangular loop nest by simply
/// checking if there is a 'negative' dependence in the memrefs present in
/// the loop nest. If yes then tiling is invalid.
static bool
checkTilingLegalityImpl(MutableArrayRef<mlir::AffineForOp> origLoops) {
assert(!origLoops.empty() && "no original loops provided");
// We first find out all dependences we intend to check.
SmallVector<Operation *, 8> loadAndStoreOps;
origLoops[0].getOperation()->walk([&](Operation *op) {
if (isa<AffineReadOpInterface, AffineWriteOpInterface>(op))
loadAndStoreOps.push_back(op);
});
unsigned numOps = loadAndStoreOps.size();
unsigned numLoops = origLoops.size();
FlatAffineConstraints dependenceConstraints;
for (unsigned d = 1; d <= numLoops + 1; ++d) {
for (unsigned i = 0; i < numOps; ++i) {
Operation *srcOp = loadAndStoreOps[i];
MemRefAccess srcAccess(srcOp);
for (unsigned j = 0; j < numOps; ++j) {
Operation *dstOp = loadAndStoreOps[j];
MemRefAccess dstAccess(dstOp);
SmallVector<DependenceComponent, 2> depComps;
dependenceConstraints.reset();
DependenceResult result = checkMemrefAccessDependence(
srcAccess, dstAccess, d, &dependenceConstraints, &depComps);
// Skip if there is no dependence in this case.
if (!hasDependence(result))
continue;
// Check whether there is any negative direction vector in the
// dependence components found above, which means that dependence is
// violated by the default hyper-rect tiling method.
LLVM_DEBUG(llvm::dbgs() << "Checking whether tiling legality violated "
"for dependence at depth: "
<< Twine(d) << " between:\n";);
LLVM_DEBUG(srcAccess.opInst->dump(););
LLVM_DEBUG(dstAccess.opInst->dump(););
for (unsigned k = 0, e = depComps.size(); k < e; k++) {
DependenceComponent depComp = depComps[k];
if (depComp.lb.hasValue() && depComp.ub.hasValue() &&
depComp.lb.getValue() < depComp.ub.getValue() &&
depComp.ub.getValue() < 0) {
LLVM_DEBUG(llvm::dbgs()
<< "Dependence component lb = "
<< Twine(depComp.lb.getValue())
<< " ub = " << Twine(depComp.ub.getValue())
<< " is negative at depth: " << Twine(d)
<< " and thus violates the legality rule.\n");
return false;
}
}
}
}
}
return true;
}
/// Checks whether hyper-rectangular loop tiling of the nest
/// represented by `origLoops` is valid. The validity condition is from Irigoin
/// and Triolet, which states that two tiles cannot depend on each other. We
/// simplify such condition to just checking whether there is any negative
/// dependence direction, since we have the prior knowledge that the tiling
/// results will be hyper-rectangles, which are scheduled in the
/// lexicographically increasing order on the vector of loop indices. This
/// function will return failure when any dependence component is negative along
/// any of `origLoops`.
LogicalResult
checkTilingLegality(MutableArrayRef<mlir::AffineForOp> origLoops) {
return success(checkTilingLegalityImpl(origLoops));
}
/// Check if the input data is valid and wheter tiled code will be legal or not.
template <typename t>
void performPreTilingChecks(MutableArrayRef<AffineForOp> input,
ArrayRef<t> tileSizes) {
// Check if the supplied for op's are all successively nested.
assert(!input.empty() && "no loops in input band");
assert(input.size() == tileSizes.size() && "Too few/many tile sizes");
assert(isPerfectlyNested(input) && "input loops not perfectly nested");
// Perform tiling legality test.
if (failed(checkTilingLegality(input)))
input[0].emitRemark("tiled code is illegal due to dependences");
}
/// Move the loop body of AffineForOp 'src' from 'src' into the specified
/// location in destination's body, ignoring the terminator.
static void moveLoopBodyImpl(AffineForOp src, AffineForOp dest,
Block::iterator loc) {
auto &ops = src.getBody()->getOperations();
dest.getBody()->getOperations().splice(loc, ops, ops.begin(),
std::prev(ops.end()));
}
/// Move the loop body of AffineForOp 'src' from 'src' to the start of dest
/// body.
void moveLoopBody(AffineForOp src, AffineForOp dest) {
moveLoopBodyImpl(src, dest, dest.getBody()->begin());
}
/// Constructs tiled loop nest, without setting the loop bounds and move the
/// body of the original loop nest to the tiled loop nest.
void constructTiledLoopNest(MutableArrayRef<AffineForOp> origLoops,
AffineForOp rootAffineForOp, unsigned width,
MutableArrayRef<AffineForOp> tiledLoops) {
Location loc = rootAffineForOp.getLoc();
// The outermost among the loops as we add more..
Operation *topLoop = rootAffineForOp.getOperation();
AffineForOp innermostPointLoop;
// Add intra-tile (or point) loops.
for (unsigned i = 0; i < width; i++) {
OpBuilder b(topLoop);
// Loop bounds will be set later.
AffineForOp pointLoop = b.create<AffineForOp>(loc, 0, 0);
pointLoop.getBody()->getOperations().splice(
pointLoop.getBody()->begin(), topLoop->getBlock()->getOperations(),
topLoop);
tiledLoops[2 * width - 1 - i] = pointLoop;
topLoop = pointLoop.getOperation();
if (i == 0)
innermostPointLoop = pointLoop;
}
// Add tile space loops;
for (unsigned i = width; i < 2 * width; i++) {
OpBuilder b(topLoop);
// Loop bounds will be set later.
AffineForOp tileSpaceLoop = b.create<AffineForOp>(loc, 0, 0);
tileSpaceLoop.getBody()->getOperations().splice(
tileSpaceLoop.getBody()->begin(), topLoop->getBlock()->getOperations(),
topLoop);
tiledLoops[2 * width - i - 1] = tileSpaceLoop;
topLoop = tileSpaceLoop.getOperation();
}
// Move the loop body of the original nest to the new one.
moveLoopBody(origLoops.back(), innermostPointLoop);
}
/// Checks whether a loop nest is hyper-rectangular or not.
LogicalResult checkIfHyperRectangular(MutableArrayRef<AffineForOp> input,
AffineForOp rootAffineForOp,
unsigned width) {
FlatAffineConstraints cst;
SmallVector<Operation *, 8> ops(input.begin(), input.end());
getIndexSet(ops, &cst);
if (!cst.isHyperRectangular(0, width)) {
rootAffineForOp.emitError("tiled code generation unimplemented for the "
"non-hyperrectangular case");
return failure();
}
return success();
}
/// Set lower and upper bounds of intra-tile loops for parametric tiling.
// TODO: Handle non-constant lower bounds.
static void setIntraTileBoundsParametric(OpBuilder &b, AffineForOp origLoop,
AffineForOp newInterTileLoop,
AffineForOp newIntraTileLoop,
Value tileSize) {
// The lower bound for the intra-tile loop is represented by an affine map
// as (%i, %t0)->((%i - %origlb) * %t0 + %origlb). Similarly, the upper bound
// for the intra-tile loop is represented by an affine map as (%i, %t0)->((%i
// - %origlb) * %t0) + (%t0 * %origLoopStep) + %origlb), where %i is loop IV
// of the corresponding inter-tile loop, %t0 is the corresponding tiling
// parameter, %origlb is lower bound and %origLoopStep is the loop step of the
// corresponding inter-tile loop.
assert(origLoop.hasConstantLowerBound() &&
"expected input loops to have constant lower bound.");
// Get lower bound of original loop as an affine expression.
AffineExpr origLowerBoundExpr;
origLowerBoundExpr =
b.getAffineConstantExpr(origLoop.getConstantLowerBound());
// Add dim operands from original lower/upper bound.
SmallVector<Value, 4> lbOperands, ubOperands;
AffineBound lb = origLoop.getLowerBound();
AffineBound ub = origLoop.getUpperBound();
lbOperands.reserve(lb.getNumOperands() + 2);
ubOperands.reserve(ub.getNumOperands() + 2);
AffineMap origLbMap = lb.getMap();
AffineMap origUbMap = ub.getMap();
for (unsigned j = 0, e = origLbMap.getNumDims(); j < e; ++j)
lbOperands.push_back(lb.getOperand(j));
for (unsigned j = 0, e = origUbMap.getNumDims(); j < e; ++j)
ubOperands.push_back(ub.getOperand(j));
// Add a new dim operand in lb/ubOperands corresponding to the origLoop
// IV.
lbOperands.push_back(newInterTileLoop.getInductionVar());
ubOperands.push_back(newInterTileLoop.getInductionVar());
// Get loop IV as an affine expression for lower/upper bound. Size of
// lb/ubOperands is guaranteed to be atleast one.
AffineExpr lbLoopIvExpr = b.getAffineDimExpr(lbOperands.size() - 1);
AffineExpr ubLoopIvExpr = b.getAffineDimExpr(ubOperands.size() - 1);
// Add symbol operands from original lower/upper bound.
for (unsigned j = 0, e = origLbMap.getNumSymbols(); j < e; ++j)
lbOperands.push_back(lb.getOperand(origLbMap.getNumDims() + j));
for (unsigned j = 0, e = origUbMap.getNumSymbols(); j < e; ++j)
ubOperands.push_back(ub.getOperand(origUbMap.getNumDims() + j));
// Add a new symbol operand which is the tile size for this loop.
lbOperands.push_back(tileSize);
ubOperands.push_back(tileSize);
SmallVector<AffineExpr, 4> lbBoundExprs;
SmallVector<AffineExpr, 4> ubBoundExprs;
lbBoundExprs.reserve(origLbMap.getNumResults());
ubBoundExprs.reserve(origUbMap.getNumResults());
// Get tiling parameter as an affine expression for lb/ub.
AffineExpr lbTileParameter = b.getAffineSymbolExpr(origLbMap.getNumSymbols());
AffineExpr ubTileParameter = b.getAffineSymbolExpr(origUbMap.getNumSymbols());
// Insert lb as inter-tile ((loop IV - origlb) * tilingParameter) + origlb.
lbBoundExprs.push_back(
((lbLoopIvExpr - origLowerBoundExpr) * lbTileParameter) +
origLowerBoundExpr);
// Get the origLoopStep as an affine expression.
AffineExpr origLoopStep = b.getAffineConstantExpr(origLoop.getStep());
// Insert ub as inter-tile ((loop IV - origlb) * tilingParameter) +
// (tilingParameter * origLoopStep) + origlb.
ubBoundExprs.push_back(
((ubLoopIvExpr - origLowerBoundExpr) * ubTileParameter) +
(ubTileParameter * origLoopStep) + origLowerBoundExpr);
ubBoundExprs.append(origUbMap.getResults().begin(),
origUbMap.getResults().end());
AffineMap lbMap =
AffineMap::get(origLbMap.getNumDims() + 1, origLbMap.getNumSymbols() + 1,
lbBoundExprs, b.getContext());
newIntraTileLoop.setLowerBound(lbOperands, lbMap);
AffineMap ubMap =
AffineMap::get(origUbMap.getNumDims() + 1, origUbMap.getNumSymbols() + 1,
ubBoundExprs, b.getContext());
newIntraTileLoop.setUpperBound(ubOperands, ubMap);
// Original loop step must be preserved.
newIntraTileLoop.setStep(origLoop.getStep());
}
/// Set lower and upper bounds of inter-tile loops for parametric tiling.
// TODO: Handle non-constant lower bounds.
static void setInterTileBoundsParametric(OpBuilder &b, AffineForOp origLoop,
AffineForOp newLoop, Value tileSize) {
OperandRange newLbOperands = origLoop.getLowerBoundOperands();
// The lower bounds for inter-tile loops are same as the correspondig lower
// bounds of original loops.
newLoop.setLowerBound(newLbOperands, origLoop.getLowerBoundMap());
// The new upper bound map for inter-tile loops, assuming constant lower
// bounds, are now originalLowerBound + ceildiv((orignalUpperBound -
// originalLowerBound), tiling paramter); where tiling parameter is the
// respective tile size for that loop. For e.g. if the original ubmap was
// ()->(1024), the new map will be
// ()[s0]->(ceildiv((1024 -lb) % s0)), where s0 is the tiling parameter.
// Therefore a new symbol operand is inserted in the map and the result
// expression is overwritten.
assert(origLoop.hasConstantLowerBound() &&
"expected input loops to have constant lower bound.");
// Get lower bound of original loop as an affine expression.
AffineExpr origLowerBoundExpr;
origLowerBoundExpr =
b.getAffineConstantExpr(origLoop.getConstantLowerBound());
// Add dim operands from original upper bound.
SmallVector<Value, 4> ubOperands;
AffineBound ub = origLoop.getUpperBound();
ubOperands.reserve(ub.getNumOperands() + 1);
AffineMap origUbMap = ub.getMap();
for (unsigned j = 0, e = origUbMap.getNumDims(); j < e; ++j)
ubOperands.push_back(ub.getOperand(j));
// Add symbol operands from original upper bound.
for (unsigned j = 0, e = origUbMap.getNumSymbols(); j < e; ++j)
ubOperands.push_back(ub.getOperand(origUbMap.getNumDims() + j));
// Add a new symbol operand which is the tile size for this loop.
ubOperands.push_back(tileSize);
// Get tiling parameter as an affine expression.
AffineExpr tileParameter = b.getAffineSymbolExpr(origUbMap.getNumSymbols());
SmallVector<AffineExpr, 4> boundExprs;
boundExprs.reserve(origUbMap.getNumResults());
int64_t origUpperBound;
AffineExpr origUpperBoundExpr;
// If upper bound for the original loop is constant, then the constant can
// be obtained as an affine expression straight away.
if (origLoop.hasConstantUpperBound()) {
origUpperBound = origLoop.getConstantUpperBound();
// Get original constant upper bound as an affine expression.
origUpperBoundExpr = b.getAffineConstantExpr(origUpperBound);
// Insert the bound as originalLowerBoundceildiv((originalUpperBound -
// originalLowerBound), tilingParameter).
boundExprs.push_back(
origLowerBoundExpr +
(origUpperBoundExpr - origLowerBoundExpr).ceilDiv(tileParameter));
} else {
// If upper bound for the original loop is not constant then two cases
// are possible, although there handeling is the same, 1.) The result of
// ubmap has only one result expression. For e.g.
// affine.for %i = 5 to %ub
//
// A symbol operand is added which represents the tiling paramater. The
// new loop bounds here will be like ()[s0, s1] -> ((s0 - 5) ceildiv s1 + 5)
// where 's0' is the original upper bound and 's1' is the tiling
// parameter. 2.) When ubMap has more than one result expression. For e.g.
// #map0 = affine_map<()[s0, s1] -> (s0, s1)
// affine.for %i = 5 to min #map0()[%s0, %s1]
//
// A symbol operand is added which represents the tiling parameter. The
// new loop bounds will be like ()[s0, s1, s2] -> ((s0 - 5) ceildiv s2 + 5,
// (s1 -5) ceildiv s2 + 5), where s2 is the tiling parameter.
// Insert the bounds as originalLowerBound + ceildiv((originalUpperBound -
// originalLowerBound), tilingParameter).
for (AffineExpr origUpperBoundExpr : origUbMap.getResults())
boundExprs.push_back(
origLowerBoundExpr +
(origUpperBoundExpr - origLowerBoundExpr).ceilDiv(tileParameter));
}
AffineMap ubMap =
AffineMap::get(origUbMap.getNumDims(), origUbMap.getNumSymbols() + 1,
boundExprs, b.getContext());
newLoop.setUpperBound(ubOperands, ubMap);
// Original loop step must be preserved.
newLoop.setStep(origLoop.getStep());
}
/// Constructs and sets new loop bounds after tiling for the case of
/// hyper-rectangular index sets, where the bounds of one dimension do not
/// depend on other dimensions and tiling parameters are captured from SSA
/// values. Bounds of each dimension can thus be treated independently,
/// and deriving the new bounds is much simpler and faster than for the case of
/// tiling arbitrary polyhedral shapes.
static void constructParametricallyTiledIndexSetHyperRect(
MutableArrayRef<AffineForOp> origLoops,
MutableArrayRef<AffineForOp> newLoops, ArrayRef<Value> tileSizes) {
assert(!origLoops.empty() && "expected atleast one loop in band");
assert(origLoops.size() == tileSizes.size() &&
"expected tiling parameter for each loop in band.");
OpBuilder b(origLoops[0].getOperation());
unsigned width = origLoops.size();
// Set bounds for tile space loops.
for (unsigned i = 0; i < width; ++i) {
setInterTileBoundsParametric(b, origLoops[i], newLoops[i], tileSizes[i]);
}
// Set bounds for intra-tile loops.
for (unsigned i = 0; i < width; ++i) {
setIntraTileBoundsParametric(b, origLoops[i], newLoops[i],
newLoops[i + width], tileSizes[i]);
}
}
/// Constructs and sets new loop bounds after tiling for the case of
/// hyper-rectangular index sets, where the bounds of one dimension do not
/// depend on other dimensions. Bounds of each dimension can thus be treated
/// independently, and deriving the new bounds is much simpler and faster
/// than for the case of tiling arbitrary polyhedral shapes.
static void
constructTiledIndexSetHyperRect(MutableArrayRef<AffineForOp> origLoops,
MutableArrayRef<AffineForOp> newLoops,
ArrayRef<unsigned> tileSizes) {
assert(!origLoops.empty());
assert(origLoops.size() == tileSizes.size());
OpBuilder b(origLoops[0].getOperation());
unsigned width = origLoops.size();
// Bounds for tile space loops.
for (unsigned i = 0; i < width; i++) {
OperandRange newLbOperands = origLoops[i].getLowerBoundOperands();
OperandRange newUbOperands = origLoops[i].getUpperBoundOperands();
newLoops[i].setLowerBound(newLbOperands, origLoops[i].getLowerBoundMap());
newLoops[i].setUpperBound(newUbOperands, origLoops[i].getUpperBoundMap());
newLoops[i].setStep(tileSizes[i]);
}
// Bounds for intra-tile loops.
for (unsigned i = 0; i < width; i++) {
int64_t largestDiv = getLargestDivisorOfTripCount(origLoops[i]);
Optional<uint64_t> mayBeConstantCount = getConstantTripCount(origLoops[i]);
// The lower bound is just the tile-space loop.
AffineMap lbMap = b.getDimIdentityMap();
newLoops[width + i].setLowerBound(
/*operands=*/newLoops[i].getInductionVar(), lbMap);
// Set the upper bound.
if (mayBeConstantCount && mayBeConstantCount.getValue() < tileSizes[i]) {
// Trip count is less than the tile size: upper bound is lower bound +
// trip count.
AffineMap ubMap =
b.getSingleDimShiftAffineMap(mayBeConstantCount.getValue());
newLoops[width + i].setUpperBound(
/*operands=*/newLoops[i].getInductionVar(), ubMap);
} else if (largestDiv % tileSizes[i] != 0) {
// Intra-tile loop ii goes from i to min(i + tileSize, ub_i).
// Construct the upper bound map; the operands are the original operands
// with 'i' (tile-space loop) appended to it. The new upper bound map is
// the original one with an additional expression i + tileSize appended.
// Add dim operands from original upper bound.
SmallVector<Value, 4> ubOperands;
AffineBound ub = origLoops[i].getUpperBound();
ubOperands.reserve(ub.getNumOperands() + 1);
AffineMap origUbMap = ub.getMap();
for (unsigned j = 0, e = origUbMap.getNumDims(); j < e; ++j)
ubOperands.push_back(ub.getOperand(j));
// Add dim operand for new loop upper bound.
ubOperands.push_back(newLoops[i].getInductionVar());
// Add symbol operands from original upper bound.
for (unsigned j = 0, e = origUbMap.getNumSymbols(); j < e; ++j)
ubOperands.push_back(ub.getOperand(origUbMap.getNumDims() + j));
SmallVector<AffineExpr, 4> boundExprs;
boundExprs.reserve(1 + origUbMap.getNumResults());
AffineExpr dim = b.getAffineDimExpr(origUbMap.getNumDims());
// The new upper bound map is the original one with an additional
// expression i + tileSize appended.
boundExprs.push_back(dim + tileSizes[i]);
boundExprs.append(origUbMap.getResults().begin(),
origUbMap.getResults().end());
AffineMap ubMap =
AffineMap::get(origUbMap.getNumDims() + 1, origUbMap.getNumSymbols(),
boundExprs, b.getContext());
newLoops[width + i].setUpperBound(/*operands=*/ubOperands, ubMap);
} else {
// No need of the min expression.
AffineExpr dim = b.getAffineDimExpr(0);
AffineMap ubMap = AffineMap::get(1, 0, dim + tileSizes[i]);
newLoops[width + i].setUpperBound(newLoops[i].getInductionVar(), ubMap);
}
}
}
/// Tiles the specified band of perfectly nested loops creating tile-space loops
/// and intra-tile loops. A band is a contiguous set of loops.
// TODO: handle non hyper-rectangular spaces.
LogicalResult
mlir::tilePerfectlyNested(MutableArrayRef<AffineForOp> input,
ArrayRef<unsigned> tileSizes,
SmallVectorImpl<AffineForOp> *tiledNest) {
performPreTilingChecks(input, tileSizes);
MutableArrayRef<AffineForOp> origLoops = input;
AffineForOp rootAffineForOp = origLoops[0];
// Note that width is at least one since band isn't empty.
unsigned width = input.size();
SmallVector<AffineForOp, 6> tiledLoops(2 * width);
// Construct a tiled loop nest without setting their bounds. Bounds are
// set later.
constructTiledLoopNest(origLoops, rootAffineForOp, width, tiledLoops);
SmallVector<Value, 8> origLoopIVs;
extractForInductionVars(input, &origLoopIVs);
if (failed(checkIfHyperRectangular(input, rootAffineForOp, width)))
return failure();
// Set loop bounds for the tiled loop nest.
constructTiledIndexSetHyperRect(origLoops, tiledLoops, tileSizes);
// Replace original IVs with intra-tile loop IVs.
for (unsigned i = 0; i < width; i++)
origLoopIVs[i].replaceAllUsesWith(tiledLoops[i + width].getInductionVar());
// Erase the old loop nest.
rootAffineForOp.erase();
if (tiledNest)
*tiledNest = std::move(tiledLoops);
return success();
}
/// Tiles the specified band of perfectly nested loops creating tile-space
/// loops and intra-tile loops, using SSA values as tiling parameters. A band
/// is a contiguous set of loops.
// TODO: handle non hyper-rectangular spaces.
LogicalResult
mlir::tilePerfectlyNestedParametric(MutableArrayRef<AffineForOp> input,
ArrayRef<Value> tileSizes,
SmallVectorImpl<AffineForOp> *tiledNest) {
performPreTilingChecks(input, tileSizes);
MutableArrayRef<AffineForOp> origLoops = input;
AffineForOp rootAffineForOp = origLoops[0];
// Note that width is at least one since band isn't empty.
unsigned width = input.size();
SmallVector<AffineForOp, 6> tiledLoops(2 * width);
// Construct a tiled loop nest without setting their bounds. Bounds are
// set later.
constructTiledLoopNest(origLoops, rootAffineForOp, width, tiledLoops);
SmallVector<Value, 8> origLoopIVs;
extractForInductionVars(input, &origLoopIVs);
if (failed(checkIfHyperRectangular(input, rootAffineForOp, width)))
return failure();
// Set loop bounds for the tiled loop nest.
constructParametricallyTiledIndexSetHyperRect(origLoops, tiledLoops,
tileSizes);
// Replace original IVs with intra-tile loop IVs.
for (unsigned i = 0; i < width; i++)
origLoopIVs[i].replaceAllUsesWith(tiledLoops[i + width].getInductionVar());
// Erase the old loop nest.
rootAffineForOp.erase();
if (tiledNest)
*tiledNest = std::move(tiledLoops);
return success();
}
/// Collect perfectly nested loops starting from `rootForOps`. Loops are
/// perfectly nested if each loop is the first and only non-terminator operation
/// in the parent loop. Collect at most `maxLoops` loops and append them to
/// `forOps`.
template <typename T>
static void getPerfectlyNestedLoopsImpl(
SmallVectorImpl<T> &forOps, T rootForOp,
unsigned maxLoops = std::numeric_limits<unsigned>::max()) {
for (unsigned i = 0; i < maxLoops; ++i) {
forOps.push_back(rootForOp);
Block &body = rootForOp.region().front();
if (body.begin() != std::prev(body.end(), 2))
return;
rootForOp = dyn_cast<T>(&body.front());
if (!rootForOp)
return;
}
}
/// Get perfectly nested sequence of loops starting at root of loop nest
/// (the first op being another AffineFor, and the second op - a terminator).
/// A loop is perfectly nested iff: the first op in the loop's body is another
/// AffineForOp, and the second op is a terminator).
void mlir::getPerfectlyNestedLoops(SmallVectorImpl<AffineForOp> &nestedLoops,
AffineForOp root) {
getPerfectlyNestedLoopsImpl(nestedLoops, root);
}
void mlir::getPerfectlyNestedLoops(SmallVectorImpl<scf::ForOp> &nestedLoops,
scf::ForOp root) {