1# RUN: %PYTHON %s | FileCheck %s
2
3from mlir.ir import *
4from mlir.dialects import builtin
5from mlir.dialects import linalg
6from mlir.dialects import std
7from mlir.dialects import arith
8
9
10def run(f):
11  print("\nTEST:", f.__name__)
12  f()
13  return f
14
15
16# CHECK-LABEL: TEST: testInitTensor
17@run
18def testInitTensor():
19  with Context() as ctx, Location.unknown():
20    module = Module.create()
21    f32 = F32Type.get()
22    with InsertionPoint(module.body):
23      # CHECK-LABEL: func @static_sizes
24      # CHECK: %0 = linalg.init_tensor [3, 4] : tensor<3x4xf32>
25      @builtin.FuncOp.from_py_func()
26      def static_sizes():
27        return linalg.InitTensorOp([3, 4], f32)
28
29      # CHECK-LABEL: func @dynamic_sizes
30      # CHECK: %0 = linalg.init_tensor [%arg0, %arg1] : tensor<?x?xf32>
31      @builtin.FuncOp.from_py_func(IndexType.get(), IndexType.get())
32      def dynamic_sizes(d0, d1):
33        return linalg.InitTensorOp([d0, d1], f32)
34
35      # CHECK-LABEL: func @zero_d
36      # CHECK: %0 = linalg.init_tensor [] : tensor<f32>
37      @builtin.FuncOp.from_py_func()
38      def zero_d():
39        return linalg.InitTensorOp([], f32)
40
41  print(module)
42
43
44# CHECK-LABEL: TEST: testInitTensorStaticSizesAttribute
45@run
46def testInitTensorStaticSizesAttribute():
47  with Context() as ctx, Location.unknown():
48    module = Module.create()
49    f32 = F32Type.get()
50    with InsertionPoint(module.body):
51      op = linalg.InitTensorOp([3, 4], f32)
52      # CHECK: [3, 4]
53      print(op.attributes["static_sizes"])
54
55
56# CHECK-LABEL: TEST: testFill
57@run
58def testFill():
59  with Context() as ctx, Location.unknown():
60    module = Module.create()
61    f32 = F32Type.get()
62    with InsertionPoint(module.body):
63      # CHECK-LABEL: func @fill_tensor
64      #  CHECK-SAME:   %[[OUT:[0-9a-z]+]]: tensor<12x?xf32>
65      #  CHECK-NEXT: %[[CST:.*]] = arith.constant 0.0{{.*}} : f32
66      #  CHECK-NEXT: %[[RES:.*]] = linalg.fill(%[[CST]], %[[OUT]]) : f32, tensor<12x?xf32> -> tensor<12x?xf32>
67      #  CHECK-NEXT: return %[[RES]] : tensor<12x?xf32>
68      @builtin.FuncOp.from_py_func(RankedTensorType.get((12, -1), f32))
69      def fill_tensor(out):
70        zero = arith.ConstantOp(value=FloatAttr.get(f32, 0.), result=f32).result
71        # TODO: FillOp.result is None. When len(results) == 1 we expect it to
72        # be results[0] as per _linalg_ops_gen.py. This seems like an
73        # orthogonal bug in the generator of _linalg_ops_gen.py.
74        return linalg.FillOp(output=out, value=zero).results[0]
75
76      # CHECK-LABEL: func @fill_buffer
77      #  CHECK-SAME:   %[[OUT:[0-9a-z]+]]: memref<12x?xf32>
78      #  CHECK-NEXT: %[[CST:.*]] = arith.constant 0.0{{.*}} : f32
79      #  CHECK-NEXT: linalg.fill(%[[CST]], %[[OUT]]) : f32, memref<12x?xf32>
80      #  CHECK-NEXT: return
81      @builtin.FuncOp.from_py_func(MemRefType.get((12, -1), f32))
82      def fill_buffer(out):
83        zero = arith.ConstantOp(value=FloatAttr.get(f32, 0.), result=f32).result
84        linalg.FillOp(output=out, value=zero)
85
86  print(module)
87
88
89# CHECK-LABEL: TEST: testStructuredOpOnTensors
90@run
91def testStructuredOpOnTensors():
92  with Context() as ctx, Location.unknown():
93    module = Module.create()
94    f32 = F32Type.get()
95    tensor_type = RankedTensorType.get((2, 3, 4), f32)
96    with InsertionPoint(module.body):
97      func = builtin.FuncOp(
98          name="matmul_test",
99          type=FunctionType.get(
100              inputs=[tensor_type, tensor_type], results=[tensor_type]))
101      with InsertionPoint(func.add_entry_block()):
102        lhs, rhs = func.entry_block.arguments
103        result = linalg.MatmulOp([lhs, rhs], results=[tensor_type]).result
104        std.ReturnOp([result])
105
106  # CHECK: %[[R:.*]] = linalg.matmul ins(%arg0, %arg1 : tensor<2x3x4xf32>, tensor<2x3x4xf32>) -> tensor<2x3x4xf32>
107  print(module)
108
109
110# CHECK-LABEL: TEST: testStructuredOpOnBuffers
111@run
112def testStructuredOpOnBuffers():
113  with Context() as ctx, Location.unknown():
114    module = Module.create()
115    f32 = F32Type.get()
116    memref_type = MemRefType.get((2, 3, 4), f32)
117    with InsertionPoint(module.body):
118      func = builtin.FuncOp(
119          name="matmul_test",
120          type=FunctionType.get(
121              inputs=[memref_type, memref_type, memref_type], results=[]))
122      with InsertionPoint(func.add_entry_block()):
123        lhs, rhs, result = func.entry_block.arguments
124        # TODO: prperly hook up the region.
125        linalg.MatmulOp([lhs, rhs], outputs=[result])
126        std.ReturnOp([])
127
128  # CHECK: linalg.matmul ins(%arg0, %arg1 : memref<2x3x4xf32>, memref<2x3x4xf32>) outs(%arg2 : memref<2x3x4xf32>)
129  print(module)
130
131
132# CHECK-LABEL: TEST: testNamedStructuredOpCustomForm
133@run
134def testNamedStructuredOpCustomForm():
135  with Context() as ctx, Location.unknown():
136    module = Module.create()
137    f32 = F32Type.get()
138    with InsertionPoint(module.body):
139
140      @builtin.FuncOp.from_py_func(
141          RankedTensorType.get((4, 16), f32), RankedTensorType.get((16, 8),
142                                                                   f32))
143      def named_form(lhs, rhs):
144        init_result = linalg.InitTensorOp([4, 8], f32)
145        # First check the named form with custom format
146        #      CHECK: linalg.matmul
147        #  CHECK-NOT: linalg.memoized_indexing_maps
148        # CHECK-SAME:    ins(%{{.*}} : tensor<4x16xf32>, tensor<16x8xf32>)
149        # CHECK-SAME:   outs(%{{.*}} : tensor<4x8xf32>)
150        # CHECK-SAME:   -> tensor<4x8xf32>
151        # CHECK-NEXT: return
152        return linalg.matmul(lhs, rhs, outs=[init_result.result])
153
154  print(module)
155
156
157# CHECK-LABEL: TEST: testNamedStructuredOpGenericForm
158@run
159def testNamedStructuredOpGenericForm():
160  with Context() as ctx, Location.unknown():
161    module = Module.create()
162    f32 = F32Type.get()
163    with InsertionPoint(module.body):
164
165      @builtin.FuncOp.from_py_func(
166          RankedTensorType.get((4, 16), f32), RankedTensorType.get((16, 8),
167                                                                   f32))
168      def named_form(lhs, rhs):
169        init_result = linalg.InitTensorOp([4, 8], f32)
170        #      CHECK: "linalg.matmul"(%{{.*}})
171        # CHECK-NEXT:  ^bb0(%{{.*}}: f32, %{{.*}}: f32, %{{.*}}: f32):
172        # CHECK-NEXT:    arith.mulf{{.*}} (f32, f32) -> f32
173        # CHECK-NEXT:    arith.addf{{.*}} (f32, f32) -> f32
174        # CHECK-NEXT:    linalg.yield{{.*}} (f32) -> ()
175        # CHECK-NEXT:    {linalg.memoized_indexing_maps{{.*}}operand_segment_sizes = dense<[2, 1]> : vector<2xi32>} :
176        # CHECK-SAME: (tensor<4x16xf32>, tensor<16x8xf32>, tensor<4x8xf32>) -> tensor<4x8xf32>
177        return linalg.matmul(lhs, rhs, outs=[init_result.result])
178
179  module.operation.print(print_generic_op_form=True)
180
181
182# CHECK-LABEL: TEST: testNamedStructuredAsGenericOp
183@run
184def testNamedStructuredAsGenericOp():
185  with Context() as ctx, Location.unknown():
186    module = Module.create()
187    f32 = F32Type.get()
188    with InsertionPoint(module.body):
189
190      @builtin.FuncOp.from_py_func(
191          RankedTensorType.get((4, 16), f32), RankedTensorType.get((16, 8),
192                                                                   f32))
193      def generic_form(lhs, rhs):
194        init_result = linalg.InitTensorOp([4, 8], f32)
195        # CHECK: linalg.generic
196        return linalg.matmul(
197            lhs, rhs, outs=[init_result.result], emit_generic=True)
198
199  print(module)
200
201
202# CHECK-LABEL: TEST: testOpResultFromOtherOp
203@run
204def testOpResultFromOtherOp():
205  with Context(), Location.unknown():
206    module = Module.create()
207    f32 = F32Type.get()
208    with InsertionPoint(module.body):
209
210      @builtin.FuncOp.from_py_func(
211          RankedTensorType.get((4, 16), f32), RankedTensorType.get((16, 8),
212                                                                   f32))
213      def pass_an_op_directly(arg0, arg1):
214        one = arith.ConstantOp(F32Type.get(), 1.0)
215        # CHECK: %[[LHS:.*]] = linalg.fill
216        lhs = linalg.FillOp(arg0, one)
217        # CHECK: %[[RHS:.*]] = linalg.fill
218        rhs = linalg.FillOp(arg1, one)
219        # CHECK: %[[INIT:.*]] = linalg.init_tensor
220        init = linalg.InitTensorOp([4, 8], f32)
221        # CHECK: linalg.matmul
222        # CHECK: ins(%[[LHS]], %[[RHS]]
223        # CHECK: outs(%[[INIT]]
224        return linalg.matmul(lhs, rhs, outs=init)
225
226  print(module)
227