1 //===- TestAffineDataCopy.cpp - Test affine data copy utility -------------===//
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 //
9 // This file implements a pass to test affine data copy utility functions and
10 // options.
11 //
12 //===----------------------------------------------------------------------===//
13
14 #include "mlir/Analysis/Utils.h"
15 #include "mlir/Dialect/Affine/IR/AffineOps.h"
16 #include "mlir/Pass/Pass.h"
17 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
18 #include "mlir/Transforms/LoopUtils.h"
19 #include "mlir/Transforms/Passes.h"
20
21 #define PASS_NAME "test-affine-data-copy"
22
23 using namespace mlir;
24
25 static llvm::cl::OptionCategory clOptionsCategory(PASS_NAME " options");
26
27 namespace {
28
29 struct TestAffineDataCopy
30 : public PassWrapper<TestAffineDataCopy, FunctionPass> {
31 TestAffineDataCopy() = default;
TestAffineDataCopy__anonf8d37b8b0111::TestAffineDataCopy32 TestAffineDataCopy(const TestAffineDataCopy &pass){};
33
34 void runOnFunction() override;
35
36 private:
37 Option<bool> clMemRefFilter{
38 *this, "memref-filter",
39 llvm::cl::desc(
40 "Enable memref filter testing in affine data copy optimization"),
41 llvm::cl::init(false)};
42 Option<bool> clTestGenerateCopyForMemRegion{
43 *this, "for-memref-region",
44 llvm::cl::desc("Test copy generation for a single memref region"),
45 llvm::cl::init(false)};
46 };
47
48 } // end anonymous namespace
49
runOnFunction()50 void TestAffineDataCopy::runOnFunction() {
51 // Gather all AffineForOps by loop depth.
52 std::vector<SmallVector<AffineForOp, 2>> depthToLoops;
53 gatherLoops(getFunction(), depthToLoops);
54 assert(depthToLoops.size() && "Loop nest not found");
55
56 // Only support tests with a single loop nest and a single innermost loop
57 // for now.
58 unsigned innermostLoopIdx = depthToLoops.size() - 1;
59 if (depthToLoops[0].size() != 1 || depthToLoops[innermostLoopIdx].size() != 1)
60 return;
61
62 auto loopNest = depthToLoops[0][0];
63 auto innermostLoop = depthToLoops[innermostLoopIdx][0];
64 AffineLoadOp load;
65 if (clMemRefFilter || clTestGenerateCopyForMemRegion) {
66 // Gather MemRef filter. For simplicity, we use the first loaded memref
67 // found in the innermost loop.
68 for (auto &op : *innermostLoop.getBody()) {
69 if (auto ld = dyn_cast<AffineLoadOp>(op)) {
70 load = ld;
71 break;
72 }
73 }
74 }
75
76 AffineCopyOptions copyOptions = {/*generateDma=*/false,
77 /*slowMemorySpace=*/0,
78 /*fastMemorySpace=*/0,
79 /*tagMemorySpace=*/0,
80 /*fastMemCapacityBytes=*/32 * 1024 * 1024UL};
81 DenseSet<Operation *> copyNests;
82 if (clMemRefFilter) {
83 affineDataCopyGenerate(loopNest, copyOptions, load.getMemRef(), copyNests);
84 } else if (clTestGenerateCopyForMemRegion) {
85 CopyGenerateResult result;
86 MemRefRegion region(loopNest.getLoc());
87 region.compute(load, /*loopDepth=*/0);
88 generateCopyForMemRegion(region, loopNest, copyOptions, result);
89 }
90
91 // Promote any single iteration loops in the copy nests and simplify
92 // load/stores.
93 SmallVector<Operation *, 4> copyOps;
94 for (auto nest : copyNests)
95 // With a post order walk, the erasure of loops does not affect
96 // continuation of the walk or the collection of load/store ops.
97 nest->walk([&](Operation *op) {
98 if (auto forOp = dyn_cast<AffineForOp>(op))
99 promoteIfSingleIteration(forOp);
100 else if (auto loadOp = dyn_cast<AffineLoadOp>(op))
101 copyOps.push_back(loadOp);
102 else if (auto storeOp = dyn_cast<AffineStoreOp>(op))
103 copyOps.push_back(storeOp);
104 });
105
106 // Promoting single iteration loops could lead to simplification of
107 // generated load's/store's, and the latter could anyway also be
108 // canonicalized.
109 OwningRewritePatternList patterns;
110 for (auto op : copyOps) {
111 patterns.clear();
112 if (isa<AffineLoadOp>(op)) {
113 AffineLoadOp::getCanonicalizationPatterns(patterns, &getContext());
114 } else {
115 assert(isa<AffineStoreOp>(op) && "expected affine store op");
116 AffineStoreOp::getCanonicalizationPatterns(patterns, &getContext());
117 }
118 applyOpPatternsAndFold(op, std::move(patterns));
119 }
120 }
121
122 namespace mlir {
registerTestAffineDataCopyPass()123 void registerTestAffineDataCopyPass() {
124 PassRegistration<TestAffineDataCopy>(
125 PASS_NAME, "Tests affine data copy utility functions.");
126 }
127 } // namespace mlir
128