llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Passes/TileAndDecomposeAttention.cpp

// 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-dialects/Dialect/LinalgExt/IR/LinalgExtDialect.h"
#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.h"
#include "iree-dialects/Dialect/LinalgExt/Passes/PassDetail.h"
#include "iree-dialects/Dialect/LinalgExt/Passes/Passes.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/Math/IR/Math.h"
#include "mlir/Dialect/SCF/Transforms/Transforms.h"
#include "mlir/Dialect/Tensor/Utils/Utils.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/Support/Debug.h"

namespace mlir {
namespace iree_compiler {
namespace IREE {
namespace LinalgExt {

namespace {

// Computes a reduction along the rows of a 2d tensor of shape MxN
// to produce a tensor of shape M
template <typename T>
static Value computeRowwiseReduction(Value a, Value output, Location loc,
                                     OpBuilder &builder) {
  SmallVector<utils::IteratorType> iteratorTypes{
      utils::IteratorType::parallel, utils::IteratorType::reduction};
  AffineMap id = AffineMap::getMultiDimIdentityMap(2, builder.getContext());
  AffineExpr d0, d1;
  bindDims(builder.getContext(), d0, d1);
  // (d0, d1) -> (d0)
  auto rowMap = AffineMap::get(2, 0, {d0}, builder.getContext());
  SmallVector<AffineMap> indexingMaps{id, rowMap};
  auto genericOp = builder.create<linalg::GenericOp>(
      loc, output.getType(), a, output, indexingMaps, iteratorTypes,
      [&](OpBuilder &b, Location loc, ValueRange args) {
        Value result = b.create<T>(loc, args[0], args[1]);
        b.create<linalg::YieldOp>(loc, result);
      });
  return genericOp.getResult(0);
}

static std::tuple<Value, Value> computeNewSum(Value oldMax, Value newMax,
                                              Value oldSum, Value currentSum,
                                              Value output, Location loc,
                                              OpBuilder &builder) {
  SmallVector<utils::IteratorType> iteratorTypes{utils::IteratorType::parallel};
  auto identityMap = AffineMap::getMultiDimIdentityMap(1, builder.getContext());
  SmallVector<AffineMap> indexingMaps(6, identityMap);
  SmallVector<Type> resultTypes(2, currentSum.getType());
  auto genericOp = builder.create<linalg::GenericOp>(
      loc, resultTypes, ValueRange{oldMax, newMax, oldSum, currentSum},
      ValueRange{output, output}, indexingMaps, iteratorTypes,
      [&](OpBuilder &b, Location loc, ValueRange args) {
        Value diff = b.create<arith::SubFOp>(loc, args[0], args[1]);
        Value weight = b.create<math::ExpOp>(loc, diff);
        Value scaledOldSum = b.create<arith::MulFOp>(loc, weight, args[2]);
        Value result = b.create<arith::AddFOp>(loc, scaledOldSum, args[3]);
        b.create<linalg::YieldOp>(loc, ValueRange{result, scaledOldSum});
      });
  return std::make_tuple(genericOp.getResult(0), genericOp.getResult(1));
}

static Value computePartialSoftmax(Value qkTranspose, Value currentMax,
                                   Value output, Location loc,
                                   OpBuilder &builder) {
  AffineMap identityMap =
      AffineMap::getMultiDimIdentityMap(2, builder.getContext());
  AffineExpr d0, d1;
  bindDims(builder.getContext(), d0, d1);
  // (d0, d1) -> (d0)
  auto rowMap = AffineMap::get(2, 0, {d0}, builder.getContext());
  SmallVector<AffineMap> indexingMaps{identityMap, rowMap, identityMap};
  SmallVector<utils::IteratorType> iteratorTypes(2,
                                                 utils::IteratorType::parallel);
  auto genericOp = builder.create<linalg::GenericOp>(
      loc, qkTranspose.getType(), ValueRange{qkTranspose, currentMax}, output,
      indexingMaps, iteratorTypes,
      [&](OpBuilder &b, Location loc, ValueRange args) {
        Value diff = b.create<arith::SubFOp>(loc, args[0], args[1]);
        Value result = b.create<math::ExpOp>(loc, diff);
        b.create<linalg::YieldOp>(loc, result);
      });
  return genericOp.getResult(0);
}

static Value scalePartialSoftmax(Value softmax, Value scaledOldSum,
                                 Value output, Location loc,
                                 OpBuilder &builder) {
  AffineMap identityMap =
      AffineMap::getMultiDimIdentityMap(2, builder.getContext());
  AffineExpr d0, d1;
  bindDims(builder.getContext(), d0, d1);
  // (d0, d1) -> (d0)
  auto rowMap = AffineMap::get(2, 0, {d0}, builder.getContext());
  SmallVector<AffineMap> indexingMaps{identityMap, rowMap, identityMap};
  SmallVector<utils::IteratorType> iteratorTypes(2,
                                                 utils::IteratorType::parallel);
  auto genericOp = builder.create<linalg::GenericOp>(
      loc, softmax.getType(), ValueRange{softmax, scaledOldSum}, output,
      indexingMaps, iteratorTypes,
      [&](OpBuilder &b, Location loc, ValueRange args) {
        Value result = b.create<arith::DivFOp>(loc, args[0], args[1]);
        b.create<linalg::YieldOp>(loc, result);
      });
  return genericOp.getResult(0);
}

static Value scaleAccumulator(Value accumulator, Value scaledOldSum,
                              Value newSum, Value output, Location loc,
                              OpBuilder &builder) {
  AffineMap identityMap =
      AffineMap::getMultiDimIdentityMap(2, builder.getContext());
  AffineExpr d0, d1;
  bindDims(builder.getContext(), d0, d1);
  // (d0, d1) -> (d0)
  auto rowMap = AffineMap::get(2, 0, {d0}, builder.getContext());
  SmallVector<AffineMap> indexingMaps{identityMap, rowMap, rowMap, identityMap};
  SmallVector<utils::IteratorType> iteratorTypes(2,
                                                 utils::IteratorType::parallel);
  auto genericOp = builder.create<linalg::GenericOp>(
      loc, accumulator.getType(), ValueRange{accumulator, scaledOldSum, newSum},
      output, indexingMaps, iteratorTypes,
      [&](OpBuilder &b, Location loc, ValueRange args) {
        Value prod = b.create<arith::DivFOp>(loc, args[1], args[2]);
        Value result = b.create<arith::MulFOp>(loc, prod, args[0]);
        b.create<linalg::YieldOp>(loc, result);
      });
  return genericOp.getResult(0);
}

static Value computeQKTranspose(Value query, Value key, Value transposedOutput,
                                Value output, Value zero,
                                RankedTensorType tensorType, Location loc,
                                OpBuilder &builder) {
  SmallVector<int64_t> perm{1, 0};
  auto transposeOp =
      builder.create<linalg::TransposeOp>(loc, key, transposedOutput, perm);
  Value acc =
      builder.create<linalg::FillOp>(loc, ValueRange{zero}, output).result();
  auto matmulOp = builder.create<linalg::MatmulOp>(
      loc, tensorType, ValueRange{query, transposeOp.getResult()[0]}, acc);
  return matmulOp.getResult(0);
}

static std::tuple<Value, Value, Value, Value>
extractSlices(Value key, Value value, Value query, Value output,
              ArrayRef<int64_t> queryShape, ArrayRef<Value> ivs,
              Value sequenceTileLength, Type elementType, Location loc,
              OpBuilder &builder) {
  auto one = builder.getIndexAttr(1);
  auto zero = builder.getIndexAttr(0);
  auto headDimension = builder.getIndexAttr(queryShape.back());
  SmallVector<OpFoldResult> strides(queryShape.size(), one);
  SmallVector<OpFoldResult> sizes(queryShape.size(), one);
  SmallVector<OpFoldResult> offsets(queryShape.size(), zero);
  sizes[1] = sequenceTileLength;
  sizes[2] = headDimension;
  offsets[0] = ivs[0];
  offsets[1] = ivs[1];
  SmallVector<int64_t> tensorShape{ShapedType::kDynamic, queryShape.back()};
  auto tensorType = RankedTensorType::get(tensorShape, elementType);
  Value keySlice = builder.create<tensor::ExtractSliceOp>(
      loc, tensorType, key, offsets, sizes, strides);
  Value valueSlice = builder.create<tensor::ExtractSliceOp>(
      loc, tensorType, value, offsets, sizes, strides);

  offsets = SmallVector<OpFoldResult>(queryShape.size(), zero);
  offsets[0] = ivs[0];
  Value querySlice = builder.create<tensor::ExtractSliceOp>(
      loc, tensorType, query, offsets, sizes, strides);
  Value outputSlice = builder.create<tensor::ExtractSliceOp>(
      loc, tensorType, output, offsets, sizes, strides);

  return std::make_tuple(keySlice, valueSlice, querySlice, outputSlice);
}

static std::tuple<Value, Value, Value>
insertSlices(Value newResult, Value result, Value newMax, Value max,
             Value newSum, Value sum, ArrayRef<int64_t> queryShape,
             ArrayRef<Value> ivs, Value sequenceTileLength, Location loc,
             OpBuilder &builder) {
  auto one = builder.getIndexAttr(1);
  auto zero = builder.getIndexAttr(0);
  auto headDimension = builder.getIndexAttr(queryShape.back());
  SmallVector<OpFoldResult> strides(queryShape.size(), one);
  SmallVector<OpFoldResult> sizes(queryShape.size(), one);
  SmallVector<OpFoldResult> offsets(queryShape.size(), zero);
  sizes[1] = sequenceTileLength;
  sizes[2] = headDimension;
  offsets[0] = ivs[0];
  Value updatedAcc = builder.create<tensor::InsertSliceOp>(
      loc, newResult, result, offsets, sizes, strides);
  offsets = SmallVector<OpFoldResult>{zero};
  sizes = SmallVector<OpFoldResult>{sequenceTileLength};
  strides = SmallVector<OpFoldResult>{one};
  Value updatedMax = builder.create<tensor::InsertSliceOp>(
      loc, newMax, max, offsets, sizes, strides);
  Value updatedSum = builder.create<tensor::InsertSliceOp>(
      loc, newSum, sum, offsets, sizes, strides);
  return std::make_tuple(updatedAcc, updatedMax, updatedSum);
}

static scf::LoopNest createLoopNest(SmallVectorImpl<Value> &ivs, Value lb,
                                    Value step, Value ub, ValueRange args,
                                    Location loc, OpBuilder &builder) {
  SmallVector<Value> lbs{lb};
  SmallVector<Value> steps{step};
  SmallVector<Value> ubs{ub};
  scf::LoopNest loopNest = scf::buildLoopNest(
      builder, loc, lbs, ubs, steps, args,
      [&](OpBuilder &nestedBuilder, Location loc, ValueRange outputIvs,
          ValueRange iterArgs) -> scf::ValueVector { return iterArgs; });
  for (scf::ForOp loop : loopNest.loops)
    ivs.push_back(loop.getInductionVar());
  return loopNest;
}

/// This is an implementation of flash attention which
/// is a tiled and fused implementation of the attention operator.
/// The attention operator computes:
/// matmul(softmax(matmul(Q, transpose(K))), V)
/// where: Q is the query matrix [B x N x d]
///        K is the key matrix   [B x N x d]
///        V is the value matrix [B x N x d]
///
/// The core algorithm is as follows:
/// For each element in B,
/// 1. Load a tile from the Q matrix of size T x d -> q
/// 2. Initialize statistics: running_sum, running_max
/// 3. for i = 0 to N with step T
///    a. Load a tile from the K matrix of size T x d -> k
///    a. Load a tile from the V matrix of size T x d -> v
///    b. Transpose(k) -> kT
///    c. Compute matmul(q, kT) -> qkT
///    d. Compute sum(qkT) along rows -> current_sum
///    e. Compute max(qkT) along rows -> current_max
///    f. Compute new max: max(current_max, running_max)
///    g. Compute new sum: alpha * running_sum + beta * current_sum
///    h. Compute curent estimate of softmax: exp(qKT - current_max) -> s
///    i. Scale softmax estimate and current value of result by
///       appropriate factors
///    j. Compute matmul(s, v) and add to accumulator
///
///
LogicalResult reifyAttentionTransform(func::FuncOp funcOp) {
  IRRewriter rewriter(funcOp.getContext());
  funcOp.walk([&](IREE::LinalgExt::AttentionOp attnOp) {
    Location loc = attnOp.getLoc();
    OpBuilder::InsertionGuard guard(rewriter);
    rewriter.setInsertionPoint(attnOp);

    Value query = attnOp.getQuery();
    ShapedType queryType = attnOp.getQueryType();
    Type elementType = queryType.getElementType();
    ArrayRef<int64_t> queryShape = queryType.getShape();
    SmallVector<OpFoldResult> queryDimValues =
        tensor::createDimValues(rewriter, loc, query);
    Value sequenceTileLength =
        getValueOrCreateConstantIndexOp(rewriter, loc, queryDimValues[1]);
    Value batchTileLength =
        getValueOrCreateConstantIndexOp(rewriter, loc, queryDimValues[0]);

    Value key = attnOp.getKey();
    Value value = attnOp.getValue();
    SmallVector<OpFoldResult> keyDimValues =
        tensor::createDimValues(rewriter, loc, key);
    Value sequenceLength =
        getValueOrCreateConstantIndexOp(rewriter, loc, keyDimValues[1]);

    // Construct first loop
    Value zeroValue = rewriter.create<arith::ConstantIndexOp>(loc, 0);
    Value oneValue = rewriter.create<arith::ConstantIndexOp>(loc, 1);
    SmallVector<Value> ivs;
    Value output = attnOp.getOutput();
    scf::LoopNest firstLoopNest =
        createLoopNest(ivs, zeroValue, oneValue, batchTileLength,
                       ValueRange({output}), loc, rewriter);
    Value iterArg = firstLoopNest.loops.back().getRegionIterArg(0);

    OpBuilder::InsertionGuard guardFirstLoop(rewriter);
    rewriter.setInsertionPointToStart(firstLoopNest.loops.back().getBody());

    // Create max and sum statistics
    SmallVector<OpFoldResult> dims{sequenceTileLength};
    Value zeroF32 = rewriter.create<arith::ConstantOp>(
        loc, rewriter.getZeroAttr(elementType));
    Value largeNegativeF32 = rewriter.create<arith::ConstantOp>(
        loc, rewriter.getFloatAttr(elementType, -1e30));
    Value max = rewriter.create<tensor::EmptyOp>(loc, dims, elementType);
    Value negativeMax =
        rewriter.create<linalg::FillOp>(loc, ValueRange{largeNegativeF32}, max)
            .result();
    Value sum = rewriter.create<tensor::EmptyOp>(loc, dims, elementType);
    Value zeroSum =
        rewriter.create<linalg::FillOp>(loc, ValueRange{zeroF32}, sum).result();

    // Construct second loop
    scf::LoopNest secondLoopNest = createLoopNest(
        ivs, zeroValue, sequenceTileLength, sequenceLength,
        ValueRange({iterArg, negativeMax, zeroSum}), loc, rewriter);

    Value iterArgResult = secondLoopNest.loops.back().getRegionIterArg(0);
    Value iterArgMax = secondLoopNest.loops.back().getRegionIterArg(1);
    Value iterArgSum = secondLoopNest.loops.back().getRegionIterArg(2);

    OpBuilder::InsertionGuard guardSecondLoop(rewriter);
    rewriter.setInsertionPointToStart(secondLoopNest.loops.back().getBody());

    auto [keySlice, valueSlice, querySlice, outputSlice] =
        extractSlices(key, value, query, iterArgResult, queryShape, ivs,
                      sequenceTileLength, elementType, loc, rewriter);

    // Compute matmul(q, transpose(k))
    auto headDimension = rewriter.getIndexAttr(queryShape.back());
    SmallVector<OpFoldResult> transposedShape{headDimension,
                                              sequenceTileLength};
    Value empty =
        rewriter.create<tensor::EmptyOp>(loc, transposedShape, elementType);
    SmallVector<OpFoldResult> resultShape{sequenceTileLength,
                                          sequenceTileLength};
    Value emptySquare =
        rewriter.create<tensor::EmptyOp>(loc, resultShape, elementType);
    auto tensorType = RankedTensorType::get(
        SmallVector<int64_t>(2, ShapedType::kDynamic), elementType);
    Value qkTranspose =
        computeQKTranspose(querySlice, keySlice, empty, emptySquare, zeroF32,
                           tensorType, loc, rewriter);

    empty = rewriter.create<tensor::EmptyOp>(
        loc, SmallVector<OpFoldResult>{sequenceTileLength}, elementType);

    // Compute current statistics
    Value newMax = computeRowwiseReduction<arith::MaxFOp>(
        qkTranspose, iterArgMax, loc, rewriter);
    Value partialSoftmax =
        computePartialSoftmax(qkTranspose, newMax, emptySquare, loc, rewriter);
    Value currentSum = computeRowwiseReduction<arith::AddFOp>(
        partialSoftmax, zeroSum, loc, rewriter);

    auto [newSum, scaledOldSum] = computeNewSum(
        iterArgMax, newMax, iterArgSum, currentSum, empty, loc, rewriter);

    // Scale partial softmax
    Value softmax =
        scalePartialSoftmax(partialSoftmax, newSum, emptySquare, loc, rewriter);

    // Update accumulator
    empty = rewriter.create<tensor::EmptyOp>(
        loc, SmallVector<OpFoldResult>{sequenceLength, headDimension},
        elementType);
    Value scaledAcc = scaleAccumulator(outputSlice, scaledOldSum, newSum, empty,
                                       loc, rewriter);

    // Compute matmul(softmax, v)
    Value result = rewriter
                       .create<linalg::MatmulOp>(
                           loc, outputSlice.getType(),
                           ValueRange{softmax, valueSlice}, scaledAcc)
                       .getResult(0);

    // Insert slices
    auto [updatedAcc, updatedMax, updatedSum] = insertSlices(
        result, iterArgResult, newMax, iterArgMax, newSum, iterArgSum,
        queryShape, ivs, sequenceTileLength, loc, rewriter);

    if (scf::YieldOp yieldOp = dyn_cast<scf::YieldOp>(
            secondLoopNest.loops.back().getBody()->getTerminator())) {
      rewriter.replaceOpWithNewOp<scf::YieldOp>(
          yieldOp, ValueRange{updatedAcc, updatedMax, updatedSum});
    }

    if (scf::YieldOp yieldOp = dyn_cast<scf::YieldOp>(
            firstLoopNest.loops.back().getBody()->getTerminator())) {
      OpBuilder::InsertionGuard yieldGuard(rewriter);
      rewriter.setInsertionPoint(yieldOp);
      rewriter.replaceOpWithNewOp<scf::YieldOp>(
          yieldOp, ValueRange{secondLoopNest.results[0]});
    }

    attnOp.getResults()[0].replaceAllUsesWith(firstLoopNest.results[0]);
    return WalkResult::advance();
  });
  return success();
}

} // namespace

namespace {
struct TileAndDecomposeAttentionPass
    : public TileAndDecomposeAttentionBase<TileAndDecomposeAttentionPass> {
  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<AffineDialect, IREE::LinalgExt::IREELinalgExtDialect,
                    linalg::LinalgDialect, scf::SCFDialect,
                    tensor::TensorDialect>();
  }

  void runOnOperation() override;
};
} // namespace

void TileAndDecomposeAttentionPass::runOnOperation() {
  MLIRContext *context = &getContext();
  IRRewriter rewriter(context);
  if (failed(reifyAttentionTransform(getOperation())))
    return signalPassFailure();
}

std::unique_ptr<Pass> createTileAndDecomposeAttentionPass() {
  return std::make_unique<TileAndDecomposeAttentionPass>();
}

} // namespace LinalgExt
} // namespace IREE
} // namespace iree_compiler
} // namespace mlir