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