//===- TileAllocation.cpp - Allocate SME ZA tiles -------------------------===//
//
// Part of the LLVM Project, 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
//
//===----------------------------------------------------------------------===//
//
// This pass allocates SME tiles at the 'func.func' op level for ArmSME
// operations. It does this using a 16-bit tile mask that has a bit for each
// 128-bit element tile (ZA0.Q-ZA15.Q), the smallest ZA tile granule.
//
// The 128-bit tiles overlap with other element tiles as follows (see section
// B2.3.2 of SME spec [1]):
//
//   Tile    Overlaps
//   ---------------------------------------------------------------------------
//   ZA0.B   ZA0.Q, ZA1.Q, ZA2.Q, ZA3.Q, ZA4.Q, ZA5.Q, ZA6.Q, ZA7.Q, ZA8.Q,
//           ZA9.Q, ZA10.Q, ZA11.Q, ZA12.Q, ZA13.Q, ZA14.Q, ZA15.Q
//   ZA0.H   ZA0.Q, ZA2.Q, ZA4.Q, ZA6.Q, ZA8.Q, ZA10.Q, ZA12.Q, ZA14.Q
//   ZA1.H   ZA1.Q, ZA3.Q, ZA5.Q, ZA7.Q, ZA9.Q, ZA11.Q, ZA13.Q, ZA15.Q
//   ZA0.S   ZA0.Q, ZA4.Q, ZA8.Q, ZA12.Q
//   ZA1.S   ZA1.Q, ZA5.Q, ZA9.Q, ZA13.Q
//   ZA2.S   ZA2.Q, ZA6.Q, ZA10.Q, ZA14.Q
//   ZA3.S   ZA3.Q, ZA7.Q, ZA11.Q, ZA15.Q
//   ZA0.D   ZA0.Q, ZA8.Q
//   ZA1.D   ZA1.Q, ZA9.Q
//   ZA2.D   ZA2.Q, ZA10.Q
//   ZA3.D   ZA3.Q, ZA11.Q
//   ZA4.D   ZA4.Q, ZA12.Q
//   ZA5.D   ZA5.Q, ZA13.Q
//   ZA6.D   ZA6.Q, ZA14.Q
//   ZA7.D   ZA7.Q, ZA15.Q
//
// The tiles in use are tracked via a function attribute 'arm_sme.tiles_in_use'
// that is initalized during the first tile allocation within a function and
// updated on each subsequent allocation.
//
// [1] https://developer.arm.com/documentation/ddi0616/aa
//
//===----------------------------------------------------------------------===//

#include "mlir/Dialect/ArmSME/IR/ArmSME.h"
#include "mlir/Dialect/ArmSME/Transforms/Passes.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/TypeSwitch.h"

#define DEBUG_TYPE "allocate-arm-sme-tiles"

namespace mlir {
namespace arm_sme {
#define GEN_PASS_DEF_TILEALLOCATION
#include "mlir/Dialect/ArmSME/Transforms/Passes.h.inc"
} // namespace arm_sme
} // namespace mlir

using namespace mlir;
using namespace mlir::arm_sme;

namespace {

static constexpr StringLiteral kTilesInUseAttr("arm_sme.tiles_in_use");
static constexpr StringLiteral
    kNextInMemoryTileIdAttr("arm_sme.next_in_memory_tile_id");

enum class TileMask : unsigned {
  // clang-format off
  kZA0B  = 0xffff, // 1111 1111 1111 1111

  kZA0H  = 0xaaaa, // 1010 1010 1010 1010
  kZA1H  = 0x5555, // 0101 0101 0101 0101

  kZA0S  = 0x8888, // 1000 1000 1000 1000
  kZA1S  = 0x4444, // 0100 0100 0100 0100
  kZA2S  = 0x2222, // 0010 0010 0010 0010
  kZA3S  = 0x1111, // 0001 0001 0001 0001

  kZA0D  = 0x8080, // 1000 0000 1000 0000
  kZA1D  = 0x4040, // 0100 0000 0100 0000
  kZA2D  = 0x2020, // 0010 0000 0010 0000
  kZA3D  = 0x1010, // 0001 0000 0001 0000
  kZA4D  = 0x808,  // 0000 1000 0000 1000
  kZA5D  = 0x404,  // 0000 0100 0000 0100
  kZA6D  = 0x202,  // 0000 0010 0000 0010
  kZA7D  = 0x101,  // 0000 0001 0000 0001

