llvm-project

1
# RUN: mlir-linalg-ods-yaml-gen %s --o-ods-decl=- | FileCheck %s --check-prefix=ODS
2
# RUN: mlir-linalg-ods-yaml-gen %s --o-impl=- | FileCheck %s --check-prefix=IMPL
3

4
# @linalg_structured_op
5
# def test1(O=TensorDef(T, S.M, S.N, output=True),
6
#           cast=TypeFnAttrDef(default=TypeFn.cast_signed)):
7
#   """Title.
8

9
#   Detailed description.
10
#   """
11
#   O[D.m, D.n] = cast(T, const(42)) + cast(T, index(D.n))
12

13
--- !LinalgOpConfig
14
metadata: !LinalgOpMetadata
15
  name: test1
16
  cpp_class_name: Test1Op
17
  doc: |-
18
    Title.
19

20
    Detailed description.
21
structured_op: !LinalgStructuredOpConfig
22
  args:
23
  - !LinalgOperandDefConfig
24
    name: O
25
    kind: output_tensor
26
    type_var: T
27
    shape_map: affine_map<()[s0, s1] -> (s0, s1)>
28
  - !LinalgOperandDefConfig
29
    name: cast
30
    kind: type_fn_attr
31
    default_fn: cast_signed
32
  indexing_maps: !LinalgIndexingMapsConfig
33
    static_indexing_maps:
34
    - affine_map<(d0, d1)[s0, s1] -> (d0, d1)>
35
  iterator_types:
36
  - parallel
37
  - parallel
38
  assignments:
39
  - !ScalarAssign
40
    arg: O
41
    value: !ScalarExpression
42
      scalar_fn:
43
        kind: binary
44
        fn_name: add
45
        operands:
46
        - !ScalarExpression
47
          scalar_fn:
48
            kind: type
49
            attr_name: cast
50
            type_var: T
51
            operands:
52
            - !ScalarExpression
53
              scalar_const: '42 : i64'
54
        - !ScalarExpression
55
          scalar_fn:
56
            kind: type
57
            attr_name: cast
58
            type_var: T
59
            operands:
60
            - !ScalarExpression
61
              scalar_index: 1
62

63
# ODS-LABEL:  def Test1Op : LinalgStructuredBase_Op<"test1"
64

65
#       ODS:  let summary = [{Title.}];
66
#  ODS-NEXT:  let description = [{Detailed description.}];
67

68
#       ODS:  let arguments =
69
#  ODS-NEXT:    Variadic<AnyType>:$inputs,
70
#  ODS-NEXT:    Variadic<AnyShaped>:$outputs,
71
#  ODS-NEXT:    DefaultValuedOptionalAttr<TypeFnAttr, "TypeFn::cast_signed">:$cast
72

73
#       ODS:  let builders =
74
#       ODS:  (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
75
#  ODS-NEXT:       "ValueRange":$outputs,
76
#  ODS-NEXT:       CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
77

78
#       ODS:  (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
79
#  ODS-NEXT:       "ValueRange":$outputs, "Attribute":$cast,
80
#  ODS-NEXT:       CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
81

82
#       ODS:    buildStructuredOp($_builder, $_state, resultTensorTypes,
83
#  ODS-NEXT:      attributes, Test1Op::getRegionBuilder())
84

85
#       ODS:    MutableOperandRange getDpsInitsMutable() {
86
#  ODS-NEXT:      return getOutputsMutable()
87
#  ODS-NEXT:    }
88

89
# IMPL-LABEL:  void Test1Op::regionBuilder(ImplicitLocOpBuilder &b,
90
#  IMPL-NEXT:    Block &block, ArrayRef<NamedAttribute> attrs)
91
#       IMPL:  TypeFn castVal = TypeFn::cast_signed;
92
#  IMPL-NEXT:  auto castIter = llvm::find_if(attrs, [&](const NamedAttribute &attr) {
93
#  IMPL-NEXT:                                return attr.getName() == "cast"; });
94
#  IMPL-NEXT:  if (castIter != attrs.end()) {
95
#  IMPL-NEXT:    if (auto attr = llvm::dyn_cast<TypeFnAttr>(castIter->getValue()))
96
#  IMPL-NEXT:      castVal = attr.getValue();
97
#  IMPL-NEXT:  }
98

