1037 lines · python
1# RUN: %PYTHON %s pybind11 | FileCheck %s2# RUN: %PYTHON %s nanobind | FileCheck %s3import sys4import typing5from typing import Union, Optional6 7from mlir.ir import *8import mlir.dialects.func as func9import mlir.dialects.python_test as test10import mlir.dialects.tensor as tensor11import mlir.dialects.arith as arith12 13if sys.argv[1] == "pybind11":14 from mlir._mlir_libs._mlirPythonTestPybind11 import (15 TestAttr,16 TestType,17 TestTensorValue,18 TestIntegerRankedTensorType,19 )20 21 test.register_python_test_dialect(get_dialect_registry(), use_nanobind=False)22elif sys.argv[1] == "nanobind":23 from mlir._mlir_libs._mlirPythonTestNanobind import (24 TestAttr,25 TestType,26 TestTensorValue,27 TestIntegerRankedTensorType,28 )29 30 test.register_python_test_dialect(get_dialect_registry(), use_nanobind=True)31else:32 raise ValueError("Expected pybind11 or nanobind as argument")33 34 35def run(f):36 print("\nTEST:", f.__name__)37 f()38 return f39 40 41# CHECK-LABEL: TEST: testAttributes42@run43def testAttributes():44 with Context() as ctx, Location.unknown():45 #46 # Check op construction with attributes.47 #48 49 i32 = IntegerType.get_signless(32)50 one = IntegerAttr.get(i32, 1)51 two = IntegerAttr.get(i32, 2)52 unit = UnitAttr.get()53 54 # CHECK: python_test.attributed_op {55 # CHECK-DAG: mandatory_i32 = 1 : i3256 # CHECK-DAG: optional_i32 = 2 : i3257 # CHECK-DAG: unit58 # CHECK: }59 op = test.AttributedOp(one, optional_i32=two, unit=unit)60 print(f"{op}")61 62 # CHECK: python_test.attributed_op {63 # CHECK: mandatory_i32 = 2 : i3264 # CHECK: }65 op2 = test.AttributedOp(two)66 print(f"{op2}")67 68 #69 # Check generic "attributes" access and mutation.70 #71 72 assert "additional" not in op.attributes73 74 # CHECK: python_test.attributed_op {75 # CHECK-DAG: additional = 1 : i3276 # CHECK-DAG: mandatory_i32 = 2 : i3277 # CHECK: }78 op2.attributes["additional"] = one79 print(f"{op2}")80 81 # CHECK: python_test.attributed_op {82 # CHECK-DAG: additional = 2 : i3283 # CHECK-DAG: mandatory_i32 = 2 : i3284 # CHECK: }85 op2.attributes["additional"] = two86 print(f"{op2}")87 88 # CHECK: python_test.attributed_op {89 # CHECK-NOT: additional = 2 : i3290 # CHECK: mandatory_i32 = 2 : i3291 # CHECK: }92 del op2.attributes["additional"]93 print(f"{op2}")94 95 try:96 print(op.attributes["additional"])97 except KeyError:98 pass99 else:100 assert False, "expected KeyError on unknown attribute key"101 102 #103 # Check accessors to defined attributes.104 #105 106 # CHECK: Mandatory: 1107 # CHECK: Optional: 2108 # CHECK: Unit: True109 print(f"Mandatory: {op.mandatory_i32.value}")110 print(f"Optional: {op.optional_i32.value}")111 print(f"Unit: {op.unit}")112 113 # CHECK: Mandatory: 2114 # CHECK: Optional: None115 # CHECK: Unit: False116 print(f"Mandatory: {op2.mandatory_i32.value}")117 print(f"Optional: {op2.optional_i32}")118 print(f"Unit: {op2.unit}")119 120 # CHECK: Mandatory: 2121 # CHECK: Optional: None122 # CHECK: Unit: False123 op.mandatory_i32 = two124 op.optional_i32 = None125 op.unit = False126 print(f"Mandatory: {op.mandatory_i32.value}")127 print(f"Optional: {op.optional_i32}")128 print(f"Unit: {op.unit}")129 assert "optional_i32" not in op.attributes130 assert "unit" not in op.attributes131 132 try:133 op.mandatory_i32 = None134 except ValueError:135 pass136 else:137 assert False, "expected ValueError on setting a mandatory attribute to None"138 139 # CHECK: Optional: 2140 op.optional_i32 = two141 print(f"Optional: {op.optional_i32.value}")142 143 # CHECK: Optional: None144 del op.optional_i32145 print(f"Optional: {op.optional_i32}")146 147 # CHECK: Unit: False148 op.unit = None149 print(f"Unit: {op.unit}")150 assert "unit" not in op.attributes151 152 # CHECK: Unit: True153 op.unit = True154 print(f"Unit: {op.unit}")155 156 # CHECK: Unit: False157 del op.unit158 print(f"Unit: {op.unit}")159 160 161# CHECK-LABEL: TEST: attrBuilder162@run163def attrBuilder():164 with Context() as ctx, Location.unknown():165 # CHECK: python_test.attributes_op166 op = test.AttributesOp(167 # CHECK-DAG: x_affinemap = affine_map<() -> (2)>168 x_affinemap=AffineMap.get_constant(2),169 # CHECK-DAG: x_affinemaparr = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>]170 x_affinemaparr=[AffineMap.get_identity(3)],171 # CHECK-DAG: x_arr = [true, "x"]172 x_arr=[BoolAttr.get(True), StringAttr.get("x")],173 x_boolarr=[False, True], # CHECK-DAG: x_boolarr = [false, true]174 x_bool=True, # CHECK-DAG: x_bool = true175 x_dboolarr=[True, False], # CHECK-DAG: x_dboolarr = array<i1: true, false>176 x_df16arr=[21, 22], # CHECK-DAG: x_df16arr = array<i16: 21, 22>177 # CHECK-DAG: x_df32arr = array<f32: 2.300000e+01, 2.400000e+01>178 x_df32arr=[23, 24],179 # CHECK-DAG: x_df64arr = array<f64: 2.500000e+01, 2.600000e+01>180 x_df64arr=[25, 26],181 x_di32arr=[0, 1], # CHECK-DAG: x_di32arr = array<i32: 0, 1>182 # CHECK-DAG: x_di64arr = array<i64: 1, 2>183 x_di64arr=[1, 2],184 x_di8arr=[2, 3], # CHECK-DAG: x_di8arr = array<i8: 2, 3>185 # CHECK-DAG: x_dictarr = [{a = false}]186 x_dictarr=[{"a": BoolAttr.get(False)}],187 x_dict={"b": BoolAttr.get(True)}, # CHECK-DAG: x_dict = {b = true}188 x_f32=-2.25, # CHECK-DAG: x_f32 = -2.250000e+00 : f32189 # CHECK-DAG: x_f32arr = [2.000000e+00 : f32, 3.000000e+00 : f32]190 x_f32arr=[2.0, 3.0],191 x_f64=4.25, # CHECK-DAG: x_f64 = 4.250000e+00 : f64192 x_f64arr=[4.0, 8.0], # CHECK-DAG: x_f64arr = [4.000000e+00, 8.000000e+00]193 # CHECK-DAG: x_f64elems = dense<[8.000000e+00, 1.600000e+01]> : tensor<2xf64>194 x_f64elems=[8.0, 16.0],195 # CHECK-DAG: x_flatsymrefarr = [@symbol1, @symbol2]196 x_flatsymrefarr=["symbol1", "symbol2"],197 x_flatsymref="symbol3", # CHECK-DAG: x_flatsymref = @symbol3198 x_i1=0, # CHECK-DAG: x_i1 = false199 x_i16=42, # CHECK-DAG: x_i16 = 42 : i16200 x_i32=6, # CHECK-DAG: x_i32 = 6 : i32201 x_i32arr=[4, 5], # CHECK-DAG: x_i32arr = [4 : i32, 5 : i32]202 x_i32elems=[5, 6], # CHECK-DAG: x_i32elems = dense<[5, 6]> : tensor<2xi32>203 x_i64=9, # CHECK-DAG: x_i64 = 9 : i64204 x_i64arr=[7, 8], # CHECK-DAG: x_i64arr = [7, 8]205 x_i64elems=[8, 9], # CHECK-DAG: x_i64elems = dense<[8, 9]> : tensor<2xi64>206 x_i64svecarr=[10, 11], # CHECK-DAG: x_i64svecarr = [10, 11]207 x_i8=11, # CHECK-DAG: x_i8 = 11 : i8208 x_idx=10, # CHECK-DAG: x_idx = 10 : index209 # CHECK-DAG: x_idxelems = dense<[11, 12]> : tensor<2xindex>210 x_idxelems=[11, 12],211 # CHECK-DAG: x_idxlistarr = [{{\[}}13], [14, 15]]212 x_idxlistarr=[[13], [14, 15]],213 x_si1=-1, # CHECK-DAG: x_si1 = -1 : si1214 x_si16=-2, # CHECK-DAG: x_si16 = -2 : si16215 x_si32=-3, # CHECK-DAG: x_si32 = -3 : si32216 x_si64=-123, # CHECK-DAG: x_si64 = -123 : si64217 x_si8=-4, # CHECK-DAG: x_si8 = -4 : si8218 x_strarr=["hello", "world"], # CHECK-DAG: x_strarr = ["hello", "world"]219 x_str="hello world!", # CHECK-DAG: x_str = "hello world!"220 # CHECK-DAG: x_symrefarr = [@flatsym, @deep::@sym]221 x_symrefarr=["flatsym", ["deep", "sym"]],222 x_symref=["deep", "sym2"], # CHECK-DAG: x_symref = @deep::@sym2223 x_sym="symbol", # CHECK-DAG: x_sym = "symbol"224 x_typearr=[F32Type.get()], # CHECK-DAG: x_typearr = [f32]225 x_type=F64Type.get(), # CHECK-DAG: x_type = f64226 x_ui1=1, # CHECK-DAG: x_ui1 = 1 : ui1227 x_ui16=2, # CHECK-DAG: x_ui16 = 2 : ui16228 x_ui32=3, # CHECK-DAG: x_ui32 = 3 : ui32229 x_ui64=4, # CHECK-DAG: x_ui64 = 4 : ui64230 x_ui8=5, # CHECK-DAG: x_ui8 = 5 : ui8231 x_unit=True, # CHECK-DAG: x_unit232 )233 op.verify()234 op.print(use_local_scope=True)235 236 # fmt: off237 assert typing.get_type_hints(test.AttributesOp.x_affinemaparr.fset)["value"] is ArrayAttr238 assert typing.get_type_hints(test.AttributesOp.x_affinemaparr.fget)["return"] is ArrayAttr239 assert type(op.x_affinemaparr) is typing.get_type_hints(test.AttributesOp.x_affinemaparr.fget)["return"]240 241 assert typing.get_type_hints(test.AttributesOp.x_affinemap.fset)["value"] is AffineMapAttr242 assert typing.get_type_hints(test.AttributesOp.x_affinemap.fget)["return"] is AffineMapAttr243 assert type(op.x_affinemap) is typing.get_type_hints(test.AttributesOp.x_affinemap.fget)["return"]244 245 assert typing.get_type_hints(test.AttributesOp.x_arr.fset)["value"] is ArrayAttr246 assert typing.get_type_hints(test.AttributesOp.x_arr.fget)["return"] is ArrayAttr247 assert type(op.x_arr) is typing.get_type_hints(test.AttributesOp.x_arr.fget)["return"]248 249 assert typing.get_type_hints(test.AttributesOp.x_boolarr.fset)["value"] is ArrayAttr250 assert typing.get_type_hints(test.AttributesOp.x_boolarr.fget)["return"] is ArrayAttr251 assert type(op.x_boolarr) is typing.get_type_hints(test.AttributesOp.x_boolarr.fget)["return"]252 253 assert typing.get_type_hints(test.AttributesOp.x_bool.fset)["value"] is BoolAttr254 assert typing.get_type_hints(test.AttributesOp.x_bool.fget)["return"] is BoolAttr255 assert type(op.x_bool) is typing.get_type_hints(test.AttributesOp.x_bool.fget)["return"]256 257 assert typing.get_type_hints(test.AttributesOp.x_dboolarr.fset)["value"] is DenseBoolArrayAttr258 assert typing.get_type_hints(test.AttributesOp.x_dboolarr.fget)["return"] is DenseBoolArrayAttr259 assert type(op.x_dboolarr) is typing.get_type_hints(test.AttributesOp.x_dboolarr.fget)["return"]260 261 assert typing.get_type_hints(test.AttributesOp.x_df32arr.fset)["value"] is DenseF32ArrayAttr262 assert typing.get_type_hints(test.AttributesOp.x_df32arr.fget)["return"] is DenseF32ArrayAttr263 assert type(op.x_df32arr) is typing.get_type_hints(test.AttributesOp.x_df32arr.fget)["return"]264 265 assert typing.get_type_hints(test.AttributesOp.x_df64arr.fset)["value"] is DenseF64ArrayAttr266 assert typing.get_type_hints(test.AttributesOp.x_df64arr.fget)["return"] is DenseF64ArrayAttr267 assert type(op.x_df64arr) is typing.get_type_hints(test.AttributesOp.x_df64arr.fget)["return"]268 269 assert typing.get_type_hints(test.AttributesOp.x_df16arr.fset)["value"] is DenseI16ArrayAttr270 assert typing.get_type_hints(test.AttributesOp.x_df16arr.fget)["return"] is DenseI16ArrayAttr271 assert type(op.x_df16arr) is typing.get_type_hints(test.AttributesOp.x_df16arr.fget)["return"]272 273 assert typing.get_type_hints(test.AttributesOp.x_di32arr.fset)["value"] is DenseI32ArrayAttr274 assert typing.get_type_hints(test.AttributesOp.x_di32arr.fget)["return"] is DenseI32ArrayAttr275 assert type(op.x_di32arr) is typing.get_type_hints(test.AttributesOp.x_di32arr.fget)["return"]276 277 assert typing.get_type_hints(test.AttributesOp.x_di64arr.fset)["value"] is DenseI64ArrayAttr278 assert typing.get_type_hints(test.AttributesOp.x_di64arr.fget)["return"] is DenseI64ArrayAttr279 assert type(op.x_di64arr) is typing.get_type_hints(test.AttributesOp.x_di64arr.fget)["return"]280 281 assert typing.get_type_hints(test.AttributesOp.x_di8arr.fset)["value"] is DenseI8ArrayAttr282 assert typing.get_type_hints(test.AttributesOp.x_di8arr.fget)["return"] is DenseI8ArrayAttr283 assert type(op.x_di8arr) is typing.get_type_hints(test.AttributesOp.x_di8arr.fget)["return"]284 285 assert typing.get_type_hints(test.AttributesOp.x_dictarr.fset)["value"] is ArrayAttr286 assert typing.get_type_hints(test.AttributesOp.x_dictarr.fget)["return"] is ArrayAttr287 assert type(op.x_dictarr) is typing.get_type_hints(test.AttributesOp.x_dictarr.fget)["return"]288 289 assert typing.get_type_hints(test.AttributesOp.x_dict.fset)["value"] is DictAttr290 assert typing.get_type_hints(test.AttributesOp.x_dict.fget)["return"] is DictAttr291 assert type(op.x_dict) is typing.get_type_hints(test.AttributesOp.x_dict.fget)["return"]292 293 assert typing.get_type_hints(test.AttributesOp.x_f32arr.fset)["value"] is ArrayAttr294 assert typing.get_type_hints(test.AttributesOp.x_f32arr.fget)["return"] is ArrayAttr295 assert type(op.x_f32arr) is typing.get_type_hints(test.AttributesOp.x_f32arr.fget)["return"]296 297 assert typing.get_type_hints(test.AttributesOp.x_f32.fset)["value"] is FloatAttr298 assert typing.get_type_hints(test.AttributesOp.x_f32.fget)["return"] is FloatAttr299 assert type(op.x_f32) is typing.get_type_hints(test.AttributesOp.x_f32.fget)["return"]300 301 assert typing.get_type_hints(test.AttributesOp.x_f64arr.fset)["value"] is ArrayAttr302 assert typing.get_type_hints(test.AttributesOp.x_f64arr.fget)["return"] is ArrayAttr303 assert type(op.x_f64arr) is typing.get_type_hints(test.AttributesOp.x_f64arr.fget)["return"]304 305 assert typing.get_type_hints(test.AttributesOp.x_f64.fset)["value"] is FloatAttr306 assert typing.get_type_hints(test.AttributesOp.x_f64.fget)["return"] is FloatAttr307 assert type(op.x_f64) is typing.get_type_hints(test.AttributesOp.x_f64.fget)["return"]308 309 assert typing.get_type_hints(test.AttributesOp.x_f64elems.fset)["value"] is DenseFPElementsAttr310 assert typing.get_type_hints(test.AttributesOp.x_f64elems.fget)["return"] is DenseFPElementsAttr311 assert type(op.x_f64elems) is typing.get_type_hints(test.AttributesOp.x_f64elems.fget)["return"]312 313 assert typing.get_type_hints(test.AttributesOp.x_flatsymrefarr.fset)["value"] is ArrayAttr314 assert typing.get_type_hints(test.AttributesOp.x_flatsymrefarr.fget)["return"] is ArrayAttr315 assert type(op.x_flatsymrefarr) is typing.get_type_hints(test.AttributesOp.x_flatsymrefarr.fget)["return"]316 317 assert typing.get_type_hints(test.AttributesOp.x_flatsymref.fset)["value"] is FlatSymbolRefAttr318 assert typing.get_type_hints(test.AttributesOp.x_flatsymref.fget)["return"] is FlatSymbolRefAttr319 assert type(op.x_flatsymref) is typing.get_type_hints(test.AttributesOp.x_flatsymref.fget)["return"]320 321 assert typing.get_type_hints(test.AttributesOp.x_i16.fset)["value"] is IntegerAttr322 assert typing.get_type_hints(test.AttributesOp.x_i16.fget)["return"] is IntegerAttr323 assert type(op.x_i16) is typing.get_type_hints(test.AttributesOp.x_i16.fget)["return"]324 325 assert typing.get_type_hints(test.AttributesOp.x_i1.fset)["value"] is BoolAttr326 assert typing.get_type_hints(test.AttributesOp.x_i1.fget)["return"] is BoolAttr327 assert type(op.x_i1) is typing.get_type_hints(test.AttributesOp.x_i1.fget)["return"]328 329 assert typing.get_type_hints(test.AttributesOp.x_i32arr.fset)["value"] is ArrayAttr330 assert typing.get_type_hints(test.AttributesOp.x_i32arr.fget)["return"] is ArrayAttr331 assert type(op.x_i32arr) is typing.get_type_hints(test.AttributesOp.x_i32arr.fget)["return"]332 333 assert typing.get_type_hints(test.AttributesOp.x_i32.fset)["value"] is IntegerAttr334 assert typing.get_type_hints(test.AttributesOp.x_i32.fget)["return"] is IntegerAttr335 assert type(op.x_i32) is typing.get_type_hints(test.AttributesOp.x_i32.fget)["return"]336 337 assert typing.get_type_hints(test.AttributesOp.x_i32elems.fset)["value"] is DenseIntElementsAttr338 assert typing.get_type_hints(test.AttributesOp.x_i32elems.fget)["return"] is DenseIntElementsAttr339 assert type(op.x_i32elems) is typing.get_type_hints(test.AttributesOp.x_i32elems.fget)["return"]340 341 assert typing.get_type_hints(test.AttributesOp.x_i64arr.fset)["value"] is ArrayAttr342 assert typing.get_type_hints(test.AttributesOp.x_i64arr.fget)["return"] is ArrayAttr343 assert type(op.x_i64arr) is typing.get_type_hints(test.AttributesOp.x_i64arr.fget)["return"]344 345 assert typing.get_type_hints(test.AttributesOp.x_i64.fset)["value"] is IntegerAttr346 assert typing.get_type_hints(test.AttributesOp.x_i64.fget)["return"] is IntegerAttr347 assert type(op.x_i64) is typing.get_type_hints(test.AttributesOp.x_i64.fget)["return"]348 349 assert typing.get_type_hints(test.AttributesOp.x_i64elems.fset)["value"] is DenseIntElementsAttr350 assert typing.get_type_hints(test.AttributesOp.x_i64elems.fget)["return"] is DenseIntElementsAttr351 assert type(op.x_i64elems) is typing.get_type_hints(test.AttributesOp.x_i64elems.fget)["return"]352 353 assert typing.get_type_hints(test.AttributesOp.x_i64svecarr.fset)["value"] is ArrayAttr354 assert typing.get_type_hints(test.AttributesOp.x_i64svecarr.fget)["return"] is ArrayAttr355 assert type(op.x_i64svecarr) is typing.get_type_hints(test.AttributesOp.x_i64svecarr.fget)["return"]356 357 assert typing.get_type_hints(test.AttributesOp.x_i8.fset)["value"] is IntegerAttr358 assert typing.get_type_hints(test.AttributesOp.x_i8.fget)["return"] is IntegerAttr359 assert type(op.x_i8) is typing.get_type_hints(test.AttributesOp.x_i8.fget)["return"]360 361 assert typing.get_type_hints(test.AttributesOp.x_idx.fset)["value"] is IntegerAttr362 assert typing.get_type_hints(test.AttributesOp.x_idx.fget)["return"] is IntegerAttr363 assert type(op.x_idx) is typing.get_type_hints(test.AttributesOp.x_idx.fget)["return"]364 365 assert typing.get_type_hints(test.AttributesOp.x_idxelems.fset)["value"] is DenseIntElementsAttr366 assert typing.get_type_hints(test.AttributesOp.x_idxelems.fget)["return"] is DenseIntElementsAttr367 assert type(op.x_idxelems) is typing.get_type_hints(test.AttributesOp.x_idxelems.fget)["return"]368 369 assert typing.get_type_hints(test.AttributesOp.x_idxlistarr.fset)["value"] is ArrayAttr370 assert typing.get_type_hints(test.AttributesOp.x_idxlistarr.fget)["return"] is ArrayAttr371 assert type(op.x_idxlistarr) is typing.get_type_hints(test.AttributesOp.x_idxlistarr.fget)["return"]372 373 assert typing.get_type_hints(test.AttributesOp.x_si16.fset)["value"] is IntegerAttr374 assert typing.get_type_hints(test.AttributesOp.x_si16.fget)["return"] is IntegerAttr375 assert type(op.x_si16) is typing.get_type_hints(test.AttributesOp.x_si16.fget)["return"]376 377 assert typing.get_type_hints(test.AttributesOp.x_si1.fset)["value"] is IntegerAttr378 assert typing.get_type_hints(test.AttributesOp.x_si1.fget)["return"] is IntegerAttr379 assert type(op.x_si1) is typing.get_type_hints(test.AttributesOp.x_si1.fget)["return"]380 381 assert typing.get_type_hints(test.AttributesOp.x_si32.fset)["value"] is IntegerAttr382 assert typing.get_type_hints(test.AttributesOp.x_si32.fget)["return"] is IntegerAttr383 assert type(op.x_si32) is typing.get_type_hints(test.AttributesOp.x_si32.fget)["return"]384 385 assert typing.get_type_hints(test.AttributesOp.x_si64.fset)["value"] is IntegerAttr386 assert typing.get_type_hints(test.AttributesOp.x_si64.fget)["return"] is IntegerAttr387 assert type(op.x_si64) is typing.get_type_hints(test.AttributesOp.x_si64.fget)["return"]388 389 assert typing.get_type_hints(test.AttributesOp.x_si8.fset)["value"] is IntegerAttr390 assert typing.get_type_hints(test.AttributesOp.x_si8.fget)["return"] is IntegerAttr391 assert type(op.x_si8) is typing.get_type_hints(test.AttributesOp.x_si8.fget)["return"]392 393 assert typing.get_type_hints(test.AttributesOp.x_strarr.fset)["value"] is ArrayAttr394 assert typing.get_type_hints(test.AttributesOp.x_strarr.fget)["return"] is ArrayAttr395 assert type(op.x_strarr) is typing.get_type_hints(test.AttributesOp.x_strarr.fget)["return"]396 397 assert typing.get_type_hints(test.AttributesOp.x_str.fset)["value"] is StringAttr398 assert typing.get_type_hints(test.AttributesOp.x_str.fget)["return"] is StringAttr399 assert type(op.x_str) is typing.get_type_hints(test.AttributesOp.x_str.fget)["return"]400 401 assert typing.get_type_hints(test.AttributesOp.x_sym.fset)["value"] is StringAttr402 assert typing.get_type_hints(test.AttributesOp.x_sym.fget)["return"] is StringAttr403 assert type(op.x_sym) is typing.get_type_hints(test.AttributesOp.x_sym.fget)["return"]404 405 assert typing.get_type_hints(test.AttributesOp.x_symrefarr.fset)["value"] is ArrayAttr406 assert typing.get_type_hints(test.AttributesOp.x_symrefarr.fget)["return"] is ArrayAttr407 assert type(op.x_symrefarr) is typing.get_type_hints(test.AttributesOp.x_symrefarr.fget)["return"]408 409 assert typing.get_type_hints(test.AttributesOp.x_symref.fset)["value"] is SymbolRefAttr410 assert typing.get_type_hints(test.AttributesOp.x_symref.fget)["return"] is SymbolRefAttr411 assert type(op.x_symref) is typing.get_type_hints(test.AttributesOp.x_symref.fget)["return"]412 413 assert typing.get_type_hints(test.AttributesOp.x_typearr.fset)["value"] is ArrayAttr414 assert typing.get_type_hints(test.AttributesOp.x_typearr.fget)["return"] is ArrayAttr415 assert type(op.x_typearr) is typing.get_type_hints(test.AttributesOp.x_typearr.fget)["return"]416 417 assert typing.get_type_hints(test.AttributesOp.x_type.fset)["value"] is TypeAttr418 assert typing.get_type_hints(test.AttributesOp.x_type.fget)["return"] is TypeAttr419 assert type(op.x_type) is typing.get_type_hints(test.AttributesOp.x_type.fget)["return"]420 421 assert typing.get_type_hints(test.AttributesOp.x_ui16.fset)["value"] is IntegerAttr422 assert typing.get_type_hints(test.AttributesOp.x_ui16.fget)["return"] is IntegerAttr423 assert type(op.x_ui16) is typing.get_type_hints(test.AttributesOp.x_ui16.fget)["return"]424 425 assert typing.get_type_hints(test.AttributesOp.x_ui1.fset)["value"] is IntegerAttr426 assert typing.get_type_hints(test.AttributesOp.x_ui1.fget)["return"] is IntegerAttr427 assert type(op.x_ui1) is typing.get_type_hints(test.AttributesOp.x_ui1.fget)["return"]428 429 assert typing.get_type_hints(test.AttributesOp.x_ui32.fset)["value"] is IntegerAttr430 assert typing.get_type_hints(test.AttributesOp.x_ui32.fget)["return"] is IntegerAttr431 assert type(op.x_ui32) is typing.get_type_hints(test.AttributesOp.x_ui32.fget)["return"]432 433 assert typing.get_type_hints(test.AttributesOp.x_ui64.fset)["value"] is IntegerAttr434 assert typing.get_type_hints(test.AttributesOp.x_ui64.fget)["return"] is IntegerAttr435 assert type(op.x_ui64) is typing.get_type_hints(test.AttributesOp.x_ui64.fget)["return"]436 437 assert typing.get_type_hints(test.AttributesOp.x_ui8.fset)["value"] is IntegerAttr438 assert typing.get_type_hints(test.AttributesOp.x_ui8.fget)["return"] is IntegerAttr439 assert type(op.x_ui8) is typing.get_type_hints(test.AttributesOp.x_ui8.fget)["return"]440 # fmt: on441 442 443# CHECK-LABEL: TEST: inferReturnTypes444@run445def inferReturnTypes():446 with Context() as ctx, Location.unknown(ctx):447 module = Module.create()448 with InsertionPoint(module.body):449 op = test.InferResultsOp()450 dummy = test.DummyOp()451 452 # CHECK: [Type(i32), Type(i64)]453 iface = InferTypeOpInterface(op)454 print(iface.inferReturnTypes())455 456 # CHECK: [Type(i32), Type(i64)]457 iface_static = InferTypeOpInterface(test.InferResultsOp)458 print(iface.inferReturnTypes())459 460 assert isinstance(iface.opview, test.InferResultsOp)461 assert iface.opview == iface.operation.opview462 463 try:464 iface_static.opview465 except TypeError:466 pass467 else:468 assert False, (469 "not expected to be able to obtain an opview from a static" " interface"470 )471 472 try:473 InferTypeOpInterface(dummy)474 except ValueError:475 pass476 else:477 assert False, "not expected dummy op to implement the interface"478 479 try:480 InferTypeOpInterface(test.DummyOp)481 except ValueError:482 pass483 else:484 assert False, "not expected dummy op class to implement the interface"485 486 487# CHECK-LABEL: TEST: resultTypesDefinedByTraits488@run489def resultTypesDefinedByTraits():490 with Context() as ctx, Location.unknown(ctx):491 module = Module.create()492 with InsertionPoint(module.body):493 inferred = test.InferResultsOp()494 495 # CHECK: i32 i64496 print(inferred.single.type, inferred.doubled.type)497 498 same = test.SameOperandAndResultTypeOp([inferred.results[0]])499 # CHECK-COUNT-2: i32500 print(same.one.type)501 print(same.two.type)502 assert (503 typing.get_type_hints(test.SameOperandAndResultTypeOp.one.fget)[504 "return"505 ]506 is OpResult507 )508 assert type(same.one) is OpResult509 510 first_type_attr = test.FirstAttrDeriveTypeAttrOp(511 inferred.results[1], TypeAttr.get(IndexType.get())512 )513 # CHECK-COUNT-2: index514 print(first_type_attr.one.type)515 print(first_type_attr.two.type)516 517 first_attr = test.FirstAttrDeriveAttrOp(FloatAttr.get(F32Type.get(), 3.14))518 # CHECK-COUNT-3: f32519 print(first_attr.one.type)520 print(first_attr.two.type)521 print(first_attr.three.type)522 523 implied = test.InferResultsImpliedOp()524 # CHECK: i32525 print(implied.integer.type)526 # CHECK: f64527 print(implied.flt.type)528 # CHECK: index529 print(implied.index.type)530 531 # provide the result types to avoid inferring them532 f64 = F64Type.get()533 no_imply = test.InferResultsImpliedOp(results=[f64, f64, f64])534 # CHECK-COUNT-3: f64535 print(no_imply.integer.type, no_imply.flt.type, no_imply.index.type)536 537 no_infer = test.InferResultsOp(results=[F32Type.get(), IndexType.get()])538 # CHECK: f32 index539 print(no_infer.single.type, no_infer.doubled.type)540 541 542# CHECK-LABEL: TEST: testOptionalOperandOp543@run544def testOptionalOperandOp():545 with Context() as ctx, Location.unknown():546 module = Module.create()547 with InsertionPoint(module.body):548 op1 = test.OptionalOperandOp()549 # CHECK: op1.input is None: True550 print(f"op1.input is None: {op1.input is None}")551 assert (552 typing.get_type_hints(test.OptionalOperandOp.input.fget)["return"]553 is Optional[Value]554 )555 assert (556 typing.get_type_hints(test.OptionalOperandOp.result.fget)["return"]557 == OpResult[IntegerType]558 )559 assert type(op1.result) is OpResult560 561 op2 = test.OptionalOperandOp(input=op1)562 # CHECK: op2.input is None: False563 print(f"op2.input is None: {op2.input is None}")564 565 566# CHECK-LABEL: TEST: testCustomAttribute567@run568def testCustomAttribute():569 with Context() as ctx, Location.unknown():570 a = TestAttr.get()571 # CHECK: #python_test.test_attr572 print(a)573 574 # CHECK: python_test.custom_attributed_op {575 # CHECK: #python_test.test_attr576 # CHECK: }577 op2 = test.CustomAttributedOp(a)578 print(f"{op2}")579 580 # CHECK: #python_test.test_attr581 print(f"{op2.test_attr}")582 583 # CHECK: TestAttr(#python_test.test_attr)584 print(repr(op2.test_attr))585 586 # The following cast must not assert.587 b = TestAttr(a)588 589 unit = UnitAttr.get()590 try:591 TestAttr(unit)592 except ValueError as e:593 assert "Cannot cast attribute to TestAttr" in str(e)594 else:595 raise596 597 # The following must trigger a TypeError from our adaptors and must not598 # crash.599 try:600 TestAttr(42)601 except TypeError as e:602 assert "Expected an MLIR object (got 42)" in str(e)603 except ValueError as e:604 assert "Cannot cast attribute to TestAttr (from 42)" in str(e)605 else:606 raise607 608 # The following must trigger a TypeError from pybind (therefore, not609 # checking its message) and must not crash.610 try:611 TestAttr(42, 56)612 except TypeError:613 pass614 else:615 raise616 617 618@run619def testCustomType():620 with Context() as ctx:621 a = TestType.get()622 # CHECK: !python_test.test_type623 print(a)624 625 # The following cast must not assert.626 b = TestType(a)627 # Instance custom types should have typeids628 assert isinstance(b.typeid, TypeID)629 # Subclasses of ir.Type should not have a static_typeid630 # CHECK: 'TestType' object has no attribute 'static_typeid'631 try:632 b.static_typeid633 except AttributeError as e:634 print(e)635 636 i8 = IntegerType.get_signless(8)637 try:638 TestType(i8)639 except ValueError as e:640 assert "Cannot cast type to TestType" in str(e)641 else:642 raise643 644 # The following must trigger a TypeError from our adaptors and must not645 # crash.646 try:647 TestType(42)648 except TypeError as e:649 assert "Expected an MLIR object (got 42)" in str(e)650 except ValueError as e:651 assert "Cannot cast type to TestType (from 42)" in str(e)652 else:653 raise654 655 # The following must trigger a TypeError from pybind (therefore, not656 # checking its message) and must not crash.657 try:658 TestType(42, 56)659 except TypeError:660 pass661 else:662 raise663 664 665@run666# CHECK-LABEL: TEST: testValue667def testValue():668 # Check that Value is a generic class at runtime.669 assert hasattr(Value, "__class_getitem__")670 671 672@run673# CHECK-LABEL: TEST: testTensorValue674def testTensorValue():675 with Context() as ctx, Location.unknown():676 i8 = IntegerType.get_signless(8)677 678 class Tensor(TestTensorValue):679 def __str__(self):680 return super().__str__().replace("Value", "Tensor")681 682 module = Module.create()683 with InsertionPoint(module.body):684 t = tensor.EmptyOp([10, 10], i8).result685 686 # CHECK: Value(%{{.*}} = tensor.empty() : tensor<10x10xi8>)687 print(Value(t))688 689 tt = Tensor(t)690 # CHECK: Tensor(%{{.*}} = tensor.empty() : tensor<10x10xi8>)691 print(tt)692 693 # CHECK: False694 print(tt.is_null())695 696 # Classes of custom types that inherit from concrete types should have697 # static_typeid698 assert isinstance(TestIntegerRankedTensorType.static_typeid, TypeID)699 # And it should be equal to the in-tree concrete type700 assert TestIntegerRankedTensorType.static_typeid == t.type.typeid701 702 d = tensor.EmptyOp([1, 2, 3], IntegerType.get_signless(5)).result703 # CHECK: Value(%{{.*}} = tensor.empty() : tensor<1x2x3xi5>)704 print(d)705 # CHECK: TestTensorValue706 print(repr(d))707 708 709# CHECK-LABEL: TEST: inferReturnTypeComponents710@run711def inferReturnTypeComponents():712 with Context() as ctx, Location.unknown(ctx):713 module = Module.create()714 i32 = IntegerType.get_signless(32)715 with InsertionPoint(module.body):716 resultType = UnrankedTensorType.get(i32)717 operandTypes = [718 RankedTensorType.get([1, 3, 10, 10], i32),719 UnrankedTensorType.get(i32),720 ]721 f = func.FuncOp(722 "test_inferReturnTypeComponents", (operandTypes, [resultType])723 )724 entry_block = Block.create_at_start(f.operation.regions[0], operandTypes)725 with InsertionPoint(entry_block):726 ranked_op = test.InferShapedTypeComponentsOp(727 resultType, entry_block.arguments[0]728 )729 unranked_op = test.InferShapedTypeComponentsOp(730 resultType, entry_block.arguments[1]731 )732 733 # CHECK: has rank: True734 # CHECK: rank: 4735 # CHECK: element type: i32736 # CHECK: shape: [1, 3, 10, 10]737 iface = InferShapedTypeOpInterface(ranked_op)738 shaped_type_components = iface.inferReturnTypeComponents(739 operands=[ranked_op.operand]740 )[0]741 print("has rank:", shaped_type_components.has_rank)742 print("rank:", shaped_type_components.rank)743 print("element type:", shaped_type_components.element_type)744 print("shape:", shaped_type_components.shape)745 746 # CHECK: has rank: False747 # CHECK: rank: None748 # CHECK: element type: i32749 # CHECK: shape: None750 iface = InferShapedTypeOpInterface(unranked_op)751 shaped_type_components = iface.inferReturnTypeComponents(752 operands=[unranked_op.operand]753 )[0]754 print("has rank:", shaped_type_components.has_rank)755 print("rank:", shaped_type_components.rank)756 print("element type:", shaped_type_components.element_type)757 print("shape:", shaped_type_components.shape)758 759 760# CHECK-LABEL: TEST: testCustomTypeTypeCaster761@run762def testCustomTypeTypeCaster():763 with Context() as ctx, Location.unknown():764 a = TestType.get()765 assert a.typeid is not None766 767 b = Type.parse("!python_test.test_type")768 # CHECK: !python_test.test_type769 print(b)770 # CHECK: TestType(!python_test.test_type)771 print(repr(b))772 773 c = TestIntegerRankedTensorType.get([10, 10], 5)774 # CHECK: tensor<10x10xi5>775 print(c)776 # CHECK: TestIntegerRankedTensorType(tensor<10x10xi5>)777 print(repr(c))778 779 # CHECK: Type caster is already registered780 try:781 782 @register_type_caster(c.typeid)783 def type_caster(pytype):784 return TestIntegerRankedTensorType(pytype)785 786 except RuntimeError as e:787 print(e)788 789 # python_test dialect registers a caster for RankedTensorType in its extension (pybind) module.790 # So this one replaces that one (successfully). And then just to be sure we restore the original caster below.791 @register_type_caster(c.typeid, replace=True)792 def type_caster(pytype):793 return RankedTensorType(pytype)794 795 d = tensor.EmptyOp([10, 10], IntegerType.get_signless(5)).result796 # CHECK: tensor<10x10xi5>797 print(d.type)798 # CHECK: ranked tensor type RankedTensorType(tensor<10x10xi5>)799 print("ranked tensor type", repr(d.type))800 801 @register_type_caster(c.typeid, replace=True)802 def type_caster(pytype):803 return TestIntegerRankedTensorType(pytype)804 805 d = tensor.EmptyOp([10, 10], IntegerType.get_signless(5)).result806 # CHECK: tensor<10x10xi5>807 print(d.type)808 # CHECK: TestIntegerRankedTensorType(tensor<10x10xi5>)809 print(repr(d.type))810 811 812# CHECK-LABEL: TEST: testInferTypeOpInterface813@run814def testInferTypeOpInterface():815 with Context() as ctx, Location.unknown(ctx):816 module = Module.create()817 with InsertionPoint(module.body):818 i64 = IntegerType.get_signless(64)819 zero = arith.ConstantOp(i64, 0)820 821 one_operand = test.InferResultsVariadicInputsOp(single=zero, doubled=None)822 # CHECK: i32823 print(one_operand.result.type)824 825 two_operands = test.InferResultsVariadicInputsOp(single=zero, doubled=zero)826 # CHECK: f32827 print(two_operands.result.type)828 829 assert (830 typing.get_type_hints(test.infer_results_variadic_inputs_op)["return"]831 is OpResult832 )833 assert (834 type(test.infer_results_variadic_inputs_op(single=zero, doubled=zero))835 is OpResult836 )837 838 839# CHECK-LABEL: TEST: testVariadicOperandAccess840@run841def testVariadicOperandAccess():842 def values(lst):843 return [str(e) for e in lst]844 845 with Context() as ctx, Location.unknown(ctx):846 module = Module.create()847 with InsertionPoint(module.body):848 i32 = IntegerType.get_signless(32)849 zero = arith.ConstantOp(i32, 0)850 one = arith.ConstantOp(i32, 1)851 two = arith.ConstantOp(i32, 2)852 three = arith.ConstantOp(i32, 3)853 four = arith.ConstantOp(i32, 4)854 855 variadic_operands = test.SameVariadicOperandSizeOp(856 [zero, one], two, [three, four]857 )858 # CHECK: Value(%{{.*}} = arith.constant 2 : i32)859 print(variadic_operands.non_variadic)860 assert (861 typing.get_type_hints(test.SameVariadicOperandSizeOp.non_variadic.fget)[862 "return"863 ]864 is Value865 )866 assert type(variadic_operands.non_variadic) is Value867 868 # CHECK: ['Value(%{{.*}} = arith.constant 0 : i32)', 'Value(%{{.*}} = arith.constant 1 : i32)']869 print(values(variadic_operands.variadic1))870 assert (871 typing.get_type_hints(test.SameVariadicOperandSizeOp.variadic1.fget)[872 "return"873 ]874 is OpOperandList875 )876 assert type(variadic_operands.variadic1) is OpOperandList877 878 # CHECK: ['Value(%{{.*}} = arith.constant 3 : i32)', 'Value(%{{.*}} = arith.constant 4 : i32)']879 print(values(variadic_operands.variadic2))880 assert type(variadic_operands.variadic2) is OpOperandList881 882 assert (883 typing.get_type_hints(test.same_variadic_operand)["return"]884 is test.SameVariadicOperandSizeOp885 )886 assert (887 type(test.same_variadic_operand([zero, one], two, [three, four]))888 is test.SameVariadicOperandSizeOp889 )890 891 892# CHECK-LABEL: TEST: testVariadicResultAccess893@run894def testVariadicResultAccess():895 def types(lst):896 return [e.type for e in lst]897 898 with Context() as ctx, Location.unknown(ctx):899 module = Module.create()900 with InsertionPoint(module.body):901 i = [IntegerType.get_signless(k) for k in range(7)]902 903 # Test Variadic-Fixed-Variadic904 op = test.SameVariadicResultSizeOpVFV([i[0], i[1]], i[2], [i[3], i[4]])905 # CHECK: i2906 print(op.non_variadic.type)907 # CHECK: [IntegerType(i0), IntegerType(i1)]908 print(types(op.variadic1))909 # CHECK: [IntegerType(i3), IntegerType(i4)]910 print(types(op.variadic2))911 912 assert (913 typing.get_type_hints(test.same_variadic_result_vfv)["return"]914 == Union[OpResult, OpResultList, test.SameVariadicResultSizeOpVFV]915 )916 assert (917 type(test.same_variadic_result_vfv([i[0], i[1]], i[2], [i[3], i[4]]))918 is OpResultList919 )920 921 # Test Variadic-Variadic-Variadic922 op = test.SameVariadicResultSizeOpVVV(923 [i[0], i[1]], [i[2], i[3]], [i[4], i[5]]924 )925 # CHECK: [IntegerType(i0), IntegerType(i1)]926 print(types(op.variadic1))927 # CHECK: [IntegerType(i2), IntegerType(i3)]928 print(types(op.variadic2))929 # CHECK: [IntegerType(i4), IntegerType(i5)]930 print(types(op.variadic3))931 932 # Test Fixed-Fixed-Variadic933 op = test.SameVariadicResultSizeOpFFV(i[0], i[1], [i[2], i[3], i[4]])934 # CHECK: i0935 print(op.non_variadic1.type)936 # CHECK: i1937 print(op.non_variadic2.type)938 # CHECK: [IntegerType(i2), IntegerType(i3), IntegerType(i4)]939 print(types(op.variadic))940 assert (941 typing.get_type_hints(test.SameVariadicResultSizeOpFFV.variadic.fget)[942 "return"943 ]944 is OpResultList945 )946 assert type(op.variadic) is OpResultList947 948 # Test Variadic-Variadic-Fixed949 op = test.SameVariadicResultSizeOpVVF(950 [i[0], i[1], i[2]], [i[3], i[4], i[5]], i[6]951 )952 # CHECK: [IntegerType(i0), IntegerType(i1), IntegerType(i2)]953 print(types(op.variadic1))954 # CHECK: [IntegerType(i3), IntegerType(i4), IntegerType(i5)]955 print(types(op.variadic2))956 # CHECK: i6957 print(op.non_variadic.type)958 959 # Test Fixed-Variadic-Fixed-Variadic-Fixed960 op = test.SameVariadicResultSizeOpFVFVF(961 i[0], [i[1], i[2]], i[3], [i[4], i[5]], i[6]962 )963 # CHECK: i0964 print(op.non_variadic1.type)965 # CHECK: [IntegerType(i1), IntegerType(i2)]966 print(types(op.variadic1))967 # CHECK: i3968 print(op.non_variadic2.type)969 # CHECK: [IntegerType(i4), IntegerType(i5)]970 print(types(op.variadic2))971 # CHECK: i6972 print(op.non_variadic3.type)973 974 # Test Fixed-Variadic-Fixed-Variadic-Fixed - Variadic group size 0975 op = test.SameVariadicResultSizeOpFVFVF(i[0], [], i[1], [], i[2])976 # CHECK: i0977 print(op.non_variadic1.type)978 # CHECK: []979 print(types(op.variadic1))980 # CHECK: i1981 print(op.non_variadic2.type)982 # CHECK: []983 print(types(op.variadic2))984 # CHECK: i2985 print(op.non_variadic3.type)986 987 # Test Fixed-Variadic-Fixed-Variadic-Fixed - Variadic group size 1988 op = test.SameVariadicResultSizeOpFVFVF(i[0], [i[1]], i[2], [i[3]], i[4])989 # CHECK: i0990 print(op.non_variadic1.type)991 # CHECK: [IntegerType(i1)]992 print(types(op.variadic1))993 # CHECK: i2994 print(op.non_variadic2.type)995 # CHECK: [IntegerType(i3)]996 print(types(op.variadic2))997 # CHECK: i4998 print(op.non_variadic3.type)999 1000 assert (1001 typing.get_type_hints(test.results_variadic)["return"]1002 == Union[OpResult, OpResultList, test.ResultsVariadicOp]1003 )1004 assert type(test.results_variadic([i[0]])) is OpResult1005 op_res_variadic = test.ResultsVariadicOp([i[0]])1006 assert (1007 typing.get_type_hints(test.ResultsVariadicOp.res.fget)["return"]1008 is OpResultList1009 )1010 assert type(op_res_variadic.res) is OpResultList1011 1012 1013# CHECK-LABEL: TEST: testVariadicAndNormalRegionOp1014@run1015def testVariadicAndNormalRegionOp():1016 with Context() as ctx, Location.unknown(ctx):1017 module = Module.create()1018 with InsertionPoint(module.body):1019 region_op = test.VariadicAndNormalRegionOp(2)1020 assert (1021 typing.get_type_hints(test.VariadicAndNormalRegionOp.region.fget)[1022 "return"1023 ]1024 is Region1025 )1026 assert type(region_op.region) is Region1027 assert (1028 typing.get_type_hints(test.VariadicAndNormalRegionOp.variadic.fget)[1029 "return"1030 ]1031 is RegionSequence1032 )1033 assert type(region_op.variadic) is RegionSequence1034 1035 assert isinstance(region_op.opview, OpView)1036 assert isinstance(region_op.operation.opview, OpView)1037