AIEInlineTraceConfig.cpp

Go to the documentation of this file.

//===- AIEInlineTraceConfig.cpp ---------------------------------*- C++ -*-===//
//
// Copyright (C) 2025 Advanced Micro Devices, Inc.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
// Pass to inline trace.start_config and generate npu.write32
//===----------------------------------------------------------------------===//
 
#include "aie/Dialect/AIE/IR/AIEDialect.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/Pass/Pass.h"
 
namespace xilinx::AIEX {
#define GEN_PASS_DEF_AIEXINLINETRACECONFIG
#include "aie/Dialect/AIEX/Transforms/AIEXPasses.h.inc"
} // namespace xilinx::AIEX
 
using namespace mlir;
using namespace xilinx;
using namespace xilinx::AIE;
using namespace xilinx::AIEX;
 
namespace {
 
struct AIEInlineTraceConfigPass
    : xilinx::AIEX::impl::AIEXInlineTraceConfigBase<AIEInlineTraceConfigPass> {
  void runOnOperation() override {
    AIE::DeviceOp device = getOperation();
    const auto &targetModel = device.getTargetModel();
 
    // Collect all trace.start_config operations
    SmallVector<TraceStartConfigOp> startConfigs;
    device.walk([&](TraceStartConfigOp startConfig) {
      startConfigs.push_back(startConfig);
    });
 
    for (auto startConfig : startConfigs) {
      OpBuilder builder(startConfig);
 
      // Lookup the trace config symbol
      auto configSymbolName = startConfig.getTraceConfig();
      auto configOp =
          dyn_cast_or_null<TraceConfigOp>(SymbolTable::lookupNearestSymbolFrom(
              device, builder.getStringAttr(configSymbolName)));
 
      if (!configOp) {
        startConfig.emitError("trace config symbol '")
            << configSymbolName << "' not found";
        return signalPassFailure();
      }
 
      // Get tile and extract col/row
      auto tile = configOp.getTile();
      auto tileOp = dyn_cast<TileOp>(tile.getDefiningOp());
      if (!tileOp) {
        startConfig.emitError("tile operand must be a TileOp");
        return signalPassFailure();
      }
 
      int col = tileOp.getCol();
      int row = tileOp.getRow();
      TileID tileID = {col, row};
 
      // Determine if we're accessing memory module from packet type
      bool isMem = false;
      if (auto packetType = configOp.getPacketType()) {
        isMem = (*packetType == TracePacketType::Mem);
      }
 
      // Process all trace.reg operations in the config
      for (auto &op : configOp.getBody().getOps()) {
        auto regOp = dyn_cast<TraceRegOp>(op);
        if (!regOp)
          continue;
 
        // After packing, field should not be present
        if (regOp.getField()) {
          regOp.emitError("aie.trace.reg still has field attribute - run "
                          "-aie-trace-pack-reg-writes pass first");
          return signalPassFailure();
        }
 
        // Look up register to get offset
        auto regName = regOp.getRegName().str();
        const RegisterInfo *regInfo =
            targetModel.lookupRegister(regName, tileID, isMem);
        if (!regInfo) {
          regOp.emitError("Register '") << regName << "' not found for tile ("
                                        << col << ", " << row << ")";
          return signalPassFailure();
        }
 
        // Extract value (mask is discarded)
        uint32_t value = 0;
        if (auto intAttr = llvm::dyn_cast<IntegerAttr>(regOp.getValue())) {
          value = intAttr.getInt();
        } else {
          regOp.emitError("value must be an integer after packing");
          return signalPassFailure();
        }
 
        // Generate aiex.npu.write32 operation with col/row
        AIEX::NpuWrite32Op::create(
            builder, regOp.getLoc(),
            AIEX::createConstantI32(builder, regOp.getLoc(), regInfo->offset),
            AIEX::createConstantI32(builder, regOp.getLoc(), value),
            nullptr,                        // buffer
            builder.getI32IntegerAttr(col), // column
            builder.getI32IntegerAttr(row)  // row
        );
      }
 
      // Remove the start_config invocation
      startConfig.erase();
    }
  }
};
 
} // namespace
 
std::unique_ptr<OperationPass<AIE::DeviceOp>>
xilinx::AIEX::createAIEXInlineTraceConfigPass() {
  return std::make_unique<AIEInlineTraceConfigPass>();
}