Skip to content

[mlir][tosa] Make TOSA RESIZE's scale, offset, border as Input #124956

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.

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 5 commits into from
Feb 19, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
7 changes: 4 additions & 3 deletions mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
Original file line number Diff line number Diff line change
Expand Up @@ -1822,9 +1822,9 @@ def Tosa_ResizeOp : Tosa_InferShapedTypeOp<"resize"> {

let arguments = (ins
Tosa_Tensor4D:$input,
Tosa_IntArrayAttr4:$scale,
Tosa_IntArrayAttr2:$offset,
Tosa_IntArrayAttr2:$border,
Rank4TosaShape:$scale,
Rank2TosaShape:$offset,
Rank2TosaShape:$border,
Tosa_ResizeTypeAttr:$mode
);

Expand All @@ -1833,6 +1833,7 @@ def Tosa_ResizeOp : Tosa_InferShapedTypeOp<"resize"> {
);

let hasFolder = 1;
let hasVerifier = 1;
}

//===----------------------------------------------------------------------===//
Expand Down
3 changes: 3 additions & 0 deletions mlir/include/mlir/Dialect/Tosa/Utils/ConversionUtils.h
Original file line number Diff line number Diff line change
Expand Up @@ -240,6 +240,9 @@ SmallVector<int64_t> convertFromMlirShape(ArrayRef<int64_t> shape);
bool getConstShapeValue(Operation *op,
llvm::SmallVector<int64_t> &result_shape);

// returns a small vector of int64_t values that attr contains
SmallVector<int64_t> convertFromIntAttr(const DenseElementsAttr &attr,
const int rank);
} // namespace tosa
} // namespace mlir

Expand Down
21 changes: 15 additions & 6 deletions mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1387,7 +1387,10 @@ class ResizeUnaryConverter : public OpRewritePattern<tosa::ResizeOp> {
return success();
}

ArrayRef<int64_t> scale = op.getScale();
SmallVector<int64_t> scale;
if (!tosa::getConstShapeValue(op.getScale().getDefiningOp(), scale)) {
return failure();
}

