| //===- 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 |