Skip to content

Commit 214ea7b

Browse files
Tai78641Tai78641
committed
[mlir][tosa] Change zero points of convolution ops to required inputs
This patch changes the input_zp and weight_zp for convolution operators to be required inputs. Convolution operators affected are: CONV2D, CONV3D, DEPTHWISE_CONV2D, and TRANSPOSE_CONV2D. Signed-off-by: Tai Ly <[email protected]> Change-Id: I7aa6e05580c83c617394c3f7fd18fba8c3f3b0ec
1 parent a1a1073 commit 214ea7b

18 files changed

+580
-493
lines changed

‎mlir/include/mlir/Dialect/Tosa/IR/TosaOps.h

Lines changed: 0 additions & 106 deletions
Original file line numberDiff line numberDiff line change
@@ -168,112 +168,6 @@ namespace tosa {
168168
std::optional<Value> createZeroPointTensor(OpBuilder &builder, Location loc,
169169
Type srcElemType, int64_t zp = 0);
170170

171-
// Get zero point value from the attribute argument.
172-
LogicalResult getZeroPoint(ElementsAttr zpAttr, int64_t &zp);
173-
174-
// Verify if zero point falls into valid range.
175-
template <typename T>
176-
LogicalResult verifyZeroPoint(Type zpElemType, int64_t zp) {
177-
if constexpr (!std::is_same_v<T, Conv2DOp> && !std::is_same_v<T, Conv3DOp> &&
178-
!std::is_same_v<T, DepthwiseConv2DOp> &&
179-
!std::is_same_v<T, TransposeConv2DOp>) {
180-
return failure();
181-
}
182-
183-
if (!zpElemType.isIntOrFloat())
184-
return failure();
185-
186-
if (!zpElemType.isInteger(8) && zp != 0)
187-
return failure();
188-
189-
if (zpElemType.isSignedInteger(8) && (zp < -128 || zp > 127))
190-
return failure();
191-
192-
if (zpElemType.isUnsignedInteger(8) && (zp < 0 || zp > 255))
193-
return failure();
194-
195-
return success();
196-
}
197-
198-
// Helper type trait to determine if an operation is a tosa convolution.
199-
template <typename Op>
200-
struct IsTosaConv : std::false_type {};
201-
202-
template <>
203-
struct IsTosaConv<tosa::Conv2DOp> : std::true_type {};
204-
template <>
205-
struct IsTosaConv<tosa::DepthwiseConv2DOp> : std::true_type {};
206-
template <>
207-
struct IsTosaConv<tosa::TransposeConv2DOp> : std::true_type {};
208-
template <>
209-
struct IsTosaConv<tosa::Conv3DOp> : std::true_type {};
210-
211-
template <typename Op>
212-
constexpr bool is_tosa_conv_v = IsTosaConv<Op>::value;
213-
214-
// Helper struct to hold the zero points of a TOSA convolution operation as
215-
// named 64-bit integer fields.
216-
struct ConvZpPair {
217-
ConvZpPair(std::int64_t inputZp, std::int64_t weightZp)
218-
: inputZp(inputZp), weightZp(weightZp) {}
219-
std::int64_t inputZp;
220-
std::int64_t weightZp;
221-
};
222-
223-
// Helper function which attempts to extract the zero points from a TOSA
224-
// convolution by matching them against defining ops which should be tosa.const
225-
// operations.
226-
//
227-
// There are three possible results:
228-
// 1. Failed to extract the zero-points i.e. they should exist and don't or they
229-
// do exist but are invalid.
230-
// 2. Succeeded in extracting zero-points.
231-
// 3. Zero points are "empty" and meaningless for this op i.e. non-quantized
232-
// convolution.
233-
using FailOrMaybeZP = llvm::FailureOr<std::optional<ConvZpPair>>;
234-
template <typename TosaConvOp>
235-
std::enable_if_t<is_tosa_conv_v<TosaConvOp>, FailOrMaybeZP>
236-
extractConvZpPair(TosaConvOp op, PatternRewriter &rewriter) {
237-
// Strictly speaking the base TOSA spec requires that for non int8 types
238-
// zero points must be zero. However, in the dialect these operands are
239-
// optional and only required for int8. They have no semantic meaning for
240-
// non-quantized types and can therefore be safely ignored. This is case 3.
241-
if (auto opElementTY =
242-
cast<ShapedType>(op->getOperand(0).getType()).getElementType();
243-
!opElementTY.isInteger(8))
244-
return FailOrMaybeZP(std::nullopt);
245-
246-
// Now we know we should have a zero point check it is valid.
247-
if (!op.getInputZp())
248-
return rewriter.notifyMatchFailure(op, "missing input zero point");
249-
250-
// Helper to extract the zero point by matching its definition against a
251-
// constant.
252-
auto extractZeroPoint = [](Value zpValue) -> std::optional<int64_t> {
253-
ElementsAttr zpAttr;
254-
if (!matchPattern(zpValue, m_Constant(&zpAttr)))
255-
return std::nullopt;
256-
257-
int64_t zp;
258-
if (tosa::getZeroPoint(zpAttr, zp).failed())
259-
return std::nullopt;
260-
261-
return std::make_optional(zp);
262-
};
263-
264-
auto maybeInputZp = extractZeroPoint(op.getInputZp());
265-
if (!maybeInputZp)
266-
return rewriter.notifyMatchFailure(op, "unable to extract input zp");
267-
268-
if (!op.getWeightZp())
269-
return rewriter.notifyMatchFailure(op, "missing weight zero point");
270-
271-
auto maybeWeightZp = extractZeroPoint(op.getWeightZp());
272-
if (!maybeWeightZp)
273-
return rewriter.notifyMatchFailure(op, "unable to extract weight zp");
274-
275-
return std::make_optional<ConvZpPair>(*maybeInputZp, *maybeWeightZp);
276-
}
277171
} // namespace tosa
278172
} // namespace mlir
279173

