156 lines · cpp
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-exception6//7//===----------------------------------------------------------------------===//8//9// This file implements a pass to test affine data copy utility functions and10// 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 TestAffineDataCopy31 : 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 ®istry) 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(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(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} // namespace64 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 loop73 // 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 loaded82 // 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(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 simplify118 // load/stores.119 SmallVector<Operation *, 4> copyOps;120 for (Operation *nest : copyNests) {121 // With a post order walk, the erasure of loops does not affect122 // continuation of the walk or the collection of load/store ops.123 nest->walk([&](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(loadOp);128 else if (auto storeOp = dyn_cast<AffineStoreOp>(op))129 copyOps.push_back(storeOp);130 });131 }132 133 // Promoting single iteration loops could lead to simplification of134 // generated load's/store's, and the latter could anyway also be135 // 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 mlir156