[mlir][CAPI][python] bind CallSiteLoc, FileLineColRange, FusedLoc, NameLoc, OpaqueLoc

Author: makslevental

mlir/include/mlir-c/IR.h

@@ -261,22 +261,96 @@ MLIR_CAPI_EXPORTED MlirLocation mlirLocationFileLineColRangeGet(
     MlirContext context, MlirStringRef filename, unsigned start_line,
     unsigned start_col, unsigned end_line, unsigned end_col);
 
+/// Getter for filename of FileLineColRange.
+MLIR_CAPI_EXPORTED MlirIdentifier
+mlirLocationFileLineColRangeGetFilename(MlirLocation location);
+
+/// Getter for start_line of FileLineColRange.
+MLIR_CAPI_EXPORTED unsigned
+mlirLocationFileLineColRangeGetStartLine(MlirLocation location);
+
+/// Getter for start_column of FileLineColRange.
+MLIR_CAPI_EXPORTED unsigned
+mlirLocationFileLineColRangeGetStartColumn(MlirLocation location);
+
+/// Getter for end_line of FileLineColRange.
+MLIR_CAPI_EXPORTED unsigned
+mlirLocationFileLineColRangeGetEndLine(MlirLocation location);
+
+/// Getter for end_column of FileLineColRange.
+MLIR_CAPI_EXPORTED unsigned
+mlirLocationFileLineColRangeGetEndColumn(MlirLocation location);
+
+/// TypeID Getter for FileLineColRange.
+MLIR_CAPI_EXPORTED MlirTypeID mlirLocationFileLineColRangeGetTypeID(void);
+
+/// Checks whether the given location is an FileLineColRange.
+MLIR_CAPI_EXPORTED bool mlirLocationIsAFileLineColRange(MlirLocation location);
+
 /// Creates a call site location with a callee and a caller.
 MLIR_CAPI_EXPORTED MlirLocation mlirLocationCallSiteGet(MlirLocation callee,
                                                         MlirLocation caller);
 
+/// Getter for callee of CallSite.
+MLIR_CAPI_EXPORTED MlirLocation
+mlirLocationCallSiteGetCallee(MlirLocation location);
+
+/// Getter for caller of CallSite.
+MLIR_CAPI_EXPORTED MlirLocation
+mlirLocationCallSiteGetCaller(MlirLocation location);
+
+/// TypeID Getter for CallSite.
+MLIR_CAPI_EXPORTED MlirTypeID mlirLocationCallSiteGetTypeID(void);
+
+/// Checks whether the given location is an CallSite.
+MLIR_CAPI_EXPORTED bool mlirLocationIsACallSite(MlirLocation location);
+
 /// Creates a fused location with an array of locations and metadata.
 MLIR_CAPI_EXPORTED MlirLocation
 mlirLocationFusedGet(MlirContext ctx, intptr_t nLocations,
                      MlirLocation const *locations, MlirAttribute metadata);
 
+/// Getter for number of locations fused together.
+MLIR_CAPI_EXPORTED unsigned
+mlirLocationFusedGetNumLocations(MlirLocation location);
+
+/// Getter for locations of Fused. Requires pre-allocated memory of
+/// #fusedLocations X sizeof(MlirLocation).
+MLIR_CAPI_EXPORTED void
+mlirLocationFusedGetLocations(MlirLocation location,
+                              MlirLocation *locationsCPtr);
+
+/// Getter for metadata of Fused.
+MLIR_CAPI_EXPORTED MlirAttribute
+mlirLocationFusedGetMetadata(MlirLocation location);
+
+/// TypeID Getter for Fused.
+MLIR_CAPI_EXPORTED MlirTypeID mlirLocationFusedGetTypeID(void);
+
+/// Checks whether the given location is an Fused.
+MLIR_CAPI_EXPORTED bool mlirLocationIsAFused(MlirLocation location);
+
 /// Creates a name location owned by the given context. Providing null location
 /// for childLoc is allowed and if childLoc is null location, then the behavior
 /// is the same as having unknown child location.
 MLIR_CAPI_EXPORTED MlirLocation mlirLocationNameGet(MlirContext context,
                                                     MlirStringRef name,
                                                     MlirLocation childLoc);
 
+/// Getter for name of Name.
+MLIR_CAPI_EXPORTED MlirIdentifier
+mlirLocationNameGetName(MlirLocation location);
+
+/// Getter for childLoc of Name.
+MLIR_CAPI_EXPORTED MlirLocation
+mlirLocationNameGetChildLoc(MlirLocation location);
+
+/// TypeID Getter for Name.
+MLIR_CAPI_EXPORTED MlirTypeID mlirLocationNameGetTypeID(void);
+
+/// Checks whether the given location is an Name.
+MLIR_CAPI_EXPORTED bool mlirLocationIsAName(MlirLocation location);
+
 /// Creates a location with unknown position owned by the given context.
 MLIR_CAPI_EXPORTED MlirLocation mlirLocationUnknownGet(MlirContext context);
 

