• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1//===- Ops.td - Toy dialect operation 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// Defines the operations of the Toy dialect.
10//
11//===----------------------------------------------------------------------===//
12
13#ifndef TOY_OPS
14#define TOY_OPS
15
16include "mlir/IR/OpBase.td"
17include "mlir/Interfaces/SideEffectInterfaces.td"
18
19// Provide a definition of the 'toy' dialect in the ODS framework so that we
20// can define our operations.
21def Toy_Dialect : Dialect {
22  let name = "toy";
23  let cppNamespace = "::mlir::toy";
24}
25
26// Base class for toy dialect operations. This operation inherits from the base
27// `Op` class in OpBase.td, and provides:
28//   * The parent dialect of the operation.
29//   * The mnemonic for the operation, or the name without the dialect prefix.
30//   * A list of traits for the operation.
31class Toy_Op<string mnemonic, list<OpTrait> traits = []> :
32    Op<Toy_Dialect, mnemonic, traits>;
33
34//===----------------------------------------------------------------------===//
35// Toy Operations
36//===----------------------------------------------------------------------===//
37
38// We define a toy operation by inheriting from our base 'Toy_Op' class above.
39// Here we provide the mnemonic and a list of traits for the operation. The
40// constant operation is marked as 'NoSideEffect' as it is a pure operation
41// and may be removed if dead.
42def ConstantOp : Toy_Op<"constant", [NoSideEffect]> {
43  // Provide a summary and description for this operation. This can be used to
44  // auto-generate documentation of the operations within our dialect.
45  let summary = "constant";
46  let description = [{
47    Constant operation turns a literal into an SSA value. The data is attached
48    to the operation as an attribute. For example:
49
50    ```mlir
51      %0 = toy.constant dense<[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]>
52                        : tensor<2x3xf64>
53    ```
54  }];
55
56  // The constant operation takes an attribute as the only input.
57  let arguments = (ins F64ElementsAttr:$value);
58
59  // The constant operation returns a single value of TensorType.
60  let results = (outs F64Tensor);
61
62  // Specify a parser and printer method.
63  let parser = [{ return ::parseConstantOp(parser, result); }];
64  let printer = [{ return ::print(p, *this); }];
65
66  // Add custom build methods for the constant operation. These method populates
67  // the `state` that MLIR uses to create operations, i.e. these are used when
68  // using `builder.create<ConstantOp>(...)`.
69  let builders = [
70    // Build a constant with a given constant tensor value.
71    OpBuilderDAG<(ins "DenseElementsAttr":$value), [{
72      build($_builder, $_state, value.getType(), value);
73    }]>,
74
75    // Build a constant with a given constant floating-point value.
76    OpBuilderDAG<(ins "double":$value)>
77  ];
78
79  // Invoke a static verify method to verify this constant operation.
80  let verifier = [{ return ::verify(*this); }];
81}
82
83def AddOp : Toy_Op<"add"> {
84  let summary = "element-wise addition operation";
85  let description = [{
86    The "add" operation performs element-wise addition between two tensors.
87    The shapes of the tensor operands are expected to match.
88  }];
89
90  let arguments = (ins F64Tensor:$lhs, F64Tensor:$rhs);
91  let results = (outs F64Tensor);
92
93  // Specify a parser and printer method.
94  let parser = [{ return ::parseBinaryOp(parser, result); }];
95  let printer = [{ return ::printBinaryOp(p, *this); }];
96
97  // Allow building an AddOp with from the two input operands.
98  let builders = [
99    OpBuilderDAG<(ins "Value":$lhs, "Value":$rhs)>
100  ];
101}
102
103def GenericCallOp : Toy_Op<"generic_call"> {
104  let summary = "generic call operation";
105  let description = [{
106    Generic calls represent calls to a user defined function that needs to
107    be specialized for the shape of its arguments. The callee name is attached
108    as a symbol reference via an attribute. The arguments list must match the
109    arguments expected by the callee. For example:
110
111    ```mlir
112     %4 = toy.generic_call @my_func(%1, %3)
113           : (tensor<2x3xf64>, tensor<2x3xf64>) -> tensor<*xf64>
114    ```
115
116    This is only valid if a function named "my_func" exists and takes two
117    arguments.
118  }];
119
120  // The generic call operation takes a symbol reference attribute as the
121  // callee, and inputs for the call.
122  let arguments = (ins FlatSymbolRefAttr:$callee, Variadic<F64Tensor>:$inputs);
123
124  // The generic call operation returns a single value of TensorType.
125  let results = (outs F64Tensor);
126
127  // Specialize assembly printing and parsing using a declarative format.
128  let assemblyFormat = [{
129    $callee `(` $inputs `)` attr-dict `:` functional-type($inputs, results)
130  }];
131
132  // Add custom build methods for the generic call operation.
133  let builders = [
134    OpBuilderDAG<(ins "StringRef":$callee, "ArrayRef<Value>":$arguments)>
135  ];
136}
137
138def MulOp : Toy_Op<"mul"> {
139  let summary = "element-wise multiplication operation";
140  let description = [{
141    The "mul" operation performs element-wise multiplication between two
142    tensors. The shapes of the tensor operands are expected to match.
143  }];
144
145  let arguments = (ins F64Tensor:$lhs, F64Tensor:$rhs);
146  let results = (outs F64Tensor);
147
148  // Specify a parser and printer method.
149  let parser = [{ return ::parseBinaryOp(parser, result); }];
150  let printer = [{ return ::printBinaryOp(p, *this); }];
151
152  // Allow building a MulOp with from the two input operands.
153  let builders = [
154    OpBuilderDAG<(ins "Value":$lhs, "Value":$rhs)>
155  ];
156}
157
158def PrintOp : Toy_Op<"print"> {
159  let summary = "print operation";
160  let description = [{
161    The "print" builtin operation prints a given input tensor, and produces
162    no results.
163  }];
164
165  // The print operation takes an input tensor to print.
166  let arguments = (ins F64Tensor:$input);
167
168  let assemblyFormat = "$input attr-dict `:` type($input)";
169}
170
171def ReshapeOp : Toy_Op<"reshape"> {
172  let summary = "tensor reshape operation";
173  let description = [{
174    Reshape operation is transforming its input tensor into a new tensor with
175    the same number of elements but different shapes. For example:
176
177    ```mlir
178       %0 = toy.reshape (%arg1 : tensor<10xf64>) to tensor<5x2xf64>
179    ```
180  }];
181
182  let arguments = (ins F64Tensor:$input);
183
184  // We expect that the reshape operation returns a statically shaped tensor.
185  let results = (outs StaticShapeTensorOf<[F64]>);
186
187  let assemblyFormat = [{
188    `(` $input `:` type($input) `)` attr-dict `to` type(results)
189  }];
190}
191
192def ReturnOp : Toy_Op<"return", [NoSideEffect, HasParent<"FuncOp">,
193                                 Terminator]> {
194  let summary = "return operation";
195  let description = [{
196    The "return" operation represents a return operation within a function.
197    The operation takes an optional tensor operand and produces no results.
198    The operand type must match the signature of the function that contains
199    the operation. For example:
200
201    ```mlir
202      func @foo() -> tensor<2xf64> {
203        ...
204        toy.return %0 : tensor<2xf64>
205      }
206    ```
207  }];
208
209  // The return operation takes an optional input operand to return. This
210  // value must match the return type of the enclosing function.
211  let arguments = (ins Variadic<F64Tensor>:$input);
212
213  // The return operation only emits the input in the format if it is present.
214  let assemblyFormat = "($input^ `:` type($input))? attr-dict ";
215
216  // Allow building a ReturnOp with no return operand.
217  let builders = [
218    OpBuilderDAG<(ins), [{ build($_builder, $_state, llvm::None); }]>
219  ];
220
221  // Provide extra utility definitions on the c++ operation class definition.
222  let extraClassDeclaration = [{
223    bool hasOperand() { return getNumOperands() != 0; }
224  }];
225
226  // Invoke a static verify method to verify this return operation.
227  let verifier = [{ return ::verify(*this); }];
228}
229
230def TransposeOp : Toy_Op<"transpose"> {
231  let summary = "transpose operation";
232
233  let arguments = (ins F64Tensor:$input);
234  let results = (outs F64Tensor);
235
236  let assemblyFormat = [{
237    `(` $input `:` type($input) `)` attr-dict `to` type(results)
238  }];
239
240  // Allow building a TransposeOp with from the input operand.
241  let builders = [
242    OpBuilderDAG<(ins "Value":$input)>
243  ];
244
245  // Invoke a static verify method to verify this transpose operation.
246  let verifier = [{ return ::verify(*this); }];
247}
248
249#endif // TOY_OPS
250