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
25using namespace mlir;
26using namespace mlir::affine;
27
28namespace {
29
30struct TestAffineDataCopy
31 : public PassWrapper<TestAffineDataCopy, OperationPass<func::FuncOp>> {
32 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestAffineDataCopy)
33
34 StringRef getArgument() const final { return PASS_NAME; }
35 StringRef getDescription() const final {
36 return "Tests affine data copy utility functions.";
37 }
38 TestAffineDataCopy() = default;
39 TestAffineDataCopy(const TestAffineDataCopy &pass) : PassWrapper(pass){};
40
41 void getDependentDialects(DialectRegistry &registry) const override {
42 registry.insert<memref::MemRefDialect>();
43 }
44 void runOnOperation() override;
45
46private:
47 Option<bool> clMemRefFilter{
48 *this, "memref-filter",
49 llvm::cl::desc(
50 "Enable memref filter testing in affine data copy optimization"),
51 llvm::cl::init(Val: false)};
52 Option<bool> clTestGenerateCopyForMemRegion{
53 *this, "for-memref-region",
54 llvm::cl::desc("Test copy generation for a single memref region"),
55 llvm::cl::init(Val: false)};
56 Option<uint64_t> clCapacityLimit{
57 *this, "capacity-kib",
58 llvm::cl::desc("Test copy generation enforcing a limit of capacity "
59 "(default: unlimited)"),
60 llvm::cl::init(UINT64_MAX)};
61};
62
63} // namespace
64
65void TestAffineDataCopy::runOnOperation() {
66 // Gather all AffineForOps by loop depth.
67 std::vector<SmallVector<AffineForOp, 2>> depthToLoops;
68 gatherLoops(getOperation(), depthToLoops);
69 if (depthToLoops.empty())
70 return;
71
72 // Only support tests with a single loop nest and a single innermost loop
73 // for now.
74 unsigned innermostLoopIdx = depthToLoops.size() - 1;
75 if (depthToLoops[0].size() != 1 || depthToLoops[innermostLoopIdx].size() != 1)
76 return;
77
78 auto loopNest = depthToLoops[0][0];
79 auto innermostLoop = depthToLoops[innermostLoopIdx][0];
80 AffineLoadOp load;
81 // For simplicity, these options are tested on the first memref being loaded
82 // from in the innermost loop.
83 if (clMemRefFilter || clTestGenerateCopyForMemRegion) {
84 for (auto &op : *innermostLoop.getBody()) {
85 if (auto ld = dyn_cast<AffineLoadOp>(op)) {
86 load = ld;
87 break;
88 }
89 }
90 if (!load)
91 return;
92 }
93
94 AffineCopyOptions copyOptions = {/*generateDma=*/false,
95 /*slowMemorySpace=*/0,
96 /*fastMemorySpace=*/0,
97 /*tagMemorySpace=*/0,
98 /*fastMemCapacityBytes=*/clCapacityLimit};
99 DenseSet<Operation *> copyNests;
100 if (clMemRefFilter) {
101 if (failed(affineDataCopyGenerate(loopNest, copyOptions, load.getMemRef(),
102 copyNests)))
103 return;
104 } else if (clTestGenerateCopyForMemRegion) {
105 CopyGenerateResult result;
106 MemRefRegion region(loopNest.getLoc());
107 if (failed(region.compute(op: load, /*loopDepth=*/0)))
108 return;
109 if (failed(generateCopyForMemRegion(region, loopNest, copyOptions, result)))
110 return;
111 } else if (failed(affineDataCopyGenerate(loopNest, copyOptions,
112 /*filterMemref=*/std::nullopt,
113 copyNests))) {
114 return;
115 }
116
117 // Promote any single iteration loops in the copy nests and simplify
118 // load/stores.
119 SmallVector<Operation *, 4> copyOps;
120 for (Operation *nest : copyNests) {
121 // With a post order walk, the erasure of loops does not affect
122 // continuation of the walk or the collection of load/store ops.
123 nest->walk(callback: [&](Operation *op) {
124 if (auto forOp = dyn_cast<AffineForOp>(op))
125 (void)promoteIfSingleIteration(forOp);
126 else if (auto loadOp = dyn_cast<AffineLoadOp>(op))
127 copyOps.push_back(Elt: loadOp);
128 else if (auto storeOp = dyn_cast<AffineStoreOp>(op))
129 copyOps.push_back(Elt: storeOp);
130 });
131 }
132
133 // Promoting single iteration loops could lead to simplification of
134 // generated load's/store's, and the latter could anyway also be
135 // canonicalized.
136 RewritePatternSet patterns(&getContext());
137 for (Operation *op : copyOps) {
138 patterns.clear();
139 if (isa<AffineLoadOp>(op)) {
140 AffineLoadOp::getCanonicalizationPatterns(patterns, &getContext());
141 } else {
142 assert(isa<AffineStoreOp>(op) && "expected affine store op");
143 AffineStoreOp::getCanonicalizationPatterns(patterns, &getContext());
144 }
145 }
146 GreedyRewriteConfig config;
147 config.setStrictness(GreedyRewriteStrictness::ExistingAndNewOps);
148 (void)applyOpPatternsGreedily(copyOps, std::move(patterns), config);
149}
150
151namespace mlir {
152void registerTestAffineDataCopyPass() {
153 PassRegistration<TestAffineDataCopy>();
154}
155} // namespace mlir
156

source code of mlir/test/lib/Dialect/Affine/TestAffineDataCopy.cpp