638 lines · python
1# RUN: %PYTHON %s | FileCheck %s2# Note that this is separate from ir_attributes.py since it depends on numpy,3# and we may want to disable if not available.4 5import gc6from mlir.ir import *7import numpy as np8import weakref9 10 11def run(f):12 print("\nTEST:", f.__name__)13 f()14 gc.collect()15 assert Context._get_live_count() == 016 return f17 18 19################################################################################20# Tests of the array/buffer .get() factory method on unsupported dtype.21################################################################################22 23 24@run25def testGetDenseElementsUnsupported():26 with Context():27 array = np.array([["hello", "goodbye"]])28 try:29 attr = DenseElementsAttr.get(array)30 except ValueError as e:31 # CHECK: unimplemented array format conversion from format:32 print(e)33 34 35# CHECK-LABEL: TEST: testGetDenseElementsUnSupportedTypeOkIfExplicitTypeProvided36@run37def testGetDenseElementsUnSupportedTypeOkIfExplicitTypeProvided():38 with Context():39 array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int64)40 # datetime64 specifically isn't important: it's just a 64-bit type that41 # doesn't have a format under the Python buffer protocol. A more42 # realistic example would be a NumPy extension type like the bfloat1643 # type from the ml_dtypes package, which isn't a dependency of this44 # test.45 attr = DenseElementsAttr.get(46 array.view(np.datetime64), type=IntegerType.get_signless(64)47 )48 # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi64>49 print(attr)50 # CHECK: {{\[}}[1 2 3]51 # CHECK: {{\[}}4 5 6]]52 print(np.array(attr))53 54 55################################################################################56# Tests of the list of attributes .get() factory method57################################################################################58 59 60# CHECK-LABEL: TEST: testGetDenseElementsFromList61@run62def testGetDenseElementsFromList():63 with Context(), Location.unknown():64 attrs = [FloatAttr.get(F64Type.get(), 1.0), FloatAttr.get(F64Type.get(), 2.0)]65 attr = DenseElementsAttr.get(attrs)66 67 # CHECK: dense<[1.000000e+00, 2.000000e+00]> : tensor<2xf64>68 print(attr)69 70 71# CHECK-LABEL: TEST: testGetDenseElementsFromListWithExplicitType72@run73def testGetDenseElementsFromListWithExplicitType():74 with Context(), Location.unknown():75 attrs = [FloatAttr.get(F64Type.get(), 1.0), FloatAttr.get(F64Type.get(), 2.0)]76 shaped_type = ShapedType(Type.parse("tensor<2xf64>"))77 attr = DenseElementsAttr.get(attrs, shaped_type)78 79 # CHECK: dense<[1.000000e+00, 2.000000e+00]> : tensor<2xf64>80 print(attr)81 82 83# CHECK-LABEL: TEST: testGetDenseElementsFromListEmptyList84@run85def testGetDenseElementsFromListEmptyList():86 with Context(), Location.unknown():87 attrs = []88 89 try:90 attr = DenseElementsAttr.get(attrs)91 except ValueError as e:92 # CHECK: Attributes list must be non-empty93 print(e)94 95 96# CHECK-LABEL: TEST: testGetDenseElementsFromListNonAttributeType97@run98def testGetDenseElementsFromListNonAttributeType():99 with Context(), Location.unknown():100 attrs = [1.0]101 102 try:103 attr = DenseElementsAttr.get(attrs)104 except RuntimeError as e:105 # CHECK: Invalid attribute when attempting to create an ArrayAttribute106 print(e)107 108 109# CHECK-LABEL: TEST: testGetDenseElementsFromListMismatchedType110@run111def testGetDenseElementsFromListMismatchedType():112 with Context(), Location.unknown():113 attrs = [FloatAttr.get(F64Type.get(), 1.0), FloatAttr.get(F64Type.get(), 2.0)]114 shaped_type = ShapedType(Type.parse("tensor<2xf32>"))115 116 try:117 attr = DenseElementsAttr.get(attrs, shaped_type)118 except ValueError as e:119 # CHECK: All attributes must be of the same type and match the type parameter120 print(e)121 122 123# CHECK-LABEL: TEST: testGetDenseElementsFromListMixedTypes124@run125def testGetDenseElementsFromListMixedTypes():126 with Context(), Location.unknown():127 attrs = [FloatAttr.get(F64Type.get(), 1.0), FloatAttr.get(F32Type.get(), 2.0)]128 129 try:130 attr = DenseElementsAttr.get(attrs)131 except ValueError as e:132 # CHECK: All attributes must be of the same type and match the type parameter133 print(e)134 135 136################################################################################137# Splats.138################################################################################139 140 141# CHECK-LABEL: TEST: testGetDenseElementsSplatInt142@run143def testGetDenseElementsSplatInt():144 with Context(), Location.unknown():145 t = IntegerType.get_signless(32)146 element = IntegerAttr.get(t, 555)147 shaped_type = RankedTensorType.get((2, 3, 4), t)148 attr = DenseElementsAttr.get_splat(shaped_type, element)149 # CHECK: dense<555> : tensor<2x3x4xi32>150 print(attr)151 # CHECK: is_splat: True152 print("is_splat:", attr.is_splat)153 154 # CHECK: splat_value: IntegerAttr(555 : i32)155 splat_value = attr.get_splat_value()156 print("splat_value:", repr(splat_value))157 assert splat_value == element158 159 160# CHECK-LABEL: TEST: testGetDenseElementsSplatFloat161@run162def testGetDenseElementsSplatFloat():163 with Context(), Location.unknown():164 t = F32Type.get()165 element = FloatAttr.get(t, 1.2)166 shaped_type = RankedTensorType.get((2, 3, 4), t)167 attr = DenseElementsAttr.get_splat(shaped_type, element)168 # CHECK: dense<1.200000e+00> : tensor<2x3x4xf32>169 print(attr)170 assert attr.get_splat_value() == element171 172 173# CHECK-LABEL: TEST: testGetDenseElementsSplatErrors174@run175def testGetDenseElementsSplatErrors():176 with Context(), Location.unknown():177 t = F32Type.get()178 other_t = F64Type.get()179 element = FloatAttr.get(t, 1.2)180 other_element = FloatAttr.get(other_t, 1.2)181 shaped_type = RankedTensorType.get((2, 3, 4), t)182 dynamic_shaped_type = UnrankedTensorType.get(t)183 non_shaped_type = t184 185 try:186 attr = DenseElementsAttr.get_splat(non_shaped_type, element)187 except ValueError as e:188 # CHECK: Expected a static ShapedType for the shaped_type parameter: Type(f32)189 print(e)190 191 try:192 attr = DenseElementsAttr.get_splat(dynamic_shaped_type, element)193 except ValueError as e:194 # CHECK: Expected a static ShapedType for the shaped_type parameter: Type(tensor<*xf32>)195 print(e)196 197 try:198 attr = DenseElementsAttr.get_splat(shaped_type, other_element)199 except ValueError as e:200 # CHECK: Shaped element type and attribute type must be equal: shaped=Type(tensor<2x3x4xf32>), element=Attribute(1.200000e+00 : f64)201 print(e)202 203 204# CHECK-LABEL: TEST: testRepeatedValuesSplat205@run206def testRepeatedValuesSplat():207 with Context():208 array = np.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]], dtype=np.float32)209 attr = DenseElementsAttr.get(array)210 # CHECK: dense<1.000000e+00> : tensor<2x3xf32>211 print(attr)212 # CHECK: is_splat: True213 print("is_splat:", attr.is_splat)214 # CHECK{LITERAL}: [[1. 1. 1.]215 # CHECK{LITERAL}: [1. 1. 1.]]216 print(np.array(attr))217 218 219# CHECK-LABEL: TEST: testNonSplat220@run221def testNonSplat():222 with Context():223 array = np.array([2.0, 1.0, 1.0], dtype=np.float32)224 attr = DenseElementsAttr.get(array)225 # CHECK: is_splat: False226 print("is_splat:", attr.is_splat)227 228 229################################################################################230# Tests of the array/buffer .get() factory method, in all of its permutations.231################################################################################232 233### explicitly provided types234 235 236@run237def testGetDenseElementsBF16():238 with Context():239 array = np.array([[2, 4, 8], [16, 32, 64]], dtype=np.uint16)240 attr = DenseElementsAttr.get(array, type=BF16Type.get())241 # Note: These values don't mean much since just bit-casting. But they242 # shouldn't change.243 # CHECK: dense<{{\[}}[1.836710e-40, 3.673420e-40, 7.346840e-40], [1.469370e-39, 2.938740e-39, 5.877470e-39]]> : tensor<2x3xbf16>244 print(attr)245 246 247@run248def testGetDenseElementsInteger4():249 with Context():250 array = np.array([[2, 4, 7], [-2, -4, -8]], dtype=np.int8)251 attr = DenseElementsAttr.get(array, type=IntegerType.get_signless(4))252 # Note: These values don't mean much since just bit-casting. But they253 # shouldn't change.254 # CHECK: dense<{{\[}}[2, 4, 7], [-2, -4, -8]]> : tensor<2x3xi4>255 print(attr)256 257 258@run259def testGetDenseElementsBool():260 with Context():261 bool_array = np.array([[1, 0, 1], [0, 1, 0]], dtype=np.bool_)262 array = np.packbits(bool_array, axis=None, bitorder="little")263 attr = DenseElementsAttr.get(264 array, type=IntegerType.get_signless(1), shape=bool_array.shape265 )266 # CHECK: dense<{{\[}}[true, false, true], [false, true, false]]> : tensor<2x3xi1>267 print(attr)268 269 270@run271def testGetDenseElementsBoolSplat():272 with Context():273 zero = np.array(0, dtype=np.uint8)274 one = np.array(255, dtype=np.uint8)275 print(one)276 # CHECK: dense<false> : tensor<4x2x5xi1>277 print(278 DenseElementsAttr.get(279 zero, type=IntegerType.get_signless(1), shape=(4, 2, 5)280 )281 )282 # CHECK: dense<true> : tensor<4x2x5xi1>283 print(284 DenseElementsAttr.get(285 one, type=IntegerType.get_signless(1), shape=(4, 2, 5)286 )287 )288 289 290### float and double arrays.291 292 293# CHECK-LABEL: TEST: testGetDenseElementsF16294@run295def testGetDenseElementsF16():296 with Context():297 array = np.array([[2.0, 4.0, 8.0], [16.0, 32.0, 64.0]], dtype=np.float16)298 attr = DenseElementsAttr.get(array)299 # CHECK: dense<{{\[}}[2.000000e+00, 4.000000e+00, 8.000000e+00], [1.600000e+01, 3.200000e+01, 6.400000e+01]]> : tensor<2x3xf16>300 print(attr)301 # CHECK: {{\[}}[ 2. 4. 8.]302 # CHECK: {{\[}}16. 32. 64.]]303 print(np.array(attr))304 305 306# CHECK-LABEL: TEST: testGetDenseElementsF32307@run308def testGetDenseElementsF32():309 with Context():310 array = np.array([[1.1, 2.2, 3.3], [4.4, 5.5, 6.6]], dtype=np.float32)311 attr = DenseElementsAttr.get(array)312 # CHECK: dense<{{\[}}[1.100000e+00, 2.200000e+00, 3.300000e+00], [4.400000e+00, 5.500000e+00, 6.600000e+00]]> : tensor<2x3xf32>313 print(attr)314 # CHECK: {{\[}}[1.1 2.2 3.3]315 # CHECK: {{\[}}4.4 5.5 6.6]]316 print(np.array(attr))317 318 319# CHECK-LABEL: TEST: testGetDenseElementsF64320@run321def testGetDenseElementsF64():322 with Context():323 array = np.array([[1.1, 2.2, 3.3], [4.4, 5.5, 6.6]], dtype=np.float64)324 attr = DenseElementsAttr.get(array)325 # CHECK: dense<{{\[}}[1.100000e+00, 2.200000e+00, 3.300000e+00], [4.400000e+00, 5.500000e+00, 6.600000e+00]]> : tensor<2x3xf64>326 print(attr)327 # CHECK: {{\[}}[1.1 2.2 3.3]328 # CHECK: {{\[}}4.4 5.5 6.6]]329 print(np.array(attr))330 331 332### 1 bit/boolean integer arrays333# CHECK-LABEL: TEST: testGetDenseElementsI1Signless334@run335def testGetDenseElementsI1Signless():336 with Context():337 array = np.array([True], dtype=np.bool_)338 attr = DenseElementsAttr.get(array)339 # CHECK: dense<true> : tensor<1xi1>340 print(attr)341 # CHECK{LITERAL}: [ True]342 print(np.array(attr))343 344 array = np.array([[True, False, True], [True, True, False]], dtype=np.bool_)345 attr = DenseElementsAttr.get(array)346 # CHECK{LITERAL}: dense<[[true, false, true], [true, true, false]]> : tensor<2x3xi1>347 print(attr)348 # CHECK{LITERAL}: [[ True False True]349 # CHECK{LITERAL}: [ True True False]]350 print(np.array(attr))351 352 array = np.array(353 [[True, True, False, False], [True, False, True, False]], dtype=np.bool_354 )355 attr = DenseElementsAttr.get(array)356 # CHECK{LITERAL}: dense<[[true, true, false, false], [true, false, true, false]]> : tensor<2x4xi1>357 print(attr)358 # CHECK{LITERAL}: [[ True True False False]359 # CHECK{LITERAL}: [ True False True False]]360 print(np.array(attr))361 362 array = np.array(363 [364 [True, True, False, False],365 [True, False, True, False],366 [False, False, False, False],367 [True, True, True, True],368 [True, False, False, True],369 ],370 dtype=np.bool_,371 )372 attr = DenseElementsAttr.get(array)373 # CHECK{LITERAL}: dense<[[true, true, false, false], [true, false, true, false], [false, false, false, false], [true, true, true, true], [true, false, false, true]]> : tensor<5x4xi1>374 print(attr)375 # CHECK{LITERAL}: [[ True True False False]376 # CHECK{LITERAL}: [ True False True False]377 # CHECK{LITERAL}: [False False False False]378 # CHECK{LITERAL}: [ True True True True]379 # CHECK{LITERAL}: [ True False False True]]380 print(np.array(attr))381 382 array = np.array(383 [384 [True, True, False, False, True, True, False, False, False],385 [False, False, False, True, False, True, True, False, True],386 ],387 dtype=np.bool_,388 )389 attr = DenseElementsAttr.get(array)390 # CHECK{LITERAL}: dense<[[true, true, false, false, true, true, false, false, false], [false, false, false, true, false, true, true, false, true]]> : tensor<2x9xi1>391 print(attr)392 # CHECK{LITERAL}: [[ True True False False True True False False False]393 # CHECK{LITERAL}: [False False False True False True True False True]]394 print(np.array(attr))395 396 array = np.array([], dtype=np.bool_)397 attr = DenseElementsAttr.get(array)398 # CHECK: dense<> : tensor<0xi1>399 print(attr)400 # CHECK{LITERAL}: []401 print(np.array(attr))402 403 404### 16 bit integer arrays405# CHECK-LABEL: TEST: testGetDenseElementsI16Signless406@run407def testGetDenseElementsI16Signless():408 with Context():409 array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int16)410 attr = DenseElementsAttr.get(array)411 # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi16>412 print(attr)413 # CHECK: {{\[}}[1 2 3]414 # CHECK: {{\[}}4 5 6]]415 print(np.array(attr))416 417 418# CHECK-LABEL: TEST: testGetDenseElementsUI16Signless419@run420def testGetDenseElementsUI16Signless():421 with Context():422 array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uint16)423 attr = DenseElementsAttr.get(array)424 # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi16>425 print(attr)426 # CHECK: {{\[}}[1 2 3]427 # CHECK: {{\[}}4 5 6]]428 print(np.array(attr))429 430 431# CHECK-LABEL: TEST: testGetDenseElementsI16432@run433def testGetDenseElementsI16():434 with Context():435 array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int16)436 attr = DenseElementsAttr.get(array, signless=False)437 # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xsi16>438 print(attr)439 # CHECK: {{\[}}[1 2 3]440 # CHECK: {{\[}}4 5 6]]441 print(np.array(attr))442 443 444# CHECK-LABEL: TEST: testGetDenseElementsUI16445@run446def testGetDenseElementsUI16():447 with Context():448 array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uint16)449 attr = DenseElementsAttr.get(array, signless=False)450 # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xui16>451 print(attr)452 # CHECK: {{\[}}[1 2 3]453 # CHECK: {{\[}}4 5 6]]454 print(np.array(attr))455 456 457### 32 bit integer arrays458# CHECK-LABEL: TEST: testGetDenseElementsI32Signless459@run460def testGetDenseElementsI32Signless():461 with Context():462 array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int32)463 attr = DenseElementsAttr.get(array)464 # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32>465 print(attr)466 # CHECK: {{\[}}[1 2 3]467 # CHECK: {{\[}}4 5 6]]468 print(np.array(attr))469 470 471# CHECK-LABEL: TEST: testGetDenseElementsUI32Signless472@run473def testGetDenseElementsUI32Signless():474 with Context():475 array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uint32)476 attr = DenseElementsAttr.get(array)477 # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32>478 print(attr)479 # CHECK: {{\[}}[1 2 3]480 # CHECK: {{\[}}4 5 6]]481 print(np.array(attr))482 483 484# CHECK-LABEL: TEST: testGetDenseElementsI32485@run486def testGetDenseElementsI32():487 with Context():488 array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int32)489 attr = DenseElementsAttr.get(array, signless=False)490 # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xsi32>491 print(attr)492 # CHECK: {{\[}}[1 2 3]493 # CHECK: {{\[}}4 5 6]]494 print(np.array(attr))495 496 497# CHECK-LABEL: TEST: testGetDenseElementsUI32498@run499def testGetDenseElementsUI32():500 with Context():501 array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uint32)502 attr = DenseElementsAttr.get(array, signless=False)503 # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xui32>504 print(attr)505 # CHECK: {{\[}}[1 2 3]506 # CHECK: {{\[}}4 5 6]]507 print(np.array(attr))508 509 510## 64bit integer arrays511# CHECK-LABEL: TEST: testGetDenseElementsI64Signless512@run513def testGetDenseElementsI64Signless():514 with Context():515 array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int64)516 attr = DenseElementsAttr.get(array)517 # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi64>518 print(attr)519 # CHECK: {{\[}}[1 2 3]520 # CHECK: {{\[}}4 5 6]]521 print(np.array(attr))522 523 524# CHECK-LABEL: TEST: testGetDenseElementsUI64Signless525@run526def testGetDenseElementsUI64Signless():527 with Context():528 array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uint64)529 attr = DenseElementsAttr.get(array)530 # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi64>531 print(attr)532 # CHECK: {{\[}}[1 2 3]533 # CHECK: {{\[}}4 5 6]]534 print(np.array(attr))535 536 537# CHECK-LABEL: TEST: testGetDenseElementsI64538@run539def testGetDenseElementsI64():540 with Context():541 array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int64)542 attr = DenseElementsAttr.get(array, signless=False)543 # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xsi64>544 print(attr)545 # CHECK: {{\[}}[1 2 3]546 # CHECK: {{\[}}4 5 6]]547 print(np.array(attr))548 549 550# CHECK-LABEL: TEST: testGetDenseElementsUI64551@run552def testGetDenseElementsUI64():553 with Context():554 array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uint64)555 attr = DenseElementsAttr.get(array, signless=False)556 # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xui64>557 print(attr)558 # CHECK: {{\[}}[1 2 3]559 # CHECK: {{\[}}4 5 6]]560 print(np.array(attr))561 562 563# CHECK-LABEL: TEST: testGetDenseElementsIndex564@run565def testGetDenseElementsIndex():566 with Context(), Location.unknown():567 idx_type = IndexType.get()568 array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int64)569 attr = DenseElementsAttr.get(array, type=idx_type)570 # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xindex>571 print(attr)572 arr = np.array(attr)573 # CHECK: {{\[}}[1 2 3]574 # CHECK: {{\[}}4 5 6]]575 print(arr)576 # CHECK: True577 print(arr.dtype == np.int64)578 array = np.array([1, 2, 3], dtype=np.int64)579 attr = DenseIntElementsAttr.get(array, type=VectorType.get([3], idx_type))580 # CHECK: [1, 2, 3]581 print(list(DenseIntElementsAttr(attr)))582 583 584# CHECK-LABEL: TEST: testGetDenseResourceElementsAttr585@run586def testGetDenseResourceElementsAttr():587 def on_delete(_):588 print("BACKING MEMORY DELETED")589 590 context = Context()591 mview = memoryview(np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int32))592 ref = weakref.ref(mview, on_delete)593 594 def test_attribute(context, mview):595 with context, Location.unknown():596 element_type = IntegerType.get_signless(32)597 tensor_type = RankedTensorType.get((2, 3), element_type)598 resource = DenseResourceElementsAttr.get_from_buffer(599 mview, "from_py", tensor_type600 )601 module = Module.parse("module {}")602 module.operation.attributes["test.resource"] = resource603 # CHECK: test.resource = dense_resource<from_py> : tensor<2x3xi32>604 # CHECK: from_py: "0x04000000010000000200000003000000040000000500000006000000"605 print(module)606 607 # Verifies type casting.608 # CHECK: dense_resource<from_py> : tensor<2x3xi32>609 print(610 DenseResourceElementsAttr(module.operation.attributes["test.resource"])611 )612 613 test_attribute(context, mview)614 mview = None615 gc.collect()616 # CHECK: FREEING CONTEXT617 print("FREEING CONTEXT")618 context = None619 gc.collect()620 # CHECK: BACKING MEMORY DELETED621 # CHECK: EXIT FUNCTION622 print("EXIT FUNCTION")623 624 625# CHECK-LABEL: TEST: testDanglingResource626print("TEST: testDanglingResource")627# see https://github.com/llvm/llvm-project/pull/149414, https://github.com/llvm/llvm-project/pull/150137, https://github.com/llvm/llvm-project/pull/150561628# This error occurs only when there is an alive context with a DenseResourceElementsAttr629# in the end of the program, so we put it here without an encapsulating function.630ctx = Context()631 632with ctx, Location.unknown():633 DenseResourceElementsAttr.get_from_buffer(634 memoryview(np.array([1, 2, 3])),635 "some_resource",636 RankedTensorType.get((3,), IntegerType.get_signed(32)),637 )638