| // Copyright 2023 The IREE Authors |
| // |
| // Licensed under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| |
| #include "iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.h" |
| #include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtDialect.h" |
| #include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtInterfaces.h" |
| #include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h" |
| #include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h" |
| #include "iree/compiler/GlobalOptimization/Passes.h" |
| #include "iree/compiler/GlobalOptimization/Utils.h" |
| #include "llvm/ADT/STLExtras.h" |
| #include "llvm/Support/Casting.h" |
| #include "mlir/Dialect/Affine/IR/AffineOps.h" |
| #include "mlir/Dialect/Arith/IR/Arith.h" |
| #include "mlir/Dialect/Linalg/IR/Linalg.h" |
| #include "mlir/Dialect/Linalg/Transforms/Transforms.h" |
| #include "mlir/Dialect/Linalg/Utils/Utils.h" |
| #include "mlir/Dialect/Math/IR/Math.h" |
| #include "mlir/Dialect/Tensor/IR/Tensor.h" |
| #include "mlir/Dialect/Utils/StaticValueUtils.h" |
| #include "mlir/IR/BuiltinAttributes.h" |
| #include "mlir/IR/PatternMatch.h" |
| #include "mlir/IR/TypeUtilities.h" |
| #include "mlir/Pass/Pass.h" |
| #include "mlir/Transforms/GreedyPatternRewriteDriver.h" |
| |
| using namespace mlir; |
| |
| namespace mlir::iree_compiler::GlobalOptimization { |
| |
| #define GEN_PASS_DEF_RAISESPECIALOPSPASS |
| #include "iree/compiler/GlobalOptimization/Passes.h.inc" |
| |
| namespace { |
| |
| //===----------------------------------------------------------------------===// |
| // Slice Raising |
| //===----------------------------------------------------------------------===// |
| |
| /// Matches a linalg.generic operation reading data from a tensor `source` using |
| /// tensor.extract, and raises the `source` tensor to an input of the linalg |
| /// operation. |
| static FailureOr<linalg::GenericOp> |
| raiseTensorExtractToInput(linalg::GenericOp linalgOp, RewriterBase &rewriter) { |
| if (!linalgOp.hasPureTensorSemantics()) { |
| return failure(); |
| } |
| if (!isElementwise(linalgOp)) { |
| return failure(); |
| } |
| if (!llvm::hasSingleElement(linalgOp.getResults())) { |
| return failure(); |
| } |
| |
| // Find a tensor.extract op in the linalgOp body. |
| auto extractOps = linalgOp.getBody()->getOps<tensor::ExtractOp>(); |
| if (!llvm::hasSingleElement(extractOps)) { |
| return failure(); |
| } |
| tensor::ExtractOp extractOp = *extractOps.begin(); |
| auto resultType = dyn_cast<TensorType>(linalgOp.getResult(0).getType()); |
| if (!resultType) { |
| return failure(); |
| } |
| |
| ArrayRef<int64_t> sourceShape = extractOp.getTensor().getType().getShape(); |
| ArrayRef<int64_t> resultShape = resultType.getShape(); |
| |
| // Raise the tensor.extract op to an input. |
| SmallVector<AffineExpr> exprs; |
| for (auto [idx, indexValue] : llvm::enumerate(extractOp.getIndices())) { |
| // For raising, the indexing value must be one of the following: |
| // 1. A constant value. |
| // 2. A linalg.index. |
| |
| // 1. Indexing value is a constant. |
| APInt constantIndex; |
| if (matchPattern(indexValue, m_ConstantInt(&constantIndex))) { |
| // Restrict to cases where the constant is 0. This is because handling |
| // constants other than 0 in indexing map, may cause problems in the |
| // lowering pipeline later. |
| if (constantIndex.getLimitedValue() != 0) { |
| return failure(); |
| } |
| exprs.push_back(getAffineConstantExpr(0, rewriter.getContext())); |
| continue; |
| } |
| // 2. The indexing value is a linalg.index. |
| if (auto indexOp = indexValue.getDefiningOp<linalg::IndexOp>()) { |
| // Make sure that for this index, the size of the input and output |
| // match and are not dynamic. We need this to maintain the op to be |
| // elementwise. |
| // TODO: This restriction can be relaxed by adding a extract_slice op |
| // on the `source` tensor. This is not same as raising the whole |
| // operation to an extract_slice, as there can be permutations and |
| // projections involved. |
| if (ShapedType::isDynamic(sourceShape[idx]) || |
| ShapedType::isDynamic(resultShape[indexOp.getDim()]) || |
| sourceShape[idx] != resultShape[indexOp.getDim()]) { |
| return failure(); |
| } |
| exprs.push_back( |
| getAffineDimExpr(indexOp.getDim(), rewriter.getContext())); |
| continue; |
| } |
| return failure(); |
| } |
| AffineMap indexingMap = AffineMap::get( |
| /*dimCount=*/linalgOp.getNumLoops(), |
| /*symbolCount=*/0, exprs, rewriter.getContext()); |
| |
| // Replace the linalgOp with a new linalgOp where the source tensor is |
| // an input with the indexing map. |
| SmallVector<Value> newInputs = linalgOp.getInputs(); |
| newInputs.insert(newInputs.begin(), extractOp.getTensor()); |
| SmallVector<Attribute> newIndexingMaps; |
| newIndexingMaps.push_back(AffineMapAttr::get(indexingMap)); |
| for (AffineMap map : linalgOp.getIndexingMapsArray()) { |
| newIndexingMaps.push_back(AffineMapAttr::get(map)); |
| } |
| |
| auto bodyBuilder = [&](OpBuilder &builder, Location loc, ValueRange args) { |
| // Create an IR mapping from old block arguments to new ones. |
| IRMapping mapper; |
| ArrayRef<BlockArgument> oldArgs = linalgOp.getBody()->getArguments(); |
| // Map i^th old argument to (i + 1)^th new argument. |
| for (unsigned i = 0; i < oldArgs.size(); ++i) { |
| mapper.map(oldArgs[i], args[i + 1]); |
| } |
| // Clone the body of the linalgOp. |
| for (Operation &op : linalgOp.getBody()->getOperations()) { |
| // Replace the extractOp with the first block argument. |
| if (&op == extractOp) { |
| mapper.map(op.getResult(0), args[0]); |
| } else { |
| builder.clone(op, mapper); |
| } |
| } |
| }; |
| |
| linalg::GenericOp newLinalgOp = linalg::GenericOp::create( |
| rewriter, linalgOp.getLoc(), linalgOp.getResultTypes(), newInputs, |
| linalgOp.getOutputs(), |
| ArrayAttr::get(linalgOp->getContext(), newIndexingMaps), |
| linalgOp.getIteratorTypesAttr(), linalgOp.getDocAttr(), |
| linalgOp.getLibraryCallAttr(), bodyBuilder); |
| |
| return newLinalgOp; |
| } |
| |
| /// Given a linalg.generic operation, and input/output tensors with their |
| /// indexing maps, tries to raise the operation to a tensor.extract_slice |
| /// operation. The tensor.extract_slice produced can be rank reducing. |
| static FailureOr<tensor::ExtractSliceOp> |
| tryRaiseToExtractSlice(AffineMap inputIndexingMap, AffineMap outputIndexingMap, |
| Value input, Value output, linalg::GenericOp linalgOp, |
| RewriterBase &rewriter) { |
| // Output shape must be smaller than input shape. |
| if (outputIndexingMap.getNumResults() >= inputIndexingMap.getNumResults()) { |
| return failure(); |
| } |
| // Output map should be identity. |
| if (!outputIndexingMap.isIdentity()) { |
| return failure(); |
| } |
| |
| auto outType = dyn_cast<RankedTensorType>(output.getType()); |
| if (!outType) { |
| return failure(); |
| } |
| ArrayRef<int64_t> outShape = outType.getShape(); |
| int64_t outRank = outShape.size(); |
| |
| // Try to match each output dimension to an input dimension, in order. |
| // If we find a constant access, we assume that dimension is supposed to be |
| // rank reduced. |
| // TODO: Support cases where the constant access matches the output dimension. |
| SmallVector<OpFoldResult> offsets; |
| SmallVector<OpFoldResult> sizes; |
| IntegerAttr zero = rewriter.getI64IntegerAttr(0); |
| IntegerAttr one = rewriter.getI64IntegerAttr(1); |
| unsigned currOutDim = 0; |
| for (auto [idx, expr] : llvm::enumerate(inputIndexingMap.getResults())) { |
| // Assume that the constant access is a rank reducing access. |
| if (isa<AffineConstantExpr>(expr)) { |
| IntegerAttr constIdx = |
| rewriter.getI64IntegerAttr(cast<AffineConstantExpr>(expr).getValue()); |
| offsets.push_back(constIdx); |
| sizes.push_back(one); |
| continue; |
| } |
| // Check if the input dimension matches the current output dimension. |
| if (currOutDim < outRank && |
| expr == outputIndexingMap.getResult(currOutDim)) { |
| offsets.push_back(zero); |
| // Get the dim size from the output tensor. |
| if (ShapedType::isDynamic(outShape[currOutDim])) { |
| auto dim = tensor::DimOp::create(rewriter, linalgOp.getLoc(), output, |
| currOutDim); |
| sizes.push_back(dim.getResult()); |
| } else { |
| sizes.push_back(rewriter.getI64IntegerAttr(outShape[currOutDim])); |
| } |
| ++currOutDim; |
| continue; |
| } |
| // Unknown access, fail. |
| return failure(); |
| } |
| |
| // All output dimensions did not match an input dimension. |
| if (currOutDim != outputIndexingMap.getNumResults()) { |
| return failure(); |
| } |
| |
| // We only support dim expr or a constant expr on the input map, so strides |
| // will always be 1. |
| SmallVector<OpFoldResult> strides(inputIndexingMap.getNumResults(), one); |
| |
| return tensor::ExtractSliceOp::create(rewriter, linalgOp.getLoc(), outType, |
| input, offsets, sizes, strides); |
| } |
| |
| /// Matches a linalg.generic operation with a single input and init output |
| /// tensor, and tries to raise it to a view-like operation on the input tensor. |
| static FailureOr<Operation *> tryRaiseToView(linalg::GenericOp linalgOp, |
| RewriterBase &rewriter) { |
| if (!linalgOp.hasPureTensorSemantics()) { |
| return failure(); |
| } |
| |
| // Assume there is only 1 input, and 1 init tensor. |
| if (linalgOp.getNumDpsInputs() != 1 || linalgOp.getNumDpsInits() != 1) { |
| return failure(); |
| } |
| OpOperand *inputOperand = linalgOp.getDpsInputOperand(0); |
| OpOperand *outputOperand = linalgOp.getDpsInitOperand(0); |
| |
| // Check if linalg.yield yields a block argument. |
| auto yieldOp = dyn_cast<linalg::YieldOp>(linalgOp.getBody()->getTerminator()); |
| if (!yieldOp) { |
| return failure(); |
| } |
| auto blockArg = dyn_cast<BlockArgument>(yieldOp.getOperand(0)); |
| if (!blockArg) { |
| return failure(); |
| } |
| // Check if the block argument is an argument of the linalgOp. |
| if (blockArg.getOwner() != linalgOp.getBody()) { |
| return failure(); |
| } |
| // Check that the block argument corresponds to the input. |
| if (blockArg.getArgNumber() != 0) { |
| return failure(); |
| } |
| |
| Value input = inputOperand->get(); |
| Value output = outputOperand->get(); |
| AffineMap inputIndexingMap = linalgOp.getMatchingIndexingMap(inputOperand); |
| AffineMap outputIndexingMap = linalgOp.getMatchingIndexingMap(outputOperand); |
| |
| // Try raising to tensor.collapse_shape. |
| return tryRaiseToExtractSlice(inputIndexingMap, outputIndexingMap, input, |
| output, linalgOp, rewriter); |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // Named GEMM-like extensions |
| //===----------------------------------------------------------------------===// |
| |
| template <typename OpTy> |
| class NamedImplicitCastOpConversion : public OpInterfaceRewritePattern<OpTy> { |
| public: |
| using OpInterfaceRewritePattern<OpTy>::OpInterfaceRewritePattern; |
| NamedImplicitCastOpConversion(MLIRContext *ctx, PatternBenefit b = 1) |
| : OpInterfaceRewritePattern<OpTy>(ctx, b) {} |
| |
| LogicalResult matchAndRewrite(OpTy namedOp, |
| PatternRewriter &rewriter) const override { |
| if (!IREE::Flow::isNonNullAndOutsideDispatch(namedOp)) { |
| return failure(); |
| } |
| |
| auto linalgOp = dyn_cast<linalg::LinalgOp>(namedOp.getOperation()); |
| if (!linalgOp || !linalgOp.hasPureTensorSemantics()) { |
| return rewriter.notifyMatchFailure( |
| namedOp, "expected an op with pure tensor semantics"); |
| } |
| |
| // Look for a producer of the given operand that does an elementwise extend |
| // and replace the operand with the source of the elementwise producer. |
| // Returns true if the operand was updated to inform the pattern rewriter |
| // of a change. |
| Type outElementType = getElementTypeOrSelf(namedOp->getResultTypes()[0]); |
| bool didChangeOperand = false; |
| { |
| OpBuilder::InsertionGuard guard(rewriter); |
| Block *block = &namedOp->getRegion(0).front(); |
| rewriter.setInsertionPointToStart(block); |
| auto replaceOperandWithTypeCast = [&](OpOperand &operand) { |
| // If the op already has implicit casting semantics for this operand, |
| // do not fuse. |
| if (getElementTypeOrSelf(operand.get().getType()) != outElementType) { |
| return false; |
| } |
| auto producer = operand.get().getDefiningOp<linalg::GenericOp>(); |
| if (!producer) { |
| return false; |
| } |
| |
| if (!llvm::all_of(producer.getIndexingMapsArray(), |
| [](AffineMap map) { return map.isIdentity(); })) { |
| return false; |
| } |
| |
| std::optional<CastOpInterface> castOp = |
| getDefiningNonI1ExtendingCastOp(operand.get()); |
| if (!castOp) { |
| return false; |
| } |
| // We only handle arith.extf/exti because truncating operations like |
| // `arith.truncf` should not be fused with consumers; it would be |
| // preferred to fuse those with producers (and the consumer fusion is |
| // arguably the less canonical form). |
| auto canFoldCast = [&]() { |
| if (isa<arith::ExtFOp>(*castOp)) { |
| return true; |
| } |
| // Signed operations can only be folded with (implicitly) signed |
| // linalg named ops |
| if (isa<arith::ExtSIOp>(*castOp)) { |
| if (auto matmul = |
| dyn_cast<linalg::MatmulOp>(namedOp.getOperation())) { |
| return matmul.getCast() != linalg::TypeFn::cast_unsigned; |
| } |
| return !isa<linalg::PoolingNhwcMaxUnsignedOp, |
| linalg::PoolingNhwcMinUnsignedOp, |
| linalg::PoolingNwcMaxUnsignedOp, |
| linalg::PoolingNwcMinUnsignedOp>(namedOp); |
| } |
| return false; |
| }; |
| if (!canFoldCast()) { |
| return false; |
| } |
| Type producerElementType = getElementTypeOrSelf( |
| producer.getDpsInputOperand(0)->get().getType()); |
| int64_t operandNumber = operand.getOperandNumber(); |
| // Set the operand to the linalg op to the smaller one. |
| namedOp->setOperand(operandNumber, producer->getOperand(0)); |
| |
| // Insert a new block argument into the body of the named op with the |
| // correct type. |
| Value blockArg = block->insertArgument( |
| operandNumber, producerElementType, namedOp.getLoc()); |
| // Create the extf/extsi. |
| IRMapping mapping; |
| mapping.map(castOp.value()->getOperand(0), blockArg); |
| Value ext = rewriter.clone(*castOp.value(), mapping)->getResult(0); |
| |
| // Replace uses of the old argument with the extended value. |
| rewriter.replaceAllUsesWith(block->getArgument(operandNumber + 1), ext); |
| // Erase the old argument. |
| block->eraseArgument(operandNumber + 1); |
| return true; |
| }; |
| |
| didChangeOperand = replaceOperandWithTypeCast(namedOp->getOpOperand(0)); |
| didChangeOperand |= replaceOperandWithTypeCast(namedOp->getOpOperand(1)); |
| } |
| return success(didChangeOperand); |
| } |
| }; |
| |
| //===----------------------------------------------------------------------===// |
| // Softmax Raising |
| //===----------------------------------------------------------------------===// |
| |
| // Recognizing a numerically-stabilized softmax means walking a small graph of |
| // `linalg.generic` ops. The helpers below match the individual nodes of that |
| // graph directly, replacing the generic `transform_ext` matcher framework that |
| // used to live in iree-dialects. The matched dataflow, working back from the |
| // root op, is: |
| // |
| // max = reduce_max(src) (reduction over the innermost dim) |
| // sub = src - broadcast(max) |
| // exp = exp(sub) |
| // sum = reduce_add(exp) (reduction over the innermost dim) |
| // root = exp * reciprocal(sum) OR exp / broadcast(sum) |
| // |
| // where each broadcast is either implicit (the consumer indexing map drops the |
| // innermost dim) or explicit (a pass-through `linalg.generic` that broadcasts |
| // along the innermost dim). |
| |
| // Returns true if every iterator of `op` is parallel. |
| static bool isAllParallel(linalg::GenericOp op) { |
| return llvm::all_of(op.getIteratorTypesArray(), [](utils::IteratorType t) { |
| return t == utils::IteratorType::parallel; |
| }); |
| } |
| |
| // Returns true if `operand`'s indexing map is the identity. |
| static bool isIdentityOperand(linalg::GenericOp op, OpOperand *operand) { |
| return op.getMatchingIndexingMap(operand).isIdentity(); |
| } |
| |
| // Returns true if `operand`'s indexing map is the projection that drops the |
| // innermost (last) loop dim, e.g. (d0, d1, d2) -> (d0, d1). |
| static bool dropsInnermostDim(linalg::GenericOp op, OpOperand *operand) { |
| AffineMap map = op.getMatchingIndexingMap(operand); |
| unsigned numLoops = op.getNumLoops(); |
| if (!map.isProjectedPermutation() || map.getNumResults() + 1 != numLoops) { |
| return false; |
| } |
| for (unsigned i = 0, e = map.getNumResults(); i < e; ++i) { |
| if (map.getDimPosition(i) != i) { |
| return false; |
| } |
| } |
| return true; |
| } |
| |
| // Returns the single compute op of `op`'s body when the body is exactly that op |
| // followed by a `linalg.yield` of its single result; otherwise nullptr. |
| static Operation *getSingleComputeOp(linalg::GenericOp op) { |
| Block *body = op.getBlock(); |
| if (body->getOperations().size() != 2) { |
| return nullptr; |
| } |
| Operation *computeOp = &body->front(); |
| auto yieldOp = cast<linalg::YieldOp>(body->getTerminator()); |
| if (computeOp->getNumResults() != 1 || yieldOp.getNumOperands() != 1 || |
| yieldOp.getOperand(0).getDefiningOp() != computeOp) { |
| return nullptr; |
| } |
| return computeOp; |
| } |
| |
| // Returns true if `v` references block argument number `argNumber` of `block`. |
| static bool isBlockArg(Value v, Block *block, unsigned argNumber) { |
| auto arg = dyn_cast<BlockArgument>(v); |
| return arg && arg.getOwner() == block && arg.getArgNumber() == argNumber; |
| } |
| |
| // Returns true if `op`'s body is a single `OpTy` whose operands are the leading |
| // block arguments in order. When `commutative`, a two-operand body may also use |
| // the arguments in swapped order. |
| template <typename OpTy> |
| static bool hasSingleOpBody(linalg::GenericOp op, bool commutative = false) { |
| Operation *computeOp = getSingleComputeOp(op); |
| if (!isa_and_nonnull<OpTy>(computeOp)) { |
| return false; |
| } |
| Block *body = op.getBlock(); |
| if (commutative && computeOp->getNumOperands() == 2) { |
| return (isBlockArg(computeOp->getOperand(0), body, 0) && |
| isBlockArg(computeOp->getOperand(1), body, 1)) || |
| (isBlockArg(computeOp->getOperand(0), body, 1) && |
| isBlockArg(computeOp->getOperand(1), body, 0)); |
| } |
| for (auto [index, operand] : llvm::enumerate(computeOp->getOperands())) { |
| if (!isBlockArg(operand, body, index)) { |
| return false; |
| } |
| } |
| return true; |
| } |
| |
| // Returns true if `op`'s body just yields its (single) block argument, i.e. it |
| // is a pure broadcast/transpose of the input. |
| static bool isPassThroughBody(linalg::GenericOp op) { |
| Block *body = op.getBlock(); |
| if (body->getOperations().size() != 1) { |
| return false; |
| } |
| auto yieldOp = cast<linalg::YieldOp>(body->getTerminator()); |
| for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) { |
| if (!isBlockArg(operand, body, index)) { |
| return false; |
| } |
| } |
| return true; |
| } |
| |
| // Returns true if `v` is a float constant satisfying `pred`. |
| static bool isFloatConstant(Value v, llvm::function_ref<bool(APFloat)> pred) { |
| auto cstOp = v.getDefiningOp<arith::ConstantFloatOp>(); |
| return cstOp && pred(cstOp.value()); |
| } |
| |
| // Returns true if `v` is a `linalg.fill` whose scalar value is a float constant |
| // satisfying `pred`. |
| static bool isFilledWith(Value v, llvm::function_ref<bool(APFloat)> pred) { |
| auto fillOp = v.getDefiningOp<linalg::FillOp>(); |
| return fillOp && isFloatConstant(fillOp.getInputs()[0], pred); |
| } |
| |
| // Matches a single-input `linalg.generic` reduction over the innermost dim with |
| // a commutative `OpTy` body and an init produced by a `linalg.fill` of a |
| // constant satisfying `fillPred`. Returns the reduced value on success. |
| template <typename OpTy> |
| static Value |
| matchInnermostReduction(Value v, llvm::function_ref<bool(APFloat)> fillPred) { |
| auto op = v.getDefiningOp<linalg::GenericOp>(); |
| if (!op || op.getNumDpsInputs() != 1 || op.getNumDpsInits() != 1) { |
| return {}; |
| } |
| SmallVector<utils::IteratorType> iterators = op.getIteratorTypesArray(); |
| if (iterators.empty() || iterators.back() != utils::IteratorType::reduction) { |
| return {}; |
| } |
| for (unsigned i = 0, e = iterators.size() - 1; i < e; ++i) { |
| if (iterators[i] != utils::IteratorType::parallel) { |
| return {}; |
| } |
| } |
| OpOperand *input = op.getDpsInputOperand(0); |
| OpOperand *init = op.getDpsInitOperand(0); |
| if (!isIdentityOperand(op, input) || !dropsInnermostDim(op, init) || |
| !hasSingleOpBody<OpTy>(op, /*commutative=*/true) || |
| !isFilledWith(init->get(), fillPred)) { |
| return {}; |
| } |
| return input->get(); |
| } |
| |
| // Matches a single-input all-parallel `linalg.generic` with identity-mapped |
| // input and output and an `OpTy` body. Returns the input value on success. |
| template <typename OpTy> |
| static Value matchUnaryElementwise(Value v) { |
| auto op = v.getDefiningOp<linalg::GenericOp>(); |
| if (!op || op.getNumDpsInputs() != 1 || op.getNumDpsInits() != 1 || |
| !isAllParallel(op)) { |
| return {}; |
| } |
| if (!isIdentityOperand(op, op.getDpsInputOperand(0)) || |
| !isIdentityOperand(op, op.getDpsInitOperand(0)) || |
| !hasSingleOpBody<OpTy>(op)) { |
| return {}; |
| } |
| return op.getDpsInputOperand(0)->get(); |
| } |
| |
| // Matches a reciprocal `linalg.generic` (body `divf 1.0, %arg0`) that is |
| // single-input, all-parallel and identity-mapped. Returns the input value. |
| static Value matchReciprocal(Value v) { |
| auto op = v.getDefiningOp<linalg::GenericOp>(); |
| if (!op || op.getNumDpsInputs() != 1 || op.getNumDpsInits() != 1 || |
| !isAllParallel(op)) { |
| return {}; |
| } |
| if (!isIdentityOperand(op, op.getDpsInputOperand(0)) || |
| !isIdentityOperand(op, op.getDpsInitOperand(0))) { |
| return {}; |
| } |
| auto divOp = dyn_cast_or_null<arith::DivFOp>(getSingleComputeOp(op)); |
| if (!divOp || |
| !isFloatConstant(divOp.getOperand(0), |
| [](APFloat f) { return f.convertToDouble() == 1.0; }) || |
| !isBlockArg(divOp.getOperand(1), op.getBlock(), 0)) { |
| return {}; |
| } |
| return op.getDpsInputOperand(0)->get(); |
| } |
| |
| // Resolves `operand` of `consumer`, which is expected to carry the broadcast of |
| // an innermost-reduced value, to that pre-broadcast value. Handles both the |
| // implicit broadcast (the operand's map drops the innermost dim) and the |
| // explicit broadcast (the operand is produced by a pass-through generic that |
| // broadcasts along the innermost dim). |
| static Value resolveBroadcastedReduction(linalg::GenericOp consumer, |
| OpOperand *operand) { |
| if (dropsInnermostDim(consumer, operand)) { |
| return operand->get(); |
| } |
| if (!isIdentityOperand(consumer, operand)) { |
| return {}; |
| } |
| auto bcast = operand->get().getDefiningOp<linalg::GenericOp>(); |
| if (!bcast || bcast.getNumDpsInputs() != 1 || bcast.getNumDpsInits() != 1 || |
| !isAllParallel(bcast) || !isPassThroughBody(bcast) || |
| !dropsInnermostDim(bcast, bcast.getDpsInputOperand(0)) || |
| !isIdentityOperand(bcast, bcast.getDpsInitOperand(0))) { |
| return {}; |
| } |
| return bcast.getDpsInputOperand(0)->get(); |
| } |
| |
| // Recognizes the softmax graph rooted at `rootOp` and returns the softmax |
| // source value on success. |
| static FailureOr<Value> matchSoftmax(linalg::LinalgOp rootOp) { |
| auto root = dyn_cast<linalg::GenericOp>(rootOp.getOperation()); |
| if (!root || root.getNumDpsInputs() != 2 || root.getNumDpsInits() != 1 || |
| !isAllParallel(root)) { |
| return failure(); |
| } |
| // The numerator (exp(...)) always flows in as input 0 with an identity map. |
| OpOperand *numerator = root.getDpsInputOperand(0); |
| OpOperand *denominator = root.getDpsInputOperand(1); |
| if (!isIdentityOperand(root, numerator)) { |
| return failure(); |
| } |
| |
| // The root normalizes either as `exp * reciprocal(sum)` or `exp / sum`. |
| Value sumValue; |
| if (hasSingleOpBody<arith::MulFOp>(root, /*commutative=*/true)) { |
| if (!dropsInnermostDim(root, denominator)) { |
| return failure(); |
| } |
| sumValue = matchReciprocal(denominator->get()); |
| } else if (hasSingleOpBody<arith::DivFOp>(root)) { |
| sumValue = resolveBroadcastedReduction(root, denominator); |
| } |
| if (!sumValue) { |
| return failure(); |
| } |
| |
| // exp = exp(sub); sum = reduce_add(exp). |
| Value expValue = numerator->get(); |
| Value subValue = matchUnaryElementwise<math::ExpOp>(expValue); |
| Value summedValue = matchInnermostReduction<arith::AddFOp>( |
| sumValue, [](APFloat f) { return f.isZero(); }); |
| if (!subValue || summedValue != expValue) { |
| return failure(); |
| } |
| |
| // sub = src - broadcast(max). |
| auto subOp = subValue.getDefiningOp<linalg::GenericOp>(); |
| if (!subOp || subOp.getNumDpsInputs() != 2 || subOp.getNumDpsInits() != 1 || |
| !isAllParallel(subOp) || |
| !isIdentityOperand(subOp, subOp.getDpsInitOperand(0)) || |
| !hasSingleOpBody<arith::SubFOp>(subOp)) { |
| return failure(); |
| } |
| OpOperand *subSource = subOp.getDpsInputOperand(0); |
| if (!isIdentityOperand(subOp, subSource)) { |
| return failure(); |
| } |
| Value maxValue = |
| resolveBroadcastedReduction(subOp, subOp.getDpsInputOperand(1)); |
| if (!maxValue) { |
| return failure(); |
| } |
| |
| // max = reduce_max(src), reducing the same source the subtraction reads. The |
| // init must be -inf or the lowest finite value. Accept both arith.maximumf |
| // (NaN-propagating, as emitted by e.g. StableHLO frontends) and arith.maxnumf |
| // (NaN-ignoring, which is what linalg.softmax itself decomposes to), since |
| // both denote the stabilizing max of a softmax. |
| auto isNegInfOrLowest = [](APFloat f) { |
| return (f.isLargest() || f.isInfinity()) && f.isNegative(); |
| }; |
| Value source = subSource->get(); |
| Value reducedValue = |
| matchInnermostReduction<arith::MaximumFOp>(maxValue, isNegInfOrLowest); |
| if (!reducedValue) { |
| reducedValue = |
| matchInnermostReduction<arith::MaxNumFOp>(maxValue, isNegInfOrLowest); |
| } |
| if (reducedValue != source) { |
| return failure(); |
| } |
| return source; |
| } |
| |
| class RaiseSoftmax : public OpInterfaceRewritePattern<linalg::LinalgOp> { |
| public: |
| using OpInterfaceRewritePattern<linalg::LinalgOp>::OpInterfaceRewritePattern; |
| RaiseSoftmax(MLIRContext *ctx, PatternBenefit b = 1) |
| : OpInterfaceRewritePattern<linalg::LinalgOp>(ctx, b) {} |
| |
| LogicalResult matchAndRewrite(linalg::LinalgOp linalgOp, |
| PatternRewriter &rewriter) const override { |
| if (!IREE::Flow::isNonNullAndOutsideDispatch(linalgOp)) { |
| return failure(); |
| } |
| |
| FailureOr<Value> src = matchSoftmax(linalgOp); |
| if (failed(src)) { |
| return rewriter.notifyMatchFailure(linalgOp, |
| "failed to match softmax root"); |
| } |
| |
| rewriter.setInsertionPoint(linalgOp); |
| rewriter.replaceOpWithNewOp<linalg::SoftmaxOp>( |
| linalgOp, linalgOp->getResultTypes(), *src, |
| linalgOp.getDpsInitOperand(0)->get(), linalgOp.getNumLoops() - 1); |
| return success(); |
| } |
| }; |
| |
| //===----------------------------------------------------------------------===// |
| // (Batch) Matmul Transpose Fusion |
| //===----------------------------------------------------------------------===// |
| |
| // Method to match a transpose operation on the two most minor dimensions of the |
| // specified rank. |
| static bool matchInner2DTranspose(linalg::LinalgOp genericOp, unsigned rank) { |
| // Only makes sense for minimum rank 2. |
| if (rank < 2) { |
| return false; |
| } |
| if (genericOp.getNumDpsInputs() != 1 || genericOp.getNumDpsInits() != 1) { |
| return false; |
| } |
| // Check only for ops of the specified rank. |
| if (genericOp.getNumLoops() != rank || |
| genericOp.getNumLoops() != genericOp.getNumParallelLoops()) { |
| return false; |
| } |
| // Check for transpose map. |
| SmallVector<AffineExpr> exprList(rank); |
| MLIRContext *context = genericOp.getContext(); |
| bindDimsList(context, MutableArrayRef{exprList}); |
| SmallVector<AffineExpr> transposeExprList(exprList); |
| std::swap(transposeExprList[rank - 1], transposeExprList[rank - 2]); |
| SmallVector<AffineMap> expectedMaps = { |
| AffineMap::get(rank, 0, exprList, context), |
| AffineMap::get(rank, 0, transposeExprList, context)}; |
| if (genericOp.getIndexingMapsArray() != expectedMaps) { |
| return false; |
| } |
| |
| Block *body = genericOp.getBlock(); |
| if (!llvm::hasSingleElement(*body)) { |
| return false; |
| } |
| auto yieldOp = cast<linalg::YieldOp>(body->getTerminator()); |
| auto blockArg = dyn_cast<BlockArgument>(yieldOp.getOperand(0)); |
| if (!blockArg || blockArg.getOwner() != body || |
| blockArg.getArgNumber() != 0) { |
| return false; |
| } |
| return true; |
| } |
| |
| // Method to match a linalg.matmul(a, linalg.transpose(b)). Returns `b` on |
| // success. |
| static std::optional<Value> matchATransposeBMatmul(linalg::LinalgOp matmulOp) { |
| if (!IREE::LinalgExt::isPureMatmul(matmulOp)) { |
| return std::nullopt; |
| } |
| auto rhs = matmulOp.getDpsInputOperand(1); |
| auto genericOp = rhs->get().getDefiningOp<linalg::GenericOp>(); |
| if (genericOp && matchInner2DTranspose(genericOp, 2)) { |
| return genericOp.getDpsInputOperand(0)->get(); |
| } |
| return std::nullopt; |
| } |
| |
| // Method to match a linalg.batch_matmul(a, linalg.transpose(b)). Returns `b` on |
| // success. |
| static std::optional<Value> |
| matchATransposeBBatchMatmul(linalg::LinalgOp bmmOp) { |
| if (!IREE::LinalgExt::isPureBatchMatmul(bmmOp)) { |
| return std::nullopt; |
| } |
| auto rhs = bmmOp.getDpsInputOperand(1); |
| auto genericOp = rhs->get().getDefiningOp<linalg::GenericOp>(); |
| if (genericOp && matchInner2DTranspose(genericOp, 3)) { |
| return genericOp.getDpsInputOperand(0)->get(); |
| } |
| return std::nullopt; |
| } |
| |
| class FuseMatmulTranspose : public OpRewritePattern<linalg::MatmulOp> { |
| public: |
| using Base::Base; |
| LogicalResult matchAndRewrite(linalg::MatmulOp matmulOp, |
| PatternRewriter &rewriter) const override { |
| if (!IREE::Flow::isNonNullAndOutsideDispatch(matmulOp)) { |
| return failure(); |
| } |
| |
| std::optional<Value> newRhs = matchATransposeBMatmul(matmulOp); |
| if (!newRhs) { |
| return rewriter.notifyMatchFailure(matmulOp, |
| "Failed to match transpose on RHS"); |
| } |
| |
| Value lhs = matmulOp.getDpsInputOperand(0)->get(); |
| Value init = matmulOp.getDpsInitOperand(0)->get(); |
| rewriter.setInsertionPoint(matmulOp); |
| SmallVector<NamedAttribute> attrs = getPrunedAttributeList(matmulOp); |
| rewriter.replaceOpWithNewOp<linalg::MatmulTransposeBOp>( |
| matmulOp, ValueRange{lhs, *newRhs}, ValueRange{init}, attrs); |
| return success(); |
| } |
| }; |
| |
| class FuseBatchMatmulTranspose |
| : public OpRewritePattern<linalg::BatchMatmulOp> { |
| public: |
| using Base::Base; |
| LogicalResult matchAndRewrite(linalg::BatchMatmulOp bmmOp, |
| PatternRewriter &rewriter) const override { |
| if (!IREE::Flow::isNonNullAndOutsideDispatch(bmmOp)) { |
| return failure(); |
| } |
| |
| std::optional<Value> newRhs = matchATransposeBBatchMatmul(bmmOp); |
| if (!newRhs) { |
| return rewriter.notifyMatchFailure(bmmOp, |
| "Failed to match transpose on RHS"); |
| } |
| |
| Value lhs = bmmOp.getDpsInputOperand(0)->get(); |
| Value init = bmmOp.getDpsInitOperand(0)->get(); |
| rewriter.setInsertionPoint(bmmOp); |
| SmallVector<NamedAttribute> attrs = getPrunedAttributeList(bmmOp); |
| rewriter.replaceOpWithNewOp<linalg::BatchMatmulTransposeBOp>( |
| bmmOp, ValueRange{lhs, *newRhs}, ValueRange{init}, attrs); |
| return success(); |
| } |
| }; |
| |
| //===----------------------------------------------------------------------===// |
| // Fill Raising |
| //===----------------------------------------------------------------------===// |
| |
| // Method to match a linalg.generic op representing a linalg.fill op. Returns |
| // the fill value (input operand to linalg.fill) on success. |
| static std::optional<Value> matchGenericFill(linalg::LinalgOp linalgOp) { |
| if (isa<linalg::GenericOp>(linalgOp.getOperation()) && |
| linalgOp.getNumDpsInputs() == 0 && linalgOp.getNumDpsInits() == 1 && |
| linalgOp.getNumParallelLoops() == linalgOp.getNumLoops() && |
| linalgOp.getIndexingMapsArray()[0].isIdentity()) { |
| // Check that the op body is only a linalg.yield op. |
| if (linalgOp.getBlock()->getOperations().size() != 1) { |
| return std::nullopt; |
| } |
| Value yieldOperand = |
| cast<linalg::YieldOp>(linalgOp.getBlock()->getOperations().begin()) |
| .getOperand(0); |
| |
| // Check that the operand of the linalg.yield op is not an argument of the |
| // linalg.generic basic block. |
| for (Value blockArg : linalgOp.getBlock()->getArguments()) { |
| if (yieldOperand == blockArg) { |
| return std::nullopt; |
| } |
| } |
| return yieldOperand; |
| } |
| return std::nullopt; |
| } |
| |
| class RaiseGenericFill : public OpRewritePattern<linalg::GenericOp> { |
| public: |
| using Base::Base; |
| LogicalResult matchAndRewrite(linalg::GenericOp genericOp, |
| PatternRewriter &rewriter) const override { |
| if (!IREE::Flow::isNonNullAndOutsideDispatch(genericOp)) { |
| return failure(); |
| } |
| |
| std::optional<Value> fillInput = matchGenericFill(genericOp); |
| if (!fillInput) { |
| return rewriter.notifyMatchFailure(genericOp, |
| "failed to match as a fill"); |
| } |
| |
| Value init = genericOp.getDpsInitOperand(0)->get(); |
| rewriter.setInsertionPoint(genericOp); |
| SmallVector<NamedAttribute> attrs = getPrunedAttributeList(genericOp); |
| rewriter.replaceOpWithNewOp<linalg::FillOp>( |
| genericOp, ValueRange{*fillInput}, ValueRange{init}, attrs); |
| |
| return success(); |
| } |
| }; |
| |
| //===----------------------------------------------------------------------===// |
| // Insert Slice to Pad Raising |
| //===----------------------------------------------------------------------===// |
| |
| class RaiseInsertSliceToPad : public OpRewritePattern<tensor::InsertSliceOp> { |
| public: |
| using Base::Base; |
| LogicalResult matchAndRewrite(tensor::InsertSliceOp sliceOp, |
| PatternRewriter &rewriter) const override { |
| if (!IREE::Flow::isNonNullAndOutsideDispatch(sliceOp)) { |
| return failure(); |
| } |
| if (sliceOp.getSourceType().getRank() != |
| sliceOp.getResultType().getRank()) { |
| return rewriter.notifyMatchFailure( |
| sliceOp, "unimplemented: rank reducing slice raising"); |
| } |
| if (!sliceOp.hasUnitStride()) { |
| return rewriter.notifyMatchFailure(sliceOp, "strided insert slice"); |
| } |
| |
| auto constantDest = sliceOp.getDest().getDefiningOp<arith::ConstantOp>(); |
| if (!constantDest) { |
| return rewriter.notifyMatchFailure(sliceOp, |
| "Non constant slice destination"); |
| } |
| |
| auto denseAttr = dyn_cast<DenseElementsAttr>(constantDest.getValue()); |
| if (!denseAttr || !denseAttr.isSplat()) { |
| return rewriter.notifyMatchFailure( |
| sliceOp, "Failed to match constant splat destination"); |
| } |
| |
| Location loc = sliceOp.getLoc(); |
| |
| AffineExpr d0, d1, d2; |
| bindDims(rewriter.getContext(), d0, d1, d2); |
| |
| SmallVector<OpFoldResult> lowPadding = sliceOp.getMixedOffsets(); |
| SmallVector<OpFoldResult> highPadding; |
| for (auto [dim, offset, size] : |
| llvm::zip_equal(denseAttr.getType().getShape(), lowPadding, |
| sliceOp.getMixedSizes())) { |
| highPadding.push_back(affine::makeComposedFoldedAffineApply( |
| rewriter, loc, d0 - d1 - d2, |
| {rewriter.getIndexAttr(dim), size, offset})); |
| } |
| |
| Value paddingValue = arith::ConstantOp::create( |
| rewriter, constantDest.getLoc(), denseAttr.getElementType(), |
| denseAttr.getSplatValue<TypedAttr>()); |
| rewriter.replaceOpWithNewOp<tensor::PadOp>(sliceOp, sliceOp.getResultType(), |
| sliceOp.getSource(), lowPadding, |
| highPadding, paddingValue); |
| return success(); |
| } |
| }; |
| |
| //===----------------------------------------------------------------------===// |
| // Concatenate Negate and Slice Raising Patterns |
| //===----------------------------------------------------------------------===// |
| |
| static bool isUnaryNegate(linalg::GenericOp op) { |
| auto block = op.getBlock(); |
| if (block->getNumArguments() != 2) { |
| return false; |
| } |
| |
| Operation *terminator = block->getTerminator(); |
| if (terminator->getNumOperands() != 1) { |
| return false; |
| } |
| |
| auto yielded = terminator->getOperand(0).getDefiningOp<arith::NegFOp>(); |
| if (!yielded) { |
| return false; |
| } |
| |
| if (yielded->getOperand(0) != block->getArgument(0)) { |
| return false; |
| } |
| |
| return true; |
| } |
| |
| // This aims to match a partial negation and reverse of a tensor and rewrite |
| // it as a single linalg generic. For example, this will match sequences like |
| // the following: |
| // |
| // %slice_0 = tensor.extract_slice %0[0, ..., 0] [..., 64] [1, ..., 1] |
| // : tensor<...x128xf16> to tensor<...x64xf16> |
| // %slice_1 = tensor.extract_slice %0[0, ..., 64] [..., 64] [1, ..., 1] |
| // : tensor<...x128xf16> to tensor<...x64xf16> |
| // %3 = linalg.generic { |
| // indexing_maps = [affine_map<(d0, ..., dn-1) -> (d0, ..., dn-1)>, |
| // affine_map<(d0, ..., dn-1) -> (d0, ..., dn-1)>], |
| // iterator_types = [n-1 x "parallel"]} |
| // ins(%slice_1 : tensor<...x64xf16>) outs(%2 : tensor<...x64xf16>) { |
| // ^bb0(%in: f16, %out: f16): |
| // %5 = arith.negf %in : f16 |
| // linalg.yield %5 : f16 |
| // } -> tensor<...x64xf16> |
| // %in_slice_0 = tensor.insert_slice %3 into |
| // %1[0, ..., 0] [..., 64] [1, ..., 1] |
| // : tensor<...x64xf16> into tensor<...x128xf16> |
| // %in_slice_1 = tensor.insert_slice %slice_0 into |
| // %in_slice_0[0, ..., 64] [..., 64] [1, ..., 1] |
| // : tensor<...x64xf16> into tensor<...x128xf16> |
| // |
| // Where the input tensor is broken down along the inner most dimension, then |
| // the bottom elements are negated and the tensor is reconstructed, reversing |
| // the order of the slices. This is then rewritten as |
| // |
| // %expanded = tensor.expand_shape %0 [[0], ..., [n-1, n]] |
| // : tensor<...x128xf16> into tensor<...x2x64xf16> |
| // %2 = linalg.generic { |
| // indexing_maps = [ |
| // affine_map<(d0, ..., dn-1, dn) -> (d0, ..., dn-1, dn)>], |
| // iterator_types = [n x "parallel"] |
| // } outs(%output : tensor<...x2x64xf16>) { |
| // ^bb0(%out: f16): |
| // %i0 = linalg.index 0 : index |
| // ... |
| // %in-1 = linalg.index n-1 : index |
| // %in = linalg.index n : index |
| // %rev = arith.subi %c1, %in-1 : index |
| // %extracted = tensor.extract %expanded[%i0, ..., %rev, %in] |
| // : tensor<...x2x64xf16> |
| // %neg = arith.negf %extracted : f16 |
| // %cmp = arith.cmpi eq, %in-1, %c1 : index |
| // %sel = arith.select %cmp, %neg, %extracted : f16 |
| // linalg.yield %sel : f16 |
| // } -> tensor<...x2x64xf16> |
| // %collapsed = tensor.collapse_shape %2 [[0], ..., [n-1, n]] |
| // : tensor<...x2x64xf16> into tensor<...x128xf16> |
| static std::optional<Value> |
| matchCatNegateAndSlice(tensor::InsertSliceOp insertOp) { |
| /// First match against the desired op chain. |
| auto topHalf = insertOp.getSource().getDefiningOp<tensor::ExtractSliceOp>(); |
| if (!topHalf) { |
| return std::nullopt; |
| } |
| auto dest = insertOp.getDest().getDefiningOp<tensor::InsertSliceOp>(); |
| if (!dest) { |
| return std::nullopt; |
| } |
| |
| /// TODO: This could be extended to other unary (or in general, elementwise |
| /// operations) if the need arises. |
| auto negate = dest.getSource().getDefiningOp<linalg::GenericOp>(); |
| if (!negate || !isUnaryNegate(negate)) { |
| return std::nullopt; |
| } |
| |
| auto bottomHalf = |
| negate.getDpsInputs()[0].getDefiningOp<tensor::ExtractSliceOp>(); |
| if (!bottomHalf) { |
| return std::nullopt; |
| } |
| Value source = topHalf.getSource(); |
| if (source != bottomHalf.getSource()) { |
| return std::nullopt; |
| } |
| |
| /// Require that the overall operation isn't changing the tensor shape. |
| if (source.getType() != dest.getType()) { |
| return std::nullopt; |
| } |
| |
| auto sourceType = dyn_cast<RankedTensorType>(source.getType()); |
| if (!sourceType || sourceType.getRank() == 0) { |
| return std::nullopt; |
| } |
| |
| /// Here we are verifying that the slice sequence inserts and extracts in the |
| /// same order. In other words, the offsets, sizes, and strides for the nth |
| /// extract slice match the nth insert slice. |
| /// TODO: We could relax this condition if we see cases where the elementwise |
| /// unary operation happens on a different slices of the tensor, but for now |
| /// don't engineer for cases that don't exist. |
| auto isSame = [](OpFoldResult a, OpFoldResult b) { return a == b; }; |
| if (!insertOp.isSameAs(bottomHalf, isSame) || |
| !dest.isSameAs(topHalf, isSame)) { |
| return std::nullopt; |
| } |
| |
| /// All slices need unit stride and the initial slice should have a zero |
| /// offset. |
| if (!insertOp.hasUnitStride() || !dest.hasUnitStride() || |
| !dest.hasZeroOffset()) { |
| return std::nullopt; |
| } |
| |
| /// Verify all slice sizes are the same. |
| for (auto [bottom, top] : |
| llvm::zip_equal(insertOp.getMixedSizes(), dest.getMixedSizes())) { |
| if (bottom != top) { |
| return std::nullopt; |
| } |
| } |
| |
| /// Verify that the final InsertSlice has zero offsets for all except the |
| /// inner most dim. |
| SmallVector<OpFoldResult> insertOffsets = insertOp.getMixedOffsets(); |
| for (int i = 0, e = insertOffsets.size() - 1; i < e; ++i) { |
| if (!isZeroInteger(insertOffsets[i])) { |
| return std::nullopt; |
| } |
| } |
| |
| /// Finally verify that the inner most offset == the inner most slice size, |
| /// and because we verified all slice sizes are the same, this necessarily |
| /// is slicing exactly half of the inner most dimension. |
| if (insertOffsets.back() != insertOp.getMixedSizes().back()) { |
| return std::nullopt; |
| } |
| |
| return source; |
| } |
| |
| // This aims to match the exact same logical operation, but instead where the |
| // final concatenation is actually represented as a concatenation. |
| // |
| // %slice_0 = tensor.extract_slice %0[0, ..., 0] [..., 64] [1, ..., 1] |
| // : tensor<...x128xf16> to tensor<...x64xf16> |
| // %slice_1 = tensor.extract_slice %0[0, ..., 64] [..., 64] [1, ..., 1] |
| // : tensor<...x128xf16> to tensor<...x64xf16> |
| // %3 = linalg.generic { |
| // indexing_maps = [affine_map<(d0, ..., dn-1) -> (d0, ..., dn-1)>, |
| // affine_map<(d0, ..., dn-1) -> (d0, ..., dn-1)>], |
| // iterator_types = [n-1 x "parallel"]} |
| // ins(%slice_1 : tensor<...x64xf16>) outs(%2 : tensor<...x64xf16>) { |
| // ^bb0(%in: f16, %out: f16): |
| // %5 = arith.negf %in : f16 |
| // linalg.yield %5 : f16 |
| // } -> tensor<...x64xf16> |
| // %in_slice_0 = tensor.concat %3, %slice_0 : ... -> tensor<...x128xf16> |
| // |
| // With the exact same rewritten IR. |
| static std::optional<Value> matchCatNegateAndSlice(tensor::ConcatOp concatOp) { |
| /// First match against the desired op chain. |
| ValueRange inputs = concatOp.getInputs(); |
| if (inputs.size() != 2) { |
| return std::nullopt; |
| } |
| |
| Value left = inputs[0]; |
| Value right = inputs[1]; |
| if (left.getType() != right.getType()) { |
| return std::nullopt; |
| } |
| |
| // The concatenation must happen along the inner most dimension. |
| auto inputType = cast<RankedTensorType>(left.getType()); |
| if (concatOp.getDim() != inputType.getRank() - 1) { |
| return std::nullopt; |
| } |
| |
| /// TODO: This could be extended to other unary (or in general, elementwise |
| /// operations) if the need arises. |
| auto negate = left.getDefiningOp<linalg::GenericOp>(); |
| if (!negate || !isUnaryNegate(negate)) { |
| return std::nullopt; |
| } |
| |
| auto topHalf = right.getDefiningOp<tensor::ExtractSliceOp>(); |
| if (!topHalf) { |
| return std::nullopt; |
| } |
| |
| auto bottomHalf = |
| negate.getDpsInputs()[0].getDefiningOp<tensor::ExtractSliceOp>(); |
| if (!bottomHalf) { |
| return std::nullopt; |
| } |
| Value source = topHalf.getSource(); |
| if (source != bottomHalf.getSource()) { |
| return std::nullopt; |
| } |
| |
| /// Require that the overall operation isn't changing the tensor shape. |
| if (source.getType() != concatOp.getType()) { |
| return std::nullopt; |
| } |
| return source; |
| } |
| |
| static Value createCatNegateAndSlice(RewriterBase &rewriter, Value outTensor, |
| Value source) { |
| Location loc = source.getLoc(); |
| /// The matcher checks that this cast is valid. |
| auto sourceType = cast<RankedTensorType>(source.getType()); |
| |
| SmallVector<int64_t> targetShape(sourceType.getShape()); |
| SmallVector<ReassociationIndices> reassoc; |
| for (int i = 0, e = targetShape.size(); i < e; i++) { |
| reassoc.push_back(ReassociationIndices{i}); |
| } |
| reassoc.back().push_back(targetShape.size()); |
| |
| /// Note that because we've verified the pair of slices is exactly half |
| /// of the whole inner most extent, the sliceSize is necessarily |
| /// divisible by 2. |
| int64_t sliceSize = targetShape.back(); |
| targetShape[targetShape.size() - 1] = 2; |
| targetShape.push_back(ShapedType::isDynamic(sliceSize) ? sliceSize |
| : sliceSize / 2); |
| Type expandedType = |
| RankedTensorType::get(targetShape, sourceType.getElementType()); |
| Value expanded = tensor::ExpandShapeOp::create(rewriter, loc, expandedType, |
| source, reassoc); |
| |
| Value expandedOutTensor = tensor::ExpandShapeOp::create( |
| rewriter, loc, expandedType, outTensor, reassoc); |
| |
| SmallVector<AffineMap> indexingMaps = { |
| rewriter.getMultiDimIdentityMap(targetShape.size())}; |
| SmallVector<utils::IteratorType> iteratorTypes(targetShape.size(), |
| utils::IteratorType::parallel); |
| |
| auto bodyBuilder = [&](OpBuilder &b, Location loc, ValueRange args) { |
| SmallVector<Value> extractionIndices; |
| for (size_t i = 0, e = targetShape.size(); i < e; ++i) { |
| extractionIndices.push_back(linalg::IndexOp::create(b, loc, i)); |
| } |
| |
| Value c1 = |
| arith::ConstantOp::create(rewriter, loc, rewriter.getIndexAttr(1)); |
| |
| // Take the reverse of the second to last iterator. Because we statically |
| // guaranteed it to be 2 it just becomes `1 - iters[-2]`. |
| Value reverseSplitIdx = arith::SubIOp::create( |
| rewriter, loc, c1, extractionIndices[targetShape.size() - 2]); |
| extractionIndices[targetShape.size() - 2] = reverseSplitIdx; |
| |
| // Extract the value from input tensor and negate the top half of the result |
| // slice (lower half of the input slice). |
| Value inputVal = |
| tensor::ExtractOp::create(b, loc, expanded, extractionIndices); |
| Value maybeNegate = arith::NegFOp::create(b, loc, inputVal); |
| |
| Value isEqual = arith::CmpIOp::create(b, loc, arith::CmpIPredicate::eq, |
| reverseSplitIdx, c1); |
| Value select = |
| arith::SelectOp::create(rewriter, loc, isEqual, maybeNegate, inputVal); |
| linalg::YieldOp::create(b, loc, select); |
| }; |
| |
| Value result = |
| linalg::GenericOp::create(rewriter, loc, expandedOutTensor.getType(), |
| ValueRange(), expandedOutTensor, indexingMaps, |
| iteratorTypes, bodyBuilder) |
| .getResult(0); |
| |
| return tensor::CollapseShapeOp::create(rewriter, loc, outTensor.getType(), |
| result, reassoc); |
| } |
| |
| static Value rewriteCatNegateAndSlice(RewriterBase &rewriter, |
| tensor::InsertSliceOp sliceOp, |
| Value source) { |
| rewriter.setInsertionPoint(sliceOp); |
| Value outTensor = |
| sliceOp.getDest().getDefiningOp<tensor::InsertSliceOp>().getDest(); |
| return createCatNegateAndSlice(rewriter, outTensor, source); |
| } |
| |
| static Value rewriteCatNegateAndSlice(RewriterBase &rewriter, |
| tensor::ConcatOp concatOp, Value source) { |
| rewriter.setInsertionPoint(concatOp); |
| Type elemType = cast<RankedTensorType>(source.getType()).getElementType(); |
| Value outTensor = tensor::EmptyOp::create( |
| rewriter, source.getLoc(), |
| tensor::getMixedSizes(rewriter, source.getLoc(), source), elemType); |
| return createCatNegateAndSlice(rewriter, outTensor, source); |
| } |
| |
| class InsertSliceNegateAndSlicePattern |
| : public OpRewritePattern<tensor::InsertSliceOp> { |
| public: |
| using Base::Base; |
| LogicalResult matchAndRewrite(tensor::InsertSliceOp sliceOp, |
| PatternRewriter &rewriter) const override { |
| if (!IREE::Flow::isNonNullAndOutsideDispatch(sliceOp)) { |
| return failure(); |
| } |
| |
| std::optional<Value> dagRoot = matchCatNegateAndSlice(sliceOp); |
| if (!dagRoot) { |
| return rewriter.notifyMatchFailure(sliceOp, "failed to match target dag"); |
| } |
| Value res = rewriteCatNegateAndSlice(rewriter, sliceOp, *dagRoot); |
| rewriter.replaceOp(sliceOp, res); |
| return success(); |
| } |
| }; |
| |
| class ConcatenateNegateAndSlicePattern |
| : public OpRewritePattern<tensor::ConcatOp> { |
| public: |
| using Base::Base; |
| LogicalResult matchAndRewrite(tensor::ConcatOp concatOp, |
| PatternRewriter &rewriter) const override { |
| if (!IREE::Flow::isNonNullAndOutsideDispatch(concatOp)) { |
| return failure(); |
| } |
| |
| std::optional<Value> dagRoot = matchCatNegateAndSlice(concatOp); |
| if (!dagRoot) { |
| return rewriter.notifyMatchFailure(concatOp, |
| "failed to match target dag"); |
| } |
| Value res = rewriteCatNegateAndSlice(rewriter, concatOp, *dagRoot); |
| rewriter.replaceOp(concatOp, res); |
| return success(); |
| } |
| }; |
| |
| //===----------------------------------------------------------------------===// |
| // Pass Implementation |
| //===----------------------------------------------------------------------===// |
| |
| struct RaiseSpecialOpsPass |
| : impl::RaiseSpecialOpsPassBase<RaiseSpecialOpsPass> { |
| void getDependentDialects(DialectRegistry ®istry) const override { |
| registry.insert<IREE::LinalgExt::IREELinalgExtDialect>(); |
| } |
| |
| void runOnOperation() override { |
| Operation *funcOp = getOperation(); |
| MLIRContext *context = &getContext(); |
| |
| // First walk the IR and try to raise any slice-like generics to tensor. |
| IRRewriter rewriter(context); |
| SmallVector<linalg::GenericOp> genericOps; |
| funcOp->walk([&](linalg::GenericOp op) { genericOps.push_back(op); }); |
| |
| for (linalg::GenericOp linalgOp : genericOps) { |
| // Try raising to tensor.extract to an input and create an linalg.generic. |
| OpBuilder::InsertionGuard guard(rewriter); |
| rewriter.setInsertionPoint(linalgOp); |
| FailureOr<linalg::GenericOp> maybeNewOp = |
| raiseTensorExtractToInput(linalgOp, rewriter); |
| if (succeeded(maybeNewOp)) { |
| rewriter.replaceOp(linalgOp, *maybeNewOp); |
| linalgOp = *maybeNewOp; |
| } |
| |
| // Try raising to a view-like operation. Replace if the op raising was |
| // successful. |
| FailureOr<Operation *> maybeRaisedView = |
| tryRaiseToView(linalgOp, rewriter); |
| if (succeeded(maybeRaisedView)) { |
| rewriter.replaceOp(linalgOp, *maybeRaisedView); |
| } |
| } |
| |
| // Next run a variety of raising patterns. |
| { |
| RewritePatternSet patterns(context); |
| patterns.insert< |
| NamedImplicitCastOpConversion<linalg::ConvolutionOpInterface>, |
| NamedImplicitCastOpConversion<linalg::ContractionOpInterface>>( |
| context); |
| patterns.insert<RaiseSoftmax>(context); |
| patterns.insert<FuseMatmulTranspose>(context); |
| patterns.insert<FuseBatchMatmulTranspose>(context); |
| patterns.insert<RaiseGenericFill>(context); |
| patterns.insert<InsertSliceNegateAndSlicePattern>(context); |
| patterns.insert<ConcatenateNegateAndSlicePattern>(context); |
| patterns.insert<RaiseInsertSliceToPad>(context); |
| if (failed(applyPatternsGreedily(funcOp, std::move(patterns)))) { |
| return signalPassFailure(); |
| } |
| } |
| } |
| }; |
| |
| } // namespace |
| } // namespace mlir::iree_compiler::GlobalOptimization |
