[Mlir-commits] [mlir] [MLIR][Python] Add a DSL for defining dialects in Python bindings (PR #169045)

[llvmlistbot at llvm.org]

================
@@ -0,0 +1,422 @@
+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from typing import (
+    Dict,
+    List,
+    Union,
+    Tuple,
+    Any,
+    Optional,
+    Callable,
+    TypeVar,
+    get_origin,
+    get_args,
+)
+from collections.abc import Sequence
+from dataclasses import dataclass
+from inspect import Parameter, Signature
+from types import UnionType
+from . import irdl
+from ._ods_common import _cext, segmented_accessor
+from .irdl import Variadicity
+from ..passmanager import PassManager
+
+ir = _cext.ir
+
+__all__ = [
+    "Dialect",
+    "Operand",
+    "Result",
+]
+
+Operand = ir.Value
+Result = ir.OpResult
+
+
+class ConstraintLoweringContext:
+    def __init__(self):
+        self._cache: Dict[str, ir.Value] = {}
+
+    def lower(self, type_) -> ir.Value:
+        if type(type_) is TypeVar:
+            if type_.__name__ in self._cache:
+                return self._cache[type_.__name__]
+            v = self._lower(type_.__bound__ or Any)
+            self._cache[type_.__name__] = v
+        else:
+            v = self._lower(type_)
+        return v
+
+    def _lower(self, type_) -> ir.Value:
+        origin = get_origin(type_)
+        if origin is UnionType or origin is Union:
+            return irdl.any_of(self.lower(arg) for arg in get_args(type_))
+        elif type_ is Any:
+            return irdl.any()
+        elif isinstance(type_, TypeVar):
+            return self.lower(type_)
+        elif origin and issubclass(origin, ir.Type):
+            t = origin.get(*get_args(type_))
+            return irdl.is_(ir.TypeAttr.get(t))
+        elif origin and issubclass(origin, ir.Attribute):
+            attr = origin.get(*get_args(type_))
+            return irdl.is_(attr)
+        elif issubclass(type_, ir.Type):
+            return irdl.base(base_name=f"!{type_.type_name}")
+        elif issubclass(type_, ir.Attribute):
+            return irdl.base(base_name=f"#{type_.attr_name}")
+
+        raise TypeError(f"unsupported type in constraints: {type_}")
+
+
+# A function to infer ir.Type from type annotation.
+# Returns a callable that returns the inferred ir.Type,
+# or None if the type cannot be inferred.
+# We use callables so that MLIR contexts are not required
+# while calling this function.
+def infer_type(type_) -> Optional[Callable[[], ir.Type]]:
+    origin = get_origin(type_)
+    if origin and issubclass(origin, ir.Type):
+        return lambda: origin.get(*get_args(type_))
+    elif isinstance(type_, TypeVar):
+        return infer_type(type_.__bound__)
----------------
PragmaTwice wrote:

I think in Python, type hints all have such semantics:
```python
class A:
  x : B
```

Here type of `x` can be a subtype of `B`, isn't it?

Also currently we only support bounded TypeVars, e.g. `X = TypeVar("X", bounds=T)` instead of `X = TypeVar("X", T)`.

I'll leave it as a TODO that we can figure out later, since it looks a corner case that user should rarely encounter this (even it is a problem).

https://github.com/llvm/llvm-project/pull/169045