Skip to content

[IE] UnrollGroupQuantize: make unrolled slice locations unique per consumer #302

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

Open
blairducrayoppat wants to merge 1 commit into
base: develop
Choose a base branch
from
Open

[IE] UnrollGroupQuantize: make unrolled slice locations unique per consumer #302

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
42 changes: 27 additions & 15 deletions src/vpux_compiler/src/dialect/IE/transforms/passes/unroll_group_quantize.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
//
// Copyright (C) 2023-2025 Intel Corporation
// Copyright (C) 2023-2026 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//

Expand Down Expand Up @@ -28,8 +28,8 @@ class GenericUnrollBase : public mlir::OpRewritePattern<ConcreteOp> {
public:
mlir::LogicalResult matchAndRewrite(ConcreteOp origOp, mlir::PatternRewriter& rewriter) const final;
mlir::Value getValue(const mlir::ValueRange values, const int64_t idx) const;
SmallVector<mlir::Value> splitValue(const mlir::Value val, const int64_t axis,
mlir::PatternRewriter& rewriter) const;
SmallVector<mlir::Value> splitValue(const mlir::Value val, const int64_t axis, mlir::PatternRewriter& rewriter,
mlir::Operation* consumerOp, StringRef operandTag) const;

private:
virtual SmallVector<mlir::Value> splitInputs(ConcreteOp origOp, const int64_t axis,
Expand Down Expand Up @@ -60,9 +60,17 @@ mlir::Value GenericUnrollBase<ConcreteOp>::getValue(const mlir::ValueRange value
// 1x3x1x16 with axis 1 will be split in three 1x1x1x16 values
// 1x1x2x16 with axis 2 will be split in two 1x1x1x16 values
// 1x3x1x16 with axis 2 won't be split and will return a vector with only one 1x3x1x16 element
//
// The location of each produced slice is derived from the consuming op being unrolled (consumerOp) plus
// operandTag (which input is being split) and the chunk index. Deriving it from the consumer rather than
// from val.getLoc() keeps the slice locations unique when the same value is shared by more than one
// unrolled consumer - e.g. a grouped-quantization weight scale that is the dequantization scale of two
// DynamicDequantize ops. Naming every slice "<val>/slice_{idx}" made the low indices collide across
// consumers and tripped StopLocationVerifierPass ("Found N duplicated names after full verification").
template <typename ConcreteOp>
SmallVector<mlir::Value> GenericUnrollBase<ConcreteOp>::splitValue(const mlir::Value val, const int64_t axis,
mlir::PatternRewriter& rewriter) const {
mlir::PatternRewriter& rewriter,
mlir::Operation* consumerOp, StringRef operandTag) const {
const auto valShape = getShape(val);
VPUX_THROW_UNLESS(axis < checked_cast<int64_t>(valShape.size()), "Cannot split shape {0} by axis {1}", valShape,
axis);
Expand All @@ -75,7 +83,7 @@ SmallVector<mlir::Value> GenericUnrollBase<ConcreteOp>::splitValue(const mlir::V
const auto staticSizesAttr = getIntArrayAttr(rewriter.getContext(), staticSizes);
SmallVector<mlir::Value> inputChunks;
for (const auto& idx : irange(groups)) {
const auto loc = appendLoc(val.getLoc(), "slice_{0}", idx);
const auto loc = takeOpLoc(consumerOp, "{0}_slice_{1}", operandTag, idx);
SmallVector<int64_t> offsets(valShape.size(), 0);
offsets[axis] = idx;
const auto offsetsAttr = getIntArrayAttr(rewriter.getContext(), offsets);
Expand Down Expand Up @@ -148,15 +156,17 @@ class UnrollFakeQuantize final : public GenericUnrollBase<IE::FakeQuantizeOp> {
// IE.FakeQuantize with data = 1x1x8x16, in_low = in_high = 1x1x1x1, out_low = out_high = 1x1x1x16
SmallVector<mlir::Value> UnrollFakeQuantize::splitInputs(IE::FakeQuantizeOp fqOp, const int64_t axis,
mlir::PatternRewriter& rewriter) const {
const auto data = splitValue(fqOp.getInput(), axis, rewriter);
const auto inLow = splitValue(fqOp.getInputLow(), axis, rewriter);
const auto inHigh = splitValue(fqOp.getInputHigh(), axis, rewriter);
const auto outLow = splitValue(fqOp.getOutputLow(), axis, rewriter);
const auto outHigh = splitValue(fqOp.getOutputHigh(), axis, rewriter);
const auto data = splitValue(fqOp.getInput(), axis, rewriter, fqOp, "data");
const auto inLow = splitValue(fqOp.getInputLow(), axis, rewriter, fqOp, "in_low");
const auto inHigh = splitValue(fqOp.getInputHigh(), axis, rewriter, fqOp, "in_high");
const auto outLow = splitValue(fqOp.getOutputLow(), axis, rewriter, fqOp, "out_low");
const auto outHigh = splitValue(fqOp.getOutputHigh(), axis, rewriter, fqOp, "out_high");
SmallVector<mlir::Value> fqResults;
const auto groups = data.size();
for (const auto& idx : irange(groups)) {
const auto loc = appendLoc(fqOp.getLoc(), "slice_{0}", idx);
// One reduced FakeQuantize per chunk on the consumer's own location - unique by construction
// (unlike the shared-input operand slices above, which carry an operand tag to stay unique).
const auto loc = takeOpLoc(fqOp, "slice_{0}", idx);
auto reducedFq = rewriter.create<IE::FakeQuantizeOp>(
loc, data[idx], getValue(inLow, idx), getValue(inHigh, idx), getValue(outLow, idx),
getValue(outHigh, idx), fqOp.getLevelsAttr(), fqOp.getLowFpTypeAttr(), fqOp.getAutoBroadcast());
Expand All @@ -183,18 +193,20 @@ class UnrollDynamicDequantize final : public GenericUnrollBase<IE::DynamicDequan

SmallVector<mlir::Value> UnrollDynamicDequantize::splitInputs(IE::DynamicDequantizeOp origOp, const int64_t axis,
mlir::PatternRewriter& rewriter) const {
const auto input = splitValue(origOp.getInput(), axis, rewriter);
const auto scale = splitValue(origOp.getScale(), axis, rewriter);
const auto input = splitValue(origOp.getInput(), axis, rewriter, origOp, "input");
const auto scale = splitValue(origOp.getScale(), axis, rewriter, origOp, "scale");
bool hasZeroPoint = false;
SmallVector<mlir::Value> zeroPoint;
if (origOp.getZp() != nullptr) {
hasZeroPoint = true;
zeroPoint = splitValue(origOp.getZp(), axis, rewriter);
zeroPoint = splitValue(origOp.getZp(), axis, rewriter, origOp, "zp");
}
SmallVector<mlir::Value> deQuantizeResults;
const auto groups = input.size();
for (const auto& idx : irange(groups)) {
const auto loc = appendLoc(origOp.getLoc(), "slice_{0}", idx);
// One reduced DynamicDequantize per chunk on the consumer's own location - unique by construction
// (unlike the shared-input operand slices above, which carry an operand tag to stay unique).
const auto loc = takeOpLoc(origOp, "slice_{0}", idx);
auto reduceDequantize = rewriter.create<IE::DynamicDequantizeOp>(
loc, getValue(input, idx), getValue(scale, idx), hasZeroPoint ? getValue(zeroPoint, idx) : nullptr,
origOp.getDstElemType());
Expand Down
65 changes: 65 additions & 0 deletions tests/lit/NPU/dialect/IE/passes/unroll_group_quantize_shared_param_locations.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
//
// Copyright (C) 2026 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//

// RUN: vpux-opt --init-compiler="platform=%platform%" --unroll-group-quantize --mlir-print-debuginfo %s | FileCheck %s
// REQUIRES: platform-NPU3720 || platform-NPU4000 || platform-NPU5010

// Regression test for a location-uniqueness violation in UnrollGroupQuantize.
//
// One quantization scale (%SCALE) is the dequantization scale of TWO IE.DynamicDequantize ops - as happens
// in grouped-INT4 attention projections, where the per-group scale feeds both the weight dequant and the
// activation/matmul dequant. UnrollGroupQuantize unrolls each consumer and slices the shared scale in both.
// Before the fix every slice was located after the shared scale value (".../slice_{idx}"), so the two
// consumers produced byte-identical slice locations and StopLocationVerifierPass aborted the compile with
// "Found N duplicated names after full verification". The fix locates each slice after its CONSUMING op plus
// an operand tag, so the slices stay unique.
//
// Note: this is a minimal reduction - both consumers here unroll the shared scale along the same axis,
// whereas the field case additionally unrolls them along different axes. The violation (slices of a shared
// value named after that value) is identical either way; the axis count is incidental. The consumer ops
// carry metadata-bearing fused locations to match importer-produced IR and exercise appendLoc's
// preserve-existing-FusedLoc branch.

!qElemType = !quant.uniform<u4:f16, 1.000000e+00>

// CHECK-LABEL: @SharedScaleTwoConsumers
func.func @SharedScaleTwoConsumers(%data0: tensor<2x1536x128x!qElemType>, %data1: tensor<2x1536x128x!qElemType>,
%scale: tensor<2x1536x1xf16>, %act0: tensor<1x256xf16>, %act1: tensor<1x256xf16>)
-> (tensor<1x1536xf16>, tensor<1x1536xf16>) {
%0 = IE.DynamicDequantize(%data0, %scale) {dstElemType = f16} : tensor<2x1536x128x!qElemType>, tensor<2x1536x1xf16> -> tensor<2x1536x128xf16> loc(fused<{name = "weight_dq", type = "DynamicDequantize"}>["weight_dq"])
%1 = IE.Transpose(%0) {order_value = affine_map<(d0, d1, d2) -> (d1, d0, d2)>} : tensor<2x1536x128xf16> -> tensor<1536x2x128xf16>
%2 = IE.AffineReshape(%1) {dim_mapping = [[0], [1], [1]], shape_value = [1536, 256]} : tensor<1536x2x128xf16> -> tensor<1536x256xf16>
%3 = IE.FullyConnected(%act0, %2) : tensor<1x256xf16>, tensor<1536x256xf16> -> tensor<1x1536xf16>

%4 = IE.DynamicDequantize(%data1, %scale) {dstElemType = f16} : tensor<2x1536x128x!qElemType>, tensor<2x1536x1xf16> -> tensor<2x1536x128xf16> loc(fused<{name = "matmul_dq", type = "DynamicDequantize"}>["matmul_dq"])
%5 = IE.Transpose(%4) {order_value = affine_map<(d0, d1, d2) -> (d1, d0, d2)>} : tensor<2x1536x128xf16> -> tensor<1536x2x128xf16>
%6 = IE.AffineReshape(%5) {dim_mapping = [[0], [1], [1]], shape_value = [1536, 256]} : tensor<1536x2x128xf16> -> tensor<1536x256xf16>
%7 = IE.FullyConnected(%act1, %6) : tensor<1x256xf16>, tensor<1536x256xf16> -> tensor<1x1536xf16>

return %3, %7 : tensor<1x1536xf16>, tensor<1x1536xf16>

// Bind each shared-scale (and data) slice to its IE.Slice op, then assert the captured location resolves
// to a fused loc rooted at the CONSUMING op + an operand tag. The two consumers' "scale_slice_0" (and
// "input_slice_0") therefore resolve to distinct locations (weight_dq vs matmul_dq roots), which is the
// invariant the fix restores. Pre-fix the suffix is the un-tagged "slice_0" rooted at the shared scale,
// so the "scale_slice_0"/"input_slice_0" alias definitions below are absent and the test fails.

// Consumer A ("weight_dq"): its data slice then its shared-scale slice.
// CHECK: IE.Slice %arg0 [0, 0, 0] [1, 1536, 128] {{.*}} loc([[A_INPUT0:#.+]])
// CHECK: IE.Slice %arg2 [0, 0, 0] [1, 1536, 1] {{.*}} loc([[A_SCALE0:#.+]])
// Consumer B ("matmul_dq"): its data slice then its shared-scale slice.
// CHECK: IE.Slice %arg1 [0, 0, 0] [1, 1536, 128] {{.*}} loc([[B_INPUT0:#.+]])
// CHECK: IE.Slice %arg2 [0, 0, 0] [1, 1536, 1] {{.*}} loc([[B_SCALE0:#.+]])

// CHECK-DAG: [[WEIGHT_DQ:#.+]] = loc("weight_dq")
// CHECK-DAG: [[MATMUL_DQ:#.+]] = loc("matmul_dq")
// CHECK-DAG: [[INPUT_SLICE0:#.+]] = loc("input_slice_0")
// CHECK-DAG: [[SCALE_SLICE0:#.+]] = loc("scale_slice_0")

// CHECK-DAG: [[A_INPUT0]] = loc(fused<{{.*}}>{{\[}}[[WEIGHT_DQ]], [[INPUT_SLICE0]]])
// CHECK-DAG: [[A_SCALE0]] = loc(fused<{{.*}}>{{\[}}[[WEIGHT_DQ]], [[SCALE_SLICE0]]])
// CHECK-DAG: [[B_INPUT0]] = loc(fused<{{.*}}>{{\[}}[[MATMUL_DQ]], [[INPUT_SLICE0]]])
// CHECK-DAG: [[B_SCALE0]] = loc(fused<{{.*}}>{{\[}}[[MATMUL_DQ]], [[SCALE_SLICE0]]])
}