387 lines · yaml
1# RUN: mlir-linalg-ods-yaml-gen %s --o-ods-decl=- | FileCheck %s --check-prefix=ODS2# RUN: mlir-linalg-ods-yaml-gen %s --o-impl=- | FileCheck %s --check-prefix=IMPL3 4# @linalg_structured_op5# 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: !LinalgOpMetadata15 name: test116 cpp_class_name: Test1Op17 doc: |-18 Title.19 20 Detailed description.21structured_op: !LinalgStructuredOpConfig22 args:23 - !LinalgOperandDefConfig24 name: O25 kind: output_tensor26 type_var: T27 shape_map: affine_map<()[s0, s1] -> (s0, s1)>28 - !LinalgOperandDefConfig29 name: cast30 kind: type_fn_attr31 default_fn: cast_signed32 indexing_maps: !LinalgIndexingMapsConfig33 static_indexing_maps:34 - affine_map<(d0, d1)[s0, s1] -> (d0, d1)>35 iterator_types:36 - parallel37 - parallel38 assignments:39 - !ScalarAssign40 arg: O41 value: !ScalarExpression42 scalar_fn:43 kind: binary44 fn_name: add45 operands:46 - !ScalarExpression47 scalar_fn:48 kind: type49 attr_name: cast50 type_var: T51 operands:52 - !ScalarExpression53 scalar_const: '42 : i64'54 - !ScalarExpression55 scalar_fn:56 kind: type57 attr_name: cast58 type_var: T59 operands:60 - !ScalarExpression61 scalar_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">:$cast72 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-NEXT: function_ref<InFlightDiagnostic()> emitError)92# IMPL: TypeFn castVal = TypeFn::cast_signed;93# IMPL-NEXT: auto castIter = llvm::find_if(attrs, [&](const NamedAttribute &attr) {94# IMPL-NEXT: return attr.getName() == "cast"; });95# IMPL-NEXT: if (castIter != attrs.end()) {96# IMPL-NEXT: if (auto attr = llvm::dyn_cast<TypeFnAttr>(castIter->getValue()))97# IMPL-NEXT: castVal = attr.getValue();98# IMPL-NEXT: }99 100# IMPL: Value [[VAL0:[a-z0-9]+]] = helper.constant("42 : i64");101# IMPL-DAG: Value [[VAL1:[a-z0-9]+]] = helper.buildTypeFn(castVal, block.getArgument(0).getType(), [[VAL0]], emitError);102# IMPL-DAG: Value [[VAL2:[a-z0-9]+]] = helper.index(1);103# IMPL-DAG: Value [[VAL3:[a-z0-9]+]] = helper.buildTypeFn(castVal, block.getArgument(0).getType(), [[VAL2]], emitError);104# IMPL-DAG: Value [[VAL4:[a-z0-9]+]] = helper.buildBinaryFn(BinaryFn::add, [[VAL1]], [[VAL3]], emitError);105 106 107# @linalg_structured_op108# def test2(I=TensorDef(T, S.M, S.N),109# O=TensorDef(T, S.M, S.N, output=True),110# strides=IndexAttrDef(S.SM, S.SN, default=[1, 2])):111# """Title.112 113# Detailed description.114# """115# O[D.m, D.n] = I[D.n * S.SM, D.m * S.SN]116 117--- !LinalgOpConfig118metadata: !LinalgOpMetadata119 name: test2120 cpp_class_name: Test2Op121 doc: |-122 Title.123 124 Detailed description.125structured_op: !LinalgStructuredOpConfig126 args:127 - !LinalgOperandDefConfig128 name: I129 kind: input_tensor130 type_var: T131 shape_map: affine_map<()[s0, s1, s2, s3] -> (s0, s1)>132 - !LinalgOperandDefConfig133 name: O134 kind: output_tensor135 type_var: T136 shape_map: affine_map<()[s0, s1, s2, s3] -> (s0, s1)>137 - !LinalgOperandDefConfig138 name: strides139 kind: index_attr140 index_attr_map: affine_map<()[s0, s1, s2, s3] -> (s2, s3)>141 default_indices:142 - 1143 - 2144 indexing_maps: !LinalgIndexingMapsConfig145 static_indexing_maps:146 - affine_map<(d0, d1)[s0, s1, s2, s3] -> (d1 * s2, d0 * s3)>147 - affine_map<(d0, d1)[s0, s1, s2, s3] -> (d0, d1)>148 iterator_types:149 - parallel150 - parallel151 assignments:152 - !ScalarAssign153 arg: O154 value: !ScalarExpression155 scalar_arg: I156 157# ODS-LABEL: def Test2Op : LinalgStructuredBase_Op<"test2"158 159# ODS: let arguments =160# ODS-NEXT: Variadic<AnyType>:$inputs,161# ODS-NEXT: Variadic<AnyShaped>:$outputs,162# ODS-NEXT: DefaultValuedOptionalAttr<RankedI64ElementsAttr<[2]>163# ODS-SAME: "{ static_cast<int64_t>(1), static_cast<int64_t>(2) }">:$strides164 165# ODS: "Attribute":$strides166# ODS: $_state.addAttribute("strides", strides);167 168# ODS: bool hasDynamicIndexingMaps();169# ODS-NEXT: LogicalResult verifyIndexingMapRequiredAttributes();170 171# IMPL: getSymbolBindings(Test2Op self)172# IMPL: cst2 = self.getStrides().getValues<int64_t>()[0];173# IMPL-NEXT: getAffineConstantExpr(cst2, context)174# IMPL: cst3 = self.getStrides().getValues<int64_t>()[1];175# IMPL-NEXT: getAffineConstantExpr(cst3, context)176 177# IMPL: Test2Op::getIndexingMaps()178# IMPL: = getSymbolBindings(*this);179# IMPL: "affine_map<(d0, d1)[s0, s1, s2, s3] -> (d1 * s2, d0 * s3)>"180# IMPL: "affine_map<(d0, d1)[s0, s1, s2, s3] -> (d0, d1)>"181 182# IMPL: Test2Op::getNumRegionArgs() { return 2; }183 184# IMPL: Test2Op::hasDynamicIndexingMaps() { return true; }185# IMPL: Test2Op::verifyIndexingMapRequiredAttributes()186# IMPL: auto attr = op->getAttrOfType<DenseElementsAttr>("strides")187# IMPL: "incorrect element type for index attribute 'strides'"188# IMPL: "incorrect shape for index attribute 'strides'"189# IMPL: void Test2Op::regionBuilder(ImplicitLocOpBuilder &b,190# IMPL-NEXT: Block &block, ArrayRef<NamedAttribute> attrs,191# IMPL-NEXT: function_ref<InFlightDiagnostic()> emitError)192# IMPL-NEXT: assert(2 > 0 && block.getNumArguments() == 2 &&193 194# IMPL: yields.push_back(block.getArgument(0));195 196# @linalg_structured_op197# def test3(value=ScalarDef(T1),198# O=TensorDef(U, output=True)):199# """Title.200 201# Detailed description.202# """203# O[None] = TypeFn.cast_signed(U, value)204 205--- !LinalgOpConfig206metadata: !LinalgOpMetadata207 name: test3208 cpp_class_name: Test3Op209 doc: |-210 Title.211 212 Detailed description.213structured_op: !LinalgStructuredOpConfig214 args:215 - !LinalgOperandDefConfig216 name: value217 kind: scalar218 type_var: T1219 - !LinalgOperandDefConfig220 name: O221 kind: output_tensor222 type_var: U223 shape_map: affine_map<() -> ()>224 indexing_maps: !LinalgIndexingMapsConfig225 static_indexing_maps:226 - affine_map<() -> ()>227 - affine_map<() -> ()>228 iterator_types: []229 assignments:230 - !ScalarAssign231 arg: O232 value: !ScalarExpression233 scalar_fn:234 kind: type235 fn_name: cast_signed236 type_var: U237 operands:238 - !ScalarExpression239 scalar_arg: value240 241# IMPL: Test3Op::getIteratorTypesArray() {242# IMPL-NEXT: int64_t rank = getRank(getDpsInitOperand(0));243 244# IMPL: Test3Op::getIndexingMaps() {245# IMPL-NEXT: MLIRContext *context = getContext();246# IMPL-NEXT: AffineMap scalarMap = AffineMap::get(getNumParallelLoops(), 0, context);247# IMPL-NEXT: AffineMap tensorMap = AffineMap::getMultiDimIdentityMap(248 249# @linalg_structured_op250# def test4(O=TensorDef(T, S.M, S.N, output=True),251# unary_fun=UnaryFnAttrDef(default=UnaryFn.exp),252# binary_fun=BinaryFnAttrDef(default=BinaryFn.add)):253# """Title.254 255# Detailed description.256# """257# O[D.m, D.n] = binary_fun(unary_fun(O[D.m, D.n]), O[D.m, D.n])258 259--- !LinalgOpConfig260metadata: !LinalgOpMetadata261 name: test4262 cpp_class_name: Test4Op263 doc: |-264 Title.265 266 Detailed description.267structured_op: !LinalgStructuredOpConfig268 args:269 - !LinalgOperandDefConfig270 name: O271 kind: output_tensor272 type_var: T273 shape_map: affine_map<()[s0, s1] -> (s0, s1)>274 - !LinalgOperandDefConfig275 name: unary_fun276 kind: unary_fn_attr277 default_fn: exp278 - !LinalgOperandDefConfig279 name: binary_fun280 kind: binary_fn_attr281 default_fn: add282 indexing_maps: !LinalgIndexingMapsConfig283 static_indexing_maps:284 - affine_map<(d0, d1)[s0, s1] -> (d0, d1)>285 iterator_types:286 - parallel287 - parallel288 assignments:289 - !ScalarAssign290 arg: O291 value: !ScalarExpression292 scalar_fn:293 kind: binary294 attr_name: binary_fun295 operands:296 - !ScalarExpression297 scalar_fn:298 kind: unary299 attr_name: unary_fun300 operands:301 - !ScalarExpression302 scalar_arg: O303 - !ScalarExpression304 scalar_arg: O305 306# ODS-LABEL: def Test4Op : LinalgStructuredBase_Op<"test4"307 308# ODS: let arguments =309# ODS-NEXT: Variadic<AnyType>:$inputs,310# ODS-NEXT: Variadic<AnyShaped>:$outputs,311# ODS-NEXT: DefaultValuedOptionalAttr<UnaryFnAttr, "UnaryFn::exp">:$unary_fun,312# ODS-NEXT: DefaultValuedOptionalAttr<BinaryFnAttr, "BinaryFn::add">:$binary_fun313 314# ODS: "Attribute":$unary_fun, "Attribute":$binary_fun,315 316# ODS: $_state.addAttribute("unary_fun", unary_fun)317# ODS-NEXT: $_state.addAttribute("binary_fun", binary_fun)318 319# IMPL-LABEL: void Test4Op::regionBuilder(ImplicitLocOpBuilder &b,320# IMPL-NEXT: Block &block, ArrayRef<NamedAttribute> attrs,321# IMPL-NEXT: function_ref<InFlightDiagnostic()> emitError)322# IMPL: UnaryFn unary_funVal = UnaryFn::exp323# IMPL: BinaryFn binary_funVal = BinaryFn::add324 325# IMPL: Value [[VAL0:[a-z0-9]+]] = helper.buildUnaryFn(unary_funVal, block.getArgument(0), emitError);326# IMPL-NEXT: if (![[VAL0]])327# IMPL-NEXT: return;328# IMPL: Value [[VAL1:[a-z0-9]+]] = helper.buildBinaryFn(binary_funVal, [[VAL0]], block.getArgument(0), emitError);329# IMPL-NEXT: if (![[VAL1]])330# IMPL-NEXT: return;331# IMPL: yields.push_back([[VAL1]])332 333# @linalg_structured_op334# def test5(value=ScalarDef(T1), O=TensorDef(U, output=True)):335# """Title.336 337# Detailed description.338# """339# implements(FillOpInterface)340# defines(Canonicalizer)341# O[None] = TypeFn.cast(U, value)342 343--- !LinalgOpConfig344metadata: !LinalgOpMetadata345 name: test5346 cpp_class_name: Test5Op347 doc: |-348 Title.349 350 Detailed description.351 implements:352 - LinalgFillOpInterface353 defines:354 - hasCanonicalizer355structured_op: !LinalgStructuredOpConfig356 args:357 - !LinalgOperandDefConfig358 name: value359 kind: scalar360 type_var: T1361 - !LinalgOperandDefConfig362 name: O363 kind: output_tensor364 type_var: U365 shape_map: affine_map<() -> ()>366 indexing_maps: !LinalgIndexingMapsConfig367 static_indexing_maps:368 - affine_map<() -> ()>369 - affine_map<() -> ()>370 iterator_types: []371 assignments:372 - !ScalarAssign373 arg: O374 value: !ScalarExpression375 scalar_fn:376 kind: type377 fn_name: cast378 type_var: U379 operands:380 - !ScalarExpression381 scalar_arg: value382 383# ODS-LABEL: def Test5Op : LinalgStructuredBase_Op<"test5"384# ODS-NEXT: /*extraInterfaces=*/[LinalgFillOpInterface])>385 386# ODS: let hasCanonicalizer = 1;387