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