diff options
Diffstat (limited to 'mlir/lib/Bindings/Python')
20 files changed, 468 insertions, 336 deletions
diff --git a/mlir/lib/Bindings/Python/DialectGPU.cpp b/mlir/lib/Bindings/Python/DialectGPU.cpp index e5045cf..a21176fff 100644 --- a/mlir/lib/Bindings/Python/DialectGPU.cpp +++ b/mlir/lib/Bindings/Python/DialectGPU.cpp @@ -9,8 +9,8 @@ #include "mlir-c/Dialect/GPU.h" #include "mlir-c/IR.h" #include "mlir-c/Support.h" -#include "mlir/Bindings/Python/NanobindAdaptors.h" #include "mlir/Bindings/Python/Nanobind.h" +#include "mlir/Bindings/Python/NanobindAdaptors.h" namespace nb = nanobind; using namespace nanobind::literals; @@ -34,7 +34,7 @@ NB_MODULE(_mlirDialectsGPU, m) { mlirGPUAsyncTokenType.def_classmethod( "get", - [](nb::object cls, MlirContext ctx) { + [](const nb::object &cls, MlirContext ctx) { return cls(mlirGPUAsyncTokenTypeGet(ctx)); }, "Gets an instance of AsyncTokenType in the same context", nb::arg("cls"), @@ -47,8 +47,9 @@ NB_MODULE(_mlirDialectsGPU, m) { mlir_attribute_subclass(m, "ObjectAttr", mlirAttributeIsAGPUObjectAttr) .def_classmethod( "get", - [](nb::object cls, MlirAttribute target, uint32_t format, - nb::bytes object, std::optional<MlirAttribute> mlirObjectProps, + [](const nb::object &cls, MlirAttribute target, uint32_t format, + const nb::bytes &object, + std::optional<MlirAttribute> mlirObjectProps, std::optional<MlirAttribute> mlirKernelsAttr) { MlirStringRef objectStrRef = mlirStringRefCreate( static_cast<char *>(const_cast<void *>(object.data())), diff --git a/mlir/lib/Bindings/Python/DialectLLVM.cpp b/mlir/lib/Bindings/Python/DialectLLVM.cpp index f211e76..ee106c0 100644 --- a/mlir/lib/Bindings/Python/DialectLLVM.cpp +++ b/mlir/lib/Bindings/Python/DialectLLVM.cpp @@ -12,8 +12,8 @@ #include "mlir-c/IR.h" #include "mlir-c/Support.h" #include "mlir/Bindings/Python/Diagnostics.h" -#include "mlir/Bindings/Python/NanobindAdaptors.h" #include "mlir/Bindings/Python/Nanobind.h" +#include "mlir/Bindings/Python/NanobindAdaptors.h" namespace nb = nanobind; @@ -24,7 +24,7 @@ using namespace mlir; using namespace mlir::python; using namespace mlir::python::nanobind_adaptors; -void populateDialectLLVMSubmodule(const nanobind::module_ &m) { +static void populateDialectLLVMSubmodule(const nanobind::module_ &m) { //===--------------------------------------------------------------------===// // StructType @@ -35,8 +35,8 @@ void populateDialectLLVMSubmodule(const nanobind::module_ &m) { llvmStructType.def_classmethod( "get_literal", - [](nb::object cls, const std::vector<MlirType> &elements, bool packed, - MlirLocation loc) { + [](const nb::object &cls, const std::vector<MlirType> &elements, + bool packed, MlirLocation loc) { CollectDiagnosticsToStringScope scope(mlirLocationGetContext(loc)); MlirType type = mlirLLVMStructTypeLiteralGetChecked( @@ -51,7 +51,7 @@ void populateDialectLLVMSubmodule(const nanobind::module_ &m) { llvmStructType.def_classmethod( "get_identified", - [](nb::object cls, const std::string &name, MlirContext context) { + [](const nb::object &cls, const std::string &name, MlirContext context) { return cls(mlirLLVMStructTypeIdentifiedGet( context, mlirStringRefCreate(name.data(), name.size()))); }, @@ -59,7 +59,7 @@ void populateDialectLLVMSubmodule(const nanobind::module_ &m) { llvmStructType.def_classmethod( "get_opaque", - [](nb::object cls, const std::string &name, MlirContext context) { + [](const nb::object &cls, const std::string &name, MlirContext context) { return cls(mlirLLVMStructTypeOpaqueGet( context, mlirStringRefCreate(name.data(), name.size()))); }, @@ -79,7 +79,7 @@ void populateDialectLLVMSubmodule(const nanobind::module_ &m) { llvmStructType.def_classmethod( "new_identified", - [](nb::object cls, const std::string &name, + [](const nb::object &cls, const std::string &name, const std::vector<MlirType> &elements, bool packed, MlirContext ctx) { return cls(mlirLLVMStructTypeIdentifiedNewGet( ctx, mlirStringRefCreate(name.data(), name.length()), @@ -123,7 +123,7 @@ void populateDialectLLVMSubmodule(const nanobind::module_ &m) { mlir_type_subclass(m, "PointerType", mlirTypeIsALLVMPointerType) .def_classmethod( "get", - [](nb::object cls, std::optional<unsigned> addressSpace, + [](const nb::object &cls, std::optional<unsigned> addressSpace, MlirContext context) { CollectDiagnosticsToStringScope scope(context); MlirType type = mlirLLVMPointerTypeGet( diff --git a/mlir/lib/Bindings/Python/DialectNVGPU.cpp b/mlir/lib/Bindings/Python/DialectNVGPU.cpp index a0d6a4b..bb3f519c 100644 --- a/mlir/lib/Bindings/Python/DialectNVGPU.cpp +++ b/mlir/lib/Bindings/Python/DialectNVGPU.cpp @@ -8,8 +8,8 @@ #include "mlir-c/Dialect/NVGPU.h" #include "mlir-c/IR.h" -#include "mlir/Bindings/Python/NanobindAdaptors.h" #include "mlir/Bindings/Python/Nanobind.h" +#include "mlir/Bindings/Python/NanobindAdaptors.h" namespace nb = nanobind; using namespace llvm; @@ -23,8 +23,8 @@ static void populateDialectNVGPUSubmodule(const nb::module_ &m) { nvgpuTensorMapDescriptorType.def_classmethod( "get", - [](nb::object cls, MlirType tensorMemrefType, int swizzle, int l2promo, - int oobFill, int interleave, MlirContext ctx) { + [](const nb::object &cls, MlirType tensorMemrefType, int swizzle, + int l2promo, int oobFill, int interleave, MlirContext ctx) { return cls(mlirNVGPUTensorMapDescriptorTypeGet( ctx, tensorMemrefType, swizzle, l2promo, oobFill, interleave)); }, diff --git a/mlir/lib/Bindings/Python/DialectPDL.cpp b/mlir/lib/Bindings/Python/DialectPDL.cpp index bcc6ff4..2acedbc 100644 --- a/mlir/lib/Bindings/Python/DialectPDL.cpp +++ b/mlir/lib/Bindings/Python/DialectPDL.cpp @@ -8,8 +8,8 @@ #include "mlir-c/Dialect/PDL.h" #include "mlir-c/IR.h" -#include "mlir/Bindings/Python/NanobindAdaptors.h" #include "mlir/Bindings/Python/Nanobind.h" +#include "mlir/Bindings/Python/NanobindAdaptors.h" namespace nb = nanobind; using namespace llvm; @@ -17,7 +17,7 @@ using namespace mlir; using namespace mlir::python; using namespace mlir::python::nanobind_adaptors; -void populateDialectPDLSubmodule(const nanobind::module_ &m) { +static void populateDialectPDLSubmodule(const nanobind::module_ &m) { //===-------------------------------------------------------------------===// // PDLType //===-------------------------------------------------------------------===// @@ -32,7 +32,7 @@ void populateDialectPDLSubmodule(const nanobind::module_ &m) { mlir_type_subclass(m, "AttributeType", mlirTypeIsAPDLAttributeType); attributeType.def_classmethod( "get", - [](nb::object cls, MlirContext ctx) { + [](const nb::object &cls, MlirContext ctx) { return cls(mlirPDLAttributeTypeGet(ctx)); }, "Get an instance of AttributeType in given context.", nb::arg("cls"), @@ -46,7 +46,7 @@ void populateDialectPDLSubmodule(const nanobind::module_ &m) { mlir_type_subclass(m, "OperationType", mlirTypeIsAPDLOperationType); operationType.def_classmethod( "get", - [](nb::object cls, MlirContext ctx) { + [](const nb::object &cls, MlirContext ctx) { return cls(mlirPDLOperationTypeGet(ctx)); }, "Get an instance of OperationType in given context.", nb::arg("cls"), @@ -59,7 +59,7 @@ void populateDialectPDLSubmodule(const nanobind::module_ &m) { auto rangeType = mlir_type_subclass(m, "RangeType", mlirTypeIsAPDLRangeType); rangeType.def_classmethod( "get", - [](nb::object cls, MlirType elementType) { + [](const nb::object &cls, MlirType elementType) { return cls(mlirPDLRangeTypeGet(elementType)); }, "Gets an instance of RangeType in the same context as the provided " @@ -77,7 +77,7 @@ void populateDialectPDLSubmodule(const nanobind::module_ &m) { auto typeType = mlir_type_subclass(m, "TypeType", mlirTypeIsAPDLTypeType); typeType.def_classmethod( "get", - [](nb::object cls, MlirContext ctx) { + [](const nb::object &cls, MlirContext ctx) { return cls(mlirPDLTypeTypeGet(ctx)); }, "Get an instance of TypeType in given context.", nb::arg("cls"), @@ -90,7 +90,7 @@ void populateDialectPDLSubmodule(const nanobind::module_ &m) { auto valueType = mlir_type_subclass(m, "ValueType", mlirTypeIsAPDLValueType); valueType.def_classmethod( "get", - [](nb::object cls, MlirContext ctx) { + [](const nb::object &cls, MlirContext ctx) { return cls(mlirPDLValueTypeGet(ctx)); }, "Get an instance of TypeType in given context.", nb::arg("cls"), diff --git a/mlir/lib/Bindings/Python/DialectQuant.cpp b/mlir/lib/Bindings/Python/DialectQuant.cpp index 55571cd..a5220fc 100644 --- a/mlir/lib/Bindings/Python/DialectQuant.cpp +++ b/mlir/lib/Bindings/Python/DialectQuant.cpp @@ -165,7 +165,7 @@ static void populateDialectQuantSubmodule(const nb::module_ &m) { quantizedType.get_class()); anyQuantizedType.def_classmethod( "get", - [](nb::object cls, unsigned flags, MlirType storageType, + [](const nb::object &cls, unsigned flags, MlirType storageType, MlirType expressedType, int64_t storageTypeMin, int64_t storageTypeMax) { return cls(mlirAnyQuantizedTypeGet(flags, storageType, expressedType, @@ -186,7 +186,7 @@ static void populateDialectQuantSubmodule(const nb::module_ &m) { quantizedType.get_class()); uniformQuantizedType.def_classmethod( "get", - [](nb::object cls, unsigned flags, MlirType storageType, + [](const nb::object &cls, unsigned flags, MlirType storageType, MlirType expressedType, double scale, int64_t zeroPoint, int64_t storageTypeMin, int64_t storageTypeMax) { return cls(mlirUniformQuantizedTypeGet(flags, storageType, @@ -221,7 +221,7 @@ static void populateDialectQuantSubmodule(const nb::module_ &m) { quantizedType.get_class()); uniformQuantizedPerAxisType.def_classmethod( "get", - [](nb::object cls, unsigned flags, MlirType storageType, + [](const nb::object &cls, unsigned flags, MlirType storageType, MlirType expressedType, std::vector<double> scales, std::vector<int64_t> zeroPoints, int32_t quantizedDimension, int64_t storageTypeMin, int64_t storageTypeMax) { @@ -293,7 +293,7 @@ static void populateDialectQuantSubmodule(const nb::module_ &m) { mlirTypeIsAUniformQuantizedSubChannelType, quantizedType.get_class()); uniformQuantizedSubChannelType.def_classmethod( "get", - [](nb::object cls, unsigned flags, MlirType storageType, + [](const nb::object &cls, unsigned flags, MlirType storageType, MlirType expressedType, MlirAttribute scales, MlirAttribute zeroPoints, std::vector<int32_t> quantizedDimensions, std::vector<int64_t> blockSizes, int64_t storageTypeMin, @@ -367,7 +367,8 @@ static void populateDialectQuantSubmodule(const nb::module_ &m) { quantizedType.get_class()); calibratedQuantizedType.def_classmethod( "get", - [](nb::object cls, MlirType expressedType, double min, double max) { + [](const nb::object &cls, MlirType expressedType, double min, + double max) { return cls(mlirCalibratedQuantizedTypeGet(expressedType, min, max)); }, "Gets an instance of CalibratedQuantizedType in the same context as the " diff --git a/mlir/lib/Bindings/Python/DialectSMT.cpp b/mlir/lib/Bindings/Python/DialectSMT.cpp index 4e76477..cab4219 100644 --- a/mlir/lib/Bindings/Python/DialectSMT.cpp +++ b/mlir/lib/Bindings/Python/DialectSMT.cpp @@ -24,7 +24,7 @@ using namespace mlir; using namespace mlir::python; using namespace mlir::python::nanobind_adaptors; -void populateDialectSMTSubmodule(nanobind::module_ &m) { +static void populateDialectSMTSubmodule(nanobind::module_ &m) { auto smtBoolType = mlir_type_subclass(m, "BoolType", mlirSMTTypeIsABool) .def_classmethod( diff --git a/mlir/lib/Bindings/Python/DialectSparseTensor.cpp b/mlir/lib/Bindings/Python/DialectSparseTensor.cpp index 97cebcc..9d7dc11 100644 --- a/mlir/lib/Bindings/Python/DialectSparseTensor.cpp +++ b/mlir/lib/Bindings/Python/DialectSparseTensor.cpp @@ -12,8 +12,8 @@ #include "mlir-c/AffineMap.h" #include "mlir-c/Dialect/SparseTensor.h" #include "mlir-c/IR.h" -#include "mlir/Bindings/Python/NanobindAdaptors.h" #include "mlir/Bindings/Python/Nanobind.h" +#include "mlir/Bindings/Python/NanobindAdaptors.h" namespace nb = nanobind; using namespace llvm; @@ -38,7 +38,8 @@ static void populateDialectSparseTensorSubmodule(const nb::module_ &m) { mlirAttributeIsASparseTensorEncodingAttr) .def_classmethod( "get", - [](nb::object cls, std::vector<MlirSparseTensorLevelType> lvlTypes, + [](const nb::object &cls, + std::vector<MlirSparseTensorLevelType> lvlTypes, std::optional<MlirAffineMap> dimToLvl, std::optional<MlirAffineMap> lvlToDim, int posWidth, int crdWidth, std::optional<MlirAttribute> explicitVal, @@ -58,7 +59,7 @@ static void populateDialectSparseTensorSubmodule(const nb::module_ &m) { "Gets a sparse_tensor.encoding from parameters.") .def_classmethod( "build_level_type", - [](nb::object cls, MlirSparseTensorLevelFormat lvlFmt, + [](const nb::object &cls, MlirSparseTensorLevelFormat lvlFmt, const std::vector<MlirSparseTensorLevelPropertyNondefault> &properties, unsigned n, unsigned m) { diff --git a/mlir/lib/Bindings/Python/DialectTransform.cpp b/mlir/lib/Bindings/Python/DialectTransform.cpp index 59a030a..1a62b06 100644 --- a/mlir/lib/Bindings/Python/DialectTransform.cpp +++ b/mlir/lib/Bindings/Python/DialectTransform.cpp @@ -11,15 +11,15 @@ #include "mlir-c/Dialect/Transform.h" #include "mlir-c/IR.h" #include "mlir-c/Support.h" -#include "mlir/Bindings/Python/NanobindAdaptors.h" #include "mlir/Bindings/Python/Nanobind.h" +#include "mlir/Bindings/Python/NanobindAdaptors.h" namespace nb = nanobind; using namespace mlir; using namespace mlir::python; using namespace mlir::python::nanobind_adaptors; -void populateDialectTransformSubmodule(const nb::module_ &m) { +static void populateDialectTransformSubmodule(const nb::module_ &m) { //===-------------------------------------------------------------------===// // AnyOpType //===-------------------------------------------------------------------===// @@ -29,7 +29,7 @@ void populateDialectTransformSubmodule(const nb::module_ &m) { mlirTransformAnyOpTypeGetTypeID); anyOpType.def_classmethod( "get", - [](nb::object cls, MlirContext ctx) { + [](const nb::object &cls, MlirContext ctx) { return cls(mlirTransformAnyOpTypeGet(ctx)); }, "Get an instance of AnyOpType in the given context.", nb::arg("cls"), @@ -44,7 +44,7 @@ void populateDialectTransformSubmodule(const nb::module_ &m) { mlirTransformAnyParamTypeGetTypeID); anyParamType.def_classmethod( "get", - [](nb::object cls, MlirContext ctx) { + [](const nb::object &cls, MlirContext ctx) { return cls(mlirTransformAnyParamTypeGet(ctx)); }, "Get an instance of AnyParamType in the given context.", nb::arg("cls"), @@ -59,7 +59,7 @@ void populateDialectTransformSubmodule(const nb::module_ &m) { mlirTransformAnyValueTypeGetTypeID); anyValueType.def_classmethod( "get", - [](nb::object cls, MlirContext ctx) { + [](const nb::object &cls, MlirContext ctx) { return cls(mlirTransformAnyValueTypeGet(ctx)); }, "Get an instance of AnyValueType in the given context.", nb::arg("cls"), @@ -74,7 +74,8 @@ void populateDialectTransformSubmodule(const nb::module_ &m) { mlirTransformOperationTypeGetTypeID); operationType.def_classmethod( "get", - [](nb::object cls, const std::string &operationName, MlirContext ctx) { + [](const nb::object &cls, const std::string &operationName, + MlirContext ctx) { MlirStringRef cOperationName = mlirStringRefCreate(operationName.data(), operationName.size()); return cls(mlirTransformOperationTypeGet(ctx, cOperationName)); @@ -101,7 +102,7 @@ void populateDialectTransformSubmodule(const nb::module_ &m) { mlirTransformParamTypeGetTypeID); paramType.def_classmethod( "get", - [](nb::object cls, MlirType type, MlirContext ctx) { + [](const nb::object &cls, MlirType type, MlirContext ctx) { return cls(mlirTransformParamTypeGet(ctx, type)); }, "Get an instance of ParamType for the given type in the given context.", diff --git a/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp b/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp index 81dada3..8bb493e 100644 --- a/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp +++ b/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp @@ -7,8 +7,8 @@ //===----------------------------------------------------------------------===// #include "mlir-c/ExecutionEngine.h" -#include "mlir/Bindings/Python/NanobindAdaptors.h" #include "mlir/Bindings/Python/Nanobind.h" +#include "mlir/Bindings/Python/NanobindAdaptors.h" namespace nb = nanobind; using namespace mlir; @@ -45,7 +45,7 @@ public: referencedObjects.push_back(obj); } - static nb::object createFromCapsule(nb::object capsule) { + static nb::object createFromCapsule(const nb::object &capsule) { MlirExecutionEngine rawPm = mlirPythonCapsuleToExecutionEngine(capsule.ptr()); if (mlirExecutionEngineIsNull(rawPm)) @@ -113,7 +113,7 @@ NB_MODULE(_mlirExecutionEngine, m) { .def( "raw_register_runtime", [](PyExecutionEngine &executionEngine, const std::string &name, - nb::object callbackObj) { + const nb::object &callbackObj) { executionEngine.addReferencedObject(callbackObj); uintptr_t rawSym = nb::cast<uintptr_t>(nb::getattr(callbackObj, "value")); @@ -125,6 +125,17 @@ NB_MODULE(_mlirExecutionEngine, m) { nb::arg("name"), nb::arg("callback"), "Register `callback` as the runtime symbol `name`.") .def( + "initialize", + [](PyExecutionEngine &executionEngine) { + mlirExecutionEngineInitialize(executionEngine.get()); + }, + "Initialize the ExecutionEngine. Global constructors specified by " + "`llvm.mlir.global_ctors` will be run. One common scenario is that " + "kernel binary compiled from `gpu.module` gets loaded during " + "initialization. Make sure all symbols are resolvable before " + "initialization by calling `register_runtime` or including " + "shared libraries.") + .def( "dump_to_object_file", [](PyExecutionEngine &executionEngine, const std::string &fileName) { mlirExecutionEngineDumpToObjectFile( diff --git a/mlir/lib/Bindings/Python/Globals.h b/mlir/lib/Bindings/Python/Globals.h index 826a34a..71a051c 100644 --- a/mlir/lib/Bindings/Python/Globals.h +++ b/mlir/lib/Bindings/Python/Globals.h @@ -10,15 +10,19 @@ #define MLIR_BINDINGS_PYTHON_GLOBALS_H #include <optional> +#include <regex> #include <string> +#include <unordered_set> #include <vector> #include "NanobindUtils.h" #include "mlir-c/IR.h" #include "mlir/CAPI/Support.h" #include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/StringExtras.h" #include "llvm/ADT/StringRef.h" #include "llvm/ADT/StringSet.h" +#include "llvm/Support/Regex.h" namespace mlir { namespace python { @@ -114,6 +118,39 @@ public: std::optional<nanobind::object> lookupOperationClass(llvm::StringRef operationName); + class TracebackLoc { + public: + bool locTracebacksEnabled(); + + void setLocTracebacksEnabled(bool value); + + size_t locTracebackFramesLimit(); + + void setLocTracebackFramesLimit(size_t value); + + void registerTracebackFileInclusion(const std::string &file); + + void registerTracebackFileExclusion(const std::string &file); + + bool isUserTracebackFilename(llvm::StringRef file); + + static constexpr size_t kMaxFrames = 512; + + private: + nanobind::ft_mutex mutex; + bool locTracebackEnabled_ = false; + size_t locTracebackFramesLimit_ = 10; + std::unordered_set<std::string> userTracebackIncludeFiles; + std::unordered_set<std::string> userTracebackExcludeFiles; + std::regex userTracebackIncludeRegex; + bool rebuildUserTracebackIncludeRegex = false; + std::regex userTracebackExcludeRegex; + bool rebuildUserTracebackExcludeRegex = false; + llvm::StringMap<bool> isUserTracebackFilenameCache; + }; + + TracebackLoc &getTracebackLoc() { return tracebackLoc; } + private: static PyGlobals *instance; @@ -134,6 +171,8 @@ private: /// Set of dialect namespaces that we have attempted to import implementation /// modules for. llvm::StringSet<> loadedDialectModules; + + TracebackLoc tracebackLoc; }; } // namespace python diff --git a/mlir/lib/Bindings/Python/IRAffine.cpp b/mlir/lib/Bindings/Python/IRAffine.cpp index 50f2a4f..a6499c9 100644 --- a/mlir/lib/Bindings/Python/IRAffine.cpp +++ b/mlir/lib/Bindings/Python/IRAffine.cpp @@ -17,9 +17,9 @@ #include "NanobindUtils.h" #include "mlir-c/AffineExpr.h" #include "mlir-c/AffineMap.h" +#include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind. #include "mlir-c/IntegerSet.h" #include "mlir/Bindings/Python/Nanobind.h" -#include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind. #include "mlir/Support/LLVM.h" #include "llvm/ADT/Hashing.h" #include "llvm/ADT/SmallVector.h" @@ -64,7 +64,7 @@ static void pyListToVector(const nb::list &list, } template <typename PermutationTy> -static bool isPermutation(std::vector<PermutationTy> permutation) { +static bool isPermutation(const std::vector<PermutationTy> &permutation) { llvm::SmallVector<bool, 8> seen(permutation.size(), false); for (auto val : permutation) { if (val < permutation.size()) { @@ -366,7 +366,7 @@ nb::object PyAffineExpr::getCapsule() { return nb::steal<nb::object>(mlirPythonAffineExprToCapsule(*this)); } -PyAffineExpr PyAffineExpr::createFromCapsule(nb::object capsule) { +PyAffineExpr PyAffineExpr::createFromCapsule(const nb::object &capsule) { MlirAffineExpr rawAffineExpr = mlirPythonCapsuleToAffineExpr(capsule.ptr()); if (mlirAffineExprIsNull(rawAffineExpr)) throw nb::python_error(); @@ -424,7 +424,7 @@ nb::object PyAffineMap::getCapsule() { return nb::steal<nb::object>(mlirPythonAffineMapToCapsule(*this)); } -PyAffineMap PyAffineMap::createFromCapsule(nb::object capsule) { +PyAffineMap PyAffineMap::createFromCapsule(const nb::object &capsule) { MlirAffineMap rawAffineMap = mlirPythonCapsuleToAffineMap(capsule.ptr()); if (mlirAffineMapIsNull(rawAffineMap)) throw nb::python_error(); @@ -500,7 +500,7 @@ nb::object PyIntegerSet::getCapsule() { return nb::steal<nb::object>(mlirPythonIntegerSetToCapsule(*this)); } -PyIntegerSet PyIntegerSet::createFromCapsule(nb::object capsule) { +PyIntegerSet PyIntegerSet::createFromCapsule(const nb::object &capsule) { MlirIntegerSet rawIntegerSet = mlirPythonCapsuleToIntegerSet(capsule.ptr()); if (mlirIntegerSetIsNull(rawIntegerSet)) throw nb::python_error(); @@ -708,7 +708,8 @@ void mlir::python::populateIRAffine(nb::module_ &m) { return static_cast<size_t>(llvm::hash_value(self.get().ptr)); }) .def_static("compress_unused_symbols", - [](nb::list affineMaps, DefaultingPyMlirContext context) { + [](const nb::list &affineMaps, + DefaultingPyMlirContext context) { SmallVector<MlirAffineMap> maps; pyListToVector<PyAffineMap, MlirAffineMap>( affineMaps, maps, "attempting to create an AffineMap"); @@ -734,7 +735,7 @@ void mlir::python::populateIRAffine(nb::module_ &m) { kDumpDocstring) .def_static( "get", - [](intptr_t dimCount, intptr_t symbolCount, nb::list exprs, + [](intptr_t dimCount, intptr_t symbolCount, const nb::list &exprs, DefaultingPyMlirContext context) { SmallVector<MlirAffineExpr> affineExprs; pyListToVector<PyAffineExpr, MlirAffineExpr>( @@ -869,7 +870,8 @@ void mlir::python::populateIRAffine(nb::module_ &m) { .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyIntegerSet::createFromCapsule) .def("__eq__", [](PyIntegerSet &self, PyIntegerSet &other) { return self == other; }) - .def("__eq__", [](PyIntegerSet &self, nb::object other) { return false; }) + .def("__eq__", + [](PyIntegerSet &self, const nb::object &other) { return false; }) .def("__str__", [](PyIntegerSet &self) { PyPrintAccumulator printAccum; @@ -898,7 +900,7 @@ void mlir::python::populateIRAffine(nb::module_ &m) { kDumpDocstring) .def_static( "get", - [](intptr_t numDims, intptr_t numSymbols, nb::list exprs, + [](intptr_t numDims, intptr_t numSymbols, const nb::list &exprs, std::vector<bool> eqFlags, DefaultingPyMlirContext context) { if (exprs.size() != eqFlags.size()) throw nb::value_error( @@ -934,8 +936,9 @@ void mlir::python::populateIRAffine(nb::module_ &m) { nb::arg("context").none() = nb::none()) .def( "get_replaced", - [](PyIntegerSet &self, nb::list dimExprs, nb::list symbolExprs, - intptr_t numResultDims, intptr_t numResultSymbols) { + [](PyIntegerSet &self, const nb::list &dimExprs, + const nb::list &symbolExprs, intptr_t numResultDims, + intptr_t numResultSymbols) { if (static_cast<intptr_t>(dimExprs.size()) != mlirIntegerSetGetNumDims(self)) throw nb::value_error( diff --git a/mlir/lib/Bindings/Python/IRAttributes.cpp b/mlir/lib/Bindings/Python/IRAttributes.cpp index db84ee1..f2eafa7 100644 --- a/mlir/lib/Bindings/Python/IRAttributes.cpp +++ b/mlir/lib/Bindings/Python/IRAttributes.cpp @@ -505,7 +505,7 @@ public: static void bindDerived(ClassTy &c) { c.def_static( "get", - [](nb::list attributes, DefaultingPyMlirContext context) { + [](const nb::list &attributes, DefaultingPyMlirContext context) { SmallVector<MlirAttribute> mlirAttributes; mlirAttributes.reserve(nb::len(attributes)); for (auto attribute : attributes) { @@ -530,7 +530,7 @@ public: .def("__iter__", [](const PyArrayAttribute &arr) { return PyArrayAttributeIterator(arr); }); - c.def("__add__", [](PyArrayAttribute arr, nb::list extras) { + c.def("__add__", [](PyArrayAttribute arr, const nb::list &extras) { std::vector<MlirAttribute> attributes; intptr_t numOldElements = mlirArrayAttrGetNumElements(arr); attributes.reserve(numOldElements + nb::len(extras)); @@ -708,7 +708,7 @@ public: static void bindDerived(ClassTy &c) { c.def_static( "get", - [](std::string value, DefaultingPyMlirContext context) { + [](const std::string &value, DefaultingPyMlirContext context) { MlirAttribute attr = mlirFlatSymbolRefAttrGet(context->get(), toMlirStringRef(value)); return PyFlatSymbolRefAttribute(context->getRef(), attr); @@ -736,8 +736,8 @@ public: static void bindDerived(ClassTy &c) { c.def_static( "get", - [](std::string dialectNamespace, nb_buffer buffer, PyType &type, - DefaultingPyMlirContext context) { + [](const std::string &dialectNamespace, const nb_buffer &buffer, + PyType &type, DefaultingPyMlirContext context) { const nb_buffer_info bufferInfo = buffer.request(); intptr_t bufferSize = bufferInfo.size; MlirAttribute attr = mlirOpaqueAttrGet( @@ -775,7 +775,7 @@ public: static void bindDerived(ClassTy &c) { c.def_static( "get", - [](std::string value, DefaultingPyMlirContext context) { + [](const std::string &value, DefaultingPyMlirContext context) { MlirAttribute attr = mlirStringAttrGet(context->get(), toMlirStringRef(value)); return PyStringAttribute(context->getRef(), attr); @@ -784,7 +784,7 @@ public: "Gets a uniqued string attribute"); c.def_static( "get", - [](nb::bytes value, DefaultingPyMlirContext context) { + [](const nb::bytes &value, DefaultingPyMlirContext context) { MlirAttribute attr = mlirStringAttrGet(context->get(), toMlirStringRef(value)); return PyStringAttribute(context->getRef(), attr); @@ -793,7 +793,7 @@ public: "Gets a uniqued string attribute"); c.def_static( "get_typed", - [](PyType &type, std::string value) { + [](PyType &type, const std::string &value) { MlirAttribute attr = mlirStringAttrTypedGet(type, toMlirStringRef(value)); return PyStringAttribute(type.getContext(), attr); @@ -826,7 +826,7 @@ public: using PyConcreteAttribute::PyConcreteAttribute; static PyDenseElementsAttribute - getFromList(nb::list attributes, std::optional<PyType> explicitType, + getFromList(const nb::list &attributes, std::optional<PyType> explicitType, DefaultingPyMlirContext contextWrapper) { const size_t numAttributes = nb::len(attributes); if (numAttributes == 0) @@ -878,8 +878,8 @@ public: } static PyDenseElementsAttribute - getFromBuffer(nb_buffer array, bool signless, - std::optional<PyType> explicitType, + getFromBuffer(const nb_buffer &array, bool signless, + const std::optional<PyType> &explicitType, std::optional<std::vector<int64_t>> explicitShape, DefaultingPyMlirContext contextWrapper) { // Request a contiguous view. In exotic cases, this will cause a copy. @@ -894,8 +894,8 @@ public: auto freeBuffer = llvm::make_scope_exit([&]() { PyBuffer_Release(&view); }); MlirContext context = contextWrapper->get(); - MlirAttribute attr = getAttributeFromBuffer(view, signless, explicitType, - explicitShape, context); + MlirAttribute attr = getAttributeFromBuffer( + view, signless, explicitType, std::move(explicitShape), context); if (mlirAttributeIsNull(attr)) { throw std::invalid_argument( "DenseElementsAttr could not be constructed from the given buffer. " @@ -1092,16 +1092,16 @@ private: "when the type is not a shaped type."); } return *bulkLoadElementType; - } else { - MlirAttribute encodingAttr = mlirAttributeGetNull(); - return mlirRankedTensorTypeGet(shape.size(), shape.data(), - *bulkLoadElementType, encodingAttr); } + MlirAttribute encodingAttr = mlirAttributeGetNull(); + return mlirRankedTensorTypeGet(shape.size(), shape.data(), + *bulkLoadElementType, encodingAttr); } static MlirAttribute getAttributeFromBuffer( Py_buffer &view, bool signless, std::optional<PyType> explicitType, - std::optional<std::vector<int64_t>> explicitShape, MlirContext &context) { + const std::optional<std::vector<int64_t>> &explicitShape, + MlirContext &context) { // Detect format codes that are suitable for bulk loading. This includes // all byte aligned integer and floating point types up to 8 bytes. // Notably, this excludes exotics types which do not have a direct @@ -1125,7 +1125,7 @@ private: bulkLoadElementType = mlirF16TypeGet(context); } else if (format == "?") { // i1 - // The i1 type needs to be bit-packed, so we will handle it seperately + // The i1 type needs to be bit-packed, so we will handle it separately return getBitpackedAttributeFromBooleanBuffer(view, explicitShape, context); } else if (isSignedIntegerFormat(format)) { @@ -1205,8 +1205,8 @@ private: packbitsFunc(nb::cast(unpackedArray), "bitorder"_a = "little"); nb_buffer_info pythonBuffer = nb::cast<nb_buffer>(packedBooleans).request(); - MlirType bitpackedType = - getShapedType(mlirIntegerTypeGet(context, 1), explicitShape, view); + MlirType bitpackedType = getShapedType(mlirIntegerTypeGet(context, 1), + std::move(explicitShape), view); assert(pythonBuffer.itemsize == 1 && "Packbits must return uint8"); // Notice that `mlirDenseElementsAttrRawBufferGet` copies the memory of // packedBooleans, hence the MlirAttribute will remain valid even when @@ -1443,9 +1443,9 @@ public: using PyConcreteAttribute::PyConcreteAttribute; static PyDenseResourceElementsAttribute - getFromBuffer(nb_buffer buffer, const std::string &name, const PyType &type, - std::optional<size_t> alignment, bool isMutable, - DefaultingPyMlirContext contextWrapper) { + getFromBuffer(const nb_buffer &buffer, const std::string &name, + const PyType &type, std::optional<size_t> alignment, + bool isMutable, DefaultingPyMlirContext contextWrapper) { if (!mlirTypeIsAShaped(type)) { throw std::invalid_argument( "Constructing a DenseResourceElementsAttr requires a ShapedType."); @@ -1534,7 +1534,7 @@ public: c.def("__len__", &PyDictAttribute::dunderLen); c.def_static( "get", - [](nb::dict attributes, DefaultingPyMlirContext context) { + [](const nb::dict &attributes, DefaultingPyMlirContext context) { SmallVector<MlirNamedAttribute> mlirNamedAttributes; mlirNamedAttributes.reserve(attributes.size()); for (std::pair<nb::handle, nb::handle> it : attributes) { @@ -1618,7 +1618,7 @@ public: static void bindDerived(ClassTy &c) { c.def_static( "get", - [](PyType value, DefaultingPyMlirContext context) { + [](const PyType &value, DefaultingPyMlirContext context) { MlirAttribute attr = mlirTypeAttrGet(value.get()); return PyTypeAttribute(context->getRef(), attr); }, @@ -1663,7 +1663,7 @@ public: static void bindDerived(ClassTy &c) { c.def_static( "get", - [](int64_t offset, const std::vector<int64_t> strides, + [](int64_t offset, const std::vector<int64_t> &strides, DefaultingPyMlirContext ctx) { MlirAttribute attr = mlirStridedLayoutAttrGet( ctx->get(), offset, strides.size(), strides.data()); diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp index 5feed95..2df2a73 100644 --- a/mlir/lib/Bindings/Python/IRCore.cpp +++ b/mlir/lib/Bindings/Python/IRCore.cpp @@ -20,11 +20,8 @@ #include "nanobind/nanobind.h" #include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/SmallVector.h" -#include "llvm/Support/raw_ostream.h" #include <optional> -#include <system_error> -#include <utility> namespace nb = nanobind; using namespace nb::literals; @@ -70,6 +67,12 @@ Returns a new MlirModule or raises an MLIRError if the parsing fails. See also: https://mlir.llvm.org/docs/LangRef/ )"; +static const char kModuleCAPICreate[] = + R"(Creates a Module from a MlirModule wrapped by a capsule (i.e. module._CAPIPtr). +Note this returns a new object BUT _clear_mlir_module(module) must be called to +prevent double-frees (of the underlying mlir::Module). +)"; + static const char kOperationCreateDocstring[] = R"(Creates a new operation. @@ -199,7 +202,7 @@ operations. /// Helper for creating an @classmethod. template <class Func, typename... Args> -nb::object classmethod(Func f, Args... args) { +static nb::object classmethod(Func f, Args... args) { nb::object cf = nb::cpp_function(f, args...); return nb::borrow<nb::object>((PyClassMethod_New(cf.ptr()))); } @@ -705,84 +708,6 @@ size_t PyMlirContext::getLiveCount() { return getLiveContexts().size(); } -size_t PyMlirContext::getLiveOperationCount() { - nb::ft_lock_guard lock(liveOperationsMutex); - return liveOperations.size(); -} - -std::vector<PyOperation *> PyMlirContext::getLiveOperationObjects() { - std::vector<PyOperation *> liveObjects; - nb::ft_lock_guard lock(liveOperationsMutex); - for (auto &entry : liveOperations) - liveObjects.push_back(entry.second.second); - return liveObjects; -} - -size_t PyMlirContext::clearLiveOperations() { - - LiveOperationMap operations; - { - nb::ft_lock_guard lock(liveOperationsMutex); - std::swap(operations, liveOperations); - } - for (auto &op : operations) - op.second.second->setInvalid(); - size_t numInvalidated = operations.size(); - return numInvalidated; -} - -void PyMlirContext::clearOperation(MlirOperation op) { - PyOperation *py_op; - { - nb::ft_lock_guard lock(liveOperationsMutex); - auto it = liveOperations.find(op.ptr); - if (it == liveOperations.end()) { - return; - } - py_op = it->second.second; - liveOperations.erase(it); - } - py_op->setInvalid(); -} - -void PyMlirContext::clearOperationsInside(PyOperationBase &op) { - typedef struct { - PyOperation &rootOp; - bool rootSeen; - } callBackData; - callBackData data{op.getOperation(), false}; - // Mark all ops below the op that the passmanager will be rooted - // at (but not op itself - note the preorder) as invalid. - MlirOperationWalkCallback invalidatingCallback = [](MlirOperation op, - void *userData) { - callBackData *data = static_cast<callBackData *>(userData); - if (LLVM_LIKELY(data->rootSeen)) - data->rootOp.getOperation().getContext()->clearOperation(op); - else - data->rootSeen = true; - return MlirWalkResult::MlirWalkResultAdvance; - }; - mlirOperationWalk(op.getOperation(), invalidatingCallback, - static_cast<void *>(&data), MlirWalkPreOrder); -} -void PyMlirContext::clearOperationsInside(MlirOperation op) { - PyOperationRef opRef = PyOperation::forOperation(getRef(), op); - clearOperationsInside(opRef->getOperation()); -} - -void PyMlirContext::clearOperationAndInside(PyOperationBase &op) { - MlirOperationWalkCallback invalidatingCallback = [](MlirOperation op, - void *userData) { - PyMlirContextRef &contextRef = *static_cast<PyMlirContextRef *>(userData); - contextRef->clearOperation(op); - return MlirWalkResult::MlirWalkResultAdvance; - }; - mlirOperationWalk(op.getOperation(), invalidatingCallback, - &op.getOperation().getContext(), MlirWalkPreOrder); -} - -size_t PyMlirContext::getLiveModuleCount() { return liveModules.size(); } - nb::object PyMlirContext::contextEnter(nb::object context) { return PyThreadContextEntry::pushContext(context); } @@ -1154,38 +1079,23 @@ PyLocation &DefaultingPyLocation::resolve() { PyModule::PyModule(PyMlirContextRef contextRef, MlirModule module) : BaseContextObject(std::move(contextRef)), module(module) {} -PyModule::~PyModule() { - nb::gil_scoped_acquire acquire; - auto &liveModules = getContext()->liveModules; - assert(liveModules.count(module.ptr) == 1 && - "destroying module not in live map"); - liveModules.erase(module.ptr); - mlirModuleDestroy(module); -} +PyModule::~PyModule() { mlirModuleDestroy(module); } PyModuleRef PyModule::forModule(MlirModule module) { MlirContext context = mlirModuleGetContext(module); PyMlirContextRef contextRef = PyMlirContext::forContext(context); - nb::gil_scoped_acquire acquire; - auto &liveModules = contextRef->liveModules; - auto it = liveModules.find(module.ptr); - if (it == liveModules.end()) { - // Create. - PyModule *unownedModule = new PyModule(std::move(contextRef), module); - // Note that the default return value policy on cast is automatic_reference, - // which does not take ownership (delete will not be called). - // Just be explicit. - nb::object pyRef = nb::cast(unownedModule, nb::rv_policy::take_ownership); - unownedModule->handle = pyRef; - liveModules[module.ptr] = - std::make_pair(unownedModule->handle, unownedModule); - return PyModuleRef(unownedModule, std::move(pyRef)); - } - // Use existing. - PyModule *existing = it->second.second; - nb::object pyRef = nb::borrow<nb::object>(it->second.first); - return PyModuleRef(existing, std::move(pyRef)); + // Create. + PyModule *unownedModule = new PyModule(std::move(contextRef), module); + // Note that the default return value policy on cast is `automatic_reference`, + // which means "does not take ownership, does not call delete/dtor". + // We use `take_ownership`, which means "Python will call the C++ destructor + // and delete operator when the Python wrapper is garbage collected", because + // MlirModule actually wraps OwningOpRef<ModuleOp> (see mlirModuleCreateParse + // etc). + nb::object pyRef = nb::cast(unownedModule, nb::rv_policy::take_ownership); + unownedModule->handle = pyRef; + return PyModuleRef(unownedModule, std::move(pyRef)); } nb::object PyModule::createFromCapsule(nb::object capsule) { @@ -1210,15 +1120,11 @@ PyOperation::~PyOperation() { // If the operation has already been invalidated there is nothing to do. if (!valid) return; - - // Otherwise, invalidate the operation and remove it from live map when it is - // attached. - if (isAttached()) { - getContext()->clearOperation(*this); - } else { - // And destroy it when it is detached, i.e. owned by Python, in which case - // all nested operations must be invalidated at removed from the live map as - // well. + // Otherwise, invalidate the operation when it is attached. + if (isAttached()) + setInvalid(); + else { + // And destroy it when it is detached, i.e. owned by Python. erase(); } } @@ -1255,35 +1161,15 @@ PyOperationRef PyOperation::createInstance(PyMlirContextRef contextRef, PyOperationRef PyOperation::forOperation(PyMlirContextRef contextRef, MlirOperation operation, nb::object parentKeepAlive) { - nb::ft_lock_guard lock(contextRef->liveOperationsMutex); - auto &liveOperations = contextRef->liveOperations; - auto it = liveOperations.find(operation.ptr); - if (it == liveOperations.end()) { - // Create. - PyOperationRef result = createInstance(std::move(contextRef), operation, - std::move(parentKeepAlive)); - liveOperations[operation.ptr] = - std::make_pair(result.getObject(), result.get()); - return result; - } - // Use existing. - PyOperation *existing = it->second.second; - nb::object pyRef = nb::borrow<nb::object>(it->second.first); - return PyOperationRef(existing, std::move(pyRef)); + return createInstance(std::move(contextRef), operation, + std::move(parentKeepAlive)); } PyOperationRef PyOperation::createDetached(PyMlirContextRef contextRef, MlirOperation operation, nb::object parentKeepAlive) { - nb::ft_lock_guard lock(contextRef->liveOperationsMutex); - auto &liveOperations = contextRef->liveOperations; - assert(liveOperations.count(operation.ptr) == 0 && - "cannot create detached operation that already exists"); - (void)liveOperations; PyOperationRef created = createInstance(std::move(contextRef), operation, std::move(parentKeepAlive)); - liveOperations[operation.ptr] = - std::make_pair(created.getObject(), created.get()); created->attached = false; return created; } @@ -1523,7 +1409,7 @@ nb::object PyOperation::create(std::string_view name, llvm::ArrayRef<MlirValue> operands, std::optional<nb::dict> attributes, std::optional<std::vector<PyBlock *>> successors, - int regions, DefaultingPyLocation location, + int regions, PyLocation &location, const nb::object &maybeIp, bool inferType) { llvm::SmallVector<MlirType, 4> mlirResults; llvm::SmallVector<MlirBlock, 4> mlirSuccessors; @@ -1627,7 +1513,7 @@ nb::object PyOperation::create(std::string_view name, if (!operation.ptr) throw nb::value_error("Operation creation failed"); PyOperationRef created = - PyOperation::createDetached(location->getContext(), operation); + PyOperation::createDetached(location.getContext(), operation); maybeInsertOperation(created, maybeIp); return created.getObject(); @@ -1655,7 +1541,7 @@ nb::object PyOperation::createOpView() { void PyOperation::erase() { checkValid(); - getContext()->clearOperationAndInside(*this); + setInvalid(); mlirOperationDestroy(operation); } @@ -1937,9 +1823,9 @@ nb::object PyOpView::buildGeneric( std::optional<nb::list> resultTypeList, nb::list operandList, std::optional<nb::dict> attributes, std::optional<std::vector<PyBlock *>> successors, - std::optional<int> regions, DefaultingPyLocation location, + std::optional<int> regions, PyLocation &location, const nb::object &maybeIp) { - PyMlirContextRef context = location->getContext(); + PyMlirContextRef context = location.getContext(); // Class level operation construction metadata. // Operand and result segment specs are either none, which does no @@ -2108,7 +1994,7 @@ nb::object PyOpView::buildGeneric( // Delegate to create. return PyOperation::create(name, /*results=*/std::move(resultTypes), - /*operands=*/std::move(operands), + /*operands=*/operands, /*attributes=*/std::move(attributes), /*successors=*/std::move(successors), /*regions=*/*regions, location, maybeIp, @@ -2789,6 +2675,156 @@ private: PyOperationRef operation; }; +// see +// https://raw.githubusercontent.com/python/pythoncapi_compat/master/pythoncapi_compat.h + +#ifndef _Py_CAST +#define _Py_CAST(type, expr) ((type)(expr)) +#endif + +// Static inline functions should use _Py_NULL rather than using directly NULL +// to prevent C++ compiler warnings. On C23 and newer and on C++11 and newer, +// _Py_NULL is defined as nullptr. +#ifndef _Py_NULL +#if (defined(__STDC_VERSION__) && __STDC_VERSION__ > 201710L) || \ + (defined(__cplusplus) && __cplusplus >= 201103) +#define _Py_NULL nullptr +#else +#define _Py_NULL NULL +#endif +#endif + +// Python 3.10.0a3 +#if PY_VERSION_HEX < 0x030A00A3 + +// bpo-42262 added Py_XNewRef() +#if !defined(Py_XNewRef) +[[maybe_unused]] PyObject *_Py_XNewRef(PyObject *obj) { + Py_XINCREF(obj); + return obj; +} +#define Py_XNewRef(obj) _Py_XNewRef(_PyObject_CAST(obj)) +#endif + +// bpo-42262 added Py_NewRef() +#if !defined(Py_NewRef) +[[maybe_unused]] PyObject *_Py_NewRef(PyObject *obj) { + Py_INCREF(obj); + return obj; +} +#define Py_NewRef(obj) _Py_NewRef(_PyObject_CAST(obj)) +#endif + +#endif // Python 3.10.0a3 + +// Python 3.9.0b1 +#if PY_VERSION_HEX < 0x030900B1 && !defined(PYPY_VERSION) + +// bpo-40429 added PyThreadState_GetFrame() +PyFrameObject *PyThreadState_GetFrame(PyThreadState *tstate) { + assert(tstate != _Py_NULL && "expected tstate != _Py_NULL"); + return _Py_CAST(PyFrameObject *, Py_XNewRef(tstate->frame)); +} + +// bpo-40421 added PyFrame_GetBack() +PyFrameObject *PyFrame_GetBack(PyFrameObject *frame) { + assert(frame != _Py_NULL && "expected frame != _Py_NULL"); + return _Py_CAST(PyFrameObject *, Py_XNewRef(frame->f_back)); +} + +// bpo-40421 added PyFrame_GetCode() +PyCodeObject *PyFrame_GetCode(PyFrameObject *frame) { + assert(frame != _Py_NULL && "expected frame != _Py_NULL"); + assert(frame->f_code != _Py_NULL && "expected frame->f_code != _Py_NULL"); + return _Py_CAST(PyCodeObject *, Py_NewRef(frame->f_code)); +} + +#endif // Python 3.9.0b1 + +MlirLocation tracebackToLocation(MlirContext ctx) { + size_t framesLimit = + PyGlobals::get().getTracebackLoc().locTracebackFramesLimit(); + // Use a thread_local here to avoid requiring a large amount of space. + thread_local std::array<MlirLocation, PyGlobals::TracebackLoc::kMaxFrames> + frames; + size_t count = 0; + + nb::gil_scoped_acquire acquire; + PyThreadState *tstate = PyThreadState_GET(); + PyFrameObject *next; + PyFrameObject *pyFrame = PyThreadState_GetFrame(tstate); + // In the increment expression: + // 1. get the next prev frame; + // 2. decrement the ref count on the current frame (in order that it can get + // gc'd, along with any objects in its closure and etc); + // 3. set current = next. + for (; pyFrame != nullptr && count < framesLimit; + next = PyFrame_GetBack(pyFrame), Py_XDECREF(pyFrame), pyFrame = next) { + PyCodeObject *code = PyFrame_GetCode(pyFrame); + auto fileNameStr = + nb::cast<std::string>(nb::borrow<nb::str>(code->co_filename)); + llvm::StringRef fileName(fileNameStr); + if (!PyGlobals::get().getTracebackLoc().isUserTracebackFilename(fileName)) + continue; + + // co_qualname and PyCode_Addr2Location added in py3.11 +#if PY_VERSION_HEX < 0x030B00F0 + std::string name = + nb::cast<std::string>(nb::borrow<nb::str>(code->co_name)); + llvm::StringRef funcName(name); + int startLine = PyFrame_GetLineNumber(pyFrame); + MlirLocation loc = + mlirLocationFileLineColGet(ctx, wrap(fileName), startLine, 0); +#else + std::string name = + nb::cast<std::string>(nb::borrow<nb::str>(code->co_qualname)); + llvm::StringRef funcName(name); + int startLine, startCol, endLine, endCol; + int lasti = PyFrame_GetLasti(pyFrame); + if (!PyCode_Addr2Location(code, lasti, &startLine, &startCol, &endLine, + &endCol)) { + throw nb::python_error(); + } + MlirLocation loc = mlirLocationFileLineColRangeGet( + ctx, wrap(fileName), startLine, startCol, endLine, endCol); +#endif + + frames[count] = mlirLocationNameGet(ctx, wrap(funcName), loc); + ++count; + } + // When the loop breaks (after the last iter), current frame (if non-null) + // is leaked without this. + Py_XDECREF(pyFrame); + + if (count == 0) + return mlirLocationUnknownGet(ctx); + + MlirLocation callee = frames[0]; + assert(!mlirLocationIsNull(callee) && "expected non-null callee location"); + if (count == 1) + return callee; + + MlirLocation caller = frames[count - 1]; + assert(!mlirLocationIsNull(caller) && "expected non-null caller location"); + for (int i = count - 2; i >= 1; i--) + caller = mlirLocationCallSiteGet(frames[i], caller); + + return mlirLocationCallSiteGet(callee, caller); +} + +PyLocation +maybeGetTracebackLocation(const std::optional<PyLocation> &location) { + if (location.has_value()) + return location.value(); + if (!PyGlobals::get().getTracebackLoc().locTracebacksEnabled()) + return DefaultingPyLocation::resolve(); + + PyMlirContext &ctx = DefaultingPyMlirContext::resolve(); + MlirLocation mlirLoc = tracebackToLocation(ctx.get()); + PyMlirContextRef ref = PyMlirContext::forContext(ctx.get()); + return {ref, mlirLoc}; +} + } // namespace //------------------------------------------------------------------------------ @@ -2876,14 +2912,6 @@ void mlir::python::populateIRCore(nb::module_ &m) { PyMlirContextRef ref = PyMlirContext::forContext(self.get()); return ref.releaseObject(); }) - .def("_get_live_operation_count", &PyMlirContext::getLiveOperationCount) - .def("_get_live_operation_objects", - &PyMlirContext::getLiveOperationObjects) - .def("_clear_live_operations", &PyMlirContext::clearLiveOperations) - .def("_clear_live_operations_inside", - nb::overload_cast<MlirOperation>( - &PyMlirContext::clearOperationsInside)) - .def("_get_live_module_count", &PyMlirContext::getLiveModuleCount) .def_prop_ro(MLIR_PYTHON_CAPI_PTR_ATTR, &PyMlirContext::getCapsule) .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyMlirContext::createFromCapsule) .def("__enter__", &PyMlirContext::contextEnter) @@ -3052,10 +3080,10 @@ void mlir::python::populateIRCore(nb::module_ &m) { .def("__eq__", [](PyLocation &self, nb::object other) { return false; }) .def_prop_ro_static( "current", - [](nb::object & /*class*/) { + [](nb::object & /*class*/) -> std::optional<PyLocation *> { auto *loc = PyThreadContextEntry::getDefaultLocation(); if (!loc) - throw nb::value_error("No current Location"); + return std::nullopt; return loc; }, "Gets the Location bound to the current thread or raises ValueError") @@ -3201,7 +3229,9 @@ void mlir::python::populateIRCore(nb::module_ &m) { //---------------------------------------------------------------------------- nb::class_<PyModule>(m, "Module", nb::is_weak_referenceable()) .def_prop_ro(MLIR_PYTHON_CAPI_PTR_ATTR, &PyModule::getCapsule) - .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyModule::createFromCapsule) + .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyModule::createFromCapsule, + kModuleCAPICreate) + .def("_clear_mlir_module", &PyModule::clearMlirModule) .def_static( "parse", [](const std::string &moduleAsm, DefaultingPyMlirContext context) { @@ -3240,8 +3270,9 @@ void mlir::python::populateIRCore(nb::module_ &m) { kModuleParseDocstring) .def_static( "create", - [](DefaultingPyLocation loc) { - MlirModule module = mlirModuleCreateEmpty(loc); + [](const std::optional<PyLocation> &loc) { + PyLocation pyLoc = maybeGetTracebackLocation(loc); + MlirModule module = mlirModuleCreateEmpty(pyLoc.get()); return PyModule::forModule(module).releaseObject(); }, nb::arg("loc").none() = nb::none(), "Creates an empty module") @@ -3280,7 +3311,13 @@ void mlir::python::populateIRCore(nb::module_ &m) { // Defer to the operation's __str__. return self.attr("operation").attr("__str__")(); }, - kOperationStrDunderDocstring); + kOperationStrDunderDocstring) + .def( + "__eq__", + [](PyModule &self, PyModule &other) { + return mlirModuleEqual(self.get(), other.get()); + }, + "other"_a); //---------------------------------------------------------------------------- // Mapping of Operation. @@ -3292,7 +3329,8 @@ void mlir::python::populateIRCore(nb::module_ &m) { }) .def("__eq__", [](PyOperationBase &self, PyOperationBase &other) { - return &self.getOperation() == &other.getOperation(); + return mlirOperationEqual(self.getOperation().get(), + other.getOperation().get()); }) .def("__eq__", [](PyOperationBase &self, nb::object other) { return false; }) @@ -3442,6 +3480,14 @@ void mlir::python::populateIRCore(nb::module_ &m) { return operation.createOpView(); }, "Detaches the operation from its parent block.") + .def_prop_ro( + "attached", + [](PyOperationBase &self) { + PyOperation &operation = self.getOperation(); + operation.checkValid(); + return operation.isAttached(); + }, + "Reports if the operation is attached to its parent block.") .def("erase", [](PyOperationBase &self) { self.getOperation().erase(); }) .def("walk", &PyOperationBase::walk, nb::arg("callback"), nb::arg("walk_order") = MlirWalkPostOrder); @@ -3454,8 +3500,8 @@ void mlir::python::populateIRCore(nb::module_ &m) { std::optional<std::vector<PyValue *>> operands, std::optional<nb::dict> attributes, std::optional<std::vector<PyBlock *>> successors, int regions, - DefaultingPyLocation location, const nb::object &maybeIp, - bool inferType) { + const std::optional<PyLocation> &location, + const nb::object &maybeIp, bool inferType) { // Unpack/validate operands. llvm::SmallVector<MlirValue, 4> mlirOperands; if (operands) { @@ -3467,8 +3513,9 @@ void mlir::python::populateIRCore(nb::module_ &m) { } } + PyLocation pyLoc = maybeGetTracebackLocation(location); return PyOperation::create(name, results, mlirOperands, attributes, - successors, regions, location, maybeIp, + successors, regions, pyLoc, maybeIp, inferType); }, nb::arg("name"), nb::arg("results").none() = nb::none(), @@ -3498,7 +3545,9 @@ void mlir::python::populateIRCore(nb::module_ &m) { [](PyOperationBase &self) { return PyOpSuccessors(self.getOperation().getRef()); }, - "Returns the list of Operation successors."); + "Returns the list of Operation successors.") + .def("_set_invalid", &PyOperation::setInvalid, + "Invalidate the operation."); auto opViewClass = nb::class_<PyOpView, PyOperationBase>(m, "OpView") @@ -3512,12 +3561,14 @@ void mlir::python::populateIRCore(nb::module_ &m) { std::optional<nb::list> resultTypeList, nb::list operandList, std::optional<nb::dict> attributes, std::optional<std::vector<PyBlock *>> successors, - std::optional<int> regions, DefaultingPyLocation location, + std::optional<int> regions, + const std::optional<PyLocation> &location, const nb::object &maybeIp) { + PyLocation pyLoc = maybeGetTracebackLocation(location); new (self) PyOpView(PyOpView::buildGeneric( name, opRegionSpec, operandSegmentSpecObj, resultSegmentSpecObj, resultTypeList, operandList, - attributes, successors, regions, location, maybeIp)); + attributes, successors, regions, pyLoc, maybeIp)); }, nb::arg("name"), nb::arg("opRegionSpec"), nb::arg("operandSegmentSpecObj").none() = nb::none(), @@ -3540,7 +3591,11 @@ void mlir::python::populateIRCore(nb::module_ &m) { [](PyOperationBase &self) { return PyOpSuccessors(self.getOperation().getRef()); }, - "Returns the list of Operation successors."); + "Returns the list of Operation successors.") + .def( + "_set_invalid", + [](PyOpView &self) { self.getOperation().setInvalid(); }, + "Invalidate the operation."); opViewClass.attr("_ODS_REGIONS") = nb::make_tuple(0, true); opViewClass.attr("_ODS_OPERAND_SEGMENTS") = nb::none(); opViewClass.attr("_ODS_RESULT_SEGMENTS") = nb::none(); @@ -3551,17 +3606,18 @@ void mlir::python::populateIRCore(nb::module_ &m) { [](nb::handle cls, std::optional<nb::list> resultTypeList, nb::list operandList, std::optional<nb::dict> attributes, std::optional<std::vector<PyBlock *>> successors, - std::optional<int> regions, DefaultingPyLocation location, + std::optional<int> regions, std::optional<PyLocation> location, const nb::object &maybeIp) { std::string name = nb::cast<std::string>(cls.attr("OPERATION_NAME")); std::tuple<int, bool> opRegionSpec = nb::cast<std::tuple<int, bool>>(cls.attr("_ODS_REGIONS")); nb::object operandSegmentSpec = cls.attr("_ODS_OPERAND_SEGMENTS"); nb::object resultSegmentSpec = cls.attr("_ODS_RESULT_SEGMENTS"); + PyLocation pyLoc = maybeGetTracebackLocation(location); return PyOpView::buildGeneric(name, opRegionSpec, operandSegmentSpec, resultSegmentSpec, resultTypeList, operandList, attributes, successors, - regions, location, maybeIp); + regions, pyLoc, maybeIp); }, nb::arg("cls"), nb::arg("results").none() = nb::none(), nb::arg("operands").none() = nb::none(), diff --git a/mlir/lib/Bindings/Python/IRModule.cpp b/mlir/lib/Bindings/Python/IRModule.cpp index e600f1b..0de2f17 100644 --- a/mlir/lib/Bindings/Python/IRModule.cpp +++ b/mlir/lib/Bindings/Python/IRModule.cpp @@ -13,9 +13,9 @@ #include "Globals.h" #include "NanobindUtils.h" +#include "mlir-c/Bindings/Python/Interop.h" #include "mlir-c/Support.h" #include "mlir/Bindings/Python/Nanobind.h" -#include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind. namespace nb = nanobind; using namespace mlir; @@ -197,3 +197,71 @@ PyGlobals::lookupOperationClass(llvm::StringRef operationName) { // Not found and loading did not yield a registration. return std::nullopt; } + +bool PyGlobals::TracebackLoc::locTracebacksEnabled() { + nanobind::ft_lock_guard lock(mutex); + return locTracebackEnabled_; +} + +void PyGlobals::TracebackLoc::setLocTracebacksEnabled(bool value) { + nanobind::ft_lock_guard lock(mutex); + locTracebackEnabled_ = value; +} + +size_t PyGlobals::TracebackLoc::locTracebackFramesLimit() { + nanobind::ft_lock_guard lock(mutex); + return locTracebackFramesLimit_; +} + +void PyGlobals::TracebackLoc::setLocTracebackFramesLimit(size_t value) { + nanobind::ft_lock_guard lock(mutex); + locTracebackFramesLimit_ = std::min(value, kMaxFrames); +} + +void PyGlobals::TracebackLoc::registerTracebackFileInclusion( + const std::string &file) { + nanobind::ft_lock_guard lock(mutex); + auto reg = "^" + llvm::Regex::escape(file); + if (userTracebackIncludeFiles.insert(reg).second) + rebuildUserTracebackIncludeRegex = true; + if (userTracebackExcludeFiles.count(reg)) { + if (userTracebackExcludeFiles.erase(reg)) + rebuildUserTracebackExcludeRegex = true; + } +} + +void PyGlobals::TracebackLoc::registerTracebackFileExclusion( + const std::string &file) { + nanobind::ft_lock_guard lock(mutex); + auto reg = "^" + llvm::Regex::escape(file); + if (userTracebackExcludeFiles.insert(reg).second) + rebuildUserTracebackExcludeRegex = true; + if (userTracebackIncludeFiles.count(reg)) { + if (userTracebackIncludeFiles.erase(reg)) + rebuildUserTracebackIncludeRegex = true; + } +} + +bool PyGlobals::TracebackLoc::isUserTracebackFilename( + const llvm::StringRef file) { + nanobind::ft_lock_guard lock(mutex); + if (rebuildUserTracebackIncludeRegex) { + userTracebackIncludeRegex.assign( + llvm::join(userTracebackIncludeFiles, "|")); + rebuildUserTracebackIncludeRegex = false; + isUserTracebackFilenameCache.clear(); + } + if (rebuildUserTracebackExcludeRegex) { + userTracebackExcludeRegex.assign( + llvm::join(userTracebackExcludeFiles, "|")); + rebuildUserTracebackExcludeRegex = false; + isUserTracebackFilenameCache.clear(); + } + if (!isUserTracebackFilenameCache.contains(file)) { + std::string fileStr = file.str(); + bool include = std::regex_search(fileStr, userTracebackIncludeRegex); + bool exclude = std::regex_search(fileStr, userTracebackExcludeRegex); + isUserTracebackFilenameCache[file] = include || !exclude; + } + return isUserTracebackFilenameCache[file]; +} diff --git a/mlir/lib/Bindings/Python/IRModule.h b/mlir/lib/Bindings/Python/IRModule.h index 9c22dea..0cc0459 100644 --- a/mlir/lib/Bindings/Python/IRModule.h +++ b/mlir/lib/Bindings/Python/IRModule.h @@ -192,16 +192,6 @@ public: PyMlirContext(const PyMlirContext &) = delete; PyMlirContext(PyMlirContext &&) = delete; - /// For the case of a python __init__ (nanobind::init) method, pybind11 is - /// quite strict about needing to return a pointer that is not yet associated - /// to an nanobind::object. Since the forContext() method acts like a pool, - /// possibly returning a recycled context, it does not satisfy this need. The - /// usual way in python to accomplish such a thing is to override __new__, but - /// that is also not supported by pybind11. Instead, we use this entry - /// point which always constructs a fresh context (which cannot alias an - /// existing one because it is fresh). - static PyMlirContext *createNewContextForInit(); - /// Returns a context reference for the singleton PyMlirContext wrapper for /// the given context. static PyMlirContextRef forContext(MlirContext context); @@ -228,40 +218,6 @@ public: /// Gets the count of live context objects. Used for testing. static size_t getLiveCount(); - /// Get a list of Python objects which are still in the live context map. - std::vector<PyOperation *> getLiveOperationObjects(); - - /// Gets the count of live operations associated with this context. - /// Used for testing. - size_t getLiveOperationCount(); - - /// Clears the live operations map, returning the number of entries which were - /// invalidated. To be used as a safety mechanism so that API end-users can't - /// corrupt by holding references they shouldn't have accessed in the first - /// place. - size_t clearLiveOperations(); - - /// Removes an operation from the live operations map and sets it invalid. - /// This is useful for when some non-bindings code destroys the operation and - /// the bindings need to made aware. For example, in the case when pass - /// manager is run. - /// - /// Note that this does *NOT* clear the nested operations. - void clearOperation(MlirOperation op); - - /// Clears all operations nested inside the given op using - /// `clearOperation(MlirOperation)`. - void clearOperationsInside(PyOperationBase &op); - void clearOperationsInside(MlirOperation op); - - /// Clears the operaiton _and_ all operations inside using - /// `clearOperation(MlirOperation)`. - void clearOperationAndInside(PyOperationBase &op); - - /// Gets the count of live modules associated with this context. - /// Used for testing. - size_t getLiveModuleCount(); - /// Enter and exit the context manager. static nanobind::object contextEnter(nanobind::object context); void contextExit(const nanobind::object &excType, @@ -288,25 +244,6 @@ private: static nanobind::ft_mutex live_contexts_mutex; static LiveContextMap &getLiveContexts(); - // Interns all live modules associated with this context. Modules tracked - // in this map are valid. When a module is invalidated, it is removed - // from this map, and while it still exists as an instance, any - // attempt to access it will raise an error. - using LiveModuleMap = - llvm::DenseMap<const void *, std::pair<nanobind::handle, PyModule *>>; - LiveModuleMap liveModules; - - // Interns all live operations associated with this context. Operations - // tracked in this map are valid. When an operation is invalidated, it is - // removed from this map, and while it still exists as an instance, any - // attempt to access it will raise an error. - using LiveOperationMap = - llvm::DenseMap<void *, std::pair<nanobind::handle, PyOperation *>>; - nanobind::ft_mutex liveOperationsMutex; - - // Guarded by liveOperationsMutex in free-threading mode. - LiveOperationMap liveOperations; - bool emitErrorDiagnostics = false; MlirContext context; @@ -558,8 +495,8 @@ class PyModule; using PyModuleRef = PyObjectRef<PyModule>; class PyModule : public BaseContextObject { public: - /// Returns a PyModule reference for the given MlirModule. This may return - /// a pre-existing or new object. + /// Returns a PyModule reference for the given MlirModule. This always returns + /// a new object. static PyModuleRef forModule(MlirModule module); PyModule(PyModule &) = delete; PyModule(PyMlirContext &&) = delete; @@ -580,11 +517,12 @@ public: nanobind::object getCapsule(); /// Creates a PyModule from the MlirModule wrapped by a capsule. - /// Note that PyModule instances are uniqued, so the returned object - /// may be a pre-existing object. Ownership of the underlying MlirModule - /// is taken by calling this function. + /// Note this returns a new object BUT clearMlirModule() must be called to + /// prevent double-frees (of the underlying mlir::Module). static nanobind::object createFromCapsule(nanobind::object capsule); + void clearMlirModule() { module = {nullptr}; } + private: PyModule(PyMlirContextRef contextRef, MlirModule module); MlirModule module; @@ -722,8 +660,7 @@ public: llvm::ArrayRef<MlirValue> operands, std::optional<nanobind::dict> attributes, std::optional<std::vector<PyBlock *>> successors, int regions, - DefaultingPyLocation location, const nanobind::object &ip, - bool inferType); + PyLocation &location, const nanobind::object &ip, bool inferType); /// Creates an OpView suitable for this operation. nanobind::object createOpView(); @@ -781,7 +718,7 @@ public: nanobind::list operandList, std::optional<nanobind::dict> attributes, std::optional<std::vector<PyBlock *>> successors, - std::optional<int> regions, DefaultingPyLocation location, + std::optional<int> regions, PyLocation &location, const nanobind::object &maybeIp); /// Construct an instance of a class deriving from OpView, bypassing its @@ -1227,7 +1164,7 @@ public: /// Note that PyAffineExpr instances are uniqued, so the returned object /// may be a pre-existing object. Ownership of the underlying MlirAffineExpr /// is taken by calling this function. - static PyAffineExpr createFromCapsule(nanobind::object capsule); + static PyAffineExpr createFromCapsule(const nanobind::object &capsule); PyAffineExpr add(const PyAffineExpr &other) const; PyAffineExpr mul(const PyAffineExpr &other) const; @@ -1254,7 +1191,7 @@ public: /// Note that PyAffineMap instances are uniqued, so the returned object /// may be a pre-existing object. Ownership of the underlying MlirAffineMap /// is taken by calling this function. - static PyAffineMap createFromCapsule(nanobind::object capsule); + static PyAffineMap createFromCapsule(const nanobind::object &capsule); private: MlirAffineMap affineMap; @@ -1274,7 +1211,7 @@ public: /// Creates a PyIntegerSet from the MlirAffineMap wrapped by a capsule. /// Note that PyIntegerSet instances may be uniqued, so the returned object /// may be a pre-existing object. Integer sets are owned by the context. - static PyIntegerSet createFromCapsule(nanobind::object capsule); + static PyIntegerSet createFromCapsule(const nanobind::object &capsule); private: MlirIntegerSet integerSet; diff --git a/mlir/lib/Bindings/Python/IRTypes.cpp b/mlir/lib/Bindings/Python/IRTypes.cpp index b11e3f7..a9b1259 100644 --- a/mlir/lib/Bindings/Python/IRTypes.cpp +++ b/mlir/lib/Bindings/Python/IRTypes.cpp @@ -963,7 +963,7 @@ public: static void bindDerived(ClassTy &c) { c.def_static( "get", - [](std::string dialectNamespace, std::string typeData, + [](const std::string &dialectNamespace, const std::string &typeData, DefaultingPyMlirContext context) { MlirType type = mlirOpaqueTypeGet(context->get(), toMlirStringRef(dialectNamespace), diff --git a/mlir/lib/Bindings/Python/MainModule.cpp b/mlir/lib/Bindings/Python/MainModule.cpp index 6f49431..278847e 100644 --- a/mlir/lib/Bindings/Python/MainModule.cpp +++ b/mlir/lib/Bindings/Python/MainModule.cpp @@ -6,7 +6,6 @@ // //===----------------------------------------------------------------------===// - #include "Globals.h" #include "IRModule.h" #include "NanobindUtils.h" @@ -44,7 +43,27 @@ NB_MODULE(_mlir, m) { .def("_register_operation_impl", &PyGlobals::registerOperationImpl, "operation_name"_a, "operation_class"_a, nb::kw_only(), "replace"_a = false, - "Testing hook for directly registering an operation"); + "Testing hook for directly registering an operation") + .def("loc_tracebacks_enabled", + [](PyGlobals &self) { + return self.getTracebackLoc().locTracebacksEnabled(); + }) + .def("set_loc_tracebacks_enabled", + [](PyGlobals &self, bool enabled) { + self.getTracebackLoc().setLocTracebacksEnabled(enabled); + }) + .def("set_loc_tracebacks_frame_limit", + [](PyGlobals &self, int n) { + self.getTracebackLoc().setLocTracebackFramesLimit(n); + }) + .def("register_traceback_file_inclusion", + [](PyGlobals &self, const std::string &filename) { + self.getTracebackLoc().registerTracebackFileInclusion(filename); + }) + .def("register_traceback_file_exclusion", + [](PyGlobals &self, const std::string &filename) { + self.getTracebackLoc().registerTracebackFileExclusion(filename); + }); // Aside from making the globals accessible to python, having python manage // it is necessary to make sure it is destroyed (and releases its python diff --git a/mlir/lib/Bindings/Python/Pass.cpp b/mlir/lib/Bindings/Python/Pass.cpp index 20017e2..88e28dc 100644 --- a/mlir/lib/Bindings/Python/Pass.cpp +++ b/mlir/lib/Bindings/Python/Pass.cpp @@ -39,7 +39,7 @@ public: return nb::steal<nb::object>(mlirPythonPassManagerToCapsule(get())); } - static nb::object createFromCapsule(nb::object capsule) { + static nb::object createFromCapsule(const nb::object &capsule) { MlirPassManager rawPm = mlirPythonCapsuleToPassManager(capsule.ptr()); if (mlirPassManagerIsNull(rawPm)) throw nb::python_error(); @@ -159,11 +159,7 @@ void mlir::python::populatePassManagerSubmodule(nb::module_ &m) { "ValueError if the pipeline can't be parsed.") .def( "run", - [](PyPassManager &passManager, PyOperationBase &op, - bool invalidateOps) { - if (invalidateOps) { - op.getOperation().getContext()->clearOperationsInside(op); - } + [](PyPassManager &passManager, PyOperationBase &op) { // Actually run the pass manager. PyMlirContext::ErrorCapture errors(op.getOperation().getContext()); MlirLogicalResult status = mlirPassManagerRunOnOp( @@ -172,7 +168,7 @@ void mlir::python::populatePassManagerSubmodule(nb::module_ &m) { throw MLIRError("Failure while executing pass pipeline", errors.take()); }, - "operation"_a, "invalidate_ops"_a = true, + "operation"_a, "Run the pass manager on the provided operation, raising an " "MLIRError on failure.") .def( diff --git a/mlir/lib/Bindings/Python/RegisterEverything.cpp b/mlir/lib/Bindings/Python/RegisterEverything.cpp index 3ba42be..3edcb09 100644 --- a/mlir/lib/Bindings/Python/RegisterEverything.cpp +++ b/mlir/lib/Bindings/Python/RegisterEverything.cpp @@ -7,8 +7,8 @@ //===----------------------------------------------------------------------===// #include "mlir-c/RegisterEverything.h" -#include "mlir/Bindings/Python/NanobindAdaptors.h" #include "mlir/Bindings/Python/Nanobind.h" +#include "mlir/Bindings/Python/NanobindAdaptors.h" NB_MODULE(_mlirRegisterEverything, m) { m.doc() = "MLIR All Upstream Dialects, Translations and Passes Registration"; diff --git a/mlir/lib/Bindings/Python/TransformInterpreter.cpp b/mlir/lib/Bindings/Python/TransformInterpreter.cpp index f9b0fed..920bca8 100644 --- a/mlir/lib/Bindings/Python/TransformInterpreter.cpp +++ b/mlir/lib/Bindings/Python/TransformInterpreter.cpp @@ -67,7 +67,6 @@ static void populateTransformInterpreterSubmodule(nb::module_ &m) { // root. This is awkward, but we don't have access to PyMlirContext // object here otherwise. nb::object obj = nb::cast(payloadRoot); - obj.attr("context").attr("_clear_live_operations_inside")(payloadRoot); MlirLogicalResult result = mlirTransformApplyNamedSequence( payloadRoot, transformRoot, transformModule, options.options); |