* `mlir.ir` now exports `_OperationBase`. It is handy to use when both
`Operation` and `OpView` are accepted.
* Added type arguments where they were missing, e.g.
`list[ir.Attribute]` instead of just `list`.
* Changed `Opview.build_generic` and `OpView.parse` to return `Self`
instead of the supertype `Type`.
* Changed the bindings generator to emit a parameterized `OpResult` when
the exact type is available.
Cleaning up XeGPU transform ops. Now that XeGPU layout propagation
works, it is sufficient to set the layouts for anchor ops (e.g.
load/store/dpas ops) only.
Changes:
* Remove `xegpu.get_desc_op` and `xegpu.set_desc_layout`. Users should
not change the layout of descriptor op's return value anymore.
* Add `xegpu.get_load_op(value)` that finds either `xegpu.load_nd` or
`xegpu.load` op in the value's producer chain. This is a useful utility
as load ops often need to be annotated with a layout.
* The generic `xegpu.set_op_layout_attr(op, ...)` is now replaced by
`xegpu.set_anchor_layout(op, ...)` that only sets layout attribute of
anchor ops. Raises an error if the given op does not support anchor
layouts.
* `xegpu.insert_prefetch` takes a load op handle instead of a value.
Previously, while users compose `irdl.any_of` and `irdl.base`, e.g.
```mlir
module {
irdl.dialect @ext_attr_in_op {
irdl.operation @op_with_attr {
%0 = irdl.base "#builtin.integer"
%1 = irdl.base "#builtin.string"
%2 = irdl.any_of(%0, %1)
irdl.attributes {"a" = %2}
}
}
}
```
The program will crash due to `llvm_unreachable("unknown IRDL
constraint")`.
This PR implements `getBases(..)` for `irdl::BaseOp` to make this work.
This make fields like `attr: IntegerAttr | StringAttr` work in
Python-defined dialects, which led to a crash previously.
Currently, Python-defined operations automatically generate an
`__init__` function to serve as the operation builder. Previously, the
parameters of this `__init__` function followed a fairly complex set of
rules. For example:
* All result fields were moved to the front to align with other op
builders.
* Fields of `Optional` type were automatically moved to the end and
treated as keyword parameters.
* If the types of all result fields could be inferred automatically,
then all result fields were removed from the parameter list.
* Other than that, the parameter order followed the field definition
order.
These rules may seem reasonable, and they have worked well in practice,
but they have one major drawback: users cannot easily tell what the
actual `__init__` parameter list will look like when writing code,
because the rules are simply too complicated. Users can inspect the
signature at runtime via `MyOp.__init__.__signature__`, but this is
clearly poor for the development experience.
After some offline discussion with Rolf, we decided to replace this with
the following rules, which is what this PR implements:
* The parameters of `__init__` now strictly follow the field definition
order.
* By default, all parameters are required. A field becomes a keyword
parameter only if it is declared with `= None`.
* By default, no type inference is performed. A field becomes a keyword
parameter and participates in type inference only if it is declared with
`= infer_type()`.
These new rules give users full control over the `__init__` parameter
list, making it easy to understand the parameter order and explicitly
control optional parameters and type inference. In addition, this makes
it much easier in the future to use `dataclass_transform` so that type
checkers can understand these automatically generated `__init__`
methods, although there are still some issues to resolve at the moment.
NOTE that this is a **breaking change**. Even so, I still believe it is
worth making, mainly for two reasons:
- the current number of users of Python-defined dialects is still quite
limited;
- this change will bring long-term benefits.
Assisted by Copilot.
When multiple dialects share td `#includes` (e.g. `affine` includes
`arith`), each dialect's `*_enum_gen.py` file registers attribute
builders under the same keys, causing "already registered" errors on the
second import; the first commit checks in such a case which currently
fails on main:
```
# | RuntimeError: Attribute builder for 'Arith_CmpFPredicateAttr' is already registered with func: <function _arith_cmpfpredicateattr at 0x78d13cbe9a80>
```
This PR implements a two-pronged fix:
1. Add `allow_existing=True` to `register_attribute_builder` (and the
underlying C++ `registerAttributeBuilder`). When set, silently skips
registration if the key already exists (first-wins semantics). This
handles `EnumInfo`-based builders which have no dialect prefix (e.g.
`AtomicRMWKindAttr`, `Arith_CmpFPredicateAttr`), which may be emitted by
every dialect whose td file includes the defining file;
2. Filter `EnumAttr` builders by `-bind-dialect` in
`EnumPythonBindingGen.cpp` and register them under dialect qualified
keys (`"dialect.AttrName"`). Update `OpPythonBindingGen.cpp` to look up
the same qualified keys for EnumAttr typed op attributes (detected via
`isSubClassOf("EnumAttr")`). Pass `-bind-dialect` from
`AddMLIRPython.cmake`.
This approach incurs no changes to `ir.py` registrations (no "builtin."
prefix), and no manual builder additions to individual dialect Python
files (unlike the previous attempt
https://github.com/llvm/llvm-project/pull/117918).
Note, this PR was "clauded" not "coded".
This patch adds support for `float8_e3m4` and `float8_e4m3` in
`np_to_memref.py` by adding the appropriate ctypes structures.
Additionally changes minimum numpy version to 2.1.0 and uses a single
ml_dtypes version of 0.5.0.
This PR adds some LLVM metadata attributes and an `llvm.named_metadata`
container op (similar to `llvm.module_flags`) for those attributes.
Summary:
- Add MLIR attributes modeling LLVM IR metadata: `#llvm.md_string`,
`#llvm.md_const`, `#llvm.md_func`, and `#llvm.md_node`;
- Add `llvm.named_metadata` container op for module-level named metadata
nodes;
- Add MLIR-to-LLVM-IR translation for the new attributes and op;
- Add C API functions (`mlirLLVMMDStringAttrGet`,
`mlirLLVMMDNodeAttrGet`, etc.);
- Add Python bindings (`llvm.MDStringAttr`, `llvm.MDConstantAttr`,
`llvm.MDFuncAttr`, `llvm.MDNodeAttr`, `llvm.FunctionType`).
This includes several changes:
- `Dialect.load(reload=False)` will fail if the dialect was already
loaded in a different context. To prevent the further program abortion.
- `Dialect.load(reload=True)` implies `replace=True` in
dialect/operation registering.
- `PyGlobals::registerDialectImpl` now has a parameter `replace`.
- `register_dialect` and `register_operation` is no longer exposed in
`mlir.dialects.ext`.
This should solve the registering problem found in writing transform
test cases by @rolfmorel.
MLIR's C++ `Operation::walk` supports type-filtered traversal (e.g.
`op->walk([](arith::AddIOp op) { ... })`), but the Python binding
`op.walk()` requires users to manually implement type filtering inside
the callback function.
This PR adds type filtering into the python binding `op.walk()`, if
users pass `op_class`, walk() will only apply callback to matching ops.
This PR also adds a common use helper in mlir/ir that collects all ops
of a given type into a list. Users can just call: `ops =
ir.get_ops_of_type(root, op_class)`.
Now we need to provide a location when call `load()`, e.g.
```python
with Context(), Location.unknown():
MyDialect.load()
```
But it's actually weird: IRDL is just an implementation details, so for
users they don't know why they need to provide a location for loading a
dialect, which is unrelated to constructing an IR module.
This PR made location optional for dialect loading.
* `_OperationBase.walk` was missing a default.
* `MLIRError` is now fully defined in C++. The monkey-patching
previously done in `_site_initialize` was opaque to type checkers.
The `op.parent` only returns the immediate parent, in which case
downstream users have to traverse the operation by themselves to find a
specific type op.
This PR adds a python function `get_parent_of_type()` to mlir.ir to
provide an API to do so.
The function mirrors the implementation here:
https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/IR/Operation.h#L257-L273.
Instead of creating a new binding, reimplement it in python using
`isinstance()` is simpler.
Currently XeGPU transform dialect does not allow the user to set the
`order` attribute of a layout in `SetDescLayoutOp` and
`SetOpLayoutAttrOp`. This PR adds `order` as an optional argument to
these transform ops.
* Removed an explicit `nb::sig` for `static_typeid`. The inferred type
would
work just fine, and unqualified `TypeID`, which was there previously,
only
really works for core types in the `ir` submodule.
* `DefaultingPyMlir*` helpers also produce qualified types, e.g.
`_mlir.ir.Location` instead of bare `Location`.
* `ir.*.__enter__` now returns a concrete type instead of `object`, e.g.
`ir.Context.__enter__` returns `Context`.
* `loc_tracebacks` uses `Generator` as the return type, since this is
what
`contextmanager` expects in typeshed.
* Changed static methods on subclasses of `DenseElementsAttribute` to
return
that concrete subclass, instead of `DenseElementsAttribute`.
---------
Co-authored-by: Maksim Levental <maksim.levental@gmail.com>
In this PR, `replace` keyword parameter is added to `Dialect.load(..)`
in `mlir.dialects.ext`. It's to replace existing registered
operations/types/attrs.
This PR is quite similiar to
https://github.com/llvm/llvm-project/pull/182805.
We added basic support of attribute definitions in Python-defined
dialects, including:
- IRDL codegen for attribute definitions
- Attr builders like `MyAttr.get(..)` and attr parameter accessors (e.g.
`my_attr.param1`)
- Use Python-defined attrs in Python-defined operations
Assisted by GitHub Copilot.
Previously, in #177782, we added support for dialect conversion and
generated an `OpAdaptor` subtype for every ODS-defined operation. In
this PR, we will also generate `OpAdaptor` subtypes for Python-defined
operations, so that they can be applied in dialect conversion as well.
Renames 'x86vector' dialect to 'x86'.
This is the first PR in series of cleanups around dialects targeting x86
platforms.
The new naming scheme is shorter, cleaner, and opens possibility of
integrating other x86-specific operations not strictly fitting pure
vector representation. For example, the generalization will allow for
future merger of AMX dialect into the x86 dialect to create one-stop x86
operations collection and boost discoverability.
Previously, we were using the static `typeid` of `DynamicType` for
checks, which is incorrect. We should instead check against the `typeid`
of `DynamicTypeDefinition` (which is a subclass of `SelfOwningTypeID`),
and register it via `register_type_caster` so that Python-defined types
can use `maybe_downcast`. (The attribute part is same.)
In this PR, we added basic support of type definitions in Python-defined
dialects, including:
- IRDL codegen for type definitions
- Type builders like `MyType.get(..)` and type parameter accessors (e.g.
`my_type.param1`)
- Use Python-defined types in Python-defined oeprations
```python
class TestType(Dialect, name="ext_type"):
pass
class Array(TestType.Type, name="array"):
elem_type: IntegerType[32] | IntegerType[64]
length: IntegerAttr
class MakeArrayOp(TestType.Operation, name="make_array"):
arr: Result[Array]
class MakeArray3Op(TestType.Operation, name="make_array3"):
arr: Result[Array[IntegerType[32], IntegerAttr[IntegerType[32], 3]]]
```
Previously, we constructed new ops using the pattern `class
MyOp(MyInt.Operation)`.
Now we’ve added a new pattern: `class MyOp(Operation, dialect=MyInt)`,
which allows more flexible composition. For example:
```python
class BinOpBase(Operation): # it can be used in any dialect!
res: Result[Any]
lhs: Operand[Any]
rhs: Operand[Any]
class MyInt(Dialect, name="myint"):
pass
class AddOp(BinOpBase, dialect=MyInt, name="add"):
...
```
Changes `transform.xegpu.set_op_layout_attr` to support xegpu anchor
layouts. By default, if `result` and `operand` bool arguments are unset,
this transform op sets the op's anchor layout, if the op supports it
(otherwise emits a silenceable failure).
In contrast to the earlier implementation, setting the operand layout
now requires setting the new `operand` argument.
Currently, when calling `.load(register=False)`, `op._attach_traits()`
isn’t executed. This PR ensures traits are attached regardless of
whether `register` is `True` or `False`.
Provides the infrastructure for implementing and late-binding
OpInterfaces from Python.
* On the mlir-c API declaration side, each `XOpInterface` has a callback
struct, with a callback for each method and a userdata member (provided
as an arg to each method), and a
`mlirXOpInterfaceAttachFallbackModel(ctx, op_name, callbacks)` func.
* This CAPI is implemented by defining a subclass of
`XOpInterface::FallbackModel` that holds the callback struct and has
each method call the corresponding callback (with userdata as an arg).
Given a callback struct, a new `FallbackModel` is created and attached,
i.e. late bound, to the named op. (MLIR's interface infrastructure is
such that the thus registered `FallbackModel` will be returned in case
the op gets cast to the `XOpInterface`.)
* On the Python side, we expose a stand-in `XOpInterface` base class
which has one (class)method: `XOpInterface.attach(cls, op_name, ctx)`.
Python users subclass this class (`class MyInterfaceImpl(XOpInterface):
...`) and implement the interface's methods (with the right names and
signatures). The user calls `attach` on the subclass
(`MyInterfaceImpl.attach("my_dialect.my_op", ctx)`) which prepares the
callbacks struct _with userdata set to the subclass_ (as we use it to
lookup methods). These callbacks (and userdata) are then registered as
an `XOpInterface::FallbackModel` by
`mlirXOpInterfaceAttachFallbackModel(...)`. From then on the Python
methods will be used to respond to calls to the interface methods
(originating in C++).
This PR enables implementing the TransformOpInterface and the
MemoryEffectsOpInterface, both of which are required for making an op
into a transform op.
Everything besides the above linked code is there to facilitate exposing
the interfaces: the right types for the arguments of the methods are
exposed as are functions/methods for manipulating these arguments (e.g.
specifying side effects on `OpOperand`s and `OpResult`s and being able
to access and set the transform handles associated with args and
results).
This will support two syntax in python-defined dialects.
First is that traits can now be declared in class parameters, e.g.
```python
class ParentIsIfTrait(DynamicOpTrait): #define a python-side trait
@staticmethod
def verify_invariants(op) -> bool:
if not isinstance(op.parent.opview, IfOp):
op.location.emit_error(
f"{op.name} should be put inside {IfOp.OPERATION_NAME}"
)
return False
return True
class YieldOp( # attach two traits: IsTerminatorTrait, ParentIsIfTrait
TestRegion.Operation, name="yield", traits=[IsTerminatorTrait, ParentIsIfTrait]
):
...
```
Second is that users can directly define
`verify_invariants`/`verify_region_invariants` methods in the operation
to add additional custom verification logic. And this is implemented via
traits.
```python
class YieldOp(TestRegion.Operation, name="yield", ...):
value: Operand[Any]
def verify_invariants(self) -> bool: # define a method directly
if self.parent.results[0].type != self.value.type:
self.location.emit_error(
"result type mismatch between YieldOp and its parent IfOp"
)
return False
return True
```
Previously we use `verify`/`verify_region` as method names (in
yesterday's PR #179705), but in this PR they are renamed to
`verify_invariants`/`verify_region_invariants` because there are
conflicts between the newly-added `verify` method and `ir.OpView.verify`
method:
- `verify_invariants` is just to attach **additional** verification
logic. but `OpView.verify` is to construct an OperationVerifer and do
full verification for an operation, so the semantics is not same between
these two. We should not shadow the `OpView.verify` method by defining a
new semantically-different `verify` method.
- it will make users confuse between these two `verify` methods, since
they have different meaning.
- if users didn't define the `verify` method in their python-defined
operation, `DynamicOpTraits.attach(opname, MyOpCls)` still do the
attaching (because `hasattr("verify")` returns `True`) and seg fault
(because we cannot attach `OpView.verify`).
---------
Co-authored-by: Rolf Morel <rolfmorel@gmail.com>
This PR adds basic support for defining regions in Python-defined
dialects. Example usage:
```python
class TestRegion(Dialect, name="ext_region"):
pass
class IfOp(TestRegion.Operation, name="if"):
cond: Operand[IntegerType[1]]
then: Region
else_: Region
```
Current limitations:
* We can’t specify region constraints yet (e.g., number of blocks or
block argument types). This will be addressed as a follow-up task.
* We can’t mark an op as a `Terminator` or `NoTerminator` yet. This
depends on `DynamicOpTraits` (#177735) and Python-side trait API
support, and will be implemented in a follow-up PR.
This is the first PR after splitting off #179032.
This is a follow-up PR of #169045.
---------
Co-authored-by: Rolf Morel <rolfmorel@gmail.com>
Python bindings for the IRDL dialect were introduced in #158488. They
are currently usable—for constructing IR and dynamically loading modules
that contain `irdl.dialect` into MLIR. However, there are still several
pain points when working with them:
* The IRDL IR-building interface is not very intuitive and tends to be
quite verbose.
* We do not yet have the corresponding `OpView` classes for IRDL-defined
operations.
To address these issues, I propose creating a wrapper (effectively a
small “DSL”) on top of the existing IRDL Python bindings. This wrapper
aims to simplify IR construction and automatically generate the
corresponding `OpView` types. A simple example is shown below.
Currently, using the IRDL bindings looks like this:
```python
m = Module.create()
with InsertionPoint(m.body):
myint = irdl.dialect("myint")
with InsertionPoint(myint.body):
constant = irdl.operation_("constant")
with InsertionPoint(constant.body):
iattr = irdl.base(base_name="#builtin.integer")
i32 = irdl.is_(TypeAttr.get(IntegerType.get_signless(32)))
irdl.attributes_([iattr], ["value"])
irdl.results_([i32], ["cst"], [irdl.Variadicity.single])
add = irdl.operation_("add")
with InsertionPoint(add.body):
i32 = irdl.is_(TypeAttr.get(IntegerType.get_signless(32)))
irdl.operands_(
[i32, i32],
["lhs", "rhs"],
[irdl.Variadicity.single, irdl.Variadicity.single],
)
irdl.results_([i32], ["res"], [irdl.Variadicity.single])
irdl.load_dialects(m)
```
With the proposed DSL (module name `mlir.dialects.ext`), the equivalent
implementation becomes:
```python
class MyInt(Dialect, name="myint"):
pass
i32 = IntegerType[32]
class ConstantOp(MyInt.Operation, name="constant"):
value: IntegerAttr
cst: Result[i32]
class AddOp(MyInt.Operation, name="add"):
lhs: Operand[i32]
rhs: Operand[i32]
res: Result[i32]
MyInt.load()
```
Compared with the current IRDL Python bindings, this DSL mainly adds the
following:
* **A more intuitive interface** for constructing IRDL definitions (as
shown in the example).
* **Automatic generation of the corresponding `OpView`
classes**—including `__init__` methods and property getters for each
defined operation. Similar to TableGen’s `ins`, operands and attributes
can be interleaved in arbitrary order. Special handling is also
implemented for optional and variadic operands/results (such as
computing segment sizes) so that they feel as natural to use as native
operations.
* **Lazy insertion of ops**: all ops are created and inserted only when
`Dialect.load()` is called, which makes it unnecessary to specify an
MLIR context immediately when defining an IRDL dialect.
* **Basic type inference** in operation builders (i.e.
`OpViewCls.__init__`) for trivial result types.
The current DSL does not yet cover all IRDL operations. Several features
are not supported at the moment:
- Defining new types or attributes
- Parametric constraints
- Adding regions to operations
---------
Co-authored-by: Rolf Morel <rolfmorel@gmail.com>
Extend linalg.pack and linalg.unpack to accept memref operands in
addition to tensors. As part of this change, we now disable all
transformations when these ops have memref semantics.
Closes https://github.com/llvm/llvm-project/issues/129004
---------
Signed-off-by: Ryutaro Okada <1015ryu88@gmail.com>
Co-authored-by: Hyunsung Lee <ita9naiwa@gmail.com>
This is a continuation of the idea from #174091 to add `match` support
for MLIR containers. In this PR the `OpAttributeMap` container is
registered as a `Mapping`, so be mapped as a "dictionary" in `match`
statements.
For this to work the `get(key, default=None)` method had to be
implemented. Those are pretty much copys of `dunderGetItemNamed` and
`dunderGetItemIndexed` with an added argument and `nb::object` as return
type, because they can now return other types than just `PyAttribute`.
Was unsure if I should refactor this to make `dunderGetItem...` use the
new `getWithDefault...` or if a separate method is preferred. Kept it as
a copy for simplicitys sake for now.
Even though the `OpAttributeMap` supports indexing by `int` and `str`,
Python does not allow to register it as a `Sequence` and a `Mapping` at
the same time. If it is registered as a Sequence it only returns the
attribute names as string, not as `NamedAttribute`. It is technically
possible to also use integer keys for the `dict`-like match, but it
doesn't provide any constraints on the number of attributes, etc., so
probably not recommended.
<details><summary>Example</summary>
```python
from mlir.ir import Context, Module, OpAttributeMap
from collections.abc import Sequence
ctx = Context()
ctx.allow_unregistered_dialects = True
module = Module.parse(
r"""
"some.op"() { some.attribute = 1 : i8,
other.attribute = 3.0,
dependent = "text" } : () -> ()
""",
ctx,
)
op = module.body.operations[0]
def test(attr):
match attr:
case [*args]:
print("matched a Sequence", args)
case _:
print("Didn't match as Sequence")
match attr:
case {"some.attribute": a, "other.attribute": b, "dependent": c}:
print("Matched as Mapping individually", a, b, c)
case _:
print("Didn't match a Mapping")
match attr:
case {0: a, 1: b}:
print("Matched as Mapping with 2 int keys", a, b)
case _:
print("Didn't match as Mapping with 2 int keys")
print("Registered as Mapping only:")
test(op.attributes)
print("\nAfter additonally registering as Sequence:")
Sequence.register(OpAttributeMap)
test(op.attributes)
```
Output:
```
Registered as Mapping only:
Didn't match as Sequence
Matched as Mapping individually 1 : i8 3.000000e+00 : f64 "text"
Matched as Mapping with 2 int keys NamedAttribute(dependent="text") NamedAttribute(other.attribute=3.000000e+00 : f64)
After additonally registering as Sequence:
matched a Sequence ['dependent', 'other.attribute', 'some.attribute']
Didn't match a Mapping
Didn't match as Mapping with 2 int keys
```
</details>
makslevental Would be great if you could take a look again ❤️
---------
Co-authored-by: Maksim Levental <maksim.levental@gmail.com>
We've been able to do `isinstance(x, Type)` for a quite a while now
(since
bfb1ba7526)
so remove `Type.isinstance` and the the special-casing
(`_is_integer_type`, `_is_floating_point_type`, `_is_index_type`) in
some places (and therefore support various `fp8`, `fp6`, `fp4` types).
This PR ports all in-tree dialect extensions to use the
`PyConcreteType`, `PyConcreteAttribute` CRTPs instead of
`mlir_pure_subclass`. After this PR we can soft deprecate
`mlir_pure_subclass`. Also API signatures are updated to use `Py*`
instead of `Mlir*` so that type "inference" and hints are improved.
This PR continues the work of
https://github.com/llvm/llvm-project/pull/171775 by moving more useful
types/attributes into MLIRPythonSupport.
You can now do
```c++
struct PyTestIntegerRankedTensorType
: mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::PyConcreteType<
PyTestIntegerRankedTensorType,
mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::PyRankedTensorType>
struct PyTestTensorValue
: mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::PyConcreteValue<
PyTestTensorValue>
```
instead of `mlir_type_subclass` and `mlir_value_subclass`;
**specifically manual registration of the "value caster" via indirection
through the Python interpreter is no longer necessary** . You can also
now freely use all such types at the nanobind API level (e.g., overload
based on `FP*`):
```c++
using mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN;
standaloneM.def("print_fp_type", [](PyF16Type &) { nb::print("this is a fp16 type"); });
standaloneM.def("print_fp_type", [](PyF32Type &) { nb::print("this is a fp32 type"); });
standaloneM.def("print_fp_type", [](PyF64Type &) { nb::print("this is a fp64 type"); });
```
Note, here we only port `PythonTestModuleNanobind` but there is a
follow-up PR that ports **all** in-tree dialect extensions
https://github.com/llvm/llvm-project/pull/174156 to use these. After
that one we can soft deprecate `mlir_pure_subclass`.
Note, depends on https://github.com/llvm/llvm-project/pull/171775
# What
This PR adds a shared library `MLIRPythonSupport` which contains all of
the CRTP classes ike `PyConcreteValue`, `PyConcreteType`,
`PyConcreteAttribute`, as well as other useful code like `Defaulting*`
and etc enabling their reuse in downstream projects. Downstream projects
can now do
```c++
struct PyTestType : mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::PyConcreteType<PyTestType> {
...
};
class PyTestAttr : public mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::PyConcreteAttribute<PyTestAttr> {
...
}
NB_MODULE(_mlirPythonTestNanobind, m) {
PyTestType::bind(m);
PyTestAttr::bind(m);
}
```
instead of using the discordant alternative
`mlir_type_subclass`/`mlir_attr_subclass` (same goes for
`PyConcreteValue`/`mlir_value_subclass`).
# Why
This PR is mostly code motion (along with CMake) but before I describe
the changes I want to state the goals/benefits:
1. Currently upstream "core" extensions and "dialect" extensions ([all
of the `Dialect*` extensions
here](d7c734b5a1/mlir/lib/Bindings/Python))
are a two-tier system;
**a**. [core
extensions](https://github.com/llvm/llvm-project/blob/main/mlir/lib/Bindings/Python/IRTypes.cpp#L361)
enjoy first class support as far as type inference[^3], type stub
generation, and ease of implementation, while dialect extensions [have
poorer support](https://reviews.llvm.org/D150927), incorrect type stub
generation much more tedious (boilerplate) implementation;
**b**. Crucially, this two-tiered system is reflected in the fact that
**the two sets of types/attributes are not in the same Python object
hierarchy**. To wit: `isinstance(..., Type)` and `isinstance(...,
Attribute)` are not supported for the dialect extensions[^2];
**c**. Since these types are not exposed in public headers, downstream
users (dialect extensions or not) cannot write functions that overload
on e.g. `PyFloat8*Type` - that's quite a [useful
feature](fdbee98df8/cpp_ext/TorchOps.cpp (L29-L69))!
2. The dialect extensions incur a sizeable performance penalty relative
to the core extensions in that every single trip across the wire (either
`python->cpp` or `cpp->python`) requires work in addition to nanobind's
own casting/construction pipeline;
**a**. When going from `python->cpp`, [we extract the capsule object
from the Python
object](https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h#L219C24-L219C46)
and then extract from the capsule the `Mlir*` opaque struct/ptr. This
side isn't so onerous;
**b**. When going from `cpp->python` we call long-hand call Python
`import` APIs and construct the Python object using `_CAPICreate`. Note,
there at least 2 `attr` calls incurred in addition to `_CAPICreate`;
this is already much more [efficiently handled by nanobind
itself](4ba51fcf79/src/nb_internals.h (L381-L382))!
3. This division blocks various features: in some configurations[^1] we
trigger a circular import bug because "dialect" types and attributes
perform an [import of the root `_mlir`
module](bd9651bf78/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h (L585))
when they are created (the types themselves, not even instances of those
types). This blocks type stub generation for dialect extensions (i.e.,
the reason we currently only generate type stubs for `_mlir`).
# How
Prior this was not done/possible because of "ODR" issues but I have
resolved those issues; the basic idea for how we solve this is "move
things we want to share into shared libraries":
1. Move IRCore (stuff like `PyConcreteValue`, `PyConcreteType`,
`PyConcreteAttribute`) into `MLIRPythonSupport`;
- Note, we move the rest of the things in `IRModule.h` (renamed to
`IRCore.h`) because `PyConcreteValue`, `PyConcreteType`,
`PyConcreteAttribute` depend on them. This makes for a bigger PR than
one would hope for but ultimately I think we should give people access
to these classes to use as they see fit (specifically inherit from, but
also liberally use in bindings signatures instead of the opaque `Mlir*`
struct wrappers).
2. Put all of this code into a nested namespace
`MLIR_BINDINGS_PYTHON_DOMAIN` which is determined by a compile time
define (and tied to `MLIR_BINDINGS_PYTHON_NB_DOMAIN`). This is necessary
in order to prevent conflicts on both symbol name **and** typeid
(necessary for nanobind to not double register binded types) between
multiple bindings libraries (e.g., `torch-mlir`, and `jax`). Note
[nanobind doesn't support `module_local` like
pybind11](https://nanobind.readthedocs.io/en/latest/porting.html#removed-features).
It does support `NB_DOMAIN` but that is not sufficient for
disambiguating typeids across projects (to wit: we currently define
`NB_DOMAIN` and it was still necessary to move everything to a nested
namespace);
3. Build the [nanobind library itself as a shared
object](https://github.com/wjakob/nanobind/blob/master/cmake/nanobind-config.cmake#L127)
(and link it to both the extensions and `MLIRPythonSupport`).
4. CMake to make this work, in-tree, out-of-tree, downstream, upstream,
etc.
# Testing
Three tests are added here
1. `PythonTestModuleNanobind` is ported to use
`PyConcreteType<PyTestType>` instead of `mlir_type_subclass` and
`PyConcreteAttribute<PyTestAttr>` instead of `mlir_atrr_subclass`,
verifying this works for non-core extensions in-tree;
2. `StandaloneExtensionNanobind` is ported to use `struct PyCustomType :
mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::PyConcreteType<PyCustomType>`
instead of `mlir_type_subclass` verifying this works for non-core
extensions out-of-tree;
3. `StandaloneExtensionNanobind`'s `smoketest` is extended to also load
another bindings package (namely `mlir`) verifying
`MLIR_BINDINGS_PYTHON_DOMAIN` successfully disambiguates symbols and
typeids.
I have also tested this downstream:
https://github.com/llvm/eudsl/pull/287 as well run the following builder
bots:
mlir-nvidia-gcc7:
https://lab.llvm.org/buildbot/#/buildrequests/6654424?redirect_to_build=true
I have also tested against IREE:
https://github.com/iree-org/iree/pull/21916
# Integration
It is highly recommended to set the CMake var
`MLIR_BINDINGS_PYTHON_NB_DOMAIN` (which will also determine
`MLIR_BINDINGS_PYTHON_DOMAIN`) to something unique for each downstream.
This can also be passed explicitly to `add_mlir_python_modules` if your
project builds multiple bindings packages. I added a `WARNING` to this
effect in `AddMLIRPython.cmake`.
[^3]: Python values being typed correctly when exiting from cpp;
[^1]: Specifically when the modules are imported using `importlib`,
which occurs with nanobind's
[stubgen](https://github.com/wjakob/nanobind/blob/master/src/stubgen.py#L965);
[^2]: The workaround we implemented was a class method for the dialect
bindings called `Class.isinstance(...)`;
This allows these containers to be used in `match` statements, which
allows extracting properties and asserting a shape at the same time.
It seems to be only possible, to match as _either_ a `Mapping` _or_ a
`Sequence`, so the `OpAttributeMap` is only a `Mapping`.
I couldn't find a way to make these C++ based types properly inherit
from `Sequence` or `Mapping`, so the Mixins are not provided (nanobind
only allows C++ parent classes, modifying `__base__` complains about
differing destructors).
`OpAttributeMap` was lacking the `get` method, so I simply copied it
from `collections.abc.Mapping`.
When writing the tests i ran into the error, that I wrote
`func.FuncOp(body=[Block(...)])` instead of
`func.FuncOp(body=Region(blocks=[Block(...)]))`. So maybe also turning
`Region` itself into a Sequence would be a good addition as well? Would
extend the Scope of this PR, though.
makslevental You suggested I make the PR, so i'm tagging you here as a
potential reviewer. I hope that is ok with you. :)
---------
Co-authored-by: Maksim Levental <maksim.levental@gmail.com>
Fixes: #164800
Ensures unsigned pooling ops in Linalg stay in the integer domain: the
lowering now rejects floating/bool inputs with a clear diagnostic, new
regression tests lock in both the error path and a valid integer
example, and transform decompositions are updated to reflect the integer
typing.
Signed-off-by: Akimasa Watanuki <mencotton0410@gmail.com>
`from ._xxx_ops_gen import _Dialect` appears in some dialect modules,
like builtin, scf, irdl.. but not all of them. This PR ensures that for
upstream dialects, `<dialect module>._Dialect` is availble, like
`arith._Dialect`.
This PR is a prerequisite for the work I’m currently doing. Later on,
I’d like to use these `_Dialect` objects via something like
`conversion_target.add_legal_dialect(arith._Dialect)` (we could of
course just use strings like `add_legal_dialect("arith")`, but compared
to using a defined symbol, I think that’s more prone to typos).
This is a follow-up of #171957 that updates the argument names of
`scf.if` Python binding to be consistent with `affine.if`. Basically,
both operations should use `has_else` to determine whether the `if`
block is presented.
cc @makslevental
Friendlier wrapper for transform.foreach.
To facilitate that friendliness, makes it so that OpResult.owner returns
the relevant OpView instead of Operation. For good measure, also changes
Value.owner to return OpView instead of Operation, thereby ensuring
consistency. That is, makes it is so that all op-returning .owner
accessors return OpView (and thereby give access to all goodies
available on registered OpViews.)
Reland of #171544 due to fixup for integration test.
Friendlier wrapper for `transform.foreach`.
To facilitate that friendliness, makes it so that `OpResult.owner`
returns the relevant `OpView` instead of `Operation`. For good measure,
also changes `Value.owner` to return `OpView` instead of `Operation`,
thereby ensuring consistency. That is, makes it is so that all
op-returning `.owner` accessors return `OpView` (and thereby give access
to all goodies available on registered `OpView`s.)
This bug was introduced by #108323, where the loc and ip were not
properly set. It may lead to errors when the operations are not linearly
asserted to the IR.
There were two bugs lurking in mlir.ir.loc_tracebacks():
1) The default None parameter was not handled correctly (passed to a
C++ function that expects ints.
2) The `yield` was incorrectly indented meaning loc_tracebacks()
could not be nested (a "generator didn't yield" exception would be
raised).
Added testing of loc_tracebacks by replacing the custom contextmanager
in the auto_location.py test with the loc_tracebacks() API.
Had to harden the test to line number differences.
---------
Co-authored-by: James Molloy <jmolloy@google.com>
