From 1ab3efac415e0d59f9686be482ea522ff0b92f99 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar@google.com>
Date: Sat, 11 Dec 2021 10:12:29 -0800
Subject: [PATCH] [mlir][python] Add fused location

---
 mlir/lib/Bindings/Python/IRCore.cpp      | 20 ++++++++++++++++++++
 mlir/python/mlir/_mlir_libs/_mlir/ir.pyi |  2 ++
 mlir/test/python/ir/location.py          | 21 +++++++++++++++++++++
 3 files changed, 43 insertions(+)

diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index cd57552..3640a15 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)";
 
@@ -2198,6 +2201,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,
              DefaultingPyMlirContext context) {
diff --git a/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi b/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
index e1a84dd..e61e34a 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 7bc7b28..1c13c48 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:
-- 
2.7.4