[GlobalOptimization] Rewrite QDQ contractions into integer math (#24902) Rewrites contractions over `iree_linalg_ext.dequantize_affine` operands into integer contractions, integer zero-point corrections, and floating-point scaling. The transformation uses indexing maps and iterator types to handle named and generic contractions, including convolutions. It reduces in integer arithmetic over dimensions where quantization parameters are invariant. For blockwise quantization, it computes integer partial results, applies each block's correction and scales, then accumulates across blocks in floating point. Zero-point correction terms are built as separate computations so constant terms can be evaluated independently. Corrections are applied before conversion to floating point. Scaling and accumulation use at least f32, preserving wider scale or output types and narrowing only the final result. The rewrite requires a plain multiply-add body, zero initialization, and positive static integer reduction extents whose products and corrections fit in signed i32. Unsupported cases retain the original computation. The pass is enabled by default; reassociation changes floating-point rounding. It can be disabled with `--iree-global-opt-enable-qdq-to-integer-math=false`. Builds on the [affine QDQ operations PR #24901](https://github.com/iree-org/iree/pull/24901) and implements the contraction rewrite portion of [Linalg(_ext) native QDQ #24862](https://github.com/iree-org/iree/issues/24862). --- <sub>Stack created with <a href="https://github.com/github/gh-stack">GitHub Stacks CLI</a> • <a href="https://gh.io/stacks-feedback">Give Feedback 💬</a></sub> --------- Signed-off-by: Thomas Ziereis <ziereis@roofline.ai> Signed-off-by: Florian Walbroel <walbroel@roofline.ai> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Florian Walbroel <walbroel@roofline.ai>
diff --git a/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel b/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel index 5b5ac17..e971773 100644 --- a/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel +++ b/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel
@@ -47,11 +47,40 @@ ) iree_compiler_cc_library( + name = "QuantizationUtils", + srcs = ["QuantizationUtils.cpp"], + hdrs = ["QuantizationUtils.h"], + deps = [ + "@llvm-project//llvm:Support", + ], +) + +iree_compiler_cc_library( + name = "QuantizedContraction", + srcs = ["QuantizedContraction.cpp"], + hdrs = ["QuantizedContraction.h"], + visibility = [ + "//compiler/src/iree/compiler/GlobalOptimization:__pkg__", + "//compiler/src/iree/compiler/GlobalOptimization/unittests:__pkg__", + ], + deps = [ + ":QuantizationUtils", + "//compiler/src/iree/compiler/Dialect/LinalgExt/IR", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:LinalgDialect", + "@llvm-project//mlir:Support", + ], +) + +iree_compiler_cc_library( name = "GlobalOptimization", srcs = [ "CleanupNumericNarrowing.cpp", "ConvertBatchMatmulToMatmul.cpp", "ConvertConv2DToImg2Col.cpp", + "ConvertQDQToIntegerMath.cpp", "ConvertStridedContractionToContraction.cpp", "DataLayoutPropagation.cpp", "DecomposeConcat.cpp", @@ -82,6 +111,8 @@ deps = [ ":PassHeaders", ":PassesIncGen", + ":QuantizationUtils", + ":QuantizedContraction", "//compiler/src/iree/compiler/Codegen/Common", "//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect", "//compiler/src/iree/compiler/Dialect/Encoding/IR",
diff --git a/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt b/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt index c3b994d..e1afc48 100644 --- a/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt +++ b/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt
@@ -35,6 +35,36 @@ iree_cc_library( NAME + QuantizationUtils + HDRS + "QuantizationUtils.h" + SRCS + "QuantizationUtils.cpp" + DEPS + LLVMSupport + PUBLIC +) + +iree_cc_library( + NAME + QuantizedContraction + HDRS + "QuantizedContraction.h" + SRCS + "QuantizedContraction.cpp" + DEPS + ::QuantizationUtils + LLVMSupport + MLIRArithDialect + MLIRIR + MLIRLinalgDialect + MLIRSupport + iree::compiler::Dialect::LinalgExt::IR + PUBLIC +) + +iree_cc_library( + NAME GlobalOptimization HDRS "Passes.h" @@ -43,6 +73,7 @@ "CleanupNumericNarrowing.cpp" "ConvertBatchMatmulToMatmul.cpp" "ConvertConv2DToImg2Col.cpp" + "ConvertQDQToIntegerMath.cpp" "ConvertStridedContractionToContraction.cpp" "DataLayoutPropagation.cpp" "DecomposeConcat.cpp" @@ -68,6 +99,8 @@ DEPS ::PassHeaders ::PassesIncGen + ::QuantizationUtils + ::QuantizedContraction LLVMSupport MLIRAffineDialect MLIRArithDialect
diff --git a/compiler/src/iree/compiler/GlobalOptimization/ConvertQDQToIntegerMath.cpp b/compiler/src/iree/compiler/GlobalOptimization/ConvertQDQToIntegerMath.cpp new file mode 100644 index 0000000..320093c --- /dev/null +++ b/compiler/src/iree/compiler/GlobalOptimization/ConvertQDQToIntegerMath.cpp
@@ -0,0 +1,602 @@ +// 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 + +// Rewrites contractions over affine-dequantized operands using +// +// sum_k (Aq - zA)*(Bq - zB) = D - zB*RA - zA*RB + N*zA*zB, +// +// where D is the integer contraction, RA/RB sum the corresponding quantized +// operand, and N is the product of the integer reduction extents. +// Quantization parameters must be invariant over the integer reduction dims. +// Other reduction dims become parallel in D and are reduced after scaling in +// the epilogue. The integer contraction preserves the original contraction's +// iteration dimension numbering; operand sums and the epilogue use their own. +// Integer corrections precede conversion to float to retain cancellation when +// D is too large to be represented exactly in the floating-point type. + +#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h" +#include "iree/compiler/GlobalOptimization/Passes.h" +#include "iree/compiler/GlobalOptimization/QuantizationUtils.h" +#include "iree/compiler/GlobalOptimization/QuantizedContraction.h" +#include "mlir/Dialect/Affine/IR/AffineOps.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/StaticValueUtils.h" +#include "mlir/IR/AffineExpr.h" +#include "mlir/IR/AffineMap.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" + +namespace mlir::iree_compiler::GlobalOptimization { + +#define GEN_PASS_DEF_CONVERTQDQTOINTEGERMATHPASS +#include "iree/compiler/GlobalOptimization/Passes.h.inc" + +using IREE::LinalgExt::DequantizeAffineOp; + +namespace { + +using detail::getExplicitExtentDimsMap; +using detail::getQuantizedContraction; +using detail::QuantizedContraction; +using detail::QuantizedOperand; + +/// Type of every integer intermediate: the contraction, the operand sums, and +/// the zero-point corrections. The legality check bounds all of them against +/// this width. +static IntegerType getAccumulatorType(MLIRContext *context) { + return IntegerType::get(context, kAccumulatorWidth); +} + +//===----------------------------------------------------------------------===// +// Rewrite +//===----------------------------------------------------------------------===// + +static Value buildZeroFilledTensor(OpBuilder &b, Location loc, + ArrayRef<OpFoldResult> sizes, + Type elementType) { + Value empty = tensor::EmptyOp::create(b, loc, sizes, elementType); + Value zero = arith::ConstantOp::create(b, loc, b.getZeroAttr(elementType)); + return linalg::FillOp::create(b, loc, zero, empty).getResult(0); +} + +/// Computes static or dynamic extents in the original contraction's loop order +/// (e.g. [M, N, K] for matmul), including output and reduction dimensions. +/// Uses quantized storage shapes and preserves the output init's extents. +static SmallVector<OpFoldResult> +buildIterationSizes(OpBuilder &b, Location loc, linalg::LinalgOp op, + const QuantizedContraction &detail) { + // Query storage shapes so tensor.dim uses do not keep dequantize ops alive + // after the rewrite. Account for permutations folded into dequantization. + SmallVector<OpFoldResult> flatShapes; + for (OpOperand &operand : op->getOpOperands()) { + Value value = operand.get(); + if (auto dequantize = value.getDefiningOp<DequantizeAffineOp>()) { + SmallVector<OpFoldResult> inputSizes = + tensor::getMixedSizes(b, loc, dequantize.getInput()); + AffineMap outputShapeMap = dequantize.getOutputMap().compose( + inversePermutation(dequantize.getInputMap())); + llvm::append_range(flatShapes, applyPermutationMap<OpFoldResult>( + outputShapeMap, inputSizes)); + } else { + llvm::append_range(flatShapes, tensor::getMixedSizes(b, loc, value)); + } + } + SmallVector<OpFoldResult> sizes = llvm::map_to_vector( + op.getShapesToLoopsMap().getResults(), [&](AffineExpr expr) { + return affine::makeComposedFoldedAffineApply(b, loc, expr, flatShapes); + }); + + // Use init extents for result dimensions to preserve the original result + // type, even when an input provides a static extent for a dynamic result dim. + SmallVector<OpFoldResult> initSizes = + tensor::getMixedSizes(b, loc, op.getDpsInits()[0]); + for (auto [position, expr] : llvm::enumerate(detail.outputMap.getResults())) { + sizes[cast<AffineDimExpr>(expr).getPosition()] = initSizes[position]; + } + return sizes; +} + +//===----------------------------------------------------------------------===// +// Integer contraction +//===----------------------------------------------------------------------===// + +static linalg::GenericOp +buildIntegerContraction(OpBuilder &b, Location loc, linalg::LinalgOp op, + const QuantizedContraction &detail, + ArrayRef<OpFoldResult> iterationSizes) { + DequantizeAffineOp lhsDequantize = detail.lhs.dequantize; + DequantizeAffineOp rhsDequantize = detail.rhs.dequantize; + Value lhs = lhsDequantize.getInput(); + Value rhs = rhsDequantize.getInput(); + bool lhsUnsigned = lhsDequantize.getInputUnsigned(); + bool rhsUnsigned = rhsDequantize.getInputUnsigned(); + + MLIRContext *context = b.getContext(); + const IntegerType accumulatorType = getAccumulatorType(context); + + // Retain a partial result for each block; the epilogue scales and reduces it. + SmallVector<AffineExpr> resultExprs(detail.outputMap.getResults()); + SmallVector<utils::IteratorType> iteratorTypes = op.getIteratorTypesArray(); + for (unsigned dim : detail.floatingReductionDims) { + resultExprs.push_back(getAffineDimExpr(dim, context)); + iteratorTypes[dim] = utils::IteratorType::parallel; + } + AffineMap resultMap = + AffineMap::get(detail.outputMap.getNumDims(), 0, resultExprs, context); + SmallVector<OpFoldResult> resultSizes = + applyPermutationMap<OpFoldResult>(resultMap, iterationSizes); + + Value init = buildZeroFilledTensor(b, loc, resultSizes, accumulatorType); + + SmallVector<AffineMap> maps{detail.lhs.inputMap, detail.rhs.inputMap, + resultMap}; + + // Scalar body: accumulate the product of quantized inputs into an i32 sum. + // a = lhsUnsigned ? extui(lhsElement) : extsi(lhsElement) + // b = rhsUnsigned ? extui(rhsElement) : extsi(rhsElement) + // yield accumulator + a * b + auto genericOp = linalg::GenericOp::create( + b, loc, init.getType(), ValueRange{lhs, rhs}, ValueRange{init}, maps, + iteratorTypes, + [&](OpBuilder &nested, Location nestedLoc, ValueRange args) { + Value lhsValue = + convertScalarToDtype(nested, nestedLoc, args[0], accumulatorType, + /*isUnsignedCast=*/lhsUnsigned); + Value rhsValue = + convertScalarToDtype(nested, nestedLoc, args[1], accumulatorType, + /*isUnsignedCast=*/rhsUnsigned); + Value product = + arith::MulIOp::create(nested, nestedLoc, lhsValue, rhsValue); + Value sum = arith::AddIOp::create(nested, nestedLoc, args[2], product); + linalg::YieldOp::create(nested, nestedLoc, sum); + }); + return genericOp; +} + +//===----------------------------------------------------------------------===// +// Operand sums +//===----------------------------------------------------------------------===// + +/// Builds an integer sum with the given input and output indexing maps. +/// `outputMap` is a projected permutation. Both maps use the domain described +/// by `iterationSizes`. Dimensions absent from both maps are dropped; those +/// used only by `inputMap` are reduced. Input values are extended to +/// `accumulatorType` according to `isUnsigned`. +static Value buildIntegerSum(OpBuilder &b, Location loc, Value input, + AffineMap inputMap, AffineMap outputMap, + ArrayRef<OpFoldResult> iterationSizes, + IntegerType accumulatorType, bool isUnsigned) { + SmallVector<AffineMap> maps{inputMap, outputMap}; + llvm::SmallBitVector unusedDims = getUnusedDimsBitVector(maps); + maps = compressUnusedDims(maps); + + SmallVector<OpFoldResult> sizes; + SmallVector<utils::IteratorType> iteratorTypes; + for (auto [dim, size] : llvm::enumerate(iterationSizes)) { + if (unusedDims.test(dim)) { + continue; + } + sizes.push_back(size); + iteratorTypes.push_back(outputMap.isFunctionOfDim(dim) + ? utils::IteratorType::parallel + : utils::IteratorType::reduction); + } + + // Window expressions such as oh * stride + kh * dilation do not expose kh's + // extent. A shape-only input supplies missing extents without loading data. + SmallVector<Value> inputs{input}; + AffineMap shapeMap = getExplicitExtentDimsMap(maps); + if (shapeMap.getNumResults() != 0) { + SmallVector<OpFoldResult> shapeSizes = + applyPermutationMap<OpFoldResult>(shapeMap, sizes); + inputs.push_back(tensor::EmptyOp::create( + b, loc, shapeSizes, getElementTypeOrSelf(input.getType()))); + maps.insert(maps.begin() + 1, shapeMap); + } + + SmallVector<OpFoldResult> resultSizes = + applyPermutationMap<OpFoldResult>(outputMap, iterationSizes); + Value init = buildZeroFilledTensor(b, loc, resultSizes, accumulatorType); + + auto genericOp = linalg::GenericOp::create( + b, loc, init.getType(), inputs, ValueRange{init}, maps, iteratorTypes, + [&](OpBuilder &nested, Location nestedLoc, ValueRange args) { + Value value = + convertScalarToDtype(nested, nestedLoc, args[0], accumulatorType, + /*isUnsignedCast=*/isUnsigned); + Value sum = + arith::AddIOp::create(nested, nestedLoc, args.back(), value); + linalg::YieldOp::create(nested, nestedLoc, sum); + }); + return genericOp.getResult(0); +} + +//===----------------------------------------------------------------------===// +// Zero-point correction +//===----------------------------------------------------------------------===// + +struct IndexedValue { + Value value; + /// Indexing map from the caller's iteration space to this value's data space. + AffineMap map; +}; + +/// Applies `buildValue` elementwise over the dimensions referenced by `inputs`. +/// Input maps share the domain described by `iterationSizes`. The result has +/// one axis per referenced dimension, in domain order; its map retains the +/// caller's domain so it can be used alongside the input maps. With no +/// referenced dimensions, evaluates the body once and returns a scalar. +static IndexedValue buildElementwise( + OpBuilder &b, Location loc, ArrayRef<IndexedValue> inputs, + ArrayRef<OpFoldResult> iterationSizes, Type elementType, + llvm::function_ref<Value(OpBuilder &, Location, ValueRange)> buildValue) { + SmallVector<Value> values; + SmallVector<AffineMap> maps; + for (IndexedValue input : inputs) { + values.push_back(input.value); + maps.push_back(input.map); + } + AffineMap resultMap = + AffineMap::getMultiDimIdentityMap(iterationSizes.size(), b.getContext()) + .dropResults(getUnusedDimsBitVector(maps)); + unsigned rank = resultMap.getNumResults(); + if (rank == 0) { + for (auto [value, map] : llvm::zip_equal(values, maps)) { + if (isa<RankedTensorType>(value.getType())) { + SmallVector<Value> indices; + for (AffineExpr expr : map.getResults()) { + indices.push_back(arith::ConstantIndexOp::create( + b, loc, cast<AffineConstantExpr>(expr).getValue())); + } + value = tensor::ExtractOp::create(b, loc, value, indices); + } + } + return {buildValue(b, loc, values), resultMap}; + } + SmallVector<OpFoldResult> sizes = + applyPermutationMap<OpFoldResult>(resultMap, iterationSizes); + maps = compressUnusedDims(maps); + maps.push_back(AffineMap::getMultiDimIdentityMap(rank, b.getContext())); + Value init = tensor::EmptyOp::create(b, loc, sizes, elementType); + auto genericOp = linalg::GenericOp::create( + b, loc, init.getType(), values, ValueRange{init}, maps, + SmallVector<utils::IteratorType>(rank, utils::IteratorType::parallel), + [&](OpBuilder &nested, Location nestedLoc, ValueRange args) { + Value result = buildValue(nested, nestedLoc, args.drop_back()); + linalg::YieldOp::create(nested, nestedLoc, result); + }); + return {genericOp.getResult(0), resultMap}; +} + +/// Computes factor * lhs * rhs after converting each input to `resultType`. +/// Widening uses each input's signedness; wider zero-point carriers are +/// narrowed under the requirement that their values fit the storage grid. +/// Broadcasting follows the input maps. +static IndexedValue buildIntegerProduct(OpBuilder &b, Location loc, + IndexedValue lhs, bool lhsUnsigned, + IndexedValue rhs, bool rhsUnsigned, + IntegerType resultType, + ArrayRef<OpFoldResult> iterationSizes, + int64_t factor = 1) { + return buildElementwise( + b, loc, {lhs, rhs}, iterationSizes, resultType, + [&](OpBuilder &nested, Location nestedLoc, ValueRange args) -> Value { + Value lhsValue = convertScalarToDtype(nested, nestedLoc, args[0], + resultType, lhsUnsigned); + Value rhsValue = convertScalarToDtype(nested, nestedLoc, args[1], + resultType, rhsUnsigned); + Value product = + arith::MulIOp::create(nested, nestedLoc, lhsValue, rhsValue); + if (factor == 1) { + return product; + } + Value multiplier = arith::ConstantOp::create( + nested, nestedLoc, nested.getIntegerAttr(resultType, factor)); + return arith::MulIOp::create(nested, nestedLoc, product, multiplier); + }); +} + +/// Builds the correction subtracted from D: zB*sum(Aq) + zA*sum(Bq) - N*zA*zB. +/// Sums span the integer reduction dims, whose extents multiply to N. Each +/// product is a separate computation so constant terms can be evaluated +/// independently of runtime terms and the integer contraction. Absent zero +/// points contribute zero. +static IndexedValue +buildZeroPointCorrection(OpBuilder &b, Location loc, + const QuantizedContraction &detail, + ArrayRef<OpFoldResult> iterationSizes) { + IntegerType accumulatorType = getAccumulatorType(b.getContext()); + AffineMap identity = + AffineMap::getMultiDimIdentityMap(iterationSizes.size(), b.getContext()); + llvm::SmallBitVector reductionDims(identity.getNumDims()); + for (unsigned dim : detail.integerReductionDims) { + reductionDims.set(dim); + } + const QuantizedOperand &lhs = detail.lhs; + DequantizeAffineOp lhsDequantize = lhs.dequantize; + const QuantizedOperand &rhs = detail.rhs; + DequantizeAffineOp rhsDequantize = rhs.dequantize; + SmallVector<IndexedValue, 2> terms; + if (detail.needsLhsSum()) { + AffineMap sumMap = identity.dropResults( + getUnusedDimsBitVector({lhs.inputMap}) | reductionDims); + Value sum = buildIntegerSum(b, loc, lhsDequantize.getInput(), lhs.inputMap, + sumMap, iterationSizes, accumulatorType, + lhsDequantize.getInputUnsigned()); + terms.push_back(buildIntegerProduct( + b, loc, {rhsDequantize.getZeroPoint(), rhs.zeroPointMap}, + rhsDequantize.getZpUnsigned(), {sum, sumMap}, /*rhsUnsigned=*/false, + accumulatorType, iterationSizes)); + } + if (detail.needsRhsSum()) { + AffineMap sumMap = identity.dropResults( + getUnusedDimsBitVector({rhs.inputMap}) | reductionDims); + Value sum = buildIntegerSum(b, loc, rhsDequantize.getInput(), rhs.inputMap, + sumMap, iterationSizes, accumulatorType, + rhsDequantize.getInputUnsigned()); + terms.push_back(buildIntegerProduct( + b, loc, {lhsDequantize.getZeroPoint(), lhs.zeroPointMap}, + lhsDequantize.getZpUnsigned(), {sum, sumMap}, /*rhsUnsigned=*/false, + accumulatorType, iterationSizes)); + } + if (terms.empty()) { + // Keep a uniform epilogue input list. Subsequent cleanup passes are + // expected to fold away this zero correction and the subtraction from D. + Value zero = + arith::ConstantOp::create(b, loc, b.getZeroAttr(accumulatorType)); + return {zero, AffineMap::get(identity.getNumDims(), 0, {}, b.getContext())}; + } + if (terms.size() == 1) { + return terms.front(); + } + + IndexedValue crossTerm = buildIntegerProduct( + b, loc, {lhsDequantize.getZeroPoint(), lhs.zeroPointMap}, + lhsDequantize.getZpUnsigned(), + {rhsDequantize.getZeroPoint(), rhs.zeroPointMap}, + rhsDequantize.getZpUnsigned(), accumulatorType, iterationSizes, + detail.reductionExtent); + return buildElementwise( + b, loc, {terms[0], terms[1], crossTerm}, iterationSizes, accumulatorType, + [&](OpBuilder &nested, Location nestedLoc, ValueRange args) -> Value { + Value sum = arith::AddIOp::create(nested, nestedLoc, args[0], args[1]); + return arith::SubIOp::create(nested, nestedLoc, sum, args[2]); + }); +} + +//===----------------------------------------------------------------------===// +// Scaling epilogue +//===----------------------------------------------------------------------===// + +/// A finite f16 result need not have an f16-representable integer accumulator +/// or scale product. Scaling and block accumulation use at least f32, +/// preserving wider scale and output types as well. +static FloatType getEpilogueComputeType(Type resultType, Type lhsScaleType, + Type rhsScaleType) { + FloatType computeType = Float32Type::get(resultType.getContext()); + for (Type type : {resultType, lhsScaleType, rhsScaleType}) { + auto floatType = cast<FloatType>(type); + if (floatType.getWidth() > computeType.getWidth()) { + computeType = floatType; + } + } + return computeType; +} + +static Value buildOutputConversion(OpBuilder &b, Location loc, Value input, + RankedTensorType outputType, + ArrayRef<OpFoldResult> sizes) { + AffineMap identity = + AffineMap::getMultiDimIdentityMap(outputType.getRank(), b.getContext()); + Type elementType = outputType.getElementType(); + Value output = tensor::EmptyOp::create(b, loc, sizes, elementType); + return linalg::GenericOp::create( + b, loc, outputType, ValueRange{input}, ValueRange{output}, + SmallVector<AffineMap>{identity, identity}, + SmallVector<utils::IteratorType>(outputType.getRank(), + utils::IteratorType::parallel), + [&](OpBuilder &nested, Location nestedLoc, ValueRange args) { + Value result = + convertScalarToDtype(nested, nestedLoc, args[0], elementType, + /*isUnsignedCast=*/false); + linalg::YieldOp::create(nested, nestedLoc, result); + }) + .getResult(0); +} + +/// Subtracts the integer correction before converting and scaling each partial. +/// Dimensions omitted by outputMap are reduced after scaling. All input maps +/// use the domain described by iterationSizes; partials.map is a projected +/// permutation covering the output and remaining reduction dimensions. +static Value buildEpilogue(OpBuilder &b, Location loc, IndexedValue partials, + IndexedValue correction, IndexedValue lhsScale, + IndexedValue rhsScale, AffineMap outputMap, + RankedTensorType outputType, + ArrayRef<OpFoldResult> iterationSizes) { + // Each integer partial becomes one epilogue iteration. The inverse maps that + // iteration back to the caller's domain, filling integer reductions with + // zero. + AffineMap epilogueMap = inverseAndBroadcastProjectedPermutation(partials.map); + unsigned rank = epilogueMap.getNumDims(); + SmallVector<Value> inputs; + SmallVector<AffineMap> maps; + for (IndexedValue input : {partials, correction, lhsScale, rhsScale}) { + inputs.push_back(input.value); + maps.push_back(input.map.compose(epilogueMap)); + } + AffineMap resultMap = outputMap.compose(epilogueMap); + maps.push_back(resultMap); + SmallVector<utils::IteratorType> iteratorTypes; + for (unsigned dim : llvm::seq<unsigned>(rank)) { + iteratorTypes.push_back(resultMap.isFunctionOfDim(dim) + ? utils::IteratorType::parallel + : utils::IteratorType::reduction); + } + + Type outputElementType = outputType.getElementType(); + FloatType computeType = getEpilogueComputeType( + outputElementType, getElementTypeOrSelf(lhsScale.value.getType()), + getElementTypeOrSelf(rhsScale.value.getType())); + bool isReduction = rank > outputType.getRank(); + Type epilogueElementType = isReduction ? computeType : outputElementType; + + // A reduction accumulates into its init; an elementwise epilogue overwrites + // it. + SmallVector<OpFoldResult> resultSizes = + applyPermutationMap<OpFoldResult>(outputMap, iterationSizes); + Value init = + isReduction + ? buildZeroFilledTensor(b, loc, resultSizes, epilogueElementType) + : tensor::EmptyOp::create(b, loc, resultSizes, epilogueElementType) + .getResult(); + + auto genericOp = linalg::GenericOp::create( + b, loc, init.getType(), inputs, ValueRange{init}, maps, iteratorTypes, + [&](OpBuilder &nested, Location nestedLoc, ValueRange args) { + Value corrected = + arith::SubIOp::create(nested, nestedLoc, args[0], args[1]); + Value real = + arith::SIToFPOp::create(nested, nestedLoc, computeType, corrected); + Value lhsScaleValue = convertScalarToDtype( + nested, nestedLoc, args[2], computeType, /*isUnsignedCast=*/false); + Value rhsScaleValue = convertScalarToDtype( + nested, nestedLoc, args[3], computeType, /*isUnsignedCast=*/false); + // Apply scales sequentially so their product cannot overflow or + // underflow before scaling the accumulator. This is still a + // reassociation of the original floating-point contraction: with + // opposing scale magnitudes, real * lhsScaleValue may overflow even + // when the original result is finite. Conversely, integer reduction + // bypasses overflow and rounding in the original dequantized inputs. + // See the pass description for the numerical limitations. + Value partial = + arith::MulFOp::create(nested, nestedLoc, real, lhsScaleValue); + Value result = + arith::MulFOp::create(nested, nestedLoc, partial, rhsScaleValue); + if (isReduction) { + result = + arith::AddFOp::create(nested, nestedLoc, args.back(), result); + } + result = + convertScalarToDtype(nested, nestedLoc, result, epilogueElementType, + /*isUnsignedCast=*/false); + linalg::YieldOp::create(nested, nestedLoc, result); + }); + if (epilogueElementType == outputElementType) { + return genericOp.getResult(0); + } + + // Blockwise accumulation stays wide; only its final result is narrowed. + return buildOutputConversion(b, loc, genericOp.getResult(0), outputType, + resultSizes); +} + +/// Rewrite contractions over affine-dequantized operands as integer +/// contractions with zero-point corrections and a scaling epilogue: +/// +/// For fixed output coordinates and fixed remaining reduction coordinates, +/// k spans the integer reduction dims and N is the product of their extents. +/// Scales sA/sB and zero points zA/zB are invariant over k. Algebraically: +/// +/// ```text +/// A[k] = sA * (Aq[k] - zA) +/// B[k] = sB * (Bq[k] - zB) +/// D = sum_k Aq[k] * Bq[k] +/// RA = sum_k Aq[k] +/// RB = sum_k Bq[k] +/// +/// C = sum_k A[k] * B[k] +/// = sA * sB * sum_k (Aq[k] - zA) * (Bq[k] - zB) +/// = sA * sB * (D - zB * RA - zA * RB + N * zA * zB) +/// ``` +/// +/// Here C is one scaled partial; the epilogue sums these partials over any +/// remaining reduction dims. With no remaining reductions, C is the output. +/// +/// Compose storage and parameter maps into contraction coordinates. Reduce +/// in integer arithmetic where parameters are invariant; retain the other +/// reduction dimensions as partial results for the floating-point epilogue. +/// Windowed operand sums use the same access expressions as the +/// contraction. +struct ConvertQDQToIntegerMath : OpInterfaceRewritePattern<linalg::LinalgOp> { + using OpInterfaceRewritePattern<linalg::LinalgOp>::OpInterfaceRewritePattern; + + LogicalResult matchAndRewrite(linalg::LinalgOp op, + PatternRewriter &rewriter) const override { + // Match a zero-initialized multiply-add contraction of two dequantized + // inputs. Separate reductions with invariant quantization parameters from + // those requiring floating-point accumulation, and prove that the integer + // intermediates fit in i32 before constructing the replacement. + FailureOr<QuantizedContraction> maybeDetail = getQuantizedContraction(op); + if (failed(maybeDetail)) { + return failure(); + } + QuantizedContraction &detail = *maybeDetail; + + // Recover the original loop extents to size the intermediate tensors. + // Query quantized storage shapes so shape computations do not keep the + // dequantization ops alive, while preserving the output init's extents. + Location loc = op.getLoc(); + SmallVector<OpFoldResult> iterationSizes = + buildIterationSizes(rewriter, loc, op, detail); + + // Compute D = sum_k Aq * Bq using i32 arithmetic, where k spans the integer + // reduction dims. Retain other reduction dims as parallel dims so each + // partial result can later receive its own scales and zero-point + // correction. + linalg::GenericOp integerContraction = + buildIntegerContraction(rewriter, loc, op, detail, iterationSizes); + + // Keep the partials' indexing map alongside their values so the epilogue + // can align them with the correction and scales in the original loop space. + IndexedValue partials{integerContraction.getResult(0), + integerContraction.getIndexingMapsArray().back()}; + + // Compute correction = zB * sum_k Aq + zA * sum_k Bq - N * zA * zB, + // where N is the product of the integer reduction extents. Subtracting it + // from D gives sum_k (Aq - zA) * (Bq - zB); absent zero points contribute + // zero. + IndexedValue correction = + buildZeroPointCorrection(rewriter, loc, detail, iterationSizes); + + // Subtract the correction in integer arithmetic before converting to float + // to preserve cancellation in large partials. Apply both scales, sum any + // remaining reduction dims in floating point, and convert to the original + // output element type after accumulation. + Value result = buildEpilogue( + rewriter, loc, partials, correction, + {detail.lhs.dequantize.getScale(), detail.lhs.scaleMap}, + {detail.rhs.dequantize.getScale(), detail.rhs.scaleMap}, + detail.outputMap, cast<RankedTensorType>(op.getDpsInits()[0].getType()), + iterationSizes); + + // Redirect uses to the completed result and erase the original contraction. + rewriter.replaceOp(op, result); + return success(); + } +}; + +struct ConvertQDQToIntegerMathPass final + : impl::ConvertQDQToIntegerMathPassBase<ConvertQDQToIntegerMathPass> { + void getDependentDialects(DialectRegistry ®istry) const override { + registry + .insert<arith::ArithDialect, linalg::LinalgDialect, + tensor::TensorDialect, IREE::LinalgExt::IREELinalgExtDialect>(); + } + + void runOnOperation() override { + MLIRContext *context = &getContext(); + RewritePatternSet patterns(context); + patterns.insert<ConvertQDQToIntegerMath>(context); + if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) { + return signalPassFailure(); + } + } +}; + +} // namespace +} // namespace mlir::iree_compiler::GlobalOptimization
diff --git a/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp b/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp index 049053d..1b556bc 100644 --- a/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp +++ b/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp
@@ -24,6 +24,13 @@ llvm::cl::desc( "Enables reassociation of quantized matmul ops (experimental)."), llvm::cl::init(false)); +static llvm::cl::opt<bool> clConvertQDQToIntegerMath( + "iree-global-opt-enable-qdq-to-integer-math", + llvm::cl::desc("Rewrites contractions over dequantized operands into " + "integer contractions with zero point corrections (" + "changes floating-point rounding and exceptional-value " + "behavior)."), + llvm::cl::init(true)); static llvm::cl::opt<bool> clEnableTransposePropagation( "iree-global-opt-propagate-transposes", llvm::cl::desc( @@ -123,6 +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. + .addPredicatedPass(clConvertQDQToIntegerMath, + createConvertQDQToIntegerMathPass) .addPredicatedPass(transformOptions.useIm2colForConvs, createConvertConv2DToImg2ColPass) .addPass(IREE::Flow::createCanonicalizePass) @@ -200,6 +212,10 @@ clEnableEdgeReshapePropagation; return createPropagateLinalgTransposePass(options); }) + // Retry for contractions whose dequantize producers were exposed by + // reshape/transpose propagation, before selecting integer layouts. + .addPredicatedPass(clConvertQDQToIntegerMath, + createConvertQDQToIntegerMathPass) .addPass(IREE::Flow::createCanonicalizePass) .addPass(mlir::createCSEPass); mainPassManager.addPass(
diff --git a/compiler/src/iree/compiler/GlobalOptimization/Passes.td b/compiler/src/iree/compiler/GlobalOptimization/Passes.td index b05e76d..448cf65 100644 --- a/compiler/src/iree/compiler/GlobalOptimization/Passes.td +++ b/compiler/src/iree/compiler/GlobalOptimization/Passes.td
@@ -101,6 +101,33 @@ let summary = "Expands tensor shape dimensions into SSA values across the program."; } +def ConvertQDQToIntegerMathPass : + InterfacePass<"iree-global-opt-convert-qdq-to-integer-math", "mlir::FunctionOpInterface"> { + let summary = "Rewrites contractions over dequantized operands into integer contractions."; + let description = [{ + Matches a contraction or convolution whose inputs are produced by + `iree_linalg_ext.dequantize_affine` and rewrites it into an integer + contraction plus a zero point correction epilogue: + + C = sA*sB*(D - zB*RA - zA*RB + N*zA*zB) + + Matches a plain floating-point multiply-add body with a zero init. Integer + reduction dimensions must have positive static extents whose products and + zero point corrections fit in signed i32. Dimensions along which a + quantization parameter varies remain floating-point reductions over partial + integer contractions. Unsupported cases retain the original computation. + + Scaling and partial-result accumulation use at least f32, preserving wider + scale and output types, before conversion to the output type. + + Numerical stability is best effort. Reassociation changes floating-point + rounding; no fixed error bound or preservation of signed zero, NaNs, + infinities, or overflow/underflow behavior is guaranteed. Disable this + transformation when reference floating-point semantics are required with + `--iree-global-opt-enable-qdq-to-integer-math=false`. + }]; +} + def FuseDequantizationMatmulPass: InterfacePass<"iree-global-opt-fuse-dequantization-matmul", "mlir::FunctionOpInterface"> { let summary = "Fuses dequantization and matmul linalg.generic ops.";
diff --git a/compiler/src/iree/compiler/GlobalOptimization/QuantizationUtils.cpp b/compiler/src/iree/compiler/GlobalOptimization/QuantizationUtils.cpp new file mode 100644 index 0000000..40de437 --- /dev/null +++ b/compiler/src/iree/compiler/GlobalOptimization/QuantizationUtils.cpp
@@ -0,0 +1,92 @@ +// 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/GlobalOptimization/QuantizationUtils.h" + +#include <algorithm> + +#include "llvm/Support/MathExtras.h" + +namespace mlir::iree_compiler::GlobalOptimization { + +int64_t getStorageMagnitude(unsigned bitWidth, bool isUnsigned) { + if (bitWidth == 0 || bitWidth >= kAccumulatorWidth) { + return 0; + } + // The signed minimum has magnitude one greater than the signed maximum. + return isUnsigned ? llvm::maxUIntN(bitWidth) : llvm::maxIntN(bitWidth) + 1; +} + +std::optional<QuantizedOperandRanges> +getQuantizedOperandRanges(unsigned storageBitWidth, bool storageIsUnsigned, + std::optional<QuantMinMax> quantMinMax, + bool isSymmetric) { + int64_t magnitude = getStorageMagnitude(storageBitWidth, storageIsUnsigned); + if (!magnitude) { + return std::nullopt; + } + QuantMinMax storage = storageIsUnsigned + ? QuantMinMax{0, magnitude} + : QuantMinMax{-magnitude, magnitude - 1}; + QuantMinMax input = quantMinMax.value_or(storage); + QuantMinMax zeroPoint = isSymmetric ? QuantMinMax{0, 0} : storage; + return QuantizedOperandRanges{input, zeroPoint}; +} + +static QuantMinMax multiplyRanges(QuantMinMax lhs, QuantMinMax rhs) { + auto products = {lhs.min * rhs.min, lhs.min * rhs.max, lhs.max * rhs.min, + lhs.max * rhs.max}; + return {*std::min_element(products.begin(), products.end()), + *std::max_element(products.begin(), products.end())}; +} + +static QuantMinMax subtractRanges(QuantMinMax lhs, QuantMinMax rhs) { + return {lhs.min - rhs.max, lhs.max - rhs.min}; +} + +int64_t getMaxReductionExtent(const QuantizedOperandRanges &lhs, + const QuantizedOperandRanges &rhs) { + QuantMinMax a = lhs.input, za = lhs.zeroPoint; + QuantMinMax b = rhs.input, zb = rhs.zeroPoint; + QuantMinMax product = multiplyRanges(a, b); + QuantMinMax lhsCorrection = multiplyRanges(zb, a); + QuantMinMax rhsCorrection = multiplyRanges(za, b); + QuantMinMax crossTerm = multiplyRanges(za, zb); + QuantMinMax correctionSum = {lhsCorrection.min + rhsCorrection.min, + lhsCorrection.max + rhsCorrection.max}; + QuantMinMax correction = subtractRanges(correctionSum, crossTerm); + // Bound the final result using (Aq-zA)*(Bq-zB), retaining the cancellation + // lost by independently bounding D - correction. + QuantMinMax corrected = + multiplyRanges(subtractRanges(a, za), subtractRanges(b, zb)); + + // Storage is at most 31 bits. Centered values have magnitude <= INT32_MAX; + // the products and correction range endpoints above all fit in int64_t. + // Include the operand sums as well as each product and correction stage. + int64_t maxMagnitude = 0; + for (QuantMinMax range : {a, b, product, lhsCorrection, rhsCorrection, + crossTerm, correctionSum, correction, corrected}) { + maxMagnitude = std::max({maxMagnitude, -range.min, range.max}); + } + return maxMagnitude > 0 ? llvm::maxIntN(kAccumulatorWidth) / maxMagnitude : 0; +} + +std::optional<int64_t> +getBoundedReductionExtent(llvm::ArrayRef<int64_t> extents, int64_t maxExtent) { + if (extents.empty() || maxExtent <= 0) { + return std::nullopt; + } + int64_t product = 1; + for (int64_t extent : extents) { + if (extent <= 0 || product > maxExtent / extent) { + return std::nullopt; + } + product *= extent; + } + return product; +} + +} // namespace mlir::iree_compiler::GlobalOptimization
diff --git a/compiler/src/iree/compiler/GlobalOptimization/QuantizationUtils.h b/compiler/src/iree/compiler/GlobalOptimization/QuantizationUtils.h new file mode 100644 index 0000000..8dc282d --- /dev/null +++ b/compiler/src/iree/compiler/GlobalOptimization/QuantizationUtils.h
@@ -0,0 +1,86 @@ +// 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 + +#ifndef IREE_COMPILER_GLOBALOPTIMIZATION_QUANTIZATIONUTILS_H_ +#define IREE_COMPILER_GLOBALOPTIMIZATION_QUANTIZATIONUTILS_H_ + +#include <cstdint> +#include <optional> + +#include "llvm/ADT/ArrayRef.h" + +namespace mlir::iree_compiler::GlobalOptimization { + +static constexpr unsigned kAccumulatorWidth = 32; + +struct QuantMinMax { + int64_t min; + int64_t max; +}; + +/// Maximum magnitude representable by the quantized storage grid. Returns zero +/// when bitWidth is zero or the grid cannot be accumulated in i32. +/// For i8 -> 128 +/// Minimum: -128 -> magnitude 128 +/// Maximum: 127 -> magnitude 127 +int64_t getStorageMagnitude(unsigned bitWidth, bool isUnsigned); + +struct QuantizedOperandRanges { + QuantMinMax input; + QuantMinMax zeroPoint; +}; + +/// Builds input and zero-point ranges from the storage grid. An optional input +/// range must be contained in that grid. Zero points use the full storage grid +/// independently of the input range and carrier type, or [0, 0] when symmetric. +/// Returns nullopt for zero-width storage or widths >= kAccumulatorWidth. +std::optional<QuantizedOperandRanges> +getQuantizedOperandRanges(unsigned storageBitWidth, bool storageIsUnsigned, + std::optional<QuantMinMax> quantMinMax, + bool isSymmetric); + +/// Bounds every integer intermediate per reduction element. Signedness enters +/// through the operand ranges; the same calculation applies to all operands. +/// +/// In general, taking |x| to mean the maximum magnitude over x's range, the +/// triangle inequality gives the sufficient condition +/// +/// N * (|Aq|*|Bq| + |zB|*|Aq| + |zA|*|Bq| + |zA|*|zB|) <= INT32_MAX, +/// +/// which factors exactly into +/// +/// N * (|Aq| + |zA|) * (|Bq| + |zB|) <= INT32_MAX. +/// +/// This bounds the expanded terms and their additions/subtractions, but loses +/// sign information. For unsigned inputs and zero points, D = sum(Aq*Bq), +/// P = zB*sum(Aq), Q = zA*sum(Bq), and R = N*zA*zB are nonnegative. Thus +/// |P+Q-R| <= max(P+Q, R), rather than P+Q+R. Also, |Aq-zA| is bounded by +/// max(|Aq|, |zA|), rather than |Aq|+|zA|. For full unsigned i8 ranges, +/// D, P, Q, R and the centered result have magnitude <= N*255^2, while +/// P+Q and the correction have magnitude <= 2*N*255^2. The sufficient bound +/// is therefore 2*N*255^2 <= INT32_MAX, versus 4*N*255^2 above. +/// +/// Signed values can have opposing signs, so those unsigned inequalities do +/// not apply in general. Propagating intervals through the actual correction +/// stages handles both cases, including mixed signedness, without a special +/// case. The centered expression separately bounds the final result. +/// +/// Operands must have ranges on supported storage grids, as produced by +/// getQuantizedOperandRanges. Returns the largest positive N allowed by these +/// bounds, or zero for an all-zero bound or when no positive N is safe. +int64_t getMaxReductionExtent(const QuantizedOperandRanges &lhs, + const QuantizedOperandRanges &rhs); + +/// Product of a nonempty list of positive static integer reduction extents, +/// provided it does not exceed maxExtent. Returns nullopt for empty lists, +/// nonpositive extents (including dynamic sentinels), or an exceeded limit. +/// Checks the limit before multiplying to avoid overflow. +std::optional<int64_t> +getBoundedReductionExtent(llvm::ArrayRef<int64_t> extents, int64_t maxExtent); + +} // namespace mlir::iree_compiler::GlobalOptimization + +#endif // IREE_COMPILER_GLOBALOPTIMIZATION_QUANTIZATIONUTILS_H_
diff --git a/compiler/src/iree/compiler/GlobalOptimization/QuantizedContraction.cpp b/compiler/src/iree/compiler/GlobalOptimization/QuantizedContraction.cpp new file mode 100644 index 0000000..09047a1 --- /dev/null +++ b/compiler/src/iree/compiler/GlobalOptimization/QuantizedContraction.cpp
@@ -0,0 +1,176 @@ +// 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/GlobalOptimization/QuantizedContraction.h" + +#include "iree/compiler/GlobalOptimization/QuantizationUtils.h" +#include "llvm/ADT/SmallBitVector.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/IR/Matchers.h" + +namespace mlir::iree_compiler::GlobalOptimization::detail { + +static std::optional<QuantizedOperandRanges> +getOperandRanges(DequantizeAffineOp dequantize) { + unsigned bitWidth = + cast<IntegerType>(dequantize.getInputType().getElementType()).getWidth(); + std::optional<QuantMinMax> quantMinMax; + if (auto quantMin = dequantize.getQuantMin()) { + quantMinMax = QuantMinMax{*quantMin, *dequantize.getQuantMax()}; + } + return getQuantizedOperandRanges(bitWidth, dequantize.getInputUnsigned(), + quantMinMax, dequantize.isSymmetric()); +} + +/// Matches exactly two scalar operations: input multiplication and accumulator +/// addition, in either operand order. Casts and other intervening operations +/// prevent matching. +static bool hasPlainFloatMulAddBody(linalg::LinalgOp op) { + Block *body = op.getBlock(); + if (body->getNumArguments() != 3 || body->getOperations().size() != 3) { + return false; + } + auto yield = dyn_cast<linalg::YieldOp>(body->getTerminator()); + if (!yield || yield.getNumOperands() != 1) { + return false; + } + auto addition = yield.getValues()[0].getDefiningOp<arith::AddFOp>(); + if (!addition) { + return false; + } + Value accumulator = body->getArgument(2); + Value product; + if (addition.getLhs() == accumulator) { + product = addition.getRhs(); + } else if (addition.getRhs() == accumulator) { + product = addition.getLhs(); + } else { + return false; + } + auto multiplication = product.getDefiningOp<arith::MulFOp>(); + if (!multiplication) { + return false; + } + Value lhs = body->getArgument(0); + Value rhs = body->getArgument(1); + return (multiplication.getLhs() == lhs && multiplication.getRhs() == rhs) || + (multiplication.getLhs() == rhs && multiplication.getRhs() == lhs); +} + +/// Rebase dequantization indexing maps onto the original contraction's +/// iteration space. The inverse output map maps dequantized data space back to +/// the dequantization iteration space, accounting for folded permutations. +static QuantizedOperand getQuantizedOperand(DequantizeAffineOp dequantize, + AffineMap operandMap) { + AffineMap toDequantizeIteration = + inversePermutation(dequantize.getOutputMap()).compose(operandMap); + return {dequantize, dequantize.getInputMap().compose(toDequantizeIteration), + dequantize.getScaleMap().compose(toDequantizeIteration), + dequantize.isSymmetric() + ? AffineMap() + : dequantize.getZeroPointMap().compose(toDequantizeIteration)}; +} + +FailureOr<QuantizedContraction> getQuantizedContraction(linalg::LinalgOp op) { + if (op.getNumDpsInputs() != 2 || op.getNumDpsInits() != 1 || + !op.hasPureTensorSemantics() || !hasPlainFloatMulAddBody(op)) { + return failure(); + } + auto fill = op.getDpsInits()[0].getDefiningOp<linalg::FillOp>(); + if (!fill || !matchPattern(fill.getInputs()[0], m_AnyZeroFloat())) { + return failure(); + } + + auto lhs = op.getDpsInputs()[0].getDefiningOp<DequantizeAffineOp>(); + auto rhs = op.getDpsInputs()[1].getDefiningOp<DequantizeAffineOp>(); + if (!lhs || !rhs) { + return failure(); + } + QuantizedContraction detail; + detail.lhs = getQuantizedOperand( + lhs, op.getMatchingIndexingMap(op.getDpsInputOperand(0))); + detail.rhs = getQuantizedOperand( + rhs, op.getMatchingIndexingMap(op.getDpsInputOperand(1))); + + detail.outputMap = op.getMatchingIndexingMap(op.getDpsInitOperand(0)); + // The output map has to be a projected permutation for the epilogue to be + // expressible over the output data space and remaining reduction dimensions. + if (!detail.outputMap.isProjectedPermutation()) { + return failure(); + } + + auto lhsRanges = getOperandRanges(lhs); + auto rhsRanges = getOperandRanges(rhs); + if (!lhsRanges || !rhsRanges) { + return failure(); + } + int64_t maxReductionExtent = getMaxReductionExtent(*lhsRanges, *rhsRanges); + if (maxReductionExtent == 0) { + return failure(); + } + + SmallVector<utils::IteratorType> iteratorTypes = op.getIteratorTypesArray(); + SmallVector<int64_t> loopRanges = op.getStaticLoopRanges(); + SmallVector<int64_t> integerReductionExtents; + for (auto [dim, iteratorType] : llvm::enumerate(iteratorTypes)) { + if (iteratorType != utils::IteratorType::reduction) { + continue; + } + // Require each reduction dim to select elements from both inputs that + // contribute to the same output element, as k does in + // C[m, n] = sum_k A[m, k] * B[k, n]. Reject other reduction structures. + if (detail.outputMap.isFunctionOfDim(dim) || + !detail.lhs.inputMap.isFunctionOfDim(dim) || + !detail.rhs.inputMap.isFunctionOfDim(dim)) { + return failure(); + } + // Integer reduction requires scales and zero points to be invariant along + // this dim, so scales can be applied after summation and zero points can + // be factored into the operand-sum corrections. For example: + // sum_k ((Aq[k] - zA) * sA) * ((Bq[k] - zB) * sB) + // = sA * sB * sum_k (Aq[k] - zA) * (Bq[k] - zB). + // If a parameter varies, retain this dim as a parallel dim in the integer + // contraction, then scale and reduce its partial results in the floating- + // point epilogue. For block quantization, the block dim takes this path, + // while reduction within each block can still use integer arithmetic. + // Symmetric operands have no zero-point map to inspect. + if (detail.lhs.scaleMap.isFunctionOfDim(dim) || + detail.rhs.scaleMap.isFunctionOfDim(dim) || + (detail.lhs.zeroPointMap && + detail.lhs.zeroPointMap.isFunctionOfDim(dim)) || + (detail.rhs.zeroPointMap && + detail.rhs.zeroPointMap.isFunctionOfDim(dim))) { + detail.floatingReductionDims.push_back(dim); + continue; + } + detail.integerReductionDims.push_back(dim); + integerReductionExtents.push_back(loopRanges[dim]); + } + auto reductionExtent = + getBoundedReductionExtent(integerReductionExtents, maxReductionExtent); + if (!reductionExtent) { + return failure(); + } + detail.reductionExtent = *reductionExtent; + return detail; +} + +AffineMap getExplicitExtentDimsMap(ArrayRef<AffineMap> maps) { + unsigned rank = maps.front().getNumDims(); + llvm::SmallBitVector namedDims(rank); + for (AffineMap map : maps) { + for (AffineExpr expr : map.getResults()) { + if (auto dimExpr = dyn_cast<AffineDimExpr>(expr)) { + namedDims.set(dimExpr.getPosition()); + } + } + } + + return AffineMap::getMultiDimIdentityMap(rank, maps.front().getContext()) + .dropResults(namedDims); +} + +} // namespace mlir::iree_compiler::GlobalOptimization::detail
diff --git a/compiler/src/iree/compiler/GlobalOptimization/QuantizedContraction.h b/compiler/src/iree/compiler/GlobalOptimization/QuantizedContraction.h new file mode 100644 index 0000000..4d08666 --- /dev/null +++ b/compiler/src/iree/compiler/GlobalOptimization/QuantizedContraction.h
@@ -0,0 +1,70 @@ +// 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 + +#ifndef IREE_COMPILER_GLOBALOPTIMIZATION_QUANTIZEDCONTRACTION_H_ +#define IREE_COMPILER_GLOBALOPTIMIZATION_QUANTIZEDCONTRACTION_H_ + +#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/IR/AffineMap.h" + +// Internal analysis support for ConvertQDQToIntegerMath. +namespace mlir::iree_compiler::GlobalOptimization::detail { + +using IREE::LinalgExt::DequantizeAffineOp; + +/// One dequantized operand. All map domains are the original contraction's +/// iteration space. +struct QuantizedOperand { + DequantizeAffineOp dequantize; + /// Range: quantized input data space. + AffineMap inputMap; + /// Range: scale data space. + AffineMap scaleMap; + /// Range: zero-point data space. Null if absent. + AffineMap zeroPointMap; +}; + +/// A matched contraction with a nonempty, statically bounded integer reduction. +/// Construction succeeds only for a plain multiply-add body and a zero init. +/// Dimension indices refer to the original contraction's iteration space. +struct QuantizedContraction { + QuantizedOperand lhs; + QuantizedOperand rhs; + /// Original contraction's indexing map for the output. + AffineMap outputMap; + /// Reduction dims over which all quantization parameters are invariant. + SmallVector<unsigned> integerReductionDims; + /// Reduction dims retained as parallel dims until the scaling epilogue. + SmallVector<unsigned> floatingReductionDims; + /// Product of integer reduction extents, bounded for the i32 accumulator. + int64_t reductionExtent = 1; + + // Each operand sum is multiplied by the opposite operand's zero point. + bool needsLhsSum() const { + auto quant = rhs.dequantize; + return !quant.isSymmetric(); + } + + bool needsRhsSum() const { + auto quant = lhs.dequantize; + return !quant.isSymmetric(); + } +}; + +/// Analyzes a verified contraction without modifying its IR. Fails when the +/// contraction cannot be rewritten using a statically bounded integer +/// reduction. +FailureOr<QuantizedContraction> getQuantizedContraction(linalg::LinalgOp op); + +/// Returns a map selecting domain dimensions that do not appear as bare results +/// in any of `maps`, a nonempty list with a shared domain. These dimensions +/// need explicit extents because Linalg cannot infer them from operand shapes. +AffineMap getExplicitExtentDimsMap(ArrayRef<AffineMap> maps); + +} // namespace mlir::iree_compiler::GlobalOptimization::detail + +#endif // IREE_COMPILER_GLOBALOPTIMIZATION_QUANTIZEDCONTRACTION_H_
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel b/compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel index a8481d7..cf1a0a0 100644 --- a/compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel +++ b/compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel
@@ -20,6 +20,11 @@ "cleanup_numeric_narrowing.mlir", "conv2d_to_img2col.mlir", "convert_batch_matmul_to_matmul.mlir", + "convert_qdq_to_integer_math_algebra.mlir", + "convert_qdq_to_integer_math_maps.mlir", + "convert_qdq_to_integer_math_negative.mlir", + "convert_qdq_to_integer_math_pipeline.mlir", + "convert_qdq_to_integer_math_scale_range.mlir", "data_layout_propagation.mlir", "demote_contraction_inputs.mlir", "detach_elementwise_from_named_ops.mlir",
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt b/compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt index 95ae4f9..de43a37 100644 --- a/compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt +++ b/compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt
@@ -17,6 +17,11 @@ "cleanup_numeric_narrowing.mlir" "conv2d_to_img2col.mlir" "convert_batch_matmul_to_matmul.mlir" + "convert_qdq_to_integer_math_algebra.mlir" + "convert_qdq_to_integer_math_maps.mlir" + "convert_qdq_to_integer_math_negative.mlir" + "convert_qdq_to_integer_math_pipeline.mlir" + "convert_qdq_to_integer_math_scale_range.mlir" "data_layout_propagation.mlir" "demote_contraction_inputs.mlir" "detach_elementwise_from_named_ops.mlir"
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/convert_qdq_to_integer_math_algebra.mlir b/compiler/src/iree/compiler/GlobalOptimization/test/convert_qdq_to_integer_math_algebra.mlir new file mode 100644 index 0000000..e2d6391 --- /dev/null +++ b/compiler/src/iree/compiler/GlobalOptimization/test/convert_qdq_to_integer_math_algebra.mlir
@@ -0,0 +1,846 @@ +// RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(func.func(iree-global-opt-convert-qdq-to-integer-math))" %s | FileCheck %s + +// Correction terms, widening, and the arithmetic of the scaling epilogue. +// Indexing lives in convert_qdq_to_integer_math_maps.mlir and the contractions +// the rewrite declines in convert_qdq_to_integer_math_negative.mlir. + +// The quantized grid bounds a zero point independently of its carrier type. +// PT2E represents scalar zero points as i64 even for i8 quantization. +func.func @wide_zero_point_carrier(%aq: tensor<4x8xi8>, %a_s: f32, + %a_z: i64, %bq: tensor<8x16xi8>, %b_s: f32) -> tensor<4x16xf32> { + %a_i = tensor.empty() : tensor<4x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s, %a_z : tensor<4x8xi8>, f32, i64) + outs(%a_i : tensor<4x8xf32>) -> tensor<4x8xf32> + %b_i = tensor.empty() : tensor<8x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s : tensor<8x16xi8>, f32) + outs(%b_i : tensor<8x16xf32>) -> tensor<8x16xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<4x16xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<4x16xf32>) -> tensor<4x16xf32> + %result = linalg.matmul ins(%a, %b : tensor<4x8xf32>, tensor<8x16xf32>) + outs(%init : tensor<4x16xf32>) -> tensor<4x16xf32> + return %result : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @wide_zero_point_carrier( +// CHECK-NOT: iree_linalg_ext.dequantize_affine +// CHECK: linalg.generic +// CHECK-SAME: outs(%{{.+}} : tensor<4x16xi32>) + +//===----------------------------------------------------------------------===// +// Zero-point corrections +// +// sum_k (Aq - zA)(Bq - zB) = D - zB*RA - zA*RB + N*zA*zB. Which terms exist +// depends on which side carries a zero point, so these four cases add one term +// at a time. The first spells out the shared epilogue arithmetic the rest +// inherit. +// +// D = sum_k Aq[k] * Bq[k] +// RA = sum_k Aq[k] +// RB = sum_k Bq[k] +// N = number of terms in the sum (reduction extent) +// +//===----------------------------------------------------------------------===// + +// Both symmetric: nothing to correct, so the epilogue only applies the scales. +// C = sA*sB*D = sA*sB*sum_k Aq[k] * Bq[k] +func.func @zp_neither(%aq: tensor<4x8xi8>, %a_s: f32, %bq: tensor<8x16xi8>, %b_s: f32) -> tensor<4x16xf32> { + %a_i = tensor.empty() : tensor<4x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s : tensor<4x8xi8>, f32) + outs(%a_i : tensor<4x8xf32>) -> tensor<4x8xf32> + %b_i = tensor.empty() : tensor<8x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s : tensor<8x16xi8>, f32) + outs(%b_i : tensor<8x16xf32>) -> tensor<8x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<4x16xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.matmul ins(%a, %b : tensor<4x8xf32>, tensor<8x16xf32>) + outs(%f : tensor<4x16xf32>) -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @zp_neither( +// CHECK-SAME: %[[AQ:[a-zA-Z0-9_]+]]: tensor<4x8xi8>, %[[AS:[a-zA-Z0-9_]+]]: f32, +// CHECK-SAME: %[[BQ:[a-zA-Z0-9_]+]]: tensor<8x16xi8>, %[[BS:[a-zA-Z0-9_]+]]: f32 +// CHECK: %[[ZERO:.+]] = arith.constant 0 : i32 +// CHECK-NOT: iree_linalg_ext.dequantize_affine +// CHECK: %[[D:[a-zA-Z0-9_]+]] = linalg.generic +// CHECK-SAME: ins(%[[AQ]], %[[BQ]] : tensor<4x8xi8>, tensor<8x16xi8>) +// CHECK-SAME: outs(%{{.+}} : tensor<4x16xi32>) +// CHECK: ^bb0(%[[L:[a-zA-Z0-9_]+]]: i8, %[[R:[a-zA-Z0-9_]+]]: i8, %[[ACC:[a-zA-Z0-9_]+]]: i32): +// CHECK: %[[LE:.+]] = arith.extsi %[[L]] : i8 to i32 +// CHECK: %[[RE:.+]] = arith.extsi %[[R]] : i8 to i32 +// CHECK: %[[P:.+]] = arith.muli %[[LE]], %[[RE]] +// CHECK: %[[SUM:.+]] = arith.addi %[[ACC]], %[[P]] : i32 +// CHECK-NEXT: linalg.yield %[[SUM]] : i32 + +// The correction input is scalar i32 zero. This is the scaling epilogue: it +// yields (float(D - 0) * sA) * sB. Subsequent scalar inlining and +// canonicalization fold away the subtraction. +// CHECK: %[[RESULT:.+]] = linalg.generic {{.*}} ins(%[[D]], %[[ZERO]], %[[AS]], %[[BS]] : tensor<4x16xi32>, i32, f32, f32) +// CHECK-SAME: outs(%{{.+}} : tensor<4x16xf32>) +// CHECK: ^bb0(%[[ED:[a-zA-Z0-9_]+]]: i32, %[[EC:[a-zA-Z0-9_]+]]: i32, %[[SA:[a-zA-Z0-9_]+]]: f32, %[[SB:[a-zA-Z0-9_]+]]: f32, %[[ACC:[a-zA-Z0-9_]+]]: f32): +// CHECK: %[[CORRECTED:.+]] = arith.subi %[[ED]], %[[EC]] : i32 +// CHECK: %[[REAL:.+]] = arith.sitofp %[[CORRECTED]] : i32 to f32 +// CHECK: %[[PARTIAL:.+]] = arith.mulf %[[REAL]], %[[SA]] : f32 +// CHECK: %[[SCALED:.+]] = arith.mulf %[[PARTIAL]], %[[SB]] : f32 +// CHECK-NEXT: linalg.yield %[[SCALED]] : f32 +// CHECK: return %[[RESULT]] : tensor<4x16xf32> + +// ----- + +// A zero point on the lhs produces the *rhs* sum and the `- zA*RB` term. +func.func @zp_lhs_only(%aq: tensor<4x8xi8>, %a_s: f32, %a_z: i8, %bq: tensor<8x16xi8>, %b_s: f32) -> tensor<4x16xf32> { + %a_i = tensor.empty() : tensor<4x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s, %a_z : tensor<4x8xi8>, f32, i8) + outs(%a_i : tensor<4x8xf32>) -> tensor<4x8xf32> + %b_i = tensor.empty() : tensor<8x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s : tensor<8x16xi8>, f32) + outs(%b_i : tensor<8x16xf32>) -> tensor<8x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<4x16xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.matmul ins(%a, %b : tensor<4x8xf32>, tensor<8x16xf32>) + outs(%f : tensor<4x16xf32>) -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @zp_lhs_only( +// CHECK-SAME: %[[AQ:[a-zA-Z0-9_]+]]: tensor<4x8xi8>, %[[AS:[a-zA-Z0-9_]+]]: f32, %[[AZ:[a-zA-Z0-9_]+]]: i8, +// CHECK-SAME: %[[BQ:[a-zA-Z0-9_]+]]: tensor<8x16xi8>, %[[BS:[a-zA-Z0-9_]+]]: f32 +// CHECK-NOT: iree_linalg_ext.dequantize_affine +// CHECK: %[[D:[a-zA-Z0-9_]+]] = linalg.generic +// CHECK-SAME: ins(%[[AQ]], %[[BQ]] : +// CHECK-SAME: outs(%{{.+}} : tensor<4x16xi32>) + +// One sum, and it reduces the *rhs*, which is what pairs with the lhs zero +// point. +// CHECK: %[[RB:[a-zA-Z0-9_]+]] = linalg.generic +// CHECK-SAME: ins(%[[BQ]] : tensor<8x16xi8>) +// CHECK-SAME: outs(%{{.+}} : tensor<16xi32>) + +// It is a plain accumulation of the widened operand, nothing more. +// CHECK: ^bb0(%[[SV:[a-zA-Z0-9_]+]]: i8, %[[SACC:[a-zA-Z0-9_]+]]: i32): +// CHECK: %[[SVE:.+]] = arith.extsi %[[SV]] : i8 to i32 +// CHECK: %[[SUM:.+]] = arith.addi %[[SACC]], %[[SVE]] +// CHECK: linalg.yield %[[SUM]] + +// The zero-point product is independently evaluable and retains the sum shape. +// CHECK: %[[TERM:.+]] = linalg.generic {{.*}} ins(%[[AZ]], %[[RB]] : i8, tensor<16xi32>) +// CHECK-SAME: outs(%{{.+}} : tensor<16xi32>) +// CHECK: ^bb0(%[[Z:[a-zA-Z0-9_]+]]: i8, %[[SUM:[a-zA-Z0-9_]+]]: i32, %{{.+}}: i32): +// CHECK: %[[WIDE:.+]] = arith.extsi %[[Z]] : i8 to i32 +// CHECK: %[[PRODUCT:.+]] = arith.muli %[[WIDE]], %[[SUM]] +// CHECK: linalg.yield %[[PRODUCT]] +// CHECK: linalg.generic {{.*}} ins(%[[D]], %[[TERM]], %[[AS]], %[[BS]] : +// CHECK-SAME: outs(%{{.+}} : tensor<4x16xf32>) +// CHECK: ^bb0(%[[ED:[a-zA-Z0-9_]+]]: i32, %[[EC:[a-zA-Z0-9_]+]]: i32, %[[SA:[a-zA-Z0-9_]+]]: f32, %[[SB:[a-zA-Z0-9_]+]]: f32, %[[ACC:[a-zA-Z0-9_]+]]: f32): +// CHECK: %[[CORRECTED:.+]] = arith.subi %[[ED]], %[[EC]] : i32 +// CHECK: %[[REAL:.+]] = arith.sitofp %[[CORRECTED]] : i32 to f32 + +// ----- + +// The mirror image: a zero point on the rhs produces the lhs sum. +func.func @zp_rhs_only(%aq: tensor<4x8xi8>, %a_s: f32, %bq: tensor<8x16xi8>, %b_s: f32, %b_z: i8) -> tensor<4x16xf32> { + %a_i = tensor.empty() : tensor<4x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s : tensor<4x8xi8>, f32) + outs(%a_i : tensor<4x8xf32>) -> tensor<4x8xf32> + %b_i = tensor.empty() : tensor<8x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s, %b_z : tensor<8x16xi8>, f32, i8) + outs(%b_i : tensor<8x16xf32>) -> tensor<8x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<4x16xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.matmul ins(%a, %b : tensor<4x8xf32>, tensor<8x16xf32>) + outs(%f : tensor<4x16xf32>) -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @zp_rhs_only( +// CHECK-SAME: %[[AQ:[a-zA-Z0-9_]+]]: tensor<4x8xi8>, %[[AS:[a-zA-Z0-9_]+]]: f32, +// CHECK-SAME: %[[BQ:[a-zA-Z0-9_]+]]: tensor<8x16xi8>, %[[BS:[a-zA-Z0-9_]+]]: f32, %[[BZ:[a-zA-Z0-9_]+]]: i8 +// CHECK: %[[D:.+]] = linalg.generic {{.*}} ins(%[[AQ]], %[[BQ]] : +// The mirror pairing uses the lhs sum and rhs zero point. +// CHECK: %[[RA:.+]] = linalg.generic {{.*}} ins(%[[AQ]] : tensor<4x8xi8>) +// CHECK: %[[TERM:.+]] = linalg.generic {{.*}} ins(%[[BZ]], %[[RA]] : i8, tensor<4xi32>) +// CHECK: linalg.generic {{.*}} ins(%[[D]], %[[TERM]], %[[AS]], %[[BS]] : + +// ----- + +// Both asymmetric: combine the two products as TA + TB - N*zA*zB. +// TA = zB * RA, TB = zA * RB +// The cross term is needed only when both zero points are present; N is 8. +func.func @zp_both(%aq: tensor<4x8xi8>, %a_s: f32, %a_z: i8, %bq: tensor<8x16xi8>, %b_s: f32, %b_z: i8) -> tensor<4x16xf32> { + %a_i = tensor.empty() : tensor<4x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s, %a_z : tensor<4x8xi8>, f32, i8) + outs(%a_i : tensor<4x8xf32>) -> tensor<4x8xf32> + %b_i = tensor.empty() : tensor<8x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s, %b_z : tensor<8x16xi8>, f32, i8) + outs(%b_i : tensor<8x16xf32>) -> tensor<8x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<4x16xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.matmul ins(%a, %b : tensor<4x8xf32>, tensor<8x16xf32>) + outs(%f : tensor<4x16xf32>) -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @zp_both( +// CHECK-SAME: %[[AQ:[a-zA-Z0-9_]+]]: tensor<4x8xi8>, %[[AS:[a-zA-Z0-9_]+]]: f32, %[[ZA:[a-zA-Z0-9_]+]]: i8, +// CHECK-SAME: %[[BQ:[a-zA-Z0-9_]+]]: tensor<8x16xi8>, %[[BS:[a-zA-Z0-9_]+]]: f32, %[[ZB:[a-zA-Z0-9_]+]]: i8 +// CHECK: %[[N:.+]] = arith.constant 8 : i32 +// CHECK: %[[D:.+]] = linalg.generic {{.*}} ins(%[[AQ]], %[[BQ]] : +// CHECK-SAME: outs(%{{.+}} : tensor<4x16xi32>) +// CHECK: %[[RA:.+]] = linalg.generic {{.*}} ins(%[[AQ]] : tensor<4x8xi8>) +// CHECK-SAME: outs(%{{.+}} : tensor<4xi32>) +// CHECK: %[[TA:.+]] = linalg.generic {{.*}} ins(%[[ZB]], %[[RA]] : i8, tensor<4xi32>) +// CHECK-SAME: outs(%{{.+}} : tensor<4xi32>) +// CHECK: %[[RB:.+]] = linalg.generic {{.*}} ins(%[[BQ]] : tensor<8x16xi8>) +// CHECK-SAME: outs(%{{.+}} : tensor<16xi32>) +// CHECK: %[[TB:.+]] = linalg.generic {{.*}} ins(%[[ZA]], %[[RB]] : i8, tensor<16xi32>) +// CHECK-SAME: outs(%{{.+}} : tensor<16xi32>) + +// The cross term stays scalar; only combining row/column terms needs M x N. +// CHECK: %[[WA:.+]] = arith.extsi %[[ZA]] : i8 to i32 +// CHECK: %[[WB:.+]] = arith.extsi %[[ZB]] : i8 to i32 +// CHECK: %[[ZZ:.+]] = arith.muli %[[WA]], %[[WB]] +// CHECK: %[[TZ:.+]] = arith.muli %[[ZZ]], %[[N]] +// CHECK: %[[CORRECTION:.+]] = linalg.generic {{.*}} ins(%[[TA]], %[[TB]], %[[TZ]] : +// CHECK-SAME: outs(%{{.+}} : tensor<4x16xi32>) +// CHECK: ^bb0(%[[A:[a-zA-Z0-9_]+]]: i32, %[[B:[a-zA-Z0-9_]+]]: i32, %[[Z:[a-zA-Z0-9_]+]]: i32, %{{.+}}: i32): +// CHECK: %[[SUM:.+]] = arith.addi %[[A]], %[[B]] +// CHECK: %[[OFFSET:.+]] = arith.subi %[[SUM]], %[[Z]] +// CHECK: linalg.yield %[[OFFSET]] +// CHECK: linalg.generic {{.*}} ins(%[[D]], %[[CORRECTION]], %[[AS]], %[[BS]] : +// CHECK-SAME: outs(%{{.+}} : tensor<4x16xf32>) + +//===----------------------------------------------------------------------===// +// Widening +// +// Storage signedness and zero-point signedness are independent per operand. The +// accumulator is signed by construction whatever the inputs were, so the final +// conversion to float is always signed. +//===----------------------------------------------------------------------===// + +// ----- + +// Unsigned storage and unsigned zero points: every widening is a zero extension. +// The final conversion stays signed, because the accumulator is signed by +// construction whatever the inputs were. +func.func @both_unsigned(%aq: tensor<4x8xi8>, %a_s: f32, %a_z: i8, %bq: tensor<8x16xi8>, %b_s: f32, %b_z: i8) -> tensor<4x16xf32> { + %a_i = tensor.empty() : tensor<4x8xf32> + %a = iree_linalg_ext.dequantize_affine + {input_unsigned, zp_unsigned, + indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s, %a_z : tensor<4x8xi8>, f32, i8) + outs(%a_i : tensor<4x8xf32>) -> tensor<4x8xf32> + %b_i = tensor.empty() : tensor<8x16xf32> + %b = iree_linalg_ext.dequantize_affine + {input_unsigned, zp_unsigned, + indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s, %b_z : tensor<8x16xi8>, f32, i8) + outs(%b_i : tensor<8x16xf32>) -> tensor<8x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<4x16xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.matmul ins(%a, %b : tensor<4x8xf32>, tensor<8x16xf32>) + outs(%f : tensor<4x16xf32>) -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @both_unsigned( +// CHECK-SAME: %[[AQ:[a-zA-Z0-9_]+]]: tensor<4x8xi8>, %[[AS:[a-zA-Z0-9_]+]]: f32, %[[AZ:[a-zA-Z0-9_]+]]: i8, +// CHECK-SAME: %[[BQ:[a-zA-Z0-9_]+]]: tensor<8x16xi8>, %[[BS:[a-zA-Z0-9_]+]]: f32, %[[BZ:[a-zA-Z0-9_]+]]: i8 +// CHECK: linalg.generic {{.*}} ins(%[[AQ]], %[[BQ]] : +// CHECK-NEXT: ^bb0(%[[L:[a-zA-Z0-9_]+]]: i8, %[[R:[a-zA-Z0-9_]+]]: i8, %{{.+}}: i32): +// CHECK-DAG: arith.extui %[[L]] : i8 to i32 +// CHECK-DAG: arith.extui %[[R]] : i8 to i32 +// CHECK: %[[RA:.+]] = linalg.generic {{.*}} ins(%[[AQ]] : +// CHECK-NEXT: ^bb0(%[[A:[a-zA-Z0-9_]+]]: i8, %{{.+}}: i32): +// CHECK: arith.extui %[[A]] : i8 to i32 +// CHECK: linalg.generic {{.*}} ins(%[[BZ]], %[[RA]] : +// CHECK-NEXT: ^bb0(%[[ZB:[a-zA-Z0-9_]+]]: i8, %{{.+}}: i32, %{{.+}}: i32): +// CHECK: arith.extui %[[ZB]] : i8 to i32 +// CHECK: %[[RB:.+]] = linalg.generic {{.*}} ins(%[[BQ]] : +// CHECK-NEXT: ^bb0(%[[B:[a-zA-Z0-9_]+]]: i8, %{{.+}}: i32): +// CHECK: arith.extui %[[B]] : i8 to i32 +// CHECK: linalg.generic {{.*}} ins(%[[AZ]], %[[RB]] : +// CHECK-NEXT: ^bb0(%[[ZA:[a-zA-Z0-9_]+]]: i8, %{{.+}}: i32, %{{.+}}: i32): +// CHECK: arith.extui %[[ZA]] : i8 to i32 +// CHECK-DAG: arith.extui %[[AZ]] : i8 to i32 +// CHECK-DAG: arith.extui %[[BZ]] : i8 to i32 +// CHECK: arith.sitofp + +// ----- + +// The combination PT2E emits on x86: unsigned activations against signed +// weights. The two operands must widen differently within the same op. +func.func @lhs_unsigned_rhs_signed(%aq: tensor<4x8xi8>, %a_s: f32, %a_z: i8, %bq: tensor<8x16xi8>, %b_s: f32, %b_z: i8) -> tensor<4x16xf32> { + %a_i = tensor.empty() : tensor<4x8xf32> + %a = iree_linalg_ext.dequantize_affine + {input_unsigned, zp_unsigned, + indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s, %a_z : tensor<4x8xi8>, f32, i8) + outs(%a_i : tensor<4x8xf32>) -> tensor<4x8xf32> + %b_i = tensor.empty() : tensor<8x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s, %b_z : tensor<8x16xi8>, f32, i8) + outs(%b_i : tensor<8x16xf32>) -> tensor<8x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<4x16xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.matmul ins(%a, %b : tensor<4x8xf32>, tensor<8x16xf32>) + outs(%f : tensor<4x16xf32>) -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @lhs_unsigned_rhs_signed( +// CHECK-SAME: %[[AQ:[a-zA-Z0-9_]+]]: tensor<4x8xi8>, %[[AS:[a-zA-Z0-9_]+]]: f32, %[[AZ:[a-zA-Z0-9_]+]]: i8, +// CHECK-SAME: %[[BQ:[a-zA-Z0-9_]+]]: tensor<8x16xi8>, %[[BS:[a-zA-Z0-9_]+]]: f32, %[[BZ:[a-zA-Z0-9_]+]]: i8 +// CHECK: linalg.generic {{.*}} ins(%[[AQ]], %[[BQ]] : +// CHECK-NEXT: ^bb0(%[[L:[a-zA-Z0-9_]+]]: i8, %[[R:[a-zA-Z0-9_]+]]: i8, %{{.+}}: i32): +// CHECK-DAG: arith.extui %[[L]] : i8 to i32 +// CHECK-DAG: arith.extsi %[[R]] : i8 to i32 +// CHECK: %[[RA:.+]] = linalg.generic {{.*}} ins(%[[AQ]] : +// CHECK-NEXT: ^bb0(%[[A:[a-zA-Z0-9_]+]]: i8, %{{.+}}: i32): +// CHECK: arith.extui %[[A]] : i8 to i32 +// CHECK: linalg.generic {{.*}} ins(%[[BZ]], %[[RA]] : +// CHECK-NEXT: ^bb0(%[[ZB:[a-zA-Z0-9_]+]]: i8, %{{.+}}: i32, %{{.+}}: i32): +// CHECK: arith.extsi %[[ZB]] : i8 to i32 +// CHECK: %[[RB:.+]] = linalg.generic {{.*}} ins(%[[BQ]] : +// CHECK-NEXT: ^bb0(%[[B:[a-zA-Z0-9_]+]]: i8, %{{.+}}: i32): +// CHECK: arith.extsi %[[B]] : i8 to i32 +// CHECK: linalg.generic {{.*}} ins(%[[AZ]], %[[RB]] : +// CHECK-NEXT: ^bb0(%[[ZA:[a-zA-Z0-9_]+]]: i8, %{{.+}}: i32, %{{.+}}: i32): +// CHECK: arith.extui %[[ZA]] : i8 to i32 +// CHECK-DAG: arith.extui %[[AZ]] : i8 to i32 +// CHECK-DAG: arith.extsi %[[BZ]] : i8 to i32 + +// ----- + +// `input_unsigned` and `zp_unsigned` are independent: an unsigned tensor may +// carry a signed zero point of a wider type, which is what ONNX and TFLite +// importers produce. The data zero extends while the zero point sign extends. +func.func @zp_signedness_differs_from_storage(%aq: tensor<4x8xi8>, %a_s: f32, %a_z: i16, %bq: tensor<8x16xi8>, %b_s: f32, %b_z: i8) -> tensor<4x16xf32> { + %a_i = tensor.empty() : tensor<4x8xf32> + %a = iree_linalg_ext.dequantize_affine + {input_unsigned, + indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s, %a_z : tensor<4x8xi8>, f32, i16) + outs(%a_i : tensor<4x8xf32>) -> tensor<4x8xf32> + %b_i = tensor.empty() : tensor<8x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s, %b_z : tensor<8x16xi8>, f32, i8) + outs(%b_i : tensor<8x16xf32>) -> tensor<8x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<4x16xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.matmul ins(%a, %b : tensor<4x8xf32>, tensor<8x16xf32>) + outs(%f : tensor<4x16xf32>) -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @zp_signedness_differs_from_storage( +// CHECK-SAME: %[[AQ:[a-zA-Z0-9_]+]]: tensor<4x8xi8>, %[[AS:[a-zA-Z0-9_]+]]: f32, %[[AZ:[a-zA-Z0-9_]+]]: i16, +// CHECK-SAME: %[[BQ:[a-zA-Z0-9_]+]]: tensor<8x16xi8>, %[[BS:[a-zA-Z0-9_]+]]: f32, %[[BZ:[a-zA-Z0-9_]+]]: i8 +// CHECK: linalg.generic {{.*}} ins(%[[AQ]], %[[BQ]] : +// CHECK-NEXT: ^bb0(%[[L:[a-zA-Z0-9_]+]]: i8, %[[R:[a-zA-Z0-9_]+]]: i8, %{{.+}}: i32): +// CHECK-DAG: arith.extui %[[L]] : i8 to i32 +// CHECK-DAG: arith.extsi %[[R]] : i8 to i32 +// CHECK: %[[RA:.+]] = linalg.generic {{.*}} ins(%[[AQ]] : +// CHECK-NEXT: ^bb0(%[[A:[a-zA-Z0-9_]+]]: i8, %{{.+}}: i32): +// CHECK: arith.extui %[[A]] : i8 to i32 +// CHECK: linalg.generic {{.*}} ins(%[[BZ]], %[[RA]] : +// CHECK-NEXT: ^bb0(%[[ZB:[a-zA-Z0-9_]+]]: i8, %{{.+}}: i32, %{{.+}}: i32): +// CHECK: arith.extsi %[[ZB]] : i8 to i32 +// CHECK: %[[RB:.+]] = linalg.generic {{.*}} ins(%[[BQ]] : +// CHECK-NEXT: ^bb0(%[[B:[a-zA-Z0-9_]+]]: i8, %{{.+}}: i32): +// CHECK: arith.extsi %[[B]] : i8 to i32 +// CHECK: linalg.generic {{.*}} ins(%[[AZ]], %[[RB]] : +// CHECK-NEXT: ^bb0(%[[ZA:[a-zA-Z0-9_]+]]: i16, %{{.+}}: i32, %{{.+}}: i32): +// CHECK: arith.extsi %[[ZA]] : i16 to i32 +// CHECK-DAG: arith.extsi %[[AZ]] : i16 to i32 +// CHECK-DAG: arith.extsi %[[BZ]] : i8 to i32 + +//===----------------------------------------------------------------------===// +// Contraction body +// +// Matching accepts either operand order in the multiply and in the add. +// Bodies that do anything else are in +// convert_qdq_to_integer_math_negative.mlir. +//===----------------------------------------------------------------------===// + +// ----- + +// Both the multiply and the add have their operands the other way round. The +// rebuilt body is in canonical order regardless of the order matched. +func.func @commuted_mul_add(%aq: tensor<1x8xi8>, %sa: f32, + %bq: tensor<8x1xi8>, %sb: f32) -> tensor<1x1xf32> { + %ai = tensor.empty() : tensor<1x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0,d1)->(d0,d1)>, affine_map<(d0,d1)->()>, affine_map<(d0,d1)->(d0,d1)>]} + ins(%aq, %sa : tensor<1x8xi8>, f32) + outs(%ai : tensor<1x8xf32>) -> tensor<1x8xf32> + %bi = tensor.empty() : tensor<8x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0,d1)->(d0,d1)>, affine_map<(d0,d1)->()>, affine_map<(d0,d1)->(d0,d1)>]} + ins(%bq, %sb : tensor<8x1xi8>, f32) + outs(%bi : tensor<8x1xf32>) -> tensor<8x1xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32> + %result = linalg.generic {indexing_maps = [affine_map<(m,n,k)->(m,k)>, affine_map<(m,n,k)->(k,n)>, affine_map<(m,n,k)->(m,n)>], iterator_types = ["parallel", "parallel", "reduction"]} + ins(%a, %b : tensor<1x8xf32>, tensor<8x1xf32>) outs(%init : tensor<1x1xf32>) { + ^bb0(%lhs: f32, %rhs: f32, %acc: f32): + %product = arith.mulf %rhs, %lhs : f32 + %sum = arith.addf %product, %acc : f32 + linalg.yield %sum : f32 + } -> tensor<1x1xf32> + return %result : tensor<1x1xf32> +} +// CHECK-LABEL: func.func @commuted_mul_add( +// CHECK: ^bb0(%[[L:[a-zA-Z0-9_]+]]: i8, %[[R:[a-zA-Z0-9_]+]]: i8, %[[ACC:[a-zA-Z0-9_]+]]: i32): +// CHECK-DAG: %[[LE:.+]] = arith.extsi %[[L]] +// CHECK-DAG: %[[RE:.+]] = arith.extsi %[[R]] +// CHECK: %[[P:.+]] = arith.muli %[[LE]], %[[RE]] +// CHECK: arith.addi %[[ACC]], %[[P]] + +//===----------------------------------------------------------------------===// +// Reduction form +// +// Reduction dims split by whether a quantization parameter varies along them. +// Only the arithmetic consequence is checked here; the resulting maps and +// iterator types belong to the maps file. +//===----------------------------------------------------------------------===// + +// ----- + +// Blockwise: K is expanded to (G, L) and the lhs parameters vary along G, so the +// contraction is integer within a block and the residual sum across blocks +// happens in floating point. The observable is the trailing addf in the +// epilogue, which no other form produces. +func.func @form_partial_reduces_in_float(%aq: tensor<4x2x4xi8>, %a_s: tensor<4x2xf32>, %a_z: tensor<4x2xi8>, + %bq: tensor<2x4x16xi8>, %b_s: f32) -> tensor<4x16xf32> { + %a_i = tensor.empty() : tensor<4x2x4xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, + affine_map<(d0, d1, d2) -> (d0, d1)>, + affine_map<(d0, d1, d2) -> (d0, d1)>, + affine_map<(d0, d1, d2) -> (d0, d1, d2)>]} + ins(%aq, %a_s, %a_z : tensor<4x2x4xi8>, tensor<4x2xf32>, tensor<4x2xi8>) + outs(%a_i : tensor<4x2x4xf32>) -> tensor<4x2x4xf32> + %b_i = tensor.empty() : tensor<2x4x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, + affine_map<(d0, d1, d2) -> ()>, + affine_map<(d0, d1, d2) -> (d0, d1, d2)>]} + ins(%bq, %b_s : tensor<2x4x16xi8>, f32) + outs(%b_i : tensor<2x4x16xf32>) -> tensor<2x4x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<4x16xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.contract + indexing_maps = [affine_map<(m, n, g, l) -> (m, g, l)>, + affine_map<(m, n, g, l) -> (g, l, n)>, + affine_map<(m, n, g, l) -> (m, n)>] + ins(%a, %b : tensor<4x2x4xf32>, tensor<2x4x16xf32>) + outs(%f : tensor<4x16xf32>) -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @form_partial_reduces_in_float( +// CHECK: %[[D:.+]] = linalg.generic {{.*}} outs(%{{.+}} : tensor<4x16x2xi32>) +// CHECK: %[[CORRECTION:.+]] = linalg.generic {{.*}} outs(%{{.+}} : tensor<4x16x2xi32>) +// CHECK: linalg.generic {{.*}} ins(%[[D]], %[[CORRECTION]], %{{[^ ,)]+}}, %{{[^ ,)]+}} : +// CHECK-SAME: outs(%{{.+}} : tensor<4x16xf32>) +// CHECK-NEXT: ^bb0(%{{.+}}: i32, %{{.+}}: i32, %{{.+}}: f32, %{{.+}}: f32, %[[ACC:[a-zA-Z0-9_]+]]: f32): +// CHECK: %[[PARTIAL:.+]] = arith.mulf +// CHECK: %[[SCALED:.+]] = arith.mulf %[[PARTIAL]], %{{.+}} +// CHECK: %[[TOTAL:.+]] = arith.addf %[[ACC]], %[[SCALED]] : f32 +// CHECK: linalg.yield %[[TOTAL]] : f32 + +//===----------------------------------------------------------------------===// +// Scale precision +// +// Scaling uses at least f32, preserving wider scale and output types. The result +// is narrowed once, after every partial has accumulated, so a low-precision +// result never rounds a partial sum. +//===----------------------------------------------------------------------===// + +// ----- + +// An f16 contraction with f32 scales, as a half precision PT2E model produces. +// Scaling stays in f32 until the final result is narrowed. +func.func @wider_scales(%aq: tensor<4x8xi8>, %a_s: f32, %bq: tensor<8x16xi8>, + %b_s: f32) -> tensor<4x16xf16> { + %a_i = tensor.empty() : tensor<4x8xf16> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s : tensor<4x8xi8>, f32) + outs(%a_i : tensor<4x8xf16>) -> tensor<4x8xf16> + %b_i = tensor.empty() : tensor<8x16xf16> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s : tensor<8x16xi8>, f32) + outs(%b_i : tensor<8x16xf16>) -> tensor<8x16xf16> + %cst = arith.constant 0.000000e+00 : f16 + %e = tensor.empty() : tensor<4x16xf16> + %f = linalg.fill ins(%cst : f16) outs(%e : tensor<4x16xf16>) -> tensor<4x16xf16> + %c = linalg.matmul ins(%a, %b : tensor<4x8xf16>, tensor<8x16xf16>) + outs(%f : tensor<4x16xf16>) -> tensor<4x16xf16> + return %c : tensor<4x16xf16> +} +// CHECK-LABEL: func.func @wider_scales( +// CHECK: linalg.generic {{.*}} outs(%{{.+}} : tensor<4x16xf16>) +// CHECK: arith.sitofp %{{.+}} : i32 to f32 +// CHECK: %[[PARTIAL:.+]] = arith.mulf %{{.+}}, %{{.+}} : f32 +// CHECK: %[[SCALED:.+]] = arith.mulf %[[PARTIAL]], %{{.+}} : f32 +// CHECK: %[[NARROW:.+]] = arith.truncf %[[SCALED]] : f32 to f16 +// CHECK: linalg.yield %[[NARROW]] : f16 + +// ----- + +// The same rule with the scale wider than the result rather than between the +// result and f32: the epilogue computes in f64 and truncates once. +func.func @double_scales(%aq: tensor<4x8xi8>, %a_s: f64, %bq: tensor<8x16xi8>, + %b_s: f64) -> tensor<4x16xf16> { + %a_i = tensor.empty() : tensor<4x8xf16> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s : tensor<4x8xi8>, f64) + outs(%a_i : tensor<4x8xf16>) -> tensor<4x8xf16> + %b_i = tensor.empty() : tensor<8x16xf16> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s : tensor<8x16xi8>, f64) + outs(%b_i : tensor<8x16xf16>) -> tensor<8x16xf16> + %cst = arith.constant 0.000000e+00 : f16 + %e = tensor.empty() : tensor<4x16xf16> + %f = linalg.fill ins(%cst : f16) outs(%e : tensor<4x16xf16>) -> tensor<4x16xf16> + %c = linalg.matmul ins(%a, %b : tensor<4x8xf16>, tensor<8x16xf16>) + outs(%f : tensor<4x16xf16>) -> tensor<4x16xf16> + return %c : tensor<4x16xf16> +} +// CHECK-LABEL: func.func @double_scales( +// CHECK: arith.sitofp %{{.+}} : i32 to f64 +// CHECK: %[[PARTIAL:.+]] = arith.mulf %{{.+}}, %{{.+}} : f64 +// CHECK: %[[SCALED:.+]] = arith.mulf %[[PARTIAL]], %{{.+}} : f64 +// CHECK: %[[NARROW:.+]] = arith.truncf %[[SCALED]] : f64 to f16 +// CHECK: linalg.yield %[[NARROW]] : f16 + +// ----- + +// Widen the half scale before multiplication and narrow only after all blocks +// have accumulated in f32. +func.func @half_partial_reduction(%aq: tensor<4x2x4xi8>, %a_s: tensor<4x2xf16>, %a_z: tensor<4x2xi8>, + %bq: tensor<2x4x16xi8>, %b_s: f32) -> tensor<4x16xf16> { + %a_i = tensor.empty() : tensor<4x2x4xf16> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, + affine_map<(d0, d1, d2) -> (d0, d1)>, + affine_map<(d0, d1, d2) -> (d0, d1)>, + affine_map<(d0, d1, d2) -> (d0, d1, d2)>]} + ins(%aq, %a_s, %a_z : tensor<4x2x4xi8>, tensor<4x2xf16>, tensor<4x2xi8>) + outs(%a_i : tensor<4x2x4xf16>) -> tensor<4x2x4xf16> + %b_i = tensor.empty() : tensor<2x4x16xf16> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, + affine_map<(d0, d1, d2) -> ()>, + affine_map<(d0, d1, d2) -> (d0, d1, d2)>]} + ins(%bq, %b_s : tensor<2x4x16xi8>, f32) + outs(%b_i : tensor<2x4x16xf16>) -> tensor<2x4x16xf16> + %cst = arith.constant 0.000000e+00 : f16 + %e = tensor.empty() : tensor<4x16xf16> + %f = linalg.fill ins(%cst : f16) outs(%e : tensor<4x16xf16>) -> tensor<4x16xf16> + %c = linalg.contract + indexing_maps = [affine_map<(m, n, g, l) -> (m, g, l)>, + affine_map<(m, n, g, l) -> (g, l, n)>, + affine_map<(m, n, g, l) -> (m, n)>] + ins(%a, %b : tensor<4x2x4xf16>, tensor<2x4x16xf16>) + outs(%f : tensor<4x16xf16>) -> tensor<4x16xf16> + return %c : tensor<4x16xf16> +} +// CHECK-LABEL: func.func @half_partial_reduction( +// This epilogue corrects and scales each integer block result, then sums all +// blocks in f32. ACCUMULATED is the complete result, with no block dimension. +// CHECK: %[[ACCUMULATED:.+]] = linalg.generic {{.*}} outs(%{{.+}} : tensor<4x16xf32>) +// CHECK-NEXT: ^bb0(%{{.+}}: i32, %{{.+}}: i32, %[[SA:[a-zA-Z0-9_]+]]: f16, %[[SB:[a-zA-Z0-9_]+]]: f32, %{{.+}}: f32): +// CHECK: %[[WIDE_SA:.+]] = arith.extf %[[SA]] : f16 to f32 +// CHECK: %[[PARTIAL:.+]] = arith.mulf %{{.+}}, %[[WIDE_SA]] : f32 +// CHECK-NEXT: arith.mulf %[[PARTIAL]], %[[SB]] : f32 +// CHECK-NOT: arith.truncf +// CHECK: %[[TOTAL:.+]] = arith.addf %{{.+}}, %{{.+}} : f32 +// CHECK-NOT: arith.truncf +// CHECK: linalg.yield %[[TOTAL]] : f32 +// A separate elementwise generic narrows each completed result to f16 once. +// CHECK: linalg.generic {{.*}} ins(%[[ACCUMULATED]] : tensor<4x16xf32>) +// CHECK-SAME: outs(%{{.+}} : tensor<4x16xf16>) +// CHECK-NEXT: ^bb0(%[[V:[a-zA-Z0-9_]+]]: f32, %{{.+}}: f16): +// CHECK: %[[NARROW:.+]] = arith.truncf %[[V]] : f32 to f16 +// CHECK: linalg.yield %[[NARROW]] : f16 + +//===----------------------------------------------------------------------===// +// Accumulator bound +// +// The bound accounts for the input and zero-point ranges, so the same shape +// can be admissible symmetric and inadmissible asymmetric. The rejected half +// of each pair below is in +// convert_qdq_to_integer_math_negative.mlir. +//===----------------------------------------------------------------------===// + +// ----- + +// i8 by i8 with K = 65536. Symmetric, so a single term, and 7+7+16+0+1 = 31 bits +// fit. Making both sides asymmetric costs two more bits and no longer fits: see +// @acc_asymmetric_too_deep. +func.func @acc_symmetric_fits_deep(%aq: tensor<4x65536xi8>, %a_s: f32, %bq: tensor<65536x16xi8>, %b_s: f32) -> tensor<4x16xf32> { + %a_i = tensor.empty() : tensor<4x65536xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s : tensor<4x65536xi8>, f32) + outs(%a_i : tensor<4x65536xf32>) -> tensor<4x65536xf32> + %b_i = tensor.empty() : tensor<65536x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s : tensor<65536x16xi8>, f32) + outs(%b_i : tensor<65536x16xf32>) -> tensor<65536x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<4x16xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.matmul ins(%a, %b : tensor<4x65536xf32>, tensor<65536x16xf32>) + outs(%f : tensor<4x16xf32>) -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @acc_symmetric_fits_deep( +// CHECK-NOT: iree_linalg_ext.dequantize_affine +// CHECK: linalg.generic + +// ----- + +// Mixed storage widths are bounded per operand rather than by the wider of the +// two: 15+7 magnitude bits plus a shallow reduction still fits. +func.func @acc_mixed_width_fits(%aq: tensor<4x8xi16>, %a_s: f32, %a_z: i16, %bq: tensor<8x16xi8>, %b_s: f32) -> tensor<4x16xf32> { + %a_i = tensor.empty() : tensor<4x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s, %a_z : tensor<4x8xi16>, f32, i16) + outs(%a_i : tensor<4x8xf32>) -> tensor<4x8xf32> + %b_i = tensor.empty() : tensor<8x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s : tensor<8x16xi8>, f32) + outs(%b_i : tensor<8x16xf32>) -> tensor<8x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<4x16xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.matmul ins(%a, %b : tensor<4x8xf32>, tensor<8x16xf32>) + outs(%f : tensor<4x16xf32>) -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @acc_mixed_width_fits( +// CHECK-NOT: iree_linalg_ext.dequantize_affine +// CHECK: linalg.generic +// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<4x8xi16>, tensor<8x16xi8>) +// CHECK-SAME: outs(%{{.+}} : tensor<4x16xi32>) + +// ----- + +// The deepest signed-i8 symmetric reduction below the positive i32 endpoint. +// One element more overflows it: see @acc_positive_endpoint. +func.func @acc_last_safe_depth(%aq: tensor<1x131071xi8>, %sa: f32, + %bq: tensor<131071x1xi8>, %sb: f32) -> tensor<1x1xf32> { + %ai = tensor.empty() : tensor<1x131071xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0,d1)->(d0,d1)>, affine_map<(d0,d1)->()>, affine_map<(d0,d1)->(d0,d1)>]} + ins(%aq, %sa : tensor<1x131071xi8>, f32) + outs(%ai : tensor<1x131071xf32>) -> tensor<1x131071xf32> + %bi = tensor.empty() : tensor<131071x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0,d1)->(d0,d1)>, affine_map<(d0,d1)->()>, affine_map<(d0,d1)->(d0,d1)>]} + ins(%bq, %sb : tensor<131071x1xi8>, f32) + outs(%bi : tensor<131071x1xf32>) -> tensor<131071x1xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32> + %result = linalg.matmul ins(%a, %b : tensor<1x131071xf32>, tensor<131071x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + return %result : tensor<1x1xf32> +} +// CHECK-LABEL: func.func @acc_last_safe_depth( +// CHECK: tensor<1x1xi32> +// CHECK-NOT: linalg.matmul + +// ----- + +// P+Q is bounded by 2*N*255^2. Accept floor(INT32_MAX/(2*255^2)), +// beyond the old four-term magnitude bound of 8256. +func.func @unsigned_reduction_16512(%aq: tensor<1x16512xi8>, %a_s: f32, %a_z: i8, %bq: tensor<16512x1xi8>, %b_s: f32, %b_z: i8) -> tensor<1x1xf32> { + %a_i = tensor.empty() : tensor<1x16512xf32> + %a = iree_linalg_ext.dequantize_affine + {input_unsigned, zp_unsigned, + indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s, %a_z : tensor<1x16512xi8>, f32, i8) + outs(%a_i : tensor<1x16512xf32>) -> tensor<1x16512xf32> + %b_i = tensor.empty() : tensor<16512x1xf32> + %b = iree_linalg_ext.dequantize_affine + {input_unsigned, zp_unsigned, + indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s, %b_z : tensor<16512x1xi8>, f32, i8) + outs(%b_i : tensor<16512x1xf32>) -> tensor<16512x1xf32> + %cst = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<1x1xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<1x1xf32>) -> tensor<1x1xf32> + %c = linalg.matmul ins(%a, %b : tensor<1x16512xf32>, tensor<16512x1xf32>) + outs(%f : tensor<1x1xf32>) -> tensor<1x1xf32> + return %c : tensor<1x1xf32> +} +// CHECK-LABEL: func.func @unsigned_reduction_16512( +// CHECK-NOT: iree_linalg_ext.dequantize_affine +// CHECK: linalg.generic +// CHECK-SAME: outs(%{{.+}} : tensor<1x1xi32>) + +// ----- + +// The centered product is in [-255^2, 255^2] and bounds all intermediates. +// Accept the last safe depth. +func.func @signed_reduction_33025(%aq: tensor<1x33025xi8>, %a_s: f32, %a_z: i8, %bq: tensor<33025x1xi8>, %b_s: f32, %b_z: i8) -> tensor<1x1xf32> { + %a_i = tensor.empty() : tensor<1x33025xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s, %a_z : tensor<1x33025xi8>, f32, i8) + outs(%a_i : tensor<1x33025xf32>) -> tensor<1x33025xf32> + %b_i = tensor.empty() : tensor<33025x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s, %b_z : tensor<33025x1xi8>, f32, i8) + outs(%b_i : tensor<33025x1xf32>) -> tensor<33025x1xf32> + %cst = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<1x1xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<1x1xf32>) -> tensor<1x1xf32> + %c = linalg.matmul ins(%a, %b : tensor<1x33025xf32>, tensor<33025x1xf32>) + outs(%f : tensor<1x1xf32>) -> tensor<1x1xf32> + return %c : tensor<1x1xf32> +} +// CHECK-LABEL: func.func @signed_reduction_33025( +// CHECK-NOT: iree_linalg_ext.dequantize_affine +// CHECK: linalg.generic +// CHECK-SAME: outs(%{{.+}} : tensor<1x1xi32>) + +// ----- + +// The correction interval is [-97665, 97410], larger than the centered product. +// Accept the last safe depth. +func.func @mixed_reduction_21988(%aq: tensor<1x21988xi8>, %a_s: f32, %a_z: i8, %bq: tensor<21988x1xi8>, %b_s: f32, %b_z: i8) -> tensor<1x1xf32> { + %a_i = tensor.empty() : tensor<1x21988xf32> + %a = iree_linalg_ext.dequantize_affine + {input_unsigned, zp_unsigned, + indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s, %a_z : tensor<1x21988xi8>, f32, i8) + outs(%a_i : tensor<1x21988xf32>) -> tensor<1x21988xf32> + %b_i = tensor.empty() : tensor<21988x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s, %b_z : tensor<21988x1xi8>, f32, i8) + outs(%b_i : tensor<21988x1xf32>) -> tensor<21988x1xf32> + %cst = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<1x1xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<1x1xf32>) -> tensor<1x1xf32> + %c = linalg.matmul ins(%a, %b : tensor<1x21988xf32>, tensor<21988x1xf32>) + outs(%f : tensor<1x1xf32>) -> tensor<1x1xf32> + return %c : tensor<1x1xf32> +} +// CHECK-LABEL: func.func @mixed_reduction_21988( +// CHECK-NOT: iree_linalg_ext.dequantize_affine +// CHECK: linalg.generic +// CHECK-SAME: outs(%{{.+}} : tensor<1x1xi32>)
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/convert_qdq_to_integer_math_maps.mlir b/compiler/src/iree/compiler/GlobalOptimization/test/convert_qdq_to_integer_math_maps.mlir new file mode 100644 index 0000000..ffabb99 --- /dev/null +++ b/compiler/src/iree/compiler/GlobalOptimization/test/convert_qdq_to_integer_math_maps.mlir
@@ -0,0 +1,907 @@ +// RUN: iree-opt --split-input-file --mlir-print-local-scope --pass-pipeline="builtin.module(func.func(iree-global-opt-convert-qdq-to-integer-math))" %s | FileCheck %s + +// Indexing maps, iterator types, and result shapes. Arithmetic lives in +// convert_qdq_to_integer_math_algebra.mlir and the contractions the rewrite +// declines in convert_qdq_to_integer_math_negative.mlir. +// +// Checks identify ops by emission order (integer contraction, optional lhs/rhs +// sums and correction terms, scaling epilogue) and `outs` type. + +//===----------------------------------------------------------------------===// +// Storage maps composed through dequantization +//===----------------------------------------------------------------------===// + +// Baseline: plain matmul, everything an identity or a projected permutation. +// The integer contraction inherits the contraction's maps and iterator types +// unchanged, and the rhs sum drops M. +func.func @matmul_baseline(%aq: tensor<4x8xi8>, %sa: tensor<f32>, %za: tensor<i8>, + %bq: tensor<8x16xi8>, %sb: tensor<16xf32>) -> tensor<4x16xf32> { + %ainit = tensor.empty() : tensor<4x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %sa, %za : tensor<4x8xi8>, tensor<f32>, tensor<i8>) + outs(%ainit : tensor<4x8xf32>) -> tensor<4x8xf32> + %binit = tensor.empty() : tensor<8x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> (d1)>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %sb : tensor<8x16xi8>, tensor<16xf32>) + outs(%binit : tensor<8x16xf32>) -> tensor<8x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %init = tensor.empty() : tensor<4x16xf32> + %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.matmul ins(%a, %b : tensor<4x8xf32>, tensor<8x16xf32>) + outs(%fill : tensor<4x16xf32>) -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @matmul_baseline( +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, +// CHECK-SAME: affine_map<(d0, d1, d2) -> (d2, d1)>, +// CHECK-SAME: affine_map<(d0, d1, d2) -> (d0, d1)>] +// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"] +// CHECK-SAME: outs(%{{.+}} : tensor<4x16xi32>) +// The rhs sum is stated over the two dims the rhs depends on, N and K, reducing +// only K. M never enters it. +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1) -> (d1, d0)>, +// CHECK-SAME: affine_map<(d0, d1) -> (d0)>] +// CHECK-SAME: iterator_types = ["parallel", "reduction"] +// CHECK-SAME: outs(%{{.+}} : tensor<16xi32>) + +// ----- + +// The rhs is stored as (N, K), with one scale per row. Its first dimension +// maps to the output column N, so the scale is constant across the K reduction. +// The integer contraction reads Bq[n, k] directly through its indexing map; +// no transpose operation is needed. +func.func @matmul_transpose_b(%aq: tensor<4x8xi8>, %sa: tensor<f32>, + %bq: tensor<16x8xi8>, %sb: tensor<16xf32>) -> tensor<4x16xf32> { + %ainit = tensor.empty() : tensor<4x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %sa : tensor<4x8xi8>, tensor<f32>) + outs(%ainit : tensor<4x8xf32>) -> tensor<4x8xf32> + %binit = tensor.empty() : tensor<16x8xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> (d0)>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %sb : tensor<16x8xi8>, tensor<16xf32>) + outs(%binit : tensor<16x8xf32>) -> tensor<16x8xf32> + %cst = arith.constant 0.000000e+00 : f32 + %init = tensor.empty() : tensor<4x16xf32> + %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, + affine_map<(d0, d1, d2) -> (d1, d2)>, + affine_map<(d0, d1, d2) -> (d0, d1)>], + iterator_types = ["parallel", "parallel", "reduction"]} + ins(%a, %b : tensor<4x8xf32>, tensor<16x8xf32>) + outs(%fill : tensor<4x16xf32>) { + ^bb0(%in: f32, %in_0: f32, %out: f32): + %m = arith.mulf %in, %in_0 : f32 + %s = arith.addf %out, %m : f32 + linalg.yield %s : f32 + } -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @matmul_transpose_b( +// CHECK: linalg.generic +// CHECK-SAME: affine_map<(d0, d1, d2) -> (d1, d2)> +// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<4x8xi8>, tensor<16x8xi8>) +// CHECK-SAME: outs(%{{.+}} : tensor<4x16xi32>) + +// ----- + +// A transpose folded into the dequantize: its output map is a non-identity +// permutation, so inverting it is what carries the contraction's operand map +// onto the 16x8 quantized data. The contraction reads the untransposed tensor +// through a permuted map, and no transpose op survives. +func.func @transposing_dequantize(%aq: tensor<4x8xi8>, %sa: tensor<f32>, %za: tensor<i8>, + %bq: tensor<16x8xi8>, %sb: tensor<16xf32>) -> tensor<4x16xf32> { + %ainit = tensor.empty() : tensor<4x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %sa, %za : tensor<4x8xi8>, tensor<f32>, tensor<i8>) + outs(%ainit : tensor<4x8xf32>) -> tensor<4x8xf32> + %binit = tensor.empty() : tensor<8x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> (d0)>, + affine_map<(d0, d1) -> (d1, d0)>]} + ins(%bq, %sb : tensor<16x8xi8>, tensor<16xf32>) + outs(%binit : tensor<8x16xf32>) -> tensor<8x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %init = tensor.empty() : tensor<4x16xf32> + %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.matmul ins(%a, %b : tensor<4x8xf32>, tensor<8x16xf32>) + outs(%fill : tensor<4x16xf32>) -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @transposing_dequantize( +// CHECK-NOT: linalg.transpose +// The contraction reads the 16x8 tensor directly, with the permutation absorbed +// into its rhs map. +// CHECK: linalg.generic +// CHECK-SAME: affine_map<(d0, d1, d2) -> (d1, d2)> +// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<4x8xi8>, tensor<16x8xi8>) +// The sum sees the same untransposed layout, so N is its leading dim. +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, +// CHECK-SAME: affine_map<(d0, d1) -> (d0)>] +// CHECK-SAME: outs(%{{.+}} : tensor<16xi32>) + +// ----- + +// A batch dim is a parallel dim like any other: the sum keeps it because the rhs +// depends on it, and drops only M. +func.func @batch_matmul(%aq: tensor<2x4x8xi8>, %sa: tensor<f32>, %za: tensor<i8>, + %bq: tensor<2x8x16xi8>, %sb: tensor<f32>) -> tensor<2x4x16xf32> { + %ainit = tensor.empty() : tensor<2x4x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, + affine_map<(d0, d1, d2) -> ()>, + affine_map<(d0, d1, d2) -> ()>, + affine_map<(d0, d1, d2) -> (d0, d1, d2)>]} + ins(%aq, %sa, %za : tensor<2x4x8xi8>, tensor<f32>, tensor<i8>) + outs(%ainit : tensor<2x4x8xf32>) -> tensor<2x4x8xf32> + %binit = tensor.empty() : tensor<2x8x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, + affine_map<(d0, d1, d2) -> ()>, + affine_map<(d0, d1, d2) -> (d0, d1, d2)>]} + ins(%bq, %sb : tensor<2x8x16xi8>, tensor<f32>) + outs(%binit : tensor<2x8x16xf32>) -> tensor<2x8x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %init = tensor.empty() : tensor<2x4x16xf32> + %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<2x4x16xf32>) -> tensor<2x4x16xf32> + %c = linalg.batch_matmul ins(%a, %b : tensor<2x4x8xf32>, tensor<2x8x16xf32>) + outs(%fill : tensor<2x4x16xf32>) -> tensor<2x4x16xf32> + return %c : tensor<2x4x16xf32> +} +// CHECK-LABEL: func.func @batch_matmul( +// CHECK: linalg.generic +// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction"] +// CHECK-SAME: outs(%{{.+}} : tensor<2x4x16xi32>) +// The sum keeps the batch dim and N, reducing only K. +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2, d1)>, +// CHECK-SAME: affine_map<(d0, d1, d2) -> (d0, d1)>] +// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"] +// CHECK-SAME: outs(%{{.+}} : tensor<2x16xi32>) + +// ----- + +// Several M, N and K dims at once. The rewrite never names M, N or K, it only +// partitions reduction dims, so a multi-dim contraction needs no special +// handling: the sum keeps both N dims and reduces both K dims. +func.func @contract_multi_dim(%aq: tensor<2x4x3x8xi8>, %sa: tensor<f32>, %za: tensor<i8>, + %bq: tensor<3x8x5x16xi8>, %sb: tensor<5x16xf32>) -> tensor<2x4x5x16xf32> { + %ainit = tensor.empty() : tensor<2x4x3x8xf32> + %a = 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(%aq, %sa, %za : tensor<2x4x3x8xi8>, tensor<f32>, tensor<i8>) + outs(%ainit : tensor<2x4x3x8xf32>) -> tensor<2x4x3x8xf32> + %binit = tensor.empty() : tensor<3x8x5x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, + affine_map<(d0, d1, d2, d3) -> (d2, d3)>, + affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>]} + ins(%bq, %sb : tensor<3x8x5x16xi8>, tensor<5x16xf32>) + outs(%binit : tensor<3x8x5x16xf32>) -> tensor<3x8x5x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %init = tensor.empty() : tensor<2x4x5x16xf32> + %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<2x4x5x16xf32>) -> tensor<2x4x5x16xf32> + %c = linalg.contract + indexing_maps = [affine_map<(m0, m1, n0, n1, k0, k1) -> (m0, m1, k0, k1)>, + affine_map<(m0, m1, n0, n1, k0, k1) -> (k0, k1, n0, n1)>, + affine_map<(m0, m1, n0, n1, k0, k1) -> (m0, m1, n0, n1)>] + ins(%a, %b : tensor<2x4x3x8xf32>, tensor<3x8x5x16xf32>) + outs(%fill : tensor<2x4x5x16xf32>) -> tensor<2x4x5x16xf32> + return %c : tensor<2x4x5x16xf32> +} +// CHECK-LABEL: func.func @contract_multi_dim( +// CHECK-SAME: %[[AQ:[a-zA-Z0-9_]+]]: tensor<2x4x3x8xi8>, +// CHECK-SAME: %[[BQ:[a-zA-Z0-9_]+]]: tensor<3x8x5x16xi8>, +// CHECK: linalg.generic +// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"] +// CHECK-SAME: ins(%[[AQ]], %[[BQ]] : tensor<2x4x3x8xi8>, tensor<3x8x5x16xi8>) +// CHECK-SAME: outs(%{{.+}} : tensor<2x4x5x16xi32>) +// Both N dims survive into the sum, both K dims are reduced, and both M dims are +// dropped. +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d2, d3, d0, d1)>, +// CHECK-SAME: affine_map<(d0, d1, d2, d3) -> (d0, d1)>] +// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction", "reduction"] +// CHECK-SAME: ins(%[[BQ]] : tensor<3x8x5x16xi8>) +// CHECK-SAME: outs(%{{.+}} : tensor<5x16xi32>) + +//===----------------------------------------------------------------------===// +// Quantization parameter maps in the epilogue +// +// The map that decided legality is reused to index the scale and zero point, +// after being renumbered into the epilogue's iteration space. +//===----------------------------------------------------------------------===// + +// ----- + +// Per-channel on both sides, along different result dims. Each parameter lands +// on its own epilogue dim, so a lhs parameter indexes M and a rhs parameter +// indexes N. Swapping them would still typecheck. +func.func @per_channel_different_dims(%aq: tensor<4x8xi8>, %sa: tensor<4xf32>, %za: tensor<4xi8>, + %bq: tensor<8x16xi8>, %sb: tensor<16xf32>, %zb: tensor<16xi8>) -> tensor<4x16xf32> { + %ainit = tensor.empty() : tensor<4x8xf32> + %a = 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(%aq, %sa, %za : tensor<4x8xi8>, tensor<4xf32>, tensor<4xi8>) + outs(%ainit : tensor<4x8xf32>) -> tensor<4x8xf32> + %binit = tensor.empty() : tensor<8x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> (d1)>, + affine_map<(d0, d1) -> (d1)>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %sb, %zb : tensor<8x16xi8>, tensor<16xf32>, tensor<16xi8>) + outs(%binit : tensor<8x16xf32>) -> tensor<8x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %init = tensor.empty() : tensor<4x16xf32> + %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.matmul ins(%a, %b : tensor<4x8xf32>, tensor<8x16xf32>) + outs(%fill : tensor<4x16xf32>) -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @per_channel_different_dims( +// CHECK-SAME: %[[AQ:[a-zA-Z0-9_]+]]: tensor<4x8xi8>, %[[SA:[a-zA-Z0-9_]+]]: tensor<4xf32>, %[[ZA:[a-zA-Z0-9_]+]]: tensor<4xi8>, +// CHECK-SAME: %[[BQ:[a-zA-Z0-9_]+]]: tensor<8x16xi8>, %[[SB:[a-zA-Z0-9_]+]]: tensor<16xf32>, %[[ZB:[a-zA-Z0-9_]+]]: tensor<16xi8> +// CHECK: %[[D:[a-zA-Z0-9_]+]] = linalg.generic +// CHECK-SAME: outs(%{{.+}} : tensor<4x16xi32>) +// The lhs sum keeps M, the rhs sum keeps N. +// CHECK: %[[RA:[a-zA-Z0-9_]+]] = linalg.generic +// CHECK-SAME: ins(%[[AQ]] : tensor<4x8xi8>) +// CHECK-SAME: outs(%{{.+}} : tensor<4xi32>) +// CHECK: %[[TA:.+]] = linalg.generic +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d0, d1)>] +// CHECK-SAME: ins(%[[ZB]], %[[RA]] : +// CHECK: %[[RB:.+]] = linalg.generic {{.*}} ins(%[[BQ]] : tensor<8x16xi8>) +// CHECK-SAME: outs(%{{.+}} : tensor<16xi32>) +// CHECK: %[[TB:.+]] = linalg.generic +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d0, d1)>] +// CHECK-SAME: ins(%[[ZA]], %[[RB]] : +// CHECK: %[[TZ:.+]] = linalg.generic +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d0, d1)>] +// CHECK-SAME: ins(%[[ZA]], %[[ZB]] : +// CHECK: %[[CORRECTION:.+]] = linalg.generic +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, +// CHECK-SAME: affine_map<(d0, d1) -> (d0, d1)>, +// CHECK-SAME: affine_map<(d0, d1) -> (d0, d1)>, +// CHECK-SAME: affine_map<(d0, d1) -> (d0, d1)>] +// CHECK-SAME: ins(%[[TA]], %[[TB]], %[[TZ]] : +// CHECK-SAME: outs(%{{.+}} : tensor<4x16xi32>) +// The complete correction uses M and N; the scales retain their own axes. +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, +// CHECK-SAME: affine_map<(d0, d1) -> (d0, d1)>, +// CHECK-SAME: affine_map<(d0, d1) -> (d0)>, +// CHECK-SAME: affine_map<(d0, d1) -> (d1)>, +// CHECK-SAME: affine_map<(d0, d1) -> (d0, d1)>] +// CHECK-SAME: iterator_types = ["parallel", "parallel"] +// CHECK-SAME: ins(%[[D]], %[[CORRECTION]], %[[SA]], %[[SB]] : +// CHECK-SAME: outs(%{{.+}} : tensor<4x16xf32>) + +//===----------------------------------------------------------------------===// +// Sum construction +// +// A sum is stated over the dims its operand depends on, compressed to their own +// numbering. Which dims survive, and in what order, is what the epilogue then +// has to match. +//===----------------------------------------------------------------------===// + +// ----- + +// Multiply a 1x128 vector by ten weight rows. The lhs zero point requires +// RB[n] = sum_k Bq[n, k], producing ten sums for the correction zA*RB[n]. +// Computing RB visits all 10x128 weights, just as the contraction does. The +// rewrite still supports dynamic weights; constant weights allow RB to be +// precomputed. This test checks the dynamic case. +func.func @gemv(%aq: tensor<1x128xi8>, %sa: f32, %za: i8, + %bq: tensor<10x128xi8>, %sb: tensor<10xf32>) -> tensor<1x10xf32> { + %ainit = tensor.empty() : tensor<1x128xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %sa, %za : tensor<1x128xi8>, f32, i8) + outs(%ainit : tensor<1x128xf32>) -> tensor<1x128xf32> + %binit = tensor.empty() : tensor<10x128xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> (d0)>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %sb : tensor<10x128xi8>, tensor<10xf32>) + outs(%binit : tensor<10x128xf32>) -> tensor<10x128xf32> + %cst = arith.constant 0.000000e+00 : f32 + %init = tensor.empty() : tensor<1x10xf32> + %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x10xf32>) -> tensor<1x10xf32> + %c = linalg.matmul + indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, + affine_map<(d0, d1, d2) -> (d1, d2)>, + affine_map<(d0, d1, d2) -> (d0, d1)>] + ins(%a, %b : tensor<1x128xf32>, tensor<10x128xf32>) + outs(%fill : tensor<1x10xf32>) -> tensor<1x10xf32> + return %c : tensor<1x10xf32> +} +// CHECK-LABEL: func.func @gemv( +// CHECK-SAME: %[[AQ:[a-zA-Z0-9_]+]]: tensor<1x128xi8>, +// CHECK-SAME: %[[BQ:[a-zA-Z0-9_]+]]: tensor<10x128xi8>, +// CHECK: %[[D:.+]] = linalg.generic +// CHECK-SAME: ins(%[[AQ]], %[[BQ]] : tensor<1x128xi8>, tensor<10x128xi8>) +// CHECK-SAME: outs(%{{.+}} : tensor<1x10xi32>) +// RB keeps the weight row N and reduces its K dimension. +// CHECK: %[[RB:.+]] = linalg.generic +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>] +// CHECK-SAME: iterator_types = ["parallel", "reduction"] +// CHECK-SAME: ins(%[[BQ]] : tensor<10x128xi8>) +// CHECK-SAME: outs(%{{.+}} : tensor<10xi32>) + +//===----------------------------------------------------------------------===// +// Shape-only witness operand +// +// linalg recovers loop extents by inverting the concatenated operand maps, which +// only recovers a dim appearing as a bare result. A convolution's window dims +// appear only inside `oh * stride + kh * dilation`, so a shape-only operand +// naming them is appended and ignored by the body. +//===----------------------------------------------------------------------===// + +// ----- + +// Convolution with asymmetric weights, which is the case needing the image sum. +// The image does not depend on F, so the sum drops it entirely and reduces C as +// well, leaving (N, OH, OW). That is what makes it 1/F the size of the +// convolution. +func.func @conv_2d_nhwc_hwcf_image_sum(%aq: tensor<1x8x8x4xi8>, %sa: tensor<f32>, + %bq: tensor<3x3x4x16xi8>, %sb: tensor<16xf32>, %zb: tensor<16xi8>) -> tensor<1x6x6x16xf32> { + %ai = tensor.empty() : tensor<1x8x8x4xf32> + %a = 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) -> (d0, d1, d2, d3)>]} + ins(%aq, %sa : tensor<1x8x8x4xi8>, tensor<f32>) + outs(%ai : tensor<1x8x8x4xf32>) -> tensor<1x8x8x4xf32> + %bi = tensor.empty() : tensor<3x3x4x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, + affine_map<(d0, d1, d2, d3) -> (d3)>, + affine_map<(d0, d1, d2, d3) -> (d3)>, + affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>]} + ins(%bq, %sb, %zb : tensor<3x3x4x16xi8>, tensor<16xf32>, tensor<16xi8>) + outs(%bi : tensor<3x3x4x16xf32>) -> tensor<3x3x4x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %init = tensor.empty() : tensor<1x6x6x16xf32> + %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x6x6x16xf32>) -> tensor<1x6x6x16xf32> + %c = linalg.conv_2d_nhwc_hwcf {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} + ins(%a, %b : tensor<1x8x8x4xf32>, tensor<3x3x4x16xf32>) + outs(%fill : tensor<1x6x6x16xf32>) -> tensor<1x6x6x16xf32> + return %c : tensor<1x6x6x16xf32> +} +// CHECK-LABEL: func.func @conv_2d_nhwc_hwcf_image_sum( +// CHECK: linalg.generic +// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction"] +// CHECK-SAME: outs(%{{.+}} : tensor<1x6x6x16xi32>) +// The window dims are named by a shape-only operand, sized from the filter. +// CHECK: %[[W:.+]] = tensor.empty() : tensor<3x3xi8> +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 + d3, d2 + d4, d5)>, +// CHECK-SAME: affine_map<(d0, d1, d2, d3, d4, d5) -> (d3, d4)>, +// CHECK-SAME: affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2)>] +// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction", "reduction", "reduction"] +// CHECK-SAME: ins(%{{.+}}, %[[W]] : tensor<1x8x8x4xi8>, tensor<3x3xi8>) +// CHECK-SAME: outs(%{{.+}} : tensor<1x6x6xi32>) + +// ----- + +// Depthwise: there is no F to drop, so the image sum keeps C and comes out the +// full output shape. Every dim of the convolution survives into the image sum, +// so this case exercises the correction without dropping a parallel dimension. +func.func @depthwise_conv_image_sum(%aq: tensor<1x8x8x4xi8>, %sa: tensor<f32>, + %bq: tensor<3x3x4xi8>, %sb: tensor<4xf32>, %zb: tensor<4xi8>) -> tensor<1x6x6x4xf32> { + %ainit = tensor.empty() : tensor<1x8x8x4xf32> + %a = 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) -> (d0, d1, d2, d3)>]} + ins(%aq, %sa : tensor<1x8x8x4xi8>, tensor<f32>) + outs(%ainit : tensor<1x8x8x4xf32>) -> tensor<1x8x8x4xf32> + %binit = tensor.empty() : tensor<3x3x4xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, + affine_map<(d0, d1, d2) -> (d2)>, + affine_map<(d0, d1, d2) -> (d2)>, + affine_map<(d0, d1, d2) -> (d0, d1, d2)>]} + ins(%bq, %sb, %zb : tensor<3x3x4xi8>, tensor<4xf32>, tensor<4xi8>) + outs(%binit : tensor<3x3x4xf32>) -> tensor<3x3x4xf32> + %cst = arith.constant 0.000000e+00 : f32 + %init = tensor.empty() : tensor<1x6x6x4xf32> + %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x6x6x4xf32>) -> tensor<1x6x6x4xf32> + %c = linalg.depthwise_conv_2d_nhwc_hwc {dilations = dense<1> : tensor<2xi64>, + strides = dense<1> : tensor<2xi64>} + ins(%a, %b : tensor<1x8x8x4xf32>, tensor<3x3x4xf32>) + outs(%fill : tensor<1x6x6x4xf32>) -> tensor<1x6x6x4xf32> + return %c : tensor<1x6x6x4xf32> +} +// CHECK-LABEL: func.func @depthwise_conv_image_sum( +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 + d4, d2 + d5, d3)>, +// CHECK-SAME: affine_map<(d0, d1, d2, d3, d4, d5) -> (d4, d5, d3)>, +// CHECK-SAME: affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>] +// CHECK-SAME: outs(%{{.+}} : tensor<1x6x6x4xi32>) +// The image sum has the same iteration space and the same result shape as the +// contraction above it, because C is kept rather than reduced: the image map and +// the result map below are the same two maps again, with only the filter operand +// replaced by the window witness. +// CHECK: %[[W:.+]] = tensor.empty() : tensor<3x3xi8> +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 + d4, d2 + d5, d3)>, +// CHECK-SAME: affine_map<(d0, d1, d2, d3, d4, d5) -> (d4, d5)>, +// CHECK-SAME: affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>] +// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"] +// CHECK-SAME: ins(%{{.+}}, %[[W]] : tensor<1x8x8x4xi8>, tensor<3x3xi8>) +// CHECK-SAME: outs(%{{.+}} : tensor<1x6x6x4xi32>) + +// ----- + +// Strides and dilations survive into the image sum's map, so it reads the same +// window the convolution does. Getting either factor wrong here would sum the +// wrong elements. +func.func @conv_strided_dilated_image_sum(%aq: tensor<1x7x7x2xi8>, %sa: tensor<f32>, + %bq: tensor<2x2x2x2xi8>, %sb: tensor<2xf32>, %zb: tensor<2xi8>) -> tensor<1x3x3x2xf32> { + %ai = tensor.empty() : tensor<1x7x7x2xf32> + %a = 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) -> (d0, d1, d2, d3)>]} + ins(%aq, %sa : tensor<1x7x7x2xi8>, tensor<f32>) + outs(%ai : tensor<1x7x7x2xf32>) -> tensor<1x7x7x2xf32> + %bi = tensor.empty() : tensor<2x2x2x2xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, + affine_map<(d0, d1, d2, d3) -> (d3)>, + affine_map<(d0, d1, d2, d3) -> (d3)>, + affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>]} + ins(%bq, %sb, %zb : tensor<2x2x2x2xi8>, tensor<2xf32>, tensor<2xi8>) + outs(%bi : tensor<2x2x2x2xf32>) -> tensor<2x2x2x2xf32> + %cst = arith.constant 0.000000e+00 : f32 + %init = tensor.empty() : tensor<1x3x3x2xf32> + %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x3x3x2xf32>) -> tensor<1x3x3x2xf32> + %c = linalg.conv_2d_nhwc_hwcf {dilations = dense<2> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} + ins(%a, %b : tensor<1x7x7x2xf32>, tensor<2x2x2x2xf32>) + outs(%fill : tensor<1x3x3x2xf32>) -> tensor<1x3x3x2xf32> + return %c : tensor<1x3x3x2xf32> +} +// CHECK-LABEL: func.func @conv_strided_dilated_image_sum( +// CHECK: linalg.generic +// CHECK-SAME: outs(%{{.+}} : tensor<1x3x3x2xi32>) +// CHECK: %[[W:.+]] = tensor.empty() : tensor<2x2xi8> +// CHECK: linalg.generic +// CHECK-SAME: affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 * 2 + d3 * 2, d2 * 2 + d4 * 2, d5)> +// CHECK-SAME: ins(%{{.+}}, %[[W]] : tensor<1x7x7x2xi8>, tensor<2x2xi8>) +// CHECK-SAME: outs(%{{.+}} : tensor<1x3x3xi32>) + +// ----- + +// A convolution whose weights are symmetric needs only the filter sum, and the +// filter map is a projected permutation, so every dim is already named as a bare +// coordinate. No witness operand is emitted. +func.func @no_witness_when_map_is_projected_permutation(%aq: tensor<1x8x8x4xi8>, %sa: tensor<f32>, %za: tensor<i8>, + %bq: tensor<3x3x4x16xi8>, %sb: tensor<16xf32>) -> tensor<1x6x6x16xf32> { + %ai = tensor.empty() : tensor<1x8x8x4xf32> + %a = 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(%aq, %sa, %za : tensor<1x8x8x4xi8>, tensor<f32>, tensor<i8>) + outs(%ai : tensor<1x8x8x4xf32>) -> tensor<1x8x8x4xf32> + %bi = tensor.empty() : tensor<3x3x4x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, + affine_map<(d0, d1, d2, d3) -> (d3)>, + affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>]} + ins(%bq, %sb : tensor<3x3x4x16xi8>, tensor<16xf32>) + outs(%bi : tensor<3x3x4x16xf32>) -> tensor<3x3x4x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %init = tensor.empty() : tensor<1x6x6x16xf32> + %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x6x6x16xf32>) -> tensor<1x6x6x16xf32> + %c = linalg.conv_2d_nhwc_hwcf {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} + ins(%a, %b : tensor<1x8x8x4xf32>, tensor<3x3x4x16xf32>) + outs(%fill : tensor<1x6x6x16xf32>) -> tensor<1x6x6x16xf32> + return %c : tensor<1x6x6x16xf32> +} +// CHECK-LABEL: func.func @no_witness_when_map_is_projected_permutation( +// CHECK-NOT: tensor.empty() : tensor<3x3xi8> +// CHECK: linalg.generic +// CHECK-SAME: outs(%{{.+}} : tensor<1x6x6x16xi32>) +// The filter sum takes one input and no witness, keeping only F. +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d1, d2, d3, d0)>, +// CHECK-SAME: affine_map<(d0, d1, d2, d3) -> (d0)>] +// CHECK-SAME: iterator_types = ["parallel", "reduction", "reduction", "reduction"] +// CHECK-SAME: ins(%{{.+}} : tensor<3x3x4x16xi8>) +// CHECK-SAME: outs(%{{.+}} : tensor<16xi32>) + +// ----- + +// The channel-first layout, to show the construction does not depend on where +// the spatial dims sit. The image map carries the window in its last two +// results rather than its middle two, and the filter sum keeps F. +func.func @conv_2d_nchw_fchw(%aq: tensor<1x4x8x8xi8>, %sa: f32, %za: i8, + %bq: tensor<16x4x3x3xi8>, %sb: tensor<16xf32>) -> tensor<1x16x6x6xf32> { + %ai = tensor.empty() : tensor<1x4x8x8xf32> + %a = iree_linalg_ext.dequantize_affine + {input_unsigned, + 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(%aq, %sa, %za : tensor<1x4x8x8xi8>, f32, i8) + outs(%ai : tensor<1x4x8x8xf32>) -> tensor<1x4x8x8xf32> + %bi = tensor.empty() : tensor<16x4x3x3xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, + affine_map<(d0, d1, d2, d3) -> (d0)>, + affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>]} + ins(%bq, %sb : tensor<16x4x3x3xi8>, tensor<16xf32>) + outs(%bi : tensor<16x4x3x3xf32>) -> tensor<16x4x3x3xf32> + %cst = arith.constant 0.000000e+00 : f32 + %init = tensor.empty() : tensor<1x16x6x6xf32> + %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x16x6x6xf32>) -> tensor<1x16x6x6xf32> + %c = linalg.conv_2d_nchw_fchw {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} + ins(%a, %b : tensor<1x4x8x8xf32>, tensor<16x4x3x3xf32>) + outs(%fill : tensor<1x16x6x6xf32>) -> tensor<1x16x6x6xf32> + return %c : tensor<1x16x6x6xf32> +} +// CHECK-LABEL: func.func @conv_2d_nchw_fchw( +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d4, d2 + d5, d3 + d6)>, +// CHECK-SAME: affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d1, d4, d5, d6)>, +// CHECK-SAME: affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)>] +// CHECK-SAME: outs(%{{.+}} : tensor<1x16x6x6xi32>) +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, +// CHECK-SAME: affine_map<(d0, d1, d2, d3) -> (d0)>] +// CHECK-SAME: outs(%{{.+}} : tensor<16xi32>) + +//===----------------------------------------------------------------------===// +// Promotion of the dims the parameters vary along +// +// In the partial form those reduction dims become parallel in the integer +// contraction, are appended to its result, and reappear as the epilogue's +// trailing reduction. +//===----------------------------------------------------------------------===// + +// ----- + +// Blockwise with an asymmetric lhs. K is expanded to (G, L), the lhs parameters +// vary along G, so G stays parallel in the integer contraction and the epilogue +// reduces across it. +func.func @blockwise_promotes_block_dim(%aq: tensor<4x2x8xi8>, %sa: tensor<4x2xf32>, %za: tensor<4x2xi8>, + %bq: tensor<2x8x16xi8>, %sb: f32) -> tensor<4x16xf32> { + %ainit = tensor.empty() : tensor<4x2x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, + affine_map<(d0, d1, d2) -> (d0, d1)>, + affine_map<(d0, d1, d2) -> (d0, d1)>, + affine_map<(d0, d1, d2) -> (d0, d1, d2)>]} + ins(%aq, %sa, %za : tensor<4x2x8xi8>, tensor<4x2xf32>, tensor<4x2xi8>) + outs(%ainit : tensor<4x2x8xf32>) -> tensor<4x2x8xf32> + %binit = tensor.empty() : tensor<2x8x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, + affine_map<(d0, d1, d2) -> ()>, + affine_map<(d0, d1, d2) -> (d0, d1, d2)>]} + ins(%bq, %sb : tensor<2x8x16xi8>, f32) + outs(%binit : tensor<2x8x16xf32>) -> tensor<2x8x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %init = tensor.empty() : tensor<4x16xf32> + %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.contract + indexing_maps = [affine_map<(m, n, g, l) -> (m, g, l)>, + affine_map<(m, n, g, l) -> (g, l, n)>, + affine_map<(m, n, g, l) -> (m, n)>] + ins(%a, %b : tensor<4x2x8xf32>, tensor<2x8x16xf32>) + outs(%fill : tensor<4x16xf32>) -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @blockwise_promotes_block_dim( +// CHECK-SAME: %[[AQ:[a-zA-Z0-9_]+]]: tensor<4x2x8xi8>, %[[SA:[a-zA-Z0-9_]+]]: tensor<4x2xf32>, %[[ZA:[a-zA-Z0-9_]+]]: tensor<4x2xi8>, +// CHECK-SAME: %[[BQ:[a-zA-Z0-9_]+]]: tensor<2x8x16xi8>, %[[SB:[a-zA-Z0-9_]+]]: f32 +// G is parallel here, and appended to the result, so the contraction yields one +// partial per block. +// CHECK: %[[D:[a-zA-Z0-9_]+]] = linalg.generic +// CHECK-SAME: affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>] +// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction"] +// CHECK-SAME: ins(%[[AQ]], %[[BQ]] : tensor<4x2x8xi8>, tensor<2x8x16xi8>) +// CHECK-SAME: outs(%{{.+}} : tensor<4x16x2xi32>) +// The rhs sum keeps N and G, reducing only L. The kept dims come out in +// iteration order, so it is (N, G) rather than (G, N), which is the trailing +// pair of the contraction's (M, N, G) and needs no transpose below. +// CHECK: %[[RB:[a-zA-Z0-9_]+]] = linalg.generic +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2, d0)>, +// CHECK-SAME: affine_map<(d0, d1, d2) -> (d0, d1)>] +// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"] +// CHECK-SAME: ins(%[[BQ]] : tensor<2x8x16xi8>) +// CHECK-SAME: outs(%{{.+}} : tensor<16x2xi32>) +// The product broadcasts the zero point over N and the sum over M, retaining G. +// CHECK: %[[CORRECTION:.+]] = linalg.generic +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, +// CHECK-SAME: affine_map<(d0, d1, d2) -> (d1, d2)>, +// CHECK-SAME: affine_map<(d0, d1, d2) -> (d0, d1, d2)>] +// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"] +// CHECK-SAME: ins(%[[ZA]], %[[RB]] : +// CHECK-SAME: outs(%{{.+}} : tensor<4x16x2xi32>) +// Only the epilogue reduces G, after applying each block's correction and scale. +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, +// CHECK-SAME: affine_map<(d0, d1, d2) -> (d0, d1, d2)>, +// CHECK-SAME: affine_map<(d0, d1, d2) -> (d0, d2)>, +// CHECK-SAME: affine_map<(d0, d1, d2) -> ()>, +// CHECK-SAME: affine_map<(d0, d1, d2) -> (d0, d1)>] +// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"] +// CHECK-SAME: ins(%[[D]], %[[CORRECTION]], %[[SA]], %[[SB]] : +// CHECK-SAME: outs(%{{.+}} : tensor<4x16xf32>) + +// ----- + +// Both operands have per-block scales and zero points. For each (M, N, G), +// correct D with zB*RA + zA*RB - L*zA*zB before scaling and summing across G. +// There are two blocks of eight elements, so the cross term uses 8, not 16. +func.func @blockwise_both_asymmetric(%aq: tensor<4x2x8xi8>, %sa: tensor<4x2xf32>, %za: tensor<4x2xi8>, + %bq: tensor<2x8x16xi8>, %sb: tensor<2x16xf32>, %zb: tensor<2x16xi8>) -> tensor<4x16xf32> { + %ainit = tensor.empty() : tensor<4x2x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(m, g, l) -> (m, g, l)>, + affine_map<(m, g, l) -> (m, g)>, + affine_map<(m, g, l) -> (m, g)>, + affine_map<(m, g, l) -> (m, g, l)>]} + ins(%aq, %sa, %za : tensor<4x2x8xi8>, tensor<4x2xf32>, tensor<4x2xi8>) + outs(%ainit : tensor<4x2x8xf32>) -> tensor<4x2x8xf32> + %binit = tensor.empty() : tensor<2x8x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(g, l, n) -> (g, l, n)>, + affine_map<(g, l, n) -> (g, n)>, + affine_map<(g, l, n) -> (g, n)>, + affine_map<(g, l, n) -> (g, l, n)>]} + ins(%bq, %sb, %zb : tensor<2x8x16xi8>, tensor<2x16xf32>, tensor<2x16xi8>) + outs(%binit : tensor<2x8x16xf32>) -> tensor<2x8x16xf32> + %zero = arith.constant 0.0 : f32 + %init = tensor.empty() : tensor<4x16xf32> + %fill = linalg.fill ins(%zero : f32) outs(%init : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.contract + indexing_maps = [affine_map<(m, n, g, l) -> (m, g, l)>, + affine_map<(m, n, g, l) -> (g, l, n)>, + affine_map<(m, n, g, l) -> (m, n)>] + ins(%a, %b : tensor<4x2x8xf32>, tensor<2x8x16xf32>) + outs(%fill : tensor<4x16xf32>) -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @blockwise_both_asymmetric( +// CHECK-SAME: %[[AQ:[a-zA-Z0-9_]+]]: tensor<4x2x8xi8>, %[[SA:[a-zA-Z0-9_]+]]: tensor<4x2xf32>, %[[ZA:[a-zA-Z0-9_]+]]: tensor<4x2xi8>, +// CHECK-SAME: %[[BQ:[a-zA-Z0-9_]+]]: tensor<2x8x16xi8>, %[[SB:[a-zA-Z0-9_]+]]: tensor<2x16xf32>, %[[ZB:[a-zA-Z0-9_]+]]: tensor<2x16xi8> +// CHECK: %[[L:.+]] = arith.constant 8 : i32 +// D retains G and reduces only L. +// CHECK: %[[D:.+]] = linalg.generic +// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction"] +// CHECK-SAME: ins(%[[AQ]], %[[BQ]] : +// CHECK-SAME: outs(%{{.+}} : tensor<4x16x2xi32>) +// RA[M, G] and RB[N, G] each reduce L; neither sums across blocks. +// CHECK: %[[RA:.+]] = linalg.generic +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>] +// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"] +// CHECK-SAME: ins(%[[AQ]] : tensor<4x2x8xi8>) +// CHECK-SAME: outs(%{{.+}} : tensor<4x2xi32>) +// CHECK: %[[TA:.+]] = linalg.generic {{.*}} ins(%[[ZB]], %[[RA]] : +// CHECK-SAME: outs(%{{.+}} : tensor<4x16x2xi32>) +// CHECK: %[[RB:.+]] = linalg.generic +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2, d0)>, affine_map<(d0, d1, d2) -> (d0, d1)>] +// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"] +// CHECK-SAME: ins(%[[BQ]] : tensor<2x8x16xi8>) +// CHECK-SAME: outs(%{{.+}} : tensor<16x2xi32>) +// CHECK: %[[TB:.+]] = linalg.generic {{.*}} ins(%[[ZA]], %[[RB]] : +// CHECK-SAME: outs(%{{.+}} : tensor<4x16x2xi32>) +// The zero points share G but index different output dimensions. +// CHECK: %[[TZ:.+]] = linalg.generic +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, +// CHECK-SAME: affine_map<(d0, d1, d2) -> (d2, d1)>, +// CHECK-SAME: affine_map<(d0, d1, d2) -> (d0, d1, d2)>] +// CHECK-SAME: ins(%[[ZA]], %[[ZB]] : +// CHECK-SAME: outs(%{{.+}} : tensor<4x16x2xi32>) +// CHECK: ^bb0(%[[AZ:[a-zA-Z0-9_]+]]: i8, %[[BZ:[a-zA-Z0-9_]+]]: i8, %{{.+}}: i32): +// CHECK: %[[AZ32:.+]] = arith.extsi %[[AZ]] : i8 to i32 +// CHECK: %[[BZ32:.+]] = arith.extsi %[[BZ]] : i8 to i32 +// CHECK: %[[ZZ:.+]] = arith.muli %[[AZ32]], %[[BZ32]] : i32 +// CHECK: %[[CROSS:.+]] = arith.muli %[[ZZ]], %[[L]] : i32 +// CHECK: linalg.yield %[[CROSS]] : i32 +// CHECK: %[[CORRECTION:.+]] = linalg.generic {{.*}} ins(%[[TA]], %[[TB]], %[[TZ]] : +// CHECK-SAME: outs(%{{.+}} : tensor<4x16x2xi32>) +// CHECK: ^bb0(%[[A:[a-zA-Z0-9_]+]]: i32, %[[B:[a-zA-Z0-9_]+]]: i32, %[[Z:[a-zA-Z0-9_]+]]: i32, %{{.+}}: i32): +// CHECK: %[[SUM:.+]] = arith.addi %[[A]], %[[B]] : i32 +// CHECK: %[[OFFSET:.+]] = arith.subi %[[SUM]], %[[Z]] : i32 +// CHECK: linalg.yield %[[OFFSET]] : i32 +// Scale each corrected block, then reduce G in floating point. +// CHECK: %[[RESULT:.+]] = linalg.generic +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, +// CHECK-SAME: affine_map<(d0, d1, d2) -> (d0, d1, d2)>, +// CHECK-SAME: affine_map<(d0, d1, d2) -> (d0, d2)>, +// CHECK-SAME: affine_map<(d0, d1, d2) -> (d2, d1)>, +// CHECK-SAME: affine_map<(d0, d1, d2) -> (d0, d1)>] +// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"] +// CHECK-SAME: ins(%[[D]], %[[CORRECTION]], %[[SA]], %[[SB]] : +// CHECK-SAME: outs(%{{.+}} : tensor<4x16xf32>) +// CHECK: ^bb0(%[[ED:[a-zA-Z0-9_]+]]: i32, %[[EC:[a-zA-Z0-9_]+]]: i32, %[[AS:[a-zA-Z0-9_]+]]: f32, %[[BS:[a-zA-Z0-9_]+]]: f32, %[[ACC:[a-zA-Z0-9_]+]]: f32): +// CHECK: %[[CORRECTED:.+]] = arith.subi %[[ED]], %[[EC]] : i32 +// CHECK: %[[REAL:.+]] = arith.sitofp %[[CORRECTED]] : i32 to f32 +// CHECK: %[[PARTIAL:.+]] = arith.mulf %[[REAL]], %[[AS]] : f32 +// CHECK: %[[SCALED:.+]] = arith.mulf %[[PARTIAL]], %[[BS]] : f32 +// CHECK: %[[TOTAL:.+]] = arith.addf %[[ACC]], %[[SCALED]] : f32 +// CHECK: linalg.yield %[[TOTAL]] : f32 +// CHECK: return %[[RESULT]] : tensor<4x16xf32> + +// ----- + +// An asymmetric lhs with the block dims stored the other way round, folded +// into the dequantize. Inverting its output map carries the contraction's +// (M, G, L) lhs map onto the (G, M, L) storage, so the integer contraction +// reads the stored order directly and the promoted group extent survives the +// permutation into every partial result. +func.func @blockwise_transposed_storage(%aq: tensor<2x4x8xi8>, %sa: tensor<4x2xf32>, %za: tensor<4x2xi8>, + %bq: tensor<2x8x16xi8>, %sb: f32) -> tensor<4x16xf32> { + %ainit = tensor.empty() : tensor<4x2x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d0, d2)>, + affine_map<(d0, d1, d2) -> (d0, d1)>, + affine_map<(d0, d1, d2) -> (d0, d1)>, + affine_map<(d0, d1, d2) -> (d0, d1, d2)>]} + ins(%aq, %sa, %za : tensor<2x4x8xi8>, tensor<4x2xf32>, tensor<4x2xi8>) + outs(%ainit : tensor<4x2x8xf32>) -> tensor<4x2x8xf32> + %binit = tensor.empty() : tensor<2x8x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, + affine_map<(d0, d1, d2) -> ()>, + affine_map<(d0, d1, d2) -> (d0, d1, d2)>]} + ins(%bq, %sb : tensor<2x8x16xi8>, f32) + outs(%binit : tensor<2x8x16xf32>) -> tensor<2x8x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %init = tensor.empty() : tensor<4x16xf32> + %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.contract + indexing_maps = [affine_map<(m, n, g, l) -> (m, g, l)>, + affine_map<(m, n, g, l) -> (g, l, n)>, + affine_map<(m, n, g, l) -> (m, n)>] + ins(%a, %b : tensor<4x2x8xf32>, tensor<2x8x16xf32>) + outs(%fill : tensor<4x16xf32>) -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @blockwise_transposed_storage( +// CHECK-SAME: %[[AQ:[a-zA-Z0-9_]+]]: tensor<2x4x8xi8> +// @blockwise_promotes_block_dim reads (d0, d2, d3) here for the same iteration +// space. +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d2, d0, d3)>, +// CHECK-SAME: ins(%[[AQ]], %{{.+}} : tensor<2x4x8xi8>, tensor<2x8x16xi8>) +// CHECK-SAME: outs(%{{.+}} : tensor<4x16x2xi32>) +// CHECK: linalg.generic {{.*}} outs(%{{.+}} : tensor<16x2xi32>) +// CHECK: linalg.generic {{.*}} outs(%{{.+}} : tensor<4x16x2xi32>) +// CHECK: linalg.generic {{.*}} outs(%{{.+}} : tensor<4x16xf32>) + +//===----------------------------------------------------------------------===// +// Dynamic shapes +// +// Extents are read from the quantized inputs rather than from the dequantized +// values, and result dims are taken from the original init so the replacement +// carries the original result type. +//===----------------------------------------------------------------------===// + +// ----- + +// A dynamic parallel dim propagates into every intermediate that carries it. +// The extents come from the quantized input, so nothing keeps the dequantize +// alive. +func.func @dynamic_parallel_dim(%aq: tensor<?x8xi8>, %sa: tensor<f32>, %za: tensor<i8>, + %bq: tensor<8x16xi8>, %sb: tensor<f32>, %zb: tensor<i8>) -> tensor<?x16xf32> { + %c0 = arith.constant 0 : index + %m = tensor.dim %aq, %c0 : tensor<?x8xi8> + %ainit = tensor.empty(%m) : tensor<?x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %sa, %za : tensor<?x8xi8>, tensor<f32>, tensor<i8>) + outs(%ainit : tensor<?x8xf32>) -> tensor<?x8xf32> + %binit = tensor.empty() : tensor<8x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %sb, %zb : tensor<8x16xi8>, tensor<f32>, tensor<i8>) + outs(%binit : tensor<8x16xf32>) -> tensor<8x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %init = tensor.empty(%m) : tensor<?x16xf32> + %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<?x16xf32>) -> tensor<?x16xf32> + %c = linalg.matmul ins(%a, %b : tensor<?x8xf32>, tensor<8x16xf32>) + outs(%fill : tensor<?x16xf32>) -> tensor<?x16xf32> + return %c : tensor<?x16xf32> +} +// CHECK-LABEL: func.func @dynamic_parallel_dim( +// CHECK-NOT: iree_linalg_ext.dequantize_affine +// M is dynamic while N and K are static, so the contraction, the lhs sum and +// its correction term, the combined correction and the epilogue are dynamic, +// and the rhs sum and its term are not. +// CHECK: linalg.generic {{.*}} outs(%{{.+}} : tensor<?x16xi32>) +// CHECK: linalg.generic {{.*}} outs(%{{.+}} : tensor<?xi32>) +// CHECK: linalg.generic {{.*}} outs(%{{.+}} : tensor<?xi32>) +// CHECK: linalg.generic {{.*}} outs(%{{.+}} : tensor<16xi32>) +// CHECK: linalg.generic {{.*}} outs(%{{.+}} : tensor<16xi32>) +// CHECK: linalg.generic {{.*}} outs(%{{.+}} : tensor<?x16xi32>) +// CHECK: linalg.generic {{.*}} outs(%{{.+}} : tensor<?x16xf32>) + +// ----- + +// The inputs know M statically but the result type does not. Sizes for result +// dims are taken from the init rather than from whichever operand mentions them +// first, so the replacement keeps the original dynamic result type. +func.func @static_input_dim_dynamic_result(%aq: tensor<4x8xi8>, %sa: tensor<f32>, %za: tensor<i8>, + %bq: tensor<8x16xi8>, %sb: tensor<f32>, %init: tensor<?x16xf32>) -> tensor<?x16xf32> { + %ainit = tensor.empty() : tensor<4x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %sa, %za : tensor<4x8xi8>, tensor<f32>, tensor<i8>) + outs(%ainit : tensor<4x8xf32>) -> tensor<4x8xf32> + %binit = tensor.empty() : tensor<8x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %sb : tensor<8x16xi8>, tensor<f32>) + outs(%binit : tensor<8x16xf32>) -> tensor<8x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<?x16xf32>) -> tensor<?x16xf32> + %c = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, + affine_map<(d0, d1, d2) -> (d2, d1)>, + affine_map<(d0, d1, d2) -> (d0, d1)>], + iterator_types = ["parallel", "parallel", "reduction"]} + ins(%a, %b : tensor<4x8xf32>, tensor<8x16xf32>) + outs(%fill : tensor<?x16xf32>) { + ^bb0(%in: f32, %in_0: f32, %out: f32): + %m = arith.mulf %in, %in_0 : f32 + %s = arith.addf %out, %m : f32 + linalg.yield %s : f32 + } -> tensor<?x16xf32> + return %c : tensor<?x16xf32> +} +// CHECK-LABEL: func.func @static_input_dim_dynamic_result( +// Both the integer contraction and the epilogue keep the dynamic result type, +// even though the lhs knows M is 4. +// CHECK: linalg.generic +// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<4x8xi8>, tensor<8x16xi8>) +// CHECK-SAME: outs(%{{.+}} : tensor<?x16xi32>) +// CHECK: linalg.generic {{.*}} outs(%{{.+}} : tensor<?x16xf32>)
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/convert_qdq_to_integer_math_negative.mlir b/compiler/src/iree/compiler/GlobalOptimization/test/convert_qdq_to_integer_math_negative.mlir new file mode 100644 index 0000000..44b9425 --- /dev/null +++ b/compiler/src/iree/compiler/GlobalOptimization/test/convert_qdq_to_integer_math_negative.mlir
@@ -0,0 +1,843 @@ +// RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(func.func(iree-global-opt-convert-qdq-to-integer-math))" %s | FileCheck %s + +// Contractions the rewrite declines. Each keeps its floating point form, which +// stays correct. The witness is the surviving dequantize together with the +// original contraction, or with the body operation that blocked the match. + +//===----------------------------------------------------------------------===// +// Structural preconditions +// +// The rewrite needs two dequantized operands feeding a zero-initialised +// contraction, and it has to be able to read the contraction out of the body. +//===----------------------------------------------------------------------===// + +// Only one operand is quantized, so there is no integer contraction to form: +// the other operand is real valued. +func.func @single_quantized_operand(%aq: tensor<4x8xi8>, %a_s: f32, %b: tensor<8x16xf32>) -> tensor<4x16xf32> { + %a_i = tensor.empty() : tensor<4x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s : tensor<4x8xi8>, f32) + outs(%a_i : tensor<4x8xf32>) -> tensor<4x8xf32> + %cst = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<4x16xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.matmul ins(%a, %b : tensor<4x8xf32>, tensor<8x16xf32>) + outs(%f : tensor<4x16xf32>) -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @single_quantized_operand( +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK: linalg.matmul + +// ----- + +// The integer contraction starts from an integer zero, so the original +// accumulator has to be a zero fill for the epilogue not to have to carry it. +// A bias initialised accumulator is not handled. +func.func @nonzero_init(%aq: tensor<4x8xi8>, %a_s: f32, %bq: tensor<8x16xi8>, %b_s: f32, %bias: tensor<4x16xf32>) -> tensor<4x16xf32> { + %a_i = tensor.empty() : tensor<4x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s : tensor<4x8xi8>, f32) + outs(%a_i : tensor<4x8xf32>) -> tensor<4x8xf32> + %b_i = tensor.empty() : tensor<8x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s : tensor<8x16xi8>, f32) + outs(%b_i : tensor<8x16xf32>) -> tensor<8x16xf32> + %c = linalg.matmul ins(%a, %b : tensor<4x8xf32>, tensor<8x16xf32>) + outs(%bias : tensor<4x16xf32>) -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @nonzero_init( +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK: linalg.matmul + +// ----- + +// The rewrite requires dequantize to be the contraction's immediate producer; +// it does not look through the tensor.pad below. The global optimization +// pipeline can expose direct producers by folding transposes into contractions, +// but moving QDQ through reshapes or padding requires separate propagation. +func.func @op_between_dequantize_and_contraction(%aq: tensor<4x6xi8>, %a_s: f32, + %bq: tensor<8x16xi8>, %b_s: f32) -> tensor<4x16xf32> { + %a_i = tensor.empty() : tensor<4x6xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s : tensor<4x6xi8>, f32) + outs(%a_i : tensor<4x6xf32>) -> tensor<4x6xf32> + %b_i = tensor.empty() : tensor<8x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s : tensor<8x16xi8>, f32) + outs(%b_i : tensor<8x16xf32>) -> tensor<8x16xf32> + %pad_cst = arith.constant 0.000000e+00 : f32 + %a_pad = tensor.pad %a low[0, 0] high[0, 2] { + ^bb0(%i: index, %j: index): + tensor.yield %pad_cst : f32 + } : tensor<4x6xf32> to tensor<4x8xf32> + %cst = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<4x16xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.matmul ins(%a_pad, %b : tensor<4x8xf32>, tensor<8x16xf32>) + outs(%f : tensor<4x16xf32>) -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @op_between_dequantize_and_contraction( +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK: linalg.matmul + +// ----- + +// The remaining cases are bodies. A body qualifies only if it is exactly three +// operations, a multiply and an add and the yield, and those three multiply the +// two inputs and add the accumulator. The five cases here fail on the count: +// each computes a contraction, or something close to one, with extra operations +// around it. This one extends f16 inputs into an f32 accumulator, which is what +// a half precision model produces. +func.func @element_type_mismatch(%aq: tensor<4x8xi8>, %a_s: f16, %bq: tensor<8x16xi8>, %b_s: f16) -> tensor<4x16xf32> { + %a_i = tensor.empty() : tensor<4x8xf16> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s : tensor<4x8xi8>, f16) + outs(%a_i : tensor<4x8xf16>) -> tensor<4x8xf16> + %b_i = tensor.empty() : tensor<8x16xf16> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s : tensor<8x16xi8>, f16) + outs(%b_i : tensor<8x16xf16>) -> tensor<8x16xf16> + %cst = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<4x16xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, + affine_map<(d0, d1, d2) -> (d2, d1)>, + affine_map<(d0, d1, d2) -> (d0, d1)>], + iterator_types = ["parallel", "parallel", "reduction"]} + ins(%a, %b : tensor<4x8xf16>, tensor<8x16xf16>) + outs(%f : tensor<4x16xf32>) { + ^bb0(%in: f16, %in_0: f16, %out: f32): + %le = arith.extf %in : f16 to f32 + %re = arith.extf %in_0 : f16 to f32 + %m = arith.mulf %le, %re : f32 + %s = arith.addf %out, %m : f32 + linalg.yield %s : f32 + } -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @element_type_mismatch( +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK: arith.extf + +// ----- + +// A round trip through f16 rounds the lhs before multiplying it. +func.func @narrowed_input(%aq: tensor<1x8xi8>, %sa: f32, + %bq: tensor<8x1xi8>, %sb: f32) -> tensor<1x1xf32> { + %ai = tensor.empty() : tensor<1x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0,d1)->(d0,d1)>, affine_map<(d0,d1)->()>, affine_map<(d0,d1)->(d0,d1)>]} + ins(%aq, %sa : tensor<1x8xi8>, f32) + outs(%ai : tensor<1x8xf32>) -> tensor<1x8xf32> + %bi = tensor.empty() : tensor<8x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0,d1)->(d0,d1)>, affine_map<(d0,d1)->()>, affine_map<(d0,d1)->(d0,d1)>]} + ins(%bq, %sb : tensor<8x1xi8>, f32) + outs(%bi : tensor<8x1xf32>) -> tensor<8x1xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32> + %result = linalg.generic {indexing_maps = [affine_map<(m,n,k)->(m,k)>, affine_map<(m,n,k)->(k,n)>, affine_map<(m,n,k)->(m,n)>], iterator_types = ["parallel", "parallel", "reduction"]} + ins(%a, %b : tensor<1x8xf32>, tensor<8x1xf32>) outs(%init : tensor<1x1xf32>) { + ^bb0(%lhs: f32, %rhs: f32, %acc: f32): + %half = arith.truncf %lhs : f32 to f16 + %rounded = arith.extf %half : f16 to f32 + %product = arith.mulf %rounded, %rhs : f32 + %sum = arith.addf %acc, %product : f32 + linalg.yield %sum : f32 + } -> tensor<1x1xf32> + return %result : tensor<1x1xf32> +} +// CHECK-LABEL: func.func @narrowed_input( +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK: arith.truncf + +// ----- + +// A negation on the way into the multiply. +func.func @negated_input(%aq: tensor<1x8xi8>, %sa: f32, + %bq: tensor<8x1xi8>, %sb: f32) -> tensor<1x1xf32> { + %ai = tensor.empty() : tensor<1x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0,d1)->(d0,d1)>, affine_map<(d0,d1)->()>, affine_map<(d0,d1)->(d0,d1)>]} + ins(%aq, %sa : tensor<1x8xi8>, f32) + outs(%ai : tensor<1x8xf32>) -> tensor<1x8xf32> + %bi = tensor.empty() : tensor<8x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0,d1)->(d0,d1)>, affine_map<(d0,d1)->()>, affine_map<(d0,d1)->(d0,d1)>]} + ins(%bq, %sb : tensor<8x1xi8>, f32) + outs(%bi : tensor<8x1xf32>) -> tensor<8x1xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32> + %result = linalg.generic {indexing_maps = [affine_map<(m,n,k)->(m,k)>, affine_map<(m,n,k)->(k,n)>, affine_map<(m,n,k)->(m,n)>], iterator_types = ["parallel", "parallel", "reduction"]} + ins(%a, %b : tensor<1x8xf32>, tensor<8x1xf32>) outs(%init : tensor<1x1xf32>) { + ^bb0(%lhs: f32, %rhs: f32, %acc: f32): + %negative = arith.negf %lhs : f32 + %product = arith.mulf %negative, %rhs : f32 + %sum = arith.addf %acc, %product : f32 + linalg.yield %sum : f32 + } -> tensor<1x1xf32> + return %result : tensor<1x1xf32> +} +// CHECK-LABEL: func.func @negated_input( +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK: arith.negf + +// ----- + +// A negation on the accumulator the add reads. +func.func @negated_accumulator(%aq: tensor<1x8xi8>, %sa: f32, + %bq: tensor<8x1xi8>, %sb: f32) -> tensor<1x1xf32> { + %ai = tensor.empty() : tensor<1x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0,d1)->(d0,d1)>, affine_map<(d0,d1)->()>, affine_map<(d0,d1)->(d0,d1)>]} + ins(%aq, %sa : tensor<1x8xi8>, f32) + outs(%ai : tensor<1x8xf32>) -> tensor<1x8xf32> + %bi = tensor.empty() : tensor<8x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0,d1)->(d0,d1)>, affine_map<(d0,d1)->()>, affine_map<(d0,d1)->(d0,d1)>]} + ins(%bq, %sb : tensor<8x1xi8>, f32) + outs(%bi : tensor<8x1xf32>) -> tensor<8x1xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32> + %result = linalg.generic {indexing_maps = [affine_map<(m,n,k)->(m,k)>, affine_map<(m,n,k)->(k,n)>, affine_map<(m,n,k)->(m,n)>], iterator_types = ["parallel", "parallel", "reduction"]} + ins(%a, %b : tensor<1x8xf32>, tensor<8x1xf32>) outs(%init : tensor<1x1xf32>) { + ^bb0(%lhs: f32, %rhs: f32, %acc: f32): + %negative = arith.negf %acc : f32 + %product = arith.mulf %lhs, %rhs : f32 + %sum = arith.addf %negative, %product : f32 + linalg.yield %sum : f32 + } -> tensor<1x1xf32> + return %result : tensor<1x1xf32> +} +// CHECK-LABEL: func.func @negated_accumulator( +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK: arith.negf + +// ----- + +// A negation between the add and the yield. +func.func @negated_result(%aq: tensor<1x8xi8>, %sa: f32, + %bq: tensor<8x1xi8>, %sb: f32) -> tensor<1x1xf32> { + %ai = tensor.empty() : tensor<1x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0,d1)->(d0,d1)>, affine_map<(d0,d1)->()>, affine_map<(d0,d1)->(d0,d1)>]} + ins(%aq, %sa : tensor<1x8xi8>, f32) + outs(%ai : tensor<1x8xf32>) -> tensor<1x8xf32> + %bi = tensor.empty() : tensor<8x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0,d1)->(d0,d1)>, affine_map<(d0,d1)->()>, affine_map<(d0,d1)->(d0,d1)>]} + ins(%bq, %sb : tensor<8x1xi8>, f32) + outs(%bi : tensor<8x1xf32>) -> tensor<8x1xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32> + %result = linalg.generic {indexing_maps = [affine_map<(m,n,k)->(m,k)>, affine_map<(m,n,k)->(k,n)>, affine_map<(m,n,k)->(m,n)>], iterator_types = ["parallel", "parallel", "reduction"]} + ins(%a, %b : tensor<1x8xf32>, tensor<8x1xf32>) outs(%init : tensor<1x1xf32>) { + ^bb0(%lhs: f32, %rhs: f32, %acc: f32): + %product = arith.mulf %lhs, %rhs : f32 + %sum = arith.addf %acc, %product : f32 + %negative = arith.negf %sum : f32 + linalg.yield %negative : f32 + } -> tensor<1x1xf32> + return %result : tensor<1x1xf32> +} +// CHECK-LABEL: func.func @negated_result( +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK: arith.negf + +// ----- + +// The four cases below are three operations wide and still do not qualify. The +// maps and iterator types alone do not make an op a contraction: this one +// maximises instead of accumulating, so the expansion the rewrite performs +// would be meaningless. +func.func @non_mul_add_body(%aq: tensor<4x8xi8>, %a_s: f32, %bq: tensor<8x16xi8>, %b_s: f32) -> tensor<4x16xf32> { + %a_i = tensor.empty() : tensor<4x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s : tensor<4x8xi8>, f32) + outs(%a_i : tensor<4x8xf32>) -> tensor<4x8xf32> + %b_i = tensor.empty() : tensor<8x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s : tensor<8x16xi8>, f32) + outs(%b_i : tensor<8x16xf32>) -> tensor<8x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<4x16xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, + affine_map<(d0, d1, d2) -> (d2, d1)>, + affine_map<(d0, d1, d2) -> (d0, d1)>], + iterator_types = ["parallel", "parallel", "reduction"]} + ins(%a, %b : tensor<4x8xf32>, tensor<8x16xf32>) + outs(%f : tensor<4x16xf32>) { + ^bb0(%in: f32, %in_0: f32, %out: f32): + %m = arith.mulf %in, %in_0 : f32 + %s = arith.maximumf %out, %m : f32 + linalg.yield %s : f32 + } -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @non_mul_add_body( +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK: arith.maximumf + +// ----- + +// The accumulated value is a difference rather than a product. +func.func @non_mul_product(%aq: tensor<1x8xi8>, %sa: f32, + %bq: tensor<8x1xi8>, %sb: f32) -> tensor<1x1xf32> { + %ai = tensor.empty() : tensor<1x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0,d1)->(d0,d1)>, affine_map<(d0,d1)->()>, affine_map<(d0,d1)->(d0,d1)>]} + ins(%aq, %sa : tensor<1x8xi8>, f32) + outs(%ai : tensor<1x8xf32>) -> tensor<1x8xf32> + %bi = tensor.empty() : tensor<8x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0,d1)->(d0,d1)>, affine_map<(d0,d1)->()>, affine_map<(d0,d1)->(d0,d1)>]} + ins(%bq, %sb : tensor<8x1xi8>, f32) + outs(%bi : tensor<8x1xf32>) -> tensor<8x1xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32> + %result = linalg.generic {indexing_maps = [affine_map<(m,n,k)->(m,k)>, affine_map<(m,n,k)->(k,n)>, affine_map<(m,n,k)->(m,n)>], iterator_types = ["parallel", "parallel", "reduction"]} + ins(%a, %b : tensor<1x8xf32>, tensor<8x1xf32>) outs(%init : tensor<1x1xf32>) { + ^bb0(%lhs: f32, %rhs: f32, %acc: f32): + %difference = arith.subf %lhs, %rhs : f32 + %sum = arith.addf %acc, %difference : f32 + linalg.yield %sum : f32 + } -> tensor<1x1xf32> + return %result : tensor<1x1xf32> +} +// CHECK-LABEL: func.func @non_mul_product( +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK: arith.subf + +// ----- + +// A product of the lhs with itself. The rhs reaches the op but not the +// multiply, so substituting the quantized operands would change what is +// computed. +func.func @multiply_ignores_rhs(%aq: tensor<1x8xi8>, %sa: f32, + %bq: tensor<8x1xi8>, %sb: f32) -> tensor<1x1xf32> { + %ai = tensor.empty() : tensor<1x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0,d1)->(d0,d1)>, affine_map<(d0,d1)->()>, affine_map<(d0,d1)->(d0,d1)>]} + ins(%aq, %sa : tensor<1x8xi8>, f32) + outs(%ai : tensor<1x8xf32>) -> tensor<1x8xf32> + %bi = tensor.empty() : tensor<8x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0,d1)->(d0,d1)>, affine_map<(d0,d1)->()>, affine_map<(d0,d1)->(d0,d1)>]} + ins(%bq, %sb : tensor<8x1xi8>, f32) + outs(%bi : tensor<8x1xf32>) -> tensor<8x1xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32> + %result = linalg.generic {indexing_maps = [affine_map<(m,n,k)->(m,k)>, affine_map<(m,n,k)->(k,n)>, affine_map<(m,n,k)->(m,n)>], iterator_types = ["parallel", "parallel", "reduction"]} + ins(%a, %b : tensor<1x8xf32>, tensor<8x1xf32>) outs(%init : tensor<1x1xf32>) { + ^bb0(%lhs: f32, %rhs: f32, %acc: f32): + %square = arith.mulf %lhs, %lhs : f32 + %sum = arith.addf %acc, %square : f32 + linalg.yield %sum : f32 + } -> tensor<1x1xf32> + return %result : tensor<1x1xf32> +} +// CHECK-LABEL: func.func @multiply_ignores_rhs( +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK: ^bb0(%[[L:[a-zA-Z0-9_]+]]: f32, %{{.+}}: f32, %{{.+}}: f32): +// CHECK: arith.mulf %[[L]], %[[L]] + +// ----- + +// The add doubles the product instead of accumulating it, so the op is not a +// reduction over the multiply however its iterator types are marked. +func.func @add_drops_accumulator(%aq: tensor<1x8xi8>, %sa: f32, + %bq: tensor<8x1xi8>, %sb: f32) -> tensor<1x1xf32> { + %ai = tensor.empty() : tensor<1x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0,d1)->(d0,d1)>, affine_map<(d0,d1)->()>, affine_map<(d0,d1)->(d0,d1)>]} + ins(%aq, %sa : tensor<1x8xi8>, f32) + outs(%ai : tensor<1x8xf32>) -> tensor<1x8xf32> + %bi = tensor.empty() : tensor<8x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0,d1)->(d0,d1)>, affine_map<(d0,d1)->()>, affine_map<(d0,d1)->(d0,d1)>]} + ins(%bq, %sb : tensor<8x1xi8>, f32) + outs(%bi : tensor<8x1xf32>) -> tensor<8x1xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32> + %result = linalg.generic {indexing_maps = [affine_map<(m,n,k)->(m,k)>, affine_map<(m,n,k)->(k,n)>, affine_map<(m,n,k)->(m,n)>], iterator_types = ["parallel", "parallel", "reduction"]} + ins(%a, %b : tensor<1x8xf32>, tensor<8x1xf32>) outs(%init : tensor<1x1xf32>) { + ^bb0(%lhs: f32, %rhs: f32, %acc: f32): + %product = arith.mulf %lhs, %rhs : f32 + %doubled = arith.addf %product, %product : f32 + linalg.yield %doubled : f32 + } -> tensor<1x1xf32> + return %result : tensor<1x1xf32> +} +// CHECK-LABEL: func.func @add_drops_accumulator( +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK: %[[P:.+]] = arith.mulf +// CHECK: arith.addf %[[P]], %[[P]] + +//===----------------------------------------------------------------------===// +// Reduction legality +// +// A reduction dim can become integer only when both inputs index it, the output +// does not, and every quantization parameter is invariant along it. Its extent +// has to be statically known and nonzero, and at least one such dim has to +// remain. +//===----------------------------------------------------------------------===// + +// ----- + +// Per-channel along K, the only reduction dim, so no sub-reduction has constant +// parameters and there is no integer contraction to form. +func.func @form_none(%aq: tensor<4x8xi8>, %a_s: tensor<8xf32>, + %bq: tensor<8x16xi8>, %b_s: f32) -> tensor<4x16xf32> { + %a_i = tensor.empty() : tensor<4x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> (d1)>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s : tensor<4x8xi8>, tensor<8xf32>) + outs(%a_i : tensor<4x8xf32>) -> tensor<4x8xf32> + %b_i = tensor.empty() : tensor<8x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s : tensor<8x16xi8>, f32) + outs(%b_i : tensor<8x16xf32>) -> tensor<8x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<4x16xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.matmul ins(%a, %b : tensor<4x8xf32>, tensor<8x16xf32>) + outs(%f : tensor<4x16xf32>) -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @form_none( +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK: linalg.matmul + +// ----- + +// An unknown reduction extent cannot be bounded against the i32 accumulator. +func.func @dynamic_reduction(%aq: tensor<4x?xi8>, %a_s: f32, %a_z: i8, + %bq: tensor<?x16xi8>, %b_s: f32, %b_z: i8) -> tensor<4x16xf32> { + %c1 = arith.constant 1 : index + %k = tensor.dim %aq, %c1 : tensor<4x?xi8> + %a_i = tensor.empty(%k) : tensor<4x?xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s, %a_z : tensor<4x?xi8>, f32, i8) + outs(%a_i : tensor<4x?xf32>) -> tensor<4x?xf32> + %b_i = tensor.empty(%k) : tensor<?x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s, %b_z : tensor<?x16xi8>, f32, i8) + outs(%b_i : tensor<?x16xf32>) -> tensor<?x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<4x16xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.matmul ins(%a, %b : tensor<4x?xf32>, tensor<?x16xf32>) + outs(%f : tensor<4x16xf32>) -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @dynamic_reduction( +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK: linalg.matmul + +// ----- + +// The same with the reduction split over two dims: the product of the extents +// is what has to be bounded, so one unknown factor is enough to decline. +func.func @dynamic_multi_dim_reduction(%aq: tensor<4x?x?xi8>, %a_s: f32, %a_z: i8, + %bq: tensor<?x?x16xi8>, %b_s: f32, %b_z: i8) -> tensor<4x16xf32> { + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %k0 = tensor.dim %aq, %c1 : tensor<4x?x?xi8> + %k1 = tensor.dim %aq, %c2 : tensor<4x?x?xi8> + %a_i = tensor.empty(%k0, %k1) : tensor<4x?x?xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, + affine_map<(d0, d1, d2) -> ()>, + affine_map<(d0, d1, d2) -> ()>, + affine_map<(d0, d1, d2) -> (d0, d1, d2)>]} + ins(%aq, %a_s, %a_z : tensor<4x?x?xi8>, f32, i8) + outs(%a_i : tensor<4x?x?xf32>) -> tensor<4x?x?xf32> + %b_i = tensor.empty(%k0, %k1) : tensor<?x?x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, + affine_map<(d0, d1, d2) -> ()>, + affine_map<(d0, d1, d2) -> ()>, + affine_map<(d0, d1, d2) -> (d0, d1, d2)>]} + ins(%bq, %b_s, %b_z : tensor<?x?x16xi8>, f32, i8) + outs(%b_i : tensor<?x?x16xf32>) -> tensor<?x?x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<4x16xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.contract + indexing_maps = [affine_map<(m, n, k0, k1) -> (m, k0, k1)>, + affine_map<(m, n, k0, k1) -> (k0, k1, n)>, + affine_map<(m, n, k0, k1) -> (m, n)>] + ins(%a, %b : tensor<4x?x?xf32>, tensor<?x?x16xf32>) + outs(%f : tensor<4x16xf32>) -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @dynamic_multi_dim_reduction( +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK: linalg.contract + +// ----- + +// An empty reduction retains its original zero result rather than acquiring one +// from a correction term the contraction never produced. +func.func @empty_reduction(%aq: tensor<1x0xi8>, %sa: f32, + %bq: tensor<0x1xi8>, %sb: f32) -> tensor<1x1xf32> { + %ai = tensor.empty() : tensor<1x0xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0,d1)->(d0,d1)>, affine_map<(d0,d1)->()>, affine_map<(d0,d1)->(d0,d1)>]} + ins(%aq, %sa : tensor<1x0xi8>, f32) + outs(%ai : tensor<1x0xf32>) -> tensor<1x0xf32> + %bi = tensor.empty() : tensor<0x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0,d1)->(d0,d1)>, affine_map<(d0,d1)->()>, affine_map<(d0,d1)->(d0,d1)>]} + ins(%bq, %sb : tensor<0x1xi8>, f32) + outs(%bi : tensor<0x1xf32>) -> tensor<0x1xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32> + %result = linalg.matmul ins(%a, %b : tensor<1x0xf32>, tensor<0x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + return %result : tensor<1x1xf32> +} +// CHECK-LABEL: func.func @empty_reduction( +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK: linalg.matmul + +// ----- + +// Linalg permits a reduction-marked dimension to index the output, but it is +// not a contraction reduction. Reject it whether its scale is invariant or +// varying: it must neither be summed away nor appended as a duplicate partial +// result axis. + +func.func @reduction_dim_in_output(%aq: tensor<2x4xi8>, %sa: f32, + %varying_scale: tensor<2xf32>, %bq: tensor<2x4xi8>, %sb: f32) + -> (tensor<2xf32>, tensor<2xf32>) { + %ai = tensor.empty() : tensor<2x4xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %sa : tensor<2x4xi8>, f32) + outs(%ai : tensor<2x4xf32>) -> tensor<2x4xf32> + %a_varying = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> (d0)>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %varying_scale : tensor<2x4xi8>, tensor<2xf32>) + outs(%ai : tensor<2x4xf32>) -> tensor<2x4xf32> + %bi = tensor.empty() : tensor<2x4xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %sb : tensor<2x4xi8>, f32) + outs(%bi : tensor<2x4xf32>) -> tensor<2x4xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<2xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<2xf32>) -> tensor<2xf32> + %constant = linalg.generic { + indexing_maps = [affine_map<(m, k) -> (m, k)>, + affine_map<(m, k) -> (m, k)>, + affine_map<(m, k) -> (m)>], + iterator_types = ["reduction", "reduction"]} + ins(%a, %b : tensor<2x4xf32>, tensor<2x4xf32>) outs(%init : tensor<2xf32>) { + ^bb0(%lhs: f32, %rhs: f32, %acc: f32): + %product = arith.mulf %lhs, %rhs : f32 + %sum = arith.addf %product, %acc : f32 + linalg.yield %sum : f32 + } -> tensor<2xf32> + %varying = linalg.generic { + indexing_maps = [affine_map<(m, k) -> (m, k)>, + affine_map<(m, k) -> (m, k)>, + affine_map<(m, k) -> (m)>], + iterator_types = ["reduction", "reduction"]} + ins(%a_varying, %b : tensor<2x4xf32>, tensor<2x4xf32>) outs(%init : tensor<2xf32>) { + ^bb0(%lhs: f32, %rhs: f32, %acc: f32): + %product = arith.mulf %lhs, %rhs : f32 + %sum = arith.addf %product, %acc : f32 + linalg.yield %sum : f32 + } -> tensor<2xf32> + return %constant, %varying : tensor<2xf32>, tensor<2xf32> +} +// CHECK-LABEL: func.func @reduction_dim_in_output( +// CHECK-COUNT-3: iree_linalg_ext.dequantize_affine +// CHECK: linalg.generic {{.*}} ins(%{{.+}}, %{{.+}} : tensor<2x4xf32>, tensor<2x4xf32>) +// CHECK: linalg.generic {{.*}} ins(%{{.+}}, %{{.+}} : tensor<2x4xf32>, tensor<2x4xf32>) + +//===----------------------------------------------------------------------===// +// Accumulator bound +// +// Every integer intermediate must fit in i32. Quantization ranges bound the +// inputs when present; otherwise the quantized storage types do. Zero points +// are points on those storage grids, regardless of their SSA carrier types. +// Ranges are propagated through the products and correction stages; the final +// result is bounded using the centered product (Aq-zA)*(Bq-zB). +// The admissible counterpart of each pair below is in +// convert_qdq_to_integer_math_algebra.mlir. +//===----------------------------------------------------------------------===// + +// ----- + +// The shape and types of @acc_symmetric_fits_deep with both sides asymmetric. +// Centered values can reach +/-255, so N*255^2 exceeds INT32_MAX. +func.func @acc_asymmetric_too_deep(%aq: tensor<4x65536xi8>, %a_s: f32, %a_z: i8, %bq: tensor<65536x16xi8>, %b_s: f32, %b_z: i8) -> tensor<4x16xf32> { + %a_i = tensor.empty() : tensor<4x65536xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s, %a_z : tensor<4x65536xi8>, f32, i8) + outs(%a_i : tensor<4x65536xf32>) -> tensor<4x65536xf32> + %b_i = tensor.empty() : tensor<65536x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s, %b_z : tensor<65536x16xi8>, f32, i8) + outs(%b_i : tensor<65536x16xf32>) -> tensor<65536x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<4x16xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.matmul ins(%a, %b : tensor<4x65536xf32>, tensor<65536x16xf32>) + outs(%f : tensor<4x16xf32>) -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @acc_asymmetric_too_deep( +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK: linalg.matmul + +// ----- + +// One element deeper than @acc_last_safe_depth: all -128 inputs sum to +2^31, +// outside the signed i32 range. +func.func @acc_positive_endpoint(%aq: tensor<1x131072xi8>, %sa: f32, + %bq: tensor<131072x1xi8>, %sb: f32) -> tensor<1x1xf32> { + %ai = tensor.empty() : tensor<1x131072xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0,d1)->(d0,d1)>, affine_map<(d0,d1)->()>, affine_map<(d0,d1)->(d0,d1)>]} + ins(%aq, %sa : tensor<1x131072xi8>, f32) + outs(%ai : tensor<1x131072xf32>) -> tensor<1x131072xf32> + %bi = tensor.empty() : tensor<131072x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0,d1)->(d0,d1)>, affine_map<(d0,d1)->()>, affine_map<(d0,d1)->(d0,d1)>]} + ins(%bq, %sb : tensor<131072x1xi8>, f32) + outs(%bi : tensor<131072x1xf32>) -> tensor<131072x1xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32> + %result = linalg.matmul ins(%a, %b : tensor<1x131072xf32>, tensor<131072x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + return %result : tensor<1x1xf32> +} +// CHECK-LABEL: func.func @acc_positive_endpoint( +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK: linalg.matmul + +// Both operands i16: 15+15 magnitude bits leaves no room for even a shallow +// reduction. +func.func @acc_wide_storage(%aq: tensor<4x8xi16>, %a_s: f32, %bq: tensor<8x16xi16>, %b_s: f32) -> tensor<4x16xf32> { + %a_i = tensor.empty() : tensor<4x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s : tensor<4x8xi16>, f32) + outs(%a_i : tensor<4x8xf32>) -> tensor<4x8xf32> + %b_i = tensor.empty() : tensor<8x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s : tensor<8x16xi16>, f32) + outs(%b_i : tensor<8x16xf32>) -> tensor<8x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<4x16xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.matmul ins(%a, %b : tensor<4x8xf32>, tensor<8x16xf32>) + outs(%f : tensor<4x16xf32>) -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @acc_wide_storage( +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK: linalg.matmul + +// ----- + +// i64 storage is at the accumulator width, so the operand is rejected on width +// before any extent is consulted and the unknown depth never matters. +func.func @acc_dynamic_wide_storage(%aq: tensor<1x?xi64>, %sa: f32, + %bq: tensor<?x1xi64>, %sb: f32) -> tensor<1x1xf32> { + %dim = arith.constant 1 : index + %k = tensor.dim %aq, %dim : tensor<1x?xi64> + %ai = tensor.empty(%k) : tensor<1x?xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0,d1)->(d0,d1)>, affine_map<(d0,d1)->()>, affine_map<(d0,d1)->(d0,d1)>]} + ins(%aq, %sa : tensor<1x?xi64>, f32) + outs(%ai : tensor<1x?xf32>) -> tensor<1x?xf32> + %bi = tensor.empty(%k) : tensor<?x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0,d1)->(d0,d1)>, affine_map<(d0,d1)->()>, affine_map<(d0,d1)->(d0,d1)>]} + ins(%bq, %sb : tensor<?x1xi64>, f32) + outs(%bi : tensor<?x1xf32>) -> tensor<?x1xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32> + %result = linalg.matmul ins(%a, %b : tensor<1x?xf32>, tensor<?x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + return %result : tensor<1x1xf32> +} +// CHECK-LABEL: func.func @acc_dynamic_wide_storage( +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK: linalg.matmul + +// ----- + +// The centered result fits, but P+Q can overflow: reject one element +// beyond floor(INT32_MAX/(2*255^2)). +func.func @unsigned_reduction_16513(%aq: tensor<1x16513xi8>, %a_s: f32, %a_z: i8, %bq: tensor<16513x1xi8>, %b_s: f32, %b_z: i8) -> tensor<1x1xf32> { + %a_i = tensor.empty() : tensor<1x16513xf32> + %a = iree_linalg_ext.dequantize_affine + {input_unsigned, zp_unsigned, + indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s, %a_z : tensor<1x16513xi8>, f32, i8) + outs(%a_i : tensor<1x16513xf32>) -> tensor<1x16513xf32> + %b_i = tensor.empty() : tensor<16513x1xf32> + %b = iree_linalg_ext.dequantize_affine + {input_unsigned, zp_unsigned, + indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s, %b_z : tensor<16513x1xi8>, f32, i8) + outs(%b_i : tensor<16513x1xf32>) -> tensor<16513x1xf32> + %cst = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<1x1xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<1x1xf32>) -> tensor<1x1xf32> + %c = linalg.matmul ins(%a, %b : tensor<1x16513xf32>, tensor<16513x1xf32>) + outs(%f : tensor<1x1xf32>) -> tensor<1x1xf32> + return %c : tensor<1x1xf32> +} +// CHECK-LABEL: func.func @unsigned_reduction_16513( +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK: linalg.matmul + +// ----- + +// The centered product is in [-255^2, 255^2] and bounds all intermediates. +// Reject one element beyond the safe depth. +func.func @signed_reduction_33026(%aq: tensor<1x33026xi8>, %a_s: f32, %a_z: i8, %bq: tensor<33026x1xi8>, %b_s: f32, %b_z: i8) -> tensor<1x1xf32> { + %a_i = tensor.empty() : tensor<1x33026xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s, %a_z : tensor<1x33026xi8>, f32, i8) + outs(%a_i : tensor<1x33026xf32>) -> tensor<1x33026xf32> + %b_i = tensor.empty() : tensor<33026x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s, %b_z : tensor<33026x1xi8>, f32, i8) + outs(%b_i : tensor<33026x1xf32>) -> tensor<33026x1xf32> + %cst = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<1x1xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<1x1xf32>) -> tensor<1x1xf32> + %c = linalg.matmul ins(%a, %b : tensor<1x33026xf32>, tensor<33026x1xf32>) + outs(%f : tensor<1x1xf32>) -> tensor<1x1xf32> + return %c : tensor<1x1xf32> +} +// CHECK-LABEL: func.func @signed_reduction_33026( +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK: linalg.matmul + +// ----- + +// The correction interval is [-97665, 97410], larger than the centered product. +// Reject one element beyond the safe depth. +func.func @mixed_reduction_21989(%aq: tensor<1x21989xi8>, %a_s: f32, %a_z: i8, %bq: tensor<21989x1xi8>, %b_s: f32, %b_z: i8) -> tensor<1x1xf32> { + %a_i = tensor.empty() : tensor<1x21989xf32> + %a = iree_linalg_ext.dequantize_affine + {input_unsigned, zp_unsigned, + indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s, %a_z : tensor<1x21989xi8>, f32, i8) + outs(%a_i : tensor<1x21989xf32>) -> tensor<1x21989xf32> + %b_i = tensor.empty() : tensor<21989x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s, %b_z : tensor<21989x1xi8>, f32, i8) + outs(%b_i : tensor<21989x1xf32>) -> tensor<21989x1xf32> + %cst = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<1x1xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<1x1xf32>) -> tensor<1x1xf32> + %c = linalg.matmul ins(%a, %b : tensor<1x21989xf32>, tensor<21989x1xf32>) + outs(%f : tensor<1x1xf32>) -> tensor<1x1xf32> + return %c : tensor<1x1xf32> +} +// CHECK-LABEL: func.func @mixed_reduction_21989( +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK: linalg.matmul
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/convert_qdq_to_integer_math_pipeline.mlir b/compiler/src/iree/compiler/GlobalOptimization/test/convert_qdq_to_integer_math_pipeline.mlir new file mode 100644 index 0000000..86b2dfc --- /dev/null +++ b/compiler/src/iree/compiler/GlobalOptimization/test/convert_qdq_to_integer_math_pipeline.mlir
@@ -0,0 +1,170 @@ +// RUN: iree-opt --split-input-file --iree-global-optimization-transformation-pipeline --iree-global-opt-enable-qdq-to-integer-math=true %s | FileCheck %s --implicit-check-not=iree_linalg_ext.dequantize_affine +// RUN: iree-opt --split-input-file --iree-global-optimization-transformation-pipeline --iree-global-opt-enable-qdq-to-integer-math=false %s | FileCheck %s --check-prefix=DISABLED +// RUN: iree-opt --split-input-file --iree-global-optimization-transformation-pipeline="use-im2col-for-convs=true" --iree-global-opt-enable-qdq-to-integer-math=true %s | FileCheck %s --implicit-check-not=iree_linalg_ext.dequantize_affine + +// Exercise the production pipeline, including unit-dimension folding and +// transpose propagation, rather than invoking the QDQ pass in isolation. + +// ----- + +// Unit batch dimensions must not hide dequantize producers before the rewrite. +util.func public @convolution_zero_point_corrections( + %aq: tensor<1x5x5x1xi8>, + %bq: tensor<2x2x1x2xi8>, + %sa: f32, + %sb: tensor<2xf32>, + %za: i8, + %zb: tensor<2xi8>) -> tensor<1x2x2x2xf32> { + %ai = tensor.empty() : tensor<1x5x5x1xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(n, h, w, c) -> (n, h, w, c)>, + affine_map<(n, h, w, c) -> ()>, + affine_map<(n, h, w, c) -> ()>, + affine_map<(n, h, w, c) -> (n, h, w, c)>]} + ins(%aq, %sa, %za : tensor<1x5x5x1xi8>, f32, i8) + outs(%ai : tensor<1x5x5x1xf32>) -> tensor<1x5x5x1xf32> + %bi = tensor.empty() : tensor<2x2x1x2xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(h, w, c, f) -> (h, w, c, f)>, + affine_map<(h, w, c, f) -> (f)>, + affine_map<(h, w, c, f) -> (f)>, + affine_map<(h, w, c, f) -> (h, w, c, f)>]} + ins(%bq, %sb, %zb : tensor<2x2x1x2xi8>, tensor<2xf32>, tensor<2xi8>) + outs(%bi : tensor<2x2x1x2xf32>) -> tensor<2x2x1x2xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x2x2x2xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32> + %result = linalg.conv_2d_nhwc_hwcf {dilations = dense<2> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} + ins(%a, %b : tensor<1x5x5x1xf32>, tensor<2x2x1x2xf32>) + outs(%init : tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32> + util.return %result : tensor<1x2x2x2xf32> +} +// CHECK-LABEL: util.func public @convolution_zero_point_corrections( +// CHECK: arith.muli %{{.+}}, %{{.+}} : i32 +// CHECK: arith.sitofp +// DISABLED-LABEL: util.func public @convolution_zero_point_corrections( +// DISABLED: iree_linalg_ext.dequantize_affine +// DISABLED: arith.mulf %{{.+}}, %{{.+}} : f32 + +// ----- + +// Unit batch dimensions must not hide dequantize producers before the rewrite. +util.func public @convolution_nchw_asymmetric( + %aq: tensor<1x2x3x4xi8>, + %bq: tensor<2x2x2x2xi8>, + %sb: tensor<2xf32>, + %zb: tensor<2xi8>, + %sa: f32, + %za: i8) -> tensor<1x2x2x3xf32> { + %ai = tensor.empty() : tensor<1x2x3x4xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(n, c, h, w) -> (n, c, h, w)>, + affine_map<(n, c, h, w) -> ()>, + affine_map<(n, c, h, w) -> ()>, + affine_map<(n, c, h, w) -> (n, c, h, w)>]} + ins(%aq, %sa, %za : tensor<1x2x3x4xi8>, f32, i8) + outs(%ai : tensor<1x2x3x4xf32>) -> tensor<1x2x3x4xf32> + %bi = tensor.empty() : tensor<2x2x2x2xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(f, c, h, w) -> (f, c, h, w)>, + affine_map<(f, c, h, w) -> (f)>, + affine_map<(f, c, h, w) -> (f)>, + affine_map<(f, c, h, w) -> (f, c, h, w)>]} + ins(%bq, %sb, %zb : tensor<2x2x2x2xi8>, tensor<2xf32>, tensor<2xi8>) + outs(%bi : tensor<2x2x2x2xf32>) -> tensor<2x2x2x2xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x2x2x3xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32> + %result = linalg.conv_2d_nchw_fchw {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} + ins(%a, %b : tensor<1x2x3x4xf32>, tensor<2x2x2x2xf32>) + outs(%init : tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32> + util.return %result : tensor<1x2x2x3xf32> +} +// CHECK-LABEL: util.func public @convolution_nchw_asymmetric( +// CHECK: arith.muli %{{.+}}, %{{.+}} : i32 +// CHECK: arith.sitofp +// DISABLED-LABEL: util.func public @convolution_nchw_asymmetric( +// DISABLED: iree_linalg_ext.dequantize_affine +// DISABLED: arith.mulf %{{.+}}, %{{.+}} : f32 + +// ----- + +// Unit batch dimensions must not hide dequantize producers before the rewrite. +util.func public @depthwise_asymmetric_strided_dilated( + %aq: tensor<1x4x5x2xi8>, + %bq: tensor<2x2x2xi8>, + %sa: tensor<2xf32>, + %za: tensor<2xi8>, + %sb: tensor<2xf32>, + %zb: tensor<2xi8>) -> tensor<1x2x2x2xf32> { + %ai = tensor.empty() : tensor<1x4x5x2xf32> + %a = iree_linalg_ext.dequantize_affine + {input_unsigned, zp_unsigned, indexing_maps = [affine_map<(n, h, w, c) -> (n, h, w, c)>, + affine_map<(n, h, w, c) -> (c)>, + affine_map<(n, h, w, c) -> (c)>, + affine_map<(n, h, w, c) -> (n, h, w, c)>]} + ins(%aq, %sa, %za : tensor<1x4x5x2xi8>, tensor<2xf32>, tensor<2xi8>) + outs(%ai : tensor<1x4x5x2xf32>) -> tensor<1x4x5x2xf32> + %bi = tensor.empty() : tensor<2x2x2xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(h, w, c) -> (h, w, c)>, + affine_map<(h, w, c) -> (c)>, + affine_map<(h, w, c) -> (c)>, + affine_map<(h, w, c) -> (h, w, c)>]} + ins(%bq, %sb, %zb : tensor<2x2x2xi8>, tensor<2xf32>, tensor<2xi8>) + outs(%bi : tensor<2x2x2xf32>) -> tensor<2x2x2xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x2x2x2xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32> + %result = linalg.depthwise_conv_2d_nhwc_hwc {dilations = dense<[2, 1]> : tensor<2xi64>, strides = dense<[1, 2]> : tensor<2xi64>} + ins(%a, %b : tensor<1x4x5x2xf32>, tensor<2x2x2xf32>) + outs(%init : tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32> + util.return %result : tensor<1x2x2x2xf32> +} +// CHECK-LABEL: util.func public @depthwise_asymmetric_strided_dilated( +// CHECK: arith.muli %{{.+}}, %{{.+}} : i32 +// CHECK: arith.sitofp +// DISABLED-LABEL: util.func public @depthwise_asymmetric_strided_dilated( +// DISABLED: iree_linalg_ext.dequantize_affine +// DISABLED: arith.mulf %{{.+}}, %{{.+}} : f32 + +// ----- + +// The early rewrite cannot see through the RHS transpose. Keep a later +// opportunity after transpose propagation exposes the dequantize producer. +util.func public @transposed_dequantized_rhs( + %aq: tensor<4x8xi8>, %bq: tensor<3x8xi8>, %sa: f32, %sb: f32, + %za: i8, %zb: i8) -> tensor<4x3xf32> { + %ai = tensor.empty() : tensor<4x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(m, k) -> (m, k)>, + affine_map<(m, k) -> ()>, + affine_map<(m, k) -> ()>, + affine_map<(m, k) -> (m, k)>]} + ins(%aq, %sa, %za : tensor<4x8xi8>, f32, i8) + outs(%ai : tensor<4x8xf32>) -> tensor<4x8xf32> + %bi = tensor.empty() : tensor<3x8xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(n, k) -> (n, k)>, + affine_map<(n, k) -> ()>, + affine_map<(n, k) -> ()>, + affine_map<(n, k) -> (n, k)>]} + ins(%bq, %sb, %zb : tensor<3x8xi8>, f32, i8) + outs(%bi : tensor<3x8xf32>) -> tensor<3x8xf32> + %ti = tensor.empty() : tensor<8x3xf32> + %bt = linalg.transpose ins(%b : tensor<3x8xf32>) + outs(%ti : tensor<8x3xf32>) permutation = [1, 0] + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<4x3xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<4x3xf32>) -> tensor<4x3xf32> + %result = linalg.matmul ins(%a, %bt : tensor<4x8xf32>, tensor<8x3xf32>) + outs(%init : tensor<4x3xf32>) -> tensor<4x3xf32> + util.return %result : tensor<4x3xf32> +} +// CHECK-LABEL: util.func public @transposed_dequantized_rhs( +// CHECK: arith.muli %{{.+}}, %{{.+}} : i32 +// CHECK: arith.sitofp +// DISABLED-LABEL: util.func public @transposed_dequantized_rhs( +// DISABLED: iree_linalg_ext.dequantize_affine +// DISABLED: linalg.matmul +// DISABLED-SAME: outs(%{{.+}} : tensor<4x3xf32>)
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/convert_qdq_to_integer_math_scale_range.mlir b/compiler/src/iree/compiler/GlobalOptimization/test/convert_qdq_to_integer_math_scale_range.mlir new file mode 100644 index 0000000..b529b56 --- /dev/null +++ b/compiler/src/iree/compiler/GlobalOptimization/test/convert_qdq_to_integer_math_scale_range.mlir
@@ -0,0 +1,289 @@ +// RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(func.func(iree-global-opt-convert-qdq-to-integer-math,linalg-inline-scalar-operands,canonicalize))" %s | FileCheck %s + +// Numerical corner cases in integer contractions and their scaling epilogues. +// Inline constant scalar scales and canonicalize to expose the arithmetic +// retained by the rewrite. FIXME checks record current behavior, not a promise +// that these numerical differences should be preserved. + +// This contraction returns zero: (1 * 1e20) * (0 * 1e20) = 0. +// Apply each scale to the accumulator in sequence; combining the scales first +// would overflow to infinity and produce NaN from 0 * infinity. +func.func @scale_product_overflow() -> tensor<1x1xf32> { + %aq = arith.constant dense<1> : tensor<1x1xi8> + %bq = arith.constant dense<0> : tensor<1x1xi8> + %scale = arith.constant 1.0e20 : f32 + %ai = tensor.empty() : tensor<1x1xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(m, k) -> (m, k)>, + affine_map<(m, k) -> ()>, + affine_map<(m, k) -> (m, k)>]} + ins(%aq, %scale : tensor<1x1xi8>, f32) + outs(%ai : tensor<1x1xf32>) -> tensor<1x1xf32> + %bi = tensor.empty() : tensor<1x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(k, n) -> (k, n)>, + affine_map<(k, n) -> ()>, + affine_map<(k, n) -> (k, n)>]} + ins(%bq, %scale : tensor<1x1xi8>, f32) + outs(%bi : tensor<1x1xf32>) -> tensor<1x1xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32> + %result = linalg.matmul ins(%a, %b : tensor<1x1xf32>, tensor<1x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + return %result : tensor<1x1xf32> +} +// CHECK-LABEL: func.func @scale_product_overflow( +// CHECK: %[[SCALE:.+]] = arith.constant 1.000000e+20 : f32 +// CHECK: linalg.generic {{.*}} outs(%{{.+}} : tensor<1x1xi32>) +// CHECK: %[[REAL:.+]] = arith.sitofp %{{.+}} : i32 to f32 +// CHECK-NEXT: %[[PARTIAL:.+]] = arith.mulf %[[REAL]], %[[SCALE]] : f32 +// CHECK-NEXT: %[[RESULT:.+]] = arith.mulf %[[PARTIAL]], %[[SCALE]] : f32 +// CHECK-NEXT: linalg.yield %[[RESULT]] : f32 + +// ----- + +// FIXME: Preserve dequantization rounding, or explicitly document that removing +// it is permitted by this pass's numerical contract. With scales of 1, f16 +// dequantization rounds 2049 to 2048, so the original dot product is zero. +// The rewrite instead computes -2048 + 2049 = 1 in i32 and rounds only the +// final result to f16. The checks below record the bypassed input rounding. +func.func @f16_dequantization_rounding() -> tensor<1x1xf16> { + %aq = arith.constant dense<[[-1, 1]]> : tensor<1x2xi8> + %bq = arith.constant dense<[[2048], [2049]]> : tensor<2x1xi16> + %scale = arith.constant 1.0 : f32 + %ai = tensor.empty() : tensor<1x2xf16> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(m, k) -> (m, k)>, + affine_map<(m, k) -> ()>, + affine_map<(m, k) -> (m, k)>]} + ins(%aq, %scale : tensor<1x2xi8>, f32) + outs(%ai : tensor<1x2xf16>) -> tensor<1x2xf16> + %bi = tensor.empty() : tensor<2x1xf16> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(k, n) -> (k, n)>, + affine_map<(k, n) -> ()>, + affine_map<(k, n) -> (k, n)>]} + ins(%bq, %scale : tensor<2x1xi16>, f32) + outs(%bi : tensor<2x1xf16>) -> tensor<2x1xf16> + %zero = arith.constant 0.0 : f16 + %empty = tensor.empty() : tensor<1x1xf16> + %init = linalg.fill ins(%zero : f16) outs(%empty : tensor<1x1xf16>) -> tensor<1x1xf16> + %result = linalg.matmul ins(%a, %b : tensor<1x2xf16>, tensor<2x1xf16>) + outs(%init : tensor<1x1xf16>) -> tensor<1x1xf16> + return %result : tensor<1x1xf16> +} +// CHECK-LABEL: func.func @f16_dequantization_rounding( +// CHECK: linalg.generic {{.*}} outs(%{{.+}} : tensor<1x1xi32>) +// CHECK: arith.extsi %{{.+}} : i16 to i32 +// CHECK: linalg.generic {{.*}} outs(%{{.+}} : tensor<1x1xf16>) +// CHECK: %[[REAL:.+]] = arith.sitofp %{{.+}} : i32 to f16 +// CHECK-NEXT: linalg.yield %[[REAL]] : f16 + +// ----- + +// FIXME: Preserve signed-zero semantics or reject a negative-zero init. The +// original computation is -0 + (0 * -1) = -0, with both scales equal to 1. +// The rewrite accepts the init but replaces it with integer zero, so its +// result is +0. The checks record the zero init and direct integer-to-float +// conversion that loses the original sign. +func.func @negative_zero_init() -> tensor<1x1xf32> { + %aq = arith.constant dense<0> : tensor<1x1xi8> + %bq = arith.constant dense<-1> : tensor<1x1xi8> + %scale = arith.constant 1.0 : f32 + %ai = tensor.empty() : tensor<1x1xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(m, k) -> (m, k)>, + affine_map<(m, k) -> ()>, + affine_map<(m, k) -> (m, k)>]} + ins(%aq, %scale : tensor<1x1xi8>, f32) + outs(%ai : tensor<1x1xf32>) -> tensor<1x1xf32> + %bi = tensor.empty() : tensor<1x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(k, n) -> (k, n)>, + affine_map<(k, n) -> ()>, + affine_map<(k, n) -> (k, n)>]} + ins(%bq, %scale : tensor<1x1xi8>, f32) + outs(%bi : tensor<1x1xf32>) -> tensor<1x1xf32> + %zero = arith.constant -0.0 : f32 + %empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32> + %result = linalg.matmul ins(%a, %b : tensor<1x1xf32>, tensor<1x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + return %result : tensor<1x1xf32> +} +// CHECK-LABEL: func.func @negative_zero_init( +// CHECK: %[[ZERO:.+]] = arith.constant 0 : i32 +// CHECK: linalg.fill ins(%[[ZERO]] : i32) +// CHECK: linalg.generic {{.*}} outs(%{{.+}} : tensor<1x1xi32>) +// CHECK: linalg.generic {{.*}} outs(%{{.+}} : tensor<1x1xf32>) +// CHECK: %[[REAL:.+]] = arith.sitofp %{{.+}} : i32 to f32 +// CHECK-NEXT: linalg.yield %[[REAL]] : f32 + +// ----- + +// FIXME: Preserve NaN propagation for infinite scales, or explicitly exclude +// non-finite scales from the supported input contract. The original lhs is +// [0 * infinity, 1 * infinity] = [NaN, infinity], so its dot product with [1, 1] +// is NaN. The rewrite computes an integer dot product of 1, then multiplies by +// infinity and returns infinity. These checks record the lost NaN propagation. +func.func @infinite_scale_nan_propagation() -> tensor<1x1xf32> { + %aq = arith.constant dense<[[0, 1]]> : tensor<1x2xi8> + %bq = arith.constant dense<1> : tensor<2x1xi8> + %sa = arith.constant 0x7F800000 : f32 + %sb = arith.constant 1.0 : f32 + %ai = tensor.empty() : tensor<1x2xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(m, k) -> (m, k)>, + affine_map<(m, k) -> ()>, + affine_map<(m, k) -> (m, k)>]} + ins(%aq, %sa : tensor<1x2xi8>, f32) + outs(%ai : tensor<1x2xf32>) -> tensor<1x2xf32> + %bi = tensor.empty() : tensor<2x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(k, n) -> (k, n)>, + affine_map<(k, n) -> ()>, + affine_map<(k, n) -> (k, n)>]} + ins(%bq, %sb : tensor<2x1xi8>, f32) + outs(%bi : tensor<2x1xf32>) -> tensor<2x1xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32> + %result = linalg.matmul ins(%a, %b : tensor<1x2xf32>, tensor<2x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + return %result : tensor<1x1xf32> +} +// CHECK-LABEL: func.func @infinite_scale_nan_propagation( +// CHECK: %[[INF:.+]] = arith.constant 0x7F800000 : f32 +// CHECK: linalg.generic {{.*}} outs(%{{.+}} : tensor<1x1xi32>) +// CHECK: %[[REAL:.+]] = arith.sitofp %{{.+}} : i32 to f32 +// CHECK-NEXT: %[[RESULT:.+]] = arith.mulf %[[REAL]], %[[INF]] : f32 +// CHECK-NEXT: linalg.yield %[[RESULT]] : f32 + +// ----- + +// This contraction returns approximately 1.07368e-37, a normal f32 value. +// Apply each 1e-23 scale to the accumulator in sequence; combining the scales +// first would underflow to zero and erase the result. +func.func @scale_product_underflow() -> tensor<1x1xf32> { + %q = arith.constant dense<32767> : tensor<1x1xi16> + %scale = arith.constant 1.0e-23 : f32 + %ai = tensor.empty() : tensor<1x1xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(m, k) -> (m, k)>, + affine_map<(m, k) -> ()>, + affine_map<(m, k) -> (m, k)>]} + ins(%q, %scale : tensor<1x1xi16>, f32) + outs(%ai : tensor<1x1xf32>) -> tensor<1x1xf32> + %bi = tensor.empty() : tensor<1x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(k, n) -> (k, n)>, + affine_map<(k, n) -> ()>, + affine_map<(k, n) -> (k, n)>]} + ins(%q, %scale : tensor<1x1xi16>, f32) + outs(%bi : tensor<1x1xf32>) -> tensor<1x1xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32> + %result = linalg.matmul ins(%a, %b : tensor<1x1xf32>, tensor<1x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + return %result : tensor<1x1xf32> +} +// CHECK-LABEL: func.func @scale_product_underflow( +// CHECK: %[[SCALE:.+]] = arith.constant 9.99999999E-24 : f32 +// CHECK: linalg.generic {{.*}} outs(%{{.+}} : tensor<1x1xi32>) +// CHECK: %[[REAL:.+]] = arith.sitofp %{{.+}} : i32 to f32 +// CHECK-NEXT: %[[PARTIAL:.+]] = arith.mulf %[[REAL]], %[[SCALE]] : f32 +// CHECK-NEXT: %[[RESULT:.+]] = arith.mulf %[[PARTIAL]], %[[SCALE]] : f32 +// CHECK-NEXT: linalg.yield %[[RESULT]] : f32 + +// ----- + +// Both symmetric operands produce a scalar zero correction in the rewrite. +// Scalar inlining exposes it to canonicalization, which folds D - 0 to D and +// removes the correction operand. Runtime scales keep the epilogue observable. +func.func @symmetric_zero_correction_folded(%aq: tensor<4x8xi8>, %a_s: f32, %bq: tensor<8x16xi8>, %b_s: f32) -> tensor<4x16xf32> { + %a_i = tensor.empty() : tensor<4x8xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%aq, %a_s : tensor<4x8xi8>, f32) + outs(%a_i : tensor<4x8xf32>) -> tensor<4x8xf32> + %b_i = tensor.empty() : tensor<8x16xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> ()>, + affine_map<(d0, d1) -> (d0, d1)>]} + ins(%bq, %b_s : tensor<8x16xi8>, f32) + outs(%b_i : tensor<8x16xf32>) -> tensor<8x16xf32> + %cst = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<4x16xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<4x16xf32>) -> tensor<4x16xf32> + %c = linalg.matmul ins(%a, %b : tensor<4x8xf32>, tensor<8x16xf32>) + outs(%f : tensor<4x16xf32>) -> tensor<4x16xf32> + return %c : tensor<4x16xf32> +} +// CHECK-LABEL: func.func @symmetric_zero_correction_folded( +// CHECK-SAME: %[[AQ:[a-zA-Z0-9_]+]]: tensor<4x8xi8>, %[[SA:[a-zA-Z0-9_]+]]: f32, +// CHECK-SAME: %[[BQ:[a-zA-Z0-9_]+]]: tensor<8x16xi8>, %[[SB:[a-zA-Z0-9_]+]]: f32 +// CHECK: %[[D:.+]] = linalg.generic {{.*}} ins(%[[AQ]], %[[BQ]] : tensor<4x8xi8>, tensor<8x16xi8>) +// CHECK-SAME: outs(%{{.+}} : tensor<4x16xi32>) +// CHECK: %[[RESULT:.+]] = linalg.generic {{.*}} ins(%[[D]] : tensor<4x16xi32>) +// CHECK-SAME: outs(%{{.+}} : tensor<4x16xf32>) +// CHECK-NEXT: ^bb0(%[[ED:[a-zA-Z0-9_]+]]: i32, %{{.+}}: f32): +// CHECK-NEXT: %[[REAL:.+]] = arith.sitofp %[[ED]] : i32 to f32 +// CHECK-NEXT: %[[PARTIAL:.+]] = arith.mulf %[[REAL]], %[[SA]] : f32 +// CHECK-NEXT: %[[SCALED:.+]] = arith.mulf %[[PARTIAL]], %[[SB]] : f32 +// CHECK-NEXT: linalg.yield %[[SCALED]] : f32 +// CHECK: return %[[RESULT]] : tensor<4x16xf32> + +// ----- + +// Cleanup also folds the zero correction for symmetric blockwise quantization +// while preserving scale widening, f32 accumulation, and final narrowing. +func.func @symmetric_half_partial_reduction(%aq: tensor<4x2x4xi8>, %a_s: tensor<4x2xf16>, + %bq: tensor<2x4x16xi8>, %b_s: f32) -> tensor<4x16xf16> { + %a_i = tensor.empty() : tensor<4x2x4xf16> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, + affine_map<(d0, d1, d2) -> (d0, d1)>, + affine_map<(d0, d1, d2) -> (d0, d1, d2)>]} + ins(%aq, %a_s : tensor<4x2x4xi8>, tensor<4x2xf16>) + outs(%a_i : tensor<4x2x4xf16>) -> tensor<4x2x4xf16> + %b_i = tensor.empty() : tensor<2x4x16xf16> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, + affine_map<(d0, d1, d2) -> ()>, + affine_map<(d0, d1, d2) -> (d0, d1, d2)>]} + ins(%bq, %b_s : tensor<2x4x16xi8>, f32) + outs(%b_i : tensor<2x4x16xf16>) -> tensor<2x4x16xf16> + %cst = arith.constant 0.000000e+00 : f16 + %e = tensor.empty() : tensor<4x16xf16> + %f = linalg.fill ins(%cst : f16) outs(%e : tensor<4x16xf16>) -> tensor<4x16xf16> + %c = linalg.contract + indexing_maps = [affine_map<(m, n, g, l) -> (m, g, l)>, + affine_map<(m, n, g, l) -> (g, l, n)>, + affine_map<(m, n, g, l) -> (m, n)>] + ins(%a, %b : tensor<4x2x4xf16>, tensor<2x4x16xf16>) + outs(%f : tensor<4x16xf16>) -> tensor<4x16xf16> + return %c : tensor<4x16xf16> +} +// CHECK-LABEL: func.func @symmetric_half_partial_reduction( +// CHECK-SAME: %{{.+}}: tensor<4x2x4xi8>, %[[SA_INPUT:[a-zA-Z0-9_]+]]: tensor<4x2xf16>, +// CHECK-SAME: %{{.+}}: tensor<2x4x16xi8>, %[[SB:[a-zA-Z0-9_]+]]: f32 +// CHECK: %[[D:.+]] = linalg.generic {{.*}} outs(%{{.+}} : tensor<4x16x2xi32>) +// CHECK: %[[PARTIALS:.+]] = linalg.generic {{.*}} ins(%[[D]], %[[SA_INPUT]] : tensor<4x16x2xi32>, tensor<4x2xf16>) +// CHECK-SAME: outs(%{{.+}} : tensor<4x16xf32>) +// CHECK-NEXT: ^bb0(%[[ED:[a-zA-Z0-9_]+]]: i32, %[[SA:[a-zA-Z0-9_]+]]: f16, %[[ACC:[a-zA-Z0-9_]+]]: f32): +// CHECK-NEXT: %[[REAL:.+]] = arith.sitofp %[[ED]] : i32 to f32 +// CHECK-NEXT: %[[WIDE_SA:.+]] = arith.extf %[[SA]] : f16 to f32 +// CHECK-NEXT: %[[PARTIAL:.+]] = arith.mulf %[[REAL]], %[[WIDE_SA]] : f32 +// CHECK-NEXT: %[[SCALED:.+]] = arith.mulf %[[PARTIAL]], %[[SB]] : f32 +// CHECK-NEXT: %[[TOTAL:.+]] = arith.addf %[[ACC]], %[[SCALED]] : f32 +// CHECK-NEXT: linalg.yield %[[TOTAL]] : f32 +// CHECK: linalg.generic {{.*}} ins(%[[PARTIALS]] : tensor<4x16xf32>) +// CHECK-SAME: outs(%{{.+}} : tensor<4x16xf16>) +// CHECK-NEXT: ^bb0(%[[V:[a-zA-Z0-9_]+]]: f32, %{{.+}}: f16): +// CHECK-NEXT: %[[NARROW:.+]] = arith.truncf %[[V]] : f32 to f16 +// CHECK-NEXT: linalg.yield %[[NARROW]] : f16
diff --git a/compiler/src/iree/compiler/GlobalOptimization/unittests/BUILD.bazel b/compiler/src/iree/compiler/GlobalOptimization/unittests/BUILD.bazel new file mode 100644 index 0000000..f9ac5c3 --- /dev/null +++ b/compiler/src/iree/compiler/GlobalOptimization/unittests/BUILD.bazel
@@ -0,0 +1,41 @@ +# 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 + +load("//build_tools/bazel:build_defs.oss.bzl", "iree_compiler_cc_test") + +package( + features = ["layering_check"], + licenses = ["notice"], # Apache 2.0 +) + +iree_compiler_cc_test( + name = "QuantizedContractionTest", + testonly = True, + srcs = ["QuantizedContractionTest.cpp"], + deps = [ + "//compiler/src/iree/compiler/Dialect/LinalgExt/IR", + "//compiler/src/iree/compiler/GlobalOptimization:QuantizedContraction", + "//compiler/src/iree/testing:gtest_main", + "@com_google_googletest//:gtest", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:FuncDialect", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:LinalgDialect", + "@llvm-project//mlir:Parser", + ], +) + +iree_compiler_cc_test( + name = "QuantizationUtilsTest", + testonly = True, + srcs = ["QuantizationUtilsTest.cpp"], + deps = [ + "//compiler/src/iree/compiler/GlobalOptimization:QuantizationUtils", + "//compiler/src/iree/testing:gtest_main", + "@com_google_googletest//:gtest", + ], +)
diff --git a/compiler/src/iree/compiler/GlobalOptimization/unittests/CMakeLists.txt b/compiler/src/iree/compiler/GlobalOptimization/unittests/CMakeLists.txt new file mode 100644 index 0000000..fbb3ebc --- /dev/null +++ b/compiler/src/iree/compiler/GlobalOptimization/unittests/CMakeLists.txt
@@ -0,0 +1,44 @@ +################################################################################ +# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from # +# compiler/src/iree/compiler/GlobalOptimization/unittests/BUILD.bazel # +# # +# Use iree_cmake_extra_content from iree/build_defs.oss.bzl to add arbitrary # +# CMake-only content. # +# # +# To disable autogeneration for this file entirely, delete this header. # +################################################################################ + +iree_add_all_subdirs() + +iree_cc_test( + NAME + QuantizedContractionTest + SRCS + "QuantizedContractionTest.cpp" + DEPS + LLVMSupport + MLIRArithDialect + MLIRFuncDialect + MLIRIR + MLIRLinalgDialect + MLIRParser + gmock + gtest + iree::compiler::Dialect::LinalgExt::IR + iree::compiler::GlobalOptimization::QuantizedContraction + iree::testing::gtest_main +) + +iree_cc_test( + NAME + QuantizationUtilsTest + SRCS + "QuantizationUtilsTest.cpp" + DEPS + gmock + gtest + iree::compiler::GlobalOptimization::QuantizationUtils + iree::testing::gtest_main +) + +### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/compiler/src/iree/compiler/GlobalOptimization/unittests/QuantizationUtilsTest.cpp b/compiler/src/iree/compiler/GlobalOptimization/unittests/QuantizationUtilsTest.cpp new file mode 100644 index 0000000..e4a8877 --- /dev/null +++ b/compiler/src/iree/compiler/GlobalOptimization/unittests/QuantizationUtilsTest.cpp
@@ -0,0 +1,213 @@ +// 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/GlobalOptimization/QuantizationUtils.h" + +#include <limits> + +#include <gtest/gtest.h> + +namespace mlir::iree_compiler::GlobalOptimization { +namespace { + +TEST(QuantizationUtilsTest, getStorageMagnitude) { + EXPECT_EQ(getStorageMagnitude(3, true), 7); + EXPECT_EQ(getStorageMagnitude(3, false), 4); + EXPECT_EQ(getStorageMagnitude(8, true), 255); + EXPECT_EQ(getStorageMagnitude(8, false), 128); + + // i31 is the widest supported storage type; i32 is rejected. + EXPECT_EQ(getStorageMagnitude(31, true), 2147483647LL); + EXPECT_EQ(getStorageMagnitude(31, false), 1073741824LL); + EXPECT_EQ(getStorageMagnitude(32, true), 0); + EXPECT_EQ(getStorageMagnitude(32, false), 0); + + // Test bit width greater than accumulator width. + EXPECT_EQ(getStorageMagnitude(33, true), 0); + EXPECT_EQ(getStorageMagnitude(33, false), 0); + + // Test zero bit width. + EXPECT_EQ(getStorageMagnitude(0, true), 0); + EXPECT_EQ(getStorageMagnitude(0, false), 0); +} + +TEST(QuantizationUtilsTest, OperandRangesStorageGrid) { + struct TestCase { + unsigned bitWidth; + bool isUnsigned; + QuantMinMax storage; + }; + for (const auto &test : + {TestCase{3, true, {0, 7}}, TestCase{3, false, {-4, 3}}, + TestCase{8, true, {0, 255}}, TestCase{8, false, {-128, 127}}, + TestCase{31, true, {0, 2147483647LL}}, + TestCase{31, false, {-1073741824LL, 1073741823LL}}}) { + SCOPED_TRACE(test.bitWidth); + SCOPED_TRACE(test.isUnsigned); + for (bool isSymmetric : {false, true}) { + auto ranges = getQuantizedOperandRanges(test.bitWidth, test.isUnsigned, + std::nullopt, isSymmetric); + ASSERT_TRUE(ranges); + EXPECT_EQ(ranges->input.min, test.storage.min); + EXPECT_EQ(ranges->input.max, test.storage.max); + EXPECT_EQ(ranges->zeroPoint.min, isSymmetric ? 0 : test.storage.min); + EXPECT_EQ(ranges->zeroPoint.max, isSymmetric ? 0 : test.storage.max); + } + } +} + +TEST(QuantizationUtilsTest, OperandRangesRestrictedInput) { + // Restricting inputs does not restrict zero points to the same range. + for (bool isUnsigned : {false, true}) { + for (bool isSymmetric : {false, true}) { + auto ranges = getQuantizedOperandRanges(8, isUnsigned, QuantMinMax{0, 7}, + isSymmetric); + ASSERT_TRUE(ranges); + EXPECT_EQ(ranges->input.min, 0); + EXPECT_EQ(ranges->input.max, 7); + EXPECT_EQ(ranges->zeroPoint.min, isSymmetric || isUnsigned ? 0 : -128); + EXPECT_EQ(ranges->zeroPoint.max, isSymmetric ? 0 + : isUnsigned ? 255 + : 127); + } + } + auto negative = + getQuantizedOperandRanges(8, false, QuantMinMax{-8, -1}, false); + ASSERT_TRUE(negative); + EXPECT_EQ(negative->input.min, -8); + EXPECT_EQ(negative->input.max, -1); + EXPECT_EQ(negative->zeroPoint.min, -128); + EXPECT_EQ(negative->zeroPoint.max, 127); +} + +TEST(QuantizationUtilsTest, OperandRangesUnsupportedWidth) { + for (unsigned bitWidth : {0, 32, 33, 64}) { + for (bool isUnsigned : {false, true}) { + for (bool isSymmetric : {false, true}) { + EXPECT_FALSE(getQuantizedOperandRanges(bitWidth, isUnsigned, + std::nullopt, isSymmetric)); + // A restricted range does not make unsupported storage legal. + EXPECT_FALSE(getQuantizedOperandRanges(bitWidth, isUnsigned, + QuantMinMax{0, 7}, isSymmetric)); + } + } + } +} + +TEST(QuantizationUtilsTest, MaxReductionExtentSymmetric) { + QuantizedOperandRanges signedI8{{-128, 127}, {0, 0}}; + QuantizedOperandRanges unsignedI8{{0, 255}, {0, 0}}; + // The raw product limits the depth, including -128 * -128 = +16384. + EXPECT_EQ(getMaxReductionExtent(signedI8, signedI8), 131071); + EXPECT_EQ(getMaxReductionExtent(unsignedI8, unsignedI8), 33025); + EXPECT_EQ(getMaxReductionExtent(signedI8, unsignedI8), 65793); + EXPECT_EQ(getMaxReductionExtent(unsignedI8, signedI8), 65793); + // A single symmetric i16 product fits; two do not. + QuantizedOperandRanges signedI16{{-32768, 32767}, {0, 0}}; + EXPECT_EQ(getMaxReductionExtent(signedI16, signedI16), 1); +} + +TEST(QuantizationUtilsTest, MaxReductionExtentAsymmetric) { + QuantizedOperandRanges signedI8{{-128, 127}, {-128, 127}}; + QuantizedOperandRanges unsignedI8{{0, 255}, {0, 255}}; + // The centered signed product is bounded by 255^2, not 256^2. + EXPECT_EQ(getMaxReductionExtent(signedI8, signedI8), 33025); + // P+Q limits unsigned depth to INT32_MAX/(2*255^2), even though the + // centered result is bounded by N*255^2. + EXPECT_EQ(getMaxReductionExtent(unsignedI8, unsignedI8), 16512); + // For mixed signs, P+Q-R has range [-97665, 97410] per element. It is + // larger than both P+Q and the centered product, so it limits the depth. + EXPECT_EQ(getMaxReductionExtent(signedI8, unsignedI8), 21988); + EXPECT_EQ(getMaxReductionExtent(unsignedI8, signedI8), 21988); +} + +TEST(QuantizationUtilsTest, MaxReductionExtentOneZeroPoint) { + QuantizedOperandRanges signedSymmetric{{-128, 127}, {0, 0}}; + QuantizedOperandRanges signedAsymmetric{{-128, 127}, {-128, 127}}; + QuantizedOperandRanges unsignedSymmetric{{0, 255}, {0, 0}}; + QuantizedOperandRanges unsignedAsymmetric{{0, 255}, {0, 255}}; + EXPECT_EQ(getMaxReductionExtent(signedSymmetric, signedAsymmetric), 65793); + EXPECT_EQ(getMaxReductionExtent(signedAsymmetric, signedSymmetric), 65793); + EXPECT_EQ(getMaxReductionExtent(unsignedSymmetric, unsignedAsymmetric), + 33025); + EXPECT_EQ(getMaxReductionExtent(unsignedAsymmetric, unsignedSymmetric), + 33025); +} + +TEST(QuantizationUtilsTest, MaxReductionExtentRestrictedInput) { + // Inputs are only three bits, but the zero points still occupy four bits. + // The cross term and centered result reach 225, exceeding P+Q's 210. + QuantizedOperandRanges restricted{{0, 7}, {0, 15}}; + EXPECT_EQ(getMaxReductionExtent(restricted, restricted), 9544371); + QuantizedOperandRanges negative{{-8, -1}, {0, 0}}; + EXPECT_EQ(getMaxReductionExtent(negative, negative), 33554431); +} + +TEST(QuantizationUtilsTest, MaxReductionExtentZeroInputs) { + QuantizedOperandRanges zero{{0, 0}, {0, 0}}; + QuantizedOperandRanges signedI8{{-128, 127}, {0, 0}}; + // Even when products vanish, the separately computed operand sums must fit. + EXPECT_EQ(getMaxReductionExtent(signedI8, zero), 16777215); + EXPECT_EQ(getMaxReductionExtent(zero, signedI8), 16777215); + EXPECT_EQ(getMaxReductionExtent(zero, zero), 0); + // Nonzero zero points still produce a cross term when both inputs are zero. + QuantizedOperandRanges zeroInput{{0, 0}, {0, 255}}; + EXPECT_EQ(getMaxReductionExtent(zeroInput, zeroInput), 33025); +} + +TEST(QuantizationUtilsTest, MaxReductionExtentI31) { + QuantizedOperandRanges one{{1, 1}, {0, 0}}; + EXPECT_EQ(getMaxReductionExtent(one, one), 2147483647LL); + for (bool isUnsigned : {false, true}) { + for (bool isSymmetric : {false, true}) { + auto ranges = + getQuantizedOperandRanges(31, isUnsigned, std::nullopt, isSymmetric); + ASSERT_TRUE(ranges); + EXPECT_EQ(getMaxReductionExtent(*ranges, one), 1); + EXPECT_EQ(getMaxReductionExtent(one, *ranges), 1); + // Even N=1 is unsafe. For unsigned asymmetric i31, P+Q approaches + // INT64_MAX; computing its bound must not overflow in the host. + EXPECT_EQ(getMaxReductionExtent(*ranges, *ranges), 0); + } + } + auto signedRanges = getQuantizedOperandRanges(31, false, std::nullopt, false); + auto unsignedRanges = + getQuantizedOperandRanges(31, true, std::nullopt, false); + ASSERT_TRUE(signedRanges); + ASSERT_TRUE(unsignedRanges); + EXPECT_EQ(getMaxReductionExtent(*signedRanges, *unsignedRanges), 0); + EXPECT_EQ(getMaxReductionExtent(*unsignedRanges, *signedRanges), 0); +} + +TEST(QuantizationUtilsTest, BoundedReductionExtentAtLimit) { + EXPECT_EQ(getBoundedReductionExtent({131071}, 131071), 131071); + EXPECT_EQ(getBoundedReductionExtent({131072}, 131071), std::nullopt); + EXPECT_EQ(getBoundedReductionExtent({3, 5, 7}, 105), 105); + EXPECT_EQ(getBoundedReductionExtent({3, 5, 7}, 104), std::nullopt); + EXPECT_EQ(getBoundedReductionExtent({7, 3, 5}, 105), 105); + EXPECT_EQ(getBoundedReductionExtent({1, 1}, 1), 1); +} + +TEST(QuantizationUtilsTest, + BoundedReductionExtentRequiresPositiveStaticExtents) { + EXPECT_EQ(getBoundedReductionExtent({}, 100), std::nullopt); + EXPECT_EQ(getBoundedReductionExtent({0}, 100), std::nullopt); + EXPECT_EQ(getBoundedReductionExtent({2, -1}, 100), std::nullopt); + EXPECT_EQ(getBoundedReductionExtent({1}, 0), std::nullopt); + EXPECT_EQ(getBoundedReductionExtent({1}, -1), std::nullopt); +} + +TEST(QuantizationUtilsTest, BoundedReductionExtentAvoidsOverflow) { + constexpr int64_t max = std::numeric_limits<int64_t>::max(); + EXPECT_EQ(getBoundedReductionExtent({max}, max), max); + EXPECT_EQ(getBoundedReductionExtent({max, 2}, max), std::nullopt); + EXPECT_EQ(getBoundedReductionExtent({2, max}, max), std::nullopt); + EXPECT_EQ(getBoundedReductionExtent({4294967296LL, 4294967296LL}, max), + std::nullopt); +} + +} // namespace +} // namespace mlir::iree_compiler::GlobalOptimization
diff --git a/compiler/src/iree/compiler/GlobalOptimization/unittests/QuantizedContractionTest.cpp b/compiler/src/iree/compiler/GlobalOptimization/unittests/QuantizedContractionTest.cpp new file mode 100644 index 0000000..838c972 --- /dev/null +++ b/compiler/src/iree/compiler/GlobalOptimization/unittests/QuantizedContractionTest.cpp
@@ -0,0 +1,914 @@ +// 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/GlobalOptimization/QuantizedContraction.h" + +#include <gmock/gmock.h> +#include <gtest/gtest.h> + +#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtDialect.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/Verifier.h" +#include "mlir/Parser/Parser.h" + +namespace mlir::iree_compiler::GlobalOptimization::detail { +namespace { + +using ::testing::ElementsAre; +using ::testing::IsEmpty; + +class QuantizedContractionTest : public ::testing::Test { +protected: + QuantizedContractionTest() { + context + .loadDialect<IREE::LinalgExt::IREELinalgExtDialect, arith::ArithDialect, + func::FuncDialect, linalg::LinalgDialect>(); + } + + bool parse(StringRef source) { + module = parseSourceString<ModuleOp>(source, &context); + return module && succeeded(verify(*module)); + } + + linalg::GenericOp contraction(StringRef id = "test") { + linalg::GenericOp result; + module->walk([&](linalg::GenericOp op) { + if (auto attr = op->getAttrOfType<StringAttr>("id")) { + if (attr.getValue() == id) { + result = op; + } + } + }); + EXPECT_TRUE(result) << "Missing contraction: " << id.str(); + return result; + } + + AffineMap map(unsigned rank, ArrayRef<unsigned> dims) { + SmallVector<AffineExpr> exprs; + for (unsigned dim : dims) { + exprs.push_back(getAffineDimExpr(dim, &context)); + } + return AffineMap::get(rank, 0, exprs, &context); + } + + MLIRContext context; + OwningOpRef<ModuleOp> module; +}; + +TEST_F(QuantizedContractionTest, SymmetricMultiDimensionalReduction) { + ASSERT_TRUE(parse(R"mlir( + !a = tensor<2x4x8xf32> + !b = tensor<4x8x3xf32> + !c = tensor<2x3xf32> + #identity = affine_map<(d0, d1, d2) -> (d0, d1, d2)> + #scalar = affine_map<(d0, d1, d2) -> ()> + #lhs = affine_map<(m, n, g, k) -> (m, g, k)> + #rhs = affine_map<(m, n, g, k) -> (g, k, n)> + #out = affine_map<(m, n, g, k) -> (m, n)> + func.func @test( + %aq: tensor<2x4x8xi8>, + %as: f32, + %ainit: !a, + %bq: tensor<4x8x3xi8>, + %bs: f32, + %binit: !b, + %cinit: !c) { + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #identity]} + ins(%aq, %as : tensor<2x4x8xi8>, f32) + outs(%ainit : !a) -> !a + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #identity]} + ins(%bq, %bs : tensor<4x8x3xi8>, f32) + outs(%binit : !b) -> !b + %zero = arith.constant 0.0 : f32 + %init = linalg.fill ins(%zero : f32) outs(%cinit : !c) -> !c + %test = linalg.generic { + id = "test", indexing_maps = [#lhs, #rhs, #out], + iterator_types = ["parallel", "parallel", "reduction", "reduction"] + } ins(%a, %b : !a, !b) outs(%init : !c) { + ^bb0(%x: f32, %y: f32, %acc: f32): + %p = arith.mulf %x, %y : f32 + %sum = arith.addf %p, %acc : f32 + linalg.yield %sum : f32 + } -> !c + return + } + )mlir")); + auto detail = getQuantizedContraction(contraction()); + ASSERT_TRUE(succeeded(detail)); + EXPECT_THAT(detail->integerReductionDims, ElementsAre(2, 3)); + EXPECT_THAT(detail->floatingReductionDims, IsEmpty()); + EXPECT_EQ(detail->reductionExtent, 32); + EXPECT_EQ(detail->lhs.inputMap, map(4, {0, 2, 3})); + EXPECT_EQ(detail->rhs.inputMap, map(4, {2, 3, 1})); + EXPECT_EQ(detail->lhs.scaleMap, map(4, {})); + EXPECT_EQ(detail->outputMap, map(4, {0, 1})); + EXPECT_FALSE(detail->lhs.zeroPointMap); + EXPECT_FALSE(detail->rhs.zeroPointMap); + EXPECT_FALSE(detail->needsLhsSum()); + EXPECT_FALSE(detail->needsRhsSum()); +} + +TEST_F(QuantizedContractionTest, LhsZeroPointNeedsRhsSum) { + ASSERT_TRUE(parse(R"mlir( + !a = tensor<2x4x8xf32> + !b = tensor<4x8x3xf32> + !c = tensor<2x3xf32> + #identity = affine_map<(d0, d1, d2) -> (d0, d1, d2)> + #scalar = affine_map<(d0, d1, d2) -> ()> + #lhs = affine_map<(m, n, g, k) -> (m, g, k)> + #rhs = affine_map<(m, n, g, k) -> (g, k, n)> + #out = affine_map<(m, n, g, k) -> (m, n)> + func.func @test( + %aq: tensor<2x4x8xi8>, + %as: f32, + %az: i8, + %ainit: !a, + %bq: tensor<4x8x3xi8>, + %bs: f32, + %binit: !b, + %cinit: !c) { + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #scalar, #identity]} + ins(%aq, %as, %az : tensor<2x4x8xi8>, f32, i8) + outs(%ainit : !a) -> !a + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #identity]} + ins(%bq, %bs : tensor<4x8x3xi8>, f32) + outs(%binit : !b) -> !b + %zero = arith.constant 0.0 : f32 + %init = linalg.fill ins(%zero : f32) outs(%cinit : !c) -> !c + %test = linalg.generic { + id = "test", indexing_maps = [#lhs, #rhs, #out], + iterator_types = ["parallel", "parallel", "reduction", "reduction"] + } ins(%a, %b : !a, !b) outs(%init : !c) { + ^bb0(%x: f32, %y: f32, %acc: f32): + %p = arith.mulf %x, %y : f32 + %sum = arith.addf %p, %acc : f32 + linalg.yield %sum : f32 + } -> !c + return + } + )mlir")); + auto detail = getQuantizedContraction(contraction()); + ASSERT_TRUE(succeeded(detail)); + EXPECT_EQ(detail->needsLhsSum(), false); + EXPECT_EQ(detail->needsRhsSum(), true); + EXPECT_EQ(detail->lhs.zeroPointMap, map(4, {})); + EXPECT_EQ(detail->rhs.zeroPointMap, AffineMap()); +} + +TEST_F(QuantizedContractionTest, RhsZeroPointNeedsLhsSum) { + ASSERT_TRUE(parse(R"mlir( + !a = tensor<2x4x8xf32> + !b = tensor<4x8x3xf32> + !c = tensor<2x3xf32> + #identity = affine_map<(d0, d1, d2) -> (d0, d1, d2)> + #scalar = affine_map<(d0, d1, d2) -> ()> + #lhs = affine_map<(m, n, g, k) -> (m, g, k)> + #rhs = affine_map<(m, n, g, k) -> (g, k, n)> + #out = affine_map<(m, n, g, k) -> (m, n)> + func.func @test( + %aq: tensor<2x4x8xi8>, + %as: f32, + %ainit: !a, + %bq: tensor<4x8x3xi8>, + %bs: f32, + %bz: i8, + %binit: !b, + %cinit: !c) { + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #identity]} + ins(%aq, %as : tensor<2x4x8xi8>, f32) + outs(%ainit : !a) -> !a + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #scalar, #identity]} + ins(%bq, %bs, %bz : tensor<4x8x3xi8>, f32, i8) + outs(%binit : !b) -> !b + %zero = arith.constant 0.0 : f32 + %init = linalg.fill ins(%zero : f32) outs(%cinit : !c) -> !c + %test = linalg.generic { + id = "test", indexing_maps = [#lhs, #rhs, #out], + iterator_types = ["parallel", "parallel", "reduction", "reduction"] + } ins(%a, %b : !a, !b) outs(%init : !c) { + ^bb0(%x: f32, %y: f32, %acc: f32): + %p = arith.mulf %x, %y : f32 + %sum = arith.addf %p, %acc : f32 + linalg.yield %sum : f32 + } -> !c + return + } + )mlir")); + auto detail = getQuantizedContraction(contraction()); + ASSERT_TRUE(succeeded(detail)); + EXPECT_EQ(detail->needsLhsSum(), true); + EXPECT_EQ(detail->needsRhsSum(), false); + EXPECT_EQ(detail->lhs.zeroPointMap, AffineMap()); + EXPECT_EQ(detail->rhs.zeroPointMap, map(4, {})); +} + +TEST_F(QuantizedContractionTest, BothZeroPointsNeedBothSums) { + ASSERT_TRUE(parse(R"mlir( + !a = tensor<2x4x8xf32> + !b = tensor<4x8x3xf32> + !c = tensor<2x3xf32> + #identity = affine_map<(d0, d1, d2) -> (d0, d1, d2)> + #scalar = affine_map<(d0, d1, d2) -> ()> + #lhs = affine_map<(m, n, g, k) -> (m, g, k)> + #rhs = affine_map<(m, n, g, k) -> (g, k, n)> + #out = affine_map<(m, n, g, k) -> (m, n)> + func.func @test( + %aq: tensor<2x4x8xi8>, + %as: f32, + %az: i8, + %ainit: !a, + %bq: tensor<4x8x3xi8>, + %bs: f32, + %bz: i8, + %binit: !b, + %cinit: !c) { + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #scalar, #identity]} + ins(%aq, %as, %az : tensor<2x4x8xi8>, f32, i8) + outs(%ainit : !a) -> !a + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #scalar, #identity]} + ins(%bq, %bs, %bz : tensor<4x8x3xi8>, f32, i8) + outs(%binit : !b) -> !b + %zero = arith.constant 0.0 : f32 + %init = linalg.fill ins(%zero : f32) outs(%cinit : !c) -> !c + %test = linalg.generic { + id = "test", indexing_maps = [#lhs, #rhs, #out], + iterator_types = ["parallel", "parallel", "reduction", "reduction"] + } ins(%a, %b : !a, !b) outs(%init : !c) { + ^bb0(%x: f32, %y: f32, %acc: f32): + %p = arith.mulf %x, %y : f32 + %sum = arith.addf %p, %acc : f32 + linalg.yield %sum : f32 + } -> !c + return + } + )mlir")); + auto detail = getQuantizedContraction(contraction()); + ASSERT_TRUE(succeeded(detail)); + EXPECT_EQ(detail->needsLhsSum(), true); + EXPECT_EQ(detail->needsRhsSum(), true); + EXPECT_EQ(detail->lhs.zeroPointMap, map(4, {})); + EXPECT_EQ(detail->rhs.zeroPointMap, map(4, {})); +} + +TEST_F(QuantizedContractionTest, LhsScaleRetainsFloatingReduction) { + ASSERT_TRUE(parse(R"mlir( + !a = tensor<2x4x8xf32> + !b = tensor<4x8x3xf32> + !c = tensor<2x3xf32> + #identity = affine_map<(d0, d1, d2) -> (d0, d1, d2)> + #scalar = affine_map<(d0, d1, d2) -> ()> + #lhs = affine_map<(m, n, g, k) -> (m, g, k)> + #rhs = affine_map<(m, n, g, k) -> (g, k, n)> + #out = affine_map<(m, n, g, k) -> (m, n)> + func.func @test( + %aq: tensor<2x4x8xi8>, + %as: tensor<4xf32>, + %ainit: !a, + %bq: tensor<4x8x3xi8>, + %bs: f32, + %binit: !b, + %cinit: !c) { + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, affine_map<(d0, d1, d2) -> (d1)>, #identity]} + ins(%aq, %as : tensor<2x4x8xi8>, tensor<4xf32>) + outs(%ainit : !a) -> !a + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #identity]} + ins(%bq, %bs : tensor<4x8x3xi8>, f32) + outs(%binit : !b) -> !b + %zero = arith.constant 0.0 : f32 + %init = linalg.fill ins(%zero : f32) outs(%cinit : !c) -> !c + %test = linalg.generic { + id = "test", indexing_maps = [#lhs, #rhs, #out], + iterator_types = ["parallel", "parallel", "reduction", "reduction"] + } ins(%a, %b : !a, !b) outs(%init : !c) { + ^bb0(%x: f32, %y: f32, %acc: f32): + %p = arith.mulf %x, %y : f32 + %sum = arith.addf %p, %acc : f32 + linalg.yield %sum : f32 + } -> !c + return + } + )mlir")); + auto detail = getQuantizedContraction(contraction()); + ASSERT_TRUE(succeeded(detail)); + EXPECT_THAT(detail->integerReductionDims, ElementsAre(3)); + EXPECT_THAT(detail->floatingReductionDims, ElementsAre(2)); + EXPECT_EQ(detail->reductionExtent, 8); +} + +TEST_F(QuantizedContractionTest, RhsScaleRetainsFloatingReduction) { + ASSERT_TRUE(parse(R"mlir( + !a = tensor<2x4x8xf32> + !b = tensor<4x8x3xf32> + !c = tensor<2x3xf32> + #identity = affine_map<(d0, d1, d2) -> (d0, d1, d2)> + #scalar = affine_map<(d0, d1, d2) -> ()> + #lhs = affine_map<(m, n, g, k) -> (m, g, k)> + #rhs = affine_map<(m, n, g, k) -> (g, k, n)> + #out = affine_map<(m, n, g, k) -> (m, n)> + func.func @test( + %aq: tensor<2x4x8xi8>, + %as: f32, + %ainit: !a, + %bq: tensor<4x8x3xi8>, + %bs: tensor<4xf32>, + %binit: !b, + %cinit: !c) { + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #identity]} + ins(%aq, %as : tensor<2x4x8xi8>, f32) + outs(%ainit : !a) -> !a + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, affine_map<(d0, d1, d2) -> (d0)>, #identity]} + ins(%bq, %bs : tensor<4x8x3xi8>, tensor<4xf32>) + outs(%binit : !b) -> !b + %zero = arith.constant 0.0 : f32 + %init = linalg.fill ins(%zero : f32) outs(%cinit : !c) -> !c + %test = linalg.generic { + id = "test", indexing_maps = [#lhs, #rhs, #out], + iterator_types = ["parallel", "parallel", "reduction", "reduction"] + } ins(%a, %b : !a, !b) outs(%init : !c) { + ^bb0(%x: f32, %y: f32, %acc: f32): + %p = arith.mulf %x, %y : f32 + %sum = arith.addf %p, %acc : f32 + linalg.yield %sum : f32 + } -> !c + return + } + )mlir")); + auto detail = getQuantizedContraction(contraction()); + ASSERT_TRUE(succeeded(detail)); + EXPECT_THAT(detail->integerReductionDims, ElementsAre(3)); + EXPECT_THAT(detail->floatingReductionDims, ElementsAre(2)); + EXPECT_EQ(detail->reductionExtent, 8); +} + +TEST_F(QuantizedContractionTest, LhsZeroPointRetainsFloatingReduction) { + ASSERT_TRUE(parse(R"mlir( + !a = tensor<2x4x8xf32> + !b = tensor<4x8x3xf32> + !c = tensor<2x3xf32> + #identity = affine_map<(d0, d1, d2) -> (d0, d1, d2)> + #scalar = affine_map<(d0, d1, d2) -> ()> + #lhs = affine_map<(m, n, g, k) -> (m, g, k)> + #rhs = affine_map<(m, n, g, k) -> (g, k, n)> + #out = affine_map<(m, n, g, k) -> (m, n)> + func.func @test( + %aq: tensor<2x4x8xi8>, + %as: f32, + %az: tensor<4xi8>, + %ainit: !a, + %bq: tensor<4x8x3xi8>, + %bs: f32, + %binit: !b, + %cinit: !c) { + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, affine_map<(d0, d1, d2) -> (d1)>, #identity]} + ins(%aq, %as, %az : tensor<2x4x8xi8>, f32, tensor<4xi8>) + outs(%ainit : !a) -> !a + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #identity]} + ins(%bq, %bs : tensor<4x8x3xi8>, f32) + outs(%binit : !b) -> !b + %zero = arith.constant 0.0 : f32 + %init = linalg.fill ins(%zero : f32) outs(%cinit : !c) -> !c + %test = linalg.generic { + id = "test", indexing_maps = [#lhs, #rhs, #out], + iterator_types = ["parallel", "parallel", "reduction", "reduction"] + } ins(%a, %b : !a, !b) outs(%init : !c) { + ^bb0(%x: f32, %y: f32, %acc: f32): + %p = arith.mulf %x, %y : f32 + %sum = arith.addf %p, %acc : f32 + linalg.yield %sum : f32 + } -> !c + return + } + )mlir")); + auto detail = getQuantizedContraction(contraction()); + ASSERT_TRUE(succeeded(detail)); + EXPECT_THAT(detail->integerReductionDims, ElementsAre(3)); + EXPECT_THAT(detail->floatingReductionDims, ElementsAre(2)); + EXPECT_EQ(detail->reductionExtent, 8); +} + +TEST_F(QuantizedContractionTest, RhsZeroPointRetainsFloatingReduction) { + ASSERT_TRUE(parse(R"mlir( + !a = tensor<2x4x8xf32> + !b = tensor<4x8x3xf32> + !c = tensor<2x3xf32> + #identity = affine_map<(d0, d1, d2) -> (d0, d1, d2)> + #scalar = affine_map<(d0, d1, d2) -> ()> + #lhs = affine_map<(m, n, g, k) -> (m, g, k)> + #rhs = affine_map<(m, n, g, k) -> (g, k, n)> + #out = affine_map<(m, n, g, k) -> (m, n)> + func.func @test( + %aq: tensor<2x4x8xi8>, + %as: f32, + %ainit: !a, + %bq: tensor<4x8x3xi8>, + %bs: f32, + %bz: tensor<4xi8>, + %binit: !b, + %cinit: !c) { + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #identity]} + ins(%aq, %as : tensor<2x4x8xi8>, f32) + outs(%ainit : !a) -> !a + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, affine_map<(d0, d1, d2) -> (d0)>, #identity]} + ins(%bq, %bs, %bz : tensor<4x8x3xi8>, f32, tensor<4xi8>) + outs(%binit : !b) -> !b + %zero = arith.constant 0.0 : f32 + %init = linalg.fill ins(%zero : f32) outs(%cinit : !c) -> !c + %test = linalg.generic { + id = "test", indexing_maps = [#lhs, #rhs, #out], + iterator_types = ["parallel", "parallel", "reduction", "reduction"] + } ins(%a, %b : !a, !b) outs(%init : !c) { + ^bb0(%x: f32, %y: f32, %acc: f32): + %p = arith.mulf %x, %y : f32 + %sum = arith.addf %p, %acc : f32 + linalg.yield %sum : f32 + } -> !c + return + } + )mlir")); + auto detail = getQuantizedContraction(contraction()); + ASSERT_TRUE(succeeded(detail)); + EXPECT_THAT(detail->integerReductionDims, ElementsAre(3)); + EXPECT_THAT(detail->floatingReductionDims, ElementsAre(2)); + EXPECT_EQ(detail->reductionExtent, 8); +} + +TEST_F(QuantizedContractionTest, RebasesTransposedDequantizationMaps) { + ASSERT_TRUE(parse(R"mlir( + !a = tensor<2x4x8xf32> + !b = tensor<4x8x3xf32> + !c = tensor<2x3xf32> + #identity = affine_map<(d0, d1, d2) -> (d0, d1, d2)> + #scalar = affine_map<(d0, d1, d2) -> ()> + #lhs = affine_map<(m, n, g, k) -> (m, g, k)> + #rhs = affine_map<(m, n, g, k) -> (g, k, n)> + #out = affine_map<(m, n, g, k) -> (m, n)> + func.func @test( + %aq: tensor<4x2x8xi8>, + %as: tensor<4xf32>, + %az: tensor<4x2xi8>, + %ainit: !a, + %bq: tensor<4x8x3xi8>, + %bs: f32, + %binit: !b, + %cinit: !c) { + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, affine_map<(d0, d1, d2) -> (d0)>, affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d1, d0, d2)>]} + ins(%aq, %as, %az : tensor<4x2x8xi8>, tensor<4xf32>, tensor<4x2xi8>) + outs(%ainit : !a) -> !a + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #identity]} + ins(%bq, %bs : tensor<4x8x3xi8>, f32) + outs(%binit : !b) -> !b + %zero = arith.constant 0.0 : f32 + %init = linalg.fill ins(%zero : f32) outs(%cinit : !c) -> !c + %test = linalg.generic { + id = "test", indexing_maps = [#lhs, #rhs, #out], + iterator_types = ["parallel", "parallel", "reduction", "reduction"] + } ins(%a, %b : !a, !b) outs(%init : !c) { + ^bb0(%x: f32, %y: f32, %acc: f32): + %p = arith.mulf %x, %y : f32 + %sum = arith.addf %p, %acc : f32 + linalg.yield %sum : f32 + } -> !c + return + } + )mlir")); + auto detail = getQuantizedContraction(contraction()); + ASSERT_TRUE(succeeded(detail)); + EXPECT_EQ(detail->lhs.inputMap, map(4, {2, 0, 3})); + EXPECT_EQ(detail->lhs.scaleMap, map(4, {2})); + EXPECT_EQ(detail->lhs.zeroPointMap, map(4, {2, 0})); + EXPECT_THAT(detail->integerReductionDims, ElementsAre(3)); + EXPECT_THAT(detail->floatingReductionDims, ElementsAre(2)); +} + +TEST_F(QuantizedContractionTest, RequiresAnIntegerReduction) { + ASSERT_TRUE(parse(R"mlir( + !a = tensor<2x4x8xf32> + !b = tensor<4x8x3xf32> + !c = tensor<2x3xf32> + #identity = affine_map<(d0, d1, d2) -> (d0, d1, d2)> + #scalar = affine_map<(d0, d1, d2) -> ()> + #lhs = affine_map<(m, n, g, k) -> (m, g, k)> + #rhs = affine_map<(m, n, g, k) -> (g, k, n)> + #out = affine_map<(m, n, g, k) -> (m, n)> + func.func @test( + %aq: tensor<2x4x8xi8>, + %as: tensor<4x8xf32>, + %ainit: !a, + %bq: tensor<4x8x3xi8>, + %bs: f32, + %binit: !b, + %cinit: !c) { + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, affine_map<(d0, d1, d2) -> (d1, d2)>, #identity]} + ins(%aq, %as : tensor<2x4x8xi8>, tensor<4x8xf32>) + outs(%ainit : !a) -> !a + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #identity]} + ins(%bq, %bs : tensor<4x8x3xi8>, f32) + outs(%binit : !b) -> !b + %zero = arith.constant 0.0 : f32 + %init = linalg.fill ins(%zero : f32) outs(%cinit : !c) -> !c + %test = linalg.generic { + id = "test", indexing_maps = [#lhs, #rhs, #out], + iterator_types = ["parallel", "parallel", "reduction", "reduction"] + } ins(%a, %b : !a, !b) outs(%init : !c) { + ^bb0(%x: f32, %y: f32, %acc: f32): + %p = arith.mulf %x, %y : f32 + %sum = arith.addf %p, %acc : f32 + linalg.yield %sum : f32 + } -> !c + return + } + )mlir")); + EXPECT_TRUE(failed(getQuantizedContraction(contraction()))); +} + +TEST_F(QuantizedContractionTest, RejectsDynamicIntegerReduction) { + ASSERT_TRUE(parse(R"mlir( + !a = tensor<2x?x8xf32> + !b = tensor<?x8x3xf32> + !c = tensor<2x3xf32> + #identity = affine_map<(d0, d1, d2) -> (d0, d1, d2)> + #scalar = affine_map<(d0, d1, d2) -> ()> + #lhs = affine_map<(m, n, g, k) -> (m, g, k)> + #rhs = affine_map<(m, n, g, k) -> (g, k, n)> + #out = affine_map<(m, n, g, k) -> (m, n)> + func.func @test( + %aq: tensor<2x?x8xi8>, + %as: f32, + %ainit: !a, + %bq: tensor<?x8x3xi8>, + %bs: f32, + %binit: !b, + %cinit: !c) { + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #identity]} + ins(%aq, %as : tensor<2x?x8xi8>, f32) + outs(%ainit : !a) -> !a + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #identity]} + ins(%bq, %bs : tensor<?x8x3xi8>, f32) + outs(%binit : !b) -> !b + %zero = arith.constant 0.0 : f32 + %init = linalg.fill ins(%zero : f32) outs(%cinit : !c) -> !c + %test = linalg.generic { + id = "test", indexing_maps = [#lhs, #rhs, #out], + iterator_types = ["parallel", "parallel", "reduction", "reduction"] + } ins(%a, %b : !a, !b) outs(%init : !c) { + ^bb0(%x: f32, %y: f32, %acc: f32): + %p = arith.mulf %x, %y : f32 + %sum = arith.addf %p, %acc : f32 + linalg.yield %sum : f32 + } -> !c + return + } + )mlir")); + EXPECT_TRUE(failed(getQuantizedContraction(contraction()))); +} + +TEST_F(QuantizedContractionTest, AcceptsDynamicFloatingReduction) { + ASSERT_TRUE(parse(R"mlir( + !a = tensor<2x?x8xf32> + !b = tensor<?x8x3xf32> + !c = tensor<2x3xf32> + #identity = affine_map<(d0, d1, d2) -> (d0, d1, d2)> + #scalar = affine_map<(d0, d1, d2) -> ()> + #lhs = affine_map<(m, n, g, k) -> (m, g, k)> + #rhs = affine_map<(m, n, g, k) -> (g, k, n)> + #out = affine_map<(m, n, g, k) -> (m, n)> + func.func @test( + %aq: tensor<2x?x8xi8>, + %as: tensor<?xf32>, + %ainit: !a, + %bq: tensor<?x8x3xi8>, + %bs: f32, + %binit: !b, + %cinit: !c) { + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, affine_map<(d0, d1, d2) -> (d1)>, #identity]} + ins(%aq, %as : tensor<2x?x8xi8>, tensor<?xf32>) + outs(%ainit : !a) -> !a + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #identity]} + ins(%bq, %bs : tensor<?x8x3xi8>, f32) + outs(%binit : !b) -> !b + %zero = arith.constant 0.0 : f32 + %init = linalg.fill ins(%zero : f32) outs(%cinit : !c) -> !c + %test = linalg.generic { + id = "test", indexing_maps = [#lhs, #rhs, #out], + iterator_types = ["parallel", "parallel", "reduction", "reduction"] + } ins(%a, %b : !a, !b) outs(%init : !c) { + ^bb0(%x: f32, %y: f32, %acc: f32): + %p = arith.mulf %x, %y : f32 + %sum = arith.addf %p, %acc : f32 + linalg.yield %sum : f32 + } -> !c + return + } + )mlir")); + auto detail = getQuantizedContraction(contraction()); + ASSERT_TRUE(succeeded(detail)); + EXPECT_EQ(detail->reductionExtent, 8); + EXPECT_THAT(detail->integerReductionDims, ElementsAre(3)); + EXPECT_THAT(detail->floatingReductionDims, ElementsAre(2)); +} + +TEST_F(QuantizedContractionTest, AcceptsSafeProductOfReductionExtents) { + ASSERT_TRUE(parse(R"mlir( + !a = tensor<2x256x511xf32> + !b = tensor<256x511x3xf32> + !c = tensor<2x3xf32> + #identity = affine_map<(d0, d1, d2) -> (d0, d1, d2)> + #scalar = affine_map<(d0, d1, d2) -> ()> + #lhs = affine_map<(m, n, g, k) -> (m, g, k)> + #rhs = affine_map<(m, n, g, k) -> (g, k, n)> + #out = affine_map<(m, n, g, k) -> (m, n)> + func.func @test( + %aq: tensor<2x256x511xi8>, + %as: f32, + %ainit: !a, + %bq: tensor<256x511x3xi8>, + %bs: f32, + %binit: !b, + %cinit: !c) { + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #identity]} + ins(%aq, %as : tensor<2x256x511xi8>, f32) + outs(%ainit : !a) -> !a + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #identity]} + ins(%bq, %bs : tensor<256x511x3xi8>, f32) + outs(%binit : !b) -> !b + %zero = arith.constant 0.0 : f32 + %init = linalg.fill ins(%zero : f32) outs(%cinit : !c) -> !c + %test = linalg.generic { + id = "test", indexing_maps = [#lhs, #rhs, #out], + iterator_types = ["parallel", "parallel", "reduction", "reduction"] + } ins(%a, %b : !a, !b) outs(%init : !c) { + ^bb0(%x: f32, %y: f32, %acc: f32): + %p = arith.mulf %x, %y : f32 + %sum = arith.addf %p, %acc : f32 + linalg.yield %sum : f32 + } -> !c + return + } + )mlir")); + auto detail = getQuantizedContraction(contraction()); + ASSERT_TRUE(succeeded(detail)); + EXPECT_EQ(detail->reductionExtent, 130816); +} + +TEST_F(QuantizedContractionTest, RejectsUnsafeProductOfReductionExtents) { + ASSERT_TRUE(parse(R"mlir( + !a = tensor<2x256x512xf32> + !b = tensor<256x512x3xf32> + !c = tensor<2x3xf32> + #identity = affine_map<(d0, d1, d2) -> (d0, d1, d2)> + #scalar = affine_map<(d0, d1, d2) -> ()> + #lhs = affine_map<(m, n, g, k) -> (m, g, k)> + #rhs = affine_map<(m, n, g, k) -> (g, k, n)> + #out = affine_map<(m, n, g, k) -> (m, n)> + func.func @test( + %aq: tensor<2x256x512xi8>, + %as: f32, + %ainit: !a, + %bq: tensor<256x512x3xi8>, + %bs: f32, + %binit: !b, + %cinit: !c) { + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #identity]} + ins(%aq, %as : tensor<2x256x512xi8>, f32) + outs(%ainit : !a) -> !a + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #identity]} + ins(%bq, %bs : tensor<256x512x3xi8>, f32) + outs(%binit : !b) -> !b + %zero = arith.constant 0.0 : f32 + %init = linalg.fill ins(%zero : f32) outs(%cinit : !c) -> !c + %test = linalg.generic { + id = "test", indexing_maps = [#lhs, #rhs, #out], + iterator_types = ["parallel", "parallel", "reduction", "reduction"] + } ins(%a, %b : !a, !b) outs(%init : !c) { + ^bb0(%x: f32, %y: f32, %acc: f32): + %p = arith.mulf %x, %y : f32 + %sum = arith.addf %p, %acc : f32 + linalg.yield %sum : f32 + } -> !c + return + } + )mlir")); + // Both extents are individually safe, but their product exceeds 131071. + EXPECT_TRUE(failed(getQuantizedContraction(contraction()))); +} + +TEST_F(QuantizedContractionTest, AcceptsCommutedMultiplyAndAdd) { + ASSERT_TRUE(parse(R"mlir( + !a = tensor<2x4x8xf32> + !b = tensor<4x8x3xf32> + !c = tensor<2x3xf32> + #identity = affine_map<(d0, d1, d2) -> (d0, d1, d2)> + #scalar = affine_map<(d0, d1, d2) -> ()> + #lhs = affine_map<(m, n, g, k) -> (m, g, k)> + #rhs = affine_map<(m, n, g, k) -> (g, k, n)> + #out = affine_map<(m, n, g, k) -> (m, n)> + func.func @test( + %aq: tensor<2x4x8xi8>, + %as: f32, + %ainit: !a, + %bq: tensor<4x8x3xi8>, + %bs: f32, + %binit: !b, + %cinit: !c) { + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #identity]} + ins(%aq, %as : tensor<2x4x8xi8>, f32) + outs(%ainit : !a) -> !a + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #identity]} + ins(%bq, %bs : tensor<4x8x3xi8>, f32) + outs(%binit : !b) -> !b + %zero = arith.constant 0.0 : f32 + %init = linalg.fill ins(%zero : f32) outs(%cinit : !c) -> !c + %plain = linalg.generic { + id = "plain", indexing_maps = [#lhs, #rhs, #out], + iterator_types = ["parallel", "parallel", "reduction", "reduction"] + } ins(%a, %b : !a, !b) outs(%init : !c) { + ^bb0(%x: f32, %y: f32, %acc: f32): + %p = arith.mulf %x, %y : f32 + %sum = arith.addf %p, %acc : f32 + linalg.yield %sum : f32 + } -> !c + %commuted_mul = linalg.generic { + id = "commuted_mul", indexing_maps = [#lhs, #rhs, #out], + iterator_types = ["parallel", "parallel", "reduction", "reduction"] + } ins(%a, %b : !a, !b) outs(%init : !c) { + ^bb0(%x: f32, %y: f32, %acc: f32): + %p = arith.mulf %y, %x : f32 + %sum = arith.addf %p, %acc : f32 + linalg.yield %sum : f32 + } -> !c + %commuted_add = linalg.generic { + id = "commuted_add", indexing_maps = [#lhs, #rhs, #out], + iterator_types = ["parallel", "parallel", "reduction", "reduction"] + } ins(%a, %b : !a, !b) outs(%init : !c) { + ^bb0(%x: f32, %y: f32, %acc: f32): + %p = arith.mulf %x, %y : f32 + %sum = arith.addf %acc, %p : f32 + linalg.yield %sum : f32 + } -> !c + %commuted_both = linalg.generic { + id = "commuted_both", indexing_maps = [#lhs, #rhs, #out], + iterator_types = ["parallel", "parallel", "reduction", "reduction"] + } ins(%a, %b : !a, !b) outs(%init : !c) { + ^bb0(%x: f32, %y: f32, %acc: f32): + %p = arith.mulf %y, %x : f32 + %sum = arith.addf %acc, %p : f32 + linalg.yield %sum : f32 + } -> !c + return + } + )mlir")); + for (StringRef id : + {"plain", "commuted_mul", "commuted_add", "commuted_both"}) { + SCOPED_TRACE(id.str()); + EXPECT_TRUE(succeeded(getQuantizedContraction(contraction(id)))); + } +} + +TEST_F(QuantizedContractionTest, RejectsBodyThatIgnoresAnInput) { + ASSERT_TRUE(parse(R"mlir( + !a = tensor<2x4x8xf32> + !b = tensor<4x8x3xf32> + !c = tensor<2x3xf32> + #identity = affine_map<(d0, d1, d2) -> (d0, d1, d2)> + #scalar = affine_map<(d0, d1, d2) -> ()> + #lhs = affine_map<(m, n, g, k) -> (m, g, k)> + #rhs = affine_map<(m, n, g, k) -> (g, k, n)> + #out = affine_map<(m, n, g, k) -> (m, n)> + func.func @test( + %aq: tensor<2x4x8xi8>, + %as: f32, + %ainit: !a, + %bq: tensor<4x8x3xi8>, + %bs: f32, + %binit: !b, + %cinit: !c) { + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #identity]} + ins(%aq, %as : tensor<2x4x8xi8>, f32) + outs(%ainit : !a) -> !a + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #identity]} + ins(%bq, %bs : tensor<4x8x3xi8>, f32) + outs(%binit : !b) -> !b + %zero = arith.constant 0.0 : f32 + %init = linalg.fill ins(%zero : f32) outs(%cinit : !c) -> !c + %test = linalg.generic { + id = "test", indexing_maps = [#lhs, #rhs, #out], + iterator_types = ["parallel", "parallel", "reduction", "reduction"] + } ins(%a, %b : !a, !b) outs(%init : !c) { + ^bb0(%x: f32, %y: f32, %acc: f32): + %p = arith.mulf %x, %x : f32 + %sum = arith.addf %p, %acc : f32 + linalg.yield %sum : f32 + } -> !c + return + } + )mlir")); + EXPECT_TRUE(failed(getQuantizedContraction(contraction()))); +} + +TEST_F(QuantizedContractionTest, RejectsNonzeroInit) { + ASSERT_TRUE(parse(R"mlir( + !a = tensor<2x4x8xf32> + !b = tensor<4x8x3xf32> + !c = tensor<2x3xf32> + #identity = affine_map<(d0, d1, d2) -> (d0, d1, d2)> + #scalar = affine_map<(d0, d1, d2) -> ()> + #lhs = affine_map<(m, n, g, k) -> (m, g, k)> + #rhs = affine_map<(m, n, g, k) -> (g, k, n)> + #out = affine_map<(m, n, g, k) -> (m, n)> + func.func @test( + %aq: tensor<2x4x8xi8>, + %as: f32, + %ainit: !a, + %bq: tensor<4x8x3xi8>, + %bs: f32, + %binit: !b, + %cinit: !c) { + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #identity]} + ins(%aq, %as : tensor<2x4x8xi8>, f32) + outs(%ainit : !a) -> !a + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [#identity, #scalar, #identity]} + ins(%bq, %bs : tensor<4x8x3xi8>, f32) + outs(%binit : !b) -> !b + %zero = arith.constant 1.0 : f32 + %init = linalg.fill ins(%zero : f32) outs(%cinit : !c) -> !c + %test = linalg.generic { + id = "test", indexing_maps = [#lhs, #rhs, #out], + iterator_types = ["parallel", "parallel", "reduction", "reduction"] + } ins(%a, %b : !a, !b) outs(%init : !c) { + ^bb0(%x: f32, %y: f32, %acc: f32): + %p = arith.mulf %x, %y : f32 + %sum = arith.addf %p, %acc : f32 + linalg.yield %sum : f32 + } -> !c + return + } + )mlir")); + EXPECT_TRUE(failed(getQuantizedContraction(contraction()))); +} + +TEST_F(QuantizedContractionTest, ExplicitExtentsCoveredAcrossMaps) { + EXPECT_EQ(getExplicitExtentDimsMap({map(3, {2, 0}), map(3, {1})}), + map(3, {})); +} + +TEST_F(QuantizedContractionTest, ExplicitExtentsPreserveDomainOrder) { + EXPECT_EQ(getExplicitExtentDimsMap({map(4, {2}), map(4, {2})}), + map(4, {0, 1, 3})); +} + +TEST_F(QuantizedContractionTest, ExplicitExtentsRequireBareDimensions) { + AffineExpr d0 = getAffineDimExpr(0, &context); + AffineExpr d1 = getAffineDimExpr(1, &context); + AffineMap window = AffineMap::get(2, 0, {2 * d0 + d1}, &context); + EXPECT_EQ(getExplicitExtentDimsMap({window, map(2, {0})}), map(2, {1})); + EXPECT_EQ(getExplicitExtentDimsMap({window}), map(2, {0, 1})); +} + +TEST_F(QuantizedContractionTest, ExplicitExtentsIgnoreConstants) { + AffineMap constant = + AffineMap::get(2, 0, {getAffineConstantExpr(0, &context)}, &context); + EXPECT_EQ(getExplicitExtentDimsMap({constant, map(2, {})}), map(2, {0, 1})); + EXPECT_EQ(getExplicitExtentDimsMap({map(0, {})}), map(0, {})); +} + +} // namespace +} // namespace mlir::iree_compiler::GlobalOptimization::detail
diff --git a/tests/e2e/quantization/BUILD.bazel b/tests/e2e/quantization/BUILD.bazel new file mode 100644 index 0000000..9c17c8c --- /dev/null +++ b/tests/e2e/quantization/BUILD.bazel
@@ -0,0 +1,41 @@ +# 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 + +load("//build_tools/bazel:iree_check_test.bzl", "iree_check_single_backend_test_suite") + +package( + features = ["layering_check"], + licenses = ["notice"], # Apache 2.0 +) + +iree_check_single_backend_test_suite( + name = "check_llvm-cpu_local-task", + srcs = [ + "contraction_scaling.mlir", + ], + compiler_flags = [ + "--iree-llvmcpu-target-cpu=generic", + "--iree-global-opt-enable-qdq-to-integer-math=true", + ], + driver = "local-task", + # Accepted numerical differences have separate reference/rewrite checks. + runner_args = ["--gtest_filter=-*_reference"], + target_backend = "llvm-cpu", +) + +iree_check_single_backend_test_suite( + name = "check_reference_llvm-cpu_local-task", + srcs = [ + "contraction_scaling.mlir", + ], + compiler_flags = [ + "--iree-global-opt-enable-qdq-to-integer-math=false", + "--iree-llvmcpu-target-cpu=generic", + ], + driver = "local-task", + runner_args = ["--gtest_filter=-*_rewritten"], + target_backend = "llvm-cpu", +)
diff --git a/tests/e2e/quantization/CMakeLists.txt b/tests/e2e/quantization/CMakeLists.txt new file mode 100644 index 0000000..eca30df --- /dev/null +++ b/tests/e2e/quantization/CMakeLists.txt
@@ -0,0 +1,45 @@ +################################################################################ +# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from # +# tests/e2e/quantization/BUILD.bazel # +# # +# Use iree_cmake_extra_content from iree/build_defs.oss.bzl to add arbitrary # +# CMake-only content. # +# # +# To disable autogeneration for this file entirely, delete this header. # +################################################################################ + +iree_add_all_subdirs() + +iree_check_single_backend_test_suite( + NAME + check_llvm-cpu_local-task + SRCS + "contraction_scaling.mlir" + TARGET_BACKEND + "llvm-cpu" + DRIVER + "local-task" + COMPILER_FLAGS + "--iree-llvmcpu-target-cpu=generic" + "--iree-global-opt-enable-qdq-to-integer-math=true" + RUNNER_ARGS + "--gtest_filter=-*_reference" +) + +iree_check_single_backend_test_suite( + NAME + check_reference_llvm-cpu_local-task + SRCS + "contraction_scaling.mlir" + TARGET_BACKEND + "llvm-cpu" + DRIVER + "local-task" + COMPILER_FLAGS + "--iree-global-opt-enable-qdq-to-integer-math=false" + "--iree-llvmcpu-target-cpu=generic" + RUNNER_ARGS + "--gtest_filter=-*_rewritten" +) + +### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/tests/e2e/quantization/contraction_scaling.mlir b/tests/e2e/quantization/contraction_scaling.mlir new file mode 100644 index 0000000..b3bdbfe --- /dev/null +++ b/tests/e2e/quantization/contraction_scaling.mlir
@@ -0,0 +1,708 @@ +// Runs with the integer rewrite enabled and disabled. +// Inputs stay unfolded so these checks execute the compiled contractions. + +// Both correction products and the cross term vary by block and output axis. +func.func @blockwise_zero_point_corrections() { + %aq = util.unfoldable_constant dense<[[[2, 3], [4, 6]], [[5, 8], [9, 10]]]> : tensor<2x2x2xi8> + %bq = util.unfoldable_constant dense<[[[2, 4, 6], [3, 5, 8]], [[6, 7, 9], [7, 9, 10]]]> : tensor<2x2x3xi8> + %sa = util.unfoldable_constant dense<[[1.0, 0.5], [2.0, 1.0]]> : tensor<2x2xf32> + %sb = util.unfoldable_constant dense<[[1.0, 2.0, 1.0], [0.5, 1.0, 2.0]]> : tensor<2x3xf32> + %za = util.unfoldable_constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi8> + %zb = util.unfoldable_constant dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi8> + %ai = tensor.empty() : tensor<2x2x2xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(m, g, k) -> (m, g, k)>, + affine_map<(m, g, k) -> (m, g)>, + affine_map<(m, g, k) -> (m, g)>, + affine_map<(m, g, k) -> (m, g, k)>]} + ins(%aq, %sa, %za : tensor<2x2x2xi8>, tensor<2x2xf32>, tensor<2x2xi8>) + outs(%ai : tensor<2x2x2xf32>) -> tensor<2x2x2xf32> + %bi = tensor.empty() : tensor<2x2x3xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(g, k, n) -> (g, k, n)>, + affine_map<(g, k, n) -> (g, n)>, + affine_map<(g, k, n) -> (g, n)>, + affine_map<(g, k, n) -> (g, k, n)>]} + ins(%bq, %sb, %zb : tensor<2x2x3xi8>, tensor<2x3xf32>, tensor<2x3xi8>) + outs(%bi : tensor<2x2x3xf32>) -> tensor<2x2x3xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<2x3xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<2x3xf32>) -> tensor<2x3xf32> + %result = linalg.contract + indexing_maps = [affine_map<(m, n, g, k) -> (m, g, k)>, + affine_map<(m, n, g, k) -> (g, k, n)>, + affine_map<(m, n, g, k) -> (m, n)>] + ins(%a, %b : tensor<2x2x2xf32>, tensor<2x2x3xf32>) + outs(%init : tensor<2x3xf32>) -> tensor<2x3xf32> + check.expect_eq_const(%result, dense<[[9.0, 26.0, 35.0], [38.0, 110.0, 140.0]]> : tensor<2x3xf32>) : tensor<2x3xf32> + return +} + +// Nonuniform image values distinguish the strided/dilated windows. Each output +// channel has a different zero point and scale but the same dequantized filter. +func.func @convolution_zero_point_corrections() { + %aq = util.unfoldable_constant dense<[[[[1], [2], [3], [4], [5]], [[6], [7], [8], [9], [10]], [[11], [12], [13], [14], [15]], [[16], [17], [18], [19], [20]], [[21], [22], [23], [24], [25]]]]> : tensor<1x5x5x1xi8> + %bq = util.unfoldable_constant dense<[[[[7, 9]], [[7, 9]]], [[[7, 9]], [[7, 9]]]]> : tensor<2x2x1x2xi8> + %sa = util.unfoldable_constant 0.5 : f32 + %sb = util.unfoldable_constant dense<[1.0, 0.5]> : tensor<2xf32> + %za = util.unfoldable_constant 3 : i8 + %zb = util.unfoldable_constant dense<[6, 7]> : tensor<2xi8> + %ai = tensor.empty() : tensor<1x5x5x1xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(n, h, w, c) -> (n, h, w, c)>, + affine_map<(n, h, w, c) -> ()>, + affine_map<(n, h, w, c) -> ()>, + affine_map<(n, h, w, c) -> (n, h, w, c)>]} + ins(%aq, %sa, %za : tensor<1x5x5x1xi8>, f32, i8) + outs(%ai : tensor<1x5x5x1xf32>) -> tensor<1x5x5x1xf32> + %bi = tensor.empty() : tensor<2x2x1x2xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(h, w, c, f) -> (h, w, c, f)>, + affine_map<(h, w, c, f) -> (f)>, + affine_map<(h, w, c, f) -> (f)>, + affine_map<(h, w, c, f) -> (h, w, c, f)>]} + ins(%bq, %sb, %zb : tensor<2x2x1x2xi8>, tensor<2xf32>, tensor<2xi8>) + outs(%bi : tensor<2x2x1x2xf32>) -> tensor<2x2x1x2xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x2x2x2xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32> + %result = linalg.conv_2d_nhwc_hwcf {dilations = dense<2> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} + ins(%a, %b : tensor<1x5x5x1xf32>, tensor<2x2x1x2xf32>) + outs(%init : tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32> + check.expect_eq_const(%result, dense<[[[[8.0, 8.0], [12.0, 12.0]], [[28.0, 28.0], [32.0, 32.0]]]]> : tensor<1x2x2x2xf32>) : tensor<1x2x2x2xf32> + return +} + +// NCHW layout with two input channels: all three filter dimensions reduce. +// Distinct signed data and per-filter parameters expose channel/layout swaps. +func.func @convolution_nchw_asymmetric() { + %aq = util.unfoldable_constant dense<[[[[-8, -5, -2, 1], [-3, 0, 3, 6], [2, 5, 8, -6]], [[3, 6, -8, -5], [8, -6, -3, 0], [-4, -1, 2, 5]]]]> : tensor<1x2x3x4xi8> + %bq = util.unfoldable_constant dense<[[[[-5, -4], [-2, -1]], [[0, 1], [3, 4]]], [[[2, 3], [5, -5]], [[-4, -3], [-1, 0]]]]> : tensor<2x2x2x2xi8> + %sb = util.unfoldable_constant dense<[1.0, 0.25]> : tensor<2xf32> + %zb = util.unfoldable_constant dense<[-1, 2]> : tensor<2xi8> + %sa = util.unfoldable_constant 0.5 : f32 + %za = util.unfoldable_constant -2 : i8 + %ai = tensor.empty() : tensor<1x2x3x4xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(n, c, h, w) -> (n, c, h, w)>, + affine_map<(n, c, h, w) -> ()>, + affine_map<(n, c, h, w) -> ()>, + affine_map<(n, c, h, w) -> (n, c, h, w)>]} + ins(%aq, %sa, %za : tensor<1x2x3x4xi8>, f32, i8) + outs(%ai : tensor<1x2x3x4xf32>) -> tensor<1x2x3x4xf32> + %bi = tensor.empty() : tensor<2x2x2x2xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(f, c, h, w) -> (f, c, h, w)>, + affine_map<(f, c, h, w) -> (f)>, + affine_map<(f, c, h, w) -> (f)>, + affine_map<(f, c, h, w) -> (f, c, h, w)>]} + ins(%bq, %sb, %zb : tensor<2x2x2x2xi8>, tensor<2xf32>, tensor<2xi8>) + outs(%bi : tensor<2x2x2x2xf32>) -> tensor<2x2x2x2xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x2x2x3xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32> + %result = linalg.conv_2d_nchw_fchw {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} + ins(%a, %b : tensor<1x2x3x4xf32>, tensor<2x2x2x2xf32>) + outs(%init : tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32> + check.expect_eq_const(%result, dense<[[[[37.5, -7.5, -10.0], [-3.5, -6.0, 0.0]], [[-14.0, -4.125, 1.5], [-8.875, -3.25, 4.5]]]]> : tensor<1x2x2x3xf32>) : tensor<1x2x2x3xf32> + return +} + +// Channels remain parallel in depthwise convolution. Both operands have +// channel-specific scales and zero points; activation bytes exercise extui. +// Different stride/dilation axes distinguish the sampled spatial windows. +func.func @depthwise_asymmetric_strided_dilated() { + %aq = util.unfoldable_constant dense<[[[[128, 133], [131, 136], [134, 139], [137, 142], [140, 145]], [[135, 140], [138, 143], [141, 146], [144, 130], [128, 133]], [[142, 128], [145, 131], [129, 134], [132, 137], [135, 140]], [[130, 135], [133, 138], [136, 141], [139, 144], [142, 128]]]]> : tensor<1x4x5x2xi8> + %bq = util.unfoldable_constant dense<[[[-3, 4], [2, -1]], [[5, 0], [-2, 6]]]> : tensor<2x2x2xi8> + %sa = util.unfoldable_constant dense<[0.5, 2.0]> : tensor<2xf32> + %za = util.unfoldable_constant dense<[132, 136]> : tensor<2xi8> + %sb = util.unfoldable_constant dense<[2.0, 0.25]> : tensor<2xf32> + %zb = util.unfoldable_constant dense<[-1, 2]> : tensor<2xi8> + %ai = tensor.empty() : tensor<1x4x5x2xf32> + %a = iree_linalg_ext.dequantize_affine + {input_unsigned, zp_unsigned, indexing_maps = [affine_map<(n, h, w, c) -> (n, h, w, c)>, + affine_map<(n, h, w, c) -> (c)>, + affine_map<(n, h, w, c) -> (c)>, + affine_map<(n, h, w, c) -> (n, h, w, c)>]} + ins(%aq, %sa, %za : tensor<1x4x5x2xi8>, tensor<2xf32>, tensor<2xi8>) + outs(%ai : tensor<1x4x5x2xf32>) -> tensor<1x4x5x2xf32> + %bi = tensor.empty() : tensor<2x2x2xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(h, w, c) -> (h, w, c)>, + affine_map<(h, w, c) -> (c)>, + affine_map<(h, w, c) -> (c)>, + affine_map<(h, w, c) -> (h, w, c)>]} + ins(%bq, %sb, %zb : tensor<2x2x2xi8>, tensor<2xf32>, tensor<2xi8>) + outs(%bi : tensor<2x2x2xf32>) -> tensor<2x2x2xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x2x2x2xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32> + %result = linalg.depthwise_conv_2d_nhwc_hwc {dilations = dense<[2, 1]> : tensor<2xi64>, strides = dense<[1, 2]> : tensor<2xi64>} + ins(%a, %b : tensor<1x4x5x2xf32>, tensor<2x2x2xf32>) + outs(%init : tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32> + check.expect_eq_const(%result, dense<[[[[52.0, -5.0], [-7.0, -2.0]], [[-1.0, -1.5], [35.0, 30.0]]]]> : tensor<1x2x2x2xf32>) : tensor<1x2x2x2xf32> + return +} + +// A custom strided window contraction with interleaved reduction/parallel +// loops, a broadcast RHS batch, and output axes ordered (n, batch, m). +// C[n,t,m] = sum_{r,c} DQ(A[t,2*m+r,c]) * DQ(B[n,c,r]). +func.func @generic_strided_permuted_contraction() { + %aq = util.unfoldable_constant dense<[[[128, 131], [132, 135], [136, 139], [140, 143], [144, 147], [148, 128], [129, 132]], [[137, 140], [141, 144], [145, 148], [149, 129], [130, 133], [134, 137], [138, 141]]]> : tensor<2x7x2xi8> + %bq = util.unfoldable_constant dense<[[[-6, -4, -2], [-3, -1, 1]], [[-1, 1, 3], [2, 4, 6]], [[4, 6, -5], [-6, -4, -2]]]> : tensor<3x2x3xi8> + %sa = util.unfoldable_constant dense<[0.5, 1.0]> : tensor<2xf32> + %za = util.unfoldable_constant dense<[130, 134]> : tensor<2xi8> + %sb = util.unfoldable_constant dense<[1.0, 0.5, 2.0]> : tensor<3xf32> + %zb = util.unfoldable_constant dense<[-2, 1, 3]> : tensor<3xi8> + %ai = tensor.empty() : tensor<2x7x2xf32> + %a = iree_linalg_ext.dequantize_affine + {input_unsigned, zp_unsigned, indexing_maps = [affine_map<(t, x, c) -> (t, x, c)>, + affine_map<(t, x, c) -> (t)>, + affine_map<(t, x, c) -> (t)>, + affine_map<(t, x, c) -> (t, x, c)>]} + ins(%aq, %sa, %za : tensor<2x7x2xi8>, tensor<2xf32>, tensor<2xi8>) + outs(%ai : tensor<2x7x2xf32>) -> tensor<2x7x2xf32> + %bi = tensor.empty() : tensor<3x2x3xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(n, c, r) -> (n, c, r)>, + affine_map<(n, c, r) -> (n)>, + affine_map<(n, c, r) -> (n)>, + affine_map<(n, c, r) -> (n, c, r)>]} + ins(%bq, %sb, %zb : tensor<3x2x3xi8>, tensor<3xf32>, tensor<3xi8>) + outs(%bi : tensor<3x2x3xf32>) -> tensor<3x2x3xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<3x2x3xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<3x2x3xf32>) -> tensor<3x2x3xf32> + %result = linalg.generic { + indexing_maps = [affine_map<(c, n, t, r, m) -> (t, m * 2 + r, c)>, + affine_map<(c, n, t, r, m) -> (n, c, r)>, + affine_map<(c, n, t, r, m) -> (n, t, m)>], + iterator_types = ["reduction", "parallel", "parallel", "reduction", "parallel"]} + ins(%a, %b : tensor<2x7x2xf32>, tensor<3x2x3xf32>) + outs(%init : tensor<3x2x3xf32>) { + ^bb0(%lhs: f32, %rhs: f32, %acc: f32): + %product = arith.mulf %lhs, %rhs : f32 + %sum = arith.addf %acc, %product : f32 + linalg.yield %sum : f32 + } -> tensor<3x2x3xf32> + check.expect_eq_const(%result, dense<[[[17.5, 5.5, -52.5], [20.0, -96.0, 41.0]], [[19.25, 37.25, -2.25], [61.0, -18.0, 29.5]], [[-133.0, -333.0, -73.0], [-516.0, 4.0, -166.0]]]> : tensor<3x2x3xf32>) : tensor<3x2x3xf32> + return +} + +// Both asymmetric operands vary by block and output axis, with different +// storage permutations folded into dequantization. Only k is reduced in i32; +// g is retained in partials and reduced after applying both block scales. +// C[n,m] = sum_{g,k} DQ(A[k,m,g]) * DQ(B[n,k,g]). +func.func @blockwise_asymmetric_permuted_storage() { + %aq = util.unfoldable_constant dense<[[[128, 131], [133, 136]], [[135, 138], [140, 143]], [[142, 145], [147, 129]]]> : tensor<3x2x2xi8> + %bq = util.unfoldable_constant dense<[[[-8, -3], [-5, 0], [-2, 3]], [[-1, 4], [2, 7], [5, -7]], [[6, -6], [-8, -3], [-5, 0]]]> : tensor<3x3x2xi8> + %sa = util.unfoldable_constant dense<[[0.5, 2.0], [1.0, 0.25]]> : tensor<2x2xf32> + %za = util.unfoldable_constant dense<[[130, 135], [134, 129]]> : tensor<2x2xi8> + %sb = util.unfoldable_constant dense<[[1.0, 0.5], [2.0, 0.25], [0.5, 2.0]]> : tensor<3x2xf32> + %zb = util.unfoldable_constant dense<[[-3, 1], [2, -4], [-1, 3]]> : tensor<3x2xi8> + %ai = tensor.empty() : tensor<2x2x3xf32> + %a = iree_linalg_ext.dequantize_affine + {input_unsigned, zp_unsigned, indexing_maps = [affine_map<(m, g, k) -> (k, m, g)>, + affine_map<(m, g, k) -> (g, m)>, + affine_map<(m, g, k) -> (g, m)>, + affine_map<(m, g, k) -> (m, g, k)>]} + ins(%aq, %sa, %za : tensor<3x2x2xi8>, tensor<2x2xf32>, tensor<2x2xi8>) + outs(%ai : tensor<2x2x3xf32>) -> tensor<2x2x3xf32> + %bi = tensor.empty() : tensor<2x3x3xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(g, k, n) -> (n, k, g)>, + affine_map<(g, k, n) -> (n, g)>, + affine_map<(g, k, n) -> (n, g)>, + affine_map<(g, k, n) -> (g, k, n)>]} + ins(%bq, %sb, %zb : tensor<3x3x2xi8>, tensor<3x2xf32>, tensor<3x2xi8>) + outs(%bi : tensor<2x3x3xf32>) -> tensor<2x3x3xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<3x2xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<3x2xf32>) -> tensor<3x2xf32> + %result = linalg.contract + indexing_maps = [affine_map<(m, n, g, k) -> (m, g, k)>, + affine_map<(m, n, g, k) -> (g, k, n)>, + affine_map<(m, n, g, k) -> (n, m)>] + ins(%a, %b : tensor<2x2x3xf32>, tensor<2x3x3xf32>) + outs(%init : tensor<3x2xf32>) -> tensor<3x2xf32> + check.expect_eq_const(%result, dense<[[21.0, 18.75], [38.75, 181.125], [-84.25, -170.5]]> : tensor<3x2xf32>) : tensor<3x2xf32> + return +} + +// Anticipated QDQ Linalg translation of a PT2E/XNNPACK bias-free Linear(16,16) +// with an 8x16 input: Q -> DQ -> matmul(input, DQ(weight).T) -> Q -> DQ. +// Input, quantized weights, scales, zero points, and expected output are from +// a seed-0 PyTorch run. The RHS indexing map incorporates aten.transpose.int. +func.func @pt2e_linear_qdq() { + %input_scale = arith.constant 0.024613387882709503 : f32 + %weight_scale = arith.constant 0.0019636377692222595 : f32 + %output_scale = arith.constant 0.010716261342167854 : f32 + %input_zp = arith.constant -33 : i64 + %weight_zp = arith.constant 0 : i64 + %output_zp = arith.constant -10 : i64 + %zero = arith.constant 0.0 : f32 + + // Prevent constant folding so the QDQ graph executes at runtime. + %input = util.unfoldable_constant dense<[ + [-0.03937254101037979, -0.8014721870422363, -0.49554431438446045, -0.36151406168937683, 0.5851132273674011, -1.1560065746307373, -0.14336487650871277, -0.19474059343338013, -0.08556340634822845, 1.3945199251174927, 0.5969000458717346, -0.4828483462333679, -0.3660986125469208, -1.3270524740219116, 1.695279598236084, 2.0654995441436768], + [-0.2339576780796051, 0.7073183059692383, 0.5800480842590332, 0.2683020830154419, -2.05893874168396, 0.5340206623077393, -0.5353949069976807, -0.863664984703064, -0.023494405671954155, 1.171669840812683, 0.3986871838569641, -0.19871552288532257, -1.1559407711029053, -0.3166695237159729, 0.9402980804443359, -1.1469545364379883], + [0.5588033199310303, 0.7917584776878357, -0.18467576801776886, -0.7317724823951721, -0.0806519165635109, -0.9800607562065125, 0.060491468757390976, -0.4889546036720276, -0.8137312531471252, 0.8199948072433472, -0.633173406124115, 1.2947547435760498, 1.4628292322158813, -0.6204336285591125, 0.988389253616333, -0.4321781396865845], + [-0.6232188940048218, -0.21624533832073212, -0.48867544531822205, 0.7869559526443481, 0.10759072750806808, -1.0714776515960693, -0.11664776504039764, -1.0169707536697388, -1.1979783773422241, 0.47843775153160095, -1.2295243740081787, -1.3700363636016846, 1.5435417890548706, -0.0332069993019104, -0.4186263680458069, -0.25559672713279724], + [-0.12923042476177216, -0.05459475517272949, 0.4083467423915863, 1.1263659000396729, 1.9350574016571045, 1.0076850652694702, 1.0046416521072388, -0.43351978063583374, -1.2425976991653442, 1.2845500707626343, 0.24377228319644928, 0.5303686261177063, -0.014530729502439499, -2.235718250274658, 1.4660301208496094, -1.2190581560134888], + [0.64423006772995, 3.9300038814544678, -0.1244242936372757, 0.29534170031547546, 0.38265419006347656, -0.549721360206604, -0.9940357208251953, 1.345936894416809, 1.9456682205200195, -1.2903639078140259, -2.3494760990142822, -2.068861961364746, 0.9094210863113403, -0.6946200728416443, 1.9594571590423584, -1.1038278341293335], + [0.5411418080329895, 1.5389583110809326, 1.0860490798950195, 1.246405005455017, 0.11507505923509598, 1.619307518005371, 0.4636935293674469, 1.3007354736328125, 0.873230516910553, 0.0651267021894455, 0.7732411026954651, -0.970138430595398, -0.8876768946647644, -0.3183162808418274, -0.3344036936759949, 0.4542836546897888], + [0.49895304441452026, 0.8779974579811096, 0.38944435119628906, 1.462517499923706, 0.47950610518455505, -0.5333999991416931, -0.03465134650468826, 0.6572969555854797, -0.3112243115901947, -0.5620035529136658, -0.4834926426410675, -1.2721126079559326, -0.17401821911334991, 0.5541167259216309, -0.18165524303913116, -0.23447339236736298] + ]> : tensor<8x16xf32> + // Signed bytes decoded from torch_tensor_16_16_torch.int8, excluding its + // four-byte resource alignment header. Rows retain Torch weight layout [N,K]. + %weight = util.unfoldable_constant dense<[ + [-1, 68, -105, -94, -49, 34, -3, 101, -11, 34, -38, -25, -122, -84, -52, 5], + [50, 76, -86, -55, 46, 106, -26, 95, -21, 13, 115, -118, -80, -32, -50, 110], + [-83, -59, -89, -119, -74, 109, 57, 62, 7, -65, 22, -119, -92, -66, 80, 75], + [-56, -5, 81, 127, 51, 17, 85, -75, 24, -99, -88, -66, 58, 51, -75, 38], + [70, -16, 5, 29, 79, 122, -98, -47, 50, 105, 111, 112, 25, -111, 12, -80], + [-119, 113, 97, -127, 24, -21, -21, -58, 49, -75, 47, 64, 91, 48, -126, -83], + [64, 27, -99, -73, 120, 86, -56, -32, -121, -2, -96, -98, -7, 19, -52, 76], + [-77, 116, 87, -107, -32, 6, 19, 30, 50, 8, -62, 60, -122, -75, -32, -62], + [-45, -104, -27, 27, -83, -7, 91, -13, 4, -11, 26, 81, 121, 81, 121, -9], + [-114, -60, 87, -1, -63, -98, -119, -107, -26, 70, 69, -123, 79, -100, -27, -52], + [-25, -25, -114, -110, -20, 2, -58, 48, -115, -9, 112, -52, 115, 46, -115, 81], + [-15, -57, 102, -103, 14, -27, 91, 36, 61, 45, -31, -27, -105, 69, 101, 87], + [-90, 6, -90, -70, -74, 44, -76, -3, 5, 82, -96, -87, -74, 89, -46, 107], + [46, 16, -1, -25, 16, -29, -1, 16, -100, -67, 103, -103, -9, 126, 46, 4], + [-110, 63, -91, -36, -43, -19, 1, 105, 16, 114, 78, -80, 57, -90, -54, 37], + [42, 96, -41, 0, 66, -123, 92, -105, 2, -22, -67, 17, 105, -37, -76, -47] + ]> : tensor<16x16xi8> + %input_empty = tensor.empty() : tensor<8x16xf32> + + %input_q_empty = tensor.empty() : tensor<8x16xi8> + %input_q = iree_linalg_ext.quantize_affine + {indexing_maps = [affine_map<(m, k) -> (m, k)>, + affine_map<(m, k) -> ()>, affine_map<(m, k) -> ()>, + affine_map<(m, k) -> (m, k)>], + quant_min = -128 : i64, quant_max = 127 : i64} + ins(%input, %input_scale, %input_zp : tensor<8x16xf32>, f32, i64) + outs(%input_q_empty : tensor<8x16xi8>) -> tensor<8x16xi8> + %input_dq = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(m, k) -> (m, k)>, + affine_map<(m, k) -> ()>, affine_map<(m, k) -> ()>, + affine_map<(m, k) -> (m, k)>]} + ins(%input_q, %input_scale, %input_zp : tensor<8x16xi8>, f32, i64) + outs(%input_empty : tensor<8x16xf32>) -> tensor<8x16xf32> + %weight_dq_empty = tensor.empty() : tensor<16x16xf32> + %weight_dq = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(n, k) -> (n, k)>, + affine_map<(n, k) -> ()>, affine_map<(n, k) -> ()>, + affine_map<(n, k) -> (n, k)>]} + ins(%weight, %weight_scale, %weight_zp : tensor<16x16xi8>, f32, i64) + outs(%weight_dq_empty : tensor<16x16xf32>) -> tensor<16x16xf32> + %init = linalg.fill ins(%zero : f32) outs(%input_empty : tensor<8x16xf32>) -> tensor<8x16xf32> + %linear = linalg.contract + indexing_maps = [affine_map<(m, n, k) -> (m, k)>, + affine_map<(m, n, k) -> (n, k)>, + affine_map<(m, n, k) -> (m, n)>] + ins(%input_dq, %weight_dq : tensor<8x16xf32>, tensor<16x16xf32>) + outs(%init : tensor<8x16xf32>) -> tensor<8x16xf32> + %output_q = iree_linalg_ext.quantize_affine + {indexing_maps = [affine_map<(m, n) -> (m, n)>, + affine_map<(m, n) -> ()>, affine_map<(m, n) -> ()>, + affine_map<(m, n) -> (m, n)>], + quant_min = -128 : i64, quant_max = 127 : i64} + ins(%linear, %output_scale, %output_zp : tensor<8x16xf32>, f32, i64) + outs(%input_q_empty : tensor<8x16xi8>) -> tensor<8x16xi8> + %output = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(m, n) -> (m, n)>, + affine_map<(m, n) -> ()>, affine_map<(m, n) -> ()>, + affine_map<(m, n) -> (m, n)>]} + ins(%output_q, %output_scale, %output_zp : tensor<8x16xi8>, f32, i64) + outs(%input_empty : tensor<8x16xf32>) -> tensor<8x16xf32> + + // Recorded torch_quantized_output. The original IREE lowering matched + // PyTorch exactly; require the same result with and without the QDQ rewrite. + check.expect_eq_const(%output, dense<[ + [0.11787887662649155, 0.4929480254650116, 0.67512446641922, -0.7179895043373108, 0.13931140303611755, -1.2323700189590454, 0.16074392199516296, -0.43936672806739807, 0.06429757177829742, 0.5465293526649475, 0.11787887662649155, 0.7930033206939697, 0.2786228060722351, -0.03214878588914871, 0.6215431690216064, -0.3214878439903259], + [0.2786228060722351, -0.2250414937734604, 0.18217644095420837, -0.6322594285011292, 0.4607992470264435, -0.18217644095420837, -0.7287057638168335, 0.6644082069396973, -0.010716261342167854, 0.7501382827758789, -0.7394220232963562, 0.02143252268433571, 0.2786228060722351, -0.15002766251564026, 0.09644635021686554, -0.6108269095420837], + [-0.17146018147468567, -0.7930033206939697, -0.7715708017349243, -0.5786781311035156, 0.33220410346984863, 0.2143252193927765, -0.05358130484819412, 0.08573009073734283, 0.4500829875469208, 0.12859514355659485, -0.04286504536867142, -0.2893390655517578, -0.5143805742263794, -0.2786228060722351, -0.010716261342167854, 0.8787334561347961], + [-0.4607992470264435, -0.6429756879806519, -0.5358130931854248, 0.7501382827758789, -0.6429756879806519, -0.07501383125782013, 0.5786781311035156, -0.7287057638168335, 0.13931140303611755, 1.103774905204773, 0.428650438785553, -0.6536919474601746, 0.2893390655517578, 0.06429757177829742, 0.2571902871131897, 0.9001659750938416], + [-0.2357577532529831, -0.2143252193927765, -0.2679065465927124, -0.12859514355659485, 1.4681278467178345, -0.7179895043373108, 0.0, 0.06429757177829742, -0.010716261342167854, 0.15002766251564026, -1.0394773483276367, -0.2786228060722351, -1.2645188570022583, -0.4929480254650116, -0.07501383125782013, 0.20360895991325378], + [0.5250968337059021, 0.17146018147468567, -0.08573009073734283, 0.2786228060722351, -0.6429756879806519, 0.2464740127325058, 0.2679065465927124, 0.8358683586120605, -0.921598494052887, -0.43936672806739807, -1.2323700189590454, -0.2893390655517578, -0.10716260969638824, -0.03214878588914871, 0.2893390655517578, 0.8787334561347961], + [0.38578540086746216, 1.1895049810409546, 0.2679065465927124, 0.33220410346984863, 0.3107715845108032, -0.2679065465927124, -0.12859514355659485, 0.40721791982650757, -0.8037195801734924, -0.6322594285011292, -0.5572456121444702, 0.12859514355659485, -0.16074392199516296, 0.06429757177829742, 0.33220410346984863, -0.5893943905830383], + [-0.15002766251564026, 0.15002766251564026, -0.4607992470264435, 0.67512446641922, -0.6215431690216064, -0.2893390655517578, 0.2250414937734604, -0.2357577532529831, -0.4607992470264435, -0.16074392199516296, -0.3107715845108032, -0.20360895991325378, -0.16074392199516296, 0.4929480254650116, -0.18217644095420837, 0.2464740127325058] + ]> : tensor<8x16xf32>) : tensor<8x16xf32> + return +} + +// Execution counterparts of numerical corner cases in +// GlobalOptimization/test/convert_qdq_to_integer_math_scale_range.mlir. +// Tests ending in _reference or _rewritten document accepted numerical +// differences. The build targets use --gtest_filter to select their respective +// expectations; tests without these suffixes run in both configurations. +// Unfolded storage operands prevent constant folding of the contractions. + +// (1 * 1e20) * (0 * 1e20) = 0. Sequential scaling avoids the overflowing +// scale product that previously caused the integer rewrite to return NaN. +func.func @scale_product_overflow() { + %aq = util.unfoldable_constant dense<1> : tensor<1x1xi8> + %bq = util.unfoldable_constant dense<0> : tensor<1x1xi8> + %scale = arith.constant 1.0e20 : f32 + %ai = tensor.empty() : tensor<1x1xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(m, k) -> (m, k)>, + affine_map<(m, k) -> ()>, + affine_map<(m, k) -> (m, k)>]} + ins(%aq, %scale : tensor<1x1xi8>, f32) + outs(%ai : tensor<1x1xf32>) -> tensor<1x1xf32> + %bi = tensor.empty() : tensor<1x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(k, n) -> (k, n)>, + affine_map<(k, n) -> ()>, + affine_map<(k, n) -> (k, n)>]} + ins(%bq, %scale : tensor<1x1xi8>, f32) + outs(%bi : tensor<1x1xf32>) -> tensor<1x1xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32> + %result = linalg.matmul ins(%a, %b : tensor<1x1xf32>, tensor<1x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + check.expect_eq_const(%result, dense<0.0> : tensor<1x1xf32>) : tensor<1x1xf32> + return +} + +// f16 dequantization rounds 2049 to 2048, so -2048 + 2048 = 0. The integer +// rewrite skips that rounding and returns 1. +func.func private @compute_f16_dequantization_rounding() -> tensor<1x1xf16> { + %aq = util.unfoldable_constant dense<[[-1, 1]]> : tensor<1x2xi8> + %bq = util.unfoldable_constant dense<[[2048], [2049]]> : tensor<2x1xi16> + %scale = arith.constant 1.0 : f32 + %ai = tensor.empty() : tensor<1x2xf16> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(m, k) -> (m, k)>, + affine_map<(m, k) -> ()>, + affine_map<(m, k) -> (m, k)>]} + ins(%aq, %scale : tensor<1x2xi8>, f32) + outs(%ai : tensor<1x2xf16>) -> tensor<1x2xf16> + %bi = tensor.empty() : tensor<2x1xf16> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(k, n) -> (k, n)>, + affine_map<(k, n) -> ()>, + affine_map<(k, n) -> (k, n)>]} + ins(%bq, %scale : tensor<2x1xi16>, f32) + outs(%bi : tensor<2x1xf16>) -> tensor<2x1xf16> + %zero = arith.constant 0.0 : f16 + %empty = tensor.empty() : tensor<1x1xf16> + %init = linalg.fill ins(%zero : f16) outs(%empty : tensor<1x1xf16>) -> tensor<1x1xf16> + %result = linalg.matmul ins(%a, %b : tensor<1x2xf16>, tensor<2x1xf16>) + outs(%init : tensor<1x1xf16>) -> tensor<1x1xf16> + return %result : tensor<1x1xf16> +} + +func.func @f16_dequantization_rounding_reference() { + %result = call @compute_f16_dequantization_rounding() : () -> tensor<1x1xf16> + check.expect_eq_const(%result, dense<0.0> : tensor<1x1xf16>) : tensor<1x1xf16> + return +} + +func.func @f16_dequantization_rounding_rewritten() { + %result = call @compute_f16_dequantization_rounding() : () -> tensor<1x1xf16> + check.expect_eq_const(%result, dense<1.0> : tensor<1x1xf16>) : tensor<1x1xf16> + return +} + +// -0 + (0 * -1) = -0. The integer rewrite discards the init sign and +// returns +0. +func.func private @compute_negative_zero_init() -> tensor<1x1xf32> { + %aq = util.unfoldable_constant dense<0> : tensor<1x1xi8> + %bq = util.unfoldable_constant dense<-1> : tensor<1x1xi8> + %scale = arith.constant 1.0 : f32 + %ai = tensor.empty() : tensor<1x1xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(m, k) -> (m, k)>, + affine_map<(m, k) -> ()>, + affine_map<(m, k) -> (m, k)>]} + ins(%aq, %scale : tensor<1x1xi8>, f32) + outs(%ai : tensor<1x1xf32>) -> tensor<1x1xf32> + %bi = tensor.empty() : tensor<1x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(k, n) -> (k, n)>, + affine_map<(k, n) -> ()>, + affine_map<(k, n) -> (k, n)>]} + ins(%bq, %scale : tensor<1x1xi8>, f32) + outs(%bi : tensor<1x1xf32>) -> tensor<1x1xf32> + %zero = arith.constant -0.0 : f32 + %empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32> + %result = linalg.matmul ins(%a, %b : tensor<1x1xf32>, tensor<1x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + return %result : tensor<1x1xf32> +} + +func.func @negative_zero_init_reference() { + %result = call @compute_negative_zero_init() : () -> tensor<1x1xf32> + // Floating-point equality cannot distinguish -0 from +0; check the bits. + %c0 = arith.constant 0 : index + %value = tensor.extract %result[%c0, %c0] : tensor<1x1xf32> + %bits = arith.bitcast %value : f32 to i32 + %negative_zero_bits = arith.constant -2147483648 : i32 + %is_negative_zero = arith.cmpi eq, %bits, %negative_zero_bits : i32 + check.expect_true(%is_negative_zero) : i1 + return +} + +func.func @negative_zero_init_rewritten() { + %result = call @compute_negative_zero_init() : () -> tensor<1x1xf32> + // Floating-point equality cannot distinguish -0 from +0; check the bits. + %c0 = arith.constant 0 : index + %value = tensor.extract %result[%c0, %c0] : tensor<1x1xf32> + %bits = arith.bitcast %value : f32 to i32 + %positive_zero_bits = arith.constant 0 : i32 + %is_positive_zero = arith.cmpi eq, %bits, %positive_zero_bits : i32 + check.expect_true(%is_positive_zero) : i1 + return +} + +// [0 * infinity, 1 * infinity] dot [1, 1] is NaN. The integer rewrite +// instead computes 1 * infinity and returns infinity. +func.func private @compute_infinite_scale_nan_propagation() -> tensor<1x1xf32> { + %aq = util.unfoldable_constant dense<[[0, 1]]> : tensor<1x2xi8> + %bq = util.unfoldable_constant dense<1> : tensor<2x1xi8> + %sa = arith.constant 0x7F800000 : f32 + %sb = arith.constant 1.0 : f32 + %ai = tensor.empty() : tensor<1x2xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(m, k) -> (m, k)>, + affine_map<(m, k) -> ()>, + affine_map<(m, k) -> (m, k)>]} + ins(%aq, %sa : tensor<1x2xi8>, f32) + outs(%ai : tensor<1x2xf32>) -> tensor<1x2xf32> + %bi = tensor.empty() : tensor<2x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(k, n) -> (k, n)>, + affine_map<(k, n) -> ()>, + affine_map<(k, n) -> (k, n)>]} + ins(%bq, %sb : tensor<2x1xi8>, f32) + outs(%bi : tensor<2x1xf32>) -> tensor<2x1xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32> + %result = linalg.matmul ins(%a, %b : tensor<1x2xf32>, tensor<2x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + return %result : tensor<1x1xf32> +} + +func.func @infinite_scale_nan_propagation_reference() { + %result = call @compute_infinite_scale_nan_propagation() : () -> tensor<1x1xf32> + // Compare unordered rather than requiring a particular NaN payload. + %c0 = arith.constant 0 : index + %value = tensor.extract %result[%c0, %c0] : tensor<1x1xf32> + %is_nan = arith.cmpf uno, %value, %value : f32 + check.expect_true(%is_nan) : i1 + return +} + +func.func @infinite_scale_nan_propagation_rewritten() { + %result = call @compute_infinite_scale_nan_propagation() : () -> tensor<1x1xf32> + check.expect_eq_const(%result, dense<0x7F800000> : tensor<1x1xf32>) : tensor<1x1xf32> + return +} + +// (32767 * 1e-23)^2 is approximately 1.07368e-37. Sequential scaling avoids +// the underflowing scale product that previously erased the result. +func.func @scale_product_underflow() { + %q = util.unfoldable_constant dense<32767> : tensor<1x1xi16> + %scale = arith.constant 1.0e-23 : f32 + %ai = tensor.empty() : tensor<1x1xf32> + %a = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(m, k) -> (m, k)>, + affine_map<(m, k) -> ()>, + affine_map<(m, k) -> (m, k)>]} + ins(%q, %scale : tensor<1x1xi16>, f32) + outs(%ai : tensor<1x1xf32>) -> tensor<1x1xf32> + %bi = tensor.empty() : tensor<1x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(k, n) -> (k, n)>, + affine_map<(k, n) -> ()>, + affine_map<(k, n) -> (k, n)>]} + ins(%q, %scale : tensor<1x1xi16>, f32) + outs(%bi : tensor<1x1xf32>) -> tensor<1x1xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32> + %result = linalg.matmul ins(%a, %b : tensor<1x1xf32>, tensor<1x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + // The dequantized operands and their product are all normal f32 values. + // Use exact equality: an absolute tolerance could incorrectly accept zero. + check.expect_eq_const(%result, dense<0x02122428> : tensor<1x1xf32>) : tensor<1x1xf32> + return +} + +// Both reference operands are finite, but D * 1e30 overflows in the epilogue. +func.func private @compute_opposing_scales_large_lhs() -> tensor<1x1xf32> { + %aq = util.unfoldable_constant dense<-1> : tensor<1x32768xi8> + %bq = util.unfoldable_constant dense<127> : tensor<32768x1xi8> + %sa = arith.constant 1.0e30 : f32 + %sb = arith.constant 1.0e-30 : f32 + %ai = tensor.empty() : tensor<1x32768xf32> + %a = iree_linalg_ext.dequantize_affine + {input_unsigned, indexing_maps = [affine_map<(m, k) -> (m, k)>, + affine_map<(m, k) -> ()>, + affine_map<(m, k) -> (m, k)>]} + ins(%aq, %sa : tensor<1x32768xi8>, f32) + outs(%ai : tensor<1x32768xf32>) -> tensor<1x32768xf32> + %bi = tensor.empty() : tensor<32768x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(k, n) -> (k, n)>, + affine_map<(k, n) -> ()>, + affine_map<(k, n) -> (k, n)>]} + ins(%bq, %sb : tensor<32768x1xi8>, f32) + outs(%bi : tensor<32768x1xf32>) -> tensor<32768x1xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32> + %result = linalg.matmul ins(%a, %b : tensor<1x32768xf32>, tensor<32768x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + return %result : tensor<1x1xf32> +} + +func.func @opposing_scales_large_lhs_reference() { + %result = call @compute_opposing_scales_large_lhs() : () -> tensor<1x1xf32> + // D = 32768 * 255 * 127 = 1061191680, within signed i32. + // Allow reduction-order rounding while requiring a finite result. + check.expect_almost_eq_const(%result, dense<1061191680.0> : tensor<1x1xf32>, atol 0.0, rtol 1.0e-4) : tensor<1x1xf32> + return +} + +func.func @opposing_scales_large_lhs_rewritten() { + %result = call @compute_opposing_scales_large_lhs() : () -> tensor<1x1xf32> + // 0x7F800000 -> INFINITY + check.expect_eq_const(%result, dense<0x7F800000> : tensor<1x1xf32>) : tensor<1x1xf32> + return +} + +// Reversing the scales avoids epilogue overflow. Floating-point reduction +// rounding differs; allow 1e-4 relative error against the exact integer sum. +func.func private @compute_opposing_scales_large_rhs() -> tensor<1x1xf32> { + %aq = util.unfoldable_constant dense<-1> : tensor<1x32768xi8> + %bq = util.unfoldable_constant dense<127> : tensor<32768x1xi8> + %sa = arith.constant 1.0e-30 : f32 + %sb = arith.constant 1.0e30 : f32 + %ai = tensor.empty() : tensor<1x32768xf32> + %a = iree_linalg_ext.dequantize_affine + {input_unsigned, indexing_maps = [affine_map<(m, k) -> (m, k)>, + affine_map<(m, k) -> ()>, + affine_map<(m, k) -> (m, k)>]} + ins(%aq, %sa : tensor<1x32768xi8>, f32) + outs(%ai : tensor<1x32768xf32>) -> tensor<1x32768xf32> + %bi = tensor.empty() : tensor<32768x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(k, n) -> (k, n)>, + affine_map<(k, n) -> ()>, + affine_map<(k, n) -> (k, n)>]} + ins(%bq, %sb : tensor<32768x1xi8>, f32) + outs(%bi : tensor<32768x1xf32>) -> tensor<32768x1xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32> + %result = linalg.matmul ins(%a, %b : tensor<1x32768xf32>, tensor<32768x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + return %result : tensor<1x1xf32> +} + +func.func @opposing_scales_large_rhs_reference() { + %result = call @compute_opposing_scales_large_rhs() : () -> tensor<1x1xf32> + // D = 32768 * 255 * 127 = 1061191680, within signed i32. + // Allow reduction-order rounding while requiring a finite result. + check.expect_almost_eq_const(%result, dense<1061191680.0> : tensor<1x1xf32>, atol 0.0, rtol 1.0e-4) : tensor<1x1xf32> + return +} + +func.func @opposing_scales_large_rhs_rewritten() { + %result = call @compute_opposing_scales_large_rhs() : () -> tensor<1x1xf32> + check.expect_eq_const(%result, dense<1061191680.0> : tensor<1x1xf32>) : tensor<1x1xf32> + return +} + +// All scales are finite, but 127 * 1e37 overflows during reference +// dequantization. Integer reduction bypasses that overflow. +func.func private @compute_finite_scale_dequantization_overflow() -> tensor<1x1xf32> { + %aq = util.unfoldable_constant dense<-1> : tensor<1x32768xi8> + %bq = util.unfoldable_constant dense<127> : tensor<32768x1xi8> + %sa = arith.constant 1.0e-37 : f32 + %sb = arith.constant 1.0e37 : f32 + %ai = tensor.empty() : tensor<1x32768xf32> + %a = iree_linalg_ext.dequantize_affine + {input_unsigned, indexing_maps = [affine_map<(m, k) -> (m, k)>, + affine_map<(m, k) -> ()>, + affine_map<(m, k) -> (m, k)>]} + ins(%aq, %sa : tensor<1x32768xi8>, f32) + outs(%ai : tensor<1x32768xf32>) -> tensor<1x32768xf32> + %bi = tensor.empty() : tensor<32768x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(k, n) -> (k, n)>, + affine_map<(k, n) -> ()>, + affine_map<(k, n) -> (k, n)>]} + ins(%bq, %sb : tensor<32768x1xi8>, f32) + outs(%bi : tensor<32768x1xf32>) -> tensor<32768x1xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32> + %result = linalg.matmul ins(%a, %b : tensor<1x32768xf32>, tensor<32768x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + return %result : tensor<1x1xf32> +} + +func.func @finite_scale_dequantization_overflow_reference() { + %result = call @compute_finite_scale_dequantization_overflow() : () -> tensor<1x1xf32> + check.expect_eq_const(%result, dense<0x7F800000> : tensor<1x1xf32>) : tensor<1x1xf32> + return +} + +func.func @finite_scale_dequantization_overflow_rewritten() { + %result = call @compute_finite_scale_dequantization_overflow() : () -> tensor<1x1xf32> + check.expect_eq_const(%result, dense<1061191616.0> : tensor<1x1xf32>) : tensor<1x1xf32> + return +} + +// With zero activations, the same overflowing reference weights yield NaN +// from 0 * infinity. Integer reduction instead yields zero. +func.func private @compute_finite_scale_dequantization_nan() -> tensor<1x1xf32> { + %aq = util.unfoldable_constant dense<0> : tensor<1x32768xi8> + %bq = util.unfoldable_constant dense<127> : tensor<32768x1xi8> + %sa = arith.constant 1.0e-37 : f32 + %sb = arith.constant 1.0e37 : f32 + %ai = tensor.empty() : tensor<1x32768xf32> + %a = iree_linalg_ext.dequantize_affine + {input_unsigned, indexing_maps = [affine_map<(m, k) -> (m, k)>, + affine_map<(m, k) -> ()>, + affine_map<(m, k) -> (m, k)>]} + ins(%aq, %sa : tensor<1x32768xi8>, f32) + outs(%ai : tensor<1x32768xf32>) -> tensor<1x32768xf32> + %bi = tensor.empty() : tensor<32768x1xf32> + %b = iree_linalg_ext.dequantize_affine + {indexing_maps = [affine_map<(k, n) -> (k, n)>, + affine_map<(k, n) -> ()>, + affine_map<(k, n) -> (k, n)>]} + ins(%bq, %sb : tensor<32768x1xi8>, f32) + outs(%bi : tensor<32768x1xf32>) -> tensor<32768x1xf32> + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32> + %result = linalg.matmul ins(%a, %b : tensor<1x32768xf32>, tensor<32768x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + return %result : tensor<1x1xf32> +} + +func.func @finite_scale_dequantization_nan_reference() { + %result = call @compute_finite_scale_dequantization_nan() : () -> tensor<1x1xf32> + // Compare unordered without requiring a specific NaN payload. + %c0 = arith.constant 0 : index + %value = tensor.extract %result[%c0, %c0] : tensor<1x1xf32> + %is_nan = arith.cmpf uno, %value, %value : f32 + check.expect_true(%is_nan) : i1 + return +} + +func.func @finite_scale_dequantization_nan_rewritten() { + %result = call @compute_finite_scale_dequantization_nan() : () -> tensor<1x1xf32> + check.expect_eq_const(%result, dense<0.0> : tensor<1x1xf32>) : tensor<1x1xf32> + return +}