llvm-project
379 строк · 11.2 Кб
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--- !LinalgOpConfig14metadata: !LinalgOpMetadata15name: test116cpp_class_name: Test1Op17doc: |-18Title.
19
20Detailed description.
21structured_op: !LinalgStructuredOpConfig22args:23- !LinalgOperandDefConfig24name: O25kind: output_tensor26type_var: T27shape_map: affine_map<()[s0, s1] -> (s0, s1)>28- !LinalgOperandDefConfig29name: cast30kind: type_fn_attr31default_fn: cast_signed32indexing_maps: !LinalgIndexingMapsConfig33static_indexing_maps:34- affine_map<(d0, d1)[s0, s1] -> (d0, d1)>35iterator_types:36- parallel37- parallel38assignments:39- !ScalarAssign40arg: O41value: !ScalarExpression42scalar_fn:43kind: binary44fn_name: add45operands:46- !ScalarExpression47scalar_fn:48kind: type49attr_name: cast50type_var: T51operands:52- !ScalarExpression53scalar_const: '42 : i64'54- !ScalarExpression55scalar_fn:56kind: type57attr_name: cast58type_var: T59operands:60- !ScalarExpression61scalar_index: 162
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--- !LinalgOpConfig117metadata: !LinalgOpMetadata118name: test2119cpp_class_name: Test2Op120doc: |-121Title.
122
123Detailed description.
124structured_op: !LinalgStructuredOpConfig125args:126- !LinalgOperandDefConfig127name: I128kind: input_tensor129type_var: T130shape_map: affine_map<()[s0, s1, s2, s3] -> (s0, s1)>131- !LinalgOperandDefConfig132name: O133kind: output_tensor134type_var: T135shape_map: affine_map<()[s0, s1, s2, s3] -> (s0, s1)>136- !LinalgOperandDefConfig137name: strides138kind: index_attr139index_attr_map: affine_map<()[s0, s1, s2, s3] -> (s2, s3)>140default_indices:141- 1142- 2143indexing_maps: !LinalgIndexingMapsConfig144static_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)>147iterator_types:148- parallel149- parallel150assignments:151- !ScalarAssign152arg: O153value: !ScalarExpression154scalar_arg: I155
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--- !LinalgOpConfig204metadata: !LinalgOpMetadata205name: test3206cpp_class_name: Test3Op207doc: |-208Title.
209
210Detailed description.
211structured_op: !LinalgStructuredOpConfig212args:213- !LinalgOperandDefConfig214name: value215kind: scalar216type_var: T1217- !LinalgOperandDefConfig218name: O219kind: output_tensor220type_var: U221shape_map: affine_map<() -> ()>222indexing_maps: !LinalgIndexingMapsConfig223static_indexing_maps:224- affine_map<() -> ()>225- affine_map<() -> ()>226iterator_types: []227assignments:228- !ScalarAssign229arg: O230value: !ScalarExpression231scalar_fn:232kind: type233fn_name: cast_signed234type_var: U235operands:236- !ScalarExpression237scalar_arg: value238
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--- !LinalgOpConfig258metadata: !LinalgOpMetadata259name: test4260cpp_class_name: Test4Op261doc: |-262Title.
263
264Detailed description.
265structured_op: !LinalgStructuredOpConfig266args:267- !LinalgOperandDefConfig268name: O269kind: output_tensor270type_var: T271shape_map: affine_map<()[s0, s1] -> (s0, s1)>272- !LinalgOperandDefConfig273name: unary_fun274kind: unary_fn_attr275default_fn: exp276- !LinalgOperandDefConfig277name: binary_fun278kind: binary_fn_attr279default_fn: add280indexing_maps: !LinalgIndexingMapsConfig281static_indexing_maps:282- affine_map<(d0, d1)[s0, s1] -> (d0, d1)>283iterator_types:284- parallel285- parallel286assignments:287- !ScalarAssign288arg: O289value: !ScalarExpression290scalar_fn:291kind: binary292attr_name: binary_fun293operands:294- !ScalarExpression295scalar_fn:296kind: unary297attr_name: unary_fun298operands:299- !ScalarExpression300scalar_arg: O301- !ScalarExpression302scalar_arg: O303
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--- !LinalgOpConfig337metadata: !LinalgOpMetadata338name: test5339cpp_class_name: Test5Op340doc: |-341Title.
342
343Detailed description.
344implements:345- LinalgFillOpInterface346defines:347- hasCanonicalizer348structured_op: !LinalgStructuredOpConfig349args:350- !LinalgOperandDefConfig351name: value352kind: scalar353type_var: T1354- !LinalgOperandDefConfig355name: O356kind: output_tensor357type_var: U358shape_map: affine_map<() -> ()>359indexing_maps: !LinalgIndexingMapsConfig360static_indexing_maps:361- affine_map<() -> ()>362- affine_map<() -> ()>363iterator_types: []364assignments:365- !ScalarAssign366arg: O367value: !ScalarExpression368scalar_fn:369kind: type370fn_name: cast371type_var: U372operands:373- !ScalarExpression374scalar_arg: value375
376# ODS-LABEL: def Test5Op : LinalgStructuredBase_Op<"test5"
377# ODS-NEXT: /*extraInterfaces=*/[LinalgFillOpInterface])>
378
379# ODS: let hasCanonicalizer = 1;
380