58 lines · cpp
1//===- TestAffineWalk.cpp - Pass to test affine walks2//----------------------===//3//4// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.5// See https://llvm.org/LICENSE.txt for license information.6// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception7//8//===----------------------------------------------------------------------===//9 10#include "mlir/Pass/Pass.h"11 12#include "mlir/IR/BuiltinOps.h"13 14using namespace mlir;15 16namespace {17/// A test pass for verifying walk interrupts.18struct TestAffineWalk19 : public PassWrapper<TestAffineWalk, OperationPass<ModuleOp>> {20 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestAffineWalk)21 22 void runOnOperation() override;23 StringRef getArgument() const final { return "test-affine-walk"; }24 StringRef getDescription() const final { return "Test affine walk method."; }25};26} // namespace27 28/// Emits a remark for the first `map`'s result expression that contains a29/// mod expression.30static void checkMod(AffineMap map, Location loc) {31 for (AffineExpr e : map.getResults()) {32 e.walk([&](AffineExpr s) {33 if (s.getKind() == mlir::AffineExprKind::Mod) {34 emitRemark(loc, "mod expression: ");35 return WalkResult::interrupt();36 }37 return WalkResult::advance();38 });39 }40}41 42void TestAffineWalk::runOnOperation() {43 auto m = getOperation();44 // Test whether the walk is being correctly interrupted.45 m.walk([](Operation *op) {46 for (NamedAttribute attr : op->getAttrs()) {47 auto mapAttr = dyn_cast<AffineMapAttr>(attr.getValue());48 if (!mapAttr)49 return;50 checkMod(mapAttr.getAffineMap(), op->getLoc());51 }52 });53}54 55namespace mlir {56void registerTestAffineWalk() { PassRegistration<TestAffineWalk>(); }57} // namespace mlir58