mirror of https://github.com/llvm/torch-mlir
135 lines
4.7 KiB
TableGen
135 lines
4.7 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 TCP_OPS
|
|
#define TCP_OPS
|
|
|
|
include "npcomp/Dialect/TCP/IR/TCPBase.td"
|
|
include "mlir/Dialect/Shape/IR/ShapeBase.td"
|
|
include "mlir/Interfaces/SideEffectInterfaces.td"
|
|
include "mlir/Interfaces/InferTypeOpInterface.td"
|
|
|
|
class TCP_Op<string mnemonic, list<OpTrait> traits = []>
|
|
: Op<TCP_Dialect, mnemonic, traits> {
|
|
}
|
|
|
|
// TODO: clarify allowed tensor element types.
|
|
// TODO: HasParent is too restrictive? can't have an island with loop.for with
|
|
// further ops inside it?
|
|
def TCP_AddOp
|
|
: TCP_Op<"add", []> {
|
|
let summary = "Adds two tensors.";
|
|
let description = [{
|
|
Adds two tensors.
|
|
}];
|
|
let arguments = (ins AnyRankedTensor:$lhs, AnyRankedTensor:$rhs);
|
|
let results = (outs AnyRankedTensor:$result);
|
|
}
|
|
|
|
def TCP_BroadcastToOp : TCP_Op<"broadcast_to"> {
|
|
let summary = "Broadcasts an operand to a given shape.";
|
|
let description = [{
|
|
Broadcasts `operand` to the shape `shape`.
|
|
|
|
It is undefined behavior if such a broadcast is not legal.
|
|
}];
|
|
let arguments = (ins AnyRankedTensor:$operand, Shape_ShapeType:$shape);
|
|
let results = (outs AnyRankedTensor:$result);
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Ops that need to be factored to a proper home.
|
|
//===----------------------------------------------------------------------===//
|
|
// TODO: Find a home for these.
|
|
|
|
// TODO: This probably doesn't belong in the tcp dialect.
|
|
def TCP_AllocMemRefOp : TCP_Op<"alloc_memref", []> {
|
|
let summary = "Allocates a memref of the given shape.";
|
|
let description = [{
|
|
Allocates a memref of the given shape.
|
|
}];
|
|
let arguments = (ins Shape_ShapeType:$shape);
|
|
let results = (outs AnyMemRef:$memref);
|
|
let assemblyFormat = "$shape attr-dict `:` type($memref)";
|
|
}
|
|
|
|
// TODO: Change to a more principled error handling mechanism.
|
|
// This op probably doesn't need to exist eventually.
|
|
// This op is also not correctly modeled right now, since it itself doesn't
|
|
// produce the error in practice. The ops like shape.broadcast itself, when
|
|
// lowered, immediately produce errors.
|
|
// TODO: This should eventually be moved to a shape dialect.
|
|
def TCP_ShapeObserveErrorOp : TCP_Op<"shape_observe_error",
|
|
[DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
|
|
let summary = "Observes the fact that a shape might be an error.";
|
|
let description = [{
|
|
This op is a structural placeholder that captures a shape such that it
|
|
is not erased. This will keep around shape computations that are later
|
|
lowered into eager error handling code.
|
|
|
|
The interaction of this op, especially with control flow and side
|
|
effecting ops, is not very well-defined, and needs to be worked
|
|
on/redesigned.
|
|
}];
|
|
let arguments = (ins Shape_ShapeType:$shape);
|
|
// TODO: ODS seems to create redeclared class members if we remove this,
|
|
// resulting in C++ compilation errors.
|
|
let results = (outs NoneType:$dummy);
|
|
}
|
|
|
|
// TODO: This probably belongs in the shape dialect.
|
|
def TCP_GetExtentOp : TCP_Op<"get_extent",
|
|
[NoSideEffect, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
|
|
let summary = "Gets the specified extent from a shape.";
|
|
let description = [{
|
|
Gets the specified extent from a shape.
|
|
|
|
This op has undefined behavior if the shape is an error.
|
|
}];
|
|
let arguments = (ins Shape_ShapeType:$shape, I64Attr:$dim);
|
|
let results = (outs Index:$extent);
|
|
let assemblyFormat = "$shape `,` $dim attr-dict";
|
|
|
|
let builders = [
|
|
// Helper to pass a simple integer instead of an integer attr.
|
|
OpBuilder<
|
|
[{
|
|
OpBuilder &builder, OperationState &result,
|
|
Value shape, int64_t dim
|
|
}],
|
|
[{
|
|
build(builder, result, shape, builder.getI64IntegerAttr(dim));
|
|
}]
|
|
>
|
|
];
|
|
}
|
|
|
|
// TODO: This op belongs in the shape dialect as `shape.from_extents`.
|
|
def TCP_ShapeFromExtentsOp : TCP_Op<"shape_from_extents",
|
|
[NoSideEffect, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
|
|
let summary = "Constructs a shape from extents";
|
|
let description = [{
|
|
Constructs a shape from the extents passed as arguments.
|
|
}];
|
|
let arguments = (ins Variadic<Index>:$extents);
|
|
let results = (outs Shape_ShapeType:$shape);
|
|
}
|
|
|
|
def TCP_AbortIfOp : TCP_Op<"abort_if"> {
|
|
let summary = "Aborts the program if the argument is true.";
|
|
let description = [{
|
|
Aborts the program if the argument is true.
|
|
|
|
TODO: Support a custom message.
|
|
}];
|
|
let arguments = (ins I1:$pred);
|
|
let results = (outs);
|
|
}
|
|
|
|
#endif // TCP_OPS
|