‎mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td

Lines changed: 40 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -115,8 +115,9 @@ def Tosa_Conv2DOp : Tosa_ConvOp<"conv2d"> {
115115
Tosa_Tensor4D:$input,
116116
TosaTensorRankOf<[Tosa_Weight], [4]>:$weight,
117117
Tosa_Tensor1D:$bias,
118-
Optional<Tosa_ScalarTensor>:$input_zp,
119-
Optional<Tosa_ScalarTensor>:$weight_zp,
118+
Tosa_ScalarTensor:$input_zp,
119+
Tosa_ScalarTensor:$weight_zp,
120+
120121
Tosa_IntArrayAttr4:$pad,
121122
Tosa_IntArrayAttr2:$stride,
122123
Tosa_IntArrayAttr2:$dilation,
@@ -133,6 +134,13 @@ def Tosa_Conv2DOp : Tosa_ConvOp<"conv2d"> {
133134
Extension<[Tosa_EXT_INT4, Tosa_EXT_INT16, Tosa_EXT_FP8E4M3, Tosa_EXT_FP8E5M2, Tosa_EXT_BF16]>,
134135
];
135136

137+
let extraClassDeclaration = [{
138+
LogicalResult getInputZeroPoint(int64_t &zp);
139+
LogicalResult getWeightZeroPoint(int64_t &zp);
140+
LogicalResult verifyInputZeroPoint(int64_t zp);
141+
LogicalResult verifyWeightZeroPoint(int64_t zp);
142+
}];
143+
136144
let builders = [Tosa_ConvOpQuantInfoBuilder];
137145
let hasVerifier = 1;
138146
}
@@ -151,8 +159,9 @@ def Tosa_Conv3DOp : Tosa_ConvOp<"conv3d"> {
151159
Tosa_Tensor5D:$input,
152160
TosaTensorRankOf<[Tosa_Weight], [5]>:$weight,
153161
Tosa_Tensor1D:$bias,
154-
Optional<Tosa_ScalarTensor>:$input_zp,
155-
Optional<Tosa_ScalarTensor>:$weight_zp,
162+
Tosa_ScalarTensor:$input_zp,
163+
Tosa_ScalarTensor:$weight_zp,
164+
156165
Tosa_IntArrayAttr6:$pad,
157166
Tosa_IntArrayAttr3:$stride,
158167
Tosa_IntArrayAttr3:$dilation,
@@ -169,6 +178,13 @@ def Tosa_Conv3DOp : Tosa_ConvOp<"conv3d"> {
169178
Extension<[Tosa_EXT_INT4, Tosa_EXT_INT16, Tosa_EXT_FP8E4M3, Tosa_EXT_FP8E5M2, Tosa_EXT_BF16]>,
170179
];
171180

181+
let extraClassDeclaration = [{
182+
LogicalResult getInputZeroPoint(int64_t &zp);
183+
LogicalResult getWeightZeroPoint(int64_t &zp);
184+
LogicalResult verifyInputZeroPoint(int64_t zp);
185+
LogicalResult verifyWeightZeroPoint(int64_t zp);
186+
}];
187+
172188
let builders = [Tosa_ConvOpQuantInfoBuilder];
173189
let hasVerifier = 1;
174190
}
@@ -188,8 +204,9 @@ def Tosa_DepthwiseConv2DOp : Tosa_ConvOp<"depthwise_conv2d"> {
188204
Tosa_Tensor4D:$input,
189205
TosaTensorRankOf<[Tosa_Weight], [4]>:$weight,
190206
Tosa_Tensor1D:$bias,
191-
Optional<Tosa_ScalarTensor>:$input_zp,
192-
Optional<Tosa_ScalarTensor>:$weight_zp,
207+
Tosa_ScalarTensor:$input_zp,
208+
Tosa_ScalarTensor:$weight_zp,
209+
193210
Tosa_IntArrayAttr4:$pad,
194211
Tosa_IntArrayAttr2:$stride,
195212
Tosa_IntArrayAttr2:$dilation,
@@ -206,6 +223,13 @@ def Tosa_DepthwiseConv2DOp : Tosa_ConvOp<"depthwise_conv2d"> {
206223
Extension<[Tosa_EXT_INT4, Tosa_EXT_INT16, Tosa_EXT_FP8E4M3, Tosa_EXT_FP8E5M2, Tosa_EXT_BF16]>,
207224
];
208225

226+
let extraClassDeclaration = [{
227+
LogicalResult getInputZeroPoint(int64_t &zp);
228+
LogicalResult getWeightZeroPoint(int64_t &zp);
229+
LogicalResult verifyInputZeroPoint(int64_t zp);
230+
LogicalResult verifyWeightZeroPoint(int64_t zp);
231+
}];
232+
209233
let builders = [Tosa_ConvOpQuantInfoBuilder];
210234
let hasVerifier = 1;
211235
}
@@ -375,8 +399,9 @@ def Tosa_TransposeConv2DOp : Tosa_ConvOp<"transpose_conv2d"> {
375399
Tosa_Tensor4D:$input,
376400
TosaTensorRankOf<[Tosa_Weight], [4]>:$weight,
377401
Tosa_Tensor1D:$bias,
378-
Optional<Tosa_ScalarTensor>:$input_zp,
379-
Optional<Tosa_ScalarTensor>:$weight_zp,
402+
Tosa_ScalarTensor:$input_zp,
403+
Tosa_ScalarTensor:$weight_zp,
404+
380405
Tosa_IntArrayAttr4:$out_pad,
381406
Tosa_IntArrayAttr2:$stride,
382407
Tosa_IntArrayAttr4:$out_shape,
@@ -393,6 +418,13 @@ def Tosa_TransposeConv2DOp : Tosa_ConvOp<"transpose_conv2d"> {
393418
Extension<[Tosa_EXT_INT4, Tosa_EXT_INT16, Tosa_EXT_FP8E4M3, Tosa_EXT_FP8E5M2, Tosa_EXT_BF16]>,
394419
];
395420

421+
let extraClassDeclaration = [{
422+
LogicalResult getInputZeroPoint(int64_t &zp);
423+
LogicalResult getWeightZeroPoint(int64_t &zp);
424+
LogicalResult verifyInputZeroPoint(int64_t zp);
425+
LogicalResult verifyWeightZeroPoint(int64_t zp);
426+
}];
427+
396428
let builders = [Tosa_TransConvOpQuantInfoBuilder];
397429
let hasVerifier = 1;
398430
}

