57 lines
1.7 KiB
C++
57 lines
1.7 KiB
C++
//===- TestAffineWalk.cpp - Pass to test affine walks
|
|
//----------------------===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Pass/Pass.h"
|
|
|
|
#include "mlir/IR/BuiltinOps.h"
|
|
|
|
using namespace mlir;
|
|
|
|
namespace {
|
|
/// A test pass for verifying walk interrupts.
|
|
struct TestAffineWalk
|
|
: public PassWrapper<TestAffineWalk, OperationPass<ModuleOp>> {
|
|
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestAffineWalk)
|
|
|
|
void runOnOperation() override;
|
|
StringRef getArgument() const final { return "test-affine-walk"; }
|
|
StringRef getDescription() const final { return "Test affine walk method."; }
|
|
};
|
|
} // namespace
|
|
|
|
/// Emits a remark for the first `map`'s result expression that contains a
|
|
/// mod expression.
|
|
static void checkMod(AffineMap map, Location loc) {
|
|
for (AffineExpr e : map.getResults()) {
|
|
e.walk([&](AffineExpr s) {
|
|
if (s.getKind() == mlir::AffineExprKind::Mod) {
|
|
emitRemark(loc, "mod expression: ");
|
|
return WalkResult::interrupt();
|
|
}
|
|
return WalkResult::advance();
|
|
});
|
|
}
|
|
}
|
|
|
|
void TestAffineWalk::runOnOperation() {
|
|
auto m = getOperation();
|
|
// Test whether the walk is being correctly interrupted.
|
|
m.walk([](Operation *op) {
|
|
for (NamedAttribute attr : op->getAttrs()) {
|
|
auto mapAttr = attr.getValue().dyn_cast<AffineMapAttr>();
|
|
if (!mapAttr)
|
|
return;
|
|
checkMod(mapAttr.getAffineMap(), op->getLoc());
|
|
}
|
|
});
|
|
}
|
|
|
|
namespace mlir {
|
|
void registerTestAffineWalk() { PassRegistration<TestAffineWalk>(); }
|
|
} // namespace mlir
|