forked from OSchip/llvm-project
[mlir][tosa] Materialize tosa.pad value and fold noop pads
Padding now can explicitly specify the padding value when non-zero is wanted. This also includes bypassing pads when the pad does nothing. Differential Revision: https://reviews.llvm.org/D113611
This commit is contained in:
parent
54eec7cafc
commit
0f1e52afa9
|
@ -1417,6 +1417,9 @@ def Tosa_PadOp : Tosa_Op<"pad", [
|
||||||
|
|
||||||
let builders = [Tosa_PadOpQuantInfoBuilder,
|
let builders = [Tosa_PadOpQuantInfoBuilder,
|
||||||
Tosa_ExplicitValuePadOpQuantInfoBuilder];
|
Tosa_ExplicitValuePadOpQuantInfoBuilder];
|
||||||
|
|
||||||
|
let hasCanonicalizer = 1;
|
||||||
|
let hasFolder = 1;
|
||||||
}
|
}
|
||||||
|
|
||||||
//===----------------------------------------------------------------------===//
|
//===----------------------------------------------------------------------===//
|
||||||
|
|
|
@ -376,6 +376,53 @@ void MulOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
|
||||||
results.insert<MulOneOptimization>(context);
|
results.insert<MulOneOptimization>(context);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
struct MaterializePadValue : public OpRewritePattern<tosa::PadOp> {
|
||||||
|
using OpRewritePattern::OpRewritePattern;
|
||||||
|
|
||||||
|
LogicalResult matchAndRewrite(tosa::PadOp op,
|
||||||
|
PatternRewriter &rewriter) const override {
|
||||||
|
if (op.pad_const())
|
||||||
|
return failure();
|
||||||
|
|
||||||
|
auto input = op.input1();
|
||||||
|
auto padding = op.padding();
|
||||||
|
|
||||||
|
ShapedType inputTy = input.getType().cast<ShapedType>();
|
||||||
|
Type elementTy = inputTy.getElementType();
|
||||||
|
|
||||||
|
Attribute constantAttr;
|
||||||
|
if (elementTy.isa<FloatType>())
|
||||||
|
constantAttr = rewriter.getFloatAttr(elementTy, 0.0);
|
||||||
|
else if (elementTy.isa<IntegerType>() && !op.quantization_info())
|
||||||
|
constantAttr = rewriter.getIntegerAttr(elementTy, 0);
|
||||||
|
else if (elementTy.isa<IntegerType>() && op.quantization_info()) {
|
||||||
|
auto value = op.quantization_info().getValue().input_zp().getValue();
|
||||||
|
constantAttr = rewriter.getIntegerAttr(elementTy, value.getZExtValue());
|
||||||
|
}
|
||||||
|
|
||||||
|
if (!constantAttr) {
|
||||||
|
return rewriter.notifyMatchFailure(
|
||||||
|
op,
|
||||||
|
"tosa.pad to linalg lowering encountered an unknown element type");
|
||||||
|
}
|
||||||
|
|
||||||
|
auto denseAttr = DenseElementsAttr::get(
|
||||||
|
RankedTensorType::get({}, elementTy), constantAttr);
|
||||||
|
auto constantVal = rewriter.create<tosa::ConstOp>(
|
||||||
|
op.getLoc(), denseAttr.getType(), denseAttr);
|
||||||
|
|
||||||
|
rewriter.replaceOpWithNewOp<tosa::PadOp>(
|
||||||
|
op, op.getType(), ValueRange{input, padding, constantVal},
|
||||||
|
op->getAttrs());
|
||||||
|
return success();
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
void PadOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
|
||||||
|
MLIRContext *context) {
|
||||||
|
results.insert<MaterializePadValue>(context);
|
||||||
|
}
|
||||||
|
|
||||||
//===----------------------------------------------------------------------===//
|
//===----------------------------------------------------------------------===//
|
||||||
// Operator Folders.
|
// Operator Folders.
|
||||||
//===----------------------------------------------------------------------===//
|
//===----------------------------------------------------------------------===//
|
||||||
|
@ -415,6 +462,18 @@ ReduceFolder(ReduceAllOp) ReduceFolder(ReduceAnyOp) ReduceFolder(ReduceMaxOp)
|
||||||
return input1();
|
return input1();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
OpFoldResult PadOp::fold(ArrayRef<Attribute> operands) {
|
||||||
|
// If the pad is all zeros we can fold this operation away.
|
||||||
|
if (operands[1]) {
|
||||||
|
auto densePad = operands[1].cast<DenseElementsAttr>();
|
||||||
|
if (densePad.isSplat() && densePad.getSplatValue<APInt>().isZero()) {
|
||||||
|
return input1();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return {};
|
||||||
|
}
|
||||||
|
|
||||||
OpFoldResult SliceOp::fold(ArrayRef<Attribute> operands) {
|
OpFoldResult SliceOp::fold(ArrayRef<Attribute> operands) {
|
||||||
auto inputTy = input().getType().dyn_cast<RankedTensorType>();
|
auto inputTy = input().getType().dyn_cast<RankedTensorType>();
|
||||||
auto outputTy = getType().dyn_cast<RankedTensorType>();
|
auto outputTy = getType().dyn_cast<RankedTensorType>();
|
||||||
|
|
|
@ -66,6 +66,49 @@ func @concat_fold_cast(%arg0: tensor<?x1xf32>) -> tensor<?x?xf32> {
|
||||||
return %0 : tensor<?x?xf32>
|
return %0 : tensor<?x?xf32>
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ----
|
||||||
|
|
||||||
|
// CHECK-LABEL: @pad_noop
|
||||||
|
func @pad_noop(%arg0: tensor<?x?xf32>) -> tensor<?x?xf32> {
|
||||||
|
// CHECK: return %arg0
|
||||||
|
%0 = "tosa.const"() { value = dense<0> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
|
||||||
|
%1 = "tosa.pad"(%arg0, %0) : (tensor<?x?xf32>, tensor<2x2xi32>) -> tensor<?x?xf32>
|
||||||
|
return %1 : tensor<?x?xf32>
|
||||||
|
}
|
||||||
|
|
||||||
|
// ----
|
||||||
|
|
||||||
|
// CHECK-LABEL: @pad_determine_val_i32
|
||||||
|
func @pad_determine_val_i32(%arg0: tensor<?x?xi32>, %arg1 : tensor<2x2xi32>) -> tensor<?x?xi32> {
|
||||||
|
// CHECK: %[[ZERO:.+]] = "tosa.const"() {value = dense<0> : tensor<i32>}
|
||||||
|
// CHECK: "tosa.pad"(%arg0, %arg1, %[[ZERO]])
|
||||||
|
%0 = "tosa.const"() { value = dense<[[1, 0], [0, 1]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
|
||||||
|
%1 = "tosa.pad"(%arg0, %arg1) : (tensor<?x?xi32>, tensor<2x2xi32>) -> tensor<?x?xi32>
|
||||||
|
return %1 : tensor<?x?xi32>
|
||||||
|
}
|
||||||
|
|
||||||
|
// ----
|
||||||
|
|
||||||
|
// CHECK-LABEL: @pad_determine_val_f32
|
||||||
|
func @pad_determine_val_f32(%arg0: tensor<?x?xf32>, %arg1 : tensor<2x2xi32>) -> tensor<?x?xf32> {
|
||||||
|
// CHECK: %[[ZERO:.+]] = "tosa.const"() {value = dense<0.000000e+00> : tensor<f32>}
|
||||||
|
// CHECK: "tosa.pad"(%arg0, %arg1, %[[ZERO]])
|
||||||
|
%0 = "tosa.const"() { value = dense<[[1, 0], [0, 1]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
|
||||||
|
%1 = "tosa.pad"(%arg0, %arg1) : (tensor<?x?xf32>, tensor<2x2xi32>) -> tensor<?x?xf32>
|
||||||
|
return %1 : tensor<?x?xf32>
|
||||||
|
}
|
||||||
|
|
||||||
|
// ----
|
||||||
|
|
||||||
|
// CHECK-LABEL: @pad_determine_val_quant
|
||||||
|
func @pad_determine_val_quant(%arg0: tensor<?x?xi32>, %arg1 : tensor<2x2xi32>) -> tensor<?x?xi32> {
|
||||||
|
// CHECK: %[[ZERO:.+]] = "tosa.const"() {value = dense<42> : tensor<i32>}
|
||||||
|
// CHECK: "tosa.pad"(%arg0, %arg1, %[[ZERO]])
|
||||||
|
%0 = "tosa.const"() { value = dense<[[1, 0], [0, 1]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
|
||||||
|
%1 = "tosa.pad"(%arg0, %arg1) { quantization_info = {input_zp = 42:i32} } : (tensor<?x?xi32>, tensor<2x2xi32>) -> tensor<?x?xi32>
|
||||||
|
return %1 : tensor<?x?xi32>
|
||||||
|
}
|
||||||
|
|
||||||
// -----
|
// -----
|
||||||
|
|
||||||
// CHECK-LABEL: @mul_one_different_shape
|
// CHECK-LABEL: @mul_one_different_shape
|
||||||
|
|
Loading…
Reference in New Issue