99
#       IMPL:  Value [[VAL0:[a-z0-9]+]] = helper.constant("42 : i64");
100
#   IMPL-DAG:  Value [[VAL1:[a-z0-9]+]] = helper.buildTypeFn(castVal, block.getArgument(0).getType(), [[VAL0]]);
101
#   IMPL-DAG:  Value [[VAL2:[a-z0-9]+]] = helper.index(1);
102
#   IMPL-DAG:  Value [[VAL3:[a-z0-9]+]] = helper.buildTypeFn(castVal, block.getArgument(0).getType(), [[VAL2]]);
103
#   IMPL-DAG:  Value [[VAL4:[a-z0-9]+]] = helper.buildBinaryFn(BinaryFn::add, [[VAL1]], [[VAL3]]);
104

105

106
# @linalg_structured_op
107
# def test2(I=TensorDef(T, S.M, S.N),
108
#           O=TensorDef(T, S.M, S.N, output=True),
109
#           strides=IndexAttrDef(S.SM, S.SN, default=[1, 2])):
110
#   """Title.
111

112
#   Detailed description.
113
#   """
114
#   O[D.m, D.n] = I[D.n * S.SM, D.m * S.SN]
115

116
--- !LinalgOpConfig
117
metadata: !LinalgOpMetadata
118
  name: test2
119
  cpp_class_name: Test2Op
120
  doc: |-
121
    Title.
122

123
    Detailed description.
124
structured_op: !LinalgStructuredOpConfig
125
  args:
126
  - !LinalgOperandDefConfig
127
    name: I
128
    kind: input_tensor
129
    type_var: T
130
    shape_map: affine_map<()[s0, s1, s2, s3] -> (s0, s1)>
131
  - !LinalgOperandDefConfig
132
    name: O
133
    kind: output_tensor
134
    type_var: T
135
    shape_map: affine_map<()[s0, s1, s2, s3] -> (s0, s1)>
136
  - !LinalgOperandDefConfig
137
    name: strides
138
    kind: index_attr
139
    index_attr_map: affine_map<()[s0, s1, s2, s3] -> (s2, s3)>
140
    default_indices:
141
    - 1
142
    - 2
143
  indexing_maps: !LinalgIndexingMapsConfig
144
    static_indexing_maps:
145
    - affine_map<(d0, d1)[s0, s1, s2, s3] -> (d1 * s2, d0 * s3)>
146
    - affine_map<(d0, d1)[s0, s1, s2, s3] -> (d0, d1)>
147
  iterator_types:
148
  - parallel
149
  - parallel
150
  assignments:
151
  - !ScalarAssign
152
    arg: O
153
    value: !ScalarExpression
154
      scalar_arg: I
155

156
# ODS-LABEL:  def Test2Op : LinalgStructuredBase_Op<"test2"
157

158
#       ODS:  let arguments =
159
#  ODS-NEXT:    Variadic<AnyType>:$inputs,
160
#  ODS-NEXT:    Variadic<AnyShaped>:$outputs,
161
#  ODS-NEXT:    DefaultValuedOptionalAttr<RankedI64ElementsAttr<[2]>
162
#  ODS-SAME:    "{ static_cast<int64_t>(1), static_cast<int64_t>(2) }">:$strides
163

164
#       ODS:  "Attribute":$strides
165
#       ODS:  $_state.addAttribute("strides", strides);
166

167
#       ODS:  bool hasDynamicIndexingMaps();
168
#  ODS-NEXT:  LogicalResult verifyIndexingMapRequiredAttributes();
169

170
#       IMPL:  getSymbolBindings(Test2Op self)
171
#       IMPL:  cst2 = self.getStrides().getValues<int64_t>()[0];
172
#  IMPL-NEXT:  getAffineConstantExpr(cst2, context)
173
#       IMPL:  cst3 = self.getStrides().getValues<int64_t>()[1];
174
#  IMPL-NEXT:  getAffineConstantExpr(cst3, context)
175

