[GlobalOptimization] Propagate affine QDQ through tensor operations (#24906) The integer contraction rewrite from #24902 needs dequantize_affine to be the direct producer of a contraction operand. This PR adds PropagateAffineQuantization, which moves the tensor ops in between onto the quantized side, where they also operate on narrower data. It runs right before each integer QDQ rewrite. It perform the following transformations: ``` quantize(collapse_shape(x)) -> collapse_shape(quantize(x)) expand_shape(dequantize(q)) -> dequantize(expand_shape(q)) pad(dequantize(q), 0.0) -> dequantize(pad(q, zp)) transpose(dequantize(q)) -> dequantize(q) with its output map permuted ``` Part of #24862 Assisted By: Claude Code --------- Signed-off-by: Thomas Ziereis <ziereis@roofline.ai> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/ReshapeFusion.cpp b/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/ReshapeFusion.cpp index 4801a5f..0631b0a 100644 --- a/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/ReshapeFusion.cpp +++ b/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/ReshapeFusion.cpp
@@ -354,39 +354,46 @@ return success(); } -template <typename AttentionOpTy> -static SmallVector<ReshapeOperandInfo> -getAttentionLikeReshapeInfo(AttentionOpTy attentionOp) { - return llvm::map_to_vector( - attentionOp->getOpOperands(), [&](OpOperand &opOperand) { - ReshapeOperandInfo operandInfo; - auto operandType = dyn_cast<ShapedType>(opOperand.get().getType()); - if (!operandType) { - assert( - attentionOp.getMatchingIndexingMap(&opOperand).getNumResults() == - 0 && - "expected non-shaped type to have no results in indexing map"); - return operandInfo; - } +/// Derives the reshape info of an op that states an indexing map for every +/// operand, where each map is a projected permutation. +template <typename OpTy> +static SmallVector<ReshapeOperandInfo> getMappedOperandsReshapeInfo(OpTy op) { + return llvm::map_to_vector(op->getOpOperands(), [&](OpOperand &opOperand) { + ReshapeOperandInfo operandInfo; + auto operandType = dyn_cast<ShapedType>(opOperand.get().getType()); + if (!operandType) { + assert(op.getMatchingIndexingMap(&opOperand).getNumResults() == 0 && + "expected non-shaped type to have no results in indexing map"); + return operandInfo; + } - operandInfo.originalShape = getDimSizes(opOperand.get()); - for (auto result : - attentionOp.getMatchingIndexingMap(&opOperand).getResults()) { - operandInfo.operandToIterationSpace.push_back( - cast<AffineDimExpr>(result).getPosition()); - } - return operandInfo; - }); + operandInfo.originalShape = getDimSizes(opOperand.get()); + for (auto result : op.getMatchingIndexingMap(&opOperand).getResults()) { + operandInfo.operandToIterationSpace.push_back( + cast<AffineDimExpr>(result).getPosition()); + } + return operandInfo; + }); } static SmallVector<ReshapeOperandInfo> getReshapeInfo(LinalgExt::AttentionOp attentionOp) { - return getAttentionLikeReshapeInfo(attentionOp); + return getMappedOperandsReshapeInfo(attentionOp); } static SmallVector<ReshapeOperandInfo> getReshapeInfo(LinalgExt::OnlineAttentionOp attentionOp) { - return getAttentionLikeReshapeInfo(attentionOp); + return getMappedOperandsReshapeInfo(attentionOp); +} + +static SmallVector<ReshapeOperandInfo> +getReshapeInfo(LinalgExt::QuantizeAffineOp quantizeOp) { + return getMappedOperandsReshapeInfo(quantizeOp); +} + +static SmallVector<ReshapeOperandInfo> +getReshapeInfo(LinalgExt::DequantizeAffineOp dequantizeOp) { + return getMappedOperandsReshapeInfo(dequantizeOp); } static SmallVector<ReshapeOperandInfo> @@ -537,7 +544,9 @@ } if constexpr (std::is_same_v<OpTy, AttentionOp> || - std::is_same_v<OpTy, OnlineAttentionOp>) { + std::is_same_v<OpTy, OnlineAttentionOp> || + std::is_same_v<OpTy, QuantizeAffineOp> || + std::is_same_v<OpTy, DequantizeAffineOp>) { auto expandedOpIndexingMaps = llvm::to_vector_of<AffineMap, 6>( llvm::map_range(op.getIndexingMapsArray(), [&](AffineMap m) { return getIndexingMapInExpandedOp(rewriter, m, info); @@ -556,7 +565,7 @@ replacements.push_back(newResult); } else { replacements.push_back(tensor::CollapseShapeOp::create( - rewriter, loc, op.getResult(i).getType(), newResult, + rewriter, loc, op->getResult(i).getType(), newResult, originalReassoc)); } } @@ -1192,6 +1201,14 @@ patterns.getContext(), controlFoldingReshapes); } +void populatePropagateAffineQuantizationReshapesPatterns( + RewritePatternSet &patterns, + const linalg::ControlFusionFn &controlFoldingReshapes) { + patterns.add<FoldWithProducerReshapeByExpansion<QuantizeAffineOp>, + FoldWithConsumerReshapeByExpansion<DequantizeAffineOp>>( + patterns.getContext(), controlFoldingReshapes); +} + SmallVector<unsigned> defaultControlDropUnitDims(Operation *op) { auto fusionOp = cast<LinalgFusionOpInterface>(op); return llvm::to_vector(llvm::seq<unsigned>(0, fusionOp.getNumLoops())); @@ -1218,49 +1235,52 @@ } } -template <typename AttentionOpType> -struct DropAttentionLikeUnitDims final : OpRewritePattern<AttentionOpType> { - DropAttentionLikeUnitDims(MLIRContext *context, - linalg::ControlDropUnitDims options, - PatternBenefit benefit = 1) - : OpRewritePattern<AttentionOpType>(context, benefit), - options(std::move(options)) {} +/// Drops unit dims from an op that states an indexing map for every operand, +/// by cloning it with the reduced operands and maps that `linalg::dropUnitDims` +/// computes. Ops carrying a body additionally need their `linalg.index` uses +/// fixed up for the dropped dims. +template <typename OpTy> +struct DropMappedOperandsUnitDims final : OpRewritePattern<OpTy> { + DropMappedOperandsUnitDims(MLIRContext *context, + linalg::ControlDropUnitDims options, + PatternBenefit benefit = 1) + : OpRewritePattern<OpTy>(context, benefit), options(std::move(options)) {} - LogicalResult matchAndRewrite(AttentionOpType attentionOp, + LogicalResult matchAndRewrite(OpTy op, PatternRewriter &rewriter) const override { linalg::DroppedUnitDimsBuilder builder = - [](Location loc, OpBuilder &b, IndexingMapOpInterface op, + [](Location loc, OpBuilder &b, IndexingMapOpInterface interfaceOp, ArrayRef<Value> newOperands, ArrayRef<AffineMap> newIndexingMaps, const llvm::SmallDenseSet<unsigned> &droppedDims) -> IndexingMapOpInterface { - auto attentionOp = cast<AttentionOpType>(op); - auto resultTypes = llvm::map_to_vector( - newOperands.take_back(attentionOp.getNumDpsInits()), - [](Value v) { return v.getType(); }); + auto op = cast<OpTy>(interfaceOp); + auto resultTypes = + llvm::map_to_vector(newOperands.take_back(op.getNumDpsInits()), + [](Value v) { return v.getType(); }); - IRMapping mapping; - for (auto [oldOperand, newOperand] : - llvm::zip(attentionOp->getOperands(), newOperands)) { - mapping.map(oldOperand, newOperand); - } - auto *cloned = b.clone(*attentionOp, mapping); - auto newOp = cast<AttentionOpType>(cloned); + // Replace operands by position: the same value can appear as several + // operands with different indexing maps and so different replacements. + auto newOp = cast<OpTy>(b.clone(*op)); + newOp->setOperands(newOperands); for (auto [result, type] : llvm::zip(newOp->getResults(), resultTypes)) { result.setType(type); } newOp.setIndexingMapsAttr(b.getAffineMapArrayAttr(newIndexingMaps)); - IRRewriter rewriter(b); - replaceUnitDimIndexOps(newOp.getBody(), droppedDims, rewriter); + if constexpr (std::is_same_v<OpTy, AttentionOp> || + std::is_same_v<OpTy, OnlineAttentionOp>) { + IRRewriter rewriter(b); + replaceUnitDimIndexOps(newOp.getBody(), droppedDims, rewriter); + } return newOp; }; FailureOr<linalg::DropUnitDimsResult> result = linalg::dropUnitDims( - rewriter, cast<IndexingMapOpInterface>(attentionOp.getOperation()), - builder, options); + rewriter, cast<IndexingMapOpInterface>(op.getOperation()), builder, + options); if (failed(result)) { return failure(); } - rewriter.replaceOp(attentionOp, result->replacements); + rewriter.replaceOp(op, result->replacements); return success(); } @@ -1271,11 +1291,12 @@ void populateFoldUnitExtentDimsPatterns( RewritePatternSet &patterns, const linalg::ControlDropUnitDims &options) { patterns.add<DropScatterUnitIndexDepth>(patterns.getContext()); - patterns - .add<DropGatherUnitDims, DropScatterUnitDims, - DropAttentionLikeUnitDims<AttentionOp>, - DropAttentionLikeUnitDims<OnlineAttentionOp>, DropMapStoreUnitDims>( - patterns.getContext(), options); + patterns.add<DropGatherUnitDims, DropScatterUnitDims, + DropMappedOperandsUnitDims<AttentionOp>, + DropMappedOperandsUnitDims<OnlineAttentionOp>, + DropMappedOperandsUnitDims<QuantizeAffineOp>, + DropMappedOperandsUnitDims<DequantizeAffineOp>, + DropMapStoreUnitDims>(patterns.getContext(), options); } } // namespace mlir::iree_compiler::IREE::LinalgExt
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/TestReshapeFusion.cpp b/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/TestReshapeFusion.cpp index f36e8c9..d8c7093 100644 --- a/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/TestReshapeFusion.cpp +++ b/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/TestReshapeFusion.cpp
@@ -31,6 +31,8 @@ RewritePatternSet patterns(context); LinalgExt::populateFoldReshapeOpsByExpansionPatterns( patterns, [](OpOperand *) { return true; }); + LinalgExt::populatePropagateAffineQuantizationReshapesPatterns( + patterns, [](OpOperand *) { return true; }); tensor::populateFoldTensorEmptyPatterns(patterns); tensor::ExpandShapeOp::getCanonicalizationPatterns(patterns, context); tensor::CollapseShapeOp::getCanonicalizationPatterns(patterns, context);
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/Transforms.h b/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/Transforms.h index af6c675..f11fbb5 100644 --- a/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/Transforms.h +++ b/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/Transforms.h
@@ -21,6 +21,12 @@ RewritePatternSet &patterns, const linalg::ControlFusionFn &controlFoldingReshapes); +/// Fold collapse_shape producers into quantize ops and expand_shape consumers +/// into dequantize ops by expanding the quantization op. +void populatePropagateAffineQuantizationReshapesPatterns( + RewritePatternSet &patterns, + const linalg::ControlFusionFn &controlFoldingReshapes); + /// Fuse transpose-like ops into LinalgExt ops (only `AttentionOp` supported). void populateFuseLinalgExtOpsWithTransposes( RewritePatternSet &patterns,
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/test/reshape_fusion.mlir b/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/test/reshape_fusion.mlir index 908e100..0b10e7f 100644 --- a/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/test/reshape_fusion.mlir +++ b/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/test/reshape_fusion.mlir
@@ -846,3 +846,84 @@ // CHECK-SAME: ins(%[[EXPANDED0]], %[[ARG1]] // CHECK-SAME: outs(%[[EXPANDED1]] // CHECK: tensor.collapse_shape %[[GATHER]] {{.*}} tensor<10x4x32xf16> into tensor<10x128xf16> + +// ----- + +// A quantize following a collapse_shape moves to the uncollapsed side, closer +// to its real-valued producer. Its per-channel scale expands with it. +util.func public @quantize_collapse_input(%x: tensor<4x8x64xf32>, %s: tensor<32xf32>) -> tensor<32x64xi8> { + %collapsed = tensor.collapse_shape %x [[0, 1], [2]] : tensor<4x8x64xf32> into tensor<32x64xf32> + %init = tensor.empty() : tensor<32x64xi8> + %q = iree_linalg_ext.quantize_affine + <indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> (d0)>, + affine_map<(d0, d1) -> (d0, d1)>], + quant_min = -128, quant_max = 127> + ins(%collapsed, %s : tensor<32x64xf32>, tensor<32xf32>) + outs(%init : tensor<32x64xi8>) -> tensor<32x64xi8> + util.return %q : tensor<32x64xi8> +} +// CHECK-LABEL: util.func public @quantize_collapse_input +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: +// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: +// CHECK-DAG: %[[SCALE:.+]] = tensor.expand_shape %[[ARG1]] {{.+}} tensor<32xf32> into tensor<4x8xf32> +// CHECK: %[[Q:.+]] = iree_linalg_ext.quantize_affine +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)> +// CHECK-SAME: affine_map<(d0, d1, d2) -> (d0, d1)> +// CHECK-SAME: affine_map<(d0, d1, d2) -> (d0, d1, d2)>] +// CHECK-SAME: ins(%[[ARG0]], %[[SCALE]] : tensor<4x8x64xf32>, tensor<4x8xf32>) +// CHECK: tensor.collapse_shape %[[Q]] {{.*}} tensor<4x8x64xi8> into tensor<32x64xi8> + +// ----- + +// A dequantize followed by expand_shape moves to the expanded side, closer to +// its real-valued consumer. Its per-channel parameters expand with it. +util.func public @dequantize_expand_result(%q: tensor<32x64xi8>, %s: tensor<32xf32>, + %z: tensor<32xi8>) -> tensor<4x8x64xf32> { + %init = tensor.empty() : tensor<32x64xf32> + %dq = iree_linalg_ext.dequantize_affine + <indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> (d0)>, + affine_map<(d0, d1) -> (d0)>, + affine_map<(d0, d1) -> (d0, d1)>]> + ins(%q, %s, %z : tensor<32x64xi8>, tensor<32xf32>, tensor<32xi8>) + outs(%init : tensor<32x64xf32>) -> tensor<32x64xf32> + %expanded = tensor.expand_shape %dq [[0, 1], [2]] output_shape [4, 8, 64] : tensor<32x64xf32> into tensor<4x8x64xf32> + util.return %expanded : tensor<4x8x64xf32> +} +// CHECK-LABEL: util.func public @dequantize_expand_result +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: +// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: +// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: +// CHECK-DAG: %[[INPUT:.+]] = tensor.expand_shape %[[ARG0]] {{.+}} tensor<32x64xi8> into tensor<4x8x64xi8> +// CHECK-DAG: %[[SCALE:.+]] = tensor.expand_shape %[[ARG1]] {{.+}} tensor<32xf32> into tensor<4x8xf32> +// CHECK-DAG: %[[ZP:.+]] = tensor.expand_shape %[[ARG2]] {{.+}} tensor<32xi8> into tensor<4x8xi8> +// CHECK: %[[DQ:.+]] = iree_linalg_ext.dequantize_affine +// CHECK-SAME: ins(%[[INPUT]], %[[SCALE]], %[[ZP]] : +// CHECK-NOT: tensor.expand_shape +// CHECK: util.return %[[DQ]] + +// ----- + +// A scalar parameter is invariant over the whole iteration space, so sinking +// dequantize through an expand_shape leaves the scalar parameter unchanged. +util.func public @dequantize_scalar_qparam_expand_result(%q: tensor<32x64xi8>, %s: f32) + -> tensor<4x8x64xf32> { + %init = tensor.empty() : tensor<32x64xf32> + %dq = iree_linalg_ext.dequantize_affine + <indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]> + ins(%q, %s : tensor<32x64xi8>, f32) + outs(%init : tensor<32x64xf32>) -> tensor<32x64xf32> + %expanded = tensor.expand_shape %dq [[0, 1], [2]] output_shape [4, 8, 64] : tensor<32x64xf32> into tensor<4x8x64xf32> + util.return %expanded : tensor<4x8x64xf32> +} +// CHECK-LABEL: util.func public @dequantize_scalar_qparam_expand_result +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: +// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: +// CHECK: %[[INPUT:.+]] = tensor.expand_shape %[[ARG0]] {{.+}} tensor<32x64xi8> into tensor<4x8x64xi8> +// CHECK: %[[DQ:.+]] = iree_linalg_ext.dequantize_affine +// CHECK-SAME: ins(%[[INPUT]], %[[ARG1]] : tensor<4x8x64xi8>, f32) +// CHECK-NOT: tensor.expand_shape +// CHECK: util.return %[[DQ]]
diff --git a/compiler/src/iree/compiler/DispatchCreation/test/fold_unit_dims.mlir b/compiler/src/iree/compiler/DispatchCreation/test/fold_unit_dims.mlir index 88a024c..a93d4a1 100644 --- a/compiler/src/iree/compiler/DispatchCreation/test/fold_unit_dims.mlir +++ b/compiler/src/iree/compiler/DispatchCreation/test/fold_unit_dims.mlir
@@ -756,3 +756,49 @@ // CHECK: %[[EXPAND2:.+]] = tensor.expand_shape %[[ATTN]]#2 // CHECK-SAME: tensor<8x64xf32> into tensor<1x8x64xf32> // CHECK: util.return %[[EXPAND0]], %[[EXPAND1]], %[[EXPAND2]] + +// ----- + +// Folding the unit batch dim shared by a dequantize and its convolution +// consumer collapses the quantized input, so the dequantize still feeds the +// convolution directly. +util.func public @dequantize_conv_unit_batch(%q: tensor<1x4x4x2xi8>, %s: tensor<f32>, + %z: tensor<i8>, %filter: tensor<2x2x2x2xf32>) -> tensor<1x3x3x2xf32> { + %cst = arith.constant 0.000000e+00 : f32 + %dqinit = tensor.empty() : tensor<1x4x4x2xf32> + %dq = iree_linalg_ext.dequantize_affine + <indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, + affine_map<(d0, d1, d2, d3) -> ()>, + affine_map<(d0, d1, d2, d3) -> ()>, + affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>]> + ins(%q, %s, %z : tensor<1x4x4x2xi8>, tensor<f32>, tensor<i8>) + outs(%dqinit : tensor<1x4x4x2xf32>) -> tensor<1x4x4x2xf32> + %init = tensor.empty() : tensor<1x3x3x2xf32> + %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x3x3x2xf32>) -> tensor<1x3x3x2xf32> + %conv = linalg.generic { + indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 + d4, d2 + d5, d6)>, + affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d4, d5, d6, d3)>, + affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)>], + iterator_types = ["parallel", "parallel", "parallel", "parallel", + "reduction", "reduction", "reduction"]} + ins(%dq, %filter : tensor<1x4x4x2xf32>, tensor<2x2x2x2xf32>) + outs(%fill : tensor<1x3x3x2xf32>) { + ^bb0(%in: f32, %in_0: f32, %out: f32): + %m = arith.mulf %in, %in_0 : f32 + %a = arith.addf %out, %m : f32 + linalg.yield %a : f32 + } -> tensor<1x3x3x2xf32> + util.return %conv : tensor<1x3x3x2xf32> +} +// CHECK-LABEL: util.func public @dequantize_conv_unit_batch +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: +// CHECK: %[[INPUT:.+]] = tensor.collapse_shape %[[ARG0]] +// CHECK-SAME: tensor<1x4x4x2xi8> into tensor<4x4x2xi8> +// CHECK: %[[DQ:.+]] = iree_linalg_ext.dequantize_affine +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)> +// CHECK-SAME: ins(%[[INPUT]], +// CHECK-SAME: -> tensor<4x4x2xf32> +// The dequantize feeds the convolution directly, with no collapse_shape in +// between. +// CHECK: linalg.generic +// CHECK-SAME: ins(%[[DQ]],
diff --git a/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel b/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel index f15a9fc..570c141 100644 --- a/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel +++ b/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel
@@ -95,6 +95,7 @@ "MaterializeHomogeneousEncodings.cpp", "OptimizeNumerics.cpp", "Passes.cpp", + "PropagateAffineQuantization.cpp", "PropagateLinalgTranspose.cpp", "QuantizedConvToConv.cpp", "QuantizedMatmulToMatmul.cpp",
diff --git a/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt b/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt index 068698a..d5387a4 100644 --- a/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt +++ b/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt
@@ -89,6 +89,7 @@ "MaterializeHomogeneousEncodings.cpp" "OptimizeNumerics.cpp" "Passes.cpp" + "PropagateAffineQuantization.cpp" "PropagateLinalgTranspose.cpp" "QuantizedConvToConv.cpp" "QuantizedMatmulToMatmul.cpp"
diff --git a/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp b/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp index 1b556bc..bd99e43 100644 --- a/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp +++ b/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp
@@ -130,9 +130,11 @@ .addPass(IREE::Util::createOptimizeIntArithmeticPass) .addPass(createLinalgQuantizedConvToConvPass) .addPass(createLinalgQuantizedMatmulToMatmulPass) - // Match direct dequantize producers before convolution lowering and - // unit-dimension folding insert reshapes between them and the - // contraction. + // Run propagation immediately before the QDQ rewrite so contractions + // consume dequantize results directly. This must happen before im2col + // inserts a packing operation between dequantize and its matmul. + .addPredicatedPass(clConvertQDQToIntegerMath, + createPropagateAffineQuantizationPass) .addPredicatedPass(clConvertQDQToIntegerMath, createConvertQDQToIntegerMathPass) .addPredicatedPass(transformOptions.useIm2colForConvs, @@ -212,8 +214,10 @@ clEnableEdgeReshapePropagation; return createPropagateLinalgTransposePass(options); }) - // Retry for contractions whose dequantize producers were exposed by - // reshape/transpose propagation, before selecting integer layouts. + // Run propagation immediately before the QDQ rewrite again: earlier + // reshape/transpose passes can separate dequantize from its consumers. + .addPredicatedPass(clConvertQDQToIntegerMath, + createPropagateAffineQuantizationPass) .addPredicatedPass(clConvertQDQToIntegerMath, createConvertQDQToIntegerMathPass) .addPass(IREE::Flow::createCanonicalizePass)
diff --git a/compiler/src/iree/compiler/GlobalOptimization/Passes.td b/compiler/src/iree/compiler/GlobalOptimization/Passes.td index 448cf65..6289471 100644 --- a/compiler/src/iree/compiler/GlobalOptimization/Passes.td +++ b/compiler/src/iree/compiler/GlobalOptimization/Passes.td
@@ -101,6 +101,21 @@ let summary = "Expands tensor shape dimensions into SSA values across the program."; } +def PropagateAffineQuantizationPass : + InterfacePass<"iree-global-opt-propagate-affine-quantization", "mlir::FunctionOpInterface"> { + let summary = "Propagates affine quantization through adjacent tensor operations."; + let description = [{ + Moves `iree_linalg_ext.quantize_affine` before a producer-side + `tensor.collapse_shape`, and moves a consumer-side `tensor.expand_shape`, + or a zero `tensor.pad` whose padded dimensions do not index the scale and + whose zero point is per-tensor, before `iree_linalg_ext.dequantize_affine`. + These shape operations then run on quantized values. A consumer-side + `linalg.transpose` is fused into the dequantize's output map instead. When + a contraction consumes the result, the dequantize becomes its direct + producer so the integer contraction rewrite can match it. + }]; +} + def ConvertQDQToIntegerMathPass : InterfacePass<"iree-global-opt-convert-qdq-to-integer-math", "mlir::FunctionOpInterface"> { let summary = "Rewrites contractions over dequantized operands into integer contractions.";
diff --git a/compiler/src/iree/compiler/GlobalOptimization/PropagateAffineQuantization.cpp b/compiler/src/iree/compiler/GlobalOptimization/PropagateAffineQuantization.cpp new file mode 100644 index 0000000..e27f88f --- /dev/null +++ b/compiler/src/iree/compiler/GlobalOptimization/PropagateAffineQuantization.cpp
@@ -0,0 +1,173 @@ +// Copyright 2026 The IREE Authors +// +// Licensed under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h" +#include "iree/compiler/Dialect/LinalgExt/Transforms/Transforms.h" +#include "iree/compiler/GlobalOptimization/Passes.h" +#include "iree/compiler/GlobalOptimization/TransposePropagationPatterns.h" +#include "llvm/ADT/SmallBitVector.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Arith/Utils/Utils.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Dialect/Utils/IndexingUtils.h" +#include "mlir/IR/AffineMap.h" +#include "mlir/IR/Matchers.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" + +namespace mlir::iree_compiler::GlobalOptimization { + +#define GEN_PASS_DEF_PROPAGATEAFFINEQUANTIZATIONPASS +#include "iree/compiler/GlobalOptimization/Passes.h.inc" + +using IREE::LinalgExt::DequantizeAffineOp; +using IREE::LinalgExt::QuantizeAffineOp; + +namespace { + +/// Bubbles a pad of a dequantized value onto the quantized side. +/// +/// The quantized pad uses the dequantize's zero point, which represents real +/// zero through `(zp - zp) * scale`. Moving the pad above the dequantize makes +/// the dequantize the immediate producer of whatever consumed the pad, and +/// pads narrower data on the way. Other users of the dequantize keep reading +/// the original one. +struct BubblePadThroughDequantize : public OpRewritePattern<tensor::PadOp> { + using OpRewritePattern<tensor::PadOp>::OpRewritePattern; + + LogicalResult matchAndRewrite(tensor::PadOp padOp, + PatternRewriter &rewriter) const override { + auto dequantizeOp = padOp.getSource().getDefiningOp<DequantizeAffineOp>(); + if (!dequantizeOp) { + return rewriter.notifyMatchFailure(padOp, "source is not a dequantize"); + } + // Only a pad with real zero can move to the quantized side, where the zero + // point represents it. + Value padValue = padOp.getConstantPaddingValue(); + if (!padValue || !matchPattern(padValue, m_AnyZeroFloat())) { + return rewriter.notifyMatchFailure(padOp, "pad value is not a zero " + "constant"); + } + + // The quantized pad must yield a value defined outside its body: + // tensor.pad tiling (tensor::bubbleUpPadSlice) supports no other padding + // value, so a zero point that varies per channel, which the pad body would + // have to read using the channel index, could not be compiled. + Value zeroPoint = dequantizeOp.getZeroPoint(); + if (zeroPoint && dequantizeOp.getZeroPointMap().getNumResults() != 0) { + return rewriter.notifyMatchFailure(padOp, "zero point is not per-tensor"); + } + if (zeroPoint && isa<ShapedType>(zeroPoint.getType()) && + !isa<RankedTensorType>(zeroPoint.getType())) { + return rewriter.notifyMatchFailure( + padOp, "zero point is neither a scalar nor a 0-d tensor"); + } + + // After the rewrite the dequantize also computes the padded positions, so + // a scale indexed by a padded dimension would be read past its end. The + // zero point is per-tensor at this point and cannot be indexed that way. + AffineMap outputToIteration = + inversePermutation(dequantizeOp.getOutputMap()); + AffineMap scaleMap = dequantizeOp.getScaleMap().compose(outputToIteration); + llvm::SmallBitVector paddedDims = padOp.getPaddedDims(); + if (llvm::any_of(paddedDims.set_bits(), [&](unsigned dim) { + return scaleMap.isFunctionOfDim(dim); + })) { + return rewriter.notifyMatchFailure( + padOp, "a padded dimension indexes the quantization parameters"); + } + + Location loc = padOp.getLoc(); + Type storageType = dequantizeOp.getInputType().getElementType(); + // A symmetric dequantize has an implicit zero point of zero. + Value quantizedPadValue; + if (zeroPoint) { + if (isa<RankedTensorType>(zeroPoint.getType())) { + zeroPoint = tensor::ExtractOp::create(rewriter, loc, zeroPoint, + /*indices=*/ValueRange{}); + } + quantizedPadValue = + convertScalarToDtype(rewriter, loc, zeroPoint, storageType, + /*isUnsignedCast=*/dequantizeOp.getZpUnsigned()); + } else { + quantizedPadValue = arith::ConstantOp::create( + rewriter, loc, rewriter.getZeroAttr(storageType)); + } + + AffineMap outputToInput = + dequantizeOp.getInputMap().compose(outputToIteration); + // Keep any static extents the original pad result has beyond what its + // source and padding amounts imply. + auto paddedInputType = RankedTensorType::get( + applyPermutationMap(outputToInput, padOp.getResultType().getShape()), + storageType); + auto paddedInput = tensor::PadOp::create( + rewriter, loc, paddedInputType, dequantizeOp.getInput(), + applyPermutationMap<OpFoldResult>(outputToInput, + padOp.getMixedLowPad()), + applyPermutationMap<OpFoldResult>(outputToInput, + padOp.getMixedHighPad()), + quantizedPadValue, padOp.getNofold()); + + Value init = tensor::EmptyOp::create( + rewriter, loc, + applyPermutationMap<OpFoldResult>( + inversePermutation(outputToInput), + tensor::getMixedSizes(rewriter, loc, paddedInput)), + padOp.getResultType().getElementType()); + SmallVector<Value> operands = dequantizeOp->getOperands(); + operands.front() = paddedInput.getResult(); + operands.back() = init; + Operation *dequantizedPad = + mlir::clone(rewriter, dequantizeOp, padOp->getResultTypes(), operands); + rewriter.replaceOp(padOp, dequantizedPad->getResults()); + return success(); + } +}; + +/// Allows a collapse_shape feeding quantize or an expand_shape consuming +/// dequantize. Both require a single-use edge so propagation removes the +/// collapse_shape or expand_shape from the real-valued side without adding one +/// for another consumer. Only the quantize's value operand counts: a +/// collapse_shape producing a quantization parameter does not separate the +/// quantize from its real-valued producer. +static bool canPropagateAffineQuantizationReshape(OpOperand *operand) { + if (auto quantize = dyn_cast<QuantizeAffineOp>(operand->getOwner())) { + return operand == &quantize.getInputMutable() && + operand->get().getDefiningOp<tensor::CollapseShapeOp>() && + operand->get().hasOneUse(); + } + return isa<tensor::ExpandShapeOp>(operand->getOwner()) && + operand->get().hasOneUse(); +} + +/// Restricts transpose rewrites to transposes of dequantized values. The +/// general transpose propagation pass owns every other transpose. +static bool isTransposeOfDequantize(OpOperand *operand) { + return isa_and_nonnull<DequantizeAffineOp>(operand->get().getDefiningOp()); +} + +struct PropagateAffineQuantizationPass + : public impl::PropagateAffineQuantizationPassBase< + PropagateAffineQuantizationPass> { + void runOnOperation() override { + MLIRContext *context = &getContext(); + RewritePatternSet patterns(context); + IREE::LinalgExt::populatePropagateAffineQuantizationReshapesPatterns( + patterns, canPropagateAffineQuantizationReshape); + populateSpecializeGenericTransposePattern(patterns, + isTransposeOfDequantize); + populateFuseTransposeIntoDequantizePattern(patterns, + isTransposeOfDequantize); + patterns.add<BubblePadThroughDequantize>(context); + if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) { + return signalPassFailure(); + } + } +}; + +} // namespace +} // namespace mlir::iree_compiler::GlobalOptimization
diff --git a/compiler/src/iree/compiler/GlobalOptimization/PropagateLinalgTranspose.cpp b/compiler/src/iree/compiler/GlobalOptimization/PropagateLinalgTranspose.cpp index 2541af8..a66947a 100644 --- a/compiler/src/iree/compiler/GlobalOptimization/PropagateLinalgTranspose.cpp +++ b/compiler/src/iree/compiler/GlobalOptimization/PropagateLinalgTranspose.cpp
@@ -261,6 +261,8 @@ populateBubbleTransposeThroughUnaryElementwiseDpsInitPattern( bubblingPatterns, isOutsideDispatch, /*benefit=*/2); populateComposeTransposesPattern(bubblingPatterns, isOutsideDispatch); + populateFuseTransposeIntoDequantizePattern(bubblingPatterns, + isOutsideDispatch); populateCommonCanonicalizationPatterns(context, bubblingPatterns); GreedyRewriteConfig config;
diff --git a/compiler/src/iree/compiler/GlobalOptimization/TransposePropagationPatterns.cpp b/compiler/src/iree/compiler/GlobalOptimization/TransposePropagationPatterns.cpp index b6c6af9..f52341b 100644 --- a/compiler/src/iree/compiler/GlobalOptimization/TransposePropagationPatterns.cpp +++ b/compiler/src/iree/compiler/GlobalOptimization/TransposePropagationPatterns.cpp
@@ -12,6 +12,7 @@ //===----------------------------------------------------------------------===// #include "iree/compiler/GlobalOptimization/TransposePropagationPatterns.h" +#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h" #include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h" #include "mlir/Dialect/Linalg/IR/Linalg.h" #include "mlir/Dialect/Linalg/Transforms/Transforms.h" @@ -23,6 +24,8 @@ namespace mlir::iree_compiler::GlobalOptimization { +using IREE::LinalgExt::DequantizeAffineOp; + //===----------------------------------------------------------------------===// // Transpose permutation helpers //===----------------------------------------------------------------------===// @@ -379,6 +382,63 @@ ControlTransposePropagationFn controlFn; }; +// Fuses a transpose into its single-use dequantize producer by permuting the +// dequantize's output map. The dequantize then writes the transposed layout +// directly, so its consumer reads it without an intervening transpose, and no +// transpose is materialized on either side. +// +// linalg.transpose follows result[i] = input[perm[i]], so the fused output +// map reads its results in the transposed order. This mirrors the output map +// update in FuseTransposeWithProducerLinalgOp. +class FuseTransposeIntoDequantize + : public OpRewritePattern<linalg::TransposeOp> { +public: + FuseTransposeIntoDequantize(MLIRContext *ctx, + ControlTransposePropagationFn controlFn, + PatternBenefit b = 1) + : OpRewritePattern<linalg::TransposeOp>(ctx, b), + controlFn(std::move(controlFn)) {} + + LogicalResult matchAndRewrite(linalg::TransposeOp transposeOp, + PatternRewriter &rewriter) const override { + auto dequantizeOp = + transposeOp.getInput().getDefiningOp<DequantizeAffineOp>(); + if (!dequantizeOp) { + return rewriter.notifyMatchFailure(transposeOp, + "input is not a dequantize"); + } + // With other consumers the untransposed dequantize has to stay, so fusing + // would duplicate the dequantization rather than remove the transpose. + if (!dequantizeOp->getResult(0).hasOneUse()) { + return rewriter.notifyMatchFailure(transposeOp, + "dequantize has other users"); + } + if (!controlFn(&transposeOp.getInputMutable())) { + return rewriter.notifyMatchFailure(transposeOp, + "rejected by control function"); + } + + SmallVector<AffineMap> maps = dequantizeOp.getIndexingMapsArray(); + AffineMap outputMap = maps.back(); + maps.back() = AffineMap::get( + outputMap.getNumDims(), outputMap.getNumSymbols(), + applyPermutation(outputMap.getResults(), transposeOp.getPermutation()), + rewriter.getContext()); + + SmallVector<Value> operands = dequantizeOp->getOperands(); + operands.back() = transposeOp.getDpsInits()[0]; + auto fusedOp = cast<DequantizeAffineOp>(mlir::clone( + rewriter, dequantizeOp, transposeOp->getResultTypes(), operands)); + fusedOp.setIndexingMapsAttr(rewriter.getAffineMapArrayAttr(maps)); + rewriter.replaceOp(transposeOp, fusedOp->getResults()); + return success(); + } + +private: + // Decides whether the transpose may fuse into the dequantize. + ControlTransposePropagationFn controlFn; +}; + } // namespace //===----------------------------------------------------------------------===// @@ -1213,6 +1273,13 @@ patterns.getContext(), std::move(controlFn), benefit); } +void populateFuseTransposeIntoDequantizePattern( + RewritePatternSet &patterns, ControlTransposePropagationFn controlFn, + PatternBenefit benefit) { + patterns.add<FuseTransposeIntoDequantize>(patterns.getContext(), + std::move(controlFn), benefit); +} + void populateComposeTransposesPattern(RewritePatternSet &patterns, ControlTransposePropagationFn controlFn, PatternBenefit benefit) {
diff --git a/compiler/src/iree/compiler/GlobalOptimization/TransposePropagationPatterns.h b/compiler/src/iree/compiler/GlobalOptimization/TransposePropagationPatterns.h index abfd656..826c775 100644 --- a/compiler/src/iree/compiler/GlobalOptimization/TransposePropagationPatterns.h +++ b/compiler/src/iree/compiler/GlobalOptimization/TransposePropagationPatterns.h
@@ -55,6 +55,12 @@ RewritePatternSet &patterns, ControlTransposePropagationFn controlFn, PatternBenefit benefit = 1); +// Fuses a transpose into its single-use iree_linalg_ext.dequantize_affine +// producer by permuting the dequantize's output map. +void populateFuseTransposeIntoDequantizePattern( + RewritePatternSet &patterns, ControlTransposePropagationFn controlFn, + PatternBenefit benefit = 1); + //===----------------------------------------------------------------------===// // Transpose sinking patterns //===----------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel b/compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel index cf1a0a0..ee4b4aa 100644 --- a/compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel +++ b/compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel
@@ -37,6 +37,7 @@ "linalg_quantized_conv_to_conv.mlir", "linalg_quantized_matmul_to_matmul.mlir", "optimize_numerics.mlir", + "propagate_affine_quantization.mlir", "propagate_linalg_transpose.mlir", "raise_special_ops.mlir", "remove_zero_extent_tensors.mlir",
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt b/compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt index de43a37..3773729 100644 --- a/compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt +++ b/compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt
@@ -34,6 +34,7 @@ "linalg_quantized_conv_to_conv.mlir" "linalg_quantized_matmul_to_matmul.mlir" "optimize_numerics.mlir" + "propagate_affine_quantization.mlir" "propagate_linalg_transpose.mlir" "raise_special_ops.mlir" "remove_zero_extent_tensors.mlir"
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/propagate_affine_quantization.mlir b/compiler/src/iree/compiler/GlobalOptimization/test/propagate_affine_quantization.mlir new file mode 100644 index 0000000..b680ef3 --- /dev/null +++ b/compiler/src/iree/compiler/GlobalOptimization/test/propagate_affine_quantization.mlir
@@ -0,0 +1,373 @@ +// RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(func.func(iree-global-opt-propagate-affine-quantization))" %s | FileCheck %s + +#id2 = affine_map<(d0, d1) -> (d0, d1)> +#scalar2 = affine_map<(d0, d1) -> ()> + +// A collapse_shape feeding quantize moves onto the quantized side, leaving the +// quantize next to its real-valued producer. +func.func @absorb_quantize_producer_collapse(%input: tensor<2x3x4xf32>, %scale: f32) + -> tensor<6x4xi8> { + %collapsed = tensor.collapse_shape %input [[0, 1], [2]] + : tensor<2x3x4xf32> into tensor<6x4xf32> + %init = tensor.empty() : tensor<6x4xi8> + %quantized = iree_linalg_ext.quantize_affine + <indexing_maps = [#id2, #scalar2, #id2], + quant_min = -128, quant_max = 127> + ins(%collapsed, %scale : tensor<6x4xf32>, f32) + outs(%init : tensor<6x4xi8>) -> tensor<6x4xi8> + return %quantized : tensor<6x4xi8> +} +// CHECK-LABEL: func.func @absorb_quantize_producer_collapse( +// CHECK-SAME: %[[INPUT:[a-zA-Z0-9_]+]]: tensor<2x3x4xf32> +// CHECK: %[[QUANTIZED:.+]] = iree_linalg_ext.quantize_affine +// CHECK-SAME: ins(%[[INPUT]] +// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[QUANTIZED]] +// CHECK: return %[[COLLAPSED]] + +// ----- + +#id2 = affine_map<(d0, d1) -> (d0, d1)> +#row2 = affine_map<(d0, d1) -> (d0)> + +// Only a collapse_shape on the quantize's value operand is folded. Folding one +// that produces the scale would add an expand_shape on the real-valued input +// rather than remove the collapse_shape. +func.func @decline_quantize_scale_collapse(%input: tensor<6x4xf32>, + %scale: tensor<2x3xf32>) -> tensor<6x4xi8> { + %collapsed = tensor.collapse_shape %scale [[0, 1]] + : tensor<2x3xf32> into tensor<6xf32> + %init = tensor.empty() : tensor<6x4xi8> + %quantized = iree_linalg_ext.quantize_affine + <indexing_maps = [#id2, #row2, #id2], + quant_min = -128, quant_max = 127> + ins(%input, %collapsed : tensor<6x4xf32>, tensor<6xf32>) + outs(%init : tensor<6x4xi8>) -> tensor<6x4xi8> + return %quantized : tensor<6x4xi8> +} +// CHECK-LABEL: func.func @decline_quantize_scale_collapse( +// CHECK-SAME: %[[INPUT:[a-zA-Z0-9_]+]]: tensor<6x4xf32> +// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape +// CHECK: iree_linalg_ext.quantize_affine +// CHECK-SAME: ins(%[[INPUT]], %[[COLLAPSED]] : + +// ----- + +#id3 = affine_map<(d0, d1, d2) -> (d0, d1, d2)> +#scalar3 = affine_map<(d0, d1, d2) -> ()> + +// An expand_shape consuming a dequantize moves onto the quantized side, so the +// consumers of the expand_shape read the dequantize directly. +func.func @absorb_consumer_expand(%aq: tensor<3x6x6xi8>, %sa: f32, %za: i64) -> tensor<1x3x6x6xf32> { + %init = tensor.empty() : tensor<3x6x6xf32> + %a = iree_linalg_ext.dequantize_affine + <indexing_maps = [#id3, #scalar3, #scalar3, #id3]> + ins(%aq, %sa, %za : tensor<3x6x6xi8>, f32, i64) + outs(%init : tensor<3x6x6xf32>) -> tensor<3x6x6xf32> + %expanded = tensor.expand_shape %a [[0, 1], [2], [3]] output_shape [1, 3, 6, 6] + : tensor<3x6x6xf32> into tensor<1x3x6x6xf32> + return %expanded : tensor<1x3x6x6xf32> +} +// CHECK-LABEL: func.func @absorb_consumer_expand( +// CHECK-SAME: %[[AQ:[a-zA-Z0-9_]+]]: tensor<3x6x6xi8> +// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[AQ]] +// CHECK-SAME: tensor<3x6x6xi8> into tensor<1x3x6x6xi8> +// CHECK: %[[DEQ:.+]] = iree_linalg_ext.dequantize_affine +// CHECK-SAME: ins(%[[EXPANDED]] +// CHECK: return %[[DEQ]] + +// ----- + +#id3 = affine_map<(d0, d1, d2) -> (d0, d1, d2)> +#scalar3 = affine_map<(d0, d1, d2) -> ()> + +// An expand_shape is not absorbed when the dequantize has another consumer. +// Expanding the dequantize would put a collapse_shape on that consumer's path +// instead of removing the expand_shape. +func.func @decline_expand_with_multiple_consumers(%aq: tensor<3x6x6xi8>, %sa: f32, %za: i64) + -> (tensor<1x3x6x6xf32>, tensor<3x6x6xf32>) { + %init = tensor.empty() : tensor<3x6x6xf32> + %a = iree_linalg_ext.dequantize_affine + <indexing_maps = [#id3, #scalar3, #scalar3, #id3]> + ins(%aq, %sa, %za : tensor<3x6x6xi8>, f32, i64) + outs(%init : tensor<3x6x6xf32>) -> tensor<3x6x6xf32> + %expanded = tensor.expand_shape %a [[0, 1], [2], [3]] output_shape [1, 3, 6, 6] + : tensor<3x6x6xf32> into tensor<1x3x6x6xf32> + return %expanded, %a : tensor<1x3x6x6xf32>, tensor<3x6x6xf32> +} +// CHECK-LABEL: func.func @decline_expand_with_multiple_consumers( +// CHECK: %[[DEQ:.+]] = iree_linalg_ext.dequantize_affine +// CHECK-SAME: outs(%{{.+}} : tensor<3x6x6xf32>) -> tensor<3x6x6xf32> +// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[DEQ]] +// CHECK: return %[[EXPANDED]], %[[DEQ]] + +// ----- + +#id2 = affine_map<(d0, d1) -> (d0, d1)> +#scalar2 = affine_map<(d0, d1) -> ()> + +// Padding a dequantized value with zero is the same as padding the quantized +// value with the zero point, because the zero point is the quantized +// representation of 0.0. Bubbling the pad above the dequantize pads i8 instead +// of f32 data and makes the dequantize the direct producer of the pad's +// consumer. +func.func @bubble_pad_through_asymmetric_dequantize(%aq: tensor<4x4xi8>, %sa: f32, %za: i64) + -> tensor<6x6xf32> { + %cst = arith.constant 0.000000e+00 : f32 + %init = tensor.empty() : tensor<4x4xf32> + %a = iree_linalg_ext.dequantize_affine + <indexing_maps = [#id2, #scalar2, #scalar2, #id2]> + ins(%aq, %sa, %za : tensor<4x4xi8>, f32, i64) + outs(%init : tensor<4x4xf32>) -> tensor<4x4xf32> + %padded = tensor.pad %a low[1, 1] high[1, 1] { + ^bb0(%i: index, %j: index): + tensor.yield %cst : f32 + } : tensor<4x4xf32> to tensor<6x6xf32> + return %padded : tensor<6x6xf32> +} +// CHECK-LABEL: func.func @bubble_pad_through_asymmetric_dequantize( +// CHECK-SAME: %[[AQ:[a-zA-Z0-9_]+]]: tensor<4x4xi8> +// CHECK-SAME: %[[ZA:[a-zA-Z0-9_]+]]: i64 +// The quantized side is padded with the zero point, truncated to the storage type. +// CHECK: %[[ZP:.+]] = arith.trunci %[[ZA]] : i64 to i8 +// CHECK: %[[PADDED:.+]] = tensor.pad %[[AQ]] low[1, 1] high[1, 1] +// CHECK: tensor.yield %[[ZP]] : i8 +// CHECK: } : tensor<4x4xi8> to tensor<6x6xi8> +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK-SAME: ins(%[[PADDED]] + +// ----- + +#id2 = affine_map<(d0, d1) -> (d0, d1)> +#scalar2 = affine_map<(d0, d1) -> ()> + +// A pad bubbles through a dequantize that has other users: the pad reads a new +// dequantize of the padded quantized value, and the other users keep the +// original dequantize. +func.func @bubble_pad_through_multi_use_dequantize(%aq: tensor<4x4xi8>, %sa: f32, %za: i64) + -> (tensor<6x6xf32>, tensor<4x4xf32>) { + %cst = arith.constant 0.000000e+00 : f32 + %init = tensor.empty() : tensor<4x4xf32> + %a = iree_linalg_ext.dequantize_affine + <indexing_maps = [#id2, #scalar2, #scalar2, #id2]> + ins(%aq, %sa, %za : tensor<4x4xi8>, f32, i64) + outs(%init : tensor<4x4xf32>) -> tensor<4x4xf32> + %padded = tensor.pad %a low[1, 1] high[1, 1] { + ^bb0(%i: index, %j: index): + tensor.yield %cst : f32 + } : tensor<4x4xf32> to tensor<6x6xf32> + return %padded, %a : tensor<6x6xf32>, tensor<4x4xf32> +} +// CHECK-LABEL: func.func @bubble_pad_through_multi_use_dequantize( +// CHECK-SAME: %[[AQ:[a-zA-Z0-9_]+]]: tensor<4x4xi8> +// CHECK-DAG: %[[UNPADDED:.+]] = iree_linalg_ext.dequantize_affine <{{.*}}> ins(%[[AQ]], +// CHECK-DAG: %[[PADDED_Q:.+]] = tensor.pad %[[AQ]] low[1, 1] high[1, 1] +// CHECK-DAG: %[[PADDED:.+]] = iree_linalg_ext.dequantize_affine <{{.*}}> ins(%[[PADDED_Q]], +// CHECK: return %[[PADDED]], %[[UNPADDED]] + +// ----- + +#id2 = affine_map<(d0, d1) -> (d0, d1)> +#scalar2 = affine_map<(d0, d1) -> ()> + +// A per-tensor zero point may also be a 0-d tensor. Its single element pads +// the quantized side. +func.func @bubble_pad_through_0d_zero_point_dequantize(%aq: tensor<4x4xi8>, + %sa: tensor<f32>, %za: tensor<i8>) -> tensor<6x6xf32> { + %cst = arith.constant 0.000000e+00 : f32 + %init = tensor.empty() : tensor<4x4xf32> + %a = iree_linalg_ext.dequantize_affine + <indexing_maps = [#id2, #scalar2, #scalar2, #id2]> + ins(%aq, %sa, %za : tensor<4x4xi8>, tensor<f32>, tensor<i8>) + outs(%init : tensor<4x4xf32>) -> tensor<4x4xf32> + %padded = tensor.pad %a low[1, 1] high[1, 1] { + ^bb0(%i: index, %j: index): + tensor.yield %cst : f32 + } : tensor<4x4xf32> to tensor<6x6xf32> + return %padded : tensor<6x6xf32> +} +// CHECK-LABEL: func.func @bubble_pad_through_0d_zero_point_dequantize( +// CHECK-SAME: %[[AQ:[a-zA-Z0-9_]+]]: tensor<4x4xi8> +// CHECK-SAME: %[[SA:[a-zA-Z0-9_]+]]: tensor<f32> +// CHECK-SAME: %[[ZA:[a-zA-Z0-9_]+]]: tensor<i8> +// CHECK: %[[ZP:.+]] = tensor.extract %[[ZA]][] : tensor<i8> +// CHECK: %[[PADDED:.+]] = tensor.pad %[[AQ]] low[1, 1] high[1, 1] +// CHECK: tensor.yield %[[ZP]] : i8 +// CHECK: } : tensor<4x4xi8> to tensor<6x6xi8> +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK-SAME: ins(%[[PADDED]], %[[SA]], %[[ZA]] : +// CHECK-SAME: -> tensor<6x6xf32> + +// ----- + +#id2 = affine_map<(d0, d1) -> (d0, d1)> +#scalar2 = affine_map<(d0, d1) -> ()> + +// A pad may have a static result type for a dynamic source; the quantized pad +// keeps that static result type. +func.func @bubble_pad_keeps_static_result_type(%aq: tensor<?x4xi8>, %sa: f32, %za: i8) + -> tensor<6x6xf32> { + %cst = arith.constant 0.000000e+00 : f32 + %c0 = arith.constant 0 : index + %dim = tensor.dim %aq, %c0 : tensor<?x4xi8> + %init = tensor.empty(%dim) : tensor<?x4xf32> + %a = iree_linalg_ext.dequantize_affine + <indexing_maps = [#id2, #scalar2, #scalar2, #id2]> + ins(%aq, %sa, %za : tensor<?x4xi8>, f32, i8) + outs(%init : tensor<?x4xf32>) -> tensor<?x4xf32> + %padded = tensor.pad %a low[1, 1] high[1, 1] { + ^bb0(%i: index, %j: index): + tensor.yield %cst : f32 + } : tensor<?x4xf32> to tensor<6x6xf32> + return %padded : tensor<6x6xf32> +} +// CHECK-LABEL: func.func @bubble_pad_keeps_static_result_type( +// CHECK-SAME: %[[AQ:[a-zA-Z0-9_]+]]: tensor<?x4xi8> +// CHECK: %[[PADDED:.+]] = tensor.pad %[[AQ]] low[1, 1] high[1, 1] +// CHECK: } : tensor<?x4xi8> to tensor<6x6xi8> +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK-SAME: ins(%[[PADDED]] +// CHECK-SAME: -> tensor<6x6xf32> + +// ----- + +#id2 = affine_map<(d0, d1) -> (d0, d1)> +#row2 = affine_map<(d0, d1) -> (d0)> + +// A symmetric dequantize has an implicit zero point of zero, so the quantized +// side is padded with a zero of the storage type. +func.func @bubble_pad_through_symmetric_dequantize(%aq: tensor<4x4xi8>, %sa: tensor<4xf32>) + -> tensor<4x6xf32> { + %cst = arith.constant 0.000000e+00 : f32 + %init = tensor.empty() : tensor<4x4xf32> + %a = iree_linalg_ext.dequantize_affine + <indexing_maps = [#id2, #row2, #id2]> + ins(%aq, %sa : tensor<4x4xi8>, tensor<4xf32>) + outs(%init : tensor<4x4xf32>) -> tensor<4x4xf32> + %padded = tensor.pad %a low[0, 1] high[0, 1] { + ^bb0(%i: index, %j: index): + tensor.yield %cst : f32 + } : tensor<4x4xf32> to tensor<4x6xf32> + return %padded : tensor<4x6xf32> +} +// CHECK-LABEL: func.func @bubble_pad_through_symmetric_dequantize( +// CHECK-SAME: %[[AQ:[a-zA-Z0-9_]+]]: tensor<4x4xi8> +// CHECK-DAG: %[[ZERO:.+]] = arith.constant 0 : i8 +// CHECK: %[[PADDED:.+]] = tensor.pad %[[AQ]] low[0, 1] high[0, 1] +// CHECK: tensor.yield %[[ZERO]] : i8 +// CHECK: } : tensor<4x4xi8> to tensor<4x6xi8> +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK-SAME: ins(%[[PADDED]] + +// ----- + +#id2 = affine_map<(d0, d1) -> (d0, d1)> +#row2 = affine_map<(d0, d1) -> (d0)> + +// The padded dimension indexes the scale, so the padded rows have no scale to +// read and the pad cannot move above the dequantize. +func.func @decline_pad_along_quantized_axis(%aq: tensor<4x4xi8>, %sa: tensor<4xf32>) + -> tensor<6x4xf32> { + %cst = arith.constant 0.000000e+00 : f32 + %init = tensor.empty() : tensor<4x4xf32> + %a = iree_linalg_ext.dequantize_affine + <indexing_maps = [#id2, #row2, #id2]> + ins(%aq, %sa : tensor<4x4xi8>, tensor<4xf32>) + outs(%init : tensor<4x4xf32>) -> tensor<4x4xf32> + %padded = tensor.pad %a low[1, 0] high[1, 0] { + ^bb0(%i: index, %j: index): + tensor.yield %cst : f32 + } : tensor<4x4xf32> to tensor<6x4xf32> + return %padded : tensor<6x4xf32> +} +// CHECK-LABEL: func.func @decline_pad_along_quantized_axis( +// CHECK: %[[DEQ:.+]] = iree_linalg_ext.dequantize_affine +// CHECK: tensor.pad %[[DEQ]] low[1, 0] high[1, 0] + +// ----- + +#id2 = affine_map<(d0, d1) -> (d0, d1)> +#scalar2 = affine_map<(d0, d1) -> ()> + +// A per-channel zero point cannot pad the quantized side: the quantized pad +// would have to read the zero point from its body using the channel index, and +// tensor.pad tiling only supports a padding value defined outside the pad. +func.func @decline_pad_with_per_channel_zero_point(%aq: tensor<4x4xi8>, %sa: tensor<4xf32>, + %za: tensor<4xi8>) -> tensor<4x6xf32> { + %cst = arith.constant 0.000000e+00 : f32 + %init = tensor.empty() : tensor<4x4xf32> + %a = iree_linalg_ext.dequantize_affine + <indexing_maps = [#id2, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d0)>, #id2]> + ins(%aq, %sa, %za : tensor<4x4xi8>, tensor<4xf32>, tensor<4xi8>) + outs(%init : tensor<4x4xf32>) -> tensor<4x4xf32> + %padded = tensor.pad %a low[0, 1] high[0, 1] { + ^bb0(%i: index, %j: index): + tensor.yield %cst : f32 + } : tensor<4x4xf32> to tensor<4x6xf32> + return %padded : tensor<4x6xf32> +} +// CHECK-LABEL: func.func @decline_pad_with_per_channel_zero_point( +// CHECK: %[[DEQ:.+]] = iree_linalg_ext.dequantize_affine +// CHECK: tensor.pad %[[DEQ]] low[0, 1] high[0, 1] + +// ----- + +// Output dimension zero corresponds to quantized input dimension one, so the +// pad on output dimension zero becomes a pad on input dimension one. The +// per-channel scale varies along the other dimension and stays unchanged. +func.func @pad_permuted_output(%q: tensor<2x3xi8>, %scale: tensor<2xf32>, + %zp: i8) -> tensor<5x2xf32> { + %init = tensor.empty() : tensor<3x2xf32> + %d = iree_linalg_ext.dequantize_affine + <indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> (d0)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d1, d0)>]> + ins(%q, %scale, %zp : tensor<2x3xi8>, tensor<2xf32>, i8) + outs(%init : tensor<3x2xf32>) -> tensor<3x2xf32> + %zero = arith.constant 0.0 : f32 + %p = tensor.pad %d low[1, 0] high[1, 0] { + ^bb0(%i: index, %j: index): + tensor.yield %zero : f32 + } : tensor<3x2xf32> to tensor<5x2xf32> + return %p : tensor<5x2xf32> +} +// CHECK-LABEL: func.func @pad_permuted_output( +// CHECK-SAME: %[[Q:[a-zA-Z0-9_]+]]: tensor<2x3xi8> +// CHECK: %[[PADDED:.+]] = tensor.pad %[[Q]] low[0, 1] high[0, 1] +// CHECK: tensor<2x3xi8> to tensor<2x5xi8> +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK-SAME: ins(%[[PADDED]] +// CHECK-SAME: outs(%{{.+}} : tensor<5x2xf32>) -> tensor<5x2xf32> + +// ----- + +// The control function admits a generic transpose of a dequantized value, so it +// is specialized and then fused into the dequantize in the same rewrite. +func.func @generic_transpose_after_dequantize(%q: tensor<2x3xi8>, %scale: f32) + -> tensor<3x2xf32> { + %init = tensor.empty() : tensor<2x3xf32> + %d = iree_linalg_ext.dequantize_affine + <indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]> + ins(%q, %scale : tensor<2x3xi8>, f32) + outs(%init : tensor<2x3xf32>) -> tensor<2x3xf32> + %transpose_init = tensor.empty() : tensor<3x2xf32> + %t = linalg.generic { + indexing_maps = [affine_map<(d0, d1) -> (d1, d0)>, + affine_map<(d0, d1) -> (d0, d1)>], + iterator_types = ["parallel", "parallel"]} + ins(%d : tensor<2x3xf32>) outs(%transpose_init : tensor<3x2xf32>) { + ^bb0(%in: f32, %out: f32): + linalg.yield %in : f32 + } -> tensor<3x2xf32> + return %t : tensor<3x2xf32> +} +// CHECK-LABEL: func.func @generic_transpose_after_dequantize( +// CHECK-NOT: linalg.generic +// CHECK-NOT: linalg.transpose +// CHECK: %[[DEQ:.+]] = iree_linalg_ext.dequantize_affine +// CHECK-SAME: -> tensor<3x2xf32> +// CHECK-NOT: linalg.generic +// CHECK-NOT: linalg.transpose +// CHECK: return %[[DEQ]]
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/propagate_linalg_transpose.mlir b/compiler/src/iree/compiler/GlobalOptimization/test/propagate_linalg_transpose.mlir index ab0d15e..3279af8 100644 --- a/compiler/src/iree/compiler/GlobalOptimization/test/propagate_linalg_transpose.mlir +++ b/compiler/src/iree/compiler/GlobalOptimization/test/propagate_linalg_transpose.mlir
@@ -6,6 +6,62 @@ // RUN: iree-opt --pass-pipeline="builtin.module(util.func(iree-global-opt-propagate-linalg-transpose{enable-edge-reshape-propagation=true}))" %s -o - --split-input-file | FileCheck %s --check-prefix=ENABLE-EDGE-PROP // RUN: iree-opt --pass-pipeline="builtin.module(util.func(iree-global-opt-propagate-linalg-transpose{enable-sink-transpose-through-pad=true}))" --split-input-file %s | FileCheck %s --check-prefix=SINK-PAD +// A transpose fuses into the dequantize's output map. The output map is +// already permuted and the transpose permutation is not its own inverse, so the +// fused map is their composition in transpose order. The input and scale maps +// are unchanged. +util.func public @fuse_transpose_into_dequantize(%q: tensor<2x3x4xi8>, + %scale: tensor<2xf32>) -> tensor<4x3x2xf32> { + %init = tensor.empty() : tensor<3x2x4xf32> + %d = iree_linalg_ext.dequantize_affine + <indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, + affine_map<(d0, d1, d2) -> (d0)>, + affine_map<(d0, d1, d2) -> (d1, d0, d2)>]> + ins(%q, %scale : tensor<2x3x4xi8>, tensor<2xf32>) + outs(%init : tensor<3x2x4xf32>) -> tensor<3x2x4xf32> + %transpose_init = tensor.empty() : tensor<4x3x2xf32> + %t = linalg.transpose ins(%d : tensor<3x2x4xf32>) + outs(%transpose_init : tensor<4x3x2xf32>) permutation = [2, 0, 1] + util.return %t : tensor<4x3x2xf32> +} +// CHECK-DAG: #[[$IN_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)> +// CHECK-DAG: #[[$SCALE_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0)> +// CHECK-DAG: #[[$OUT_MAP:.+]] = affine_map<(d0, d1, d2) -> (d2, d1, d0)> +// CHECK-LABEL: util.func public @fuse_transpose_into_dequantize( +// CHECK-SAME: %[[Q:[a-zA-Z0-9_]+]]: tensor<2x3x4xi8> +// CHECK-NOT: linalg.transpose +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK-SAME: indexing_maps = [#[[$IN_MAP]], #[[$SCALE_MAP]], #[[$OUT_MAP]]] +// CHECK-SAME: ins(%[[Q]] +// CHECK-SAME: -> tensor<4x3x2xf32> +// CHECK-NOT: linalg.transpose + +// ----- + +// A transpose does not fuse when the dequantize has another consumer, since the +// untransposed dequantize would have to stay for that consumer. +util.func public @decline_transpose_into_multi_use_dequantize(%q: tensor<2x3xi8>, + %scale: f32) -> (tensor<3x2xf32>, tensor<2x3xf32>) { + %init = tensor.empty() : tensor<2x3xf32> + %d = iree_linalg_ext.dequantize_affine + <indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]> + ins(%q, %scale : tensor<2x3xi8>, f32) + outs(%init : tensor<2x3xf32>) -> tensor<2x3xf32> + %transpose_init = tensor.empty() : tensor<3x2xf32> + %t = linalg.transpose ins(%d : tensor<2x3xf32>) + outs(%transpose_init : tensor<3x2xf32>) permutation = [1, 0] + util.return %t, %d : tensor<3x2xf32>, tensor<2x3xf32> +} +// CHECK-LABEL: util.func public @decline_transpose_into_multi_use_dequantize( +// CHECK: %[[DEQ:.+]] = iree_linalg_ext.dequantize_affine +// CHECK-SAME: -> tensor<2x3xf32> +// CHECK: %[[T:.+]] = linalg.transpose ins(%[[DEQ]] +// CHECK: util.return %[[T]], %[[DEQ]] + +// ----- + util.func public @specialize_transpose_op(%arg0 : tensor<1x2x3xf32>, %empty : tensor<3x2x1xf32>) -> tensor<3x2x1xf32> { %transposed = linalg.generic {indexing_maps = [