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