colin
/
onnx-mlir
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'upstream/master' into shapeinference-pad

This commit is contained in:
chentong 2020-02-17 08:27:43 -05:00
parent bbdf4e3b4d 3c505ae31d
commit ec43fadc3b
4 changed files with 116 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
src/dialect/onnx/onnx.td
View File

@ -111,6 +111,7 @@ def ONNXGemmNoBiasOp: ONNX_Op<"GemmNoBias",
def ONNXConvNoBiasOp:ONNX_Op<"ConvNoBias",
[NoSideEffect, DeclareOpInterfaceMethods<ShapeInferenceOpInterface>]> {
let hasCanonicalizer = 1;
let summary = "ONNX Conv operation with no Bias operand.";
let description = [{
"The convolution operator consumes an input tensor and a filter, and"

103
src/pass/onnx_rewrite.cpp
View File

@ -263,6 +263,103 @@ struct ReduceSumSquareOpPattern : public RewritePattern {
return matchSuccess();
};
};
//===----------------------------------------------------------------------===//
// Rewrite:
// %0 = onnx.ConvNoBiasOp(%D : tensor<DShape>, %K)
// {pads = [b0, b1, ... bK, e0, e1, ..., eK]} ->
// tensor<OutShape>
//
// as:
// %0 = onnx.PadConstantValuePasOp(%D)
// {pads = [0, 0, b0, b1, ... bK, 0, 0, e0, e1, ..., eK]} ->
// tensor<DPaddedShape>
// %1 = onnx.ConvNoBias(%0 : tensor<DPaddedShape>, %K) {pads = [0, ..., 0]} ->
// tensor<OutShape>
//===----------------------------------------------------------------------===//
struct SplitConvOpPattern : public RewritePattern {
SplitConvOpPattern(MLIRContext *context)
: RewritePattern(ONNXConvNoBiasOp::getOperationName(),
{ONNXPadConstantValuePadOp::getOperationName(),
ONNXConvNoBiasOp::getOperationName()},
1, context) {}
PatternMatchResult matchAndRewrite(Operation *op,
PatternRewriter &rewriter) const override {
auto loc = op->getLoc();
// If convolution does not use padding then no rewrite is required.
ONNXConvNoBiasOp convOp = llvm::dyn_cast<ONNXConvNoBiasOp>(op);
auto padsAttribute = convOp.padsAttr();
if (!padsAttribute)
return matchFailure();
// If auto_pad is VALID then no padding happens and no rewrite isrequired.
auto autoPad = convOp.auto_pad();
if (autoPad == "VALID")
return matchFailure();
auto data = op->getOperands()[0];
auto inputShape = data.getType().cast<TensorType>().getShape();
// Dimensionality of the input:
// inputRank
// |----------------------|
// D : (N x C x D1 x D2 x ... DK)
// |______________|
// inputDims
//
int64_t inputRank = inputShape.size();
int64_t inputDims = inputRank - 2;
// If all pads values are equal to zero then no rewrite is required.
bool allZeros = true;
for (auto padsValue : padsAttribute.getValue()) {
if (padsValue.cast<IntegerAttr>().getInt() > 0) {
allZeros = false;
break;
}
}
if (allZeros)
return matchFailure();
// Create padding vector for the explicit padding op attribute.
SmallVector<int64_t, 4> pads(2 * inputRank, 0);
SmallVector<int64_t, 4> outPaddedShape(inputRank, 0);
outPaddedShape[0] = inputShape[0];
outPaddedShape[1] = inputShape[1];
for (int i = 0; i < inputDims; ++i) {
int64_t beginPad =
padsAttribute.getValue()[i].cast<IntegerAttr>().getInt();
int64_t endPad =
padsAttribute.getValue()[inputDims + i].cast<IntegerAttr>().getInt();
pads[i + 2] = beginPad;
pads[inputRank + i + 2] = endPad;
outPaddedShape[i + 2] += beginPad + inputShape[i + 2] + endPad;
}
// Create padding operation.
auto inputElemType = data.getType().cast<TensorType>().getElementType();
ONNXPadConstantValuePadOp paddingOp =
rewriter.create<ONNXPadConstantValuePadOp>(
loc, RankedTensorType::get(outPaddedShape, inputElemType), data,
rewriter.getI64ArrayAttr(pads), FloatAttr::get(inputElemType, 0),
StringAttr::get("constant", loc->getContext()));
SmallVector<int64_t, 4> newConvPads(2 * inputDims, 0);
auto tensorType = (*op->result_type_begin()).cast<TensorType>();
ONNXConvNoBiasOp newConvOp = rewriter.create<ONNXConvNoBiasOp>(
loc, tensorType, paddingOp.getResult(), convOp.getOperands()[1],
convOp.auto_padAttr(), convOp.dilationsAttr(),
convOp.groupAttr(), convOp.kernel_shapeAttr(),
rewriter.getI64ArrayAttr(newConvPads),
convOp.stridesAttr());
rewriter.replaceOp(op, newConvOp.getResult());
return matchSuccess();
};
};
} // end anonymous namespace
/// on the ONNXReduceL1Op.
@ -293,3 +390,9 @@ void ONNXReduceSumSquareOp::getCanonicalizationPatterns(
OwningRewritePatternList &results, MLIRContext *context) {
results.insert<ReduceSumSquareOpPattern>(context);
}
/// on the ONNXReduceSumSquareOp.
void ONNXConvNoBiasOp::getCanonicalizationPatterns(
OwningRewritePatternList &results, MLIRContext *context) {
results.insert<SplitConvOpPattern>(context);
}

9
test/mlir/onnx/onnx_canonicalization.mlir
View File

@ -92,3 +92,12 @@ func @test_constant_pad(%arg0 : tensor<?x?xf32>) -> tensor<*xf32> {
%2 = "onnx.PadConstantValue"(%arg0, %0) {constant_value=0. : f32, mode = "constant"} : (tensor<?x?xf32>, tensor<?xi64>)-> tensor<*xf32>
"std.return"(%2) : (tensor<*xf32>) -> ()
}
// CHECK-LABEL: @test_conv_split(%{{.*}}: tensor<1x9x32x64xf32>, %{{.*}}: tensor<5x9x6x7xf32>) -> tensor<*xf32> {
func @test_conv_split(%arg0 : tensor<1x9x32x64xf32>, %arg1 : tensor<5x9x6x7xf32>) -> tensor<*xf32> {
%0 = "onnx.ConvNoBias"(%arg0, %arg1) {auto_pad = "NOTSET", group = 1 : i64, pads = [2, 3, 4, 5]} : (tensor<1x9x32x64xf32>, tensor<5x9x6x7xf32>) -> tensor<*xf32>
"std.return"(%0) : (tensor<*xf32>) -> ()
// CHECK-NEXT: %0 = "onnx.PadConstatValuePad"(%arg0) {constant_value = 0.000000e+00 : f32, mode = "constant", pads = [0, 0, 2, 3, 0, 0, 4, 5]} : (tensor<1x9x32x64xf32>) -> tensor<1x9x38x72xf32>
// CHECK-NEXT: %1 = "onnx.ConvNoBias"(%0, %arg1) {auto_pad = "NOTSET", group = 1 : i64, pads = [0, 0, 0, 0]} : (tensor<1x9x38x72xf32>, tensor<5x9x6x7xf32>) -> tensor<*xf32>
// CHECK-NEXT: return %1 : tensor<*xf32>
}