AIECtrlPacketToDma.cpp

Go to the documentation of this file.

//===- AIECtrlPacketToDma.cpp -----------------------------------*- C++ -*-===//
//
// Copyright (C) 2024 Advanced Micro Devices, Inc.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
 
#include "aie/Dialect/AIE/IR/AIEDialect.h"
#include "aie/Dialect/AIE/Transforms/AIEGenerateColumnControlOverlay.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/IRMapping.h"
#include "mlir/Pass/Pass.h"
 
#include "llvm/ADT/TypeSwitch.h"
 
namespace xilinx::AIEX {
#define GEN_PASS_DEF_AIECTRLPACKETTODMA
#define GEN_PASS_DEF_AIECTRLPACKETINFERTILES
#include "aie/Dialect/AIEX/Transforms/AIEXPasses.h.inc"
} // namespace xilinx::AIEX
 
#define DEBUG_TYPE "aie-ctrl-packet-to-dma"
 
using namespace mlir;
using namespace xilinx;
using namespace xilinx::AIE;
using namespace xilinx::AIEX;
 
struct AIECtrlPacketInferTilesPass
    : xilinx::AIEX::impl::AIECtrlPacketInferTilesBase<
          AIECtrlPacketInferTilesPass> {
  void runOnOperation() override {
    DeviceOp device = getOperation();
    const auto &targetModel = device.getTargetModel();
    OpBuilder devBuilder = OpBuilder::atBlockBegin(device.getBody());
 
    auto sequenceOps = device.getOps<AIE::RuntimeSequenceOp>();
    for (auto f : sequenceOps) {
      auto ctrlPktOps = f.getOps<AIEX::NpuControlPacketOp>();
      for (auto ctrlPktOp : ctrlPktOps) {
        auto tOp = TileOp::getOrCreate(devBuilder, device,
                                       (int)ctrlPktOp.getColumnFromAddr(),
                                       (int)ctrlPktOp.getRowFromAddr());
        // Assign controller id
        auto tileIDMap = getTileToControllerIdMap(true, targetModel);
        if (tOp->hasAttr("controller_id"))
          continue;
        auto pktInfoAttr = AIE::PacketInfoAttr::get(
            tOp->getContext(), /*pkt_type*/ 0,
            /*pkt_id*/ tileIDMap[{tOp.colIndex(), tOp.rowIndex()}]);
        tOp->setAttr("controller_id", pktInfoAttr);
      }
    }
  }
};
 
struct AIECtrlPacketToDmaPass
    : xilinx::AIEX::impl::AIECtrlPacketToDmaBase<AIECtrlPacketToDmaPass> {
  void runOnOperation() override {
    DeviceOp device = getOperation();
    const auto &targetModel = device.getTargetModel();
    auto *ctx = device->getContext();
    auto loc = device->getLoc();
 
    if (targetModel.getTargetArch() == AIEArch::AIE1)
      return; // Disable this pass for AIE1; AIE1 support NYI.
 
    SmallVector<Operation *> erased;
    auto sequenceOps = device.getOps<AIE::RuntimeSequenceOp>();
    for (auto f : sequenceOps) {
 
      auto controlPacketOps = f.getOps<AIEX::NpuControlPacketOp>();
      if (controlPacketOps.empty())
        continue;
 
      OpBuilder builder(f);
 
      IRMapping mapping;
 
      auto newSeq = AIE::RuntimeSequenceOp::create(
          builder, loc, f.getSymNameAttr(), BoolAttr{}, f.getTraceBufferAttr(),
          f.getTraceSlicesAttr(), f.getScratchpadParametersAttr());
      newSeq.getBody().push_back(new Block);
 
      // Copy the arguments from the old sequence to the new one.
      for (auto arg : f.getBody().getArguments()) {
        // Add the argument to the new sequence.
        auto newArg = newSeq.getBody().addArgument(arg.getType(), arg.getLoc());
        // Replace all uses of the old argument with the new one.
        arg.replaceAllUsesWith(newArg);
        // Add the mapping for the argument.
        mapping.map(arg, newArg);
      }
 
      // Using dynamic shape for ctrl pkt stream.
      auto ctrlPktMemrefType = MemRefType::get(
          ShapedType::kDynamic, IntegerType::get(ctx, 32), nullptr, nullptr);
      auto newBlockArg = newSeq.getBody().addArgument(ctrlPktMemrefType, loc);
 
      builder.setInsertionPointToStart(&newSeq.getBody().front());
 
      // Collect all npu.control_packet ops, grouped by location in 'batches'
      struct BatchInfo {
        TileID tileId;
        int64_t startOffset;
        int64_t totalSize;
        std::string shimDmaAllocName;
        int shimChan;
        Operation *first;
      };
      std::vector<BatchInfo> batches;
 
      int64_t ddrOffset = 0;
      Block &entry = f.getBody().front();
 
      // First pass: collect and batch control packet operations
      bool new_batch = true;
      for (Operation &o : entry) {
        auto ctrlPktOp = dyn_cast<NpuControlPacketOp>(&o);
 
        // A non-control_packet op ends the current batch
        if (!ctrlPktOp) {
          new_batch = true;
          continue;
        }
        int col = ctrlPktOp.getColumnFromAddr();
        int row = ctrlPktOp.getRowFromAddr();
 
        // Calculate control packet size
        int64_t ctrlPktSize = 0;
        auto data = ctrlPktOp.getData();
        auto length = ctrlPktOp.getLength();
        if (data)
          ctrlPktSize = data->size();
        else if (length)
          ctrlPktSize = *length;
        ctrlPktSize++; // Ctrl info word
        ctrlPktSize++; // Packet header
 
        // Check if we can batch with the previous packet
        if (targetModel.getTargetArch() == AIEArch::AIE2p && !new_batch &&
            batches.back().tileId == TileID{col, row}) {
          // Add to existing batch
          batches.back().totalSize += ctrlPktSize;
        } else {
          // Start new batch
          auto rowToShimChanMap =
              getRowToShimChanMap(targetModel, WireBundle::DMA);
          int shimChan = rowToShimChanMap[row];
 
          std::string shimDmaAllocName = "ctrlpkt";
          shimDmaAllocName += "_col" + std::to_string(col);
          shimDmaAllocName += "_mm2s";
          shimDmaAllocName += "_chan" + std::to_string(shimChan);
 
          batches.push_back({TileID{col, row}, ddrOffset, ctrlPktSize,
                             shimDmaAllocName, shimChan, &o});
          new_batch = false;
        }
        ddrOffset += ctrlPktSize;
      }
 
      // Second pass: emit batched operations in original order
      auto batchIt = batches.begin();
 
      for (Operation &o : entry) {
        auto ctrlPktOp = dyn_cast<NpuControlPacketOp>(&o);
        if (!ctrlPktOp) {
          builder.clone(o, mapping);
          continue;
        }
 
        // There are no more control packet batches to emit
        if (batchIt == batches.end())
          continue;
 
        // Check if this is the first packet of a new batch, otherwise skip it.
        if (batchIt->first != &o)
          continue;
 
        int col = ctrlPktOp.getColumnFromAddr();
 
        // Emit the batched DMA operation for this (col, row) pair
        const std::vector<int64_t> staticOffsets = {0, 0, 0,
                                                    batchIt->startOffset};
        const std::vector<int64_t> staticSizes = {1, 1, 1, batchIt->totalSize};
        const std::vector<int64_t> staticStrides = {0, 0, 0, 1};
 
        SymbolRefAttr metadata =
            SymbolRefAttr::get(builder.getContext(), batchIt->shimDmaAllocName);
        NpuDmaMemcpyNdOp::create(builder, loc, newBlockArg, staticOffsets,
                                 staticSizes, staticStrides, metadata,
                                 /*id=*/0, /*issue_token=*/true);
 
        Value shimRow = AIEX::createConstantI32(builder, loc, 0);
        Value shimCol = AIEX::createConstantI32(builder, loc, col);
        Value dir = AIEX::createConstantI32(builder, loc, 1); // MM2S
        Value chan = AIEX::createConstantI32(builder, loc, batchIt->shimChan);
        Value col_num = AIEX::createConstantI32(builder, loc, 1);
        Value row_num = AIEX::createConstantI32(builder, loc, 1);
        AIEX::NpuSyncOp::create(builder, loc, shimCol, shimRow, dir, chan,
                                col_num, row_num);
        ++batchIt;
      }
 
      erased.push_back(f);
    }
 
    for (auto *e : erased)
      e->erase();
  }
};
 
std::unique_ptr<OperationPass<DeviceOp>>
AIEX::createAIECtrlPacketInferTilesPass() {
  return std::make_unique<AIECtrlPacketInferTilesPass>();
}
std::unique_ptr<OperationPass<DeviceOp>> AIEX::createAIECtrlPacketToDmaPass() {
  return std::make_unique<AIECtrlPacketToDmaPass>();
}