mlir/lib/Bindings/Python/IRCore.cpp

@@ -12,14 +12,14 @@
 #include "Globals.h"
 #include "IRModule.h"
 #include "NanobindUtils.h"
+#include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind.
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/Debug.h"
 #include "mlir-c/Diagnostics.h"
 #include "mlir-c/IR.h"
 #include "mlir-c/Support.h"
 #include "mlir/Bindings/Python/Nanobind.h"
 #include "mlir/Bindings/Python/NanobindAdaptors.h"
-#include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind.
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
 
@@ -299,7 +299,7 @@ struct PyAttrBuilderMap {
     return *builder;
   }
   static void dunderSetItemNamed(const std::string &attributeKind,
-                                nb::callable func, bool replace) {
+                                 nb::callable func, bool replace) {
     PyGlobals::get().registerAttributeBuilder(attributeKind, std::move(func),
                                               replace);
   }
@@ -2933,6 +2933,9 @@ void mlir::python::populateIRCore(nb::module_ &m) {
           nb::arg("callee"), nb::arg("frames"),
           nb::arg("context").none() = nb::none(),
           kContextGetCallSiteLocationDocstring)
+      .def("is_a_callsite", mlirLocationIsACallSite)
+      .def("get_callee", mlirLocationCallSiteGetCallee)
+      .def("get_caller", mlirLocationCallSiteGetCaller)
       .def_static(
           "file",
           [](std::string filename, int line, int col,
@@ -2957,6 +2960,12 @@ void mlir::python::populateIRCore(nb::module_ &m) {
           nb::arg("filename"), nb::arg("start_line"), nb::arg("start_col"),
           nb::arg("end_line"), nb::arg("end_col"),
           nb::arg("context").none() = nb::none(), kContextGetFileRangeDocstring)
+      .def("is_a_file", mlirLocationIsAFileLineColRange)
+      .def("get_file_name", mlirLocationFileLineColRangeGetFilename)
+      .def("get_file_start_line", mlirLocationFileLineColRangeGetStartLine)
+      .def("get_file_start_col", mlirLocationFileLineColRangeGetStartColumn)
+      .def("get_file_end_line", mlirLocationFileLineColRangeGetStartLine)
+      .def("get_file_end_col", mlirLocationFileLineColRangeGetStartColumn)
       .def_static(
           "fused",
           [](const std::vector<PyLocation> &pyLocations,
@@ -2974,6 +2983,14 @@ void mlir::python::populateIRCore(nb::module_ &m) {
           nb::arg("locations"), nb::arg("metadata").none() = nb::none(),
           nb::arg("context").none() = nb::none(),
           kContextGetFusedLocationDocstring)
+      .def("is_a_fused", mlirLocationIsAFused)
+      .def("get_fused_locations",
+           [](MlirLocation loc) {
+             unsigned numLocations = mlirLocationFusedGetNumLocations(loc);
+             std::vector<MlirLocation> locations(numLocations);
+             mlirLocationFusedGetLocations(loc, locations.data());
+             return locations;
+           })
       .def_static(
           "name",
           [](std::string name, std::optional<PyLocation> childLoc,
@@ -2988,6 +3005,9 @@ void mlir::python::populateIRCore(nb::module_ &m) {
           nb::arg("name"), nb::arg("childLoc").none() = nb::none(),
           nb::arg("context").none() = nb::none(),
           kContextGetNameLocationDocString)
+      .def("is_a_name", mlirLocationIsAName)
+      .def("get_name_str", mlirLocationNameGetName)
+      .def("get_name_child_loc", mlirLocationNameGetChildLoc)
       .def_static(
           "from_attr",
           [](PyAttribute &attribute, DefaultingPyMlirContext context) {

mlir/lib/CAPI/IR/IR.cpp

@@ -273,10 +273,54 @@ mlirLocationFileLineColRangeGet(MlirContext context, MlirStringRef filename,
                                      startLine, startCol, endLine, endCol)));
 }
 
+MlirIdentifier mlirLocationFileLineColRangeGetFilename(MlirLocation location) {
+  return wrap((llvm::cast<FileLineColRange>(unwrap(location)).getFilename()));
+}
+
+unsigned mlirLocationFileLineColRangeGetStartLine(MlirLocation location) {
+  return llvm::cast<FileLineColRange>(unwrap(location)).getStartLine();
+}
+
+unsigned mlirLocationFileLineColRangeGetStartColumn(MlirLocation location) {
+  return llvm::cast<FileLineColRange>(unwrap(location)).getStartColumn();
+}
+
+unsigned mlirLocationFileLineColRangeGetEndLine(MlirLocation location) {
+  return llvm::cast<FileLineColRange>(unwrap(location)).getEndLine();
+}
+
+unsigned mlirLocationFileLineColRangeGetEndColumn(MlirLocation location) {
+  return llvm::cast<FileLineColRange>(unwrap(location)).getEndColumn();
+}
+
+MlirTypeID mlirLocationFileLineColRangeGetTypeID() {
+  return wrap(FileLineColRange::getTypeID());
+}
+
+bool mlirLocationIsAFileLineColRange(MlirLocation location) {
+  return isa<FileLineColRange>(unwrap(location));
+}
+
 MlirLocation mlirLocationCallSiteGet(MlirLocation callee, MlirLocation caller) {
   return wrap(Location(CallSiteLoc::get(unwrap(callee), unwrap(caller))));
 }
 
+MlirLocation mlirLocationCallSiteGetCallee(MlirLocation location) {
+  return wrap(Location(llvm::cast<CallSiteLoc>(unwrap(location)).getCallee()));
+}
+
+MlirLocation mlirLocationCallSiteGetCaller(MlirLocation location) {
+  return wrap(Location(llvm::cast<CallSiteLoc>(unwrap(location)).getCaller()));
+}
+
+MlirTypeID mlirLocationCallSiteGetTypeID() {
+  return wrap(CallSiteLoc::getTypeID());
+}
+
+bool mlirLocationIsACallSite(MlirLocation location) {
+  return isa<CallSiteLoc>(unwrap(location));
+}
+
 MlirLocation mlirLocationFusedGet(MlirContext ctx, intptr_t nLocations,
                                   MlirLocation const *locations,
                                   MlirAttribute metadata) {
@@ -285,6 +329,30 @@ MlirLocation mlirLocationFusedGet(MlirContext ctx, intptr_t nLocations,
   return wrap(FusedLoc::get(unwrappedLocs, unwrap(metadata), unwrap(ctx)));
 }
 
+unsigned mlirLocationFusedGetNumLocations(MlirLocation location) {
+  llvm::ArrayRef<Location> locationsArrRef =
+      llvm::cast<FusedLoc>(unwrap(location)).getLocations();
+  return locationsArrRef.size();
+}
+
+void mlirLocationFusedGetLocations(MlirLocation location,
+                                   MlirLocation *locationsCPtr) {
+  llvm::ArrayRef<Location> locationsArrRef =
+      llvm::cast<FusedLoc>(unwrap(location)).getLocations();
+  for (auto [i, location] : llvm::enumerate(locationsArrRef))
+    locationsCPtr[i] = wrap(location);
+}
+
+MlirAttribute mlirLocationFusedGetMetadata(MlirLocation location) {
+  return wrap(llvm::cast<FusedLoc>(unwrap(location)).getMetadata());
+}
+
+MlirTypeID mlirLocationFusedGetTypeID() { return wrap(FusedLoc::getTypeID()); }
+
+bool mlirLocationIsAFused(MlirLocation location) {
+  return isa<FusedLoc>(unwrap(location));
+}
+
 MlirLocation mlirLocationNameGet(MlirContext context, MlirStringRef name,
                                  MlirLocation childLoc) {
   if (mlirLocationIsNull(childLoc))
@@ -294,6 +362,20 @@ MlirLocation mlirLocationNameGet(MlirContext context, MlirStringRef name,
       StringAttr::get(unwrap(context), unwrap(name)), unwrap(childLoc))));
 }
 
+MlirIdentifier mlirLocationNameGetName(MlirLocation location) {
+  return wrap((llvm::cast<NameLoc>(unwrap(location)).getName()));
+}
+
+MlirLocation mlirLocationNameGetChildLoc(MlirLocation location) {
+  return wrap(Location(llvm::cast<NameLoc>(unwrap(location)).getChildLoc()));
+}
+
+MlirTypeID mlirLocationNameGetTypeID() { return wrap(NameLoc::getTypeID()); }
+
+bool mlirLocationIsAName(MlirLocation location) {
+  return isa<NameLoc>(unwrap(location));
+}
+
 MlirLocation mlirLocationUnknownGet(MlirContext context) {
   return wrap(Location(UnknownLoc::get(unwrap(context))));
 }