252 lines
8.6 KiB
TableGen
252 lines
8.6 KiB
TableGen
|
//===- Ops.td - Toy dialect operation definitions ----------*- tablegen -*-===//
|
||
|
//
|
||
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
||
|
// See https://llvm.org/LICENSE.txt for license information.
|
||
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||
|
//
|
||
|
//===----------------------------------------------------------------------===//
|
||
|
//
|
||
|
// Defines the operations of the Toy dialect.
|
||
|
//
|
||
|
//===----------------------------------------------------------------------===//
|
||
|
|
||
|
#ifndef TOY_OPS
|
||
|
#define TOY_OPS
|
||
|
|
||
|
include "mlir/IR/OpBase.td"
|
||
|
include "mlir/Interfaces/SideEffectInterfaces.td"
|
||
|
|
||
|
// Provide a definition of the 'toy' dialect in the ODS framework so that we
|
||
|
// can define our operations.
|
||
|
def Toy_Dialect : Dialect {
|
||
|
let name = "toy";
|
||
|
let cppNamespace = "toy";
|
||
|
}
|
||
|
|
||
|
// Base class for toy dialect operations. This operation inherits from the base
|
||
|
// `Op` class in OpBase.td, and provides:
|
||
|
// * The parent dialect of the operation.
|
||
|
// * The mnemonic for the operation, or the name without the dialect prefix.
|
||
|
// * A list of traits for the operation.
|
||
|
class Toy_Op<string mnemonic, list<OpTrait> traits = []> :
|
||
|
Op<Toy_Dialect, mnemonic, traits>;
|
||
|
|
||
|
//===----------------------------------------------------------------------===//
|
||
|
// Toy Operations
|
||
|
//===----------------------------------------------------------------------===//
|
||
|
|
||
|
// We define a toy operation by inheriting from our base 'Toy_Op' class above.
|
||
|
// Here we provide the mnemonic and a list of traits for the operation. The
|
||
|
// constant operation is marked as 'NoSideEffect' as it is a pure operation
|
||
|
// and may be removed if dead.
|
||
|
def ConstantOp : Toy_Op<"constant", [NoSideEffect]> {
|
||
|
// Provide a summary and description for this operation. This can be used to
|
||
|
// auto-generate documentation of the operations within our dialect.
|
||
|
let summary = "constant";
|
||
|
let description = [{
|
||
|
Constant operation turns a literal into an SSA value. The data is attached
|
||
|
to the operation as an attribute. For example:
|
||
|
|
||
|
```mlir
|
||
|
%0 = toy.constant dense<[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]>
|
||
|
: tensor<2x3xf64>
|
||
|
```
|
||
|
}];
|
||
|
|
||
|
// The constant operation takes an attribute as the only input.
|
||
|
let arguments = (ins F64ElementsAttr:$value);
|
||
|
|
||
|
// The constant operation returns a single value of TensorType.
|
||
|
let results = (outs F64Tensor);
|
||
|
|
||
|
// Specify a parser and printer method.
|
||
|
let parser = [{ return ::parseConstantOp(parser, result); }];
|
||
|
let printer = [{ return ::print(p, *this); }];
|
||
|
|
||
|
// Add custom build methods for the constant operation. These method populates
|
||
|
// the `state` that MLIR uses to create operations, i.e. these are used when
|
||
|
// using `builder.create<ConstantOp>(...)`.
|
||
|
let builders = [
|
||
|
// Build a constant with a given constant tensor value.
|
||
|
OpBuilder<"OpBuilder &builder, OperationState &state, "
|
||
|
"DenseElementsAttr value", [{
|
||
|
build(builder, state, value.getType(), value);
|
||
|
}]>,
|
||
|
|
||
|
// Build a constant with a given constant floating-point value.
|
||
|
OpBuilder<"OpBuilder &builder, OperationState &state, double value">
|
||
|
];
|
||
|
|
||
|
// Invoke a static verify method to verify this constant operation.
|
||
|
let verifier = [{ return ::verify(*this); }];
|
||
|
}
|
||
|
|
||
|
def AddOp : Toy_Op<"add"> {
|
||
|
let summary = "element-wise addition operation";
|
||
|
let description = [{
|
||
|
The "add" operation performs element-wise addition between two tensors.
|
||
|
The shapes of the tensor operands are expected to match.
|
||
|
}];
|
||
|
|
||
|
let arguments = (ins F64Tensor:$lhs, F64Tensor:$rhs);
|
||
|
let results = (outs F64Tensor);
|
||
|
|
||
|
// Specify a parser and printer method.
|
||
|
let parser = [{ return ::parseBinaryOp(parser, result); }];
|
||
|
let printer = [{ return ::printBinaryOp(p, *this); }];
|
||
|
|
||
|
// Allow building an AddOp with from the two input operands.
|
||
|
let builders = [
|
||
|
OpBuilder<"OpBuilder &b, OperationState &state, Value lhs, Value rhs">
|
||
|
];
|
||
|
}
|
||
|
|
||
|
def GenericCallOp : Toy_Op<"generic_call"> {
|
||
|
let summary = "generic call operation";
|
||
|
let description = [{
|
||
|
Generic calls represent calls to a user defined function that needs to
|
||
|
be specialized for the shape of its arguments. The callee name is attached
|
||
|
as a symbol reference via an attribute. The arguments list must match the
|
||
|
arguments expected by the callee. For example:
|
||
|
|
||
|
```mlir
|
||
|
%4 = toy.generic_call @my_func(%1, %3)
|
||
|
: (tensor<2x3xf64>, tensor<2x3xf64>) -> tensor<*xf64>
|
||
|
```
|
||
|
|
||
|
This is only valid if a function named "my_func" exists and takes two
|
||
|
arguments.
|
||
|
}];
|
||
|
|
||
|
// The generic call operation takes a symbol reference attribute as the
|
||
|
// callee, and inputs for the call.
|
||
|
let arguments = (ins FlatSymbolRefAttr:$callee, Variadic<F64Tensor>:$inputs);
|
||
|
|
||
|
// The generic call operation returns a single value of TensorType.
|
||
|
let results = (outs F64Tensor);
|
||
|
|
||
|
// Specialize assembly printing and parsing using a declarative format.
|
||
|
let assemblyFormat = [{
|
||
|
$callee `(` $inputs `)` attr-dict `:` functional-type($inputs, results)
|
||
|
}];
|
||
|
|
||
|
// Add custom build methods for the generic call operation.
|
||
|
let builders = [
|
||
|
OpBuilder<"OpBuilder &builder, OperationState &state, "
|
||
|
"StringRef callee, ArrayRef<Value> arguments">
|
||
|
];
|
||
|
}
|
||
|
|
||
|
def MulOp : Toy_Op<"mul"> {
|
||
|
let summary = "element-wise multiplication operation";
|
||
|
let description = [{
|
||
|
The "mul" operation performs element-wise multiplication between two
|
||
|
tensors. The shapes of the tensor operands are expected to match.
|
||
|
}];
|
||
|
|
||
|
let arguments = (ins F64Tensor:$lhs, F64Tensor:$rhs);
|
||
|
let results = (outs F64Tensor);
|
||
|
|
||
|
// Specify a parser and printer method.
|
||
|
let parser = [{ return ::parseBinaryOp(parser, result); }];
|
||
|
let printer = [{ return ::printBinaryOp(p, *this); }];
|
||
|
|
||
|
// Allow building a MulOp with from the two input operands.
|
||
|
let builders = [
|
||
|
OpBuilder<"OpBuilder &b, OperationState &state, Value lhs, Value rhs">
|
||
|
];
|
||
|
}
|
||
|
|
||
|
def PrintOp : Toy_Op<"print"> {
|
||
|
let summary = "print operation";
|
||
|
let description = [{
|
||
|
The "print" builtin operation prints a given input tensor, and produces
|
||
|
no results.
|
||
|
}];
|
||
|
|
||
|
// The print operation takes an input tensor to print.
|
||
|
let arguments = (ins F64Tensor:$input);
|
||
|
|
||
|
let assemblyFormat = "$input attr-dict `:` type($input)";
|
||
|
}
|
||
|
|
||
|
def ReshapeOp : Toy_Op<"reshape"> {
|
||
|
let summary = "tensor reshape operation";
|
||
|
let description = [{
|
||
|
Reshape operation is transforming its input tensor into a new tensor with
|
||
|
the same number of elements but different shapes. For example:
|
||
|
|
||
|
```mlir
|
||
|
%0 = toy.reshape (%arg1 : tensor<10xf64>) to tensor<5x2xf64>
|
||
|
```
|
||
|
}];
|
||
|
|
||
|
let arguments = (ins F64Tensor:$input);
|
||
|
|
||
|
// We expect that the reshape operation returns a statically shaped tensor.
|
||
|
let results = (outs StaticShapeTensorOf<[F64]>);
|
||
|
|
||
|
let assemblyFormat = [{
|
||
|
`(` $input `:` type($input) `)` attr-dict `to` type(results)
|
||
|
}];
|
||
|
}
|
||
|
|
||
|
def ReturnOp : Toy_Op<"return", [NoSideEffect, HasParent<"FuncOp">,
|
||
|
Terminator]> {
|
||
|
let summary = "return operation";
|
||
|
let description = [{
|
||
|
The "return" operation represents a return operation within a function.
|
||
|
The operation takes an optional tensor operand and produces no results.
|
||
|
The operand type must match the signature of the function that contains
|
||
|
the operation. For example:
|
||
|
|
||
|
```mlir
|
||
|
func @foo() -> tensor<2xf64> {
|
||
|
...
|
||
|
toy.return %0 : tensor<2xf64>
|
||
|
}
|
||
|
```
|
||
|
}];
|
||
|
|
||
|
// The return operation takes an optional input operand to return. This
|
||
|
// value must match the return type of the enclosing function.
|
||
|
let arguments = (ins Variadic<F64Tensor>:$input);
|
||
|
|
||
|
// The return operation only emits the input in the format if it is present.
|
||
|
let assemblyFormat = "($input^ `:` type($input))? attr-dict ";
|
||
|
|
||
|
// Allow building a ReturnOp with no return operand.
|
||
|
let builders = [OpBuilder<
|
||
|
"OpBuilder &b, OperationState &state", [{ build(b, state, llvm::None); }]
|
||
|
>];
|
||
|
|
||
|
// Provide extra utility definitions on the c++ operation class definition.
|
||
|
let extraClassDeclaration = [{
|
||
|
bool hasOperand() { return getNumOperands() != 0; }
|
||
|
}];
|
||
|
|
||
|
// Invoke a static verify method to verify this return operation.
|
||
|
let verifier = [{ return ::verify(*this); }];
|
||
|
}
|
||
|
|
||
|
def TransposeOp : Toy_Op<"transpose"> {
|
||
|
let summary = "transpose operation";
|
||
|
|
||
|
let arguments = (ins F64Tensor:$input);
|
||
|
let results = (outs F64Tensor);
|
||
|
|
||
|
let assemblyFormat = [{
|
||
|
`(` $input `:` type($input) `)` attr-dict `to` type(results)
|
||
|
}];
|
||
|
|
||
|
// Allow building a TransposeOp with from the input operand.
|
||
|
let builders = [
|
||
|
OpBuilder<"OpBuilder &b, OperationState &state, Value input">
|
||
|
];
|
||
|
|
||
|
// Invoke a static verify method to verify this transpose operation.
|
||
|
let verifier = [{ return ::verify(*this); }];
|
||
|
}
|
||
|
|
||
|
#endif // TOY_OPS
|