llvm-project/mlir/lib/Dialect/Shape/IR/ShapeCanonicalization.td
Adrian Kuegel 2bb208ddfd [mlir] Don't allow dynamic extent tensor types for ConstShapeOp.
ConstShapeOp has a constant shape, so its type can always be static.
We still allow it to have ShapeType though.

Differential Revision: https://reviews.llvm.org/D111139
2021-10-07 10:56:16 +02:00

46 lines
1.4 KiB
TableGen

include "mlir/Dialect/Shape/IR/ShapeOps.td"
include "mlir/Dialect/StandardOps/IR/Ops.td"
include "mlir/Dialect/Tensor/IR/TensorOps.td"
def AllInputShapesEq : Constraint<CPred< [{
llvm::all_of($0, [&](mlir::Value val) {
return $0[0] == val;
})
}]>>;
def HasSingleElement : Constraint<CPred< [{
$0.size() == 1
}]>>;
def HasStaticShape : Constraint<CPred< [{
$0.getType().dyn_cast<ShapedType>().hasStaticShape()
}]>>;
// Canonicalization patterns.
def AssumingAllOneOp : Pat<(Shape_AssumingAllOp $args),
(replaceWithValue $args),
[(HasSingleElement $args)]>;
def CstrBroadcastableEqOps : Pat<(Shape_CstrBroadcastableOp:$op $shapes),
(Shape_ConstWitnessOp ConstBoolAttrTrue),
[(AllInputShapesEq $shapes)]>;
def CstrEqEqOps : Pat<(Shape_CstrEqOp:$op $shapes),
(Shape_ConstWitnessOp ConstBoolAttrTrue),
[(AllInputShapesEq $shapes)]>;
def IndexToSizeToIndexCanonicalization : Pat<
(Shape_SizeToIndexOp (Shape_IndexToSizeOp $arg)),
(replaceWithValue $arg)>;
def SizeToIndexToSizeCanonicalization : Pat<
(Shape_IndexToSizeOp (Shape_SizeToIndexOp $arg)),
(replaceWithValue $arg)>;
// Fold tensor.cast(const_shape) to const_shape. This changes the type of
// const_shape to the destination type of the cast.
def TensorCastConstShape : Pat <
(Tensor_CastOp:$res (Shape_ConstShapeOp $arg)), (Shape_ConstShapeOp $arg),
[(HasStaticShape $res)]>;