TestAffineDataCopy.cpp
4.42 KB
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
//===- TestAffineDataCopy.cpp - Test affine data copy utility -------------===//
//
// 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 a pass to test affine data copy utility functions and
// options.
//
//===----------------------------------------------------------------------===//
#include "mlir/Analysis/Utils.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/LoopUtils.h"
#include "mlir/Transforms/Passes.h"
#define PASS_NAME "test-affine-data-copy"
using namespace mlir;
static llvm::cl::OptionCategory clOptionsCategory(PASS_NAME " options");
namespace {
struct TestAffineDataCopy
: public PassWrapper<TestAffineDataCopy, FunctionPass> {
TestAffineDataCopy() = default;
TestAffineDataCopy(const TestAffineDataCopy &pass){};
void runOnFunction() override;
private:
Option<bool> clMemRefFilter{
*this, "memref-filter",
llvm::cl::desc(
"Enable memref filter testing in affine data copy optimization"),
llvm::cl::init(false)};
Option<bool> clTestGenerateCopyForMemRegion{
*this, "for-memref-region",
llvm::cl::desc("Test copy generation for a single memref region"),
llvm::cl::init(false)};
};
} // end anonymous namespace
void TestAffineDataCopy::runOnFunction() {
// Gather all AffineForOps by loop depth.
std::vector<SmallVector<AffineForOp, 2>> depthToLoops;
gatherLoops(getFunction(), depthToLoops);
assert(depthToLoops.size() && "Loop nest not found");
// Only support tests with a single loop nest and a single innermost loop
// for now.
unsigned innermostLoopIdx = depthToLoops.size() - 1;
if (depthToLoops[0].size() != 1 || depthToLoops[innermostLoopIdx].size() != 1)
return;
auto loopNest = depthToLoops[0][0];
auto innermostLoop = depthToLoops[innermostLoopIdx][0];
AffineLoadOp load;
if (clMemRefFilter || clTestGenerateCopyForMemRegion) {
// Gather MemRef filter. For simplicity, we use the first loaded memref
// found in the innermost loop.
for (auto &op : *innermostLoop.getBody()) {
if (auto ld = dyn_cast<AffineLoadOp>(op)) {
load = ld;
break;
}
}
}
AffineCopyOptions copyOptions = {/*generateDma=*/false,
/*slowMemorySpace=*/0,
/*fastMemorySpace=*/0,
/*tagMemorySpace=*/0,
/*fastMemCapacityBytes=*/32 * 1024 * 1024UL};
DenseSet<Operation *> copyNests;
if (clMemRefFilter) {
affineDataCopyGenerate(loopNest, copyOptions, load.getMemRef(), copyNests);
} else if (clTestGenerateCopyForMemRegion) {
CopyGenerateResult result;
MemRefRegion region(loopNest.getLoc());
region.compute(load, /*loopDepth=*/0);
generateCopyForMemRegion(region, loopNest, copyOptions, result);
}
// Promote any single iteration loops in the copy nests and simplify
// load/stores.
SmallVector<Operation *, 4> copyOps;
for (auto nest : copyNests)
// With a post order walk, the erasure of loops does not affect
// continuation of the walk or the collection of load/store ops.
nest->walk([&](Operation *op) {
if (auto forOp = dyn_cast<AffineForOp>(op))
promoteIfSingleIteration(forOp);
else if (auto loadOp = dyn_cast<AffineLoadOp>(op))
copyOps.push_back(loadOp);
else if (auto storeOp = dyn_cast<AffineStoreOp>(op))
copyOps.push_back(storeOp);
});
// Promoting single iteration loops could lead to simplification of
// generated load's/store's, and the latter could anyway also be
// canonicalized.
OwningRewritePatternList patterns;
for (auto op : copyOps) {
patterns.clear();
if (isa<AffineLoadOp>(op)) {
AffineLoadOp::getCanonicalizationPatterns(patterns, &getContext());
} else {
assert(isa<AffineStoreOp>(op) && "expected affine store op");
AffineStoreOp::getCanonicalizationPatterns(patterns, &getContext());
}
applyOpPatternsAndFold(op, std::move(patterns));
}
}
namespace mlir {
void registerTestAffineDataCopyPass() {
PassRegistration<TestAffineDataCopy>(
PASS_NAME, "Tests affine data copy utility functions.");
}
} // namespace mlir