Cold-Air-Networks/mirror-llvm-project
2
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 7

[MLIR][Python] Support type definitions in Python-defined dialects (#182805)

Browse source 
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]]]
```
This commit is contained in:
Twice 2026-02-24 10:34:58 +08:00 committed by GitHub
parent ccfd59a03e
commit 3b2c1db870
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 200 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

134
mlir/python/mlir/dialects/ext.py
View file

@ -31,6 +31,7 @@ __all__ = [
"Operand",
"Result",
"Region",
"Type",
"register_dialect",
"register_operation",
]
@ -43,6 +44,20 @@ register_dialect = _cext.register_dialect
register_operation = _cext.register_operation
def construct_instance(origin, args):
# `origin.get` is to construct an instance of MLIR type or attribute.
return origin.get(
*(
(
construct_instance(get_origin(arg), get_args(arg))
if get_origin(arg)
else arg
)
for arg in args
)
)
class ConstraintLoweringContext:
def __init__(self):
self._cache: Dict[str, ir.Value] = {}
@ -70,14 +85,19 @@ class ConstraintLoweringContext:
elif isinstance(type_, TypeVar):
return self.lower(type_)
elif origin and issubclass(origin, ir.Type):
# `origin.get` is to construct an instance of MLIR type.
t = origin.get(*get_args(type_))
if issubclass(origin, Type):
return irdl.parametric(
base_type=[origin._dialect_name, origin._name],
args=[self.lower(arg) for arg in get_args(type_)],
)
t = construct_instance(origin, get_args(type_))
return irdl.is_(ir.TypeAttr.get(t))
elif origin and issubclass(origin, ir.Attribute):
# `origin.get` is to construct an instance of MLIR attribute.
attr = origin.get(*get_args(type_))
attr = construct_instance(origin, get_args(type_))
return irdl.is_(attr)
elif issubclass(type_, ir.Type):
if issubclass(type_, Type):
return irdl.base(base_ref=[type_._dialect_name, type_._name])
return irdl.base(base_name=f"!{type_.type_name}")
elif issubclass(type_, ir.Attribute):
return irdl.base(base_name=f"#{type_.attr_name}")
@ -96,8 +116,7 @@ def infer_type(type_) -> Optional[Callable[[], ir.Type]]:
origin = get_origin(type_)
if origin and issubclass(origin, ir.Type):
# `origin.get` is to construct an instance of MLIR type/attribute.
return lambda: origin.get(*get_args(type_))
return lambda: construct_instance(origin, get_args(type_))
elif isinstance(type_, TypeVar):
return infer_type(type_.__bound__)
return None
@ -202,7 +221,7 @@ def match_optional(type_) -> Optional[Any]:
class Operation(ir.OpView):
"""
Base class of Python-defined operation.
Base class of Python-defined operations.
The following example shows two ways to define operations via this class:
```python
@ -488,6 +507,98 @@ class Operation(ir.OpView):
)
@dataclass
class ParamDef:
name: str
constraint: Any
class Type(ir.DynamicType):
"""
Base class of Python-defined types.
The following example shows two ways to define types via this class:
```python
class MyType(MyDialect.Type, name=..):
...
class MyType(Type, dialect=MyDialect, name=..):
...
```
"""
@classmethod
def __init_subclass__(
cls,
*,
name: str | None = None,
dialect: type | None = None,
**kwargs,
):
super().__init_subclass__(**kwargs)
fields = []
for base in cls.__bases__:
if hasattr(base, "_fields"):
fields.extend(base._fields)
for key, value in cls.__annotations__.items():
field = ParamDef(key, value)
fields.append(field)
cls._fields = fields
if dialect:
if hasattr(cls, "_dialect_obj"):
raise RuntimeError(
f"This type has already been attached to dialect '{cls._dialect_obj.DIALECT_NAMESPACE}'."
)
cls._dialect_obj = dialect
# for subclasses without "name" parameter,
# just treat them as normal classes
if not name:
return
if not hasattr(cls, "_dialect_obj"):
raise RuntimeError(
"Type subclasses must either inherit from a Dialect's Type subclass "
"or provide the dialect as a class keyword argument."
)
cls._name = name
cls._dialect_name = cls._dialect_obj.DIALECT_NAMESPACE
cls.type_name = f"{cls._dialect_name}.{name}"
for i, field in enumerate(cls._fields):
setattr(
cls,
field.name,
property(lambda self, i=i: self.params[i]),
)
cls._dialect_obj.types.append(cls)
@classmethod
def get(cls, *args, context=None):
args = [
ir.TypeAttr.get(arg, context) if isinstance(arg, ir.Type) else arg
for arg in args
]
return cls(ir.DynamicType.get(cls.type_name, args, context=context))
@classmethod
def _emit_type(cls) -> None:
ctx = ConstraintLoweringContext()
t = irdl.type_(cls._name)
with ir.InsertionPoint(t.body):
irdl.parameters(
[ctx.lower(f.constraint) for f in cls._fields],
[f.name for f in cls._fields],
)
class Dialect(ir.Dialect):
"""
Base class of a Python-defined dialect.
@ -521,11 +632,20 @@ class Dialect(ir.Dialect):
dict(),
dialect=cls,
)
cls.types = []
cls.Type = type(
"Type",
(Type,),
dict(),
dialect=cls,
)
@classmethod
def _emit_dialect(cls) -> None:
d = irdl.dialect(cls.name)
with ir.InsertionPoint(d.body):
for type_ in cls.types:
type_._emit_type()
for op in cls.operations:
op._emit_operation()

73
mlir/test/python/dialects/ext.py
View file

@ -519,3 +519,76 @@ def testExtDialectWithRegion():
# CHECK: Verification failed:
# CHECK: result type mismatch
print(e)
# CHECK: TEST: testExtDialectWithType
@run
def testExtDialectWithType():
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]]]
with Context(), Location.unknown():
TestType.load()
# CHECK: irdl.dialect @ext_type {
# CHECK: irdl.type @array {
# CHECK: %0 = irdl.is i32
# CHECK: %1 = irdl.is i64
# CHECK: %2 = irdl.any_of(%0, %1)
# CHECK: %3 = irdl.base "#builtin.integer"
# CHECK: irdl.parameters(elem_type: %2, length: %3)
# CHECK: }
# CHECK: irdl.operation @make_array {
# CHECK: %0 = irdl.base @ext_type::@array
# CHECK: irdl.results(arr: %0)
# CHECK: }
# CHECK: irdl.operation @make_array3 {
# CHECK: %0 = irdl.is i32
# CHECK: %1 = irdl.is 3 : i32
# CHECK: %2 = irdl.parametric @ext_type::@array<%0, %1>
# CHECK: irdl.results(arr: %2)
# CHECK: }
# CHECK: }
print(TestType._mlir_module)
# CHECK: ext_type.array
print(Array.type_name)
i32 = IntegerType.get_signless(32)
i64 = IntegerType.get_signless(64)
a4 = Array.get(i32, IntegerAttr.get(i32, 4))
a6 = Array.get(i64, IntegerAttr.get(i32, 6))
# CHECK: !ext_type.array<i32, 4 : i32>
print(a4)
# CHECK: !ext_type.array<i64, 6 : i32>
print(a6)
# CHECK: i32
print(a4.elem_type)
# CHECK: 4 : i32
print(a4.length)
# CHECK: i64
print(a6.elem_type)
# CHECK: 6 : i32
print(a6.length)
module = Module.create()
with InsertionPoint(module.body):
MakeArrayOp(a4)
MakeArrayOp(a6)
MakeArray3Op()
# CHECK: %0 = "ext_type.make_array"() : () -> !ext_type.array<i32, 4 : i32>
# CHECK: %1 = "ext_type.make_array"() : () -> !ext_type.array<i64, 6 : i32>
# CHECK: %2 = "ext_type.make_array3"() : () -> !ext_type.array<i32, 3 : i32>
assert module.operation.verify()
print(module)