136 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3from mlir.ir import *4from mlir.dialects import arith5from mlir.dialects import builtin6from mlir.dialects import func7 8 9def constructAndPrintInModule(f):10 print("\nTEST:", f.__name__)11 with Context(), Location.unknown():12 module = Module.create()13 with InsertionPoint(module.body):14 f()15 print(module)16 return f17 18 19# CHECK-LABEL: TEST: testConstantOp20 21 22@constructAndPrintInModule23def testConstantOp():24 c1 = arith.ConstantOp(IntegerType.get_signless(32), 42)25 c2 = arith.ConstantOp(IntegerType.get_signless(64), 100)26 c3 = arith.ConstantOp(F32Type.get(), 3.14)27 c4 = arith.ConstantOp(F64Type.get(), 1.23)28 # CHECK: 4229 print(c1.literal_value)30 31 # CHECK: 10032 print(c2.literal_value)33 34 # CHECK: 3.14000010490417535 print(c3.literal_value)36 37 # CHECK: 1.2338 print(c4.literal_value)39 40 41# CHECK: = arith.constant 42 : i3242# CHECK: = arith.constant 100 : i6443# CHECK: = arith.constant 3.140000e+00 : f3244# CHECK: = arith.constant 1.230000e+00 : f6445 46 47# CHECK-LABEL: TEST: testVectorConstantOp48@constructAndPrintInModule49def testVectorConstantOp():50 int_type = IntegerType.get_signless(32)51 vec_type = VectorType.get([2, 2], int_type)52 c1 = arith.ConstantOp(53 vec_type, DenseElementsAttr.get_splat(vec_type, IntegerAttr.get(int_type, 42))54 )55 try:56 print(c1.literal_value)57 except ValueError as e:58 assert "only integer and float constants have literal values" in str(e)59 else:60 assert False61 62 63# CHECK: = arith.constant dense<42> : vector<2x2xi32>64 65 66# CHECK-LABEL: TEST: testConstantIndexOp67@constructAndPrintInModule68def testConstantIndexOp():69 c1 = arith.ConstantOp.create_index(10)70 # CHECK: 1071 print(c1.literal_value)72 73 74# CHECK: = arith.constant 10 : index75 76 77# CHECK-LABEL: TEST: testFunctionCalls78@constructAndPrintInModule79def testFunctionCalls():80 foo = func.FuncOp("foo", ([], []))81 foo.sym_visibility = StringAttr.get("private")82 bar = func.FuncOp("bar", ([], [IndexType.get()]))83 bar.sym_visibility = StringAttr.get("private")84 qux = func.FuncOp("qux", ([], [F32Type.get()]))85 qux.sym_visibility = StringAttr.get("private")86 87 con = func.ConstantOp(qux.type, qux.sym_name.value)88 assert con.type == qux.type89 90 with InsertionPoint(func.FuncOp("caller", ([], [])).add_entry_block()):91 func.CallOp(foo, [])92 func.CallOp([IndexType.get()], "bar", [])93 func.CallOp([F32Type.get()], FlatSymbolRefAttr.get("qux"), [])94 func.ReturnOp([])95 96 97# CHECK: func private @foo()98# CHECK: func private @bar() -> index99# CHECK: func private @qux() -> f32100# CHECK: %f = func.constant @qux : () -> f32101# CHECK: func @caller() {102# CHECK: call @foo() : () -> ()103# CHECK: %0 = call @bar() : () -> index104# CHECK: %1 = call @qux() : () -> f32105# CHECK: return106# CHECK: }107 108 109# CHECK-LABEL: TEST: testFunctionArgAttrs110@constructAndPrintInModule111def testFunctionArgAttrs():112 foo = func.FuncOp("foo", ([F32Type.get()], []))113 foo.sym_visibility = StringAttr.get("private")114 foo2 = func.FuncOp("foo2", ([F32Type.get(), F32Type.get()], []))115 foo2.sym_visibility = StringAttr.get("private")116 117 empty_attr = DictAttr.get({})118 test_attr = DictAttr.get({"test.foo": StringAttr.get("bar")})119 test_attr2 = DictAttr.get({"test.baz": StringAttr.get("qux")})120 121 assert len(foo.arg_attrs) == 1122 assert foo.arg_attrs[0] == empty_attr123 124 foo.arg_attrs = [test_attr]125 assert foo.arg_attrs[0]["test.foo"] == StringAttr.get("bar")126 127 assert len(foo2.arg_attrs) == 2128 assert foo2.arg_attrs == ArrayAttr.get([empty_attr, empty_attr])129 130 foo2.arg_attrs = [empty_attr, test_attr2]131 assert foo2.arg_attrs == ArrayAttr.get([empty_attr, test_attr2])132 133 134# CHECK: func private @foo(f32 {test.foo = "bar"})135# CHECK: func private @foo2(f32, f32 {test.baz = "qux"})136