// Collapse the unit width and height away.
SmallVector<ReassociationExprs, 4> reassociationMap(2);
Expand Down Expand Up @@ -1488,8 +1491,9 @@ class MaterializeResizeBroadcast : public OpRewritePattern<tosa::ResizeOp> {
resizeShape.push_back(channels);

auto resizeTy = resultTy.clone(resizeShape);
auto resize =
builder.create<tosa::ResizeOp>(resizeTy, input, op->getAttrs());
auto resize = builder.create<tosa::ResizeOp>(resizeTy, input, op.getScale(),
op.getOffset(), op.getBorder(),
op.getMode());

// Collapse an unit result dims.
SmallVector<ReassociationExprs, 4> reassociationMap(2);
Expand Down Expand Up @@ -1604,9 +1608,14 @@ class GenericResizeConverter : public OpRewritePattern<tosa::ResizeOp> {
Value inY = b.create<arith::IndexCastOp>(b.getI32Type(), y);
Value inX = b.create<arith::IndexCastOp>(b.getI32Type(), x);

ArrayRef<int64_t> offset = op.getOffset();
ArrayRef<int64_t> border = op.getBorder();
ArrayRef<int64_t> scale = op.getScale();
SmallVector<int64_t> scale, offset, border;
if (!tosa::getConstShapeValue(op.getScale().getDefiningOp(), scale) ||
!tosa::getConstShapeValue(op.getOffset().getDefiningOp(), offset) ||
!tosa::getConstShapeValue(op.getBorder().getDefiningOp(), border)) {
return rewriter.notifyMatchFailure(
op, "tosa.resize scale/offset/border should have compile time "
"constant values.");
}

Value yScaleN, yScaleD, xScaleN, xScaleD;
yScaleN = b.create<arith::ConstantOp>(b.getI32IntegerAttr(scale[0]));
Expand Down
19 changes: 16 additions & 3 deletions mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1034,9 +1034,22 @@ OpFoldResult PadOp::fold(FoldAdaptor adaptor) {
// Fold away cases where a tosa.resize operation returns a copy
// of the input image.
OpFoldResult ResizeOp::fold(FoldAdaptor adaptor) {
ArrayRef<int64_t> offset = getOffset();
ArrayRef<int64_t> border = getBorder();
ArrayRef<int64_t> scale = getScale();
auto scaleAttr =
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getScale());
auto offsetAttr =
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getOffset());
auto borderAttr =
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getBorder());
if (!scaleAttr || !offsetAttr || !borderAttr) {
return {};
}

auto scale = tosa::convertFromIntAttr(scaleAttr, /* rank = */ 4);
auto offset = tosa::convertFromIntAttr(offsetAttr, /* rank = */ 2);
auto border = tosa::convertFromIntAttr(borderAttr, /* rank = */ 2);
if (scale.size() != 4 || offset.size() != 2 || border.size() != 2) {
return {};
}

// Check unit scaling.
if (scale[0] != scale[1] || scale[2] != scale[3]) {
Expand Down
103 changes: 100 additions & 3 deletions mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1685,9 +1685,14 @@ LogicalResult tosa::ResizeOp::inferReturnTypeComponents(
(inputWidth == ShapedType::kDynamic))
return failure();

llvm::ArrayRef<int64_t> scaleInt = adaptor.getScale();
llvm::ArrayRef<int64_t> offsetInt = adaptor.getOffset();
llvm::ArrayRef<int64_t> borderInt = adaptor.getBorder();
SmallVector<int64_t> scaleInt, offsetInt, borderInt;
if (!tosa::getConstShapeValue(adaptor.getScale().getDefiningOp(), scaleInt) ||
!tosa::getConstShapeValue(adaptor.getOffset().getDefiningOp(),
offsetInt) ||
!tosa::getConstShapeValue(adaptor.getBorder().getDefiningOp(),
borderInt)) {
return failure();
}

// Compute the output shape based on attributes: scale, offset, and border.
outputShape[1] =
Expand All @@ -1704,6 +1709,98 @@ LogicalResult tosa::ResizeOp::inferReturnTypeComponents(
return success();
}

LogicalResult tosa::ResizeOp::verify() {
const Value input = getInput();
const Value output = getOutput();
const RankedTensorType inputType =
llvm::dyn_cast<RankedTensorType>(input.getType());
const RankedTensorType outputType =
llvm::dyn_cast<RankedTensorType>(output.getType());

if (!inputType)
return emitOpError("expect a ranked input tensor");
if (!outputType)
return emitOpError("expect a ranked output tensor");

const int64_t oh = outputType.getDimSize(1);
const int64_t ow = outputType.getDimSize(2);
const int64_t ih = inputType.getDimSize(1);
const int64_t iw = inputType.getDimSize(2);

SmallVector<int64_t> scaleValues;
SmallVector<int64_t> offsetValues;
SmallVector<int64_t> borderValues;
if (!tosa::getConstShapeValue(getScale().getDefiningOp(), scaleValues) ||
!tosa::getConstShapeValue(getOffset().getDefiningOp(), offsetValues) ||
!tosa::getConstShapeValue(getBorder().getDefiningOp(), borderValues)) {
// Skip following checks if shape is not constant
return success();
}

if (llvm::any_of(scaleValues, [](int64_t s) { return s <= 0; }))
return emitOpError("expect all scale values to be > 0, got ")
<< scaleValues;

const int64_t scaleYN = scaleValues[0];
const int64_t scaleYD = scaleValues[1];
const int64_t scaleXN = scaleValues[2];
const int64_t scaleXD = scaleValues[3];

const int64_t offsetY = offsetValues[0];
const int64_t offsetX = offsetValues[1];

const int64_t borderY = borderValues[0];
const int64_t borderX = borderValues[1];

auto idivCheck = [](const int64_t lhs,
const int64_t rhs) -> std::optional<int64_t> {
if (lhs % rhs != 0)
return std::nullopt;
return lhs / rhs;
};

// Don't check with input height that could be broadcast (ih != 1)
// since Linalg, a consumer of TOSA, expects broadcasting support
// in resize to be available. Taking the cautious approach for now,
// we can consider removing support for broadcasting later.
if (ih != ShapedType::kDynamic && ih != 1) {
const std::optional<int64_t> calculatedOutHeightMinusOne =
idivCheck((ih - 1) * scaleYN - offsetY + borderY, scaleYD);
if (!calculatedOutHeightMinusOne.has_value())
return emitOpError("expected (input_height - 1) * scale_y_n - offset_y + "
"border_y ")
<< "to be wholly divisible by scale_y_d, got ((" << ih
<< " - 1) * " << scaleYN << " - " << offsetY << " + " << borderY
<< ") / " << scaleYD;
const int64_t calculatedOutHeight = calculatedOutHeightMinusOne.value() + 1;
if (oh != ShapedType::kDynamic && calculatedOutHeight != oh)
return emitOpError("calculated output height did not match expected: ")
<< "calculated=" << calculatedOutHeight << ", expected=" << oh;
}

// Don't check with input width that could be broadcast (iw != 1)
// since Linalg, a consumer of TOSA, expects broadcasting support
// in resize to be available. Taking the cautious approach for now,
// we can consider removing support for broadcasting later.
if (iw != ShapedType::kDynamic && iw != 1) {
const int64_t scaledInWidth = (iw - 1) * scaleXN - offsetX + borderX;
const std::optional<int64_t> calculatedOutWidthMinusOne =
idivCheck(scaledInWidth, scaleXD);
if (!calculatedOutWidthMinusOne.has_value())
return emitOpError("expected (input_width - 1) * scale_x_n - offset_x + "
"border_x ")
<< "to be wholly divisible by scale_x_d, got ((" << iw
<< " - 1) * " << scaleXN << " - " << offsetX << " + " << borderX
<< ") / " << scaleXD;
const int64_t calculatedOutWidth = calculatedOutWidthMinusOne.value() + 1;
if (ow != ShapedType::kDynamic && calculatedOutWidth != ow)
return emitOpError("calculated output width did not match expected: ")
<< "calculated=" << calculatedOutWidth << ", expected=" << ow;
}

return success();
}

LogicalResult tosa::ScatterOp::inferReturnTypeComponents(
MLIRContext *context, ::std::optional<Location> location,
ScatterOp::Adaptor adaptor,
Expand Down
Loading