Skip to content

Commit

Permalink
Add scale and zero_point to xten_nn.quantize/dequantize
Browse files Browse the repository at this point in the history
  • Loading branch information
@jorickert
jorickert committed Feb 20, 2025
1 parent eca83cf commit 9728c35
Show file tree
Hide file tree
Showing 8 changed files with 406 additions and 106 deletions.
46 changes: 32 additions & 14 deletions include/xten/Dialect/XTenNN/IR/XTenNNOps.td
View file
Original file line number Diff line number Diff line change
Expand Up @@ -113,59 +113,77 @@ def XTenNN_QuantizeOp: XTenNN_Op<"quantize", [
Elementwise,
Pure,
SameOperandsAndResultShape]> {
let summary = "Quantizes a float32 tensor to a signless or unsigned integer tensor of given width.";
let summary = "Quantizes a float tensor to a signless or unsigned integer tensor of given width.";
let description = [{
Quantizes a given float32 tensor into a signless or unsigned integer tensor of given width.
Quantizes a given float tensor into a signless or unsigned integer tensor of given width.
Since tosa is using signless/unsigned types currently, we also consider signless integer types for
signed ones when the type is not unsigned until tosa support signed integers.

Applies the following linear quantization to the input tensor x:
y = round( x / 2^shift )
y = round((x / scale) + zero_point)

Where 2^shift is equal to the scale of the quantize operation and
the shift is an attribute of the operation in si32.
Iff log2(scale) is representable as si32 and zero_point == 0, shift is set to log2(scale).
In this case the quantization is equal to:
y = round( x / 2^shift )

Round will saturate to the range of the output type and the rounding mode is set to half
to nearest even.
}];

let arguments = (ins
F32Tensor:$input,
SI32Attr:$shift
XTenNN_AnyFloatTensor:$input,
OptionalAttr<SI32Attr>:$shift,
F32Attr:$scale, // Restricted to F32 for now, but may be relaxed in the future
XTenNN_AnyIntegerAttr:$zero_point
);

let results = (outs XTenNN_AnySignlessOrUnsignedIntegerTensor:$output);
let builders = [
OpBuilder<(ins "::mlir::Type":$output, "::mlir::Value":$input, "::mlir::FloatAttr":$scale, "::mlir::IntegerAttr":$zeroPoint)>,
OpBuilder<(ins "::mlir::Type":$output, "::mlir::Value":$input, "int32_t":$shift)>
];


let assemblyFormat = [{ `(`$input `:` type($input)`)` attr-dict `->` type($output) }];

let hasFolder = 1;
let hasVerifier = 1;
}

def XTenNN_DequantizeOp: XTenNN_Op<"dequantize", [
Elementwise,
Pure,
SameOperandsAndResultShape]> {
let summary = "Dequantizes a signless/unsigned integer tensor of given bitwidth to a float32 tensor.";
let summary = "Dequantizes a signless/unsigned integer tensor of given bitwidth to a float tensor.";
let description = [{
Dequantizes a signless/unsigned integer tensor of given bitwidth to a float32 tensor.
Dequantizes a signless/unsigned integer tensor of given bitwidth to a float tensor.
Since tosa is using signless/unsigned types currently, we also consider signless integer types for
signed ones when the type is not unsigned until tosa support signed integers.

Applies the following linear dequantization to the input tensor x:
y = x * ( 2^shift )
y = (x - zero_point) * scale

Where 2^shift is equal to scale of the dequantize operation and
the shift is an attribute of the operation in si32.
Iff log2(scale) is representable as si32 and zero_point == 0, shift is set to log2(scale).
In this case the dequantization is equal to:
y = x * ( 2^shift )
}];

let arguments = (ins
XTenNN_AnySignlessOrUnsignedIntegerTensor:$input,
SI32Attr:$shift
OptionalAttr<SI32Attr>:$shift,
F32Attr:$scale, // Restricted to F32 for now, but may be relaxed in the future
XTenNN_AnyIntegerAttr:$zero_point
);

let results = (outs F32Tensor:$output);
let results = (outs XTenNN_AnyFloatTensor:$output);

let builders = [
OpBuilder<(ins "::mlir::Type":$output, "::mlir::Value":$input, "::mlir::FloatAttr":$scale, "::mlir::IntegerAttr":$zeroPoint)>,
OpBuilder<(ins "::mlir::Type":$output, "::mlir::Value":$input, "int32_t":$shift)>
];

let assemblyFormat = [{ `(`$input `:` type($input)`)` attr-dict `->` type($output) }];
let hasVerifier = 1;
}

