Rewrite shape and size OP (#285)

* add shape inference

* Revert "add shape inference"

This reverts commit f9d42f39e68e14b5648abccfc8617fff00244d16.

* add rewrite rules

* test cases

* format

* add constraint

* response to review

* response to review
This commit is contained in:
chentong319 2020-09-10 13:46:00 -05:00 committed by GitHub
parent 5e11429d77
commit ac67900baf
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 113 additions and 1 deletions

View File

@ -4729,6 +4729,7 @@ def ONNXSequenceLengthOp:ONNX_Op<"SequenceLength",
def ONNXShapeOp:ONNX_Op<"Shape", def ONNXShapeOp:ONNX_Op<"Shape",
[NoSideEffect, DeclareOpInterfaceMethods<ShapeInferenceOpInterface>]> { [NoSideEffect, DeclareOpInterfaceMethods<ShapeInferenceOpInterface>]> {
let hasCanonicalizer = 1;
let summary = "ONNX Shape operation"; let summary = "ONNX Shape operation";
let description = [{ let description = [{
"Takes a tensor as input and outputs an 1D int64 tensor containing the shape of the input tensor." "Takes a tensor as input and outputs an 1D int64 tensor containing the shape of the input tensor."
@ -4863,6 +4864,7 @@ def ONNXSinhOp:ONNX_Op<"Sinh",
def ONNXSizeOp:ONNX_Op<"Size", def ONNXSizeOp:ONNX_Op<"Size",
[NoSideEffect]> { [NoSideEffect]> {
let hasCanonicalizer = 1;
let summary = "ONNX Size operation"; let summary = "ONNX Size operation";
let description = [{ let description = [{
"Takes a tensor as input and outputs a int64 scalar that equals to the total number of elements of the input tensor." "Takes a tensor as input and outputs a int64 scalar that equals to the total number of elements of the input tensor."

View File

@ -27,6 +27,29 @@ DenseElementsAttr createDenseElementsAttrFromFloatAttr(
return mlir::DenseElementsAttr::get(tensorType, llvm::makeArrayRef(values)); return mlir::DenseElementsAttr::get(tensorType, llvm::makeArrayRef(values));
} }
// Create a DenseElementsAttr based on the shape of type.
DenseElementsAttr createDenseElementsAttrFromShape(
PatternRewriter &rewriter, Value value) {
auto inType = value.getType().cast<ShapedType>();
auto shape = inType.getShape();
SmallVector<int64_t, 1> dims = {inType.getRank()};
SmallVector<int64_t, 4> values(shape.begin(), shape.end());
auto tensorType =
mlir::RankedTensorType::get(dims, rewriter.getIntegerType(64));
return mlir::DenseElementsAttr::get(tensorType, llvm::makeArrayRef(values));
}
// Create a DenseElementsAttr based on the size of type.
DenseElementsAttr createDenseElementsAttrFromSize(
PatternRewriter &rewriter, Value value) {
auto inType = value.getType().cast<ShapedType>();
SmallVector<int64_t, 1> dims(1, 1);
SmallVector<int64_t, 1> values = {inType.getNumElements()};
auto tensorType =
mlir::RankedTensorType::get(dims, rewriter.getIntegerType(64));
return mlir::DenseElementsAttr::get(tensorType, llvm::makeArrayRef(values));
}
// If 'lhs' is not NoneType, return 'lhs - rhs'. // If 'lhs' is not NoneType, return 'lhs - rhs'.
// Otherwise, return '-rhs'. // Otherwise, return '-rhs'.
Value subtractOrNeg( Value subtractOrNeg(
@ -128,3 +151,15 @@ void ONNXBatchNormalizationTestModeOp::getCanonicalizationPatterns(
OwningRewritePatternList &results, MLIRContext *context) { OwningRewritePatternList &results, MLIRContext *context) {
results.insert<FuseBatchNormTestModeConvPattern>(context); results.insert<FuseBatchNormTestModeConvPattern>(context);
} }
/// on the ONNXShapeOp.
void ONNXShapeOp::getCanonicalizationPatterns(
OwningRewritePatternList &results, MLIRContext *context) {
results.insert<ShapeToConstantPattern>(context);
}
/// on the ONNXSizeOp.
void ONNXSizeOp::getCanonicalizationPatterns(
OwningRewritePatternList &results, MLIRContext *context) {
results.insert<SizeToConstantPattern>(context);
}

View File

@ -28,6 +28,14 @@ include "src/Dialect/ONNX/ONNXOps.td"
def createDenseElementsAttrFromFloatAttr : NativeCodeCall< def createDenseElementsAttrFromFloatAttr : NativeCodeCall<
"createDenseElementsAttrFromFloatAttr($_builder, $0.getType().cast<ShapedType>().getElementType(), $1)">; "createDenseElementsAttrFromFloatAttr($_builder, $0.getType().cast<ShapedType>().getElementType(), $1)">;
// Create a DenseElementsAttr from the shape of the type of a value.
def createDenseElementsAttrFromShape : NativeCodeCall<
"createDenseElementsAttrFromShape($_builder, $0)">;
// Create a DenseElementsAttr from the size of the type of a value.
def createDenseElementsAttrFromSize : NativeCodeCall<
"createDenseElementsAttrFromSize($_builder, $0)">;
// If '$1' is not NoneType, do subtraction '$1 - $2'. // If '$1' is not NoneType, do subtraction '$1 - $2'.
// Otherwise, take the negative of '$2'. // Otherwise, take the negative of '$2'.
def subtractOrNeg: NativeCodeCall< def subtractOrNeg: NativeCodeCall<
@ -172,4 +180,25 @@ def FuseBatchNormTestModeConvPattern: Pat<
$auto_pad, $dilation, $group, $kernel_shape, $pads, $strides) $auto_pad, $dilation, $group, $kernel_shape, $pads, $strides)
>; >;
def IsStaticShapeTensor:
Constraint<
CPred<
"$_self.getType().cast<::mlir::ShapedType>().hasStaticShape()">,
"hasStaticShape">;
def ShapeToConstantPattern: Pat<
(ONNXShapeOp $A),
(ONNXConstantOp
(GetNullAttr),
(createDenseElementsAttrFromShape $A)),
[(IsStaticShapeTensor:$A)]
>;
def SizeToConstantPattern: Pat<
(ONNXSizeOp $A),
(ONNXConstantOp
(GetNullAttr),
(createDenseElementsAttrFromSize $A)),
[(IsStaticShapeTensor:$A)]
>;
#endif // ONNX_REWRITE #endif // ONNX_REWRITE

View File

@ -222,3 +222,49 @@ func @test_transpose_fusion_removal(%arg0: tensor<10x11x12x13xf32>) -> tensor<10
// CHECK-NEXT: return %arg0 : tensor<10x11x12x13xf32> // CHECK-NEXT: return %arg0 : tensor<10x11x12x13xf32>
"std.return"(%1) : (tensor<10x11x12x13xf32>) -> () "std.return"(%1) : (tensor<10x11x12x13xf32>) -> ()
} }
// -----
func @test_shape1(%arg0 : tensor<2x4x8x16xf32>) -> tensor<*xi64> {
%0 = "onnx.Shape"(%arg0) : (tensor<2x4x8x16xf32>) -> tensor<*xi64>
return %0 : tensor<*xi64>
// CHECK-LABEL: @test_shape1
// CHECK-NEXT: %0 = "onnx.Constant"() {value = dense<[2, 4, 8, 16]> : tensor<4xi64>} : () -> tensor<*xi64>
// CHECK-NEXT: %0 : tensor<*xi64>
}
// -----
func @test_shape2(%arg0 : tensor<?x4x8x16xf32>) -> tensor<*xi64> {
%0 = "onnx.Shape"(%arg0) : (tensor<?x4x8x16xf32>) -> tensor<*xi64>
return %0 : tensor<*xi64>
// CHECK-LABEL: @test_shape2
// CHECK-NEXT: %0 = "onnx.Shape"(%arg0) : (tensor<?x4x8x16xf32>) -> tensor<*xi64>
// CHECK-NEXT: return %0 : tensor<*xi64>
}
// -----
func @test_size1(%arg0 : tensor<2x4x8x16xf32>) -> tensor<*xi64> {
%0 = "onnx.Size"(%arg0) : (tensor<2x4x8x16xf32>) -> tensor<*xi64>
return %0 : tensor<*xi64>
// CHECK-LABEL: @test_size1
// CHECK-NEXT: %0 = "onnx.Constant"() {value = dense<1024> : tensor<1xi64>} : () -> tensor<*xi64>
// CHECK-NEXT: %0 : tensor<*xi64>
}
// -----
func @test_size2(%arg0 : tensor<*xf32>) -> tensor<*xi64> {
%0 = "onnx.Size"(%arg0) : (tensor<*xf32>) -> tensor<*xi64>
return %0 : tensor<*xi64>
// CHECK-LABEL: @test_size2
// CHECK-NEXT: %0 = "onnx.Size"(%arg0) : (tensor<*xf32>) -> tensor<*xi64>
// CHECK-NEXT: return %0 : tensor<*xi64>
}

View File

@ -321,7 +321,7 @@ OpsWithShapeInference=[
] ]
# Operations supporting canonicalization. # Operations supporting canonicalization.
OpsWithCanonicalizer = ['Add', 'Identity', 'Gemm', 'Conv', 'Cast', 'Transpose', 'Dropout'] OpsWithCanonicalizer = ['Add', 'Identity', 'Gemm', 'Conv', 'Cast', 'Transpose', 'Dropout', 'Shape', 'Size']
# Operations who have operands that, if produced by constant operations, should # Operations who have operands that, if produced by constant operations, should
# be promoted to become an attribute (via attribute promotion). # be promoted to become an attribute (via attribute promotion).