mirror of https://github.com/llvm/torch-mlir
120 lines
3.8 KiB
TableGen
120 lines
3.8 KiB
TableGen
//===-------------------------------------------------------*- 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
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#ifndef TCF_OPS
|
|
#define TCF_OPS
|
|
|
|
include "npcomp/Dialect/TCF/IR/TCFBase.td"
|
|
|
|
class TCF_Op<string mnemonic, list<OpTrait> traits = []>
|
|
: Op<TCF_Dialect, mnemonic, traits> {
|
|
}
|
|
|
|
// TODO: investigate effects framework for defining error semantics
|
|
// TODO: define in a general way across the dialect what "encounters an error" means.
|
|
|
|
class BinaryArithmeticOp<string mnemonic, list<OpTrait> traits = []> :
|
|
TCF_Op<mnemonic, traits> {
|
|
let arguments = (ins AnyTensor:$lhs, AnyTensor:$rhs);
|
|
let results = (outs AnyTensor:$result);
|
|
let assemblyFormat = "$lhs `,` $rhs attr-dict `:` functional-type(operands, results)";
|
|
}
|
|
|
|
def TCF_AddOp : BinaryArithmeticOp<"add"> {
|
|
let summary = "Addition of two tensors.";
|
|
let description = [{
|
|
Addition of two tensors.
|
|
|
|
Numpy-style broadcasting is allowed.
|
|
}];
|
|
}
|
|
|
|
def TCF_MaxOp : BinaryArithmeticOp<"max"> {
|
|
let summary = "Maximum of two tensors.";
|
|
let description = [{
|
|
Maximum of two tensors.
|
|
|
|
Numpy-style broadcasting is allowed.
|
|
}];
|
|
}
|
|
|
|
def TCF_MulOp : BinaryArithmeticOp<"mul"> {
|
|
let summary = "Multiply an input tensor by a scalar tensor.";
|
|
let description = [{
|
|
Multiplies each element of the input `input` with the scalar `other` and returns a new resulting tensor. The tensor types must match and shapes must be broadcastable.
|
|
}];
|
|
}
|
|
|
|
class UnaryArithmeticOp<string mnemonic, list<OpTrait> traits = []> :
|
|
TCF_Op<mnemonic,
|
|
!listconcat(traits, [AllTypesMatch<["operand", "result"]>])>,
|
|
AllTypesMatch<["operand", "result"]> {
|
|
let arguments = (ins AnyTensor:$operand);
|
|
let results = (outs AnyTensor:$result);
|
|
let assemblyFormat = "$operand attr-dict `:` type($operand)";
|
|
}
|
|
|
|
def TCF_ExpOp : UnaryArithmeticOp<"exp"> {
|
|
let summary = "base-e exponential";
|
|
let description = [{
|
|
See std.exp for more details.
|
|
}];
|
|
}
|
|
|
|
def TCF_TanhOp : UnaryArithmeticOp<"tanh"> {
|
|
let summary = "hyperbolic tangent";
|
|
let description = [{
|
|
See std.tanh for more details.
|
|
}];
|
|
}
|
|
|
|
// TODO: Generalize this op appropriately and add more verification.
|
|
// For example, an unranked operand probably should be allowed and verified
|
|
// dynamically in TCF->TCP lowering if needed.
|
|
def TCF_MatmulOp : TCF_Op<"matmul"> {
|
|
let summary = "Performs a matrix multiplication";
|
|
let description = [{
|
|
Performs a matrix multiplication.
|
|
|
|
The tensors have dimensions:
|
|
- lhs: [M, K]
|
|
- rhs: [K, N]
|
|
- result: [M, N]
|
|
|
|
If the `K` dimension mismatches between the operands, this op aborts the
|
|
program.
|
|
}];
|
|
let arguments = (ins 2DTensorOf<[F32]>:$lhs, 2DTensorOf<[F32]>:$rhs);
|
|
let results = (outs 2DTensorOf<[F32]>:$result);
|
|
|
|
let assemblyFormat = "$lhs `,` $rhs attr-dict `:` functional-type(operands, results)";
|
|
}
|
|
|
|
def TCF_ConvNCHWOp : TCF_Op<"conv_2d_nchw"> {
|
|
let summary = "2-D convolution";
|
|
let description = [{
|
|
Performs 2-D convolution. This op is inspired by PyTorch's Conv2d layer (https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html).
|
|
|
|
The tensors have dimensions:
|
|
- in: [N, Cin, H, W]
|
|
- filter: [Cout, Cin, KH, KW]
|
|
- result: [N, Cout, Hout, Wout]
|
|
|
|
The tensors must meet the following conditions; otherwise, this op aborts the program.
|
|
- H is greater than or equal to KH
|
|
- W is greater than or equal to KW
|
|
- Cin matches between in and filter
|
|
}];
|
|
let arguments = (ins 4DTensorOf<[F32]>:$in, 4DTensorOf<[F32]>:$filter);
|
|
let results = (outs 4DTensorOf<[F32]>:$result);
|
|
|
|
let assemblyFormat = "$in `,` $filter attr-dict `:` functional-type(operands, results)";
|
|
}
|
|
|
|
#endif // #ifndef TCF_OPS
|