Enable reconfiguration through lowerings for aiex.configure and aiex.run (#2775)

Author: andrej

include/aie/Conversion/AIEToConfiguration/AIEToConfiguration.h

@@ -19,6 +19,16 @@ namespace xilinx::AIE {
 
 class DeviceOp;
 
+// --------------------------------------------------------------------------
+// Device configuration
+// --------------------------------------------------------------------------
+
+// an enum to represent the output type of the transaction binary
+enum AIEToConfigurationOutputType {
+  Transaction,
+  ControlPacket,
+};
+
 std::unique_ptr<mlir::OperationPass<xilinx::AIE::DeviceOp>>
 createConvertAIEToTransactionPass();
 
@@ -29,6 +39,58 @@ std::optional<mlir::ModuleOp>
 convertTransactionBinaryToMLIR(mlir::MLIRContext *ctx,
                                std::vector<uint8_t> &binary);
 
+// Generate transaction binary and insert configuration operations at the
+// current insertion point
+mlir::LogicalResult
+generateAndInsertConfigOps(mlir::OpBuilder &builder,
+                           xilinx::AIE::DeviceOp device,
+                           llvm::StringRef clElfDir = "",
+                           AIEToConfigurationOutputType outputType =
+                               AIEToConfigurationOutputType::Transaction);
+
+// --------------------------------------------------------------------------
+// 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 at the current insertion point
+mlir::LogicalResult generateAndInsertResetOps(
+    mlir::OpBuilder &builder, xilinx::AIE::DeviceOp device,
+    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

@@ -2160,6 +2160,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/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 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,8 @@ std::unique_ptr<mlir::OperationPass<AIE::DeviceOp>>
 createAIETxnToControlPacketPass();
 std::unique_ptr<mlir::OperationPass<AIE::DeviceOp>>
 createAIELegalizeControlPacketPass();
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+createAIEExpandLoadPdiPass();
 
 /// 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,24 @@ 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_N`.
+    2. Replace the load_pdi with `load_pdi { device_ref = @empty_N }` -- this causes a device reset.
+    3. 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.
+  }];
+
+  let constructor = "xilinx::AIEX::createAIEExpandLoadPdiPass()";
+  let dependentDialects = [
+    "mlir::memref::MemRefDialect",
+    "xilinx::AIE::AIEDialect",
+    "xilinx::AIEX::AIEXDialect",
+  ];
+}
+
 #endif

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<uint8_t> 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;