  kZA0Q  = 0x8000, // 1000 0000 0000 0000
  kZA1Q  = 0x4000, // 0100 0000 0000 0000
  kZA2Q  = 0x2000, // 0010 0000 0000 0000
  kZA3Q  = 0x1000, // 0001 0000 0000 0000
  kZA4Q  = 0x800,  // 0000 1000 0000 0000
  kZA5Q  = 0x400,  // 0000 0100 0000 0000
  kZA6Q  = 0x200,  // 0000 0010 0000 0000
  kZA7Q  = 0x100,  // 0000 0001 0000 0000
  kZA8Q  = 0x80,   // 0000 0000 1000 0000
  kZA9Q  = 0x40,   // 0000 0000 0100 0000
  kZA10Q = 0x20,   // 0000 0000 0010 0000
  kZA11Q = 0x10,   // 0000 0000 0001 0000
  kZA12Q = 0x8,    // 0000 0000 0000 1000
  kZA13Q = 0x4,    // 0000 0000 0000 0100
  kZA14Q = 0x2,    // 0000 0000 0000 0010
  kZA15Q = 0x1,    // 0000 0000 0000 0001

  kNone = 0x0,     // 0000 0000 0000 0000
  // clang-format on

  LLVM_MARK_AS_BITMASK_ENUM(kZA0B)
};

/// Returns the set of masks relevant for the given type.
static ArrayRef<TileMask> getMasks(ArmSMETileType type) {
  static constexpr std::array ZA_B_MASKS = {TileMask::kZA0B};
  static constexpr std::array ZA_H_MASKS = {TileMask::kZA0H, TileMask::kZA1H};
  static constexpr std::array ZA_S_MASKS = {TileMask::kZA0S, TileMask::kZA1S,
                                            TileMask::kZA2S, TileMask::kZA3S};
  static constexpr std::array ZA_D_MASKS = {
      TileMask::kZA0D, TileMask::kZA1D, TileMask::kZA2D, TileMask::kZA3D,
      TileMask::kZA4D, TileMask::kZA5D, TileMask::kZA6D, TileMask::kZA7D};
  static constexpr std::array ZA_Q_MASKS = {
      TileMask::kZA0Q,  TileMask::kZA1Q,  TileMask::kZA2Q,  TileMask::kZA3Q,
      TileMask::kZA4Q,  TileMask::kZA5Q,  TileMask::kZA6Q,  TileMask::kZA7Q,
      TileMask::kZA8Q,  TileMask::kZA9Q,  TileMask::kZA10Q, TileMask::kZA11Q,
      TileMask::kZA12Q, TileMask::kZA13Q, TileMask::kZA14Q, TileMask::kZA15Q};
  switch (type) {
  case ArmSMETileType::ZAB:
    return ZA_B_MASKS;
  case ArmSMETileType::ZAH:
    return ZA_H_MASKS;
  case ArmSMETileType::ZAS:
    return ZA_S_MASKS;
  case ArmSMETileType::ZAD:
    return ZA_D_MASKS;
  case ArmSMETileType::ZAQ:
    return ZA_Q_MASKS;
  }
}

/// Allocates and returns a tile ID. Returns an error if there are no tiles
/// left.
static FailureOr<unsigned> allocateTileId(ArmSMETileType tileType,
                                          TileMask &tilesInUse) {
  auto masks = getMasks(tileType);
  for (auto [tileId, tileMask] : llvm::enumerate(masks)) {
    if ((tilesInUse & tileMask) == TileMask::kNone) {
      tilesInUse |= tileMask;
      return tileId;
    }
  }
  return failure();
}

/// Collects transitive uses of a root value through control flow. This can
/// handle basic SCF constructs, along with control flow (br and cond_br).
/// Simple loops work at the SCF level, while more complex control flow can be
/// dealt with after lowering to CF. This is used to implement basic tile
/// allocation.
static void findDependantOps(Value rootValue,
                             SetVector<Operation *> &dependantOps) {
  auto traverseCorrespondingValues = [&](auto inputValues, auto exitValues) {
    for (auto [idx, value] : llvm::enumerate(inputValues)) {
      if (value == rootValue)
        findDependantOps(exitValues[idx], dependantOps);
    }
  };
  for (Operation *user : rootValue.getUsers()) {
    if (dependantOps.contains(user))
      continue;
    dependantOps.insert(user);
    TypeSwitch<Operation *>(user)
        .Case<cf::BranchOp>([&](auto branchOp) {
          // (CF) Follow branch.
          traverseCorrespondingValues(branchOp.getDestOperands(),
                                      branchOp.getDest()->getArguments());
        })
        .Case<cf::CondBranchOp>([&](auto condBranchOp) {
          // (CF) Follow true branch.
          traverseCorrespondingValues(
              condBranchOp.getTrueOperands(),
              condBranchOp.getTrueDest()->getArguments());
          // (CF) Follow false branch.
          traverseCorrespondingValues(
              condBranchOp.getFalseOperands(),
              condBranchOp.getFalseDest()->getArguments());
        })
        .Case<LoopLikeOpInterface>([&](auto loopOp) {
          // (SCF) Follow iter_args of (basic) loops (e.g. for loops).
          traverseCorrespondingValues(loopOp.getInits(),
                                      loopOp.getRegionIterArgs());
        })
        .Case<scf::YieldOp>([&](auto yieldOp) {
          // (SCF) Follow yields of (basic) control flow (e.g. for loops).
          auto parent = user->getParentOp();
          traverseCorrespondingValues(user->getOperands(),
                                      parent->getResults());
        })
        .Default([&](auto) {
          // Otherwise, assume users of _any_ result are dependant.
          for (Value result : user->getResults())
            findDependantOps(result, dependantOps);
        });
  }
}
struct AssignTileIDsPattern
    : public OpInterfaceRewritePattern<ArmSMETileOpInterface> {
  using OpInterfaceRewritePattern::OpInterfaceRewritePattern;
  LogicalResult matchAndRewrite(ArmSMETileOpInterface tileOp,
                                PatternRewriter &rewriter) const override {
    if (tileOp.getTileId())
      return failure();

    auto func = tileOp->getParentOfType<FunctionOpInterface>();
    auto getDiscardableIntAttr = [&](StringRef name, unsigned defaultVal = 0) {
      if (auto attr = llvm::dyn_cast_or_null<IntegerAttr>(
              func->getDiscardableAttr(name)))
        return unsigned(attr.getInt());
      return defaultVal;
    };
    auto setDiscardableIntAttr = [&](StringRef name, auto value) {
      rewriter.modifyOpInPlace(tileOp, [&] {
        func->setDiscardableAttr(name,
                                 rewriter.getI32IntegerAttr((unsigned)value));
      });
    };

    std::optional<ArmSMETileType> tileType = tileOp.getAllocatedTileType();
    if (!tileType)
      return rewriter.notifyMatchFailure(tileOp, "op does not allocate a tile");

    TileMask tilesInUse =
        static_cast<TileMask>(getDiscardableIntAttr(kTilesInUseAttr));
    auto tileId = allocateTileId(*tileType, tilesInUse);
    bool tileIsInMemory = failed(tileId);
    if (tileIsInMemory) {
      // If we could not find a real tile ID, use an in-memory tile ID (ID >=
      // 16). A later pass will insert the necessary spills and reloads.
      tileId =
          getDiscardableIntAttr(kNextInMemoryTileIdAttr, kInMemoryTileIdBase);
      tileOp->emitWarning(
          "failed to allocate SME virtual tile to operation, all tile "
          "operations will go through memory, expect degraded performance");
    }

    // Set all operations dependent on `tileOp` to use the same tile ID.
    // This is a naive tile allocation scheme, but works for common cases. For
    // example, as this only allocates tile IDs to existing ops, it can't solve
    // cases like this (%tileA and %tileB come from different root operations):
    //
    // %tile = scf.if %some_cond -> vector<[4]x[4]xi32> {
    //   scf.yield %tileA {tile_id = 0} : vector<[4]x[4]xi32>
    // } else {
    //   scf.yield %tileB {tile_id = 1} : vector<[4]x[4]xi32>
    // }
    //
    // This case would require allocating a new tile for the result of the
    // scf.if, and moving the contents of %tileA or %tileB to result tile (based
    // on the %some_cond).
    // Find all the ops that (transitively) depend on this tile.
    SetVector<Operation *> dependantOps;
    findDependantOps(tileOp->getResult(0), dependantOps);
    auto tileIDAttr = rewriter.getI32IntegerAttr(*tileId);
    for (auto *op : dependantOps) {
      if (auto dependantTileOp = llvm::dyn_cast<ArmSMETileOpInterface>(op)) {
        auto currentTileId = dependantTileOp.getTileId();
        if (currentTileId && unsigned(currentTileId.getInt()) != tileId)
          return dependantTileOp.emitOpError(
              "already assigned different SME virtual tile!");
      }
    }

    // Rewrite IR.
    if (!tileIsInMemory)
      setDiscardableIntAttr(kTilesInUseAttr, tilesInUse);
    else
      setDiscardableIntAttr(kNextInMemoryTileIdAttr, *tileId + 1);
    rewriter.modifyOpInPlace(tileOp, [&] { tileOp.setTileId(tileIDAttr); });
    for (auto *op : dependantOps) {
      if (auto dependantTileOp = llvm::dyn_cast<ArmSMETileOpInterface>(op)) {
        rewriter.modifyOpInPlace(
            dependantTileOp, [&] { dependantTileOp.setTileId(tileIDAttr); });
      }
    }

    return success();
  }
};

struct TileAllocationPass
    : public arm_sme::impl::TileAllocationBase<TileAllocationPass> {
  void runOnOperation() override {
    RewritePatternSet patterns(&getContext());
    patterns.add<AssignTileIDsPattern>(patterns.getContext());
    GreedyRewriteConfig config;
    // Setting useTopDownTraversal ensures tiles are allocated in program
    // order.
    config.useTopDownTraversal = true;
    if (mlir::failed(mlir::applyPatternsAndFoldGreedily(
            getOperation(), std::move(patterns), config))) {
      signalPassFailure();
    }
  }
};
} // namespace

std::unique_ptr<Pass> mlir::arm_sme::createTileAllocationPass() {
  return std::make_unique<TileAllocationPass>();
}