AIETargetShared.cpp

Go to the documentation of this file.

//===- AIETargetShared.cpp --------------------------------------*- C++ -*-===//
//
// Copyright (C) 2021-2022 Xilinx, Inc.
// Copyright (C) 2022-2023 Advanced Micro Devices, Inc.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//
//===----------------------------------------------------------------------===//
 
#include "aie/Targets/AIETargetShared.h"
 
#include "aie/Dialect/AIE/IR/AIEDialect.h"
#include "aie/Dialect/AIEX/IR/AIEXDialect.h"
#include "aie/Targets/AIETargets.h"
 
#include "mlir/Target/LLVMIR/Import.h"
 
#include "llvm/ADT/StringExtras.h"
 
using namespace mlir;
using namespace xilinx;
using namespace xilinx::AIE;
using namespace xilinx::AIEX;
 
namespace xilinx::AIE {
 
std::string tileLocStr(StringRef col, StringRef row) {
  std::string str;
  llvm::raw_string_ostream rss(str);
  rss << "XAie_TileLoc(" << col << "," << row << ")";
  return str;
}
 
std::string tileLocStr(int col, int row) {
  return tileLocStr(std::to_string(col), std::to_string(row));
}
 
std::string tileDMAInstStr(StringRef col, StringRef row, StringRef bdNum) {
  std::string str;
  llvm::raw_string_ostream rss(str);
  rss << "dma_tile" << col << row << "_bd" << bdNum;
  return str;
}
 
std::string tileDMAInstStr(int col, int row, int bdNum) {
  return tileDMAInstStr(std::to_string(col), std::to_string(row),
                        std::to_string(bdNum));
}
 
std::string tileDMAInstRefStr(StringRef col, StringRef row, StringRef bdNum) {
  std::string str;
  llvm::raw_string_ostream rss(str);
  rss << "&(" << tileDMAInstStr(col, row, bdNum) << ")";
  return str;
}
 
std::string tileDMAInstRefStr(int col, int row, int bdNum) {
  return tileDMAInstRefStr(std::to_string(col), std::to_string(row),
                           std::to_string(bdNum));
}
 
std::string packetStr(StringRef id, StringRef type) {
  std::string str;
  llvm::raw_string_ostream rss(str);
  rss << "XAie_PacketInit(" << id << "," << type << ")";
  return str;
}
 
std::string packetStr(int id, int type) {
  return packetStr(std::to_string(id), std::to_string(type));
}
 
static std::string tileDMATensorStr(StringRef col, StringRef row,
                                    StringRef bdNum) {
  std::string str;
  llvm::raw_string_ostream rss(str);
  rss << "dma_tile_" << col << "_" << row << "_bd_" << bdNum << "_tensor";
  return str;
}
 
static std::string tileDMATensorStr(int col, int row, int bdNum) {
  return tileDMATensorStr(std::to_string(col), std::to_string(row),
                          std::to_string(bdNum));
}
 
void generateXAieDmaSetMultiDimAddr(raw_ostream &output, int ndims,
                                    ArrayRef<BDDimLayoutAttr> dims, int col,
                                    int row, int bdNum, int baseAddrA,
                                    int64_t offsetA, uint32_t lenA,
                                    int elementWidthInBytes,
                                    const char *errorRetval) {
  // libxaie requires stride in multiples of 32b
  double elementWidthIn32bWords =
      static_cast<double>(elementWidthInBytes) / 4.0;
  std::string tensor = tileDMATensorStr(col, row, bdNum);
  output << "XAie_DmaTensor " << tensor << " = {};\n";
  output << tensor << ".NumDim = " << std::to_string(ndims) << ";\n";
  output << tensor
         << ".Dim ="
            "__mlir_aie_alloc_dim_desc("
         << std::to_string(ndims) << ");\n";
  output << "if(NULL == " << tensor << ".Dim){\n"
         << "  return " << errorRetval << ";\n"
         << "}\n";
  for (size_t i = 0; i < dims.size(); i++) {
    uint16_t size;
    uint32_t stride;
    // Pass down dimensions in reverse order; in the MLIR, this allows us
    // to specify strides/sizes in the same order as we would access a
    // multi-dim C array, with the highest dimension first.
    int j = dims.size() - i - 1;
    if (j > 0) {
      stride =
          static_cast<uint32_t>(dims[i].getStride() * elementWidthIn32bWords);
      size = dims[i].getSize();
    } else {
      stride = dims[i].getStride();
      size = static_cast<uint16_t>(dims[i].getSize() * elementWidthIn32bWords);
    }
    stride = stride > 0 ? stride : 1;
    // Assume AIE-ML architecture; we assert this above
    output << tensor << ".Dim[" << std::to_string(j) << "].AieMlDimDesc"
           << " = { /* Stride */ " << std::to_string(stride) << ", /* Size */ "
           << std::to_string(size) << "};\n";
  }
  if ((baseAddrA + offsetA) % 4) {
    llvm::report_fatal_error("bd address must be 4B (32b) aligned");
  }
  output << "__mlir_aie_try(XAie_DmaSetMultiDimAddr("
         << tileDMAInstRefStr(col, row, bdNum) << ", " << "&" << tensor << ", "
         << "0x" << llvm::utohexstr(baseAddrA + offsetA) << ", "
         << " /* len */ " << lenA << "));\n";
}
 
// Traverse through a chain of blocks.
llvm::SetVector<Block *> getOrderedChainOfBlocks(Region *region) {
  // Get the region's entry block, then start traversing through the chain of
  // blocks.
  llvm::SetVector<Block *> blockVector;
  SmallVector<Block *, 16> worklist;
  Block *firstBD = &region->front();
  blockVector.insert(firstBD);
  worklist.push_back(firstBD);
  while (!worklist.empty()) {
    Block *block = worklist.pop_back_val();
    if (block->empty()) {
      continue;
    }
    auto successors = block->getTerminator()->getSuccessors();
    for (auto *i : successors) {
      if (!blockVector.contains(i)) {
        blockVector.insert(i);
        worklist.push_back(i);
      }
    }
  }
  return blockVector;
}
 
llvm::SmallPtrSet<Block *, 8>
collectUnchainedBlocks(const llvm::SetVector<Block *> &blockVector) {
  // Scoped to this channel only.
  llvm::SmallPtrSet<Block *, 8> channelHeads;
  for (Block *block : blockVector) {
    for (auto startOp : block->getOps<DMAStartOp>()) {
      channelHeads.insert(startOp.getDest());
    }
  }
 
  llvm::SmallPtrSet<Block *, 8> unchainedBlocks;
  for (Block *block : blockVector) {
    for (auto startOp : block->getOps<DMAStartOp>()) {
      if (!startOp.getOutOfOrder() && !startOp.getCorePushes()) {
        continue;
      }
      Block *b = startOp.getDest();
      while (b && unchainedBlocks.insert(b).second &&
             b->getNumSuccessors() > 0) {
        Block *next = b->getSuccessor(0);
        // Another channel's head or this own channel's head.
        if (channelHeads.contains(next)) {
          break;
        }
        b = next;
      }
    }
  }
  return unchainedBlocks;
}
 
std::optional<std::vector<char>>
denseAttrToBytes(mlir::DenseElementsAttr denseInit) {
  mlir::Type elemType = denseInit.getElementType();
  if (!elemType.isIntOrIndex() && !llvm::isa<mlir::FloatType>(elemType)) {
    return std::nullopt;
  }
 
  std::vector<char> bytes;
  // A float reaches its byte image through APInt too, so one loop covers both.
  // The APInt object aliases its own value only for a width that fits the
  // inline union, so this reads getRawData() instead of &value.
  auto append = [&bytes](const llvm::APInt &value) {
    size_t byteSize = (value.getBitWidth() + 7) / 8;
    const auto *first = reinterpret_cast<const char *>(value.getRawData());
    bytes.insert(bytes.end(), first, first + byteSize);
  };
  if (elemType.isIntOrIndex()) {
    for (const llvm::APInt &intVal : denseInit.getValues<llvm::APInt>()) {
      append(intVal);
    }
  } else {
    for (const llvm::APFloat &floatVal : denseInit.getValues<llvm::APFloat>()) {
      append(floatVal.bitcastToAPInt());
    }
  }
  return bytes;
}
 
MemoryRun coreDataRegion(TileOp tile, llvm::ArrayRef<BufferOp> buffers) {
  BufferOp coreData;
  for (auto buf : buffers) {
    if (buf.getCoreData()) {
      coreData = buf;
    }
  }
  if (coreData) {
    if (std::optional<int32_t> addr = coreData.getAddress()) {
      return MemoryRun{*addr, (int64_t)coreData.getAllocationSize()};
    }
  }
 
  const auto &targetModel = getTargetModel(tile);
  CoreOp core = tile.getCoreOp();
  llvm::SmallVector<std::pair<int64_t, int64_t>> occupied;
  MemoryRun stackRun = core ? core.getStackRun() : MemoryRun{};
  occupied.emplace_back(stackRun.start, stackRun.end());
  for (auto buf : buffers) {
    int64_t base = getBufferBaseAddress(buf);
    occupied.emplace_back(base, base + buf.getAllocationSize());
  }
  return largestFreeRun(
      targetModel.getLocalMemorySize(), std::move(occupied),
      std::max<int64_t>(targetModel.getComputeTileMaxVectorAlignBits() / 8, 1));
}
 
llvm::SmallVector<MemoryRun> coreBankRegions(TileOp tile,
                                             llvm::ArrayRef<BufferOp> buffers,
                                             MemoryRun dataRun) {
  const auto &targetModel = getTargetModel(tile);
  int numBanks = targetModel.getNumBanks(tile.getCol(), tile.getRow());
  llvm::SmallVector<MemoryRun> regions;
  if (numBanks <= 0) {
    return regions;
  }
  int64_t memSize = targetModel.getLocalMemorySize();
  int64_t bankSize = memSize / numBanks;
  int64_t align =
      std::max<int64_t>(targetModel.getComputeTileMaxVectorAlignBits() / 8, 1);
 
  CoreOp core = tile.getCoreOp();
  llvm::SmallVector<std::pair<int64_t, int64_t>> occupied;
  MemoryRun stackRun = core ? core.getStackRun() : MemoryRun{};
  occupied.emplace_back(stackRun.start, stackRun.end());
  // A reservation held for a bank *is* that bank's region: the allocator set it
  // aside for precisely the sections this function is sizing a region for.
  // Counting it as occupied would hand the region back the space it was told to
  // keep, and hand the section a hole somewhere else -- so it may be the
  // largest free run in the bank while the reservation sits in a smaller one.
  llvm::SmallVector<std::optional<MemoryRun>> reserved(numBanks, std::nullopt);
  for (auto buf : buffers) {
    int64_t base = getBufferBaseAddress(buf);
    if (auto bank = buf.getMemBank();
        buf.getBankReserved() && bank && *bank < numBanks) {
      reserved[*bank] = MemoryRun{base, buf.getAllocationSize()};
      continue;
    }
    occupied.emplace_back(base, base + buf.getAllocationSize());
  }
  // The unpinned region is spoken for, so a bank only offers what it leaves.
  occupied.emplace_back(dataRun.start, dataRun.end());
 
  for (int bank = 0; bank < numBanks; ++bank) {
    if (auto region = reserved[bank]) {
      regions.push_back(*region);
      continue;
    }
    MemoryRun window{bankSize * bank, bankSize};
    MemoryRun run = largestFreeRunIn(window, occupied, align);
    // A bank with nothing spare still has to name its own base. Left at the
    // default the region would report address 0, which is inside bank 0.
    if (run.size == 0) {
      run.start = window.start;
    }
    regions.push_back(run);
  }
  return regions;
}
 
} // namespace xilinx::AIE