[mlir][python] Add fused location

This commit is contained in:
Jacques Pienaar 2021-12-11 10:12:29 -08:00
parent 6bcf1f9181
commit 1ab3efac41
3 changed files with 43 additions and 0 deletions

View File

@ -47,6 +47,9 @@ static const char kContextGetCallSiteLocationDocstring[] =
static const char kContextGetFileLocationDocstring[] =
R"(Gets a Location representing a file, line and column)";
static const char kContextGetFusedLocationDocstring[] =
R"(Gets a Location representing a fused location with optional metadata)";
static const char kContextGetNameLocationDocString[] =
R"(Gets a Location representing a named location with optional child location)";
@ -2197,6 +2200,23 @@ void mlir::python::populateIRCore(py::module &m) {
},
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);
},
py::arg("locations"), py::arg("metadata") = py::none(),
py::arg("context") = py::none(), kContextGetFusedLocationDocstring)
.def_static(
"name",
[](std::string name, llvm::Optional<PyLocation> childLoc,

View File

@ -658,6 +658,8 @@ class Location:
@staticmethod
def file(filename: str, line: int, col: int, context: Optional["Context"] = None) -> "Location": ...
@staticmethod
def fused(locations: Sequence["Location"], metadata: Optional["Attribute"] = None, context: Optional["Context"] = None) -> "Location": ...
@staticmethod
def name(name: str, childLoc: Optional["Location"] = None, context: Optional["Context"] = None) -> "Location": ...
@staticmethod
def unknown(context: Optional["Context"] = None) -> Any: ...

View File

@ -75,6 +75,27 @@ def testCallSite():
run(testCallSite)
# CHECK-LABEL: TEST: testFused
def testFused():
with Context() as ctx:
loc = Location.fused(
[Location.name("apple"), Location.name("banana")])
attr = Attribute.parse('"sauteed"')
loc_attr = Location.fused([Location.name("carrot"),
Location.name("potatoes")], attr)
ctx = None
# 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))
run(testFused)
# CHECK-LABEL: TEST: testLocationCapsule
def testLocationCapsule():
with Context() as ctx: