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/IR/PatternMatch.h"
17 #include "mlir/Pass/Pass.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;
32   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 
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.
92   for (auto nest : copyNests)
93     nest->walk([](AffineForOp forOp) { promoteIfSingleIteration(forOp); });
94 
95   // Promoting single iteration loops could lead to simplification
96   // of load's/store's. We will run the canonicalization patterns again.
97   OwningRewritePatternList patterns;
98   AffineLoadOp::getCanonicalizationPatterns(patterns, &getContext());
99   AffineStoreOp::getCanonicalizationPatterns(patterns, &getContext());
100   applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
101 }
102 
103 namespace mlir {
104 void registerTestAffineDataCopyPass() {
105   PassRegistration<TestAffineDataCopy>(
106       PASS_NAME, "Tests affine data copy utility functions.");
107 }
108 } // namespace mlir
109