[Mlir-commits] [mlir] 1ab3efa - [mlir][python] Add fused location

[Jacques Pienaar]

Author: Jacques Pienaar
Date: 2021-12-11T10:16:13-08:00
New Revision: 1ab3efac415e0d59f9686be482ea522ff0b92f99

URL: https://github.com/llvm/llvm-project/commit/1ab3efac415e0d59f9686be482ea522ff0b92f99
DIFF: https://github.com/llvm/llvm-project/commit/1ab3efac415e0d59f9686be482ea522ff0b92f99.diff

LOG: [mlir][python] Add fused location

Added: 
    

Modified: 
    mlir/lib/Bindings/Python/IRCore.cpp
    mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
    mlir/test/python/ir/location.py

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index cd5755248ea2e..3640a15e3407b 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -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,

diff  --git a/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi b/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
index e1a84ddfd7655..e61e34a176b03 100644
--- a/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
+++ b/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
@@ -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: ...

diff  --git a/mlir/test/python/ir/location.py b/mlir/test/python/ir/location.py
index 7bc7b28ae6a4a..1c13c4870cbe0 100644
--- a/mlir/test/python/ir/location.py
+++ b/mlir/test/python/ir/location.py
@@ -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: