AIECreateLocks.cpp

Go to the documentation of this file.

//===- AIECreateLocks.cpp ---------------------------------------*- C++ -*-===//
//
// Copyright (C) 2019-2022 Xilinx, Inc.
// Copyright (C) 2022-2026 Advanced Micro Devices, Inc.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
 
#include "aie/Dialect/AIE/IR/AIEDialect.h"
#include "aie/Dialect/AIEX/AIETokenAnalysis.h"
#include "aie/Dialect/AIEX/IR/AIEXDialect.h"
#include "aie/Dialect/AIEX/Transforms/AIEXPasses.h"
 
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Tools/mlir-translate/MlirTranslateMain.h"
#include "mlir/Transforms/DialectConversion.h"
 
namespace xilinx::AIEX {
#define GEN_PASS_DEF_AIECREATELOCKS
#include "aie/Dialect/AIEX/Transforms/AIEXPasses.h.inc"
} // namespace xilinx::AIEX
 
#define DEBUG_TYPE "aie-create-locks"
 
using namespace mlir;
using namespace xilinx;
using namespace xilinx::AIE;
using namespace xilinx::AIEX;
 
struct Token2LockLowering : public OpConversionPattern<UseTokenOp> {
  using OpConversionPattern<UseTokenOp>::OpConversionPattern;
  DenseMap<Operation *, std::vector<std::pair<Value, int>>> &acqLocks;
  DenseMap<Operation *, std::vector<std::pair<Value, int>>> &relLocks;
  DenseMap<std::pair<Operation *, Operation *>, std::pair<Value, int>>
      &lockChains;
 
  Token2LockLowering(
      MLIRContext *context,
      DenseMap<Operation *, std::vector<std::pair<Value, int>>> &acqLocks,
      DenseMap<Operation *, std::vector<std::pair<Value, int>>> &relLocks,
      DenseMap<std::pair<Operation *, Operation *>, std::pair<Value, int>>
          &lockChains,
      PatternBenefit benefit = 1)
      : OpConversionPattern<UseTokenOp>(context, benefit), acqLocks(acqLocks),
        relLocks(relLocks), lockChains(lockChains) {}
 
  LogicalResult
  matchAndRewrite(UseTokenOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Operation *Op = op.getOperation();
    Operation *parentOp = op->getParentOp();
 
    if (!isa<CoreOp, MemOp, ShimDMAOp>(parentOp))
      llvm_unreachable("A parent operation of UseTokenOp must be either CoreOp "
                       "or MemOp or ShimDMAOp");
 
    if (op.acquire()) {
      // Acquire lock from pair
      LLVM_DEBUG(llvm::dbgs() << "Replacing Acquire: " << op << "\n");
      for (auto acqLock : acqLocks[op]) {
        Value lockFromPair = acqLock.first;
        int lockValueFromPair = acqLock.second;
        UseLockOp::create(rewriter, op.getLoc(), lockFromPair,
                          LockAction::Acquire, lockValueFromPair);
        LLVM_DEBUG(llvm::dbgs() << "Acquire from pair " << lockFromPair
                                << " with " << lockValueFromPair << "\n");
      }
 
      // Acquire lock from chain
      for (auto map : lockChains) {
        Value lockFromChain = map.second.first;
        int lockValueFromChain = map.second.second;
        Operation *acqOp = map.first.second;
        if (acqOp != Op)
          continue;
 
        UseLockOp::create(rewriter, op.getLoc(), lockFromChain,
                          LockAction::Acquire, lockValueFromChain);
        LLVM_DEBUG(llvm::dbgs() << "Acquire from chain " << lockFromChain
                                << " with " << lockValueFromChain << "\n");
      }
    } else if (op.release()) {
      // Release lock from pair
      LLVM_DEBUG(llvm::dbgs() << "Replacing Release: " << op << "\n");
      for (auto relLock : relLocks[op]) {
        Value lockFromPair = relLock.first;
        int lockValueFromPair = relLock.second;
        UseLockOp::create(rewriter, op.getLoc(), lockFromPair,
                          LockAction::Release, lockValueFromPair);
        LLVM_DEBUG(llvm::dbgs() << "Release from pair " << lockFromPair
                                << " with " << lockValueFromPair << "\n");
      }
 
      // Release lock from chain
      for (auto map : lockChains) {
        Value lockFromChain = map.second.first;
        int lockValueFromChain = map.second.second;
        Operation *relOp = map.first.first;
        if (relOp != Op)
          continue;
 
        UseLockOp::create(rewriter, op.getLoc(), lockFromChain,
                          LockAction::Release, lockValueFromChain);
        LLVM_DEBUG(llvm::dbgs() << "Release from chain " << lockFromChain
                                << " with " << lockValueFromChain << "\n");
      }
    }
 
    rewriter.eraseOp(Op);
 
    return success();
  }
};
 
