[mlir] Retain metadata for single loc fusedloc

[jpienaar]

Summary

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.

[rriddle]

mlir/lib/IR/Location.cpp

108–109

Might want to add a comment here about not dropping metadata.

112–113

Can you not just pass it directly?

[jpienaar]

mlir/lib/IR/Location.cpp

112–113

No, funnily that led to a heap use after free asan failure (I was surprised by that too, didn't dig though)

[mehdi_amini]

mlir/lib/IR/Location.cpp

112–113

That seems fishy, can you shared the ASAN failure?

[Chia-hungDuan]

mlir/lib/IR/Location.cpp

112–113

Also curious about this case, I applied this patch and passed UnknownLoc::get() directly. Didn't see the problem while running the tests under llvm/llvm-project/mlir/test/IR

Saw other failures when run with asan but it doesn't matter to passed UnknownLoc::get() directly or not.

[jpienaar]

mlir/lib/IR/Location.cpp

112–113

Sure,

==60765==ERROR: AddressSanitizer: stack-use-after-scope on address 0xA at pc 0xA bp 0xA sp 0xA
READ of size 8 at 0xA thread T0
    #0 0xA in get_hashable_data<mlir::Location> //llvm-project/llvm/include/llvm/ADT/Hashing.h:379:21
    #1 0xA in llvm::hash_code llvm::hashing::detail::hash_combine_range_impl<mlir::Location const*>(mlir::Location const*, mlir::Location 
const*) //llvm-project/llvm/include/llvm/ADT/Hashing.h:412:45
    #2 0xA in hash_combine_range<const mlir::Location *> //llvm-project/llvm/include/llvm/ADT/Hashing.h:484:10
    #3 0xA in hash_value<mlir::Location> //llvm-project/llvm/include/llvm/ADT/ArrayRef.h:569:12
    #4 0xA in get_hashable_data<llvm::ArrayRef<mlir::Location> > //llvm-project/llvm/include/llvm/ADT/Hashing.h:379:10
    #5 0xA in combine<llvm::ArrayRef<mlir::Location>, mlir::Attribute> //llvm-project/llvm/include/llvm/ADT/Hashing.h:560:63
    #6 0xA in hash_combine<llvm::ArrayRef<mlir::Location>, mlir::Attribute> //llvm-project/llvm/include/llvm/ADT/Hashing.h:608:17
    #7 0xA in hashKey build_asan/tools/mlir/include/mlir/IR/BuiltinLocationAttributes.cpp.inc:123:12
    #8 0xA in getHash<mlir::detail::FusedLocAttrStorage, std::tuple<llvm::ArrayRef<mlir::Location>, mlir::Attribute> > //llvm/mlir/include/mlir/Support/StorageUniquer.h:317:12
    #9 0xA in mlir::detail::FusedLocAttrStorage* mlir::StorageUniquer::get<mlir::detail::FusedLocAttrStorage, mlir::UnknownLoc&, mlir::Attribute&>(llvm::function_ref<void (mlir::detail::FusedLocAttrStorage*)>, mlir::TypeID, mlir::UnknownLoc&, mlir::Attribute&) //llvm/mlir/include/mlir/Support/StorageUniquer.h:194:26
    #10 0xA in get<mlir::FusedLoc, mlir::UnknownLoc &, mlir::Attribute &> //llvm-project/mlir/include/mlir/IR/AttributeSupport.h:192:39
    #11 0xA in get<mlir::UnknownLoc, mlir::Attribute> //llvm-project/mlir/include/mlir/IR/StorageUniquerSupport.h:141:12
    #12 0xA in mlir::FusedLoc::get(llvm::ArrayRef<mlir::Location>, mlir::Attribute, mlir::MLIRContext*) //llvm/mlir/lib/IR/Location.cpp:114:12
    #13 0xA in mlirLocationFusedGet //llvm-project/mlir/lib/CAPI/IR/IR.cpp:141:15
    #14 0xA  (build_asan/tools/mlir/python_packages/mlir_core/mlir/_mlir_libs/_mlir.cpython-39-x86_64-linux-gnu.so+0x29a6a9)