176
#       IMPL:  Test2Op::getIndexingMaps()
177
#       IMPL:  = getSymbolBindings(*this);
178
#       IMPL:  "affine_map<(d0, d1)[s0, s1, s2, s3] -> (d1 * s2, d0 * s3)>"
179
#       IMPL:  "affine_map<(d0, d1)[s0, s1, s2, s3] -> (d0, d1)>"
180

181
#       IMPL:  Test2Op::getNumRegionArgs() { return 2; }
182

183
#       IMPL:  Test2Op::hasDynamicIndexingMaps() { return true; }
184
#       IMPL:  Test2Op::verifyIndexingMapRequiredAttributes()
185
#       IMPL:  auto attr = op->getAttrOfType<DenseElementsAttr>("strides")
186
#       IMPL:  "incorrect element type for index attribute 'strides'"
187
#       IMPL:  "incorrect shape for index attribute 'strides'"
188
#       IMPL:  void Test2Op::regionBuilder(ImplicitLocOpBuilder &b,
189
#  IMPL-NEXT:    Block &block, ArrayRef<NamedAttribute> attrs)
190
#  IMPL-NEXT:    assert(2 > 0 && block.getNumArguments() == 2 &&
191

192
#       IMPL:   yields.push_back(block.getArgument(0));
193

194
# @linalg_structured_op
195
# def test3(value=ScalarDef(T1),
196
#           O=TensorDef(U, output=True)):
197
#   """Title.
198

199
#   Detailed description.
200
#   """
201
#   O[None] = TypeFn.cast_signed(U, value)
202

203
--- !LinalgOpConfig
204
metadata: !LinalgOpMetadata
205
  name: test3
206
  cpp_class_name: Test3Op
207
  doc: |-
208
    Title.
209

210
    Detailed description.
211
structured_op: !LinalgStructuredOpConfig
212
  args:
213
  - !LinalgOperandDefConfig
214
    name: value
215
    kind: scalar
216
    type_var: T1
217
  - !LinalgOperandDefConfig
218
    name: O
219
    kind: output_tensor
220
    type_var: U
221
    shape_map: affine_map<() -> ()>
222
  indexing_maps: !LinalgIndexingMapsConfig
223
    static_indexing_maps:
224
    - affine_map<() -> ()>
225
    - affine_map<() -> ()>
226
  iterator_types: []
227
  assignments:
228
  - !ScalarAssign
229
    arg: O
230
    value: !ScalarExpression
231
      scalar_fn:
232
        kind: type
233
        fn_name: cast_signed
234
        type_var: U
235
        operands:
236
        - !ScalarExpression
237
          scalar_arg: value
238

239
#       IMPL:  Test3Op::getIteratorTypesArray() {
240
#  IMPL-NEXT:    int64_t rank = getRank(getDpsInitOperand(0));
241

242
#       IMPL:  Test3Op::getIndexingMaps() {
243
#  IMPL-NEXT:    MLIRContext *context = getContext();
244
#  IMPL-NEXT:    AffineMap scalarMap = AffineMap::get(getNumParallelLoops(), 0, context);
245
#  IMPL-NEXT:    AffineMap tensorMap = AffineMap::getMultiDimIdentityMap(
246

247
# @linalg_structured_op
248
# def test4(O=TensorDef(T, S.M, S.N, output=True),
249
#           unary_fun=UnaryFnAttrDef(default=UnaryFn.exp),
250
#           binary_fun=BinaryFnAttrDef(default=BinaryFn.add)):
251
#   """Title.
252

253
#   Detailed description.
254
#   """
255
#   O[D.m, D.n] = binary_fun(unary_fun(O[D.m, D.n]), O[D.m, D.n])
256

257
--- !LinalgOpConfig
258
metadata: !LinalgOpMetadata
259
  name: test4
260
  cpp_class_name: Test4Op
261
  doc: |-
262
    Title.
263

264
    Detailed description.
265
structured_op: !LinalgStructuredOpConfig
266
  args:
267
  - !LinalgOperandDefConfig
268
    name: O
269
    kind: output_tensor
270
    type_var: T
271
    shape_map: affine_map<()[s0, s1] -> (s0, s1)>
272
  - !LinalgOperandDefConfig
273
    name: unary_fun
274
    kind: unary_fn_attr
275
    default_fn: exp
276
  - !LinalgOperandDefConfig
277
    name: binary_fun
278
    kind: binary_fn_attr
