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

Lower constant padding operation to KRNL dialect (#27)

This commit is contained in:
chentong319 2020-03-11 16:54:07 -04:00 committed by GitHub
parent e8a0b47e10
commit 391f565a66
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 138 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
src/CMakeLists.txt
View File

@ -88,6 +88,7 @@ add_library(onnf_lower_frontend
conversion/onnx_to_krnl/nn/pooling.cpp
conversion/onnx_to_krnl/tensor/identity.cpp
conversion/onnx_to_krnl/tensor/reshape.cpp
conversion/onnx_to_krnl/tensor/padconstantvaluepad.cpp
conversion/onnx_to_krnl/tensor/transpose.cpp
conversion/onnx_to_krnl/tensor/unsqueeze.cpp
conversion/onnx_to_krnl/convert_onnx_to_krnl.cpp)

1
src/conversion/onnx_to_krnl/convert_onnx_to_krnl.cpp
View File

@ -93,6 +93,7 @@ void FrontendToKrnlLoweringPass::runOnModule() {
populateLoweringONNXMatMulOpPattern(patterns, &getContext());
// Tensor
populateLoweringONNXReshapeOpPattern(patterns, &getContext());
populateLoweringONNXPadConstantValuePadOpPattern(patterns, &getContext());
populateLoweringONNXUnsqueezeOpPattern(patterns, &getContext());
populateLoweringONNXTransposeOpPattern(patterns, &getContext());
populateLoweringONNXIdentityOpPattern(patterns, &getContext());

3
src/conversion/onnx_to_krnl/onnx_to_krnl_common.hpp
View File

@ -234,6 +234,9 @@ void populateLoweringONNXUnsqueezeOpPattern(
void populateLoweringONNXTransposeOpPattern(
OwningRewritePatternList &patterns, MLIRContext *ctx);
void populateLoweringONNXPadConstantValuePadOpPattern(
OwningRewritePatternList &patterns, MLIRContext *ctx);
void populateLoweringONNXReshapeOpPattern(
OwningRewritePatternList &patterns, MLIRContext *ctx);

108
src/conversion/onnx_to_krnl/tensor/padconstantvaluepad.cpp Normal file
View File

@ -0,0 +1,108 @@
//===----padconstantvaluepad.cpp - Lowering PadConstantValuePad Op --------===//
//
// Copyright 2019 The IBM Research Authors.
//
// =============================================================================
//
// This file lowers the ONNX PadConstantValuePad Operator to Krnl dialect.
//
//===----------------------------------------------------------------------===//
#include "src/conversion/onnx_to_krnl/onnx_to_krnl_common.hpp"
using namespace mlir;
struct ONNXPadConstantValuePadOpLowering : public ConversionPattern {
ONNXPadConstantValuePadOpLowering(MLIRContext *ctx)
: ConversionPattern(mlir::ONNXPadConstantValuePadOp::getOperationName(),
1, ctx) {}
PatternMatchResult
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
ConversionPatternRewriter &rewriter) const final {
auto tensorType = (*op->result_type_begin());
auto loc = op->getLoc();
// Only constant padding is supported now.
auto padMode = llvm::dyn_cast<ONNXPadConstantValuePadOp>(op).mode();
if (padMode != "constant")
emitError(loc, "unsupported mode for PadConstantValuePad");
auto constantValAttr =
llvm::dyn_cast<ONNXPadConstantValuePadOp>(op).constant_valueAttr();
// Insert an allocation and deallocation for the result of this operation.
auto memRefType = convertToMemRefType(tensorType);
Value alloc;
bool insertDealloc = checkInsertDealloc(op);
if (hasAllConstantDimensions(memRefType))
alloc = insertAllocAndDealloc(memRefType, loc, rewriter, insertDealloc);
else
emitError(loc, "unexpected output has non-Constant shape");
// Number of loops
auto memRefShape = memRefType.getShape();
int64_t rank = memRefShape.size();
// Iterate over the loop nest using the output shape.
BuildKrnlLoop padLoops(rewriter, loc, rank);
padLoops.createDefineAndOptimizeOp();
for (int i = 0; i < rank; ++i)
padLoops.pushBounds(0, alloc, i);
padLoops.createIterateOp();
// Iterate over the loop nest using the input shape.
BuildKrnlLoop valueLoops(rewriter, loc, rank);
valueLoops.createDefineAndOptimizeOp();
for (int i = 0; i < rank; ++i)
valueLoops.pushBounds(0, operands[0], i);
valueLoops.createIterateOp();
// Copy the input data into the output.
rewriter.setInsertionPointToStart(valueLoops.getIterateBlock());
SmallVector<Value, 4> inLoopIVs;
for (int i = 0; i < rank; ++i)
inLoopIVs.emplace_back(valueLoops.getInductionVar(i));
auto pads = llvm::dyn_cast<ONNXPadConstantValuePadOp>(op).pads();
SmallVector<int64_t, 4> pad_begin;
for (int i = 0; i < pads.size()/2; ++i) {
pad_begin.emplace_back(pads.getValue()[i].cast<IntegerAttr>().getInt());
}
SmallVector<Value, 4> outLoopIVs;
for (int i = 0; i < rank; ++i) {
// Calculate the index for the load and store.
if (pad_begin[i] == 0) {
outLoopIVs.emplace_back(valueLoops.getInductionVar(i));
} else {
auto outIV = rewriter.create<AddIOp>(
loc, rewriter.create<ConstantIndexOp>(loc, pad_begin[i]),
valueLoops.getInductionVar(i));
outLoopIVs.emplace_back(outIV);
}
}
auto inVal = rewriter.create<LoadOp>(loc, operands[0], inLoopIVs);
rewriter.create<StoreOp>(loc, inVal, alloc, outLoopIVs);
rewriter.setInsertionPointToStart(padLoops.getIterateBlock());
SmallVector<Value, 4> outLoopIVs1;
for (int i = 0; i < rank; ++i)
outLoopIVs1.emplace_back(padLoops.getInductionVar(i));
auto inVal1 = rewriter.create<ConstantOp>(loc, constantValAttr);
rewriter.create<StoreOp>(loc, inVal1, alloc, outLoopIVs1);
// Replace the original op with the generated code.
rewriter.replaceOp(op, alloc);
return matchSuccess();
}
};
void populateLoweringONNXPadConstantValuePadOpPattern(
OwningRewritePatternList &patterns, MLIRContext *ctx) {
patterns.insert<ONNXPadConstantValuePadOpLowering>(ctx);
}

25
test/mlir/onnx/onnx_lowering.mlir
View File

@ -1510,3 +1510,28 @@ func @test_maxpooling_singleout_no_pad_w_strides_w_ceil_mode_w_unknown_dims(%arg
// CHECK: }
// CHECK: return [[RES]] : memref<?x3x?x16xf32>
}
func @test_constant_pad1(%arg0: tensor<16x16xf32>) -> tensor<18x20xf32> {
%0 = "onnx.PadConstantValuePad"(%arg0) {constant_value = 0.000000e+00 : f32, mode = "constant", pads = [0, 3, 2, 1]} : (tensor<16x16xf32>) -> tensor<18x20xf32>
return %0 : tensor<18x20xf32>
// CHECK-LABEL: test_constant_pad1
// CHECK: [[RES:%.+]] = alloc() : memref<18x20xf32>
// CHECK: [[DEF_LOOPS1:%.+]]:2 = krnl.define_loops 2
// CHECK: [[OPT_LOOPS1:%.+]]:2 = krnl.optimize_loops {
// CHECK: krnl.return_loops [[DEF_LOOPS1]]#0, [[DEF_LOOPS1]]#1
// CHECK: } : () -> (!krnl.loop, !krnl.loop)
// CHECK: krnl.iterate([[OPT_LOOPS1]]#0, [[OPT_LOOPS1]]#1) with ([[DEF_LOOPS1]]#0 -> %arg1 = 0 to 18, [[DEF_LOOPS1]]#1 -> %arg2 = 0 to 20) {
// CHECK: [[CST:%.+]] = constant 0.000000e+00 : f32
// CHECK: store [[CST]], [[RES]][%arg1, %arg2] : memref<18x20xf32>
// CHECK: }
// CHECK: [[DEF_LOOPS2:%.+]]:2 = krnl.define_loops 2
// CHECK: [[OPT_LOOPS2:%.+]]:2 = krnl.optimize_loops {
// CHECK: krnl.return_loops [[DEF_LOOPS2]]#0, [[DEF_LOOPS2]]#1
// CHECK: } : () -> (!krnl.loop, !krnl.loop)
// CHECK: krnl.iterate([[OPT_LOOPS2]]#0, [[OPT_LOOPS2]]#1) with ([[DEF_LOOPS2]]#0 -> %arg1 = 0 to 16, [[DEF_LOOPS2]]#1 -> %arg2 = 0 to 16) {
// CHECK: [[CST1:%.+]] = constant 3 : index
// CHECK: [[ADD:%.+]] = addi [[CST1]], %arg2 : index
// CHECK: [[LOAD:%.+]] = load %arg0[%arg1, %arg2] : memref<16x16xf32>
// CHECK: store [[LOAD]], [[RES]][%arg1, [[ADD]]] : memref<18x20xf32>
// CHECK: }
}