llvm-capstone/mlir/unittests/IR/SymbolTableTest.cpp
Ingo Müller 479057887f
[mlir] Make overloads of SymbolTable::replaceAllSymbolUses consistent. (#68320)
This function has several overloads that allow to specify the symbol
that should be renamed and the scope for that renaming in different
ways. The overloads were inconsistent in the following way (quoted
strings are `StringAttr`s, other variables are `Operation *`):

* `replaceAllSymbolUses(symbolOp, "new_symbol", scopeOp)` would traverse
into the nested regions of `scopeOp` and hence rename the symbol inside
of `scopeOp`.
* `replaceAllSymbolUses("symbol", "new_symbol", scopeOp)` would *not*
traverse into the nested regions of `scopeOp` and hence *not* rename the
symbol.

The underlying behavior was spread over different places and is somewhat
hard to understand. The two overloads above mainly differed by what
`collectSymbolScopes` computed, which is itself overloaded. If `scopeOp`
is a top-level module, then the overload on `(Operation *, Operation
*)`, which is used in the first of the above cases, computes a scope
where the body region of the module is the `limit`; however, the
overload on `(StringAttr, Operation *)` computed the module op itself as
the `limit`. Later, `walkSymbolTable` would walk the body of the module
if it was given as a region but it would *not* enter the regions of the
module op because that op has a symbol table (which was assumed to be a
*different* scope).

The fix in this commit is change the behavior of `collectSymbolScopes`
such that the `(StringAttr, Operation *)` overload returns a scope for
each region in the `limit` argument.
2023-10-10 07:47:08 +02:00

137 lines
4.8 KiB
C++

//===- SymbolTableTest.cpp - SymbolTable unit tests -----------------------===//
//
// 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/IR/SymbolTable.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/Verifier.h"
#include "mlir/Interfaces/CallInterfaces.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Parser/Parser.h"
#include "gtest/gtest.h"
using namespace mlir;
namespace test {
void registerTestDialect(DialectRegistry &);
} // namespace test
class ReplaceAllSymbolUsesTest : public ::testing::Test {
protected:
using ReplaceFnType = llvm::function_ref<LogicalResult(
SymbolTable, ModuleOp, Operation *, Operation *)>;
void SetUp() override {
::test::registerTestDialect(registry);
context = std::make_unique<MLIRContext>(registry);
}
void testReplaceAllSymbolUses(ReplaceFnType replaceFn) {
// Set up IR and find func ops.
OwningOpRef<ModuleOp> module =
parseSourceString<ModuleOp>(kInput, context.get());
SymbolTable symbolTable(module.get());
auto opIterator = module->getBody(0)->getOperations().begin();
auto fooOp = cast<FunctionOpInterface>(opIterator++);
auto barOp = cast<FunctionOpInterface>(opIterator++);
ASSERT_EQ(fooOp.getNameAttr(), "foo");
ASSERT_EQ(barOp.getNameAttr(), "bar");
// Call test function that does symbol replacement.
LogicalResult res = replaceFn(symbolTable, module.get(), fooOp, barOp);
ASSERT_TRUE(succeeded(res));
ASSERT_TRUE(succeeded(verify(module.get())));
// Check that it got renamed.
bool calleeFound = false;
fooOp->walk([&](CallOpInterface callOp) {
StringAttr callee = callOp.getCallableForCallee()
.dyn_cast<SymbolRefAttr>()
.getLeafReference();
EXPECT_EQ(callee, "baz");
calleeFound = true;
});
EXPECT_TRUE(calleeFound);
}
std::unique_ptr<MLIRContext> context;
private:
constexpr static llvm::StringLiteral kInput = R"MLIR(
module {
test.conversion_func_op private @foo() {
"test.conversion_call_op"() { callee=@bar } : () -> ()
"test.return"() : () -> ()
}
test.conversion_func_op private @bar()
}
)MLIR";
DialectRegistry registry;
};
namespace {
TEST_F(ReplaceAllSymbolUsesTest, OperationInModuleOp) {
// Symbol as `Operation *`, rename within module.
testReplaceAllSymbolUses([&](auto symbolTable, auto module, auto fooOp,
auto barOp) -> LogicalResult {
return symbolTable.replaceAllSymbolUses(
barOp, StringAttr::get(context.get(), "baz"), module);
});
}
TEST_F(ReplaceAllSymbolUsesTest, StringAttrInModuleOp) {
// Symbol as `StringAttr`, rename within module.
testReplaceAllSymbolUses([&](auto symbolTable, auto module, auto fooOp,
auto barOp) -> LogicalResult {
return symbolTable.replaceAllSymbolUses(
StringAttr::get(context.get(), "bar"),
StringAttr::get(context.get(), "baz"), module);
});
}
TEST_F(ReplaceAllSymbolUsesTest, OperationInModuleBody) {
// Symbol as `Operation *`, rename within module body.
testReplaceAllSymbolUses([&](auto symbolTable, auto module, auto fooOp,
auto barOp) -> LogicalResult {
return symbolTable.replaceAllSymbolUses(
barOp, StringAttr::get(context.get(), "baz"), &module->getRegion(0));
});
}
TEST_F(ReplaceAllSymbolUsesTest, StringAttrInModuleBody) {
// Symbol as `StringAttr`, rename within module body.
testReplaceAllSymbolUses([&](auto symbolTable, auto module, auto fooOp,
auto barOp) -> LogicalResult {
return symbolTable.replaceAllSymbolUses(
StringAttr::get(context.get(), "bar"),
StringAttr::get(context.get(), "baz"), &module->getRegion(0));
});
}
TEST_F(ReplaceAllSymbolUsesTest, OperationInFuncOp) {
// Symbol as `Operation *`, rename within function.
testReplaceAllSymbolUses([&](auto symbolTable, auto module, auto fooOp,
auto barOp) -> LogicalResult {
return symbolTable.replaceAllSymbolUses(
barOp, StringAttr::get(context.get(), "baz"), fooOp);
});
}
TEST_F(ReplaceAllSymbolUsesTest, StringAttrInFuncOp) {
// Symbol as `StringAttr`, rename within function.
testReplaceAllSymbolUses([&](auto symbolTable, auto module, auto fooOp,
auto barOp) -> LogicalResult {
return symbolTable.replaceAllSymbolUses(
StringAttr::get(context.get(), "bar"),
StringAttr::get(context.get(), "baz"), fooOp);
});
}
} // namespace