1# RUN: %PYTHON %s 2>&1 | FileCheck %s
2
3import sys
4from mlir.ir import *
5from mlir.dialects import builtin
6from mlir.dialects import linalg
7from mlir.dialects import std
8from mlir.passmanager import *
9from mlir.execution_engine import *
10
11
12# Log everything to stderr and flush so that we have a unified stream to match
13# errors/info emitted by MLIR to stderr.
14def log(*args):
15  print(*args, file=sys.stderr)
16  sys.stderr.flush()
17
18
19matmul_boiler = """
20func @main() -> f32 attributes {llvm.emit_c_interface} {
21  %v0 = constant 0.0 : f32
22  %v1 = constant 1.0 : f32
23  %v2 = constant 2.0 : f32
24
25  %A = memref.alloc() : memref<4x16xf32>
26  %B = memref.alloc() : memref<16x8xf32>
27  %C = memref.alloc() : memref<4x8xf32>
28  linalg.fill(%v1, %A) : f32, memref<4x16xf32>
29  linalg.fill(%v2, %B) : f32, memref<16x8xf32>
30  linalg.fill(%v0, %C) : f32, memref<4x8xf32>
31
32  call @matmul_on_buffers(%A, %B, %C) :
33    (memref<4x16xf32>, memref<16x8xf32>, memref<4x8xf32>) -> ()
34
35  %c0 = constant 0 : index
36  %0 = memref.load %C[%c0, %c0] : memref<4x8xf32>
37
38  // TODO: FFI-based solution to allow testing and printing with python code.
39  return %0 : f32
40}
41"""
42
43fill_boiler = """
44func @main() -> i32 attributes {llvm.emit_c_interface} {
45  %O = memref.alloc() : memref<4x16xi32>
46  %min = constant -1000.0 : f64
47  %max = constant 1000.0 : f64
48  %seed = constant 42 : i32
49
50  call @fill_on_buffers(%min, %max, %seed, %O) :
51    (f64, f64, i32, memref<4x16xi32>) -> ()
52
53  %c0 = constant 0 : index
54  %0 = memref.load %O[%c0, %c0] : memref<4x16xi32>
55
56  // TODO: FFI-based solution to allow testing and printing with python code.
57  return %0 : i32
58}
59"""
60
61conv_boiler = """
62func @main() -> i32 attributes {llvm.emit_c_interface} {
63  %v0 = constant 0 : i32
64  %v1 = constant 1.0 : f64
65  %v2 = constant 2.0 : f64
66
67  %input = memref.alloc() : memref<1x4x16x1xf64>
68  %filter = memref.alloc() : memref<2x2x1xf64>
69  %output = memref.alloc() : memref<1x2x4x1xi32>
70  linalg.fill(%v1, %input) : f64, memref<1x4x16x1xf64>
71  linalg.fill(%v2, %filter) : f64, memref<2x2x1xf64>
72  linalg.fill(%v0, %output) : i32, memref<1x2x4x1xi32>
73
74  call @conv_on_buffers(%input, %filter, %output) :
75    (memref<1x4x16x1xf64>, memref<2x2x1xf64>, memref<1x2x4x1xi32>) -> ()
76
77  %c0 = constant 0 : index
78  %0 = memref.load %output[%c0, %c0, %c0, %c0] : memref<1x2x4x1xi32>
79
80  // TODO: FFI-based solution to allow testing and printing with python code.
81  return %0 : i32
82}
83"""
84
85pooling_boiler = """
86func @main() -> i32 attributes {llvm.emit_c_interface} {
87  %v0 = constant 0 : i32
88  %v42 = constant 42.0 : f64
89  %v77 = constant 77.0 : f64
90  %v-13 = constant -13.0 : f64
91  %v1 = constant 1.0 : f64
92
93  %input = memref.alloc() : memref<1x4x16x1xf64>
94  %shape = memref.alloc() : memref<2x2xf64>
95  %output = memref.alloc() : memref<1x2x4x1xi32>
96  linalg.fill(%v1, %input) : f64, memref<1x4x16x1xf64>
97  linalg.fill(%v1, %shape) : f64, memref<2x2xf64>
98  linalg.fill(%v0, %output) : i32, memref<1x2x4x1xi32>
99
100  %c0 = constant 0 : index
101  %c1 = constant 1 : index
102  %c2 = constant 2 : index
103  memref.store %v42, %input[%c0, %c0, %c0, %c0] : memref<1x4x16x1xf64>
104  memref.store %v77, %input[%c0, %c0, %c1, %c0] : memref<1x4x16x1xf64>
105  memref.store %v-13, %input[%c0, %c0, %c2, %c0] : memref<1x4x16x1xf64>
106
107  call @pooling_on_buffers(%input, %shape, %output) :
108    (memref<1x4x16x1xf64>, memref<2x2xf64>, memref<1x2x4x1xi32>) -> ()
109
110  %0 = memref.load %output[%c0, %c0, %c0, %c0] : memref<1x2x4x1xi32>
111
112  // TODO: FFI-based solution to allow testing and printing with python code.
113  return %0 : i32
114}
115"""
116
117
118def transform(module, boilerplate):
119  import mlir.conversions
120  import mlir.dialects.linalg.passes
121  import mlir.transforms
122
123  # TODO: Allow cloning functions from one module to another.
124  # Atm we have to resort to string concatenation.
125  mod = Module.parse(
126      str(module.operation.regions[0].blocks[0].operations[0].operation) +
127      boilerplate)
128  pm = PassManager.parse("func(convert-linalg-to-loops, lower-affine, " +
129                         "convert-scf-to-std), convert-vector-to-llvm," +
130                         "convert-std-to-llvm")
131  pm.run(mod)
132  return mod
133
134
135def test_matmul_builtin():
136  with Context() as ctx, Location.unknown():
137    module = Module.create()
138    f32 = F32Type.get()
139    with InsertionPoint(module.body):
140
141      @builtin.FuncOp.from_py_func(
142          MemRefType.get((4, 16), f32), MemRefType.get((16, 8), f32),
143          MemRefType.get((4, 8), f32))
144      def matmul_on_buffers(lhs, rhs, out):
145        linalg.matmul(lhs, rhs, outs=[out])
146
147    execution_engine = ExecutionEngine(transform(module, matmul_boiler))
148
149    # TODO: FFI-based solution to allow testing and printing with python code.
150    # Prepare arguments: one result f32.
151    # Arguments must be passed as pointers.
152    c_float_p = ctypes.c_float * 1
153    res = c_float_p(-1.)
154    execution_engine.invoke("main", res)
155
156    log("RESULT: ", res[0])
157    # CHECK: RESULT: 32.0
158
159
160test_matmul_builtin()
161
162
163def test_matmul_generic():
164  with Context() as ctx, Location.unknown():
165    module = Module.create()
166    f32 = F32Type.get()
167    with InsertionPoint(module.body):
168
169      @builtin.FuncOp.from_py_func(
170          MemRefType.get((4, 16), f32), MemRefType.get((16, 8), f32),
171          MemRefType.get((4, 8), f32))
172      def matmul_on_buffers(lhs, rhs, out):
173        linalg.matmul(lhs, rhs, outs=[out], emit_generic=True)
174
175    execution_engine = ExecutionEngine(transform(module, matmul_boiler))
176
177    # TODO: FFI-based solution to allow testing and printing with python code.
178    # Prepare arguments: one result f32.
179    # Arguments must be passed as pointers.
180    c_float_p = ctypes.c_float * 1
181    res = c_float_p(-1.)
182    execution_engine.invoke("main", res)
183
184    log("RESULT: ", res[0])
185    # CHECK: RESULT: 32.0
186
187
188test_matmul_generic()
189
190
191def test_fill_builtin():
192  with Context() as ctx, Location.unknown():
193    module = Module.create()
194    f64 = F64Type.get()
195    i32 = IntegerType.get_signless(32)
196    with InsertionPoint(module.body):
197
198      @builtin.FuncOp.from_py_func(f64, f64, i32, MemRefType.get((4, 16), i32))
199      def fill_on_buffers(min, max, seed, out):
200        linalg.fill_rng_2d(min, max, seed, outs=[out])
201
202    execution_engine = ExecutionEngine(transform(module, fill_boiler))
203
204    # TODO: FFI-based solution to allow testing and printing with python code.
205    # Prepare arguments: one result i32.
206    # Arguments must be passed as pointers.
207    c_int_p = ctypes.c_int * 1
208    res = c_int_p(-1)
209    execution_engine.invoke("main", res)
210
211    log("RESULT: ", res[0])
212    # CHECK: RESULT: -480
213
214
215test_fill_builtin()
216
217
218def test_fill_generic():
219  with Context() as ctx, Location.unknown():
220    module = Module.create()
221    f64 = F64Type.get()
222    i32 = IntegerType.get_signless(32)
223    with InsertionPoint(module.body):
224
225      @builtin.FuncOp.from_py_func(f64, f64, i32, MemRefType.get((4, 16), i32))
226      def fill_on_buffers(min, max, seed, out):
227        linalg.fill_rng_2d(min, max, seed, outs=[out], emit_generic=True)
228
229    execution_engine = ExecutionEngine(transform(module, fill_boiler))
230
231    # TODO: FFI-based solution to allow testing and printing with python code.
232    # Prepare arguments: one result i32.
233    # Arguments must be passed as pointers.
234    c_int_p = ctypes.c_int * 1
235    res = c_int_p(-1)
236    execution_engine.invoke("main", res)
237
238    log("RESULT: ", res[0])
239    # CHECK: RESULT: -480
240
241
242test_fill_generic()
243
244
245def test_conv_builtin():
246  with Context() as ctx, Location.unknown():
247    module = Module.create()
248    f64 = F64Type.get()
249    i32 = IntegerType.get_signless(32)
250    with InsertionPoint(module.body):
251
252      @builtin.FuncOp.from_py_func(
253          MemRefType.get((1, 4, 16, 1), f64), MemRefType.get((2, 2, 1), f64),
254          MemRefType.get((1, 2, 4, 1), i32))
255      def conv_on_buffers(input, filter, output):
256        linalg.depthwise_conv_2d_input_nhwc_filter_hwc_poly(
257            input, filter, outs=[output], strides=[2, 4], dilations=[1, 2])
258
259    execution_engine = ExecutionEngine(transform(module, conv_boiler))
260
261    # TODO: FFI-based solution to allow testing and printing with python code.
262    # Prepare arguments: one result i32.
263    # Arguments must be passed as pointers.
264    c_int_p = ctypes.c_int * 1
265    res = c_int_p(-1)
266    execution_engine.invoke("main", res)
267
268    log("RESULT: ", res[0])
269    # CHECK: RESULT: 8
270
271
272test_conv_builtin()
273
274
275def test_conv_generic():
276  with Context() as ctx, Location.unknown():
277    module = Module.create()
278    f64 = F64Type.get()
279    i32 = IntegerType.get_signless(32)
280    with InsertionPoint(module.body):
281
282      @builtin.FuncOp.from_py_func(
283          MemRefType.get((1, 4, 16, 1), f64), MemRefType.get((2, 2, 1), f64),
284          MemRefType.get((1, 2, 4, 1), i32))
285      def conv_on_buffers(input, filter, output):
286        linalg.depthwise_conv_2d_input_nhwc_filter_hwc_poly(
287            input,
288            filter,
289            outs=[output],
290            strides=[2, 4],
291            dilations=[1, 2],
292            emit_generic=True)
293
294    execution_engine = ExecutionEngine(transform(module, conv_boiler))
295
296    # TODO: FFI-based solution to allow testing and printing with python code.
297    # Prepare arguments: one result i32.
298    # Arguments must be passed as pointers.
299    c_int_p = ctypes.c_int * 1
300    res = c_int_p(-1)
301    execution_engine.invoke("main", res)
302
303    log("RESULT: ", res[0])
304    # CHECK: RESULT: 8
305
306
307test_conv_generic()
308
309
310def test_max_pooling_builtin():
311  with Context() as ctx, Location.unknown():
312    module = Module.create()
313    f64 = F64Type.get()
314    i32 = IntegerType.get_signless(32)
315    with InsertionPoint(module.body):
316
317      @builtin.FuncOp.from_py_func(
318          MemRefType.get((1, 4, 16, 1), f64), MemRefType.get((2, 2), f64),
319          MemRefType.get((1, 2, 4, 1), i32))
320      def pooling_on_buffers(input, shape, output):
321        linalg.pooling_nhwc_max_poly(
322            input, shape, outs=[output], strides=[2, 4], dilations=[1, 2])
323
324    execution_engine = ExecutionEngine(transform(module, pooling_boiler))
325
326    # TODO: FFI-based solution to allow testing and printing with python code.
327    # Prepare arguments: one result i32.
328    # Arguments must be passed as pointers.
329    c_int_p = ctypes.c_int * 1
330    res = c_int_p(-1)
331    execution_engine.invoke("main", res)
332
333    log("RESULT: ", res[0])
334    # 77 is not selected due to the dilation 2 in the second dimension.
335    # CHECK: RESULT: 42
336
337
338test_max_pooling_builtin()
339
340
341def test_max_pooling_generic():
342  with Context() as ctx, Location.unknown():
343    module = Module.create()
344    f64 = F64Type.get()
345    i32 = IntegerType.get_signless(32)
346    with InsertionPoint(module.body):
347
348      @builtin.FuncOp.from_py_func(
349          MemRefType.get((1, 4, 16, 1), f64), MemRefType.get((2, 2), f64),
350          MemRefType.get((1, 2, 4, 1), i32))
351      def pooling_on_buffers(input, shape, output):
352        linalg.pooling_nhwc_max_poly(
353            input,
354            shape,
355            outs=[output],
356            strides=[2, 4],
357            dilations=[1, 2],
358            emit_generic=True)
359
360    execution_engine = ExecutionEngine(transform(module, pooling_boiler))
361
362    # TODO: FFI-based solution to allow testing and printing with python code.
363    # Prepare arguments: one result i32.
364    # Arguments must be passed as pointers.
365    c_int_p = ctypes.c_int * 1
366    res = c_int_p(-1)
367    execution_engine.invoke("main", res)
368
369    log("RESULT: ", res[0])
370    # 77 is not selected due to the dilation 2 in the second dimension.
371    # CHECK: RESULT: 42
372
373
374test_max_pooling_generic()
375
376
377def test_min_pooling_builtin():
378  with Context() as ctx, Location.unknown():
379    module = Module.create()
380    f64 = F64Type.get()
381    i32 = IntegerType.get_signless(32)
382    with InsertionPoint(module.body):
383
384      @builtin.FuncOp.from_py_func(
385          MemRefType.get((1, 4, 16, 1), f64), MemRefType.get((2, 2), f64),
386          MemRefType.get((1, 2, 4, 1), i32))
387      def pooling_on_buffers(input, shape, output):
388        linalg.pooling_nhwc_min_poly(
389            input, shape, outs=[output], strides=[2, 4], dilations=[1, 2])
390
391    execution_engine = ExecutionEngine(transform(module, pooling_boiler))
392
393    # TODO: FFI-based solution to allow testing and printing with python code.
394    # Prepare arguments: one result i32.
395    # Arguments must be passed as pointers.
396    c_int_p = ctypes.c_int * 1
397    res = c_int_p(-1)
398    execution_engine.invoke("main", res)
399
400    log("RESULT: ", res[0])
401    # CHECK: RESULT: -13
402
403
404test_min_pooling_builtin()
405
406
407def test_min_pooling_generic():
408  with Context() as ctx, Location.unknown():
409    module = Module.create()
410    f64 = F64Type.get()
411    i32 = IntegerType.get_signless(32)
412    with InsertionPoint(module.body):
413
414      @builtin.FuncOp.from_py_func(
415          MemRefType.get((1, 4, 16, 1), f64), MemRefType.get((2, 2), f64),
416          MemRefType.get((1, 2, 4, 1), i32))
417      def pooling_on_buffers(input, shape, output):
418        linalg.pooling_nhwc_min_poly(
419            input,
420            shape,
421            outs=[output],
422            strides=[2, 4],
423            dilations=[1, 2],
424            emit_generic=True)
425
426    execution_engine = ExecutionEngine(transform(module, pooling_boiler))
427
428    # TODO: FFI-based solution to allow testing and printing with python code.
429    # Prepare arguments: one result i32.
430    # Arguments must be passed as pointers.
431    c_int_p = ctypes.c_int * 1
432    res = c_int_p(-1)
433    execution_engine.invoke("main", res)
434
435    log("RESULT: ", res[0])
436    # CHECK: RESULT: -13
437
438
439test_min_pooling_generic()
440