static int getLockID(DenseMap<std::pair<Operation *, int>, int> &locks,
                     Operation *op) {
  auto tileOp = cast<TileOp>(op);
  const auto &targetModel = xilinx::AIE::getTargetModel(op);
  for (unsigned i = 0;
       i < targetModel.getNumLocks(tileOp.getCol(), tileOp.getRow()); i++) {
    int usageCnt = locks[std::make_pair(tileOp, i)];
    if (usageCnt == 0) {
      locks[std::make_pair(tileOp, i)] = 1;
      return i;
    }
  }
 
  return -1;
}
 
struct AIECreateLocksPass
    : public xilinx::AIEX::impl::AIECreateLocksBase<AIECreateLocksPass> {
  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<arith::ArithDialect>();
  }
  void runOnOperation() override {
 
    DeviceOp device = getOperation();
    OpBuilder builder = OpBuilder::atBlockTerminator(device.getBody());
 
    TokenAnalysis TA(device);
    TA.runAnalysis();
    LLVM_DEBUG(TA.print(llvm::dbgs()));
    SmallVector<std::pair<Operation *, Operation *>, 4> tokenChains(
        TA.getTokenChains());
    SmallVector<std::pair<Operation *, Operation *>, 4> tokenPairs(
        TA.getTokenPairs());
    DenseMap<TileID, Operation *> tiles(TA.getTiles());
 
    DenseMap<std::pair<Operation *, int>, int> locks;
    DenseMap<std::pair<Operation *, Operation *>, std::pair<Value, int>>
        lockChains;
    DenseMap<Operation *, std::vector<std::pair<Value, int>>> acqLocks;
    DenseMap<Operation *, std::vector<std::pair<Value, int>>> relLocks;
 
    for (auto map : tokenChains) {
      Operation *release = map.first;
      Operation *acquire = map.second;
 
      Operation *relUser = TA.getTokenUserOp(release);
      Operation *acqUser = TA.getTokenUserOp(acquire);
      bool IsRelUserCore = isa<CoreOp>(relUser);
      bool IsAcqUserCore = isa<CoreOp>(acqUser);
 
      Operation *tileOp = TA.getShareableTileOp(relUser, acqUser);
 
      LLVM_DEBUG(llvm::dbgs() << "\n=== CHECKING TOKEN CHAIN ===\n";
                 release->print(llvm::dbgs());
                 if (IsRelUserCore) llvm::dbgs() << " @Core";
                 else llvm::dbgs() << " @DMA";
                 llvm::dbgs() << " (" << TA.getCoord(relUser).col << ", "
                              << TA.getCoord(relUser).row << ")" << '\n';
                 acquire->print(llvm::dbgs());
                 if (IsAcqUserCore) llvm::dbgs() << " @Core";
                 else llvm::dbgs() << " @DMA";
                 llvm::dbgs() << " (" << TA.getCoord(acqUser).col << ", "
                              << TA.getCoord(acqUser).row << ")" << '\n';);
 
      // ignore chain that involves a MemOp (DMA) user and CoreOp user and they
      // don't have a shareable tile. This might be caused by MemcpyOp lowering
      // -- there are two MemOps that use the same token and the same lock
      // action + value. Therefore, TokenAnalysis accidentally chains one MemOp
      // to a Core (from the MemcpyOp relationship) that does not have memory
      // affinity with it
      // TODO: verify if it is actually safe to ignore this case
      if (!tileOp && ((!IsRelUserCore && IsAcqUserCore) ||
                      (!IsAcqUserCore && IsRelUserCore)))
        continue;
      assert(tileOp &&
             "the lock users of this chain do not have a common lock");
 
      TileOp tile = cast<TileOp>(tileOp);
      int lockID = getLockID(locks, tileOp);
      assert(lockID >= 0 && "No more locks to allocate!");
      LLVM_DEBUG(llvm::dbgs() << "Shared tile \n"; tileOp->print(llvm::dbgs()));
      LLVM_DEBUG(llvm::dbgs() << " LockID: " << lockID << '\n');
      builder.setInsertionPointAfter(tileOp);
      LockOp lock = LockOp::create(builder, tile.getLoc(), tile, lockID, 0);
 
      lockChains[std::make_pair(release, acquire)] = std::make_pair(lock, 1);
 
      for (auto pair : tokenPairs) {
        Operation *acqFromPair = pair.first;
        Operation *relFromPair = pair.second;
 
        if (relFromPair == release)
          acqLocks[acqFromPair].emplace_back(lock, 0);
 
        if (acqFromPair == acquire)
          relLocks[relFromPair].emplace_back(lock, 0);
      }
    }
 
    ConversionTarget target(getContext());
    target.addLegalOp<UseLockOp>();
 
    RewritePatternSet patterns(&getContext());
    patterns.insert<Token2LockLowering>(device.getContext(), acqLocks, relLocks,
                                        lockChains);
 
    if (failed(applyPartialConversion(device, target, std::move(patterns))))
      signalPassFailure();
  }
};
 
std::unique_ptr<OperationPass<DeviceOp>>
xilinx::AIEX::createAIECreateLocksPass() {
  return std::make_unique<AIECreateLocksPass>();
}