xref: /llvm-project/mlir/test/python/python_test_ops.td (revision 3766ba44a8945681f4c52acb0331efcff66ef7b1) 
1//===-- python_test_ops.td - Python test Op definitions ----*- tablegen -*-===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8
9#ifndef PYTHON_TEST_OPS
10#define PYTHON_TEST_OPS
11
12include "mlir/IR/AttrTypeBase.td"
13include "mlir/IR/OpBase.td"
14include "mlir/Interfaces/InferTypeOpInterface.td"
15
16def Python_Test_Dialect : Dialect {
17  let name = "python_test";
18  let cppNamespace = "python_test";
19
20  let useDefaultTypePrinterParser = 1;
21  let useDefaultAttributePrinterParser = 1;
22}
23
24class TestType<string name, string typeMnemonic>
25    : TypeDef<Python_Test_Dialect, name> {
26  let mnemonic = typeMnemonic;
27}
28
29class TestAttr<string name, string attrMnemonic>
30    : AttrDef<Python_Test_Dialect, name> {
31  let mnemonic = attrMnemonic;
32}
33
34class TestOp<string mnemonic, list<Trait> traits = []>
35    : Op<Python_Test_Dialect, mnemonic, traits> {
36  let assemblyFormat = "operands attr-dict functional-type(operands, results)";
37}
38
39//===----------------------------------------------------------------------===//
40// Type definitions.
41//===----------------------------------------------------------------------===//
42
43def TestType : TestType<"TestType", "test_type">;
44
45//===----------------------------------------------------------------------===//
46// Attribute definitions.
47//===----------------------------------------------------------------------===//
48
49def TestAttr : TestAttr<"TestAttr", "test_attr">;
50
51//===----------------------------------------------------------------------===//
52// Operation definitions.
53//===----------------------------------------------------------------------===//
54
55def AttributedOp : TestOp<"attributed_op"> {
56  let arguments = (ins I32Attr:$mandatory_i32,
57                   OptionalAttr<I32Attr>:$optional_i32,
58                   UnitAttr:$unit);
59}
60
61def CustomAttributedOp : TestOp<"custom_attributed_op"> {
62  let arguments = (ins TestAttr:$test_attr);
63}
64
65def AttributesOp : TestOp<"attributes_op"> {
66  let arguments = (ins
67                   AffineMapArrayAttr:$x_affinemaparr,
68                   AffineMapAttr:$x_affinemap,
69                   ArrayAttr:$x_arr,
70                   BoolArrayAttr:$x_boolarr,
71                   BoolAttr:$x_bool,
72                   DenseBoolArrayAttr:$x_dboolarr,
73                   DenseF32ArrayAttr:$x_df32arr,
74                   DenseF64ArrayAttr:$x_df64arr,
75                   DenseI16ArrayAttr:$x_df16arr,
76                   DenseI32ArrayAttr:$x_di32arr,
77                   DenseI64ArrayAttr:$x_di64arr,
78                   DenseI8ArrayAttr:$x_di8arr,
79                   DictArrayAttr:$x_dictarr,
80                   DictionaryAttr:$x_dict,
81                   F32ArrayAttr:$x_f32arr,
82                   F32Attr:$x_f32,
83                   F64ArrayAttr:$x_f64arr,
84                   F64Attr:$x_f64,
85                   F64ElementsAttr:$x_f64elems,
86                   FlatSymbolRefArrayAttr:$x_flatsymrefarr,
87                   FlatSymbolRefAttr:$x_flatsymref,
88                   I16Attr:$x_i16,
89                   I1Attr:$x_i1,
90                   I32ArrayAttr:$x_i32arr,
91                   I32Attr:$x_i32,
92                   I32ElementsAttr:$x_i32elems,
93                   I64ArrayAttr:$x_i64arr,
94                   I64Attr:$x_i64,
95                   I64ElementsAttr:$x_i64elems,
96                   I64SmallVectorArrayAttr:$x_i64svecarr,
97                   I8Attr:$x_i8,
98                   IndexAttr:$x_idx,
99                   IndexElementsAttr:$x_idxelems,
100                   IndexListArrayAttr:$x_idxlistarr,
101                   SI16Attr:$x_si16,
102                   SI1Attr:$x_si1,
103                   SI32Attr:$x_si32,
104                   SI64Attr:$x_si64,
105                   SI8Attr:$x_si8,
106                   StrArrayAttr:$x_strarr,
107                   StrAttr:$x_str,
108                   SymbolNameAttr:$x_sym,
109                   SymbolRefArrayAttr:$x_symrefarr,
110                   SymbolRefAttr:$x_symref,
111                   TypeArrayAttr:$x_typearr,
112                   TypeAttr:$x_type,
113                   UI16Attr:$x_ui16,
114                   UI1Attr:$x_ui1,
115                   UI32Attr:$x_ui32,
116                   UI64Attr:$x_ui64,
117                   UI8Attr:$x_ui8,
118                   UnitAttr:$x_unit
119                   );
120}
121
122def PropertyOp : TestOp<"property_op"> {
123  let arguments = (ins I32Attr:$property,
124                   I32:$idx);
125}
126
127def DummyOp : TestOp<"dummy_op"> {
128}
129
130def InferResultsOp : TestOp<"infer_results_op", [InferTypeOpInterface]> {
131  let arguments = (ins);
132  let results = (outs AnyInteger:$single, AnyInteger:$doubled);
133
134  let extraClassDeclaration = [{
135    static ::llvm::LogicalResult inferReturnTypes(
136      ::mlir::MLIRContext *context, ::std::optional<::mlir::Location> location,
137      ::mlir::ValueRange operands, ::mlir::DictionaryAttr attributes,
138      ::mlir::OpaqueProperties,
139      ::mlir::RegionRange regions,
140      ::llvm::SmallVectorImpl<::mlir::Type> &inferredReturnTypes) {
141      ::mlir::Builder b(context);
142      inferredReturnTypes.push_back(b.getI32Type());
143      inferredReturnTypes.push_back(b.getI64Type());
144      return ::mlir::success();
145    }
146  }];
147}
148
149def I32OrF32 : TypeConstraint<Or<[I32.predicate, F32.predicate]>,
150                                 "i32 or f32">;
151
152def InferResultsVariadicInputsOp : TestOp<"infer_results_variadic_inputs_op",
153    [InferTypeOpInterface, AttrSizedOperandSegments]> {
154  let arguments = (ins Optional<I64>:$single, Optional<I64>:$doubled);
155  let results = (outs I32OrF32:$res);
156
157  let extraClassDeclaration = [{
158    static ::llvm::LogicalResult inferReturnTypes(
159      ::mlir::MLIRContext *context, ::std::optional<::mlir::Location> location,
160      ::mlir::ValueRange operands, ::mlir::DictionaryAttr attributes,
161      ::mlir::OpaqueProperties,
162      ::mlir::RegionRange regions,
163      ::llvm::SmallVectorImpl<::mlir::Type> &inferredReturnTypes) {
164      ::mlir::Builder b(context);
165      if (operands.size() == 1)
166          inferredReturnTypes.push_back(b.getI32Type());
167      else if (operands.size() == 2)
168          inferredReturnTypes.push_back(b.getF32Type());
169      return ::mlir::success();
170    }
171  }];
172}
173
174// If all result types are buildable, the InferTypeOpInterface is implied and is
175// autogenerated by C++ ODS.
176def InferResultsImpliedOp : TestOp<"infer_results_implied_op"> {
177  let results = (outs I32:$integer, F64:$flt, Index:$index);
178}
179
180def InferShapedTypeComponentsOp : TestOp<"infer_shaped_type_components_op",
181  [DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
182                             ["inferReturnTypeComponents"]>]> {
183  let arguments = (ins AnyTensor:$operand);
184  let results = (outs AnyTensor:$result);
185
186  let extraClassDefinition = [{
187    ::llvm::LogicalResult $cppClass::inferReturnTypeComponents(
188      ::mlir::MLIRContext *context, ::std::optional<::mlir::Location> location,
189      ::mlir::ValueShapeRange operands, ::mlir::DictionaryAttr attributes,
190      ::mlir::OpaqueProperties properties, ::mlir::RegionRange regions,
191      ::llvm::SmallVectorImpl<
192        ::mlir::ShapedTypeComponents>& inferredShapedTypeComponents) {
193      $cppClass::Adaptor adaptor(operands, attributes, properties, regions);
194      auto operandType =
195          ::llvm::cast<::mlir::ShapedType>(adaptor.getOperand().getType());
196      if (operandType.hasRank()) {
197        inferredShapedTypeComponents.emplace_back(operandType.getShape(),
198            operandType.getElementType());
199      } else {
200        inferredShapedTypeComponents.emplace_back(operandType.getElementType());
201      }
202      return ::mlir::success();
203    }
204  }];
205}
206
207def SameOperandAndResultTypeOp : TestOp<"same_operand_and_result_type_op",
208                                        [SameOperandsAndResultType]> {
209  let arguments = (ins Variadic<AnyType>);
210  let results = (outs AnyType:$one, AnyType:$two);
211}
212
213def FirstAttrDeriveTypeAttrOp : TestOp<"first_attr_derive_type_attr_op",
214                               [FirstAttrDerivedResultType]> {
215  let arguments = (ins AnyType:$input, TypeAttr:$type);
216  let results = (outs AnyType:$one, AnyType:$two);
217}
218
219def FirstAttrDeriveAttrOp : TestOp<"first_attr_derive_attr_op",
220                               [FirstAttrDerivedResultType]> {
221  let arguments = (ins AnyAttr:$iattr);
222  let results = (outs AnyType:$one, AnyType:$two, AnyType:$three);
223}
224
225def OptionalOperandOp : TestOp<"optional_operand_op"> {
226  let arguments = (ins Optional<AnyType>:$input);
227  let results = (outs I32:$result);
228}
229
230def SameVariadicOperandSizeOp : TestOp<"same_variadic_operand",
231                                       [SameVariadicOperandSize]> {
232  let arguments = (ins Variadic<AnyType>:$variadic1, AnyType:$non_variadic,
233                   Variadic<AnyType>:$variadic2);
234}
235
236// Check different arrangements of variadic groups
237def SameVariadicResultSizeOpVFV : TestOp<"same_variadic_result_vfv",
238                                        [SameVariadicResultSize]> {
239  let results = (outs Variadic<AnyType>:$variadic1, AnyType:$non_variadic,
240                 Variadic<AnyType>:$variadic2);
241}
242
243def SameVariadicResultSizeOpVVV : TestOp<"same_variadic_result_vvv",
244                                        [SameVariadicResultSize]> {
245  let results = (outs Variadic<AnyType>:$variadic1, Variadic<AnyType>:$variadic2,
246                 Variadic<AnyType>:$variadic3);
247}
248
249def SameVariadicResultSizeOpFFV : TestOp<"same_variadic_result_ffv",
250                                        [SameVariadicResultSize]> {
251  let results = (outs AnyType:$non_variadic1, AnyType:$non_variadic2,
252                 Variadic<AnyType>:$variadic);
253}
254
255def SameVariadicResultSizeOpVVF : TestOp<"same_variadic_result_vvf",
256                                        [SameVariadicResultSize]> {
257  let results = (outs Variadic<AnyType>:$variadic1, Variadic<AnyType>:$variadic2,
258                 AnyType:$non_variadic);
259}
260
261def SameVariadicResultSizeOpFVFVF : TestOp<"same_variadic_result_fvfvf",
262                                          [SameVariadicResultSize]> {
263  let results = (outs AnyType:$non_variadic1, Variadic<AnyType>:$variadic1,
264                 AnyType:$non_variadic2, Variadic<AnyType>:$variadic2,
265                 AnyType:$non_variadic3);
266}
267
268#endif // PYTHON_TEST_OPS
269