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