diff --git a/include/aie/Conversion/AIEToConfiguration/AIEToConfiguration.h b/include/aie/Conversion/AIEToConfiguration/AIEToConfiguration.h index f07f0fa0aa0..ce58a4cfa72 100644 --- a/include/aie/Conversion/AIEToConfiguration/AIEToConfiguration.h +++ b/include/aie/Conversion/AIEToConfiguration/AIEToConfiguration.h @@ -19,6 +19,10 @@ namespace xilinx::AIE { class DeviceOp; +// -------------------------------------------------------------------------- +// Device configuration +// -------------------------------------------------------------------------- + std::unique_ptr> createConvertAIEToTransactionPass(); @@ -29,6 +33,55 @@ std::optional convertTransactionBinaryToMLIR(mlir::MLIRContext *ctx, std::vector &binary); +// Generate transaction binary and insert configuration operations at a specific +// point +mlir::LogicalResult generateAndInsertConfigOps(xilinx::AIE::DeviceOp device, + mlir::Operation *insertionPoint, + llvm::StringRef clElfDir = ""); + +// -------------------------------------------------------------------------- +// Device reset +// -------------------------------------------------------------------------- + +// Enum for specifying which tile types to reset +enum class ResetTileType : unsigned { + None = 0, + ShimNOC = 1 << 0, + MemTile = 1 << 1, + CoreTile = 1 << 2, + All = ShimNOC | MemTile | CoreTile +}; + +inline bool hasFlag(ResetTileType value, ResetTileType flag) { + return (static_cast(value) & static_cast(flag)) != 0; +} + +// Enum for specifying when to reset +enum class ResetMode { + Never, // Never perform reset + IfUsed, // Reset only if the tile is used in the device + IfUsedFineGrained, // Reset only individual locks/connections that are used + IfChanged, // Reset only if the tile configuration changed from previous + IfChangedFineGrained, // Reset only individual locks/connections that changed + Always // Reset all tiles of the specified type +}; + +// Configuration for different reset operations +struct ResetConfig { + ResetTileType tileType; + ResetMode mode; + + ResetConfig(ResetTileType tt = ResetTileType::None, + ResetMode m = ResetMode::Never) + : tileType(tt), mode(m) {} +}; + +// Insert reset operations +mlir::LogicalResult generateAndInsertResetOps( + xilinx::AIE::DeviceOp device, mlir::Operation *insertionPoint, + ResetConfig dmaConfig, ResetConfig switchConfig, ResetConfig lockConfig, + ResetConfig coreConfig, xilinx::AIE::DeviceOp previousDevice); + } // namespace xilinx::AIE #endif // AIE_CONVERSION_AIETOCONFIGURATION_AIETOCONFIGURATION_H diff --git a/include/aie/Dialect/AIE/IR/AIEOps.td b/include/aie/Dialect/AIE/IR/AIEOps.td index d00361b35fb..9b13faf178d 100644 --- a/include/aie/Dialect/AIE/IR/AIEOps.td +++ b/include/aie/Dialect/AIE/IR/AIEOps.td @@ -76,6 +76,11 @@ def AIE_DeviceOp: AIE_Op<"device", [ static xilinx::AIE::DeviceOp getForSymbolInModule(mlir::ModuleOp module, llvm::StringRef symbol); static xilinx::AIE::DeviceOp getForSymbolInModuleOrError(mlir::ModuleOp module, llvm::StringRef symbol); static llvm::StringRef getDefaultDeviceName() { return "main"; } + + // Helper methods to find operations by tile coordinates, if they already exist (nullptr otherwise) + xilinx::AIE::SwitchboxOp lookupSwitchbox(int col, int row); + xilinx::AIE::CoreOp lookupCore(int col, int row); + mlir::Operation* lookupDMA(int col, int row); }]; } @@ -212,6 +217,7 @@ def AIE_SwitchboxOp: AIE_Op<"switchbox", [ int colIndex(); int rowIndex(); TileOp getTileOp(); + bool isEquivalent(SwitchboxOp other); void getAsmResultNames( llvm::function_ref setNameFn) { ::xilinx::AIE::TileElement::Trait::getAsmResultNames(setNameFn); @@ -322,6 +328,7 @@ def AIE_ShimDMAOp: AIE_Op<"shim_dma", [ int colIndex(); int rowIndex(); TileOp getTileOp(); + bool isEquivalent(mlir::Operation *other); void getAsmResultNames( llvm::function_ref setNameFn) { ::xilinx::AIE::TileElement::Trait::getAsmResultNames(setNameFn); @@ -389,6 +396,9 @@ def AIE_CoreOp: AIE_Op<"core", [ bool isMemWest() { return ((rowIndex() % 2) == 0); }; bool isEmpty(); TileOp getTileOp(); + // Check whether another core has equivalent program memory. + // Note that this currently only checks the ELF attribute, so only works after compiling the program memory. + bool isEquivalent(CoreOp other); void getAsmResultNames( llvm::function_ref setNameFn) { ::xilinx::AIE::TileElement::Trait::getAsmResultNames(setNameFn); @@ -1135,6 +1145,7 @@ def AIE_MemOp: AIE_Op<"mem", [ int colIndex(); int rowIndex(); TileOp getTileOp(); + bool isEquivalent(mlir::Operation *other); void getAsmResultNames( llvm::function_ref setNameFn) { ::xilinx::AIE::TileElement::Trait::getAsmResultNames(setNameFn); @@ -1182,6 +1193,7 @@ def AIE_MemTileDMAOp: AIE_Op<"memtile_dma", [ int colIndex(); int rowIndex(); TileOp getTileOp(); + bool isEquivalent(mlir::Operation *other); void getAsmResultNames( llvm::function_ref setNameFn) { ::xilinx::AIE::TileElement::Trait::getAsmResultNames(setNameFn); @@ -2160,6 +2172,7 @@ def AIE_RuntimeSequenceOp : AIE_Op<"runtime_sequence", [ static llvm::StringRef getDefaultRuntimeSequenceName() { return "sequence"; } static RuntimeSequenceOp getForSymbolInDevice(DeviceOp module, llvm::StringRef symbol); static RuntimeSequenceOp getForSymbolInDeviceOrError(DeviceOp module, llvm::StringRef symbol); + llvm::LogicalResult verifyBeforeMaterialization(); }]; let extraClassDefinition = [{ }]; diff --git a/include/aie/Dialect/AIE/IR/AIETargetModel.h b/include/aie/Dialect/AIE/IR/AIETargetModel.h index 201d7da658b..72ea0d2cd12 100644 --- a/include/aie/Dialect/AIE/IR/AIETargetModel.h +++ b/include/aie/Dialect/AIE/IR/AIETargetModel.h @@ -325,6 +325,28 @@ class AIETargetModel { // Returns true if the target model supports the given block format. virtual bool isSupportedBlockFormat(std::string const &format) const; + + /// Check if a register at the given address is a "special register" that + /// should not be reordered. Special registers are those that trigger + /// side-effects like device reset, start processing, etc. + /// \param absoluteAddress The full 32-bit address (including tile col/row) + /// \return true if this is a special register that must not be reordered + virtual bool isSpecialRegister(uint32_t absoluteAddress) const { + return false; // Default: no special registers + } + + /// Check if a register at the given tile and offset is a "special register" + /// \param col The column of the tile + /// \param row The row of the tile + /// \param offset The offset within the tile's address space + /// \return true if this is a special register that must not be reordered + virtual bool isSpecialRegister(int col, int row, uint32_t offset) const { + uint32_t absoluteAddress = + ((col & 0xff) << getColumnShift()) | + ((row & 0xff) << getRowShift()) | + (offset & 0xfffff); + return isSpecialRegister(absoluteAddress); + } }; class AIE1TargetModel : public AIETargetModel { @@ -711,6 +733,51 @@ class BaseNPU2TargetModel : public AIE2TargetModel { bool isSupportedBlockFormat(std::string const &format) const override; + /// Check if a register is a special register that should not be reordered + /// for AIE2P (NPU2) devices + bool isSpecialRegister(uint32_t absoluteAddress) const override { + // Extract tile coordinates from absolute address + // Address format: (col << 25) | (row << 20) | offset + int col = (absoluteAddress >> getColumnShift()) & 0xff; + int row = (absoluteAddress >> getRowShift()) & 0x1f; + uint32_t offset = absoluteAddress & 0xfffff; + return isSpecialRegister(col, row, offset); + } + + bool isSpecialRegister(int col, int row, uint32_t offset) const override { + if (isCoreTile(col, row)) { + // Core tile special registers + return offset == 0x32000 || // Core_Control + offset == 0x36070 || // Memory_Control + offset == 0x60000 || // Module_Clock_Control + offset == 0x60010 || // Module_Reset_Control + offset == 0x60020 || // Tile_Control + offset == 0x1DE04 || offset == 0x1DE0C || + offset == 0x1DE14 || offset == 0x1DE1C; + } else if (isShimNOCTile(col, row)) { + // Shim tile special registers + return (offset >= 0x15000 && offset <= 0x15010) || + offset == 0x1D204 || offset == 0x1D20C || + offset == 0x1D214 || offset == 0x1D21C || + offset == 0x1D200 || offset == 0x1D208 || + offset == 0x1D210 || offset == 0x1D218 || + offset == 0x1F000 || offset == 0x1F004 || + offset == 0x40000; + } else if (isMemTile(col, row)) { + // Memory tile special registers + return offset == 0x94008 || offset == 0x96048 || + offset == 0xA0604 || offset == 0xA060C || + offset == 0xA0614 || offset == 0xA061C || + offset == 0xA0624 || offset == 0xA062C || + offset == 0xA0634 || offset == 0xA063C || + offset == 0xA0644 || offset == 0xA064C || + offset == 0xA0654 || offset == 0xA065C || + offset == 0xFFF00 || offset == 0xFFF10 || + offset == 0xFFF20; + } + return false; + } + static bool classof(const AIETargetModel *model) { return model->getKind() >= TK_AIE2_NPU2 && model->getKind() < TK_AIE2_NPU2_Last; diff --git a/include/aie/Dialect/AIEX/AIEUtils.h b/include/aie/Dialect/AIEX/AIEUtils.h index 35e96e9a401..c3eddbbb4dc 100644 --- a/include/aie/Dialect/AIEX/AIEUtils.h +++ b/include/aie/Dialect/AIEX/AIEUtils.h @@ -10,6 +10,7 @@ #include "aie/Dialect/AIE/IR/AIEDialect.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/IR/Value.h" using namespace mlir; @@ -20,5 +21,20 @@ memref::GlobalOp getOrCreateDataMemref(OpBuilder &builder, AIE::DeviceOp dev, mlir::Location loc, ArrayRef words); +// Result of tracing through subview/cast operations to a block argument for +// traceSubviewToBlockArgument function. +struct SubviewTraceResult { + BlockArgument rootArg; + int64_t offsetInBytes; +}; + +// Trace through memref.subview, memref.cast, and memref.reinterpret_cast +// operations until the referenced SSA value is a block argument. +// +// Returns the root block argument and cumulative byte offset, or std::nullopt +// if the chain doesn't lead to a block argument or contains unsupported ops. +// +// This function supports checks that all subviews remain static and contiguous. +std::optional traceSubviewToBlockArgument(Value value); } } // namespace xilinx \ No newline at end of file diff --git a/include/aie/Dialect/AIEX/IR/AIEX.td b/include/aie/Dialect/AIEX/IR/AIEX.td index c7eceffd48b..46f34a80cb4 100644 --- a/include/aie/Dialect/AIEX/IR/AIEX.td +++ b/include/aie/Dialect/AIEX/IR/AIEX.td @@ -534,8 +534,7 @@ def AIE_SelectOp: AIEX_Op<"select", []>, Results<(outs Index)> { def AIE_ConfigureOp: AIEX_Op<"configure", [ HasParent<"AIE::RuntimeSequenceOp">, NoTerminator -]>, - Results<(outs Index:$result)> +]> { let summary = "Set up a configuration (program memories, stream switches, etc.) on the NPU device."; let arguments = ( diff --git a/include/aie/Dialect/AIEX/Transforms/AIEXPasses.h b/include/aie/Dialect/AIEX/Transforms/AIEXPasses.h index dc48a1df963..005b58d2c0d 100644 --- a/include/aie/Dialect/AIEX/Transforms/AIEXPasses.h +++ b/include/aie/Dialect/AIEX/Transforms/AIEXPasses.h @@ -32,6 +32,8 @@ std::unique_ptr> createAIEDmaToNpuPass(); std::unique_ptr> createAIEXToStandardPass(); std::unique_ptr> createAIEMaterializeBDChainsPass(); +std::unique_ptr> +createAIEMaterializeRuntimeSequencesPass(); std::unique_ptr> createAIEAssignRuntimeSequenceBDIDsPass(); std::unique_ptr> @@ -50,6 +52,10 @@ std::unique_ptr> createAIETxnToControlPacketPass(); std::unique_ptr> createAIELegalizeControlPacketPass(); +std::unique_ptr> +createAIEExpandLoadPdiPass(); +std::unique_ptr> +createAIECoalesceWrite32sPass(); /// Generate the code for registering passes. #define GEN_PASS_REGISTRATION diff --git a/include/aie/Dialect/AIEX/Transforms/AIEXPasses.td b/include/aie/Dialect/AIEX/Transforms/AIEXPasses.td index 990f45af4a5..aa872ca6818 100644 --- a/include/aie/Dialect/AIEX/Transforms/AIEXPasses.td +++ b/include/aie/Dialect/AIEX/Transforms/AIEXPasses.td @@ -160,6 +160,23 @@ def AIEMaterializeBDChains : Pass<"aie-materialize-bd-chains", "AIE::DeviceOp"> ]; } +def AIEMaterializeRuntimeSequences : Pass<"aie-materialize-runtime-sequences", "mlir::ModuleOp"> { + let summary = "Turn `aiex.configure` into configuration primitives and inline `aiex.run` calls"; + let description = [{ + Turns top-level `aiex.configure` into load PDI firmware instructions that set up the requested configuration. + + Inlines calls to other runtime sequences using `aiex.run` into the calling runtime sequence, resolving references within them to the device containing the definition. + + This pass requires ObjectFIFOs to be lowered first; apply the ObjectFIFO stateful transform before applying this pass. + }]; + + let constructor = "xilinx::AIEX::createAIEMaterializeRuntimeSequencesPass()"; + let dependentDialects = [ + "xilinx::AIE::AIEDialect", + "xilinx::AIEX::AIEXDialect", + ]; +} + def AIEAssignRuntimeSequenceBDIDs : Pass<"aie-assign-runtime-sequence-bd-ids", "AIE::DeviceOp"> { let summary = "Assign IDs to Buffer Descriptors Configured in the Runtime Sequence"; @@ -254,4 +271,88 @@ def AIELegalizeControlPacket : Pass<"aie-legalize-ctrl-packet", "AIE::DeviceOp"> "xilinx::AIEX::AIEXDialect", ]; } + +def AIEExpandLoadPdi : Pass<"aie-expand-load-pdi", "mlir::ModuleOp"> { + let summary = "Expand load_pdi operations to explicit configuration sequences"; + let description = [{ + This pass transforms `load_pdi { device_ref = @device }` operations as follows: + 1. Create an empty device `@empty` if it doesn't exist + 2. Replace the load_pdi with `load_pdi { device_ref = @empty }` -- this causes a device reset. + 3. Insert additional reset instructions (e.g., reset switchboxes, based on command-line arguments passed). + 4. Configure the device for the selected design by inserting write32s and blockwrites. + + The configuration operations are generated by converting the referenced device's configuration to a transaction binary and then disassembling it into MLIR operations. + + Example command-line usage: + ``` + aie-opt --aie-expand-load-pdi \ + --reset-switches-tiles=shim,mem,core \ + --reset-switches-mode=ifused \ + --reset-locks-tiles=mem,core \ + --reset-locks-mode=always \ + input.mlir + ``` + }]; + + let constructor = "xilinx::AIEX::createAIEExpandLoadPdiPass()"; + let dependentDialects = [ + "mlir::memref::MemRefDialect", + "xilinx::AIE::AIEDialect", + "xilinx::AIEX::AIEXDialect", + ]; + + let options = [ + Option<"resetDmasTiles", "reset-dmas-tiles", "std::string", "\"\"", + "Comma-separated list of tile types to reset DMAs for (shim,mem,core)">, + Option<"resetDmasMode", "reset-dmas-mode", "std::string", "\"never\"", + "When to reset DMAs: never, ifused, ifchanged, always">, + Option<"resetSwitchesTiles", "reset-switches-tiles", "std::string", "\"\"", + "Comma-separated list of tile types to reset switches for (shim,mem,core)">, + Option<"resetSwitchesMode", "reset-switches-mode", "std::string", "\"never\"", + "When to reset switches: never, ifused, ifusedfinegrained, ifchanged, ifchangedfinegrained, always">, + Option<"resetLocksTiles", "reset-locks-tiles", "std::string", "\"\"", + "Comma-separated list of tile types to reset locks for (mem,core)">, + Option<"resetLocksMode", "reset-locks-mode", "std::string", "\"never\"", + "When to reset locks: never, ifused, ifusedfinegrained, ifchanged, ifchangedfinegrained, always">, + Option<"resetCoresTiles", "reset-cores-tiles", "std::string", "\"\"", + "Comma-separated list of tile types to reset cores for (core)">, + Option<"resetCoresMode", "reset-cores-mode", "std::string", "\"never\"", + "When to reset cores: never, ifused, ifchanged, always">, + ]; +} + +def AIECoalesceWrite32s : Pass<"aie-coalesce-write32s", "AIE::DeviceOp"> { + let summary = "Coalesce consecutive write32 operations into blockwrite operations"; + let description = [{ + This pass optimizes sequences of consecutive npu.write32 operations by + coalescing them into npu.blockwrite operations when their addresses are + contiguous (4-byte increments). + + The pass can reorder non-special register writes within slices bounded by + special register writes. Special registers (e.g., control registers) act + as barriers and cannot be reordered. + + For each sequence of N consecutive write32 operations at addresses + [addr, addr+4, addr+8, ...], the pass: + 1. Eliminates duplicate writes (keeps only the last value for each address) + 2. Reorders writes by address within safe boundaries + 3. Creates a memref.global with the values + 4. Replaces the write32 sequence with a single npu.blockwrite + + This reduces instruction count and can improve performance. + }]; + + let options = [ + Option<"minWritesToCoalesce", "min-writes-to-coalesce", "unsigned", + /*default=*/"2", + "Minimum number of contiguous writes required to coalesce into a blockwrite"> + ]; + + let constructor = "xilinx::AIEX::createAIECoalesceWrite32sPass()"; + let dependentDialects = [ + "mlir::memref::MemRefDialect", + "xilinx::AIE::AIEDialect", + "xilinx::AIEX::AIEXDialect", + ]; +} #endif diff --git a/include/aie/Targets/AIERT.h b/include/aie/Targets/AIERT.h index e36eace870a..4f4cdf9f5cd 100644 --- a/include/aie/Targets/AIERT.h +++ b/include/aie/Targets/AIERT.h @@ -46,6 +46,18 @@ struct AIERTControl { void startTransaction(); void dmaUpdateBdAddr(int col, int row, size_t addr, size_t bdId); std::vector exportSerializedTransaction(); + mlir::LogicalResult resetPartition(); + mlir::LogicalResult resetDMA(int col, int row, bool on); + mlir::LogicalResult resetCore(int col, int row); + mlir::LogicalResult resetCoreUnreset(int col, int row); + mlir::LogicalResult resetSwitch(int col, int row); + mlir::LogicalResult resetLock(int col, int row, int lockId); + mlir::LogicalResult resetSwitchConnection(int col, int row, + WireBundle sourceBundle, + int sourceChannel, + WireBundle destBundle, + int destChannel); + mlir::LogicalResult resetPerfCounters(int col, int row); private: const AIETargetModel &targetModel; diff --git a/lib/Conversion/AIEToConfiguration/AIEToConfiguration.cpp b/lib/Conversion/AIEToConfiguration/AIEToConfiguration.cpp index f66e71a8f56..b2574e5c083 100644 --- a/lib/Conversion/AIEToConfiguration/AIEToConfiguration.cpp +++ b/lib/Conversion/AIEToConfiguration/AIEToConfiguration.cpp @@ -629,9 +629,11 @@ enum OutputType { ControlPacket, }; -static LogicalResult convertTransactionOpsToMLIR( - OpBuilder builder, AIE::DeviceOp device, OutputType outputType, - std::vector &operations) { +static LogicalResult +convertTransactionOpsToMLIR(OpBuilder builder, AIE::DeviceOp device, + OutputType outputType, + std::vector &operations, + Operation *insertionPoint = nullptr) { auto loc = builder.getUnknownLoc(); @@ -647,9 +649,9 @@ static LogicalResult convertTransactionOpsToMLIR( std::vector data32(d, d + size); int id = 0; - std::string name = "blockwrite_data"; + std::string name = "config_blockwrite_data"; while (device.lookupSymbol(name)) - name = "blockwrite_data_" + std::to_string(id++); + name = "config_blockwrite_data_" + std::to_string(id++); MemRefType memrefType = MemRefType::get({size}, builder.getI32Type()); TensorType tensorType = RankedTensorType::get({size}, builder.getI32Type()); @@ -659,27 +661,32 @@ static LogicalResult convertTransactionOpsToMLIR( global_data.push_back(global); } - // search for aiex.configure ops in runtime sequences by walking the device - // and collect them in a vector. If there are none, create a new runtime - // sequence. Otherwise assume the insertion point is the first aiex.configure - // op. + // If an explicit insertion point is provided, use it for the config ops SmallVector configureOps; - device.walk([&](AIEX::ConfigureOp op) { configureOps.push_back(op); }); - - if (configureOps.empty()) { - - // create aiex.runtime_sequence - int id = 0; - std::string seq_name = "configure"; - while (device.lookupSymbol(seq_name)) - seq_name = "configure" + std::to_string(id++); - StringAttr seq_sym_name = builder.getStringAttr(seq_name); - auto seq = AIE::RuntimeSequenceOp::create(builder, loc, seq_sym_name); - seq.getBody().push_back(new Block); - - builder.setInsertionPointToStart(&seq.getBody().front()); + if (insertionPoint) { + builder.setInsertionPointAfter(insertionPoint); } else { - builder.setInsertionPoint(configureOps.front()); + // search for aiex.configure ops in runtime sequences by walking the device + // and collect them in a vector. If there are none, create a new runtime + // sequence. Otherwise assume the insertion point is the first + // aiex.configure op. + device.walk([&](AIEX::ConfigureOp op) { configureOps.push_back(op); }); + + if (configureOps.empty()) { + + // create aiex.runtime_sequence + int id = 0; + std::string seq_name = "configure"; + while (device.lookupSymbol(seq_name)) + seq_name = "configure" + std::to_string(id++); + StringAttr seq_sym_name = builder.getStringAttr(seq_name); + auto seq = AIE::RuntimeSequenceOp::create(builder, loc, seq_sym_name); + seq.getBody().push_back(new Block); + + builder.setInsertionPointToStart(&seq.getBody().front()); + } else { + builder.setInsertionPoint(configureOps.front()); + } } // create the txn ops @@ -696,7 +703,7 @@ static LogicalResult convertTransactionOpsToMLIR( llvm_unreachable("bad output type"); } - if (!configureOps.empty()) { + if (!configureOps.empty() && !insertionPoint) { // splice the body into the current insertion point builder.getBlock()->getOperations().splice( builder.getInsertionPoint(), @@ -850,3 +857,408 @@ std::unique_ptr> xilinx::AIE::createConvertAIEToControlPacketsPass() { return std::make_unique(); } + +mlir::LogicalResult +xilinx::AIE::generateAndInsertConfigOps(xilinx::AIE::DeviceOp device, + mlir::Operation *insertionPoint, + llvm::StringRef clElfDir) { + const AIETargetModel &targetModel = + (const AIETargetModel &)device.getTargetModel(); + + if (!targetModel.hasProperty(AIETargetModel::IsNPU)) + return failure(); + + bool aieSim = false; + bool xaieDebug = false; + + AIERTControl ctl(targetModel); + if (failed(ctl.setIOBackend(aieSim, xaieDebug))) + return failure(); + + // start collecting transactions + ctl.startTransaction(); + + bool generateElfs = true; + if (failed(generateTransactions(ctl, clElfDir, device, aieSim, generateElfs, + true, true))) + return failure(); + + // Export the transactions to a binary buffer + std::vector txn_data = ctl.exportSerializedTransaction(); + + // parse the binary data + std::vector operations; + if (!parseTransactionBinary(txn_data, operations)) { + llvm::errs() << "Failed to parse binary\n"; + return failure(); + } + + // Get the parent device for the insertion point + auto parentDevice = insertionPoint->getParentOfType(); + if (!parentDevice) { + llvm::errs() << "Insertion point must be within a DeviceOp\n"; + return failure(); + } + + OpBuilder builder(parentDevice.getBodyRegion()); + + // convert the parsed ops to MLIR, inserting after the provided point + if (failed(convertTransactionOpsToMLIR(builder, parentDevice, + OutputType::Transaction, operations, + insertionPoint))) + return failure(); + + return success(); +} + +// Helper function to check DMA equivalence for different operation types +static bool areDMAsEquivalent(Operation *dma1, Operation *dma2) { + if (!dma1 || !dma2) + return dma1 == dma2; + if (auto mem1 = dyn_cast(dma1)) { + return mem1.isEquivalent(dma2); + } else if (auto memTile1 = dyn_cast(dma1)) { + return memTile1.isEquivalent(dma2); + } else if (auto shim1 = dyn_cast(dma1)) { + return shim1.isEquivalent(dma2); + } + return false; +} + +// Reset cores +static LogicalResult resetCores(DeviceOp device, DeviceOp previousDevice, + ResetConfig coreConfig, AIERTControl &ctl) { + const AIETargetModel &targetModel = device.getTargetModel(); + + if (coreConfig.mode == ResetMode::Never || + !hasFlag(coreConfig.tileType, ResetTileType::CoreTile)) + return success(); + + int numCols = targetModel.columns(); + int numRows = targetModel.rows(); + + for (int col = 0; col < numCols; col++) { + for (int row = 0; row < numRows; row++) { + if (!targetModel.isCoreTile(col, row)) + continue; + + bool shouldReset = false; + + if (coreConfig.mode == ResetMode::Always) { + shouldReset = true; + } else if (coreConfig.mode == ResetMode::IfUsed || + coreConfig.mode == ResetMode::IfUsedFineGrained) { + shouldReset = device.lookupCore(col, row) != nullptr; + } else if (coreConfig.mode == ResetMode::IfChanged || + coreConfig.mode == ResetMode::IfChangedFineGrained) { + auto currentCore = device.lookupCore(col, row); + auto prevCore = + previousDevice ? previousDevice.lookupCore(col, row) : nullptr; + shouldReset = currentCore && !currentCore.isEquivalent(prevCore); + } + + if (shouldReset && failed(ctl.resetCore(col, row))) + return failure(); + } + } + + return success(); +} + +// Reset DMAs based on the provided configuration +static LogicalResult resetDMAs(DeviceOp device, DeviceOp previousDevice, + ResetConfig dmaConfig, AIERTControl &ctl) { + const AIETargetModel &targetModel = device.getTargetModel(); + + if (dmaConfig.mode == ResetMode::Never) + return success(); + + int numCols = targetModel.columns(); + int numRows = targetModel.rows(); + + for (int col = 0; col < numCols; col++) { + for (int row = 0; row < numRows; row++) { + if ((targetModel.isMemTile(col, row) && + !hasFlag(dmaConfig.tileType, ResetTileType::MemTile)) || + (targetModel.isCoreTile(col, row) && + !hasFlag(dmaConfig.tileType, ResetTileType::CoreTile)) || + (targetModel.isShimNOCTile(col, row) && + !hasFlag(dmaConfig.tileType, ResetTileType::ShimNOC))) { + continue; + } + + bool shouldReset = false; + + if (dmaConfig.mode == ResetMode::Always) { + shouldReset = true; + } else if (dmaConfig.mode == ResetMode::IfUsed || + dmaConfig.mode == ResetMode::IfUsedFineGrained) { + shouldReset = nullptr != device.lookupDMA(col, row); + } else if (dmaConfig.mode == ResetMode::IfChanged || + dmaConfig.mode == ResetMode::IfChangedFineGrained) { + Operation *currentDMA = device.lookupDMA(col, row); + Operation *prevDMA = + previousDevice ? previousDevice.lookupDMA(col, row) : nullptr; + shouldReset = currentDMA && !areDMAsEquivalent(currentDMA, prevDMA); + } + + if (shouldReset && failed(ctl.resetDMA(col, row, false))) + return failure(); + } + } + + return success(); +} + +// Reset switches based on the provided configuration +static LogicalResult resetSwitches(DeviceOp device, DeviceOp previousDevice, + ResetConfig switchConfig, + AIERTControl &ctl) { + const AIETargetModel &targetModel = device.getTargetModel(); + + if (switchConfig.mode == ResetMode::Never) + return success(); + + int numCols = targetModel.columns(); + int numRows = targetModel.rows(); + + for (int col = 0; col < numCols; col++) { + for (int row = 0; row < numRows; row++) { + if ((targetModel.isMemTile(col, row) && + !hasFlag(switchConfig.tileType, ResetTileType::MemTile)) || + (targetModel.isCoreTile(col, row) && + !hasFlag(switchConfig.tileType, ResetTileType::CoreTile)) || + (targetModel.isShimNOCTile(col, row) && + !hasFlag(switchConfig.tileType, ResetTileType::ShimNOC))) { + continue; + } + + bool shouldReset = false; + + if (switchConfig.mode == ResetMode::Always) { + shouldReset = true; + } else if (switchConfig.mode == ResetMode::IfUsed) { + shouldReset = nullptr != device.lookupSwitchbox(col, row); + } else if (switchConfig.mode == ResetMode::IfChanged) { + SwitchboxOp currentSB = device.lookupSwitchbox(col, row); + SwitchboxOp prevSB = + previousDevice ? previousDevice.lookupSwitchbox(col, row) : nullptr; + shouldReset = currentSB && !currentSB.isEquivalent(prevSB); + } + + if (shouldReset && failed(ctl.resetSwitch(col, row))) + return failure(); + } + } + + return success(); +} + +// Reset locks based on the provided configuration +static LogicalResult resetLocks(DeviceOp device, DeviceOp previousDevice, + ResetConfig lockConfig, AIERTControl &ctl) { + const AIETargetModel &targetModel = device.getTargetModel(); + + if (lockConfig.mode == ResetMode::Never) + return success(); + + int numCols = targetModel.columns(); + int numRows = targetModel.rows(); + + for (int col = 0; col < numCols; col++) { + for (int row = 0; row < numRows; row++) { + if ((targetModel.isMemTile(col, row) && + !hasFlag(lockConfig.tileType, ResetTileType::MemTile)) || + (targetModel.isCoreTile(col, row) && + !hasFlag(lockConfig.tileType, ResetTileType::CoreTile))) { + continue; + } + + if (lockConfig.mode == ResetMode::Always) { + // Reset all locks on this tile; + int numLocks = targetModel.getNumLocks(col, row); + for (int lockID = 0; lockID < numLocks; lockID++) { + if (failed(ctl.resetLock(col, row, lockID))) + return failure(); + } + } else { + // treat "ifUsed" and "ifChanged" as the same -- if a tile defined some + // locks, it probably also toggled them, so we should reset them + for (auto lockOp : device.getOps()) { + auto tileOp = lockOp.getTileOp(); + if (tileOp.getCol() == col && tileOp.getRow() == row) { + if (auto lockId = lockOp.getLockID()) { + if (failed(ctl.resetLock(col, row, *lockId))) { + return failure(); + } + } + } + } + } + } + } + + return success(); +} + +// Reset fine-grained switch connections based on the provided configuration +static LogicalResult resetSwitchConnectionsFineGrained(DeviceOp device, + DeviceOp previousDevice, + ResetConfig switchConfig, + AIERTControl &ctl) { + const AIETargetModel &targetModel = device.getTargetModel(); + + if (switchConfig.mode != ResetMode::IfUsedFineGrained && + switchConfig.mode != ResetMode::IfChangedFineGrained) + return success(); + + int numCols = targetModel.columns(); + int numRows = targetModel.rows(); + + for (int col = 0; col < numCols; col++) { + for (int row = 0; row < numRows; row++) { + if ((targetModel.isMemTile(col, row) && + !hasFlag(switchConfig.tileType, ResetTileType::MemTile)) || + (targetModel.isCoreTile(col, row) && + !hasFlag(switchConfig.tileType, ResetTileType::CoreTile)) || + (targetModel.isShimNOCTile(col, row) && + !hasFlag(switchConfig.tileType, ResetTileType::ShimNOC))) { + continue; + } + + auto currentSB = device.lookupSwitchbox(col, row); + if (!currentSB) + continue; + + // Collect unique sources and active connections + DenseSet> uniqueSources; + DenseSet> activeConns; + for (auto connectOp : currentSB.getOps()) { + WireBundle srcBundle = connectOp.getSourceBundle(); + int srcChan = connectOp.getSourceChannel(); + WireBundle dstBundle = connectOp.getDestBundle(); + int dstChan = connectOp.getDestChannel(); + uniqueSources.insert({srcBundle, srcChan}); + activeConns.insert({srcBundle, srcChan, dstBundle, dstChan}); + } + DenseSet> prevConnsSet; + auto prevSB = + previousDevice ? previousDevice.lookupSwitchbox(col, row) : nullptr; + if (prevSB) { + for (auto connectOp : prevSB.getOps()) { + prevConnsSet.insert( + {connectOp.getSourceBundle(), connectOp.getSourceChannel(), + connectOp.getDestBundle(), connectOp.getDestChannel()}); + } + } + + // Get all possible destinations + llvm::SmallVector> possibleDests; + for (auto bundle : + {WireBundle::Core, WireBundle::DMA, WireBundle::FIFO, + WireBundle::South, WireBundle::West, WireBundle::North, + WireBundle::East, WireBundle::PLIO, WireBundle::NOC, + WireBundle::Trace, WireBundle::TileControl}) { + uint32_t numDestChans = + targetModel.getNumDestSwitchboxConnections(col, row, bundle); + if (numDestChans > 0) { + possibleDests.push_back(std::make_pair(bundle, numDestChans)); + } + } + + // Iterate over all sources of connections in the current device + for (auto [srcBundle, srcChan] : uniqueSources) { + // Iterate over all legal destinations from this source + for (auto [altDstBundle, numDestChans] : possibleDests) { + for (uint32_t altDstChan = 0; altDstChan < numDestChans; + altDstChan++) { + if (!targetModel.isLegalTileConnection(col, row, srcBundle, srcChan, + altDstBundle, altDstChan)) { + continue; + } + // Skip if this is an active connection in the new configuration + if (activeConns.contains( + {srcBundle, srcChan, altDstBundle, altDstChan})) { + continue; + } + + if (switchConfig.mode == ResetMode::IfUsedFineGrained || + prevConnsSet.contains( + {srcBundle, srcChan, altDstBundle, altDstChan})) { + // Reset any legal destination that's not the active one + if (failed(ctl.resetSwitchConnection(col, row, srcBundle, srcChan, + altDstBundle, altDstChan))) { + return failure(); + } + } + } + } + } + } + } + + return success(); +} + +// Generate reset operations for all tiles used in the device +mlir::LogicalResult xilinx::AIE::generateAndInsertResetOps( + xilinx::AIE::DeviceOp device, mlir::Operation *insertionPoint, + ResetConfig dmaConfig, ResetConfig switchConfig, ResetConfig lockConfig, + ResetConfig coreConfig, xilinx::AIE::DeviceOp previousDevice) { + const AIETargetModel &targetModel = device.getTargetModel(); + + if (!targetModel.hasProperty(AIETargetModel::IsNPU)) + return failure(); + + AIERTControl ctl(targetModel); + + bool aieSim = false; + bool xaieDebug = false; + if (failed(ctl.setIOBackend(aieSim, xaieDebug))) + return failure(); + + ctl.startTransaction(); + + if (failed(resetCores(device, previousDevice, coreConfig, ctl))) + return failure(); + + if (failed(resetDMAs(device, previousDevice, dmaConfig, ctl))) + return failure(); + + if (failed(resetSwitches(device, previousDevice, switchConfig, ctl))) + return failure(); + + if (failed(resetSwitchConnectionsFineGrained(device, previousDevice, + switchConfig, ctl))) + return failure(); + + if (failed(resetLocks(device, previousDevice, lockConfig, ctl))) + return failure(); + + // Export the reset transactions + std::vector txn_data = ctl.exportSerializedTransaction(); + + // Parse the binary data + std::vector operations; + if (!parseTransactionBinary(txn_data, operations)) { + llvm::errs() << "Failed to parse reset transaction binary\n"; + return failure(); + } + + // Get the parent device for the insertion point + auto parentDevice = insertionPoint->getParentOfType(); + if (!parentDevice) { + llvm::errs() << "Insertion point must be within a DeviceOp\n"; + return failure(); + } + + OpBuilder builder(parentDevice.getBodyRegion()); + + // Convert the parsed reset ops to MLIR, inserting after the provided point + if (failed(convertTransactionOpsToMLIR(builder, parentDevice, + OutputType::Transaction, operations, + insertionPoint))) + return failure(); + + return success(); +} diff --git a/lib/Dialect/AIE/IR/AIEDialect.cpp b/lib/Dialect/AIE/IR/AIEDialect.cpp index ec618b85b24..1c5b215a0b9 100644 --- a/lib/Dialect/AIE/IR/AIEDialect.cpp +++ b/lib/Dialect/AIE/IR/AIEDialect.cpp @@ -1093,6 +1093,45 @@ DeviceOp::getForSymbolInModuleOrError(mlir::ModuleOp module, return deviceOp; } +SwitchboxOp DeviceOp::lookupSwitchbox(int col, int row) { + for (auto sb : getOps()) { + if (sb.colIndex() == col && sb.rowIndex() == row) + return sb; + } + return nullptr; +} + +CoreOp DeviceOp::lookupCore(int col, int row) { + for (auto core : getOps()) { + if (core.colIndex() == col && core.rowIndex() == row) + return core; + } + return nullptr; +} + +Operation *DeviceOp::lookupDMA(int col, int row) { + const AIETargetModel &targetModel = getTargetModel(); + + // Use target model to determine which type of DMA op to look for + if (targetModel.isMemTile(col, row)) { + for (auto mem : getOps()) { + if (mem.colIndex() == col && mem.rowIndex() == row) + return mem.getOperation(); + } + } else if (targetModel.isCoreTile(col, row)) { + for (auto mem : getOps()) { + if (mem.colIndex() == col && mem.rowIndex() == row) + return mem.getOperation(); + } + } else if (targetModel.isShimNOCTile(col, row)) { + for (auto shim : getOps()) { + if (shim.colIndex() == col && shim.rowIndex() == row) + return shim.getOperation(); + } + } + return nullptr; +} + //===----------------------------------------------------------------------===// // TileOp //===----------------------------------------------------------------------===// @@ -1298,6 +1337,56 @@ int ShimMuxOp::colIndex() { return getTileOp().colIndex(); } int ShimMuxOp::rowIndex() { return getTileOp().rowIndex(); } +// Helper function to compare DMA configuration for equivalence +static bool areDMAsEquivalent(Operation *dma1, Operation *dma2) { + if (!dma1 || !dma2) + return dma1 == dma2; + + // Compare their contained operations + // This is a structural comparison including attributes + auto ®ion1 = dma1->getRegion(0); + auto ®ion2 = dma2->getRegion(0); + + if (region1.empty() != region2.empty()) + return false; + + if (region1.empty()) + return true; + + auto &block1 = region1.front(); + auto &block2 = region2.front(); + + auto ops1 = block1.without_terminator(); + auto ops2 = block2.without_terminator(); + + auto it1 = ops1.begin(); + auto it2 = ops2.begin(); + + while (it1 != ops1.end() && it2 != ops2.end()) { + // Check operation names match + if (it1->getName() != it2->getName()) + return false; + + // Check all attributes match + auto attrs1 = it1->getAttrs(); + auto attrs2 = it2->getAttrs(); + + if (attrs1.size() != attrs2.size()) + return false; + + for (auto attr1 : attrs1) { + auto attr2 = it2->getAttr(attr1.getName()); + if (!attr2 || attr1.getValue() != attr2) + return false; + } + + ++it1; + ++it2; + } + + return it1 == ops1.end() && it2 == ops2.end(); +} + //===----------------------------------------------------------------------===// // ShimDMAOp //===----------------------------------------------------------------------===// @@ -1320,6 +1409,10 @@ int ShimDMAOp::colIndex() { return getTileOp().colIndex(); } int ShimDMAOp::rowIndex() { return getTileOp().rowIndex(); } +bool ShimDMAOp::isEquivalent(Operation *other) { + return areDMAsEquivalent(getOperation(), other); +} + LogicalResult PacketRulesOp::verify() { if (Region &body = getRules(); body.empty()) return emitOpError("should have non-empty body"); @@ -1376,6 +1469,20 @@ bool CoreOp::isEmpty() { llvm::isa(body.front().front())); } +bool CoreOp::isEquivalent(CoreOp other) { + if (!other) + return false; + + // Compare program memory ELF + auto elf1 = getElfFileAttr(); + auto elf2 = other.getElfFileAttr(); + + if (elf1 && elf2) + return elf1 == elf2; + + return !elf1 && !elf2; +} + //===----------------------------------------------------------------------===// // BufferOp //===----------------------------------------------------------------------===// @@ -1477,6 +1584,10 @@ int MemOp::colIndex() { return getTileOp().colIndex(); } int MemOp::rowIndex() { return getTileOp().rowIndex(); } +bool MemOp::isEquivalent(Operation *other) { + return areDMAsEquivalent(getOperation(), other); +} + //===----------------------------------------------------------------------===// // MemTileDMAOp //===----------------------------------------------------------------------===// @@ -1899,6 +2010,10 @@ int MemTileDMAOp::colIndex() { return getTileOp().colIndex(); } int MemTileDMAOp::rowIndex() { return getTileOp().rowIndex(); } +bool MemTileDMAOp::isEquivalent(Operation *other) { + return areDMAsEquivalent(getOperation(), other); +} + //===----------------------------------------------------------------------===// // DMAStartOp //===----------------------------------------------------------------------===// @@ -2114,6 +2229,47 @@ int SwitchboxOp::colIndex() { return getTileOp().colIndex(); } int SwitchboxOp::rowIndex() { return getTileOp().rowIndex(); } +bool SwitchboxOp::isEquivalent(SwitchboxOp other) { + if (!other) + return false; + + // Compare all connection ops + auto conns1 = getOps(); + auto conns2 = other.getOps(); + + llvm::SmallVector vec1(conns1.begin(), conns1.end()); + llvm::SmallVector vec2(conns2.begin(), conns2.end()); + + if (vec1.size() != vec2.size()) + return false; + + // Sort by source/dest to make comparison order-independent + auto connComparator = [](ConnectOp a, ConnectOp b) { + if (a.getSourceBundle() != b.getSourceBundle()) + return static_cast(a.getSourceBundle()) < + static_cast(b.getSourceBundle()); + if (a.getSourceChannel() != b.getSourceChannel()) + return a.getSourceChannel() < b.getSourceChannel(); + if (a.getDestBundle() != b.getDestBundle()) + return static_cast(a.getDestBundle()) < + static_cast(b.getDestBundle()); + return a.getDestChannel() < b.getDestChannel(); + }; + + llvm::sort(vec1, connComparator); + llvm::sort(vec2, connComparator); + + for (size_t i = 0; i < vec1.size(); ++i) { + if (vec1[i].getSourceBundle() != vec2[i].getSourceBundle() || + vec1[i].getSourceChannel() != vec2[i].getSourceChannel() || + vec1[i].getDestBundle() != vec2[i].getDestBundle() || + vec1[i].getDestChannel() != vec2[i].getDestChannel()) + return false; + } + + return true; +} + template struct HasSomeParent { static LogicalResult verifyTrait(Operation *op) { @@ -2479,6 +2635,42 @@ RuntimeSequenceOp::getForSymbolInDeviceOrError(DeviceOp deviceOp, return runtimeSequenceOp; } +LogicalResult RuntimeSequenceOp::verifyBeforeMaterialization() { + // Check that all symbol references within the runtime sequence + // are either to ShimDMAAllocationOp, DeviceOp or another RuntimeSequenceOp; + // these are the only symbols that can be lowered with the NPU passes + auto result = (*this)->walk([&](Operation *op) { + for (NamedAttribute namedAttr : op->getAttrs()) { + Attribute attr = namedAttr.getValue(); + auto walkResult = attr.walk([&](SymbolRefAttr symbolRef) { + Operation *symbolDefOp = + SymbolTable::lookupNearestSymbolFrom(*this, symbolRef); + if (symbolDefOp && !llvm::isa(symbolDefOp) && + !llvm::isa(symbolDefOp) && + !llvm::isa(symbolDefOp)) { + op->emitOpError() << "references symbol '" + << symbolRef.getRootReference().getValue() + << "' which must be either a ShimDMAAllocationOp, " + "DeviceOp or RuntimeSequenceOp, but got: " + << symbolDefOp->getName().getStringRef(); + return WalkResult::interrupt(); + } + return WalkResult::advance(); + }); + if (walkResult.wasInterrupted()) { + return WalkResult::interrupt(); + } + } + return WalkResult::advance(); + }); + + if (result.wasInterrupted()) { + return failure(); + } + + return success(); +} + // Include implementations for custom attributes #define GET_ATTRDEF_CLASSES #include "aie/Dialect/AIE/IR/AIEAttrs.cpp.inc" diff --git a/lib/Dialect/AIEX/IR/AIEXDialect.cpp b/lib/Dialect/AIEX/IR/AIEXDialect.cpp index 025680ff8e2..ef2a92ecf0e 100644 --- a/lib/Dialect/AIEX/IR/AIEXDialect.cpp +++ b/lib/Dialect/AIEX/IR/AIEXDialect.cpp @@ -50,9 +50,7 @@ void AIEXDialect::initialize() { } // namespace xilinx::AIEX #define GET_OP_CLASSES -#include "aie/Dialect/AIEX/IR/AIEX.cpp.inc" - -/* Return the correct values to write to the hardware registers to configure +#include "aie/Dialect/AIEX/IR/AIEX.cpp.inc" /* Return the correct values to write to the hardware registers to configure strides and wraps given the input user-facing strides and wraps. In the IR, we express strides in units of element data type, but the hardware @@ -980,8 +978,47 @@ AIE::RuntimeSequenceOp AIEX::RunOp::getCalleeRuntimeSequenceOp() { } LogicalResult AIEX::RunOp::verify() { - if (getCalleeDeviceOp() && getCalleeRuntimeSequenceOp()) { - return success(); + AIE::DeviceOp calleeDevice = getCalleeDeviceOp(); + if (!calleeDevice) { + return emitOpError() << "cannot find callee device for runtime sequence '" + << getRuntimeSequenceSymbol() << "'"; } - return failure(); + + AIE::RuntimeSequenceOp calleeRuntimeSequence = getCalleeRuntimeSequenceOp(); + if (!calleeRuntimeSequence) { + return emitOpError() << "cannot find runtime sequence '" + << getRuntimeSequenceSymbol() << "'"; + } + + // Validate argument types match the callee's parameters + Block &calleeBody = calleeRuntimeSequence.getBody().front(); + ValueRange values = getArgs(); + + if (calleeBody.getNumArguments() != values.size()) { + return emitOpError() << "argument count mismatch"; + } + + for (unsigned i = 0, n = calleeBody.getNumArguments(); i < n; i++) { + BlockArgument arg = calleeBody.getArgument(i); + Value val = values[i]; + + // For memref types, check compatibility (same shape and element type) + auto argType = dyn_cast(arg.getType()); + auto valType = dyn_cast(val.getType()); + + if (argType && valType) { + if (argType.getShape() != valType.getShape() || + argType.getElementType() != valType.getElementType()) { + return emitOpError() + << "argument " << i << " type mismatch: " << "expected " + << argType << " but got " << valType; + } + } else if (arg.getType() != val.getType()) { + return emitOpError() << "argument " << i + << " type mismatch: " << "expected " << arg.getType() + << " but got " << val.getType(); + } + } + + return success(); } diff --git a/lib/Dialect/AIEX/Transforms/AIECoalesceWrite32s.cpp b/lib/Dialect/AIEX/Transforms/AIECoalesceWrite32s.cpp new file mode 100644 index 00000000000..7621aafeaef --- /dev/null +++ b/lib/Dialect/AIEX/Transforms/AIECoalesceWrite32s.cpp @@ -0,0 +1,433 @@ +//===- AIECoalesceWrite32s.cpp ----------------------------------*- C++ -*-===// +// +// This file is licensed 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 +// +// (c) Copyright 2025 Advanced Micro Devices Inc. +// +//===----------------------------------------------------------------------===// +// +// This pass coalesces consecutive npu.write32 operations into npu.blockwrite +// operations when their addresses are contiguous (4-byte increments). +// +// The pass can reorder non-special register writes within slices. Special +// register writes act as barriers and cannot be reordered. Duplicate writes +// to the same address are eliminated (keeping the last value). +// +// Masked writes (npu.maskwrite32) are handled as follows: +// - If a maskwrite is done to an address after a previous write32/blockwrite, +// the maskwrite supersedes the earlier write (earlier write is removed) +// - If a write32/blockwrite is done after a maskwrite to the same address, +// the later write supersedes the maskwrite (maskwrite is removed) +// - Maskwrite operations cannot be coalesced into blockwrites and act as +// barriers for coalescing (similar to special register writes) +// +// A configurable threshold controls the minimum number of contiguous writes +// required for coalescing into a blockwrite. +// +//===----------------------------------------------------------------------===// + +#include "aie/Dialect/AIE/IR/AIEDialect.h" +#include "aie/Dialect/AIEX/IR/AIEXDialect.h" +#include "aie/Dialect/AIEX/Transforms/AIEXPasses.h" + +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" + +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/SmallVector.h" + +#define DEBUG_TYPE "aie-coalesce-write32s" + +using namespace mlir; +using namespace xilinx; +using namespace xilinx::AIEX; +using namespace xilinx::AIE; + +namespace { + +// Represents a single 32-bit write to memory +struct WriteWord { + uint32_t address; + uint32_t value; + std::optional mask; // If present, this is a masked write + Operation *op; // The original operation that created this write + + bool operator<(const WriteWord &other) const { + return address < other.address; + } + + bool isMaskWrite() const { return mask.has_value(); } +}; + +// Represents a slice of execution between special register barriers +// All writes within a slice can be reordered +struct WriteSlice { + SmallVector writes; // All writes in this slice (can be reordered) + SmallVector specialOps; // Special register writes (cannot be reordered) + SmallVector otherOps; // Non-write operations that create barriers + + bool isEmpty() const { + return writes.empty() && specialOps.empty() && otherOps.empty(); + } +}; + +struct AIECoalesceWrite32sPass + : public AIECoalesceWrite32sBase { + + void runOnOperation() override { + AIE::DeviceOp deviceOp = getOperation(); + + // Collect all runtime sequences in the device + SmallVector runtimeSeqs; + deviceOp.walk([&](RuntimeSequenceOp seqOp) { + runtimeSeqs.push_back(seqOp); + }); + + // Process each runtime sequence + for (auto seqOp : runtimeSeqs) { + coalesceWrite32sInSequence(seqOp); + } + } + +private: + void coalesceWrite32sInSequence(RuntimeSequenceOp seqOp) { + OpBuilder builder(seqOp.getContext()); + + // Get the device and target model + auto deviceOp = seqOp->getParentOfType(); + if (!deviceOp) return; + + const auto &tm = deviceOp.getTargetModel(); + + // Walk through the sequence and partition into slices + for (Block &block : seqOp.getBody()) { + SmallVector slices; + WriteSlice currentSlice; + + for (Operation &op : block) { + // Skip get_global operations - they don't create barriers + if (isa(&op)) { + continue; + } + + bool isSpecialWrite = false; + + // Check if this is a write32 operation + if (auto write32Op = dyn_cast(&op)) { + // Only handle writes with absolute addresses (no buffer/col/row) + if (!write32Op.getBuffer() && !write32Op.getColumn() && + !write32Op.getRow()) { + uint32_t addr = write32Op.getAddress(); + isSpecialWrite = tm.isSpecialRegister(addr); + + if (isSpecialWrite) { + currentSlice.specialOps.push_back(&op); + } else { + WriteWord word{addr, write32Op.getValue(), std::nullopt, &op}; + currentSlice.writes.push_back(word); + } + } else { + // Symbolic write creates a barrier + if (!currentSlice.isEmpty()) { + slices.push_back(std::move(currentSlice)); + currentSlice = WriteSlice(); + } + currentSlice.otherOps.push_back(&op); + } + } + // Check if this is a maskwrite operation + else if (auto maskWriteOp = dyn_cast(&op)) { + // Only handle writes with absolute addresses (no buffer/col/row) + if (!maskWriteOp.getBuffer() && !maskWriteOp.getColumn() && + !maskWriteOp.getRow()) { + uint32_t addr = maskWriteOp.getAddress(); + isSpecialWrite = tm.isSpecialRegister(addr); + + if (isSpecialWrite) { + currentSlice.specialOps.push_back(&op); + } else { + WriteWord word{addr, maskWriteOp.getValue(), maskWriteOp.getMask(), &op}; + currentSlice.writes.push_back(word); + } + } else { + // Symbolic write creates a barrier + if (!currentSlice.isEmpty()) { + slices.push_back(std::move(currentSlice)); + currentSlice = WriteSlice(); + } + currentSlice.otherOps.push_back(&op); + } + } + // Check if this is a blockwrite operation + else if (auto blockWriteOp = dyn_cast(&op)) { + // Only handle writes with absolute addresses + if (!blockWriteOp.getBuffer() && !blockWriteOp.getColumn() && + !blockWriteOp.getRow()) { + auto addr = blockWriteOp.getAbsoluteAddress(); + auto dataWords = blockWriteOp.getDataWords(); + + if (addr && dataWords) { + uint32_t baseAddr = *addr; + + // Check if any word writes to a special register + int64_t numWords = dataWords.size(); + for (int64_t i = 0; i < numWords; ++i) { + if (tm.isSpecialRegister(baseAddr + i * 4)) { + isSpecialWrite = true; + break; + } + } + + if (isSpecialWrite) { + currentSlice.specialOps.push_back(&op); + } else { + // Add each word in the blockwrite as a separate WriteWord + int64_t idx = 0; + for (auto val : dataWords.getValues()) { + WriteWord word{ + static_cast(baseAddr + idx * 4), + static_cast(val.getZExtValue()), + std::nullopt, // blockwrites don't have masks + &op + }; + currentSlice.writes.push_back(word); + ++idx; + } + } + } else { + // Invalid blockwrite creates a barrier + if (!currentSlice.isEmpty()) { + slices.push_back(std::move(currentSlice)); + currentSlice = WriteSlice(); + } + currentSlice.otherOps.push_back(&op); + } + } else { + // Symbolic write creates a barrier + if (!currentSlice.isEmpty()) { + slices.push_back(std::move(currentSlice)); + currentSlice = WriteSlice(); + } + currentSlice.otherOps.push_back(&op); + } + } + // Any other operation creates a barrier + else { + if (!currentSlice.isEmpty()) { + slices.push_back(std::move(currentSlice)); + currentSlice = WriteSlice(); + } + currentSlice.otherOps.push_back(&op); + } + + // Special register writes create slice boundaries + if (isSpecialWrite) { + slices.push_back(std::move(currentSlice)); + currentSlice = WriteSlice(); + } + } + + // Don't forget the last slice + if (!currentSlice.isEmpty()) { + slices.push_back(std::move(currentSlice)); + } + + // Now process each slice: sort writes and coalesce + for (auto &slice : slices) { + processSlice(builder, slice, deviceOp); + } + } + } + + void processSlice(OpBuilder &builder, WriteSlice &slice, + AIE::DeviceOp deviceOp) { + // Non-write operations (barriers) are kept in their original positions + // Nothing to do here - they're already in the IR + + // If there are no writes to coalesce, nothing to do + if (slice.writes.empty()) { + return; + } + + // Eliminate duplicate writes - keep only the last write to each address + // Map from address to the index of the last write to that address + DenseMap lastWriteIndex; + DenseSet supersededOps; // Operations that are superseded by later writes + + for (size_t i = 0; i < slice.writes.size(); ++i) { + uint32_t addr = slice.writes[i].address; + + // If there was a previous write to this address, mark it as superseded + if (lastWriteIndex.count(addr)) { + supersededOps.insert(slice.writes[lastWriteIndex[addr]].op); + } + + lastWriteIndex[addr] = i; + } + + // Filter writes to keep only the last write to each address + SmallVector uniqueWrites; + for (size_t i = 0; i < slice.writes.size(); ++i) { + if (lastWriteIndex[slice.writes[i].address] == i) { + uniqueWrites.push_back(slice.writes[i]); + } + } + + // Update slice.writes with unique writes + slice.writes = std::move(uniqueWrites); + + // Sort writes by address to enable coalescing + llvm::sort(slice.writes); + + // Separate mask writes from regular writes + // Mask writes act as barriers for coalescing + DenseSet maskWriteAddresses; + for (auto &write : slice.writes) { + if (write.isMaskWrite()) { + maskWriteAddresses.insert(write.address); + } + } + + // Find contiguous sequences of non-masked writes and coalesce them + // Masked writes cannot be coalesced and act as barriers + SmallVector> sequences; + SmallVector currentSeq; + + for (auto &write : slice.writes) { + // Mask writes cannot be part of coalesced sequences + if (write.isMaskWrite()) { + // Save current sequence if it meets threshold + if (currentSeq.size() >= minWritesToCoalesce) { + sequences.push_back(currentSeq); + } + currentSeq.clear(); + continue; + } + + if (currentSeq.empty()) { + currentSeq.push_back(write); + } else if (write.address == currentSeq.back().address + 4 && + !maskWriteAddresses.count(write.address)) { + // Contiguous with previous write and not a mask write address + currentSeq.push_back(write); + } else { + // Not contiguous, start new sequence + // Only save sequences that meet the minimum threshold + if (currentSeq.size() >= minWritesToCoalesce) { + sequences.push_back(currentSeq); + } + currentSeq.clear(); + currentSeq.push_back(write); + } + } + + // Don't forget the last sequence (also check threshold) + if (currentSeq.size() >= minWritesToCoalesce) { + sequences.push_back(currentSeq); + } + + // Track which operations to erase + DenseSet toErase = supersededOps; // Start with superseded ops + + // Create blockwrites for each sequence + for (auto &seq : sequences) { + createBlockWrite(builder, seq, deviceOp); + + // Mark operations for erasure + for (auto &write : seq) { + toErase.insert(write.op); + } + } + + // For writes that couldn't be coalesced, keep them as write32 ops + for (auto &write : slice.writes) { + if (!toErase.contains(write.op)) { + // This write wasn't coalesced - keep it but possibly reorder + // The write is already in the IR, we just need to move it if needed + } + } + + // Special register writes stay in their original positions + // Nothing to do here - they're already in the IR + + // Erase operations that were coalesced or superseded + for (auto *op : toErase) { + op->erase(); + } + } + + void createBlockWrite(OpBuilder &builder, + const SmallVector &sequence, + AIE::DeviceOp deviceOp) { + if (sequence.empty()) return; + + MLIRContext *ctx = builder.getContext(); + uint32_t startAddr = sequence[0].address; + + // Generate a unique name for the global memref + static unsigned globalCounter = 0; + std::string globalName = "coalesced_write32_" + + std::to_string(startAddr) + "_" + + std::to_string(globalCounter++); + + // Collect values + SmallVector values; + for (auto &write : sequence) { + values.push_back(static_cast(write.value)); + } + + // Create memref type + auto i32Type = IntegerType::get(ctx, 32); + auto memrefType = MemRefType::get({static_cast(values.size())}, + i32Type); + + // Create DenseIntElementsAttr for initial values + auto tensorType = RankedTensorType::get({static_cast(values.size())}, + i32Type); + auto valuesAttr = DenseElementsAttr::get( + tensorType, ArrayRef(values)); + + // Get location from first operation + Location loc = sequence[0].op->getLoc(); + + // Insert the global at the beginning of the device + builder.setInsertionPointToStart(deviceOp.getBody()); + memref::GlobalOp::create( + builder, + loc, + globalName, + /*sym_visibility=*/builder.getStringAttr("private"), + /*type=*/memrefType, + /*initial_value=*/valuesAttr, + /*constant=*/true, + /*alignment=*/nullptr); + + // Create a GetGlobalOp and BlockWriteOp at the position of the first write + builder.setInsertionPoint(sequence[0].op); + auto getGlobalOp = memref::GetGlobalOp::create( + builder, loc, memrefType, globalName); + + NpuBlockWriteOp::create( + builder, + loc, + startAddr, + getGlobalOp.getResult(), + nullptr, // buffer + nullptr, // column + nullptr // row + ); + } +}; + +} // namespace + +namespace xilinx::AIEX { +std::unique_ptr> createAIECoalesceWrite32sPass() { + return std::make_unique(); +} +} // namespace xilinx::AIEX diff --git a/lib/Dialect/AIEX/Transforms/AIEDMATasksToNPU.cpp b/lib/Dialect/AIEX/Transforms/AIEDMATasksToNPU.cpp index 18aaed93724..5d669cee39c 100644 --- a/lib/Dialect/AIEX/Transforms/AIEDMATasksToNPU.cpp +++ b/lib/Dialect/AIEX/Transforms/AIEDMATasksToNPU.cpp @@ -12,9 +12,11 @@ #include #include "aie/Dialect/AIE/IR/AIEDialect.h" +#include "aie/Dialect/AIEX/AIEUtils.h" #include "aie/Dialect/AIEX/IR/AIEXDialect.h" #include "aie/Dialect/AIEX/Transforms/AIEXPasses.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Pass/Pass.h" #include "mlir/Transforms/DialectConversion.h" #include "llvm/ADT/TypeSwitch.h" @@ -184,14 +186,30 @@ struct AIEDMATasksToNPUPass : AIEDMATasksToNPUBase { uint64_t register_addr = target_model.getDmaBdAddress(tile.getCol(), tile.getRow(), bd_id) + target_model.getDmaBdAddressOffset(tile.getCol(), tile.getRow()); - if (mlir::BlockArgument buf_arg = - llvm::dyn_cast(buf)) { + + // A buffer descriptor can refer to a statically allocated aie.buffer, or to + // a DDR buffer which will be passed as a runtime arguument (block + // argument). Try to find the root block argument, either directly or + // through subviews/casts. + mlir::BlockArgument buf_arg = nullptr; + int64_t offset = 0; + + if (auto directArg = llvm::dyn_cast(buf)) { + buf_arg = directArg; + offset = 0; + } else if (auto traceResult = traceSubviewToBlockArgument(buf)) { + buf_arg = traceResult->rootArg; + offset = traceResult->offsetInBytes; + } + + if (buf_arg) { if (!target_model.isShimNOCTile(tile.getCol(), tile.getRow())) { return bd_op->emitOpError("DDR memory (runtime input arguments) can " "only be referred to on shim tiles."); } + unsigned arg_idx = buf_arg.getArgNumber(); - int64_t offset = bd_op.getOffsetInBytes(); + offset += bd_op.getOffsetInBytes(); NpuAddressPatchOp::create(builder, bd_op.getLoc(), /*addr*/ register_addr, /*arg_idx*/ arg_idx, @@ -209,9 +227,10 @@ struct AIEDMATasksToNPUPass : AIEDMATasksToNPUBase { NpuWrite32Op::create(builder, bd_op.getLoc(), register_addr, buf_addr, nullptr, nullptr, nullptr); } else { - return bd_op->emitOpError("Buffer argument must be either a constant " - "aie.buffer or a runtime " - "sequence input argument."); + return bd_op->emitOpError( + "Buffer argument must be a constant aie.buffer, a runtime sequence " + "input argument, or a (chain of) subview(s) or cast(s) of a block " + "argument with constant offsets and strides equal to one."); } return success(); } diff --git a/lib/Dialect/AIEX/Transforms/AIEDmaToNpu.cpp b/lib/Dialect/AIEX/Transforms/AIEDmaToNpu.cpp index 16794a3f3fe..ce091b331a5 100644 --- a/lib/Dialect/AIEX/Transforms/AIEDmaToNpu.cpp +++ b/lib/Dialect/AIEX/Transforms/AIEDmaToNpu.cpp @@ -14,6 +14,7 @@ #include "aie/Dialect/AIEX/IR/AIEXDialect.h" #include "aie/Dialect/AIEX/Transforms/AIEXPasses.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Pass/Pass.h" #include "mlir/Transforms/DialectConversion.h" #include @@ -284,7 +285,7 @@ struct DmaToNpuPattern : OpConversionPattern { return failure(); } - // arg_idx + // arg_idx and offset for block arguments AIE::RuntimeSequenceOp seq_op = op->getParentOfType(); if (!seq_op) { @@ -292,10 +293,26 @@ struct DmaToNpuPattern : OpConversionPattern { "time of lowering."); return failure(); } + + mlir::Value rootMemref = memref; + int64_t subviewOffset = 0; + + // Trace through memref.subview and memref.reinterpret_cast chain, if any, + // to find root block argument + auto traceResult = traceSubviewToBlockArgument(memref); + if (!traceResult) { + return op->emitOpError( + "memref must be a block argument or subview/cast/reinterpret_cast of " + "a block argument with static offsets, sizes, and strides"); + } + rootMemref = traceResult->rootArg; + subviewOffset = traceResult->offsetInBytes; + + // Find the argument index of the root memref Block &entryBB = seq_op.getBody().front(); int arg_idx = -1; for (int i = 0, e = entryBB.getNumArguments(); i < e; i++) { - if (entryBB.getArgument(i) == memref) { + if (entryBB.getArgument(i) == rootMemref) { arg_idx = i; break; } @@ -303,6 +320,8 @@ struct DmaToNpuPattern : OpConversionPattern { if (arg_idx < 0) return failure(); + offset += subviewOffset; + // bd_id bd_id = IntegerAttr::get(i32ty, op.getId()); diff --git a/lib/Dialect/AIEX/Transforms/AIEExpandLoadPdi.cpp b/lib/Dialect/AIEX/Transforms/AIEExpandLoadPdi.cpp new file mode 100644 index 00000000000..d14ab0ad954 --- /dev/null +++ b/lib/Dialect/AIEX/Transforms/AIEExpandLoadPdi.cpp @@ -0,0 +1,261 @@ +//===- AIEExpandLoadPdi.cpp -------------------------------------*- C++ -*-===// +// +// This file is licensed 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 +// +// (c) Copyright 2025 Advanced Micro Devices Inc. +// +//===----------------------------------------------------------------------===// +// +// This pass expands `npu.load_pdi` operations that reference a device into: +// 1. An empty device PDI load (causes firmware to reset the device) +// 2. Optional reset operations based on command-line configuration +// 3. Explicit configuration operations (write32/blockwrite) +// +// Example command-line usage: +// aie-opt --aie-expand-load-pdi \ +// --reset-switches-tiles=shim,mem,core \ +// --reset-switches-mode=ifused \ +// --reset-locks-tiles=mem,core \ +// --reset-locks-mode=always \ +// input.mlir +// +//===----------------------------------------------------------------------===// + +#include "aie/Conversion/AIEToConfiguration/AIEToConfiguration.h" +#include "aie/Dialect/AIE/IR/AIEDialect.h" +#include "aie/Dialect/AIEX/IR/AIEXDialect.h" +#include "aie/Dialect/AIEX/Transforms/AIEXPasses.h" + +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/DialectConversion.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" + +#include "llvm/ADT/StringRef.h" +#include "llvm/Support/raw_ostream.h" + +#define DEBUG_TYPE "aie-expand-load-pdi" + +using namespace mlir; +using namespace xilinx; +using namespace xilinx::AIEX; +using namespace xilinx::AIE; + +namespace { + +// Helper function to parse tile types from comma-separated string +static ResetTileType parseTileTypes(llvm::StringRef str) { + if (str.empty()) + return ResetTileType::None; + + unsigned result = 0; + llvm::SmallVector tokens; + str.split(tokens, ',', -1, false); + + for (auto token : tokens) { + token = token.trim(); + if (token == "shim" || token == "shimnoc") + result |= static_cast(ResetTileType::ShimNOC); + else if (token == "mem" || token == "memtile") + result |= static_cast(ResetTileType::MemTile); + else if (token == "core" || token == "coretile") + result |= static_cast(ResetTileType::CoreTile); + else if (token == "all") + return ResetTileType::All; + } + + return static_cast(result); +} + +// Helper function to parse reset mode from string +static ResetMode parseResetMode(llvm::StringRef str) { + auto trimmed = str.trim().lower(); + if (trimmed == "never") + return ResetMode::Never; + else if (trimmed == "ifused" || trimmed == "if-used") + return ResetMode::IfUsed; + else if (trimmed == "ifusedfinegrained" || trimmed == "if-used-fine-grained") + return ResetMode::IfUsedFineGrained; + else if (trimmed == "ifchanged" || trimmed == "if-changed") + return ResetMode::IfChanged; + else if (trimmed == "ifchangedfinegrained" || + trimmed == "if-changed-fine-grained") + return ResetMode::IfChangedFineGrained; + else if (trimmed == "always" || trimmed == "all") + return ResetMode::Always; + else + return ResetMode::Never; +} + +struct ExpandLoadPdiPattern : public OpRewritePattern { + const ResetConfig dmaConfig; + const ResetConfig switchConfig; + const ResetConfig lockConfig; + const ResetConfig coreConfig; + mutable AIE::DeviceOp previousDevice; + + ExpandLoadPdiPattern(MLIRContext *context, ResetConfig dma, ResetConfig sw, + ResetConfig lock, ResetConfig core) + : OpRewritePattern(context), dmaConfig(dma), + switchConfig(sw), lockConfig(lock), coreConfig(core), + previousDevice(nullptr) {} + + LogicalResult matchAndRewrite(NpuLoadPdiOp loadPdiOp, + PatternRewriter &rewriter) const override { + // Only process load_pdi ops that reference a device + auto deviceRefAttr = loadPdiOp.getDeviceRefAttr(); + if (!deviceRefAttr) + return failure(); + + // Look up the referenced device + auto moduleOp = loadPdiOp->getParentOfType(); + if (!moduleOp) + return failure(); + + auto referencedDevice = moduleOp.lookupSymbol(deviceRefAttr); + if (!referencedDevice) { + loadPdiOp.emitError("Referenced symbol '") + << deviceRefAttr.getValue() << "' is not a device"; + return failure(); + } + + // If already referencing an empty device, don't transform + if (referencedDevice.getSymName().starts_with("empty")) + return failure(); + + // Determine if this is the first load_pdi (no previous device) + bool isFirstLoadPdi = !previousDevice; + + // For first load_pdi with no reset modes requiring comparison, skip empty + // device + bool needsEmptyDevice = !isFirstLoadPdi || + dmaConfig.mode != ResetMode::Never || + switchConfig.mode != ResetMode::Never || + lockConfig.mode != ResetMode::Never; + + AIE::DeviceOp emptyDevice; + if (needsEmptyDevice) { + // Create a unique empty device for this reset to avoid PDI address + // caching + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPointToStart(moduleOp.getBody()); + + // Find a unique name for the empty device + std::string emptyName; + int emptyIndex = 0; + do { + emptyName = "empty_" + std::to_string(emptyIndex++); + } while (moduleOp.lookupSymbol(emptyName)); + + auto deviceType = referencedDevice.getDevice(); + auto loc = rewriter.getUnknownLoc(); + emptyDevice = AIE::DeviceOp::create(rewriter, loc, deviceType, + rewriter.getStringAttr(emptyName)); + emptyDevice.getRegion().emplaceBlock(); + + Block *deviceBlock = &emptyDevice.getRegion().front(); + rewriter.setInsertionPointToEnd(deviceBlock); + AIE::EndOp::create(rewriter, loc); + } + + // Create new load_pdi operation + rewriter.setInsertionPoint(loadPdiOp); + NpuLoadPdiOp newLoadPdi; + if (needsEmptyDevice) { + newLoadPdi = NpuLoadPdiOp::create( + rewriter, loadPdiOp.getLoc(), + FlatSymbolRefAttr::get(emptyDevice.getSymNameAttr()), + loadPdiOp.getIdAttr(), loadPdiOp.getSizeAttr(), + loadPdiOp.getAddressAttr()); + } else { + // First load_pdi with no resets - keep original device reference + newLoadPdi = NpuLoadPdiOp::create( + rewriter, loadPdiOp.getLoc(), loadPdiOp.getDeviceRefAttr(), + loadPdiOp.getIdAttr(), loadPdiOp.getSizeAttr(), + loadPdiOp.getAddressAttr()); + } + + // Generate and insert reset operations (skip for first load_pdi) + if (!isFirstLoadPdi && (dmaConfig.mode != ResetMode::Never || + switchConfig.mode != ResetMode::Never || + lockConfig.mode != ResetMode::Never || + coreConfig.mode != ResetMode::Never)) { + if (failed(xilinx::AIE::generateAndInsertResetOps( + referencedDevice, newLoadPdi.getOperation(), dmaConfig, + switchConfig, lockConfig, coreConfig, previousDevice))) { + loadPdiOp.emitError("Failed to generate reset operations"); + return failure(); + } + } + + // Generate and insert the configuration operations after the reset ops + // Find the last operation inserted (which should be from reset ops) + Operation *lastResetOp = newLoadPdi.getOperation(); + Operation *nextOp = newLoadPdi->getNextNode(); + while (nextOp && nextOp != loadPdiOp.getOperation()) { + lastResetOp = nextOp; + nextOp = nextOp->getNextNode(); + } + + if (failed(xilinx::AIE::generateAndInsertConfigOps(referencedDevice, + lastResetOp, ""))) { + loadPdiOp.emitError("Failed to generate configuration operations"); + return failure(); + } + + // Remove the original load_pdi operation + rewriter.eraseOp(loadPdiOp); + + // Update previous device for next load_pdi + previousDevice = referencedDevice; + + return success(); + } +}; + +struct AIEExpandLoadPdiPass + : public AIEExpandLoadPdiBase { + void getDependentDialects(DialectRegistry ®istry) const override { + registry + .insert(); + } + + void runOnOperation() override { + auto module = getOperation(); + + // Parse reset configurations from command-line options + ResetTileType dmaTiles = parseTileTypes(resetDmasTiles); + ResetMode dmaMode = parseResetMode(resetDmasMode); + ResetConfig dmaConfig(dmaTiles, dmaMode); + + ResetTileType switchTiles = parseTileTypes(resetSwitchesTiles); + ResetMode switchMode = parseResetMode(resetSwitchesMode); + ResetConfig switchConfig(switchTiles, switchMode); + + ResetTileType lockTiles = parseTileTypes(resetLocksTiles); + ResetMode lockMode = parseResetMode(resetLocksMode); + ResetConfig lockConfig(lockTiles, lockMode); + + ResetTileType coreTiles = parseTileTypes(resetCoresTiles); + ResetMode coreMode = parseResetMode(resetCoresMode); + ResetConfig coreConfig(coreTiles, coreMode); + + RewritePatternSet patterns(&getContext()); + patterns.add(&getContext(), dmaConfig, switchConfig, + lockConfig, coreConfig); + + if (failed(applyPatternsGreedily(module, std::move(patterns)))) { + signalPassFailure(); + } + } +}; + +} // namespace + +std::unique_ptr> +xilinx::AIEX::createAIEExpandLoadPdiPass() { + return std::make_unique(); +} diff --git a/lib/Dialect/AIEX/Transforms/AIEMaterializeRuntimeSequences.cpp b/lib/Dialect/AIEX/Transforms/AIEMaterializeRuntimeSequences.cpp new file mode 100644 index 00000000000..bb15f4c4591 --- /dev/null +++ b/lib/Dialect/AIEX/Transforms/AIEMaterializeRuntimeSequences.cpp @@ -0,0 +1,540 @@ +//===- AIEMaterializeRuntimeSequence.cpp -------------------------*- C++ +//-*-===// +// +// This file is licensed 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 +// +// (c) Copyright 2025 Advanced Micro Devices Inc. +// +//===----------------------------------------------------------------------===// + +#include "aie/Conversion/AIEToConfiguration/AIEToConfiguration.h" +#include "aie/Dialect/AIE/IR/AIEDialect.h" +#include "aie/Dialect/AIEX/AIEUtils.h" +#include "aie/Dialect/AIEX/IR/AIEXDialect.h" +#include "aie/Dialect/AIEX/Transforms/AIEXPasses.h" +#include "aie/Targets/AIERT.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/IR/IRMapping.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/DialectConversion.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "mlir/Transforms/WalkPatternRewriteDriver.h" + +#define DEBUG_TYPE "aie-materialize-runtime-sequence" + +using namespace mlir; +using namespace xilinx; +using namespace xilinx::AIEX; + +struct RuntimeCallGraphCyclicityAnalysis { + AnalysisManager &analysisManager; + + // if invalid, analysis failed and results should not be considered + bool isValid = false; + + // Call graph is cyclic + bool isCyclic = false; + + RuntimeCallGraphCyclicityAnalysis(Operation *op, AnalysisManager &am) + : analysisManager(am) { + AIE::RuntimeSequenceOp runtimeSequenceOp = + llvm::dyn_cast(op); + if (!runtimeSequenceOp) { + op->emitError("RuntimeCallGraphCyclicityAnalysis can only be called on " + "aiex.runtime_sequence operations."); + return; + } + + // Use DFS with a stack to detect cycles + // A cycle exists if we encounter a sequence already on the current path + llvm::DenseSet callStack; + llvm::DenseSet visited; + + std::function hasCycle = + [&](AIE::RuntimeSequenceOp seq) -> bool { + if (callStack.contains(seq)) { + return true; // Found a cycle + } + if (visited.contains(seq)) { + return false; // Already checked this sequence + } + + callStack.insert(seq); + visited.insert(seq); + + // Check all sequences called by this one + bool foundCycle = false; + seq.walk([&](RunOp runOp) { + if (AIE::RuntimeSequenceOp callee = + runOp.getCalleeRuntimeSequenceOp()) { + if (hasCycle(callee)) { + foundCycle = true; + return WalkResult::interrupt(); + } + } + return WalkResult::advance(); + }); + + callStack.erase(seq); + return foundCycle; + }; + + if (hasCycle(runtimeSequenceOp)) { + isCyclic = true; + isValid = true; + return; + } + isCyclic = false; + isValid = true; + } +}; + +// Turn aie.configure @device into aie.run %.. @configure +// TODO: add check that liveness of two aie.configures do not overlap +// (i.e., when we configure A, then configure B, cannot call runtime sequence of +// A after configuring B) +// TODO: add code to remove repeated @configure ops +struct InsertLoadPdiForConfigurePattern : RewritePattern { + + InsertLoadPdiForConfigurePattern(MLIRContext *context, + PatternBenefit benefit = 1) + : RewritePattern(ConfigureOp::getOperationName(), benefit, context) {} + + LogicalResult matchAndRewrite(Operation *op, + PatternRewriter &rewriter) const override { + ConfigureOp configureOp = llvm::dyn_cast(op); + if (!configureOp) { + return failure(); + } + + // LoadPDI resets the whole device, hence cannot do partial reconfiguration; + // therefore, this only supports top-level configure ops + if (!llvm::isa(configureOp->getParentOp())) { + return failure(); + } + + AIE::DeviceOp referencedDevice = configureOp.getReferencedDeviceOp(); + if (!referencedDevice) { + configureOp.emitError("Referenced symbol is not a device"); + return failure(); + } + + Block *configureBlock; + if (configureOp.getBody().empty()) { + configureBlock = rewriter.createBlock(&configureOp.getBody()); + } else { + configureBlock = &configureOp.getBody().front(); + } + rewriter.setInsertionPointToStart(configureBlock); + AIEX::NpuLoadPdiOp::create( + rewriter, configureOp.getLoc(), + FlatSymbolRefAttr::get(referencedDevice.getSymNameAttr())); + + return success(); + } +}; + +// Collects all external SSA values referenced by an operation (and its nested +// operations). +// 1. Collects SSA values from the operation's operands. +// 2. Recursively walks through all operations in the operation's regions. +// 3. For each nested operation, collects SSA values from its operands. +// 4. Skips values that are already in argMap or defined within the operation. +// 5. For memref.subview operations, traces to the root block argument +static void +collectReferencedSSAValues(Operation *op, const IRMapping &argMap, + llvm::SetVector &referencedValues) { + + auto processValue = [&](Value operand) { + if (argMap.contains(operand)) { + return; + } + + // If this is a subview, trace to the root block argument + if (auto traceResult = traceSubviewToBlockArgument(operand)) { + // Check if the root argument is already mapped + if (!argMap.contains(traceResult->rootArg)) { + referencedValues.insert(traceResult->rootArg); + } + return; + } + + // Not a subview chain leading to block arg, add as-is + referencedValues.insert(operand); + }; + + // Collect SSA values from the operation's direct operands. + for (Value operand : op->getOperands()) { + processValue(operand); + } + + // Recursively collect SSA values from nested operations in all regions. + for (Region ®ion : op->getRegions()) { + region.walk([&](Operation *nestedOp) { + for (Value operand : nestedOp->getOperands()) { + if (argMap.contains(operand)) { + return; + } + + // Check if defined within the parent operation + Operation *defOp = operand.getDefiningOp(); + if (defOp && op->isProperAncestor(defOp)) { + return; + } + + processValue(operand); + } + }); + } +} + +// Copies SSA value definitions into the caller device. +// Currently, only `aie.tile` operations are supported. +// Updates argMap to map old values to new/existing values. +static LogicalResult +copyReferencedSSAValues(PatternRewriter &rewriter, + const llvm::SetVector &referencedValues, + AIE::DeviceOp callerDevice, IRMapping &argMap, + mlir::OpBuilder::InsertPoint &clonedSSAInsertionPoint, + Operation *errorReportOp) { + + for (Value referencedValue : referencedValues) { + Operation *definingOp = referencedValue.getDefiningOp(); + if (!definingOp) { + return errorReportOp->emitError() + << "Referenced value is not defined by an operation"; + } + + if (auto tileOp = llvm::dyn_cast(definingOp)) { + int col = tileOp.getCol(); + int row = tileOp.getRow(); + + // Check if a tile with matching col/row already exists in the caller + // device + AIE::TileOp existingTile = nullptr; + for (AIE::TileOp tile : callerDevice.getOps()) { + if (tile.getCol() == col && tile.getRow() == row) { + existingTile = tile; + break; + } + } + + if (existingTile) { + // Verify that all attributes match + if (tileOp->getAttrDictionary() != existingTile->getAttrDictionary()) { + // Filter out result type attributes and symbol attributes for + // comparison + auto filterAttrs = [](DictionaryAttr dict) -> DictionaryAttr { + SmallVector filteredAttrs; + for (auto namedAttr : dict) { + StringRef name = namedAttr.getName().getValue(); + if (name != "col" && name != "row") { + filteredAttrs.push_back(namedAttr); + } + } + return DictionaryAttr::get(dict.getContext(), filteredAttrs); + }; + + DictionaryAttr tileAttrs = filterAttrs(tileOp->getAttrDictionary()); + DictionaryAttr existingAttrs = + filterAttrs(existingTile->getAttrDictionary()); + + if (tileAttrs != existingAttrs) { + return errorReportOp->emitError() + << "aie.tile(" << col << ", " << row + << ") already exists in the device with different " + "attributes"; + } + } + + // Use the existing tile + rewriter.restoreInsertionPoint(clonedSSAInsertionPoint); + if (!existingTile->isBeforeInBlock(&*rewriter.getInsertionPoint())) { + rewriter.moveOpBefore(existingTile, + clonedSSAInsertionPoint.getBlock(), + clonedSSAInsertionPoint.getPoint()); + rewriter.setInsertionPointAfter(existingTile); + } + argMap.map(referencedValue, existingTile.getResult()); + } else { + // Clone the tile operation into the caller device + rewriter.restoreInsertionPoint(clonedSSAInsertionPoint); + Operation *clonedTile = rewriter.clone(*tileOp); + argMap.map(referencedValue, clonedTile->getResult(0)); + clonedSSAInsertionPoint = rewriter.saveInsertionPoint(); + } + + } else { + return errorReportOp->emitError() + << "Referenced SSA value defined by unsupported operation type: " + << definingOp->getName().getStringRef() + << ". Currently only aie.tile operations are supported."; + } + } + + return success(); +} + +// Inlines the definitions of all symbols referenced in the given operation +// at the current insertion point in the given rewriter, unless the symbol +// definition is in the "previouslyInlinedSymbolMap" map. While inlining, +// symbols will be renamed to have a unique name. +// Also copies in SSA values referenced by the inlined symbol definitions. +static LogicalResult inlineReferencedSymbolDefinitions( + PatternRewriter &rewriter, Operation *op, Operation *lookupFrom, + IRMapping argMap, + llvm::DenseMap &previouslyInlinedSymbolMap, + AIE::DeviceOp callerDevice, + mlir::OpBuilder::InsertPoint &clonedDefOpsInsertionPoint) { + MLIRContext *ctx = op->getContext(); + for (NamedAttribute namedAttr : op->getAttrs()) { + Attribute attr = namedAttr.getValue(); + auto newAttr = attr.replace([&](SymbolRefAttr oldSymbolRef) { + SymbolRefAttr newSymbolRef; + if (!previouslyInlinedSymbolMap.count(oldSymbolRef)) { + llvm::StringRef oldName = oldSymbolRef.getRootReference().getValue(); + std::string uniqueName = oldName.str(); + unsigned uniquingCounter = 0; + while (SymbolTable::lookupNearestSymbolFrom( + op, StringAttr::get(ctx, uniqueName))) { + uniqueName = oldName.str() + "_" + std::to_string(uniquingCounter); + uniquingCounter++; + } + newSymbolRef = SymbolRefAttr::get(ctx, uniqueName); + previouslyInlinedSymbolMap[oldSymbolRef] = newSymbolRef; + + // Add the new symbol definition + Operation *symbolDefOp = + SymbolTable::lookupNearestSymbolFrom(lookupFrom, oldSymbolRef); + if (!symbolDefOp) { + return std::make_pair(newSymbolRef, WalkResult::interrupt()); + } + + // Collect SSA values referenced by the symbol definition operation + llvm::SetVector symbolReferencedValues; + collectReferencedSSAValues(symbolDefOp, argMap, symbolReferencedValues); + + // Copy SSA values referenced by the symbol definition + if (failed(copyReferencedSSAValues(rewriter, symbolReferencedValues, + callerDevice, argMap, + clonedDefOpsInsertionPoint, op))) { + return std::make_pair(newSymbolRef, WalkResult::interrupt()); + } + + Operation *clonedSymbolDefOp = rewriter.clone(*symbolDefOp, argMap); + clonedSymbolDefOp->setAttr(SymbolTable::getSymbolAttrName(), + StringAttr::get(ctx, uniqueName)); + } else { + newSymbolRef = previouslyInlinedSymbolMap[oldSymbolRef]; + } + + return std::make_pair(newSymbolRef, WalkResult::advance()); + }); + if (!newAttr) { + return failure(); + } + op->setAttr(namedAttr.getName(), newAttr); + } + return success(); +} + +struct InlineRuntimeCallsPattern : RewritePattern { + + InlineRuntimeCallsPattern(MLIRContext *ctx) + : RewritePattern(RunOp::getOperationName(), PatternBenefit(1), ctx) {} + + LogicalResult matchAndRewrite(Operation *op, + PatternRewriter &rewriter) const override { + // matching logic + RunOp runOp = llvm::dyn_cast(op); + if (!runOp) { + return failure(); + } + AIE::DeviceOp calleeDevice = runOp.getCalleeDeviceOp(); + AIE::RuntimeSequenceOp calleeRuntimeSequence = + runOp.getCalleeRuntimeSequenceOp(); + if (!calleeDevice || !calleeRuntimeSequence) { + return failure(); + } + if (!calleeRuntimeSequence.getOps().empty()) { + return failure(); + } + + // rewrite logic + + // Get caller and callee bodies. The callee body will be inlined into the + // caller body at the point of the RunOp. + Region &calleeBody = calleeRuntimeSequence.getBody(); + AIE::DeviceOp callerDevice = op->getParentOfType(); + if (!callerDevice) { + runOp.emitError() << "needs to be in a DeviceOp"; + return failure(); + } + Region &callerDeviceBody = callerDevice.getBodyRegion(); + + // The argMap maps callee arguments to caller SSA values. + IRMapping argMap; + ValueRange values = runOp.getArgs(); + for (unsigned i = 0, n = calleeBody.getNumArguments(); i < n; i++) { + BlockArgument arg = calleeBody.getArgument(i); + Value val = values[i]; + argMap.map(arg, val); + } + + // The callee body may reference SSA values and symbols that are defined + // in the callee device (outside the callee runtime sequence). We will + // inline a supported set of these and error otherwise. + + // Collect SSA values referenced in the callee not defined by the callee and + // and not in the argMap. + llvm::SetVector referencedValues; + for (Operation &op : calleeBody.getOps()) { + collectReferencedSSAValues(&op, argMap, referencedValues); + } + llvm::SetVector filteredValues; + for (Value val : referencedValues) { + if (val.getParentRegion() != &calleeBody) { + filteredValues.insert(val); + } + } + referencedValues = std::move(filteredValues); + + // Copy the operations that define these SSA values into the caller device + mlir::Block &callerDeviceBodyFirstBlock = callerDeviceBody.front(); + mlir::OpBuilder::InsertPoint clonedDefOpsInsertPoint( + &callerDeviceBodyFirstBlock, callerDeviceBodyFirstBlock.begin()); + if (failed(copyReferencedSSAValues(rewriter, referencedValues, callerDevice, + argMap, clonedDefOpsInsertPoint, + runOp))) { + return failure(); + } + + // Now, also inline symbol definitions referenced in the callee body; + // this may pull in additional SSA values referenced by the symbol + // definitions. + rewriter.setInsertionPoint(runOp); + mlir::OpBuilder::InsertPoint clonedOpInsertionPoint = + rewriter.saveInsertionPoint(); + llvm::DenseMap previouslyInlinedSymbolMap; + for (Operation &op : calleeBody.getOps()) { + rewriter.restoreInsertionPoint(clonedOpInsertionPoint); + Operation *clonedOp = rewriter.clone(op, argMap); + clonedOpInsertionPoint = rewriter.saveInsertionPoint(); + + if (failed(inlineReferencedSymbolDefinitions( + rewriter, clonedOp, calleeRuntimeSequence.getOperation(), argMap, + previouslyInlinedSymbolMap, callerDevice, + clonedDefOpsInsertPoint))) { + return failure(); + } + } + + // The aiex.run op has been inlined; erase it. + rewriter.eraseOp(runOp); + + return success(); + } +}; + +struct AIEMaterializeRuntimeSequencesPass + : AIEMaterializeRuntimeSequencesBase { + void runOnOperation() override { + ModuleOp moduleOp = getOperation(); + + // Process each device in the module + for (AIE::DeviceOp deviceOp : moduleOp.getOps()) { + + // Verify all runtime sequences before materialization + for (AIE::RuntimeSequenceOp runtimeSequenceOp : + deviceOp.getOps()) { + if (failed(runtimeSequenceOp.verifyBeforeMaterialization())) { + return signalPassFailure(); + } + } + + // Check for cycles in runtime sequence calls + for (AIE::RuntimeSequenceOp runtimeSequenceOp : + deviceOp.getOps()) { + AnalysisManager am = + getAnalysisManager().nest(deviceOp).nest(runtimeSequenceOp); + RuntimeCallGraphCyclicityAnalysis cyclicity = + am.getAnalysis(); + if (!cyclicity.isValid) { + return signalPassFailure(); + } + if (cyclicity.isCyclic) { + runtimeSequenceOp.emitError( + "Runtime sequence call graph contains a cycle"); + return signalPassFailure(); + } + } + + // Greedily inline all runtime sequences that can be inlined; + // this will start with runtime sequences that do not call other runtime + // sequences (leaves); once their callers inline them, the callers can + // be inlined as well, and so on + MLIRContext *ctx = &getContext(); + GreedyRewriteConfig rewriter_config = GreedyRewriteConfig(); + rewriter_config.setRegionSimplificationLevel( + GreedySimplifyRegionLevel::Disabled); + + RewritePatternSet patterns_0(ctx); + patterns_0.insert(ctx); + if (failed(applyPatternsGreedily(deviceOp, std::move(patterns_0), + rewriter_config))) { + return signalPassFailure(); + } + + // Insert LoadPDI ops for each aiex.configure op + for (AIE::RuntimeSequenceOp runtimeSequenceOp : + deviceOp.getOps()) { + RewritePatternSet patterns_1(ctx); + patterns_1.insert(ctx); + walkAndApplyPatterns(runtimeSequenceOp, std::move(patterns_1)); + } + + // Flatten the IR: hoist all operations inside aiex.configure to be direct + // children of the runtime sequence, preserving order + for (AIE::RuntimeSequenceOp runtimeSequenceOp : + deviceOp.getOps()) { + SmallVector configureOps; + + for (ConfigureOp configureOp : + runtimeSequenceOp.getOps()) { + configureOps.push_back(configureOp); + } + + IRRewriter rewriter(ctx); + for (ConfigureOp configureOp : configureOps) { + Block &configureBlock = configureOp.getBody().front(); + + // Collect all operations in the configure block + SmallVector opsToHoist; + for (Operation &op : configureBlock) { + opsToHoist.push_back(&op); + } + + // Hoist operations to be right before the configure op + rewriter.setInsertionPoint(configureOp); + for (Operation *op : opsToHoist) { + op->moveBefore(configureOp); + } + + // Erase the now-empty configure op + rewriter.eraseOp(configureOp); + } + } + + } // end for each device + } +}; + +std::unique_ptr> +AIEX::createAIEMaterializeRuntimeSequencesPass() { + return std::make_unique(); +} diff --git a/lib/Dialect/AIEX/Transforms/CMakeLists.txt b/lib/Dialect/AIEX/Transforms/CMakeLists.txt index aebe62d7732..4679734f2e2 100644 --- a/lib/Dialect/AIEX/Transforms/CMakeLists.txt +++ b/lib/Dialect/AIEX/Transforms/CMakeLists.txt @@ -15,6 +15,7 @@ add_mlir_dialect_library(AIEXTransforms AIELowerMemcpy.cpp AIEDmaToNpu.cpp AIEMaterializeBDChains.cpp + AIEMaterializeRuntimeSequences.cpp AIEAssignRuntimeSequenceBDIDs.cpp AIEDMATasksToNPU.cpp AIESubstituteShimDMAAllocations.cpp @@ -22,6 +23,8 @@ add_mlir_dialect_library(AIEXTransforms AIELowerSetLock.cpp AIETransformBfpTypes.cpp AIETxnToControlPacket.cpp + AIEExpandLoadPdi.cpp + AIECoalesceWrite32s.cpp ADDITIONAL_HEADER_DIRS ${AIE_BINARY_DIR}/include @@ -33,6 +36,7 @@ add_mlir_dialect_library(AIEXTransforms LINK_LIBS PUBLIC AIE + MLIRAIEToConfiguration AIEXUtils MLIRIR MLIRPass diff --git a/lib/Dialect/AIEX/Utils/AIEUtils.cpp b/lib/Dialect/AIEX/Utils/AIEUtils.cpp index 1a7b7638c9e..1a2676ec468 100644 --- a/lib/Dialect/AIEX/Utils/AIEUtils.cpp +++ b/lib/Dialect/AIEX/Utils/AIEUtils.cpp @@ -15,6 +15,101 @@ using namespace xilinx; static unsigned cachedId = 0; +std::optional +AIEX::traceSubviewToBlockArgument(Value value) { + int64_t offsetInBytes = 0; + Value current = value; + + // Walk through the chain of operations until we reach a block argument + while (current) { + // Check if we've reached a block argument + if (auto blockArg = dyn_cast(current)) { + return SubviewTraceResult{blockArg, offsetInBytes}; + } + + Operation *defOp = current.getDefiningOp(); + if (!defOp) { + return std::nullopt; + } + + // Handle memref.cast (just pass through) + if (auto castOp = dyn_cast(defOp)) { + current = castOp.getSource(); + continue; + } + + // Handle memref.reinterpret_cast (validate and pass through) + if (auto reinterpretOp = dyn_cast(defOp)) { + auto sourceType = + dyn_cast(reinterpretOp.getSource().getType()); + if (!sourceType) { + return std::nullopt; + } + + // Validate that source is contiguous (all strides must be 1) + if (auto strided = dyn_cast(sourceType.getLayout())) { + for (int64_t stride : strided.getStrides()) { + if (stride != 1) { + return std::nullopt; // Non-contiguous memory, cannot safely + // reinterpret + } + } + } + + current = reinterpretOp.getSource(); + continue; + } + + // Handle memref.subview (accumulate offset and validate) + if (auto subviewOp = dyn_cast(defOp)) { + // Verify static offsets, sizes, strides + if (!subviewOp.getStaticOffsets().empty() && + subviewOp.getStaticOffsets()[0] == ShapedType::kDynamic) { + return std::nullopt; + } + if (!subviewOp.getStaticSizes().empty() && + subviewOp.getStaticSizes()[0] == ShapedType::kDynamic) { + return std::nullopt; + } + if (!subviewOp.getStaticStrides().empty() && + subviewOp.getStaticStrides()[0] == ShapedType::kDynamic) { + return std::nullopt; + } + + // Only support rank-1 subviews + if (subviewOp.getSourceType().getRank() != 1 || + subviewOp.getType().getRank() != 1) { + return std::nullopt; + } + + // Only support stride of 1 (contiguous) + if (!subviewOp.getStaticStrides().empty() && + subviewOp.getStaticStrides()[0] != 1) { + return std::nullopt; + } + + // Calculate and accumulate offset in bytes + auto sourceType = subviewOp.getSourceType(); + unsigned elemSizeInBits = + sourceType.getElementType().getIntOrFloatBitWidth(); + if (elemSizeInBits % 8 != 0) { + return std::nullopt; + } + unsigned elemSizeInBytes = elemSizeInBits / 8; + int64_t offsetInElements = subviewOp.getStaticOffsets()[0]; + offsetInBytes += offsetInElements * elemSizeInBytes; + + current = subviewOp.getSource(); + continue; + } + + // Encountered an unsupported operation + return std::nullopt; + } + + return std::nullopt; +} + memref::GlobalOp AIEX::getOrCreateDataMemref(OpBuilder &builder, AIE::DeviceOp dev, mlir::Location loc, diff --git a/lib/Targets/AIERT.cpp b/lib/Targets/AIERT.cpp index f62ec0cc9a5..32b53b096ad 100644 --- a/lib/Targets/AIERT.cpp +++ b/lib/Targets/AIERT.cpp @@ -20,7 +20,9 @@ extern "C" { #include "xaiengine/xaie_interrupt.h" #include "xaiengine/xaie_locks.h" #include "xaiengine/xaie_mem.h" +#include "xaiengine/xaie_perfcnt.h" #include "xaiengine/xaie_plif.h" +#include "xaiengine/xaie_reset.h" #include "xaiengine/xaie_ss.h" #include "xaiengine/xaie_txn.h" #include "xaiengine/xaiegbl.h" @@ -843,6 +845,137 @@ LogicalResult xilinx::AIE::AIERTControl::addAieElf(uint8_t col, uint8_t row, return success(); } +LogicalResult xilinx::AIE::AIERTControl::resetPartition() { + TRY_XAIE_API_LOGICAL_RESULT(XAie_ResetPartition, &aiert->devInst); + return success(); +} + +LogicalResult xilinx::AIE::AIERTControl::resetDMA(int col, int row, bool on) { + auto tileLoc = XAie_TileLoc(col, row); + XAie_DmaDesc dmaTileBd; + TRY_XAIE_API_LOGICAL_RESULT(XAie_DmaDescInit, &aiert->devInst, &dmaTileBd, + tileLoc); + TRY_XAIE_API_LOGICAL_RESULT(XAie_DmaDisableBd, &dmaTileBd); + TRY_XAIE_API_LOGICAL_RESULT(XAie_DmaChannelResetAll, &aiert->devInst, tileLoc, + on ? XAie_DmaChReset::DMA_CHANNEL_UNRESET + : XAie_DmaChReset::DMA_CHANNEL_RESET); + return success(); +} + +LogicalResult xilinx::AIE::AIERTControl::resetCore(int col, int row) { + auto tileLoc = XAie_TileLoc(col, row); + TRY_XAIE_API_LOGICAL_RESULT(XAie_CoreReset, &aiert->devInst, tileLoc); + return success(); +} + +LogicalResult xilinx::AIE::AIERTControl::resetSwitch(int col, int row) { + XAie_LocType tileLoc = XAie_TileLoc(col, row); + + // Reset all combinations of input/output routing in the switchbox + for (auto endpoint_a : WIRE_BUNDLE_TO_STRM_SW_PORT_TYPE) { + for (auto endpoint_b : WIRE_BUNDLE_TO_STRM_SW_PORT_TYPE) { + unsigned n_a_connections = targetModel.getNumSourceSwitchboxConnections( + col, row, endpoint_a.first); + unsigned n_b_connections = targetModel.getNumDestSwitchboxConnections( + col, row, endpoint_b.first); + for (unsigned a_index = 0; a_index < n_a_connections; a_index++) { + for (unsigned b_index = 0; b_index < n_b_connections; b_index++) { + if (!targetModel.isLegalTileConnection(col, row, endpoint_a.first, + a_index, endpoint_b.first, + b_index)) { + continue; + } + TRY_XAIE_API_FATAL_ERROR(XAie_StrmConnCctDisable, &aiert->devInst, + tileLoc, endpoint_a.second, a_index, + endpoint_b.second, b_index); + } + } + } + } + + return success(); +} + +LogicalResult xilinx::AIE::AIERTControl::resetCoreUnreset(int col, int row) { + auto tileLoc = XAie_TileLoc(col, row); + TRY_XAIE_API_LOGICAL_RESULT(XAie_CoreUnreset, &aiert->devInst, tileLoc); + return success(); +} + +LogicalResult xilinx::AIE::AIERTControl::resetLock(int col, int row, + int lockId) { + auto tileLoc = XAie_TileLoc(col, row); + // Reset a single lock to value 0 + XAie_Lock lock; + lock.LockId = lockId; + lock.LockVal = 0; + TRY_XAIE_API_LOGICAL_RESULT(XAie_LockSetValue, &aiert->devInst, tileLoc, + lock); + return success(); +} + +LogicalResult xilinx::AIE::AIERTControl::resetSwitchConnection( + int col, int row, WireBundle sourceBundle, int sourceChannel, + WireBundle destBundle, int destChannel) { + auto tileLoc = XAie_TileLoc(col, row); + + // Helper lambda to map WireBundle to StrmSwPortType + auto mapBundle = [](WireBundle bundle) -> StrmSwPortType { + switch (bundle) { + case WireBundle::Core: + return CORE; + case WireBundle::DMA: + return DMA; + case WireBundle::FIFO: + return FIFO; + case WireBundle::South: + return SOUTH; + case WireBundle::West: + return WEST; + case WireBundle::North: + return NORTH; + case WireBundle::East: + return EAST; + case WireBundle::Trace: + return TRACE; + default: + return SOUTH; + } + }; + + StrmSwPortType sourcePortType = mapBundle(sourceBundle); + StrmSwPortType destPortType = mapBundle(destBundle); + + // Disconnect the specific connection from source to destination + TRY_XAIE_API_LOGICAL_RESULT(XAie_StrmConnCctDisable, &aiert->devInst, tileLoc, + sourcePortType, sourceChannel, destPortType, + destChannel); + + return success(); +} + +LogicalResult xilinx::AIE::AIERTControl::resetPerfCounters(int col, int row) { + auto tileLoc = XAie_TileLoc(col, row); + // Reset performance counters in all modules + // Try core module counters (if applicable) + for (int counterId = 0; counterId < 4; counterId++) { + // Ignore errors as not all tiles have all counter types + (void)XAie_PerfCounterReset(&aiert->devInst, tileLoc, XAIE_CORE_MOD, + counterId); + } + // Try mem module counters + for (int counterId = 0; counterId < 4; counterId++) { + (void)XAie_PerfCounterReset(&aiert->devInst, tileLoc, XAIE_MEM_MOD, + counterId); + } + // Try PL module counters (for shim tiles) + for (int counterId = 0; counterId < 4; counterId++) { + (void)XAie_PerfCounterReset(&aiert->devInst, tileLoc, XAIE_PL_MOD, + counterId); + } + return success(); +} + LogicalResult xilinx::AIE::AIERTControl::addAieElfs(DeviceOp &targetOp, const StringRef elfPath, bool aieSim) { diff --git a/python/compiler/aiecc/cl_arguments.py b/python/compiler/aiecc/cl_arguments.py index 24657b49f32..7ebbbce84ff 100644 --- a/python/compiler/aiecc/cl_arguments.py +++ b/python/compiler/aiecc/cl_arguments.py @@ -170,6 +170,66 @@ def parse_args(args=None): action="store_true", help="Use packet switched flows when lowering object fifos", ) + parser.add_argument( + "--reset-tiles", + dest="reset_tiles", + default="shim,mem,core", + help="Default comma-separated list of tile types to reset (shim,mem,core). Overridden by specific --reset-*-tiles options.", + ) + parser.add_argument( + "--reset-mode", + dest="reset_mode", + default="ifchangedfinegrained", + help="Default reset mode: never, ifused, ifusedfinegrained, ifchanged, ifchangedfinegrained, always. Overridden by specific --reset-*-mode options.", + ) + parser.add_argument( + "--reset-dmas-tiles", + dest="reset_dmas_tiles", + default=None, + help="Comma-separated list of tile types to reset DMAs for (shim,mem,core). Defaults to --reset-tiles if not specified.", + ) + parser.add_argument( + "--reset-dmas-mode", + dest="reset_dmas_mode", + default=None, + help="When to reset DMAs: never, ifused, ifchanged, always. Defaults to --reset-mode if not specified.", + ) + parser.add_argument( + "--reset-switches-tiles", + dest="reset_switches_tiles", + default=None, + help="Comma-separated list of tile types to reset switches for (shim,mem,core). Defaults to --reset-tiles if not specified.", + ) + parser.add_argument( + "--reset-switches-mode", + dest="reset_switches_mode", + default=None, + help="When to reset switches: never, ifused, ifusedfinegrained, ifchanged, ifchangedfinegrained, always. Defaults to --reset-mode if not specified.", + ) + parser.add_argument( + "--reset-locks-tiles", + dest="reset_locks_tiles", + default=None, + help="Comma-separated list of tile types to reset locks for (mem,core). Defaults to --reset-tiles if not specified.", + ) + parser.add_argument( + "--reset-locks-mode", + dest="reset_locks_mode", + default=None, + help="When to reset locks: never, ifused, ifusedfinegrained, ifchanged, ifchangedfinegrained, always. Defaults to --reset-mode if not specified.", + ) + parser.add_argument( + "--reset-cores-tiles", + dest="reset_cores_tiles", + default=None, + help="Comma-separated list of tile types to reset cores for (core). Defaults to --reset-tiles if not specified.", + ) + parser.add_argument( + "--reset-cores-mode", + dest="reset_cores_mode", + default=None, + help="When to reset cores: never, ifused, ifchanged, always. Defaults to --reset-mode if not specified.", + ) parser.add_argument( "host_args", action="store", @@ -387,8 +447,52 @@ def parse_args(args=None): default="aie.elf", help="Output filename for full ELF (default: aie.elf)", ) + parser.add_argument( + "--expand-load-pdis", + dest="expand_load_pdis", + default=False, + action="store_true", + help="Expand load_pdi operations into explicit device reset and configuration sequences", + ) + parser.add_argument( + "--no-coalesce-write32s", + dest="coalesce_write32s", + default=True, + action="store_const", + const=False, + help="Do not coalesce write32 operations into blockwrites in the NPU instruction sequence", + ) + parser.add_argument( + "--min-writes-to-coalesce", + dest="min_writes_to_coalesce", + default=2, + type=int, + help="Minimum number of consecutive write32 operations required to coalesce into a blockwrite (default: 2)", + ) opts = parser.parse_args(args) + + # Set defaults for specific reset options if not specified + if opts.reset_dmas_tiles is None: + opts.reset_dmas_tiles = opts.reset_tiles + if opts.reset_dmas_mode is None: + opts.reset_dmas_mode = opts.reset_mode + + if opts.reset_switches_tiles is None: + opts.reset_switches_tiles = opts.reset_tiles + if opts.reset_switches_mode is None: + opts.reset_switches_mode = opts.reset_mode + + if opts.reset_locks_tiles is None: + opts.reset_locks_tiles = opts.reset_tiles + if opts.reset_locks_mode is None: + opts.reset_locks_mode = opts.reset_mode + + if opts.reset_cores_tiles is None: + opts.reset_cores_tiles = opts.reset_tiles + if opts.reset_cores_mode is None: + opts.reset_cores_mode = opts.reset_mode + return opts diff --git a/python/compiler/aiecc/main.py b/python/compiler/aiecc/main.py index a4475e32e73..2bfd76408b6 100644 --- a/python/compiler/aiecc/main.py +++ b/python/compiler/aiecc/main.py @@ -147,17 +147,50 @@ def _create_aie_lower_to_llvm_pipeline( AIE_LOWER_TO_LLVM = _create_aie_lower_to_llvm_pipeline + # pipeline to lower and legalize runtime sequence for NPU -NPU_LOWERING_PIPELINE = Pipeline().Nested( - "aie.device", - Pipeline() - .add_pass("aie-materialize-bd-chains") - .add_pass("aie-substitute-shim-dma-allocations") - .add_pass("aie-assign-runtime-sequence-bd-ids") - .add_pass("aie-dma-tasks-to-npu") - .add_pass("aie-dma-to-npu") - .add_pass("aie-lower-set-lock"), -) +def _create_npu_lowering_pipeline( + expand_load_pdis=False, + reset_dmas_tiles="", + reset_dmas_mode="never", + reset_switches_tiles="", + reset_switches_mode="never", + reset_locks_tiles="", + reset_locks_mode="never", + coalesce_write32s=True, + min_writes_to_coalesce=2, +): + device_pipeline_1 = () + pipeline = ( + Pipeline() + .add_pass("aie-materialize-runtime-sequences") + .Nested( + "aie.device", + Pipeline() + .add_pass("aie-materialize-bd-chains") + .add_pass("aie-substitute-shim-dma-allocations") + .add_pass("aie-assign-runtime-sequence-bd-ids") + .add_pass("aie-dma-tasks-to-npu") + .add_pass("aie-dma-to-npu") + .add_pass("aie-lower-set-lock"), + ) + ) + if expand_load_pdis: + pipeline = pipeline.add_pass( + "aie-expand-load-pdi", + reset_dmas_tiles=reset_dmas_tiles, + reset_dmas_mode=reset_dmas_mode, + reset_switches_tiles=reset_switches_tiles, + reset_switches_mode=reset_switches_mode, + reset_locks_tiles=reset_locks_tiles, + reset_locks_mode=reset_locks_mode, + ) + if coalesce_write32s: + pipeline = pipeline.Nested( + "aie.device", + Pipeline().add_pass("aie-coalesce-write32s", min_writes_to_coalesce=min_writes_to_coalesce), + ) + return pipeline async def read_file_async(file_path: str) -> str: @@ -373,11 +406,9 @@ def create_device_id_mapping(devices): return device_to_id -def assign_load_pdi_ids(mlir_module_str, device_to_id_mapping): +def assign_load_pdi_ids(module, device_to_id_mapping): """Transform symbolic aiex.npu.load_pdi references to numeric IDs""" - with Context() as context, Location.unknown(): - module = Module.parse(mlir_module_str) - + with module.context as context, Location.unknown(): for runtime_seq in find_ops( module.operation, lambda o: isinstance(o.operation.opview, aiexdialect.RuntimeSequenceOp), @@ -399,8 +430,6 @@ def assign_load_pdi_ids(mlir_module_str, device_to_id_mapping): IntegerType.get_signless(32, context=context), pdi_id ) - return str(module) - def set_elf_file_for_core(core, path): with InsertionPoint.at_block_terminator( @@ -964,6 +993,10 @@ async def generate_full_elf_config_json( for device_op, device_name in devices: sequences = generate_runtime_sequences_list(device_op) + # Skip devices with no runtime sequences (e.g., @empty device) + if not sequences: + continue + max_arg_count = max( len(seq_op.body.blocks[0].arguments) for seq_op, seq_name in sequences ) @@ -979,9 +1012,11 @@ async def generate_full_elf_config_json( "PDIs": [], } - pdi_id = device_to_id_mapping[device_name] - pdi_filename = self.pdi_file_name(device_name) - kernel_entry["PDIs"].append({"id": pdi_id, "PDI_file": pdi_filename}) + for other_device_name, other_pdi_id in device_to_id_mapping.items(): + pdi_filename = self.pdi_file_name(other_device_name) + kernel_entry["PDIs"].append( + {"id": other_pdi_id, "PDI_file": pdi_filename} + ) for seq_op, seq_name in sequences: insts_filename = self.npu_insts_file_name(device_name, seq_name) @@ -1607,15 +1642,6 @@ async def run_flow(self): sys.exit(1) aie_target, aie_peano_target = aie_targets[0], aie_peano_targets[0] - # Handle full ELF generation configuration - if opts.generate_full_elf: - device_to_id_mapping = create_device_id_mapping(devices) - self.mlir_module_str = assign_load_pdi_ids( - self.mlir_module_str, device_to_id_mapping - ) - transformed_mlir_path = self.prepend_tmp("input_with_pdi_ids.mlir") - await write_file_async(self.mlir_module_str, transformed_mlir_path) - pass_pipeline = INPUT_WITH_ADDRESSES_PIPELINE( opts.alloc_scheme, opts.dynamic_objFifos, @@ -1706,7 +1732,18 @@ async def run_flow(self): input_physical_with_elfs_module = Module.parse( await read_file_async(input_physical_with_elfs) ) - pass_pipeline = NPU_LOWERING_PIPELINE.materialize(module=True) + npu_pipeline = _create_npu_lowering_pipeline( + opts.expand_load_pdis, + reset_dmas_tiles=opts.reset_dmas_tiles, + reset_dmas_mode=opts.reset_dmas_mode, + reset_switches_tiles=opts.reset_switches_tiles, + reset_switches_mode=opts.reset_switches_mode, + reset_locks_tiles=opts.reset_locks_tiles, + reset_locks_mode=opts.reset_locks_mode, + coalesce_write32s=opts.coalesce_write32s, + min_writes_to_coalesce=opts.min_writes_to_coalesce + ) + pass_pipeline = npu_pipeline.materialize(module=True) npu_insts_file = self.prepend_tmp(f"npu_insts.mlir") self.progress_bar.update( task3, advance=1, command=pass_pipeline[0:30] @@ -1717,6 +1754,31 @@ async def run_flow(self): npu_insts_file, self.opts.verbose, ) + + # If expand_load_pdis is enabled, the pass may have created new devices + # (e.g., @empty), so we need to regenerate the device list from the transformed module + if opts.expand_load_pdis: + devices = generate_devices_list(npu_insts_module) + input_physical_with_expanded = self.prepend_tmp( + "input_physical_with_expanded.mlir" + ) + await write_file_async( + str(npu_insts_module), input_physical_with_expanded + ) + # Update both input_physical and input_physical_with_elfs to point to the file with expanded devices + input_physical = input_physical_with_expanded + input_physical_with_elfs = input_physical_with_expanded + + if opts.generate_full_elf: + device_to_id_mapping = create_device_id_mapping(devices) + assign_load_pdi_ids(npu_insts_module, device_to_id_mapping) + transformed_mlir_path = self.prepend_tmp( + "npu_insts_with_pid_ids.mlir" + ) + await write_file_async( + str(npu_insts_module), transformed_mlir_path + ) + self.progress_bar.update(task3, advance=1) # 4.) Generate compilation artifacts for each device diff --git a/test/Conversion/DmaToNpu/dma_to_npu_subview.mlir b/test/Conversion/DmaToNpu/dma_to_npu_subview.mlir new file mode 100644 index 00000000000..971389f0afb --- /dev/null +++ b/test/Conversion/DmaToNpu/dma_to_npu_subview.mlir @@ -0,0 +1,31 @@ +//===- dma_to_npu_subview.mlir ---------------------------------*- MLIR -*-===// +// +// This file is licensed 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 +// +// Copyright (C) 2025, Advanced Micro Devices, Inc. +// +//===----------------------------------------------------------------------===// + +// RUN: aie-opt --aie-dma-to-npu %s | FileCheck %s + +// Test that memref.subview is correctly traced through to adjust address patches + +// CHECK-LABEL: module +// CHECK: aiex.npu.address_patch +// CHECK-SAME: arg_idx = 0 : i32 +// CHECK-SAME: arg_plus = 512 : i32 +module { + aie.device(npu1) { + aie.runtime_sequence(%arg0: memref<1024xi32>) { + // Create a subview at offset 128 elements (512 bytes for i32) + %subview = memref.subview %arg0[128] [256] [1] : memref<1024xi32> to memref<256xi32, strided<[1], offset: 128>> + %reinterpret = memref.reinterpret_cast %subview to offset: [0], sizes: [256], strides: [1] : memref<256xi32, strided<[1], offset: 128>> to memref<256xi32> + // DMA should use the reinterpreted memref, and the pass should trace back to %arg0 with offset + aiex.npu.dma_memcpy_nd (%reinterpret[0, 0, 0, 0][1, 1, 1, 256][0, 0, 0, 1]) { metadata = @buffer, id = 0 : i64 } : memref<256xi32> + } + %tile_0_0 = aie.tile(0, 0) + aie.shim_dma_allocation @buffer (%tile_0_0, S2MM, 0) + } +} diff --git a/test/Passes/materialize-runtime-sequences/error_arg_type_mismatch.mlir b/test/Passes/materialize-runtime-sequences/error_arg_type_mismatch.mlir new file mode 100644 index 00000000000..fe5d6396955 --- /dev/null +++ b/test/Passes/materialize-runtime-sequences/error_arg_type_mismatch.mlir @@ -0,0 +1,35 @@ +//===- error_arg_type_mismatch.mlir ----------------------------*- MLIR -*-===// +// +// This file is licensed 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 +// +// Copyright (C) 2025, Advanced Micro Devices, Inc. +// +//===----------------------------------------------------------------------===// + +// RUN: aie-opt --verify-diagnostics --aie-materialize-runtime-sequences %s + +// Test detection of argument type mismatches when calling runtime sequences + +module { + aie.device(npu2) @main { + %tile00 = aie.tile(0, 0) + + aie.runtime_sequence @main_seq(%arg0: memref<16xi32>) { + aiex.configure @config_wrong_type { + // expected-error @+1 {{argument 0 type mismatch}} + aiex.run @seq_expects_i16(%arg0) : (memref<16xi32>) + } + } + } + + aie.device(npu2) @config_wrong_type { + %tile10 = aie.tile(1, 0) + + // This expects memref<16xi16> but we're passing memref<16xi32> + aie.runtime_sequence @seq_expects_i16(%arg0: memref<16xi16>) { + aiex.npu.write32 {address = 100 : ui32, column = 1 : i32, row = 0 : i32, value = 42 : ui32} + } + } +} diff --git a/test/Passes/materialize-runtime-sequences/error_cyclic.mlir b/test/Passes/materialize-runtime-sequences/error_cyclic.mlir new file mode 100644 index 00000000000..d2e0ce541fb --- /dev/null +++ b/test/Passes/materialize-runtime-sequences/error_cyclic.mlir @@ -0,0 +1,64 @@ +//===- error_cyclic.mlir ----------------------------------------*- MLIR -*-===// +// +// This file is licensed 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 +// +// Copyright (C) 2025, Advanced Micro Devices, Inc. +// +//===----------------------------------------------------------------------===// + +// RUN: aie-opt --verify-diagnostics --split-input-file --aie-materialize-runtime-sequences %s + +// Test detection of self-recursive runtime sequence calls + +module { + aie.device(npu2) @main { + %tile00 = aie.tile(0, 0) + + // expected-error @+1 {{Runtime sequence call graph contains a cycle}} + aie.runtime_sequence @main_seq(%arg0: memref<16xi32>) { + aiex.configure @main { + aiex.run @main_seq(%arg0) : (memref<16xi32>) + } + } + } +} + +// ----- + +// Test detection of cycles through intermediate runtime sequences + +module { + aie.device(npu2) @main { + %tile00 = aie.tile(0, 0) + + // expected-error @+1 {{Runtime sequence call graph contains a cycle}} + aie.runtime_sequence @main_seq(%arg0: memref<16xi32>) { + aiex.configure @config_a { + aiex.run @seq_a(%arg0) : (memref<16xi32>) + } + } + } + + aie.device(npu2) @config_a { + %tile10 = aie.tile(1, 0) + + aie.runtime_sequence @seq_a(%arg0: memref<16xi32>) { + aiex.configure @config_b { + aiex.run @seq_b(%arg0) : (memref<16xi32>) + } + } + } + + aie.device(npu2) @config_b { + %tile20 = aie.tile(2, 0) + + aie.runtime_sequence @seq_b(%arg0: memref<16xi32>) { + // This creates a cycle: main -> a -> b -> a + aiex.configure @config_a { + aiex.run @seq_a(%arg0) : (memref<16xi32>) + } + } + } +} diff --git a/test/Passes/materialize-runtime-sequences/inline_and_load_pdi.mlir b/test/Passes/materialize-runtime-sequences/inline_and_load_pdi.mlir new file mode 100644 index 00000000000..be67b086e09 --- /dev/null +++ b/test/Passes/materialize-runtime-sequences/inline_and_load_pdi.mlir @@ -0,0 +1,44 @@ +//===- inline_and_load_pdi.mlir --------------------------------*- MLIR -*-===// +// +// This file is licensed 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 +// +// Copyright (C) 2025, Advanced Micro Devices, Inc. +// +//===----------------------------------------------------------------------===// + +// RUN: aie-opt --aie-materialize-runtime-sequences %s | FileCheck %s + +// Test that aiex.configure blocks are properly transformed into load_pdi operations +// and that aiex.run calls are inlined + +module { + // CHECK-LABEL: aie.device(npu2) { + aie.device(npu2) @main { + %tile00 = aie.tile(0, 0) + + // CHECK-LABEL: aie.runtime_sequence @main_seq + // CHECK-SAME: (%[[ARG0:.*]]: memref<16xi32>) + aie.runtime_sequence @main_seq(%arg0: memref<16xi32>) { + // CHECK: aiex.npu.load_pdi {device_ref = @config_a} + // CHECK-NOT: aiex.configure + // CHECK-NOT: aiex.run + // CHECK: aiex.npu.write32 {address = 100 : ui32, column = 1 : i32, row = 0 : i32, value = 42 : ui32} + aiex.configure @config_a { + aiex.run @seq_a(%arg0) : (memref<16xi32>) + } + // CHECK: aiex.npu.write32 {address = 200 : ui32, column = 0 : i32, row = 0 : i32, value = 99 : ui32} + aiex.npu.write32 {address = 200 : ui32, column = 0 : i32, row = 0 : i32, value = 99 : ui32} + } + } + + // CHECK: aie.device(npu2) @config_a + aie.device(npu2) @config_a { + %tile10 = aie.tile(1, 0) + + aie.runtime_sequence @seq_a(%arg0: memref<16xi32>) { + aiex.npu.write32 {address = 100 : ui32, column = 1 : i32, row = 0 : i32, value = 42 : ui32} + } + } +} diff --git a/test/Passes/materialize-runtime-sequences/inline_symbols.mlir b/test/Passes/materialize-runtime-sequences/inline_symbols.mlir new file mode 100644 index 00000000000..162b32ebcd7 --- /dev/null +++ b/test/Passes/materialize-runtime-sequences/inline_symbols.mlir @@ -0,0 +1,41 @@ +//===- inline_symbols.mlir -------------------------------------*- MLIR -*-===// +// +// This file is licensed 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 +// +// Copyright (C) 2025, Advanced Micro Devices, Inc. +// +//===----------------------------------------------------------------------===// + +// RUN: aie-opt --aie-materialize-runtime-sequences %s | FileCheck %s + +// Test that symbol definitions (like shim_dma_allocation) are properly preserved +// when inlining runtime sequences from other devices + +module { + aie.device(npu2) @main { + %tile00 = aie.tile(0, 0) + + // CHECK-LABEL: aie.runtime_sequence @main_seq + aie.runtime_sequence @main_seq(%arg0: memref<64xi32>) { + // CHECK: aiex.npu.load_pdi {device_ref = @config_with_symbols} + // CHECK: aiex.npu.dma_memcpy_nd(%[[ARG0:.*]][0, 0, 0, 0][1, 1, 1, 64][0, 0, 0, 1]) + // CHECK-SAME: metadata = @buffer_in + aiex.configure @config_with_symbols { + aiex.run @seq_with_dma(%arg0) : (memref<64xi32>) + } + } + } + + // CHECK: aie.device(npu2) @config_with_symbols + aie.device(npu2) @config_with_symbols { + %tile20 = aie.tile(2, 0) + // CHECK: aie.shim_dma_allocation @buffer_in + aie.shim_dma_allocation @buffer_in(%tile20, S2MM, 0) + + aie.runtime_sequence @seq_with_dma(%arg0: memref<64xi32>) { + aiex.npu.dma_memcpy_nd(%arg0[0, 0, 0, 0][1, 1, 1, 64][0, 0, 0, 1]) { metadata = @buffer_in, id = 0 : i64 } : memref<64xi32> + } + } +} diff --git a/test/bd-chains-and-dma-tasks/dma-tasks-to-npu/good_subview.mlir b/test/bd-chains-and-dma-tasks/dma-tasks-to-npu/good_subview.mlir new file mode 100644 index 00000000000..7318856e20d --- /dev/null +++ b/test/bd-chains-and-dma-tasks/dma-tasks-to-npu/good_subview.mlir @@ -0,0 +1,33 @@ +// +// This file is licensed 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 +// +// Copyright (C) 2025 AMD Inc. + +// RUN: aie-opt --aie-dma-tasks-to-npu %s | FileCheck %s + +// Test that memref.subview is correctly traced through when configuring DMA tasks + +module { + aie.device(npu1) { + %tile_0_0 = aie.tile(0, 0) + + aie.runtime_sequence(%arg0: memref<1024xi16>) { + // Create a subview at offset 64 elements (128 bytes for i16) + // CHECK: aiex.npu.writebd + // CHECK: aiex.npu.address_patch {addr = {{.*}}, arg_idx = 0 : i32, arg_plus = 128 : i32} + %subview = memref.subview %arg0[64] [128] [1] : memref<1024xi16> to memref<128xi16, strided<[1], offset: 64>> + %reinterpret = memref.reinterpret_cast %subview to offset: [0], sizes: [128], strides: [1] : memref<128xi16, strided<[1], offset: 64>> to memref<128xi16> + %t1 = aiex.dma_configure_task(%tile_0_0, MM2S, 0) { + aie.dma_bd(%reinterpret : memref<128xi16>, 0, 128) {bd_id = 7 : i32} + aie.end + } {issue_token = true} + + // CHECK: aiex.npu.push_queue(0, 0, MM2S : 0) {bd_id = 7 : i32, issue_token = true, repeat_count = 0 : i32} + aiex.dma_start_task(%t1) + // CHECK: aiex.npu.sync {channel = 0 : i32, column = 0 : i32, column_num = 1 : i32, direction = 1 : i32, row = 0 : i32, row_num = 1 : i32} + aiex.dma_await_task(%t1) + } + } +} diff --git a/test/npu-xrt/reconfigure_loadpdi/aie.mlir b/test/npu-xrt/reconfigure_loadpdi/aie.mlir new file mode 100644 index 00000000000..7a59536f904 --- /dev/null +++ b/test/npu-xrt/reconfigure_loadpdi/aie.mlir @@ -0,0 +1,144 @@ +// (c) Copyright 2025 Advanced Micro Devices, Inc. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +// This test: +// 1. Calls the @add_two design on the first four i32 elements of the input. +// 2. Calls the @add_two design on the last four i32 elements of the input. +// 3. Calls the @add_three design on i32 elements 4-7 of the input. +// 4. Calls the @add_three_design on the last four i32 elements of the input. +// Elements 8-11 remain untouched by any design. +// This tests, end-to-end: +// - Inlining ('calling') runtime sequences multiple times. +// - Slicing input argument buffers to the runtime sequence in different ways. +// - Reconfiguration between different designs. + +module { + + aie.device(npu2) @main { + + aie.runtime_sequence @sequence(%arg : memref<512xi32>) { + + %c3_i32 = arith.constant 3 : i32 + + aiex.configure @add_two { + %arg1_subview = memref.subview %arg[0] [4] [1] : memref<512xi32> to memref<4xi32, strided<[1], offset: 0>> + %arg1 = memref.reinterpret_cast %arg1_subview to offset: [0], sizes: [4], strides: [1] : memref<4xi32, strided<[1], offset: 0>> to memref<4xi32> + aiex.run @add_two_sequence (%arg1) : (memref<4xi32>) + %arg2_subview = memref.subview %arg[12] [4] [1] : memref<512xi32> to memref<4xi32, strided<[1], offset: 12>> + %arg2 = memref.reinterpret_cast %arg2_subview to offset: [0], sizes: [4], strides: [1] : memref<4xi32, strided<[1], offset: 12>> to memref<4xi32> + aiex.run @add_two_sequence (%arg2) : (memref<4xi32>) + } + + aiex.configure @add_three { + %arg1_subview = memref.subview %arg[4] [4] [1] : memref<512xi32> to memref<4xi32, strided<[1], offset: 4>> + %arg1 = memref.reinterpret_cast %arg1_subview to offset: [0], sizes: [4], strides: [1] : memref<4xi32, strided<[1], offset: 4>> to memref<4xi32> + aiex.run @add_three_sequence (%arg1) : (memref<4xi32>) + %arg2_subview = memref.subview %arg[12] [4] [1] : memref<512xi32> to memref<4xi32, strided<[1], offset: 12>> + %arg2 = memref.reinterpret_cast %arg2_subview to offset: [0], sizes: [4], strides: [1] : memref<4xi32, strided<[1], offset: 12>> to memref<4xi32> + aiex.run @add_three_sequence (%arg2) : (memref<4xi32>) + } + + } + + } + + aie.device(npu2) @add_two { + + %t00 = aie.tile(0, 0) + %t02 = aie.tile(0, 2) + + aie.objectfifo @objfifo_in (%t00, {%t02}, 1 : i32) : !aie.objectfifo> + aie.objectfifo @objfifo_out(%t02, {%t00}, 1 : i32) : !aie.objectfifo> + + aie.core(%t02) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %c2_i32 = arith.constant 2 : i32 + %c8 = arith.constant 8 : index + %c4 = arith.constant 4 : index + %c_intmax = arith.constant 0xFFFFFE : index + + scf.for %niter = %c0 to %c_intmax step %c1 { + %subview_in = aie.objectfifo.acquire @objfifo_in (Consume, 1) : !aie.objectfifosubview> + %subview_out = aie.objectfifo.acquire @objfifo_out(Produce, 1) : !aie.objectfifosubview> + %elem_in = aie.objectfifo.subview.access %subview_in [0] : !aie.objectfifosubview> -> memref<4xi32> + %elem_out = aie.objectfifo.subview.access %subview_out[0] : !aie.objectfifosubview> -> memref<4xi32> + scf.for %i = %c0 to %c4 step %c1 { + %0 = memref.load %elem_in[%i] : memref<4xi32> + %1 = arith.addi %0, %c2_i32 : i32 + memref.store %1, %elem_out[%i] : memref<4xi32> + } + aie.objectfifo.release @objfifo_in (Consume, 1) + aie.objectfifo.release @objfifo_out(Produce, 1) + } + aie.end + } + + aie.runtime_sequence @add_two_sequence(%a : memref<4xi32>) { + + %t_in = aiex.dma_configure_task_for @objfifo_in { + aie.dma_bd(%a : memref<4xi32>, 0, 4) + aie.end + } + %t_out = aiex.dma_configure_task_for @objfifo_out { + aie.dma_bd(%a: memref<4xi32>, 0, 4) + aie.end + } {issue_token = true} + aiex.dma_start_task(%t_in) + aiex.dma_start_task(%t_out) + aiex.dma_await_task(%t_out) + } + + } + + aie.device(npu2) @add_three { + + %t00 = aie.tile(0, 0) + %t02 = aie.tile(0, 2) + + aie.objectfifo @objfifo_in (%t00, {%t02}, 1 : i32) : !aie.objectfifo> + aie.objectfifo @objfifo_out(%t02, {%t00}, 1 : i32) : !aie.objectfifo> + + aie.core(%t02) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %c3_i32 = arith.constant 3 : i32 + %c8 = arith.constant 8 : index + %c4 = arith.constant 4 : index + %c_intmax = arith.constant 0xFFFFFE : index + + scf.for %niter = %c0 to %c_intmax step %c1 { + %subview_in = aie.objectfifo.acquire @objfifo_in (Consume, 1) : !aie.objectfifosubview> + %subview_out = aie.objectfifo.acquire @objfifo_out(Produce, 1) : !aie.objectfifosubview> + %elem_in = aie.objectfifo.subview.access %subview_in [0] : !aie.objectfifosubview> -> memref<4xi32> + %elem_out = aie.objectfifo.subview.access %subview_out[0] : !aie.objectfifosubview> -> memref<4xi32> + scf.for %i = %c0 to %c4 step %c1 { + %0 = memref.load %elem_in[%i] : memref<4xi32> + %1 = arith.addi %0, %c3_i32 : i32 + memref.store %1, %elem_out[%i] : memref<4xi32> + } + aie.objectfifo.release @objfifo_in (Consume, 1) + aie.objectfifo.release @objfifo_out(Produce, 1) + } + aie.end + } + + aie.runtime_sequence @add_three_sequence(%a : memref<4xi32>) { + + %t_in = aiex.dma_configure_task_for @objfifo_in { + aie.dma_bd(%a : memref<4xi32>, 0, 4) + aie.end + } + %t_out = aiex.dma_configure_task_for @objfifo_out { + aie.dma_bd(%a: memref<4xi32>, 0, 4) + aie.end + } {issue_token = true} + aiex.dma_start_task(%t_in) + aiex.dma_start_task(%t_out) + aiex.dma_await_task(%t_out) + } + + } +} diff --git a/test/npu-xrt/reconfigure_loadpdi/run_loadpdi.lit b/test/npu-xrt/reconfigure_loadpdi/run_loadpdi.lit new file mode 100644 index 00000000000..e335c44c7bd --- /dev/null +++ b/test/npu-xrt/reconfigure_loadpdi/run_loadpdi.lit @@ -0,0 +1,8 @@ +// (c) Copyright 2025 Advanced Micro Devices, Inc. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +// REQUIRES: ryzen_ai_npu2, peano +// +// RUN: %PYTHON aiecc.py -v --generate-full-elf --no-xchesscc --no-xbridge %S/aie.mlir +// RUN: clang -o test %S/test.cpp -std=c++17 -lstdc++ %xrt_flags +// RUN: ./test diff --git a/test/npu-xrt/reconfigure_loadpdi/run_write32s.lit b/test/npu-xrt/reconfigure_loadpdi/run_write32s.lit new file mode 100644 index 00000000000..1ee87787843 --- /dev/null +++ b/test/npu-xrt/reconfigure_loadpdi/run_write32s.lit @@ -0,0 +1,8 @@ +// (c) Copyright 2025 Advanced Micro Devices, Inc. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +// REQUIRES: ryzen_ai_npu2, peano +// +// RUN: %PYTHON aiecc.py -v --generate-full-elf --expand-load-pdis --no-xchesscc --no-xbridge %S/aie.mlir +// RUN: clang -o test %S/test.cpp -std=c++17 -lstdc++ %xrt_flags +// RUN: ./test diff --git a/test/npu-xrt/reconfigure_loadpdi/test.cpp b/test/npu-xrt/reconfigure_loadpdi/test.cpp new file mode 100644 index 00000000000..ae4c70fe131 --- /dev/null +++ b/test/npu-xrt/reconfigure_loadpdi/test.cpp @@ -0,0 +1,93 @@ +// (c) Copyright 2025 Advanced Micro Devices, Inc. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#define DTYPE int32_t + +constexpr size_t DATA_COUNT = 16; +constexpr size_t BUF_SIZE = DATA_COUNT * sizeof(DTYPE); + +int main(int argc, const char *argv[]) { + // Set up input data + srand(1726250518); + std::vector vec_in(DATA_COUNT); + for (int i = 0; i < vec_in.size(); i++) { + vec_in[i] = DTYPE(i); + } + + // Set up XRT + unsigned int device_index = 0; + auto device = xrt::device(device_index); + // The name format here is : from the config.json + std::string kernelName = "main:sequence"; + xrt::elf ctx_elf{"aie.elf"}; + xrt::hw_context context = xrt::hw_context(device, ctx_elf); + auto kernel = xrt::ext::kernel(context, kernelName); + xrt::bo bo_inout = xrt::ext::bo{device, BUF_SIZE}; + + // Set up kernel run + char *buf_inout = bo_inout.map(); + memcpy(buf_inout, vec_in.data(), BUF_SIZE); + bo_inout.sync(XCL_BO_SYNC_BO_TO_DEVICE); + auto run = xrt::run(kernel); + run.set_arg(0, bo_inout); + + // Run + auto t_start = std::chrono::high_resolution_clock::now(); + run.start(); + run.wait2(); + auto t_stop = std::chrono::high_resolution_clock::now(); + bo_inout.sync(XCL_BO_SYNC_BO_FROM_DEVICE); + + // Validate output + std::vector vec_out(DATA_COUNT); + std::vector vec_ref(DATA_COUNT); + memcpy(vec_out.data(), buf_inout, BUF_SIZE); + for (int i = 0; i < 4; i++) { + vec_ref[i] = vec_in[i] + 2; + } + for (int i = 4; i < 8; i++) { + vec_ref[i] = vec_in[i] + 3; + } + for (int i = 8; i < 12; i++) { + vec_ref[i] = vec_in[i]; + } + for (int i = 12; i < 16; i++) { + vec_ref[i] = vec_in[i] + 2 + 3; + } + bool outputs_correct = (vec_out == vec_ref); + + // Report results + float time = + std::chrono::duration_cast(t_stop - t_start) + .count(); + std::cout << "Elapsed time for all kernel executions: " << std::fixed + << std::setprecision(0) << std::setw(8) << time << " μs" + << std::endl; + if (outputs_correct) { + std::cout << "PASS!" << std::endl; + } else { + for (int i = 0; i < DATA_COUNT; i++) { + std::cout << "in: " << std::setw(12) << vec_in[i] << ", " + << "out: " << std::setw(12) << vec_out[i] + << ", ref: " << std::setw(12) << vec_ref[i] << std::endl; + } + std::cout << "Fail." << std::endl; + } + + return (outputs_correct ? 0 : 1); +}