Skip to content

[mlir][tosa] Remove optional for pad_const and remove input_zp attr for PadOp #129336

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Mar 5, 2025
Merged

[mlir][tosa] Remove optional for pad_const and remove input_zp attr for PadOp #129336

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 0 additions & 8 deletions mlir/include/mlir/Dialect/Tosa/IR/TosaOpBase.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -197,14 +197,6 @@ def Tosa_PadOpQuantInfoBuilder : OpBuilder<
input, paddings);
}]>;

def Tosa_ExplicitValuePadOpQuantInfoBuilder : OpBuilder<
(ins "Type":$outputType, "Value":$input, "Value":$paddings,
"Value":$pad_value),
[{
buildExplicitValuePadOpWithQuantInfo($_builder, $_state, outputType,
input, paddings, pad_value);
}]>;

// Wrapper over base I32EnumAttr to set common fields.
class Tosa_I32Enum<string name, string description, list<I32EnumAttrCase> cases>
: I32EnumAttr<name, description, cases> {
Expand Down
4 changes: 4 additions & 0 deletions mlir/include/mlir/Dialect/Tosa/IR/TosaOps.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -168,6 +168,10 @@ namespace tosa {
std::optional<Value> createZeroPointTensor(OpBuilder &builder, Location loc,
Type srcElemType, int64_t zp = 0);

// Create a pad-const const tensor with value of `val` of required data-type
Value createPadConstTensor(OpBuilder &builder, Location loc, Value src,
int32_t val = 0);

} // namespace tosa
} // namespace mlir

