Add a basic model to set tile sizes + some cleanup
- compute tile sizes based on a simple model that looks at memory footprints
(instead of using the hardcoded default value)
- adjust tile sizes to make them factors of trip counts based on an option
- update loop fusion CL options to allow setting maximal fusion at pass creation
- change an emitError to emitWarning (since it's not a hard error unless the client
treats it that way, in which case, it can emit one)
$ mlir-opt -debug-only=loop-tile -loop-tile test/Transforms/loop-tiling.mlir
test/Transforms/loop-tiling.mlir:81:3: note: using tile sizes [4 4 5 ]
for %i = 0 to 256 {
for %i0 = 0 to 256 step 4 {
for %i1 = 0 to 256 step 4 {
for %i2 = 0 to 250 step 5 {
for %i3 = #map4(%i0) to #map11(%i0) {
for %i4 = #map4(%i1) to #map11(%i1) {
for %i5 = #map4(%i2) to #map12(%i2) {
%0 = load %arg0[%i3, %i5] : memref<8x8xvector<64xf32>>
%1 = load %arg1[%i5, %i4] : memref<8x8xvector<64xf32>>
%2 = load %arg2[%i3, %i4] : memref<8x8xvector<64xf32>>
%3 = mulf %0, %1 : vector<64xf32>
%4 = addf %2, %3 : vector<64xf32>
store %4, %arg2[%i3, %i4] : memref<8x8xvector<64xf32>>
}
}
}
}
}
}
PiperOrigin-RevId: 237461836
diff --git a/mlir/lib/Transforms/LoopFusion.cpp b/mlir/lib/Transforms/LoopFusion.cpp
index 3786177..4e619ef 100644
--- a/mlir/lib/Transforms/LoopFusion.cpp
+++ b/mlir/lib/Transforms/LoopFusion.cpp
@@ -48,7 +48,9 @@
static llvm::cl::OptionCategory clOptionsCategory(DEBUG_TYPE " options");
-/// Disables fusion profitability check and fuses if valid.
+/// Disables fusion profitability check and fuses if valid. Ignore any
+/// additional (redundant) computation tolerance threshold
+/// that would have prevented fusion.
static llvm::cl::opt<bool>
clMaximalLoopFusion("fusion-maximal",
llvm::cl::desc("Enables maximal loop fusion"),
@@ -66,8 +68,8 @@
llvm::cl::desc("Faster memory space number to promote fusion buffers to"),
llvm::cl::cat(clOptionsCategory));
-// A local buffer of size less than or equal to this size is promoted to fast
-// memory.
+// A local buffer of size less than or equal to this size is automatically
+// promoted to fast memory after producer-consumer fusion.
static llvm::cl::opt<unsigned long long> clFusionLocalBufThreshold(
"fusion-local-buf-threshold",
llvm::cl::desc("Threshold size (KiB) for promoting local buffers to fast "
@@ -86,9 +88,10 @@
// and add support for more general loop fusion algorithms.
struct LoopFusion : public FunctionPass<LoopFusion> {
- LoopFusion(unsigned fastMemorySpace = 0, uint64_t localBufSizeThreshold = 0)
+ LoopFusion(unsigned fastMemorySpace = 0, uint64_t localBufSizeThreshold = 0,
+ bool maximalFusion = false)
: localBufSizeThreshold(localBufSizeThreshold),
- fastMemorySpace(fastMemorySpace) {}
+ fastMemorySpace(fastMemorySpace), maximalFusion(maximalFusion) {}
void runOnFunction() override;
@@ -96,6 +99,9 @@
// `fastMemorySpace` if provided.
uint64_t localBufSizeThreshold;
Optional<unsigned> fastMemorySpace = None;
+ // If true, ignore any additional (redundant) computation tolerance threshold
+ // that would have prevented fusion.
+ bool maximalFusion;
// The amount of additional computation that is tolerated while fusing
// pair-wise as a fraction of the total computation.
@@ -105,8 +111,9 @@
} // end anonymous namespace
FunctionPassBase *mlir::createLoopFusionPass(unsigned fastMemorySpace,
- uint64_t localBufSizeThreshold) {
- return new LoopFusion(fastMemorySpace, localBufSizeThreshold);
+ uint64_t localBufSizeThreshold,
+ bool maximalFusion) {
+ return new LoopFusion(fastMemorySpace, localBufSizeThreshold, maximalFusion);
}
namespace {
@@ -1411,7 +1418,7 @@
ArrayRef<Instruction *> dstLoadOpInsts,
ArrayRef<Instruction *> dstStoreOpInsts,
ComputationSliceState *sliceState,
- unsigned *dstLoopDepth) {
+ unsigned *dstLoopDepth, bool maximalFusion) {
LLVM_DEBUG({
llvm::dbgs() << "Checking whether fusion is profitable between:\n";
llvm::dbgs() << " " << *srcOpInst << " and \n";
@@ -1620,7 +1627,7 @@
// (as per computeToleranceThreshold), we will simply pick the one that
// reduces the intermediary size the most.
if ((storageReduction > maxStorageReduction) &&
- (clMaximalLoopFusion ||
+ (maximalFusion ||
(additionalComputeFraction < computeToleranceThreshold))) {
maxStorageReduction = storageReduction;
bestDstLoopDepth = i;
@@ -1632,7 +1639,7 @@
// A simple cost model: fuse if it reduces the memory footprint. If
// -maximal-fusion is set, fuse nevertheless.
- if (!clMaximalLoopFusion && !bestDstLoopDepth.hasValue()) {
+ if (!maximalFusion && !bestDstLoopDepth.hasValue()) {
LLVM_DEBUG(
llvm::dbgs()
<< "All fusion choices involve more than the threshold amount of "
@@ -1661,7 +1668,7 @@
Optional<double> storageReduction = None;
- if (!clMaximalLoopFusion) {
+ if (!maximalFusion) {
if (!dstMemSize.hasValue() || !srcMemSize.hasValue()) {
LLVM_DEBUG(
llvm::dbgs()
@@ -1785,13 +1792,16 @@
unsigned localBufSizeThreshold;
// Parameter for fast memory space.
Optional<unsigned> fastMemorySpace;
+ // If true, ignore any additional (redundant) computation tolerance threshold
+ // that would have prevented fusion.
+ bool maximalFusion;
using Node = MemRefDependenceGraph::Node;
GreedyFusion(MemRefDependenceGraph *mdg, unsigned localBufSizeThreshold,
- Optional<unsigned> fastMemorySpace)
+ Optional<unsigned> fastMemorySpace, bool maximalFusion)
: mdg(mdg), localBufSizeThreshold(localBufSizeThreshold),
- fastMemorySpace(fastMemorySpace) {}
+ fastMemorySpace(fastMemorySpace), maximalFusion(maximalFusion) {}
// Initializes 'worklist' with nodes from 'mdg'
void init() {
@@ -1917,7 +1927,7 @@
// Check if fusion would be profitable.
if (!isFusionProfitable(srcStoreOpInst, srcStoreOpInst,
dstLoadOpInsts, dstStoreOpInsts, &sliceState,
- &bestDstLoopDepth))
+ &bestDstLoopDepth, maximalFusion))
continue;
// Fuse computation slice of 'srcLoopNest' into 'dstLoopNest'.
@@ -2076,7 +2086,8 @@
// Check if fusion would be profitable.
if (!isFusionProfitable(sibLoadOpInst, sibStoreOpInst, dstLoadOpInsts,
- dstStoreOpInsts, &sliceState, &bestDstLoopDepth))
+ dstStoreOpInsts, &sliceState, &bestDstLoopDepth,
+ maximalFusion))
continue;
// Fuse computation slice of 'sibLoopNest' into 'dstLoopNest'.
@@ -2231,9 +2242,13 @@
localBufSizeThreshold = clFusionLocalBufThreshold * 1024;
}
+ if (clMaximalLoopFusion.getNumOccurrences() > 0)
+ maximalFusion = clMaximalLoopFusion;
+
MemRefDependenceGraph g;
if (g.init(getFunction()))
- GreedyFusion(&g, localBufSizeThreshold, fastMemorySpace).run();
+ GreedyFusion(&g, localBufSizeThreshold, fastMemorySpace, maximalFusion)
+ .run();
}
static PassRegistration<LoopFusion> pass("loop-fusion", "Fuse loop nests");
