bfb1ba7526
depends on D150839
This diff uses `MlirTypeID` to register `TypeCaster`s (i.e., `[](PyType pyType) -> DerivedTy { return pyType; }`) for all concrete types (i.e., `PyConcrete<...>`) that are then queried for (by `MlirTypeID`) and called in `struct type_caster<MlirType>::cast`. The result is that anywhere an `MlirType mlirType` is returned from a python binding, that `mlirType` is automatically cast to the correct concrete type. For example:
```
c0 = arith.ConstantOp(f32, 0.0)
# CHECK: F32Type(f32)
print(repr(c0.result.type))
unranked_tensor_type = UnrankedTensorType.get(f32)
unranked_tensor = tensor.FromElementsOp(unranked_tensor_type, [c0]).result
# CHECK: UnrankedTensorType
print(type(unranked_tensor.type).__name__)
# CHECK: UnrankedTensorType(tensor<*xf32>)
print(repr(unranked_tensor.type))
```
This functionality immediately extends to typed attributes (i.e., `attr.type`).
The diff also implements similar functionality for `mlir_type_subclass`es but in a slightly different way - for such types (which have no cpp corresponding `class` or `struct`) the user must provide a type caster in python (similar to how `AttrBuilder` works) or in cpp as a `py::cpp_function`.
Reviewed By: ftynse
Differential Revision: https://reviews.llvm.org/D150927
470 lines
14 KiB
Python
470 lines
14 KiB
Python
# RUN: %PYTHON %s | FileCheck %s
|
|
|
|
from mlir.ir import *
|
|
import mlir.dialects.func as func
|
|
import mlir.dialects.python_test as test
|
|
import mlir.dialects.tensor as tensor
|
|
|
|
|
|
def run(f):
|
|
print("\nTEST:", f.__name__)
|
|
f()
|
|
return f
|
|
|
|
|
|
# CHECK-LABEL: TEST: testAttributes
|
|
@run
|
|
def testAttributes():
|
|
with Context() as ctx, Location.unknown():
|
|
ctx.allow_unregistered_dialects = True
|
|
|
|
#
|
|
# Check op construction with attributes.
|
|
#
|
|
|
|
i32 = IntegerType.get_signless(32)
|
|
one = IntegerAttr.get(i32, 1)
|
|
two = IntegerAttr.get(i32, 2)
|
|
unit = UnitAttr.get()
|
|
|
|
# CHECK: "python_test.attributed_op"() {
|
|
# CHECK-DAG: mandatory_i32 = 1 : i32
|
|
# CHECK-DAG: optional_i32 = 2 : i32
|
|
# CHECK-DAG: unit
|
|
# CHECK: }
|
|
op = test.AttributedOp(one, optional_i32=two, unit=unit)
|
|
print(f"{op}")
|
|
|
|
# CHECK: "python_test.attributed_op"() {
|
|
# CHECK: mandatory_i32 = 2 : i32
|
|
# CHECK: }
|
|
op2 = test.AttributedOp(two)
|
|
print(f"{op2}")
|
|
|
|
#
|
|
# Check generic "attributes" access and mutation.
|
|
#
|
|
|
|
assert "additional" not in op.attributes
|
|
|
|
# CHECK: "python_test.attributed_op"() {
|
|
# CHECK-DAG: additional = 1 : i32
|
|
# CHECK-DAG: mandatory_i32 = 2 : i32
|
|
# CHECK: }
|
|
op2.attributes["additional"] = one
|
|
print(f"{op2}")
|
|
|
|
# CHECK: "python_test.attributed_op"() {
|
|
# CHECK-DAG: additional = 2 : i32
|
|
# CHECK-DAG: mandatory_i32 = 2 : i32
|
|
# CHECK: }
|
|
op2.attributes["additional"] = two
|
|
print(f"{op2}")
|
|
|
|
# CHECK: "python_test.attributed_op"() {
|
|
# CHECK-NOT: additional = 2 : i32
|
|
# CHECK: mandatory_i32 = 2 : i32
|
|
# CHECK: }
|
|
del op2.attributes["additional"]
|
|
print(f"{op2}")
|
|
|
|
try:
|
|
print(op.attributes["additional"])
|
|
except KeyError:
|
|
pass
|
|
else:
|
|
assert False, "expected KeyError on unknown attribute key"
|
|
|
|
#
|
|
# Check accessors to defined attributes.
|
|
#
|
|
|
|
# CHECK: Mandatory: 1
|
|
# CHECK: Optional: 2
|
|
# CHECK: Unit: True
|
|
print(f"Mandatory: {op.mandatory_i32.value}")
|
|
print(f"Optional: {op.optional_i32.value}")
|
|
print(f"Unit: {op.unit}")
|
|
|
|
# CHECK: Mandatory: 2
|
|
# CHECK: Optional: None
|
|
# CHECK: Unit: False
|
|
print(f"Mandatory: {op2.mandatory_i32.value}")
|
|
print(f"Optional: {op2.optional_i32}")
|
|
print(f"Unit: {op2.unit}")
|
|
|
|
# CHECK: Mandatory: 2
|
|
# CHECK: Optional: None
|
|
# CHECK: Unit: False
|
|
op.mandatory_i32 = two
|
|
op.optional_i32 = None
|
|
op.unit = False
|
|
print(f"Mandatory: {op.mandatory_i32.value}")
|
|
print(f"Optional: {op.optional_i32}")
|
|
print(f"Unit: {op.unit}")
|
|
assert "optional_i32" not in op.attributes
|
|
assert "unit" not in op.attributes
|
|
|
|
try:
|
|
op.mandatory_i32 = None
|
|
except ValueError:
|
|
pass
|
|
else:
|
|
assert False, "expected ValueError on setting a mandatory attribute to None"
|
|
|
|
# CHECK: Optional: 2
|
|
op.optional_i32 = two
|
|
print(f"Optional: {op.optional_i32.value}")
|
|
|
|
# CHECK: Optional: None
|
|
del op.optional_i32
|
|
print(f"Optional: {op.optional_i32}")
|
|
|
|
# CHECK: Unit: False
|
|
op.unit = None
|
|
print(f"Unit: {op.unit}")
|
|
assert "unit" not in op.attributes
|
|
|
|
# CHECK: Unit: True
|
|
op.unit = True
|
|
print(f"Unit: {op.unit}")
|
|
|
|
# CHECK: Unit: False
|
|
del op.unit
|
|
print(f"Unit: {op.unit}")
|
|
|
|
|
|
# CHECK-LABEL: TEST: attrBuilder
|
|
@run
|
|
def attrBuilder():
|
|
with Context() as ctx, Location.unknown():
|
|
ctx.allow_unregistered_dialects = True
|
|
op = test.AttributesOp(
|
|
x_bool=True,
|
|
x_i16=1,
|
|
x_i32=2,
|
|
x_i64=3,
|
|
x_si16=-1,
|
|
x_si32=-2,
|
|
x_f32=1.5,
|
|
x_f64=2.5,
|
|
x_str="x_str",
|
|
x_i32_array=[1, 2, 3],
|
|
x_i64_array=[4, 5, 6],
|
|
x_f32_array=[1.5, -2.5, 3.5],
|
|
x_f64_array=[4.5, 5.5, -6.5],
|
|
x_i64_dense=[1, 2, 3, 4, 5, 6],
|
|
)
|
|
print(op)
|
|
|
|
|
|
# CHECK-LABEL: TEST: inferReturnTypes
|
|
@run
|
|
def inferReturnTypes():
|
|
with Context() as ctx, Location.unknown(ctx):
|
|
test.register_python_test_dialect(ctx)
|
|
module = Module.create()
|
|
with InsertionPoint(module.body):
|
|
op = test.InferResultsOp()
|
|
dummy = test.DummyOp()
|
|
|
|
# CHECK: [Type(i32), Type(i64)]
|
|
iface = InferTypeOpInterface(op)
|
|
print(iface.inferReturnTypes())
|
|
|
|
# CHECK: [Type(i32), Type(i64)]
|
|
iface_static = InferTypeOpInterface(test.InferResultsOp)
|
|
print(iface.inferReturnTypes())
|
|
|
|
assert isinstance(iface.opview, test.InferResultsOp)
|
|
assert iface.opview == iface.operation.opview
|
|
|
|
try:
|
|
iface_static.opview
|
|
except TypeError:
|
|
pass
|
|
else:
|
|
assert False, (
|
|
"not expected to be able to obtain an opview from a static" " interface"
|
|
)
|
|
|
|
try:
|
|
InferTypeOpInterface(dummy)
|
|
except ValueError:
|
|
pass
|
|
else:
|
|
assert False, "not expected dummy op to implement the interface"
|
|
|
|
try:
|
|
InferTypeOpInterface(test.DummyOp)
|
|
except ValueError:
|
|
pass
|
|
else:
|
|
assert False, "not expected dummy op class to implement the interface"
|
|
|
|
|
|
# CHECK-LABEL: TEST: resultTypesDefinedByTraits
|
|
@run
|
|
def resultTypesDefinedByTraits():
|
|
with Context() as ctx, Location.unknown(ctx):
|
|
test.register_python_test_dialect(ctx)
|
|
module = Module.create()
|
|
with InsertionPoint(module.body):
|
|
inferred = test.InferResultsOp()
|
|
same = test.SameOperandAndResultTypeOp([inferred.results[0]])
|
|
# CHECK-COUNT-2: i32
|
|
print(same.one.type)
|
|
print(same.two.type)
|
|
|
|
first_type_attr = test.FirstAttrDeriveTypeAttrOp(
|
|
inferred.results[1], TypeAttr.get(IndexType.get())
|
|
)
|
|
# CHECK-COUNT-2: index
|
|
print(first_type_attr.one.type)
|
|
print(first_type_attr.two.type)
|
|
|
|
first_attr = test.FirstAttrDeriveAttrOp(FloatAttr.get(F32Type.get(), 3.14))
|
|
# CHECK-COUNT-3: f32
|
|
print(first_attr.one.type)
|
|
print(first_attr.two.type)
|
|
print(first_attr.three.type)
|
|
|
|
implied = test.InferResultsImpliedOp()
|
|
# CHECK: i32
|
|
print(implied.integer.type)
|
|
# CHECK: f64
|
|
print(implied.flt.type)
|
|
# CHECK: index
|
|
print(implied.index.type)
|
|
|
|
|
|
# CHECK-LABEL: TEST: testOptionalOperandOp
|
|
@run
|
|
def testOptionalOperandOp():
|
|
with Context() as ctx, Location.unknown():
|
|
test.register_python_test_dialect(ctx)
|
|
|
|
module = Module.create()
|
|
with InsertionPoint(module.body):
|
|
|
|
op1 = test.OptionalOperandOp()
|
|
# CHECK: op1.input is None: True
|
|
print(f"op1.input is None: {op1.input is None}")
|
|
|
|
op2 = test.OptionalOperandOp(input=op1)
|
|
# CHECK: op2.input is None: False
|
|
print(f"op2.input is None: {op2.input is None}")
|
|
|
|
|
|
# CHECK-LABEL: TEST: testCustomAttribute
|
|
@run
|
|
def testCustomAttribute():
|
|
with Context() as ctx:
|
|
test.register_python_test_dialect(ctx)
|
|
a = test.TestAttr.get()
|
|
# CHECK: #python_test.test_attr
|
|
print(a)
|
|
|
|
# The following cast must not assert.
|
|
b = test.TestAttr(a)
|
|
|
|
unit = UnitAttr.get()
|
|
try:
|
|
test.TestAttr(unit)
|
|
except ValueError as e:
|
|
assert "Cannot cast attribute to TestAttr" in str(e)
|
|
else:
|
|
raise
|
|
|
|
# The following must trigger a TypeError from our adaptors and must not
|
|
# crash.
|
|
try:
|
|
test.TestAttr(42)
|
|
except TypeError as e:
|
|
assert "Expected an MLIR object" in str(e)
|
|
else:
|
|
raise
|
|
|
|
# The following must trigger a TypeError from pybind (therefore, not
|
|
# checking its message) and must not crash.
|
|
try:
|
|
test.TestAttr(42, 56)
|
|
except TypeError:
|
|
pass
|
|
else:
|
|
raise
|
|
|
|
|
|
@run
|
|
def testCustomType():
|
|
with Context() as ctx:
|
|
test.register_python_test_dialect(ctx)
|
|
a = test.TestType.get()
|
|
# CHECK: !python_test.test_type
|
|
print(a)
|
|
|
|
# The following cast must not assert.
|
|
b = test.TestType(a)
|
|
# Instance custom types should have typeids
|
|
assert isinstance(b.typeid, TypeID)
|
|
# Subclasses of ir.Type should not have a static_typeid
|
|
# CHECK: 'TestType' object has no attribute 'static_typeid'
|
|
try:
|
|
b.static_typeid
|
|
except AttributeError as e:
|
|
print(e)
|
|
|
|
i8 = IntegerType.get_signless(8)
|
|
try:
|
|
test.TestType(i8)
|
|
except ValueError as e:
|
|
assert "Cannot cast type to TestType" in str(e)
|
|
else:
|
|
raise
|
|
|
|
# The following must trigger a TypeError from our adaptors and must not
|
|
# crash.
|
|
try:
|
|
test.TestType(42)
|
|
except TypeError as e:
|
|
assert "Expected an MLIR object" in str(e)
|
|
else:
|
|
raise
|
|
|
|
# The following must trigger a TypeError from pybind (therefore, not
|
|
# checking its message) and must not crash.
|
|
try:
|
|
test.TestType(42, 56)
|
|
except TypeError:
|
|
pass
|
|
else:
|
|
raise
|
|
|
|
|
|
@run
|
|
# CHECK-LABEL: TEST: testTensorValue
|
|
def testTensorValue():
|
|
with Context() as ctx, Location.unknown():
|
|
test.register_python_test_dialect(ctx)
|
|
|
|
i8 = IntegerType.get_signless(8)
|
|
|
|
class Tensor(test.TestTensorValue):
|
|
def __str__(self):
|
|
return super().__str__().replace("Value", "Tensor")
|
|
|
|
module = Module.create()
|
|
with InsertionPoint(module.body):
|
|
t = tensor.EmptyOp([10, 10], i8).result
|
|
|
|
# CHECK: Value(%{{.*}} = tensor.empty() : tensor<10x10xi8>)
|
|
print(Value(t))
|
|
|
|
tt = Tensor(t)
|
|
# CHECK: Tensor(%{{.*}} = tensor.empty() : tensor<10x10xi8>)
|
|
print(tt)
|
|
|
|
# CHECK: False
|
|
print(tt.is_null())
|
|
|
|
# Classes of custom types that inherit from concrete types should have
|
|
# static_typeid
|
|
assert isinstance(test.TestIntegerRankedTensorType.static_typeid, TypeID)
|
|
# And it should be equal to the in-tree concrete type
|
|
assert test.TestIntegerRankedTensorType.static_typeid == t.type.typeid
|
|
|
|
|
|
# CHECK-LABEL: TEST: inferReturnTypeComponents
|
|
@run
|
|
def inferReturnTypeComponents():
|
|
with Context() as ctx, Location.unknown(ctx):
|
|
test.register_python_test_dialect(ctx)
|
|
module = Module.create()
|
|
i32 = IntegerType.get_signless(32)
|
|
with InsertionPoint(module.body):
|
|
resultType = UnrankedTensorType.get(i32)
|
|
operandTypes = [
|
|
RankedTensorType.get([1, 3, 10, 10], i32),
|
|
UnrankedTensorType.get(i32),
|
|
]
|
|
f = func.FuncOp(
|
|
"test_inferReturnTypeComponents", (operandTypes, [resultType])
|
|
)
|
|
entry_block = Block.create_at_start(f.operation.regions[0], operandTypes)
|
|
with InsertionPoint(entry_block):
|
|
ranked_op = test.InferShapedTypeComponentsOp(
|
|
resultType, entry_block.arguments[0]
|
|
)
|
|
unranked_op = test.InferShapedTypeComponentsOp(
|
|
resultType, entry_block.arguments[1]
|
|
)
|
|
|
|
# CHECK: has rank: True
|
|
# CHECK: rank: 4
|
|
# CHECK: element type: i32
|
|
# CHECK: shape: [1, 3, 10, 10]
|
|
iface = InferShapedTypeOpInterface(ranked_op)
|
|
shaped_type_components = iface.inferReturnTypeComponents(
|
|
operands=[ranked_op.operand]
|
|
)[0]
|
|
print("has rank:", shaped_type_components.has_rank)
|
|
print("rank:", shaped_type_components.rank)
|
|
print("element type:", shaped_type_components.element_type)
|
|
print("shape:", shaped_type_components.shape)
|
|
|
|
# CHECK: has rank: False
|
|
# CHECK: rank: None
|
|
# CHECK: element type: i32
|
|
# CHECK: shape: None
|
|
iface = InferShapedTypeOpInterface(unranked_op)
|
|
shaped_type_components = iface.inferReturnTypeComponents(
|
|
operands=[unranked_op.operand]
|
|
)[0]
|
|
print("has rank:", shaped_type_components.has_rank)
|
|
print("rank:", shaped_type_components.rank)
|
|
print("element type:", shaped_type_components.element_type)
|
|
print("shape:", shaped_type_components.shape)
|
|
|
|
|
|
# CHECK-LABEL: TEST: testCustomTypeTypeCaster
|
|
@run
|
|
def testCustomTypeTypeCaster():
|
|
with Context() as ctx, Location.unknown():
|
|
test.register_python_test_dialect(ctx)
|
|
|
|
a = test.TestType.get()
|
|
assert a.typeid is not None
|
|
|
|
b = Type.parse("!python_test.test_type")
|
|
# CHECK: !python_test.test_type
|
|
print(b)
|
|
# CHECK: TestType(!python_test.test_type)
|
|
print(repr(b))
|
|
|
|
c = test.TestIntegerRankedTensorType.get([10, 10], 5)
|
|
# CHECK: tensor<10x10xi5>
|
|
print(c)
|
|
# CHECK: TestIntegerRankedTensorType(tensor<10x10xi5>)
|
|
print(repr(c))
|
|
|
|
# CHECK: Type caster is already registered
|
|
try:
|
|
|
|
def type_caster(pytype):
|
|
return test.TestIntegerRankedTensorType(pytype)
|
|
|
|
register_type_caster(c.typeid, type_caster)
|
|
except RuntimeError as e:
|
|
print(e)
|
|
|
|
def type_caster(pytype):
|
|
return test.TestIntegerRankedTensorType(pytype)
|
|
|
|
register_type_caster(c.typeid, type_caster, replace=True)
|
|
|
|
d = tensor.EmptyOp([10, 10], IntegerType.get_signless(5)).result
|
|
# CHECK: tensor<10x10xi5>
|
|
print(d.type)
|
|
# CHECK: TestIntegerRankedTensorType(tensor<10x10xi5>)
|
|
print(repr(d.type))
|