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/Pass/Pass.h"
17 #include "mlir/Transforms/LoopUtils.h"
18 #include "mlir/Transforms/Passes.h"
19 
20 #define PASS_NAME "test-affine-data-copy"
21 
22 using namespace mlir;
23 
24 static llvm::cl::OptionCategory clOptionsCategory(PASS_NAME " options");
25 
26 namespace {
27 
28 struct TestAffineDataCopy : public FunctionPass<TestAffineDataCopy> {
29   TestAffineDataCopy() = default;
30   TestAffineDataCopy(const TestAffineDataCopy &pass){};
31 
32   void runOnFunction() override;
33 
34 private:
35   Option<bool> clMemRefFilter{
36       *this, "memref-filter",
37       llvm::cl::desc(
38           "Enable memref filter testing in affine data copy optimization"),
39       llvm::cl::init(false)};
40   Option<bool> clTestGenerateCopyForMemRegion{
41       *this, "for-memref-region",
42       llvm::cl::desc("Test copy generation for a single memref region"),
43       llvm::cl::init(false)};
44 };
45 
46 } // end anonymous namespace
47 
48 void TestAffineDataCopy::runOnFunction() {
49   // Gather all AffineForOps by loop depth.
50   std::vector<SmallVector<AffineForOp, 2>> depthToLoops;
51   gatherLoops(getFunction(), depthToLoops);
52   assert(depthToLoops.size() && "Loop nest not found");
53 
54   // Only support tests with a single loop nest and a single innermost loop
55   // for now.
56   unsigned innermostLoopIdx = depthToLoops.size() - 1;
57   if (depthToLoops[0].size() != 1 || depthToLoops[innermostLoopIdx].size() != 1)
58     return;
59 
60   auto loopNest = depthToLoops[0][0];
61   auto innermostLoop = depthToLoops[innermostLoopIdx][0];
62   AffineLoadOp load;
63   if (clMemRefFilter || clTestGenerateCopyForMemRegion) {
64     // Gather MemRef filter. For simplicity, we use the first loaded memref
65     // found in the innermost loop.
66     for (auto &op : *innermostLoop.getBody()) {
67       if (auto ld = dyn_cast<AffineLoadOp>(op)) {
68         load = ld;
69         break;
70       }
71     }
72   }
73 
74   AffineCopyOptions copyOptions = {/*generateDma=*/false,
75                                    /*slowMemorySpace=*/0,
76                                    /*fastMemorySpace=*/0,
77                                    /*tagMemorySpace=*/0,
78                                    /*fastMemCapacityBytes=*/32 * 1024 * 1024UL};
79   if (clMemRefFilter) {
80     DenseSet<Operation *> copyNests;
81     affineDataCopyGenerate(loopNest, copyOptions, load.getMemRef(), copyNests);
82   } else if (clTestGenerateCopyForMemRegion) {
83     CopyGenerateResult result;
84     MemRefRegion region(loopNest.getLoc());
85     region.compute(load, /*loopDepth=*/0);
86     generateCopyForMemRegion(region, loopNest, copyOptions, result);
87   }
88 }
89 
90 namespace mlir {
91 void registerTestAffineDataCopyPass() {
92   PassRegistration<TestAffineDataCopy>(
93       PASS_NAME, "Tests affine data copy utility functions.");
94 }
95 } // namespace mlir
96