Skip to content

Commit 129a1c6

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 e7bf54d commit 129a1c6

File tree

15 files changed

+560
-481
lines changed

15 files changed

+560
-481
lines changed

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

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

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

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

Lines changed: 40 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -105,8 +105,9 @@ def Tosa_Conv2DOp : Tosa_ConvOp<"conv2d"> {
105105
Tosa_Tensor4D:$input,
106106
TosaTensorRankOf<[Tosa_Weight], [4]>:$weight,
107107
Tosa_Tensor1D:$bias,
108-
Optional<Tosa_ScalarTensor>:$input_zp,
109-
Optional<Tosa_ScalarTensor>:$weight_zp,
108+
Tosa_ScalarTensor:$input_zp,
109+
Tosa_ScalarTensor:$weight_zp,
110+
110111
Tosa_IntArrayAttr4:$pad,
111112
Tosa_IntArrayAttr2:$stride,
112113
Tosa_IntArrayAttr2:$dilation,
@@ -118,6 +119,13 @@ def Tosa_Conv2DOp : Tosa_ConvOp<"conv2d"> {
118119
Tosa_Tensor4D:$output
119120
);
120121

122+
let extraClassDeclaration = [{
123+
LogicalResult getInputZeroPoint(int64_t &zp);
124+
LogicalResult getWeightZeroPoint(int64_t &zp);
125+
LogicalResult verifyInputZeroPoint(int64_t zp);
126+
LogicalResult verifyWeightZeroPoint(int64_t zp);
127+
}];
128+
121129
let builders = [Tosa_ConvOpQuantInfoBuilder];
122130
let hasVerifier = 1;
123131
}
@@ -136,8 +144,9 @@ def Tosa_Conv3DOp : Tosa_ConvOp<"conv3d"> {
136144
Tosa_Tensor5D:$input,
137145
TosaTensorRankOf<[Tosa_Weight], [5]>:$weight,
138146
Tosa_Tensor1D:$bias,
139-
Optional<Tosa_ScalarTensor>:$input_zp,
140-
Optional<Tosa_ScalarTensor>:$weight_zp,
147+
Tosa_ScalarTensor:$input_zp,
148+
Tosa_ScalarTensor:$weight_zp,
149+
141150
Tosa_IntArrayAttr6:$pad,
142151
Tosa_IntArrayAttr3:$stride,
143152
Tosa_IntArrayAttr3:$dilation,
@@ -149,6 +158,13 @@ def Tosa_Conv3DOp : Tosa_ConvOp<"conv3d"> {
149158
Tosa_Tensor5D:$output
150159
);
151160

161+
let extraClassDeclaration = [{
162+
LogicalResult getInputZeroPoint(int64_t &zp);
163+
LogicalResult getWeightZeroPoint(int64_t &zp);
164+
LogicalResult verifyInputZeroPoint(int64_t zp);
165+
LogicalResult verifyWeightZeroPoint(int64_t zp);
166+
}];
167+
152168
let builders = [Tosa_ConvOpQuantInfoBuilder];
153169
let hasVerifier = 1;
154170
}
@@ -168,8 +184,9 @@ def Tosa_DepthwiseConv2DOp : Tosa_ConvOp<"depthwise_conv2d"> {
168184
Tosa_Tensor4D:$input,
169185
TosaTensorRankOf<[Tosa_Weight], [4]>:$weight,
170186
Tosa_Tensor1D:$bias,
171-
Optional<Tosa_ScalarTensor>:$input_zp,
172-
Optional<Tosa_ScalarTensor>:$weight_zp,
187+
Tosa_ScalarTensor:$input_zp,
188+
Tosa_ScalarTensor:$weight_zp,
189+
173190
Tosa_IntArrayAttr4:$pad,
174191
Tosa_IntArrayAttr2:$stride,
175192
Tosa_IntArrayAttr2:$dilation,
@@ -181,6 +198,13 @@ def Tosa_DepthwiseConv2DOp : Tosa_ConvOp<"depthwise_conv2d"> {
181198
Tosa_Tensor4D:$output
182199
);
183200

201+
let extraClassDeclaration = [{
202+
LogicalResult getInputZeroPoint(int64_t &zp);
203+
LogicalResult getWeightZeroPoint(int64_t &zp);
204+
LogicalResult verifyInputZeroPoint(int64_t zp);
205+
LogicalResult verifyWeightZeroPoint(int64_t zp);
206+
}];
207+
184208
let builders = [Tosa_ConvOpQuantInfoBuilder];
185209
let hasVerifier = 1;
186210
}
@@ -330,8 +354,9 @@ def Tosa_TransposeConv2DOp : Tosa_ConvOp<"transpose_conv2d"> {
330354
Tosa_Tensor4D:$input,
331355
TosaTensorRankOf<[Tosa_Weight], [4]>:$weight,
332356
Tosa_Tensor1D:$bias,
333-
Optional<Tosa_ScalarTensor>:$input_zp,
334-
Optional<Tosa_ScalarTensor>:$weight_zp,
357+
Tosa_ScalarTensor:$input_zp,
358+
Tosa_ScalarTensor:$weight_zp,
359+
335360
Tosa_IntArrayAttr4:$out_pad,
336361
Tosa_IntArrayAttr2:$stride,
337362
Tosa_IntArrayAttr4:$out_shape,
@@ -343,6 +368,13 @@ def Tosa_TransposeConv2DOp : Tosa_ConvOp<"transpose_conv2d"> {
343368
Tosa_Tensor4D:$output
344369
);
345370

371+
let extraClassDeclaration = [{
372+
LogicalResult getInputZeroPoint(int64_t &zp);
373+
LogicalResult getWeightZeroPoint(int64_t &zp);
374+
LogicalResult verifyInputZeroPoint(int64_t zp);
375+
LogicalResult verifyWeightZeroPoint(int64_t zp);
376+
}];
377+
346378
let builders = [Tosa_TransConvOpQuantInfoBuilder];
347379
let hasVerifier = 1;
348380
}

‎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)