Expand Down
7 changes: 2 additions & 5 deletions mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1891,8 +1891,7 @@ def Tosa_PadOp : Tosa_InferShapedTypeOp<"pad"> {
let arguments = (ins
Tosa_RankedTensor:$input1,
Tosa_Shape:$padding,
Optional<Tosa_ScalarTensor>:$pad_const,
OptionalAttr<I32Attr>:$input_zp
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should this change be split into a separate commit, it seems orthognal to making the padding constant non optional.

Copy link
Member Author

@Jerry-Ge Jerry-Ge Mar 4, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The upstreaming plan is to combine patches related to each operator into a single patch. Those two are the last for PadOp to do some removal.

Tosa_ScalarTensor:$pad_const
);

let results = (outs
Expand All @@ -1904,10 +1903,8 @@ def Tosa_PadOp : Tosa_InferShapedTypeOp<"pad"> {
Extension<[Tosa_EXT_FP8E4M3, Tosa_EXT_FP8E5M2, Tosa_EXT_BF16]>,
];

let builders = [Tosa_PadOpQuantInfoBuilder,
Tosa_ExplicitValuePadOpQuantInfoBuilder];
let builders = [Tosa_PadOpQuantInfoBuilder];

let hasCanonicalizer = 1;
let hasFolder = 1;
let hasVerifier = 1;
}
Expand Down
92 changes: 41 additions & 51 deletions mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,10 +36,10 @@ TensorType inferReshapeInputType(TypedValue<TensorType> input,
return input.getType();

// The input type must be cast into a tensor with the same rank and all static
// dimensions set to 1. This prevents the generation of a tensor.collapse_shape
// op that converts a dynamically shaped tensor into a 0D tensor. While such
// construct is not incorrect on its own, bufferization cannot properly handle
// it at the moment, so we avoid it.
// dimensions set to 1. This prevents the generation of a
// tensor.collapse_shape op that converts a dynamically shaped tensor into a
// 0D tensor. While such construct is not incorrect on its own, bufferization
// cannot properly handle it at the moment, so we avoid it.
SmallVector<int64_t> shape(input.getType().getRank(), 1);
return input.getType().clone(shape);
}
Expand All @@ -58,29 +58,31 @@ TensorType inferReshapeExpandedType(TensorType inputType,
int64_t totalSize = inputIsStatic ? inputType.getNumElements() : -1;

// Compute result shape
auto resultShape = llvm::map_to_vector(newShape, [&](int64_t size) -> int64_t {
// If this is not a placeholder, do not change it
if (size >= 0)
return size;

// If we do not know the total size of the tensor, keep this dimension
// dynamic in the result shape.
if (!inputIsStatic)
return ShapedType::kDynamic;

// Calculate the product of all elements in 'newShape' except for the -1
// placeholder, which we discard by negating the result.
int64_t totalSizeNoPlaceholder = -std::accumulate(
newShape.begin(), newShape.end(), 1, std::multiplies<int64_t>());

// If there is a 0 component in 'newShape', resolve the placeholder as 0.
if (totalSizeNoPlaceholder == 0)
return 0;

// Resolve the placeholder as the quotient between the total tensor size and
// the product of all other sizes.
return totalSize / totalSizeNoPlaceholder;
});
auto resultShape =
llvm::map_to_vector(newShape, [&](int64_t size) -> int64_t {
// If this is not a placeholder, do not change it.
if (size >= 0)
return size;

// If we do not know the total size of the tensor, keep this dimension
// dynamic in the result shape.
if (!inputIsStatic)
return ShapedType::kDynamic;

// Calculate the product of all elements in 'newShape' except for the -1
// placeholder, which we discard by negating the result.
int64_t totalSizeNoPlaceholder = -std::accumulate(
newShape.begin(), newShape.end(), 1, std::multiplies<int64_t>());

// If there is a 0 component in 'newShape', resolve the placeholder as
// 0.
if (totalSizeNoPlaceholder == 0)
return 0;

// Resolve the placeholder as the quotient between the total tensor size
// and the product of all other sizes.
return totalSize / totalSizeNoPlaceholder;
});

bool resultIsStatic = !ShapedType::isDynamicShape(resultShape);

Expand Down Expand Up @@ -108,7 +110,8 @@ TensorType inferReshapeCollapsedType(TensorType lhsType, TensorType rhsType) {
if (lhsShape.empty() || rhsShape.empty())
return lhsType.clone(ArrayRef<int64_t>{});

if (ShapedType::isDynamicShape(lhsShape) || ShapedType::isDynamicShape(rhsShape))
if (ShapedType::isDynamicShape(lhsShape) ||
ShapedType::isDynamicShape(rhsShape))
return lhsType.clone({ShapedType::kDynamic});

SmallVector<int64_t> intermediateShape;
Expand Down Expand Up @@ -150,14 +153,16 @@ TensorType inferReshapeCollapsedType(TensorType lhsType, TensorType rhsType) {
}

SmallVector<ReassociationExprs>
createReassociationMapForCollapse(OpBuilder &builder, Type srcType, Type dstType) {
createReassociationMapForCollapse(OpBuilder &builder, Type srcType,
Type dstType) {
auto srcShape = cast<TensorType>(srcType).getShape();
auto dstShape = cast<TensorType>(dstType).getShape();

if (srcShape.empty() || dstShape.empty())
return {};

if (ShapedType::isDynamicShape(srcShape) || ShapedType::isDynamicShape(dstShape)) {
if (ShapedType::isDynamicShape(srcShape) ||
ShapedType::isDynamicShape(dstShape)) {
assert(dstShape.size() == 1);
SmallVector<AffineExpr, 2> exprs;
for (auto i : llvm::seq<int64_t>(srcShape.size()))
Expand Down Expand Up @@ -249,14 +254,16 @@ class ReshapeConverter : public OpConversionPattern<tosa::ReshapeOp> {
auto collapsedType = inferReshapeCollapsedType(inputType, expandedType);

// Cast input if needed
auto castInput = rewriter.createOrFold<tensor::CastOp>(loc, inputType, input);
auto castInput =
rewriter.createOrFold<tensor::CastOp>(loc, inputType, input);

// Emit collaspe-expand pair
auto collapsed = createCollapse(rewriter, loc, collapsedType, castInput);
auto expanded = createExpand(rewriter, loc, expandedType, collapsed);

// Cast to final result type if needed
auto result = rewriter.createOrFold<tensor::CastOp>(loc, resultType, expanded);
auto result =
rewriter.createOrFold<tensor::CastOp>(loc, resultType, expanded);
rewriter.replaceOp(reshape, result);
return success();
}
Expand Down Expand Up @@ -350,29 +357,12 @@ class PadConverter : public OpConversionPattern<tosa::PadOp> {
}

ShapedType inputTy = cast<ShapedType>(input.getType());
Type elementTy = inputTy.getElementType();
int64_t rank = inputTy.getRank();

// Setup the default constantAttr.

Value padConstant;

if (padOp.getPadConst()) {
padConstant = rewriter.createOrFold<tensor::ExtractOp>(
loc, padOp.getPadConst(), ValueRange({}));
} else {
TypedAttr constantAttr;
if (isa<FloatType>(elementTy)) {
constantAttr = rewriter.getFloatAttr(elementTy, 0.0);
} else if (isa<IntegerType>(elementTy) && !padOp.getInputZpAttr()) {
constantAttr = rewriter.getIntegerAttr(elementTy, 0);
} else if (isa<IntegerType>(elementTy) && padOp.getInputZpAttr()) {
int64_t value = padOp.getInputZpAttr().getInt();
constantAttr = rewriter.getIntegerAttr(elementTy, value);
}
if (constantAttr)
padConstant = rewriter.create<arith::ConstantOp>(loc, constantAttr);
}
Value padConstant = rewriter.createOrFold<tensor::ExtractOp>(
loc, padOp.getPadConst(), ValueRange({}));

if (!padConstant) {
return rewriter.notifyMatchFailure(
Expand Down
47 changes: 0 additions & 47 deletions mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -175,53 +175,6 @@ void TransposeOp::getCanonicalizationPatterns(RewritePatternSet &results,
results.add<ConsolidateTransposeOptimization, TransposeIsReshape>(context);
}

struct MaterializePadValue : public OpRewritePattern<tosa::PadOp> {
using OpRewritePattern::OpRewritePattern;

LogicalResult matchAndRewrite(tosa::PadOp op,
PatternRewriter &rewriter) const override {
if (op.getPadConst())
return failure();

auto input = op.getInput1();
auto padding = op.getPadding();

ShapedType inputTy = llvm::cast<ShapedType>(input.getType());
Type elementTy = inputTy.getElementType();

Attribute constantAttr;
if (llvm::isa<FloatType>(elementTy)) {
constantAttr = rewriter.getFloatAttr(elementTy, 0.0);
} else if (llvm::isa<IntegerType>(elementTy) && !op.getInputZpAttr()) {
constantAttr = rewriter.getIntegerAttr(elementTy, 0);
} else if (llvm::isa<IntegerType>(elementTy) && op.getInputZpAttr()) {
int64_t value = op.getInputZpAttr().getInt();
constantAttr = rewriter.getIntegerAttr(elementTy, value);
}

if (!constantAttr) {
return rewriter.notifyMatchFailure(
op,
"tosa.pad to linalg lowering encountered an unknown element type");
}

auto denseAttr = DenseElementsAttr::get(
RankedTensorType::get({1}, 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(RewritePatternSet &results,
MLIRContext *context) {
results.add<MaterializePadValue>(context);
}

struct MaxPool2dIsNoOp : public OpRewritePattern<tosa::MaxPool2dOp> {
using OpRewritePattern::OpRewritePattern;

Expand Down
44 changes: 22 additions & 22 deletions mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -214,6 +214,22 @@ void mlir::tosa::printTypeOrAttr(OpAsmPrinter &p, Operation *op, TypeAttr type,
}
}

// Create a pad-const const tensor with value of `val` of required data-type
Value mlir::tosa::createPadConstTensor(OpBuilder &builder, Location loc,
Value src, int32_t val) {
const auto srcType = getElementTypeOrSelf(src);
const auto srcElemType = getElementTypeOrSelf(src);
const auto padConstType = mlir::RankedTensorType::get({1}, srcType);
const auto padConstEType = mlir::RankedTensorType::get({1}, srcElemType);
const auto padConstAttr{
llvm::isa<FloatType>(srcElemType)
? DenseElementsAttr::get(padConstEType,
builder.getFloatAttr(srcElemType, val))
: DenseElementsAttr::get(padConstEType,
builder.getIntegerAttr(srcElemType, val))};
return builder.create<tosa::ConstOp>(loc, padConstType, padConstAttr);
}

//===----------------------------------------------------------------------===//
// Tosa utilities.
//===----------------------------------------------------------------------===//
Expand Down Expand Up @@ -679,30 +695,14 @@ static void buildUnaryOpWithQuantInfo(OpBuilder &builder,
static void buildPadOpWithQuantInfo(OpBuilder &builder, OperationState &result,
Type outputType, Value input,
Value paddings) {
result.addOperands({input, paddings});
auto quantAttr = buildPadOpQuantizationAttr(builder, input);
const Location loc{result.location};
int32_t zp{0};
const auto quantAttr = buildPadOpQuantizationAttr(builder, input);
if (quantAttr) {
result.addAttribute("input_zp",
builder.getI32IntegerAttr(
static_cast<int32_t>(quantAttr.getInputZp())));
}
result.types.push_back(outputType);
}

/// This builder is called on TOSA pad operator when an explicit pad_const
/// value is passed in. It also optionally constructs quantization_attr.
static void buildExplicitValuePadOpWithQuantInfo(OpBuilder &builder,
OperationState &result,
Type outputType, Value input,
Value paddings,
Value padConst) {
result.addOperands({input, paddings, padConst});
auto quantAttr = buildPadOpQuantizationAttr(builder, input);
if (quantAttr) {
result.addAttribute("input_zp",
builder.getI32IntegerAttr(
static_cast<int32_t>(quantAttr.getInputZp())));
zp = static_cast<int32_t>(quantAttr.getInputZp());
}
const auto padConstOp{createPadConstTensor(builder, loc, input, zp)};
result.addOperands({input, paddings, padConstOp});
result.types.push_back(outputType);
}

Expand Down
33 changes: 13 additions & 20 deletions mlir/lib/Dialect/Tosa/Transforms/TosaDecomposeTransposeConv.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,16 +148,16 @@ class TransposeConvStridedConverter
return rewriter.notifyMatchFailure(
op, "zero point must be zero for non-int8 integer types");

if (weightZpVal != 0) {
weight = CreateOpAndInferShape<tosa::PadOp>(
rewriter, loc, UnrankedTensorType::get(weightETy), weight,
weightPaddingVal, nullptr, rewriter.getI32IntegerAttr(weightZpVal));

} else {
weight = CreateOpAndInferShape<tosa::PadOp>(
rewriter, loc, UnrankedTensorType::get(weightETy), weight,
weightPaddingVal);
}
// construct pad_const values from zp values
ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
const Value inputPadConst =
createPadConstTensor(builder, op->getLoc(), input, inputZpVal);
const Value weightPadConst =
createPadConstTensor(builder, op->getLoc(), input, weightZpVal);

weight = CreateOpAndInferShape<tosa::PadOp>(
rewriter, loc, UnrankedTensorType::get(weightETy), weight,
weightPaddingVal, weightPadConst);

weightTy = cast<ShapedType>(weight.getType());
weightHeight = weightTy.getDimSize(1);
Expand All @@ -169,7 +169,6 @@ class TransposeConvStridedConverter
stride[0], weightWidth / stride[1],
stride[1], inputChannels};

ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
weight = CreateOpAndInferShape<tosa::ReshapeOp>(
builder, UnrankedTensorType::get(weightETy), weight,
getTosaConstShape(rewriter, loc, weightReshapeDims0));
Expand Down Expand Up @@ -206,15 +205,9 @@ class TransposeConvStridedConverter
Value inputPaddingVal =
getTosaConstShape(rewriter, op->getLoc(), inputPadding);

if (inputZpVal != 0) {
input = CreateOpAndInferShape<tosa::PadOp>(
rewriter, loc, UnrankedTensorType::get(inputETy), input,
inputPaddingVal, nullptr, rewriter.getI32IntegerAttr(inputZpVal));
} else {
input = CreateOpAndInferShape<tosa::PadOp>(
rewriter, loc, UnrankedTensorType::get(inputETy), input,
inputPaddingVal);
}
input = CreateOpAndInferShape<tosa::PadOp>(
rewriter, loc, UnrankedTensorType::get(inputETy), input,
inputPaddingVal, inputPadConst);

// We use a zero bias as we need to broadcast the bias.
auto zeroBias = rewriter.create<tosa::ConstOp>(
Expand Down
Loading