1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
|
# RUN: %PYTHON %s | FileCheck %s
from mlir.ir import *
from mlir.dialects import arith
from mlir.dialects import builtin
from mlir.dialects import func
def constructAndPrintInModule(f):
print("\nTEST:", f.__name__)
with Context(), Location.unknown():
module = Module.create()
with InsertionPoint(module.body):
f()
print(module)
return f
# CHECK-LABEL: TEST: testConstantOp
@constructAndPrintInModule
def testConstantOp():
c1 = arith.ConstantOp(IntegerType.get_signless(32), 42)
c2 = arith.ConstantOp(IntegerType.get_signless(64), 100)
c3 = arith.ConstantOp(F32Type.get(), 3.14)
c4 = arith.ConstantOp(F64Type.get(), 1.23)
# CHECK: 42
print(c1.literal_value)
# CHECK: 100
print(c2.literal_value)
# CHECK: 3.140000104904175
print(c3.literal_value)
# CHECK: 1.23
print(c4.literal_value)
# CHECK: = arith.constant 42 : i32
# CHECK: = arith.constant 100 : i64
# CHECK: = arith.constant 3.140000e+00 : f32
# CHECK: = arith.constant 1.230000e+00 : f64
# CHECK-LABEL: TEST: testVectorConstantOp
@constructAndPrintInModule
def testVectorConstantOp():
int_type = IntegerType.get_signless(32)
vec_type = VectorType.get([2, 2], int_type)
c1 = arith.ConstantOp(
vec_type, DenseElementsAttr.get_splat(vec_type, IntegerAttr.get(int_type, 42))
)
try:
print(c1.literal_value)
except ValueError as e:
assert "only integer and float constants have literal values" in str(e)
else:
assert False
# CHECK: = arith.constant dense<42> : vector<2x2xi32>
# CHECK-LABEL: TEST: testConstantIndexOp
@constructAndPrintInModule
def testConstantIndexOp():
c1 = arith.ConstantOp.create_index(10)
# CHECK: 10
print(c1.literal_value)
# CHECK: = arith.constant 10 : index
# CHECK-LABEL: TEST: testFunctionCalls
@constructAndPrintInModule
def testFunctionCalls():
foo = func.FuncOp("foo", ([], []))
foo.sym_visibility = StringAttr.get("private")
bar = func.FuncOp("bar", ([], [IndexType.get()]))
bar.sym_visibility = StringAttr.get("private")
qux = func.FuncOp("qux", ([], [F32Type.get()]))
qux.sym_visibility = StringAttr.get("private")
con = func.ConstantOp(qux.type, qux.sym_name.value)
assert con.type == qux.type
with InsertionPoint(func.FuncOp("caller", ([], [])).add_entry_block()):
func.CallOp(foo, [])
func.CallOp([IndexType.get()], "bar", [])
func.CallOp([F32Type.get()], FlatSymbolRefAttr.get("qux"), [])
func.ReturnOp([])
# CHECK: func private @foo()
# CHECK: func private @bar() -> index
# CHECK: func private @qux() -> f32
# CHECK: %f = func.constant @qux : () -> f32
# CHECK: func @caller() {
# CHECK: call @foo() : () -> ()
# CHECK: %0 = call @bar() : () -> index
# CHECK: %1 = call @qux() : () -> f32
# CHECK: return
# CHECK: }
# CHECK-LABEL: TEST: testFunctionArgAttrs
@constructAndPrintInModule
def testFunctionArgAttrs():
foo = func.FuncOp("foo", ([F32Type.get()], []))
foo.sym_visibility = StringAttr.get("private")
foo2 = func.FuncOp("foo2", ([F32Type.get(), F32Type.get()], []))
foo2.sym_visibility = StringAttr.get("private")
empty_attr = DictAttr.get({})
test_attr = DictAttr.get({"test.foo": StringAttr.get("bar")})
test_attr2 = DictAttr.get({"test.baz": StringAttr.get("qux")})
assert len(foo.arg_attrs) == 1
assert foo.arg_attrs[0] == empty_attr
foo.arg_attrs = [test_attr]
assert foo.arg_attrs[0]["test.foo"] == StringAttr.get("bar")
assert len(foo2.arg_attrs) == 2
assert foo2.arg_attrs == ArrayAttr.get([empty_attr, empty_attr])
foo2.arg_attrs = [empty_attr, test_attr2]
assert foo2.arg_attrs == ArrayAttr.get([empty_attr, test_attr2])
# CHECK: func private @foo(f32 {test.foo = "bar"})
# CHECK: func private @foo2(f32, f32 {test.baz = "qux"})
|