‎mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp

Lines changed: 41 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -259,11 +259,21 @@ class ConvConverter : public OpConversionPattern<TosaConvOp> {
259259
DenseI64ArrayAttr strideTosaAttr = op.getStrideAttr();
260260
DenseI64ArrayAttr dilationTosaAttr = op.getDilationAttr();
261261

262-
auto failureOrMaybeZps = extractConvZpPair(op, rewriter);
263-
if (llvm::failed(failureOrMaybeZps))
264-
return failure();
262+
// Get and verify zero points.
263+
int64_t inputZpVal;
264+
int64_t weightZpVal;
265+
266+
if (op.getInputZeroPoint(inputZpVal).failed() ||
267+
op.getWeightZeroPoint(weightZpVal).failed())
268+
return rewriter.notifyMatchFailure(
269+
op, "bail out if zero points cannot statically be determined");
270+
271+
if (op.verifyInputZeroPoint(inputZpVal).failed() ||
272+
op.verifyWeightZeroPoint(weightZpVal).failed())
273+
return rewriter.notifyMatchFailure(
274+
op, "zero point must be zero for non-int8 integer types");
265275

266-
auto maybeZps = failureOrMaybeZps.value();
276+
bool hasZp = (inputZpVal != 0) || (weightZpVal != 0);
267277

268278
if (!weightTy.hasStaticShape() || !biasTy.hasStaticShape())
269279
return rewriter.notifyMatchFailure(
@@ -289,19 +299,19 @@ class ConvConverter : public OpConversionPattern<TosaConvOp> {
289299

290300
// Apply padding as necessary.
291301
TypedAttr zeroAttr = rewriter.getZeroAttr(inputETy);
292-
if (maybeZps) {
302+
if (hasZp) {
293303
int64_t intMin =
294304
APInt::getSignedMinValue(inputETy.getIntOrFloatBitWidth())
295305
.getSExtValue();
296306
int64_t intMax =
297307
APInt::getSignedMaxValue(inputETy.getIntOrFloatBitWidth())
298308
.getSExtValue();
299309

300-
if (maybeZps->inputZp < intMin || maybeZps->inputZp > intMax)
310+
if (inputZpVal < intMin || inputZpVal > intMax)
301311
return rewriter.notifyMatchFailure(
302312
op, "tosa.conv op quantization has zp outside of input range");
303313

304-
zeroAttr = rewriter.getIntegerAttr(inputETy, maybeZps->inputZp);
314+
zeroAttr = rewriter.getIntegerAttr(inputETy, inputZpVal);
305315
}
306316

307317
llvm::SmallVector<int64_t> pad;
@@ -314,8 +324,8 @@ class ConvConverter : public OpConversionPattern<TosaConvOp> {
314324
// For 2D convolutions, we need to check if the target convolution op
315325
// wants a HWCF kernel layout.
316326
bool wantHwcf =
317-
maybeZps ? std::is_same_v<LinalgConvQOp, linalg::Conv2DNhwcHwcfQOp>
318-
: std::is_same_v<LinalgConvOp, linalg::Conv2DNhwcHwcfOp>;
327+
hasZp ? std::is_same_v<LinalgConvQOp, linalg::Conv2DNhwcHwcfQOp>
328+
: std::is_same_v<LinalgConvOp, linalg::Conv2DNhwcHwcfOp>;
319329
if (wantHwcf) {
320330
// Transpose the kernel to match dimension ordering of the linalg
321331
// convolution operation.
@@ -376,9 +386,9 @@ class ConvConverter : public OpConversionPattern<TosaConvOp> {
376386
Value broadcastBias =
377387
linalgBroadcastAndMaybeExtSI(rewriter, loc, bias, biasEmptyTensor);
378388

379-
if (maybeZps) {
380-
auto iZp = rewriter.getI32IntegerAttr(maybeZps->inputZp);
381-
auto kZp = rewriter.getI32IntegerAttr(maybeZps->weightZp);
389+
if (hasZp) {
390+
auto iZp = rewriter.getI32IntegerAttr(inputZpVal);
391+
auto kZp = rewriter.getI32IntegerAttr(weightZpVal);
382392

383393
auto iZpVal = rewriter.create<arith::ConstantOp>(loc, iZp);
384394
auto kZpVal = rewriter.create<arith::ConstantOp>(loc, kZp);
@@ -441,31 +451,40 @@ class DepthwiseConvConverter
441451
/*inputSizeDims=*/{1, 2},
442452
/*kernelSizeDims=*/{0, 1}, rewriter);
443453

444-
auto failureOrMaybeZps = extractConvZpPair(op, rewriter);
445-
if (llvm::failed(failureOrMaybeZps))
446-
return failure();
454+
// Get and verify zero points.
455+
int64_t inputZpVal;
456+
int64_t weightZpVal;
457+
458+
if (op.getInputZeroPoint(inputZpVal).failed() ||
459+
op.getWeightZeroPoint(weightZpVal).failed())
460+
return rewriter.notifyMatchFailure(
461+
op, "bail out if zero points cannot statically be determined");
447462

448-
auto maybeZps = failureOrMaybeZps.value();
463+
if (op.verifyInputZeroPoint(inputZpVal).failed() ||
464+
op.verifyWeightZeroPoint(weightZpVal).failed())
465+
return rewriter.notifyMatchFailure(
466+
op, "zero point must be zero for non-int8 integer types");
449467

468+
bool hasZp = (inputZpVal != 0) || (weightZpVal != 0);
450469
auto weightShape = weightTy.getShape();
451470
auto resultShape = resultTy.getShape();
452471

453472
// Apply padding as necessary.
454473
TypedAttr zeroAttr = rewriter.getZeroAttr(inputETy);
455-
if (maybeZps) {
474+
if (hasZp) {
456475
int64_t intMin =
457476
APInt::getSignedMinValue(inputETy.getIntOrFloatBitWidth())
458477
.getSExtValue();
459478
int64_t intMax =
460479
APInt::getSignedMaxValue(inputETy.getIntOrFloatBitWidth())
461480
.getSExtValue();
462481

463-
if (maybeZps->inputZp < intMin || maybeZps->inputZp > intMax)
482+
if (inputZpVal < intMin || inputZpVal > intMax)
464483
return rewriter.notifyMatchFailure(
465484
op, "tosa.depthwise_conv op quantization has zp outside of input "
466485
"range");
467486

468-
zeroAttr = rewriter.getIntegerAttr(inputETy, maybeZps->inputZp);
487+
zeroAttr = rewriter.getIntegerAttr(inputETy, inputZpVal);
469488
}
470489

471490
llvm::SmallVector<int64_t> pad;
@@ -505,7 +524,7 @@ class DepthwiseConvConverter
505524
indexingMaps.push_back(rewriter.getMultiDimIdentityMap(resultRank));
506525
indexingMaps.push_back(rewriter.getMultiDimIdentityMap(resultRank));
507526

508-
if (!maybeZps) {
527+
if (!hasZp) {
509528
Value conv = rewriter
510529
.create<linalg::DepthwiseConv2DNhwcHwcmOp>(
511530
loc, linalgConvTy, ValueRange{input, weight},
@@ -532,8 +551,8 @@ class DepthwiseConvConverter
532551
.getResult(0);
533552
rewriter.replaceOp(op, result);
534553
} else {
535-
IntegerAttr iZp = rewriter.getI32IntegerAttr(maybeZps->inputZp);
536-
IntegerAttr wZp = rewriter.getI32IntegerAttr(maybeZps->weightZp);
554+
IntegerAttr iZp = rewriter.getI32IntegerAttr(inputZpVal);
555+
IntegerAttr wZp = rewriter.getI32IntegerAttr(weightZpVal);
537556
auto iZpVal = rewriter.create<arith::ConstantOp>(loc, iZp);
538557
auto kZpVal = rewriter.create<arith::ConstantOp>(loc, wZp);
539558
Value conv =

0 commit comments

Comments
 (0)