Page MenuHomePhabricator
Differential D75965 Diff 249528 mlir/test/lib/Transforms/TestAffineDataCopy.cpp

Changeset View

Standalone View

View Options

mlir/test/lib/Transforms/TestAffineDataCopy.cpp

//===- TestAffineDataCopy.cpp - Test affine data copy utility -------------===// //===- TestAffineDataCopy.cpp - Test affine data copy utility -------------===//
// //
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information. // See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// //
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// //
// This file implements a pass to test affine data copy utility functions and // This file implements a pass to test affine data copy utility functions and
// options. // options.
// //
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
#include "mlir/Analysis/Passes.h" #include "mlir/Analysis/Passes.h"
#include "mlir/Analysis/Utils.h"
#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/AffineOps/AffineOps.h"
#include "mlir/Pass/Pass.h" #include "mlir/Pass/Pass.h"
#include "mlir/Transforms/LoopUtils.h" #include "mlir/Transforms/LoopUtils.h"
#include "mlir/Transforms/Passes.h" #include "mlir/Transforms/Passes.h"
#define PASS_NAME "test-affine-data-copy" #define PASS_NAME "test-affine-data-copy"
using namespace mlir; using namespace mlir;
Show All 9 Linesstruct TestAffineDataCopy : public FunctionPass<TestAffineDataCopy> {
void runOnFunction() override; void runOnFunction() override;
private: private:
Option<bool> clMemRefFilter{ Option<bool> clMemRefFilter{
*this, "memref-filter", *this, "memref-filter",
llvm::cl::desc( llvm::cl::desc(
"Enable memref filter testing in affine data copy optimization"), "Enable memref filter testing in affine data copy optimization"),
llvm::cl::init(false)}; llvm::cl::init(false)};
Option<bool> clTestGenerateDataCopyAroundOp{
*this, "test-generate-data-copy-around-op",
llvm::cl::desc("Test through generateCopyAroundOp"),
bondhugulaUnsubmitted
Done
ReplyInline Actions

Test copy generation for a single op / memref region?

bondhugula: Test copy generation for a single op / memref region?
llvm::cl::init(false)};
}; };
} // end anonymous namespace } // end anonymous namespace
void TestAffineDataCopy::runOnFunction() { void TestAffineDataCopy::runOnFunction() {
// Gather all AffineForOps by loop depth. // Gather all AffineForOps by loop depth.
std::vector<SmallVector<AffineForOp, 2>> depthToLoops; std::vector<SmallVector<AffineForOp, 2>> depthToLoops;
gatherLoops(getFunction(), depthToLoops); gatherLoops(getFunction(), depthToLoops);
assert(depthToLoops.size() && "Loop nest not found"); assert(depthToLoops.size() && "Loop nest not found");
// Only support tests with a single loop nest and a single innermost loop // Only support tests with a single loop nest and a single innermost loop
// for now. // for now.
unsigned innermostLoopIdx = depthToLoops.size() - 1; unsigned innermostLoopIdx = depthToLoops.size() - 1;
if (depthToLoops[0].size() != 1 || depthToLoops[innermostLoopIdx].size() != 1) if (depthToLoops[0].size() != 1 || depthToLoops[innermostLoopIdx].size() != 1)
return; return;
auto loopNest = depthToLoops[0][0]; auto loopNest = depthToLoops[0][0];
auto innermostLoop = depthToLoops[innermostLoopIdx][0]; auto innermostLoop = depthToLoops[innermostLoopIdx][0];
Optional<Value> memrefFilter; AffineLoadOp load;
if (clMemRefFilter) { if (clMemRefFilter || clTestGenerateDataCopyAroundOp) {
// Gather MemRef filter. For simplicity, we use the first loaded memref // Gather MemRef filter. For simplicity, we use the first loaded memref
// found in the innermost loop. // found in the innermost loop.
for (auto &op : *innermostLoop.getBody()) { for (auto &op : *innermostLoop.getBody()) {
if (auto load = dyn_cast<AffineLoadOp>(op)) { if (auto ld = dyn_cast<AffineLoadOp>(op)) {
memrefFilter = load.getMemRef(); load = ld;
break; break;
} }
} }
} }
AffineCopyOptions copyOptions = {/*generateDma=*/false, AffineCopyOptions copyOptions = {/*generateDma=*/false,
/*slowMemorySpace=*/0, /*slowMemorySpace=*/0,
/*fastMemorySpace=*/0, /*fastMemorySpace=*/0,
/*tagMemorySpace=*/0, /*tagMemorySpace=*/0,
/*fastMemCapacityBytes=*/32 * 1024 * 1024UL}; /*fastMemCapacityBytes=*/32 * 1024 * 1024UL};
if (clMemRefFilter) {
DenseSet<Operation *> copyNests; DenseSet<Operation *> copyNests;
affineDataCopyGenerate(loopNest, copyOptions, memrefFilter, copyNests); affineDataCopyGenerate(loopNest, copyOptions, load.getMemRef(), copyNests);
} else if (clTestGenerateDataCopyAroundOp) {
CopyGenerateResult result;
MemRefRegion region(loopNest.getLoc());
region.compute(load, 0);
bondhugulaUnsubmitted
Done
ReplyInline Actions

/*loopDepth = */0

bondhugula: /*loopDepth = */0
generateDataCopyAroundOp(region, loopNest, copyOptions, result);
}
} }
namespace mlir { namespace mlir {
void registerTestAffineDataCopyPass() { void registerTestAffineDataCopyPass() {
PassRegistration<TestAffineDataCopy>( PassRegistration<TestAffineDataCopy>(
PASS_NAME, "Tests affine data copy utility functions."); PASS_NAME, "Tests affine data copy utility functions.");
} }
} // namespace mlir } // namespace mlir