def XTenNN_GroupQuantizeOp: XTenNN_Op<"group_quantize", [
Expand Down
9 changes: 9 additions & 0 deletions include/xten/Dialect/XTenNN/IR/XTenNNTypes.td
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,15 @@ def XTenNN_AnySignlessOrUnsignedIntegerTensor : TensorOf<

def XTenNN_AnyFloatTensor : TensorOf<[AnyFloat]>;

def XTenNN_AnyIntegerAttr:
TypedAttrBase<
AnyInteger, "IntegerAttr",
CPred<"::llvm::isa<::mlir::IntegerAttr>($_self)">,
"Any integer attr"> {
let returnType = [{ ::llvm::APInt }];
let constBuilderCall = ?;
}

def XTenNN_AnyIntegerOrFloat : AnyTypeOf<[AnyInteger, AnyFloat], "Integer or Float">;

#endif // XTENNN_TYPES
127 changes: 122 additions & 5 deletions lib/Dialect/XTenNN/IR/XTenNNOps.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
#include "xten/Dialect/XTenNN/IR/XTenNNBase.h"
#include "xten/Dialect/XTenNN/Interfaces/EnclaveOpInterfaces.h"

#include <cfenv>
#include <cstdint>

using namespace mlir;
Expand Down Expand Up @@ -444,25 +445,141 @@ LogicalResult SubgraphOp::inferReturnTypeComponents(
// XTenNNDialect
//===----------------------------------------------------------------------===//

namespace {
std::optional<int32_t> getShiftValue(float constValue) {
const float log2Value = std::log2f(constValue);

// The log2 of the value must not have fractions.
if (std::roundf(log2Value) != log2Value)
return {};

return static_cast<int32_t>(log2Value);
}

float getScaleFromShift(int32_t shift) {
std::feclearexcept(FE_ALL_EXCEPT);
errno = 0;
const float scale = exp2f(shift);
assert(!std::fetestexcept(FE_OVERFLOW));
assert(errno == 0);
return scale;
}
} // namespace

void amd::xten_nn::QuantizeOp::build(mlir::OpBuilder &odsBuilder,
mlir::OperationState &odsState,
mlir::Type output, mlir::Value input,
mlir::FloatAttr scale,
mlir::IntegerAttr zeroPoint) {
const bool zeroPointIsZero = zeroPoint.getValue().isZero();
assert(scale.getType().isF32());
const auto shiftValue = getShiftValue(scale.getValue().convertToFloat());
if (zeroPointIsZero && shiftValue) {
return build(odsBuilder, odsState, output, input,
odsBuilder.getSI32IntegerAttr(*shiftValue), scale, zeroPoint);
}
return build(odsBuilder, odsState, output, input, nullptr, scale, zeroPoint);
}

void amd::xten_nn::QuantizeOp::build(mlir::OpBuilder &odsBuilder,
mlir::OperationState &odsState,
mlir::Type output, mlir::Value input,
int32_t shift) {
const auto outputElemType = cast<TensorType>(output).getElementType();
return build(odsBuilder, odsState, output, input,
odsBuilder.getSI32IntegerAttr(shift),
odsBuilder.getF32FloatAttr(getScaleFromShift(shift)),
odsBuilder.getIntegerAttr(outputElemType, 0));
}

LogicalResult amd::xten_nn::QuantizeOp::verify() {
if (getResult().getType().getElementType() != getZeroPointAttr().getType()) {
return emitOpError("Result elem type needs to match match zero point type");
}
// if shift is set, zero point needs to be zero and scale needs to match
if (getShift()) {
const auto computedShift = getShiftValue(getScale().convertToFloat());
if (!computedShift || computedShift != *getShift()) {
return emitOpError(
"Shift set, but does not match shift calculated from scale");
}
if (!getZeroPoint().isZero()) {
return emitOpError("Shift set, but zero_point not zero");
}
}

return success();
}

OpFoldResult amd::xten_nn::QuantizeOp::fold(FoldAdaptor adaptor) {
// Fold away cases where a xten_nn.quantize is preceeded by xten_nn.dequantize
// that uses the same shift factor and has same types.
// Fold away cases where a xten_nn.quantize is preceeded by
// xten_nn.dequantize that uses the same scale factor, zeroPoint and has same
// types.

auto dequantizeOp =
dyn_cast_or_null<amd::xten_nn::DequantizeOp>(getInput().getDefiningOp());
if (!dequantizeOp)
return {};

if (!dequantizeOp->hasOneUse() || dequantizeOp.getShift() != getShift())
return {};

auto dequantizeInput = dequantizeOp.getInput();
if (dequantizeInput.getType() != getType())
return {};

if (!dequantizeOp->hasOneUse() || dequantizeOp.getScale() != getScale() ||
dequantizeOp.getZeroPoint() != getZeroPoint())
return {};

return dequantizeInput;
}

void amd::xten_nn::DequantizeOp::build(mlir::OpBuilder &odsBuilder,
mlir::OperationState &odsState,
mlir::Type output, mlir::Value input,
mlir::FloatAttr scale,
mlir::IntegerAttr zeroPoint) {
const bool zeroPointIsZero = zeroPoint.getValue().isZero();
assert(scale.getType().isF32());
const auto shiftValue = getShiftValue(scale.getValue().convertToFloat());
if (zeroPointIsZero && shiftValue) {
return build(odsBuilder, odsState, output, input,
odsBuilder.getSI32IntegerAttr(*shiftValue), scale, zeroPoint);
}
return build(odsBuilder, odsState, output, input, nullptr, scale, zeroPoint);
}

void amd::xten_nn::DequantizeOp::build(mlir::OpBuilder &odsBuilder,
mlir::OperationState &odsState,
mlir::Type output, mlir::Value input,
int32_t shift) {
const auto inputElemType = cast<TensorType>(input.getType()).getElementType();
return build(odsBuilder, odsState, output, input,
odsBuilder.getSI32IntegerAttr(shift),
odsBuilder.getF32FloatAttr(getScaleFromShift(shift)),
odsBuilder.getIntegerAttr(inputElemType, 0));
}

LogicalResult amd::xten_nn::DequantizeOp::verify() {
// Input elem type should match zero point type
if (cast<TensorType>(getOperand().getType()).getElementType() !=
getZeroPointAttr().getType()) {
return emitOpError(
"Operand elem type needs to match match zero point type");
}
// if shift is set, zero point needs to be zero and scale needs to match
if (getShift()) {
const auto computedShift = getShiftValue(getScale().convertToFloat());
if (!computedShift || computedShift != *getShift()) {
return emitOpError(
"Shift set, but does not match shift calculated from scale");
}
if (!getZeroPoint().isZero()) {
return emitOpError("Shift set, but zero_point not zero");
}
}

return success();
}

OpFoldResult amd::xten_nn::GroupQuantizeOp::fold(FoldAdaptor adaptor) {
// Fold away cases where a xten_nn.group_quantize is preceeded by
// xten_nn.group_dequantize that uses the same shift factor and has same
Expand Down
11 changes: 7 additions & 4 deletions lib/Dialect/XTenNN/Transforms/QDQConcat.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,17 +36,20 @@ struct RemoveQDQBetweenConcat : public OpRewritePattern<DequantizeOp> {
PatternRewriter &rewriter) const override {
// Match concat->QDQ->concat and remove QDQ, if concats would be foldable.
// Removing a QDQ is already destructive. Try to be a little-less
// destructive by checking that the QDQ nodes have the same shift.
// destructive by checking that the QDQ nodes have the same scale +
// zeropoint.
auto quantize =
llvm::dyn_cast_or_null<QuantizeOp>(op.getInput().getDefiningOp());
if (!quantize) {
return rewriter.notifyMatchFailure(
op, "DequantizeOp input not produced by QuantizeOp.");
}

if (quantize.getShift() != op.getShift()) {
return rewriter.notifyMatchFailure(
op, "DequantizeOp and QuantizeOp do not share the same shift value.");
if (quantize.getScale() != op.getScale() ||
quantize.getZeroPoint() != op.getZeroPoint()) {
return rewriter.notifyMatchFailure(op,
"DequantizeOp and QuantizeOp do not "
"share the same scale + zero_point.");
}

// Try to match an incoming concat
Expand Down
Loading

0 comments on commit 9728c35

Please sign in to comment.