xref: /llvm-project/mlir/test/python/dialects/index_dialect.py (revision eb861acd49e4c6777e2fe09dd2004ac1f4731bba)
1# RUN: %PYTHON %s | FileCheck %s
2
3from mlir.ir import *
4from mlir.dialects import index, arith
5
6
7def run(f):
8    print("\nTEST:", f.__name__)
9    with Context() as ctx, Location.unknown():
10        module = Module.create()
11        with InsertionPoint(module.body):
12            f(ctx)
13        print(module)
14
15
16# CHECK-LABEL: TEST: testConstantOp
17@run
18def testConstantOp(ctx):
19    a = index.ConstantOp(value=42)
20    # CHECK: %{{.*}} = index.constant 42
21
22
23# CHECK-LABEL: TEST: testBoolConstantOp
24@run
25def testBoolConstantOp(ctx):
26    a = index.BoolConstantOp(value=True)
27    # CHECK: %{{.*}} = index.bool.constant true
28
29
30# CHECK-LABEL: TEST: testAndOp
31@run
32def testAndOp(ctx):
33    a = index.ConstantOp(value=42)
34    r = index.AndOp(a, a)
35    # CHECK: %{{.*}} = index.and %{{.*}}, %{{.*}}
36
37
38# CHECK-LABEL: TEST: testOrOp
39@run
40def testOrOp(ctx):
41    a = index.ConstantOp(value=42)
42    r = index.OrOp(a, a)
43    # CHECK: %{{.*}} = index.or %{{.*}}, %{{.*}}
44
45
46# CHECK-LABEL: TEST: testXOrOp
47@run
48def testXOrOp(ctx):
49    a = index.ConstantOp(value=42)
50    r = index.XOrOp(a, a)
51    # CHECK: %{{.*}} = index.xor %{{.*}}, %{{.*}}
52
53
54# CHECK-LABEL: TEST: testCastSOp
55@run
56def testCastSOp(ctx):
57    a = index.ConstantOp(value=42)
58    b = arith.ConstantOp(value=23, result=IntegerType.get_signless(64))
59    c = index.CastSOp(input=a, output=IntegerType.get_signless(32))
60    d = index.CastSOp(input=b, output=IndexType.get())
61    # CHECK: %{{.*}} = index.casts %{{.*}} : index to i32
62    # CHECK: %{{.*}} = index.casts %{{.*}} : i64 to index
63
64
65# CHECK-LABEL: TEST: testCastUOp
66@run
67def testCastUOp(ctx):
68    a = index.ConstantOp(value=42)
69    b = arith.ConstantOp(value=23, result=IntegerType.get_signless(64))
70    c = index.CastUOp(input=a, output=IntegerType.get_signless(32))
71    d = index.CastUOp(input=b, output=IndexType.get())
72    # CHECK: %{{.*}} = index.castu %{{.*}} : index to i32
73    # CHECK: %{{.*}} = index.castu %{{.*}} : i64 to index
74
75
76# CHECK-LABEL: TEST: testCeilDivSOp
77@run
78def testCeilDivSOp(ctx):
79    a = index.ConstantOp(value=42)
80    r = index.CeilDivSOp(a, a)
81    # CHECK: %{{.*}} = index.ceildivs %{{.*}}, %{{.*}}
82
83
84# CHECK-LABEL: TEST: testCeilDivUOp
85@run
86def testCeilDivUOp(ctx):
87    a = index.ConstantOp(value=42)
88    r = index.CeilDivUOp(a, a)
89    # CHECK: %{{.*}} = index.ceildivu %{{.*}}, %{{.*}}
90
91
92# CHECK-LABEL: TEST: testCmpOp
93@run
94def testCmpOp(ctx):
95    a = index.ConstantOp(value=42)
96    b = index.ConstantOp(value=23)
97    pred = AttrBuilder.get("IndexCmpPredicateAttr")("slt", context=ctx)
98    r = index.CmpOp(pred, lhs=a, rhs=b)
99    # CHECK: %{{.*}} = index.cmp slt(%{{.*}}, %{{.*}})
100
101
102# CHECK-LABEL: TEST: testAddOp
103@run
104def testAddOp(ctx):
105    a = index.ConstantOp(value=42)
106    r = index.AddOp(a, a)
107    # CHECK: %{{.*}} = index.add %{{.*}}, %{{.*}}
108
109
110# CHECK-LABEL: TEST: testSubOp
111@run
112def testSubOp(ctx):
113    a = index.ConstantOp(value=42)
114    r = index.SubOp(a, a)
115    # CHECK: %{{.*}} = index.sub %{{.*}}, %{{.*}}
116
117
118# CHECK-LABEL: TEST: testMulOp
119@run
120def testMulOp(ctx):
121    a = index.ConstantOp(value=42)
122    r = index.MulOp(a, a)
123    # CHECK: %{{.*}} = index.mul %{{.*}}, %{{.*}}
124
125
126# CHECK-LABEL: TEST: testDivSOp
127@run
128def testDivSOp(ctx):
129    a = index.ConstantOp(value=42)
130    r = index.DivSOp(a, a)
131    # CHECK: %{{.*}} = index.divs %{{.*}}, %{{.*}}
132
133
134# CHECK-LABEL: TEST: testDivUOp
135@run
136def testDivUOp(ctx):
137    a = index.ConstantOp(value=42)
138    r = index.DivUOp(a, a)
139    # CHECK: %{{.*}} = index.divu %{{.*}}, %{{.*}}
140
141
142# CHECK-LABEL: TEST: testFloorDivSOp
143@run
144def testFloorDivSOp(ctx):
145    a = index.ConstantOp(value=42)
146    r = index.FloorDivSOp(a, a)
147    # CHECK: %{{.*}} = index.floordivs %{{.*}}, %{{.*}}
148
149
150# CHECK-LABEL: TEST: testMaxSOp
151@run
152def testMaxSOp(ctx):
153    a = index.ConstantOp(value=42)
154    b = index.ConstantOp(value=23)
155    r = index.MaxSOp(a, b)
156    # CHECK: %{{.*}} = index.maxs %{{.*}}, %{{.*}}
157
158
159# CHECK-LABEL: TEST: testMaxUOp
160@run
161def testMaxUOp(ctx):
162    a = index.ConstantOp(value=42)
163    b = index.ConstantOp(value=23)
164    r = index.MaxUOp(a, b)
165    # CHECK: %{{.*}} = index.maxu %{{.*}}, %{{.*}}
166
167
168# CHECK-LABEL: TEST: testMinSOp
169@run
170def testMinSOp(ctx):
171    a = index.ConstantOp(value=42)
172    b = index.ConstantOp(value=23)
173    r = index.MinSOp(a, b)
174    # CHECK: %{{.*}} = index.mins %{{.*}}, %{{.*}}
175
176
177# CHECK-LABEL: TEST: testMinUOp
178@run
179def testMinUOp(ctx):
180    a = index.ConstantOp(value=42)
181    b = index.ConstantOp(value=23)
182    r = index.MinUOp(a, b)
183    # CHECK: %{{.*}} = index.minu %{{.*}}, %{{.*}}
184
185
186# CHECK-LABEL: TEST: testRemSOp
187@run
188def testRemSOp(ctx):
189    a = index.ConstantOp(value=42)
190    b = index.ConstantOp(value=23)
191    r = index.RemSOp(a, b)
192    # CHECK: %{{.*}} = index.rems %{{.*}}, %{{.*}}
193
194
195# CHECK-LABEL: TEST: testRemUOp
196@run
197def testRemUOp(ctx):
198    a = index.ConstantOp(value=42)
199    b = index.ConstantOp(value=23)
200    r = index.RemUOp(a, b)
201    # CHECK: %{{.*}} = index.remu %{{.*}}, %{{.*}}
202
203
204# CHECK-LABEL: TEST: testShlOp
205@run
206def testShlOp(ctx):
207    a = index.ConstantOp(value=42)
208    b = index.ConstantOp(value=3)
209    r = index.ShlOp(a, b)
210    # CHECK: %{{.*}} = index.shl %{{.*}}, %{{.*}}
211
212
213# CHECK-LABEL: TEST: testShrSOp
214@run
215def testShrSOp(ctx):
216    a = index.ConstantOp(value=42)
217    b = index.ConstantOp(value=3)
218    r = index.ShrSOp(a, b)
219    # CHECK: %{{.*}} = index.shrs %{{.*}}, %{{.*}}
220
221
222# CHECK-LABEL: TEST: testShrUOp
223@run
224def testShrUOp(ctx):
225    a = index.ConstantOp(value=42)
226    b = index.ConstantOp(value=3)
227    r = index.ShrUOp(a, b)
228    # CHECK: %{{.*}} = index.shru %{{.*}}, %{{.*}}
229
230
231# CHECK-LABEL: TEST: testSizeOfOp
232@run
233def testSizeOfOp(ctx):
234    r = index.SizeOfOp()
235    # CHECK: %{{.*}} = index.sizeof
236