[mlir] Retain metadata for single loc fusedloc

If a fusedloc is created with a single location then no fusedloc
was previously created and single location returned instead. In the case
where there is a metadata associated with the location this results in
discarding the metadata. Instead only canonicalize where there is no
loss of information.

Differential Revision: https://reviews.llvm.org/D115605
This commit is contained in:
Jacques Pienaar 2022-01-04 15:37:33 -08:00
parent d496abbe2a
commit 9e365fe326
4 changed files with 36 additions and 6 deletions

View File

@ -2363,8 +2363,6 @@ void mlir::python::populateIRCore(py::module &m) {
[](const std::vector<PyLocation> &pyLocations,
llvm::Optional<PyAttribute> metadata,
DefaultingPyMlirContext context) {
if (pyLocations.empty())
throw py::value_error("No locations provided");
llvm::SmallVector<MlirLocation, 4> locations;
locations.reserve(pyLocations.size());
for (auto &pyLocation : pyLocations)

View File

@ -106,10 +106,18 @@ Location FusedLoc::get(ArrayRef<Location> locs, Attribute metadata,
}
locs = decomposedLocs.getArrayRef();
// Handle the simple cases of less than two locations.
if (locs.empty())
return UnknownLoc::get(context);
if (locs.size() == 1)
// Handle the simple cases of less than two locations. Ensure the metadata (if
// provided) is not dropped.
if (locs.empty()) {
if (!metadata)
return UnknownLoc::get(context);
// TODO: Investigate ASAN failure when using implicit conversion from
// Location to ArrayRef<Location> below.
return Base::get(context, ArrayRef<Location>{UnknownLoc::get(context)},
metadata);
}
if (locs.size() == 1 && !metadata)
return locs.front();
return Base::get(context, locs, metadata);
}

View File

@ -21,6 +21,10 @@ func @inline_notation() -> i32 {
affine.if #set0(%2) {
} loc(fused<"myPass">["foo", "foo2"])
// CHECK: } loc(fused<"myPass">["foo"])
affine.if #set0(%2) {
} loc(fused<"myPass">["foo"])
// CHECK: return %0 : i32 loc(unknown)
return %1 : i32 loc(unknown)
}

View File

@ -78,12 +78,20 @@ run(testCallSite)
# CHECK-LABEL: TEST: testFused
def testFused():
with Context() as ctx:
loc_single = Location.fused([Location.name("apple")])
loc = Location.fused(
[Location.name("apple"), Location.name("banana")])
attr = Attribute.parse('"sauteed"')
loc_attr = Location.fused([Location.name("carrot"),
Location.name("potatoes")], attr)
loc_empty = Location.fused([])
loc_empty_attr = Location.fused([], attr)
loc_single_attr = Location.fused([Location.name("apple")], attr)
ctx = None
# CHECK: file str: loc("apple")
print("file str:", str(loc_single))
# CHECK: file repr: loc("apple")
print("file repr:", repr(loc_single))
# CHECK: file str: loc(fused["apple", "banana"])
print("file str:", str(loc))
# CHECK: file repr: loc(fused["apple", "banana"])
@ -92,6 +100,18 @@ def testFused():
print("file str:", str(loc_attr))
# CHECK: file repr: loc(fused<"sauteed">["carrot", "potatoes"])
print("file repr:", repr(loc_attr))
# CHECK: file str: loc(unknown)
print("file str:", str(loc_empty))
# CHECK: file repr: loc(unknown)
print("file repr:", repr(loc_empty))
# CHECK: file str: loc(fused<"sauteed">[unknown])
print("file str:", str(loc_empty_attr))
# CHECK: file repr: loc(fused<"sauteed">[unknown])
print("file repr:", repr(loc_empty_attr))
# CHECK: file str: loc(fused<"sauteed">["apple"])
print("file str:", str(loc_single_attr))
# CHECK: file repr: loc(fused<"sauteed">["apple"])
print("file repr:", repr(loc_single_attr))
run(testFused)