279
    default_fn: add
280
  indexing_maps: !LinalgIndexingMapsConfig
281
    static_indexing_maps:
282
    - affine_map<(d0, d1)[s0, s1] -> (d0, d1)>
283
  iterator_types:
284
  - parallel
285
  - parallel
286
  assignments:
287
  - !ScalarAssign
288
    arg: O
289
    value: !ScalarExpression
290
      scalar_fn:
291
        kind: binary
292
        attr_name: binary_fun
293
        operands:
294
        - !ScalarExpression
295
          scalar_fn:
296
            kind: unary
297
            attr_name: unary_fun
298
            operands:
299
            - !ScalarExpression
300
              scalar_arg: O
301
        - !ScalarExpression
302
          scalar_arg: O
303

304
# ODS-LABEL:  def Test4Op : LinalgStructuredBase_Op<"test4"
305

306
#       ODS:  let arguments =
307
#  ODS-NEXT:    Variadic<AnyType>:$inputs,
308
#  ODS-NEXT:    Variadic<AnyShaped>:$outputs,
309
#  ODS-NEXT:    DefaultValuedOptionalAttr<UnaryFnAttr, "UnaryFn::exp">:$unary_fun,
310
#  ODS-NEXT:    DefaultValuedOptionalAttr<BinaryFnAttr, "BinaryFn::add">:$binary_fun
311

312
#       ODS:    "Attribute":$unary_fun, "Attribute":$binary_fun,
313

314
#       ODS:    $_state.addAttribute("unary_fun", unary_fun)
315
#  ODS-NEXT:    $_state.addAttribute("binary_fun", binary_fun)
316

317
# IMPL-LABEL:  void Test4Op::regionBuilder(ImplicitLocOpBuilder &b,
318
#  IMPL-NEXT:    Block &block, ArrayRef<NamedAttribute> attrs)
319
#       IMPL:  UnaryFn unary_funVal = UnaryFn::exp
320
#       IMPL:  BinaryFn binary_funVal = BinaryFn::add
321

322
#       IMPL:  Value [[VAL0:[a-z0-9]+]] = helper.buildUnaryFn(unary_funVal, block.getArgument(0))
323
#  IMPL-NEXT:  Value [[VAL1:[a-z0-9]+]] = helper.buildBinaryFn(binary_funVal, [[VAL0]], block.getArgument(0))
324
#  IMPL-NEXT:  yields.push_back([[VAL1]])
325

326
# @linalg_structured_op
327
# def test5(value=ScalarDef(T1), O=TensorDef(U, output=True)):
328
#   """Title.
329

330
#   Detailed description.
331
#   """
332
#   implements(FillOpInterface)
333
#   defines(Canonicalizer)
334
#   O[None] = TypeFn.cast(U, value)
335

336
--- !LinalgOpConfig
337
metadata: !LinalgOpMetadata
338
  name: test5
339
  cpp_class_name: Test5Op
340
  doc: |-
341
    Title.
342

343
    Detailed description.
344
  implements:
345
  - LinalgFillOpInterface
346
  defines:
347
  - hasCanonicalizer
348
structured_op: !LinalgStructuredOpConfig
349
  args:
350
  - !LinalgOperandDefConfig
351
    name: value
352
    kind: scalar
353
    type_var: T1
354
  - !LinalgOperandDefConfig
355
    name: O
356
    kind: output_tensor
357
    type_var: U
358
    shape_map: affine_map<() -> ()>
359
  indexing_maps: !LinalgIndexingMapsConfig
360
    static_indexing_maps:
361
    - affine_map<() -> ()>
362
    - affine_map<() -> ()>
363
  iterator_types: []
364
  assignments:
365
  - !ScalarAssign
366
    arg: O
367
    value: !ScalarExpression
368
      scalar_fn:
369
        kind: type
370
        fn_name: cast
371
        type_var: U
372
        operands:
373
        - !ScalarExpression
374
          scalar_arg: value
375

376
# ODS-LABEL:  def Test5Op : LinalgStructuredBase_Op<"test5"
377
#  ODS-NEXT:  /*extraInterfaces=*/[LinalgFillOpInterface])>
378

379
#       ODS:  let hasCanonicalizer = 1;
380