this is with

+ if (locs.empty()) {
+ if (!metadata)
+ return UnknownLoc::get(context);
+ // std::array<Location, 1> loc({UnknownLoc::get(context)});
+ return Base::get(context, UnknownLoc::get(context), metadata);
+ }
+ if (locs.size() == 1 && !metadata)

[jpienaar]

Update comment & add TODO

[mehdi_amini]

mlir/lib/IR/Location.cpp

112–113

I'm worried this points at a bug in:

#9 0xA in mlir::detail::FusedLocAttrStorage* mlir::StorageUniquer::get<mlir::detail::FusedLocAttrStorage, mlir::UnknownLoc&, mlir::Attribute&>(llvm::function_ref<void (mlir::detail::FusedLocAttrStorage*)>, mlir::TypeID, mlir::UnknownLoc&, mlir::Attribute&) //llvm/mlir/include/mlir/Support/StorageUniquer.h:194:26

Where:

// Construct a value of the derived key type.
auto derivedKey = getKey<Storage>(std::forward<Args>(args)...);

Has to return std::tuple<llvm::ArrayRef<mlir::Location>, mlir::Attribute>> but the arguments are mlir::UnknownLoc&, mlir::Attribute& and it'll internally create an ArrayRef from a local copy of the mlir::UnknownLoc&.

River is best to pin point where this happens exactly I think.

mlir/lib/Bindings/Python/IRCore.cpp

py::class_<PyLocation>(m, "Location", py::module_local())
           },
           py::arg("filename"), py::arg("line"), py::arg("col"),
           py::arg("context") = py::none(), kContextGetFileLocationDocstring)
       .def_static(
           "fused",
           [](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)
               locations.push_back(pyLocation.get());
             MlirLocation location = mlirLocationFusedGet(
                 context->get(), locations.size(), locations.data(),
                 metadata ? metadata->get() : MlirAttribute{0});
             return PyLocation(context->getRef(), location);

mlir/lib/IR/Location.cpp

if (auto fusedLoc = loc.dyn_cast<FusedLoc>()) {
         continue;
       }
     }
     // Otherwise, only add known locations to the set.
     if (!loc.isa<UnknownLoc>())
       decomposedLocs.insert(loc);
   }
   locs = decomposedLocs.getArrayRef();
 
-  // Handle the simple cases of less than two locations.
-  if (locs.empty())
+  // 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);
-  if (locs.size() == 1)
+    // 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);
 }

mlir/test/IR/locations.mlir

func @inline_notation() -> i32 {
   // CHECK: } loc(fused["foo", "mysource.cc":10:8])
   affine.for %i0 = 0 to 8 {
   } loc(fused["foo", "mysource.cc":10:8])
 
   // CHECK: } loc(fused<"myPass">["foo", "foo2"])
   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)
 }
 
 // CHECK-LABEL: func private @loc_attr(i1 {foo.loc_attr = loc(callsite("foo" at "mysource.cc":10:8))})
 func private @loc_attr(i1 {foo.loc_attr = loc(callsite("foo" at "mysource.cc":10:8))})
 
   // Check that locations get properly escaped.

mlir/test/python/ir/location.py

def testCallSite():
   print("file repr:", repr(loc))
 
 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"])
   print("file repr:", repr(loc))
   # CHECK: file str: loc(fused<"sauteed">["carrot", "potatoes"])
   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)
 
 
 # CHECK-LABEL: TEST: testLocationCapsule
 def testLocationCapsule():
   with Context() as ctx:
     loc1 = Location.file("foo.txt", 123, 56)
   # CHECK: mlir.ir.Location._CAPIPtr
   loc_capsule = loc1._CAPIPtr
   print(loc_capsule)
   loc2 = Location._CAPICreate(loc_capsule)
   assert loc2 == loc1
   assert loc2.context is ctx
 
 run(testLocationCapsule)