Add pass to coalesce write32s/blockwrites

include/aie/Conversion/AIEToConfiguration/AIEToConfiguration.h

@@ -19,6 +19,10 @@ namespace xilinx::AIE {
 
 class DeviceOp;
 
+// --------------------------------------------------------------------------
+// Device configuration
+// --------------------------------------------------------------------------
+
 std::unique_ptr<mlir::OperationPass<xilinx::AIE::DeviceOp>>
 createConvertAIEToTransactionPass();
 
@@ -29,6 +33,55 @@ std::optional<mlir::ModuleOp>
 convertTransactionBinaryToMLIR(mlir::MLIRContext *ctx,
                                std::vector<uint8_t> &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<unsigned>(value) & static_cast<unsigned>(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

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<void(mlir::Value, llvm::StringRef)> setNameFn) {
       ::xilinx::AIE::TileElement::Trait<SwitchboxOp>::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<void(mlir::Value, llvm::StringRef)> setNameFn) {
       ::xilinx::AIE::TileElement::Trait<ShimDMAOp>::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<void(mlir::Value, llvm::StringRef)> setNameFn) {
       ::xilinx::AIE::TileElement::Trait<CoreOp>::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<void(mlir::Value, llvm::StringRef)> setNameFn) {
       ::xilinx::AIE::TileElement::Trait<MemOp>::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<void(mlir::Value, llvm::StringRef)> setNameFn) {
       ::xilinx::AIE::TileElement::Trait<MemTileDMAOp>::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 = [{
   }];

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;

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<uint32_t> 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<SubviewTraceResult> traceSubviewToBlockArgument(Value value);
 }
 } // namespace xilinx

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 = (

include/aie/Dialect/AIEX/Transforms/AIEXPasses.h

@@ -32,6 +32,8 @@ std::unique_ptr<mlir::OperationPass<AIE::DeviceOp>> createAIEDmaToNpuPass();
 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> createAIEXToStandardPass();
 std::unique_ptr<mlir::OperationPass<AIE::DeviceOp>>
 createAIEMaterializeBDChainsPass();
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+createAIEMaterializeRuntimeSequencesPass();
 std::unique_ptr<mlir::OperationPass<AIE::DeviceOp>>
 createAIEAssignRuntimeSequenceBDIDsPass();
 std::unique_ptr<mlir::OperationPass<AIE::DeviceOp>>
@@ -50,6 +52,10 @@ std::unique_ptr<mlir::OperationPass<AIE::DeviceOp>>
 createAIETxnToControlPacketPass();
 std::unique_ptr<mlir::OperationPass<AIE::DeviceOp>>
 createAIELegalizeControlPacketPass();
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+createAIEExpandLoadPdiPass();
+std::unique_ptr<mlir::OperationPass<AIE::DeviceOp>>
+createAIECoalesceWrite32sPass();
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION

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